From acc2251dc6ad6b264c1c822041a67e9d838f6d1a Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Thu, 9 Jul 2026 20:01:39 +0000
Subject: [PATCH 1/6] Bound retained tool result images

---
 packages/agent/src/agent.ts       |  44 ++++
 packages/agent/test/agent.test.ts | 353 +++++++++++++++++++++++++++++-
 2 files changed, 396 insertions(+), 1 deletion(-)

diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index 38ca37b..ecd0367 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -2,6 +2,7 @@ import {
 	Agent,
 	AgentHarness,
 	type AgentHarnessOptions,
+	type AgentMessage,
 	type AgentOptions,
 	type AgentState,
 	type AgentTool,
@@ -36,6 +37,45 @@ import { InternalComputerTranslator, type KernelBrowser } from "./translator/tra
  */
 type CuaRuntimeInput = CuaModelRef | Model<Api>;
 
+const TOOL_RESULT_IMAGE_LIMIT = 4;
+const OMITTED_TOOL_RESULT_IMAGES = "[stale tool-result images omitted]";
+
+function projectToolResultImages(messages: AgentMessage[]): AgentMessage[] {
+	let imageCount = 0;
+	for (const message of messages) {
+		if (message.role === "toolResult") {
+			imageCount += message.content.filter((block) => block.type === "image").length;
+		}
+	}
+	if (imageCount <= TOOL_RESULT_IMAGE_LIMIT) return messages;
+
+	const firstRetainedImage = Math.floor((imageCount - 1) / TOOL_RESULT_IMAGE_LIMIT) * TOOL_RESULT_IMAGE_LIMIT;
+	let imageOrdinal = 0;
+	return messages.map((message) => {
+		if (message.role !== "toolResult") return message;
+
+		let changed = false;
+		let markerInserted = false;
+		const content: typeof message.content = [];
+		for (const block of message.content) {
+			if (block.type !== "image") {
+				content.push(block);
+				continue;
+			}
+			if (imageOrdinal++ >= firstRetainedImage) {
+				content.push(block);
+				continue;
+			}
+			changed = true;
+			if (!markerInserted) {
+				content.push({ type: "text", text: OMITTED_TOOL_RESULT_IMAGES });
+				markerInserted = true;
+			}
+		}
+		return changed ? { ...message, content } : message;
+	});
+}
+
 /**
  * Agent state exposed by {@link CuaAgent}.
  *
@@ -292,6 +332,7 @@ export class CuaAgent extends Agent {
 			onPayload,
 			streamFn,
 			prepareNextTurn,
+			transformContext,
 			extraTools,
 			mode,
 			nativeTool,
@@ -321,6 +362,8 @@ export class CuaAgent extends Agent {
 			...agentOptions,
 			getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey,
 			streamFn: wrappedStreamFn,
+			transformContext: async (messages, signal) =>
+				projectToolResultImages(transformContext ? await transformContext(messages, signal) : messages),
 			initialState: {
 				...initialState,
 				model: runtime.model,
@@ -471,6 +514,7 @@ export class CuaAgentHarness<
 
 		this.runtime = runtime;
 		this.requestedActiveToolNames = activeToolNames;
+		this.on("context", ({ messages }) => ({ messages: projectToolResultImages(messages) }));
 		this.on("before_provider_payload", async ({ model, payload }: { model: Model<Api>; payload: unknown }) => {
 			const onPayload = this.runtime.onPayload();
 			if (!onPayload) return { payload };
diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts
index d7d2273..8ed3805 100644
--- a/packages/agent/test/agent.test.ts
+++ b/packages/agent/test/agent.test.ts
@@ -1,6 +1,6 @@
 import { describe, expect, it, vi } from "vitest";
 import { createAssistantMessageEventStream, type AssistantMessage } from "@earendil-works/pi-ai";
-import { resolveCuaRuntimeSpec } from "@onkernel/cua-ai";
+import { createCuaModels, resolveCuaRuntimeSpec } from "@onkernel/cua-ai";
 import type Kernel from "@onkernel/sdk";
 import {
 	Agent,
@@ -9,6 +9,7 @@ import {
 	CuaAgentHarness,
 	InMemorySessionRepo,
 	NodeExecutionEnv,
+	type AgentMessage,
 	type AgentTool,
 	type KernelBrowser,
 	type StreamFn,
@@ -394,9 +395,359 @@ describe("CuaAgent", () => {
 		expect(fedBack, "second provider request should carry the tool result").toBeDefined();
 		expect(fedBack!.content.some((block) => block.type === "image" && block.mimeType === "image/png")).toBe(true);
 	});
+
+	it("applies screenshot projection after a caller context transform", async () => {
+		const history: AgentMessage[] = [];
+		for (let index = 1; index <= 5; index += 1) {
+			const assistant = createAssistantMessage({ api: "anthropic-messages", provider: "anthropic", id: "test" });
+			assistant.content = [{ type: "toolCall", id: `call-${index}`, name: "custom_image", arguments: {} }];
+			assistant.stopReason = "toolUse";
+			const images = index < 5 ? [index] : [5, 6, 7, 8, 9];
+			history.push(assistant, {
+				role: "toolResult",
+				toolCallId: `call-${index}`,
+				toolName: "custom_image",
+				content: [
+					{ type: "text", text: `result-${index}` },
+					...images.map((image) => ({
+						type: "image" as const,
+						data: Buffer.from(`image-${image}`).toString("base64"),
+						mimeType: "image/png",
+					})),
+				],
+				details: { index },
+				isError: false,
+				timestamp: index,
+			});
+		}
+
+		let providerContext: Parameters<StreamFn>[1] | undefined;
+		const streamFn: StreamFn = (model, context) => {
+			providerContext = context;
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			message.content = [{ type: "text", text: "done" }];
+			stream.push({ type: "start", partial: message });
+			stream.push({ type: "done", reason: "stop", message });
+			stream.end(message);
+			return stream;
+		};
+		const agent = new CuaAgent({
+			browser,
+			client,
+			streamFn,
+			transformContext: async (messages) =>
+				messages.map((message) =>
+					message.role === "user" && Array.isArray(message.content)
+						? {
+								...message,
+								content: message.content.map((block) =>
+									block.type === "text" ? { ...block, text: `caller:${block.text}` } : block,
+								),
+							}
+						: message,
+				),
+			initialState: { model: "anthropic:claude-haiku-4-5", messages: history },
+		});
+
+		await agent.prompt("next");
+
+		const projected = providerContext!.messages;
+		const projectedImages = projected.flatMap((message) => message.content).filter((block) => block.type === "image");
+		expect(projectedImages).toHaveLength(1);
+		expect(projectedImages[0]!.data).toBe(Buffer.from("image-9").toString("base64"));
+		const finalResult = projected.find(
+			(message) => message.role === "toolResult" && message.toolCallId === "call-5",
+		);
+		expect(finalResult?.content).toContainEqual({ type: "text", text: "result-5" });
+		expect(
+			finalResult?.content.filter(
+				(block) => block.type === "text" && block.text === "[stale tool-result images omitted]",
+			),
+		).toHaveLength(1);
+		expect(
+			projected.flatMap((message) => message.content).some((block) => block.type === "text" && block.text === "caller:next"),
+		).toBe(true);
+		expect(
+			projected.filter((message) => message.role === "toolResult").map(({ toolCallId, toolName, details }) => ({
+				toolCallId,
+				toolName,
+				details,
+			})),
+		).toEqual(
+			Array.from({ length: 5 }, (_, index) => ({
+				toolCallId: `call-${index + 1}`,
+				toolName: "custom_image",
+				details: { index: index + 1 },
+			})),
+		);
+		expect(agent.state.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(9);
+	});
 });
 
 describe("CuaAgentHarness", () => {
+	it("bounds screenshots retained across long stateless CUA trajectories", async () => {
+		const contexts: Parameters<StreamFn>[1][] = [];
+		let providerCalls = 0;
+		const models = createCuaModels();
+		const streamFn: StreamFn = (model, context) => {
+			contexts.push({ ...context, messages: structuredClone(context.messages) });
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			if (providerCalls++ % 2 === 0) {
+				message.content = [{ type: "toolCall", id: `screenshot-${providerCalls}`, name: "screenshot", arguments: {} }];
+				message.stopReason = "toolUse";
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", reason: "toolUse", message });
+			} else {
+				message.content = [{ type: "text", text: "continue" }];
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", reason: "stop", message });
+			}
+			stream.end(message);
+			return stream;
+		};
+		models.setProvider({
+			id: "anthropic",
+			name: "stateless test provider",
+			auth: { apiKey: { name: "test key", resolve: async () => ({ auth: { apiKey: "test-key" } }) } },
+			getModels: () => [],
+			stream: streamFn,
+			streamSimple: streamFn,
+		});
+
+		let screenshots = 0;
+		const screenshotClient = {
+			browsers: {
+				computer: { captureScreenshot: async () => new Response(Buffer.from(`screenshot-${++screenshots}`)) },
+			},
+		} as unknown as Kernel;
+		const services = await createHarnessServices();
+		const harness = new CuaAgentHarness({
+			...services,
+			browser,
+			client: screenshotClient,
+			model: "anthropic:claude-haiku-4-5",
+			models,
+		});
+
+		for (let turn = 0; turn < 9; turn += 1) await harness.prompt(`turn ${turn}`);
+
+		const providerImages = contexts.map((context) =>
+			context.messages
+				.flatMap((message) => message.content)
+				.filter((block) => block.type === "image")
+				.map((block) => Buffer.from(block.data, "base64").toString()),
+		);
+		expect(providerImages).toEqual([
+			[],
+			["screenshot-1"],
+			["screenshot-1"],
+			["screenshot-1", "screenshot-2"],
+			["screenshot-1", "screenshot-2"],
+			["screenshot-1", "screenshot-2", "screenshot-3"],
+			["screenshot-1", "screenshot-2", "screenshot-3"],
+			["screenshot-1", "screenshot-2", "screenshot-3", "screenshot-4"],
+			["screenshot-1", "screenshot-2", "screenshot-3", "screenshot-4"],
+			["screenshot-5"],
+			["screenshot-5"],
+			["screenshot-5", "screenshot-6"],
+			["screenshot-5", "screenshot-6"],
+			["screenshot-5", "screenshot-6", "screenshot-7"],
+			["screenshot-5", "screenshot-6", "screenshot-7"],
+			["screenshot-5", "screenshot-6", "screenshot-7", "screenshot-8"],
+			["screenshot-5", "screenshot-6", "screenshot-7", "screenshot-8"],
+			["screenshot-9"],
+		]);
+		expect(Math.max(...providerImages.map((images) => images.length))).toBeLessThanOrEqual(4);
+
+		// Appending within a chunk does not rewrite the already projected request prefix.
+		for (const [earlier, later] of [[9, 10], [9, 11], [11, 12], [11, 13], [13, 14], [13, 15]] as const) {
+			expect(contexts[later]!.messages.slice(0, contexts[earlier]!.messages.length)).toEqual(
+				contexts[earlier]!.messages,
+			);
+		}
+
+		const omissionCounts = contexts.map((context) =>
+			context.messages
+				.flatMap((message) => message.content)
+				.filter((block) => block.type === "text" && block.text === "[stale tool-result images omitted]").length,
+		);
+		expect(omissionCounts[9]).toBe(4);
+		expect(omissionCounts[17]).toBe(8);
+
+		const stored = await services.session.buildContext();
+		expect(
+			stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image"),
+		).toHaveLength(9);
+	});
+
+	it("projects stale images returned by extra tools without removing their result shell", async () => {
+		const contexts: Parameters<StreamFn>[1][] = [];
+		const tools = ["image_extra", "screenshot", "screenshot", "screenshot", "screenshot"];
+		let providerCalls = 0;
+		const models = createCuaModels();
+		const streamFn: StreamFn = (model, context) => {
+			contexts.push({ ...context, messages: structuredClone(context.messages) });
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			if (providerCalls % 2 === 0) {
+				const turn = providerCalls / 2;
+				message.content = [{ type: "toolCall", id: `call-${turn}`, name: tools[turn]!, arguments: {} }];
+				message.stopReason = "toolUse";
+			} else {
+				message.content = [{ type: "text", text: "continue" }];
+			}
+			providerCalls += 1;
+			stream.push({ type: "start", partial: message });
+			stream.push({ type: "done", reason: message.stopReason, message });
+			stream.end(message);
+			return stream;
+		};
+		models.setProvider({
+			id: "anthropic",
+			name: "stateless test provider",
+			auth: { apiKey: { name: "test key", resolve: async () => ({ auth: { apiKey: "test-key" } }) } },
+			getModels: () => [],
+			stream: streamFn,
+			streamSimple: streamFn,
+		});
+		let screenshots = 0;
+		const services = await createHarnessServices();
+		const harness = new CuaAgentHarness({
+			...services,
+			browser,
+			client: {
+				browsers: { computer: { captureScreenshot: async () => new Response(Buffer.from(`builtin-${++screenshots}`)) } },
+			} as unknown as Kernel,
+			model: "anthropic:claude-haiku-4-5",
+			models,
+			extraTools: [{
+				...createCustomTool("image_extra"),
+				async execute() {
+					return {
+						content: [
+							{ type: "text" as const, text: "extra text" },
+							{ type: "image" as const, data: Buffer.from("extra-image").toString("base64"), mimeType: "image/png" },
+						],
+						details: { source: "extra" },
+					};
+				},
+			}],
+		});
+
+		for (let turn = 0; turn < tools.length; turn += 1) await harness.prompt(`turn ${turn}`);
+
+		const finalContext = contexts.at(-1)!;
+		const extraResult = finalContext.messages.find(
+			(message) => message.role === "toolResult" && message.toolCallId === "call-0",
+		);
+		expect(extraResult).toMatchObject({
+			role: "toolResult",
+			toolCallId: "call-0",
+			toolName: "image_extra",
+			details: { source: "extra" },
+			isError: false,
+		});
+		expect(extraResult?.content).toEqual([
+			{ type: "text", text: "extra text" },
+			{ type: "text", text: "[stale tool-result images omitted]" },
+		]);
+		expect(finalContext.messages.some((message) =>
+			message.role === "assistant" && message.content.some((block) => block.type === "toolCall" && block.id === "call-0"),
+		)).toBe(true);
+		expect(
+			finalContext.messages.flatMap((message) => message.content).filter((block) => block.type === "image").map((block) => Buffer.from(block.data, "base64").toString()),
+		).toEqual(["builtin-4"]);
+		const stored = await services.session.buildContext();
+		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
+	});
+
+	it("caps a multi-screenshot computer_batch result by image ordinal", async () => {
+		const contexts: Parameters<StreamFn>[1][] = [];
+		let providerCalls = 0;
+		const models = createCuaModels();
+		const streamFn: StreamFn = (model, context) => {
+			contexts.push({ ...context, messages: structuredClone(context.messages) });
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			if (providerCalls++ === 0) {
+				message.content = [{
+					type: "toolCall",
+					id: "batch-1",
+					name: ANTHROPIC_BATCH_TOOL_NAME,
+					arguments: { actions: [
+						{ type: "screenshot" }, { type: "screenshot" }, { type: "screenshot" },
+						{ type: "screenshot" }, { type: "cursor_position" },
+						{ type: "screenshot" }, { type: "screenshot" },
+					] },
+				}];
+				message.stopReason = "toolUse";
+			} else {
+				message.content = [{ type: "text", text: "done" }];
+			}
+			stream.push({ type: "start", partial: message });
+			stream.push({ type: "done", reason: message.stopReason, message });
+			stream.end(message);
+			return stream;
+		};
+		models.setProvider({
+			id: "anthropic",
+			name: "stateless test provider",
+			auth: { apiKey: { name: "test key", resolve: async () => ({ auth: { apiKey: "test-key" } }) } },
+			getModels: () => [],
+			stream: streamFn,
+			streamSimple: streamFn,
+		});
+		let screenshots = 0;
+		const services = await createHarnessServices();
+		const harness = new CuaAgentHarness({
+			...services,
+			browser,
+			client: { browsers: { computer: {
+				captureScreenshot: async () => new Response(Buffer.from(`batch-${++screenshots}`)),
+				getMousePosition: async () => ({ x: 17, y: 23 }),
+			} } } as unknown as Kernel,
+			model: "anthropic:claude-haiku-4-5",
+			models,
+		});
+
+		await harness.prompt("inspect several frames");
+
+		const followUp = contexts[1]!;
+		const result = followUp.messages.find(
+			(message) => message.role === "toolResult" && message.toolCallId === "batch-1",
+		);
+		expect(result).toMatchObject({
+			role: "toolResult",
+			toolCallId: "batch-1",
+			toolName: ANTHROPIC_BATCH_TOOL_NAME,
+			isError: false,
+			details: {
+				statusText: "Actions executed successfully.",
+				readResults: [
+					{ type: "screenshot", bytes: 7 }, { type: "screenshot", bytes: 7 },
+					{ type: "screenshot", bytes: 7 }, { type: "screenshot", bytes: 7 },
+					{ type: "cursor_position", x: 17, y: 23 },
+					{ type: "screenshot", bytes: 7 }, { type: "screenshot", bytes: 7 },
+				],
+			},
+		});
+		expect(result?.content).toEqual([
+			{ type: "text", text: "[stale tool-result images omitted]" },
+			{ type: "text", text: "cursor_position(): 17,23" },
+			{ type: "image", data: Buffer.from("batch-5").toString("base64"), mimeType: "image/png" },
+			{ type: "image", data: Buffer.from("batch-6").toString("base64"), mimeType: "image/png" },
+		]);
+		expect(result?.content.filter((block) => block.type === "text" && block.text === "[stale tool-result images omitted]")).toHaveLength(1);
+		expect(followUp.messages.some((message) =>
+			message.role === "assistant" && message.content.some((block) => block.type === "toolCall" && block.id === "batch-1"),
+		)).toBe(true);
+		expect(followUp.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(2);
+		const stored = await services.session.buildContext();
+		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(6);
+	});
+
 	it("extends pi AgentHarness and resolves model refs", async () => {
 		const harness = new CuaAgentHarness({
 			...(await createHarnessServices()),

From 8514beec0253dd3124ecc7ca5cf7c22626d87c84 Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Thu, 9 Jul 2026 21:16:20 +0000
Subject: [PATCH 2/6] Configure tool result image replay

---
 packages/agent/README.md          |  24 ++++
 packages/agent/src/agent.ts       |  35 ++++-
 packages/agent/test/agent.test.ts | 217 +++++++++++++++++++++++++++---
 3 files changed, 249 insertions(+), 27 deletions(-)

diff --git a/packages/agent/README.md b/packages/agent/README.md
index 5371d01..d96d162 100644
--- a/packages/agent/README.md
+++ b/packages/agent/README.md
@@ -99,6 +99,30 @@ Both classes mirror pi constructor shapes and behavior, with minimal additions:
 - `extraTools` to add your own pi tools alongside the built-in browser tools
 - `playwright: true` to let the model run Playwright/TypeScript against the
   live browser session
+- `toolResultImageReplayLimit` to control how many recent tool-result images
+  are replayed from local history on each provider request
+
+### Tool-result image replay
+
+Both constructors default `toolResultImageReplayLimit` to `4`. Set a different
+non-negative integer to retain that many recent image blocks, `0` to omit every
+tool-result image from local replay, or `false` to disable this projection.
+The limit is global across built-in tools, custom tools, and multi-image tool
+results. Each changed tool result retains its text, metadata, and ordering and
+gets one `[stale tool-result images omitted]` marker at its first removed image.
+
+This only projects the request context. `CuaAgent.state.messages`, harness
+session entries, and `session.buildContext()` retain the original images.
+Harness `context` handlers are not a transform pipeline: every handler receives
+the original event and the last non-`undefined` result wins. A later handler
+that returns `{ messages }` can therefore replace this built-in projection.
+
+OpenAI and Tzafon normally chain requests with `previous_response_id`, so older
+images may remain in provider-side response state after they leave local
+replay. `CUA_DISABLE_RESPONSE_THREADING=1` disables that chaining process-wide,
+so each request uses locally replayed history instead of prior provider response
+state. It is independent of `toolResultImageReplayLimit` and is not a
+constructor option.
 
 If auth callbacks are omitted, both classes default to CUA env var conventions:
 - OpenAI: `OPENAI_API_KEY`
diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index ecd0367..14be1a9 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -37,19 +37,31 @@ import { InternalComputerTranslator, type KernelBrowser } from "./translator/tra
  */
 type CuaRuntimeInput = CuaModelRef | Model<Api>;
 
-const TOOL_RESULT_IMAGE_LIMIT = 4;
+const DEFAULT_TOOL_RESULT_IMAGE_REPLAY_LIMIT = 4;
 const OMITTED_TOOL_RESULT_IMAGES = "[stale tool-result images omitted]";
 
-function projectToolResultImages(messages: AgentMessage[]): AgentMessage[] {
+export type ToolResultImageReplayLimit = number | false;
+
+function resolveToolResultImageReplayLimit(limit: ToolResultImageReplayLimit | undefined): ToolResultImageReplayLimit {
+	if (limit === undefined) return DEFAULT_TOOL_RESULT_IMAGE_REPLAY_LIMIT;
+	if (limit !== false && (!Number.isFinite(limit) || !Number.isInteger(limit) || limit < 0)) {
+		throw new TypeError("toolResultImageReplayLimit must be a finite non-negative integer or false");
+	}
+	return limit;
+}
+
+function projectToolResultImages(messages: AgentMessage[], limit: ToolResultImageReplayLimit): AgentMessage[] {
+	if (limit === false) return messages;
+
 	let imageCount = 0;
 	for (const message of messages) {
 		if (message.role === "toolResult") {
 			imageCount += message.content.filter((block) => block.type === "image").length;
 		}
 	}
-	if (imageCount <= TOOL_RESULT_IMAGE_LIMIT) return messages;
+	if (imageCount <= limit) return messages;
 
-	const firstRetainedImage = Math.floor((imageCount - 1) / TOOL_RESULT_IMAGE_LIMIT) * TOOL_RESULT_IMAGE_LIMIT;
+	const firstRetainedImage = Math.max(0, imageCount - limit);
 	let imageOrdinal = 0;
 	return messages.map((message) => {
 		if (message.role !== "toolResult") return message;
@@ -118,6 +130,8 @@ export type CuaAgentOptions = Omit<AgentOptions, "initialState"> & {
 	nativeTool?: CuaNativeToolSpec;
 	/** Expose a tool that runs Playwright code against the browser session. */
 	playwright?: boolean;
+	/** Maximum tool-result images replayed from local history per provider request. Defaults to 4; false disables projection. */
+	toolResultImageReplayLimit?: ToolResultImageReplayLimit;
 };
 
 /**
@@ -156,6 +170,8 @@ export type CuaAgentHarnessOptions<
 	playwright?: boolean;
 	/** Optional payload hook composed after the provider-specific CUA payload hook. */
 	onPayload?: SimpleStreamOptions["onPayload"];
+	/** Maximum tool-result images replayed from local history per provider request. Defaults to 4; false disables projection. */
+	toolResultImageReplayLimit?: ToolResultImageReplayLimit;
 };
 
 /**
@@ -337,8 +353,10 @@ export class CuaAgent extends Agent {
 			mode,
 			nativeTool,
 			playwright,
+			toolResultImageReplayLimit,
 			...agentOptions
 		} = options;
+		const imageReplayLimit = resolveToolResultImageReplayLimit(toolResultImageReplayLimit);
 		const runtime = new CuaRuntimeController({
 			browser,
 			client,
@@ -363,7 +381,10 @@ export class CuaAgent extends Agent {
 			getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey,
 			streamFn: wrappedStreamFn,
 			transformContext: async (messages, signal) =>
-				projectToolResultImages(transformContext ? await transformContext(messages, signal) : messages),
+				projectToolResultImages(
+					transformContext ? await transformContext(messages, signal) : messages,
+					imageReplayLimit,
+				),
 			initialState: {
 				...initialState,
 				model: runtime.model,
@@ -489,8 +510,10 @@ export class CuaAgentHarness<
 			systemPrompt,
 			onPayload,
 			activeToolNames,
+			toolResultImageReplayLimit,
 			...harnessOptions
 		} = options;
+		const imageReplayLimit = resolveToolResultImageReplayLimit(toolResultImageReplayLimit);
 		const runtime = new CuaRuntimeController({
 			browser,
 			client,
@@ -514,7 +537,7 @@ export class CuaAgentHarness<
 
 		this.runtime = runtime;
 		this.requestedActiveToolNames = activeToolNames;
-		this.on("context", ({ messages }) => ({ messages: projectToolResultImages(messages) }));
+		this.on("context", ({ messages }) => ({ messages: projectToolResultImages(messages, imageReplayLimit) }));
 		this.on("before_provider_payload", async ({ model, payload }: { model: Model<Api>; payload: unknown }) => {
 			const onPayload = this.runtime.onPayload();
 			if (!onPayload) return { payload };
diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts
index 8ed3805..b449fed 100644
--- a/packages/agent/test/agent.test.ts
+++ b/packages/agent/test/agent.test.ts
@@ -454,8 +454,12 @@ describe("CuaAgent", () => {
 
 		const projected = providerContext!.messages;
 		const projectedImages = projected.flatMap((message) => message.content).filter((block) => block.type === "image");
-		expect(projectedImages).toHaveLength(1);
-		expect(projectedImages[0]!.data).toBe(Buffer.from("image-9").toString("base64"));
+		expect(projectedImages.map((block) => Buffer.from(block.data, "base64").toString())).toEqual([
+			"image-6",
+			"image-7",
+			"image-8",
+			"image-9",
+		]);
 		const finalResult = projected.find(
 			(message) => message.role === "toolResult" && message.toolCallId === "call-5",
 		);
@@ -469,23 +473,199 @@ describe("CuaAgent", () => {
 			projected.flatMap((message) => message.content).some((block) => block.type === "text" && block.text === "caller:next"),
 		).toBe(true);
 		expect(
-			projected.filter((message) => message.role === "toolResult").map(({ toolCallId, toolName, details }) => ({
+			projected.filter((message) => message.role === "toolResult").map(({ toolCallId, toolName, details, isError, timestamp }) => ({
 				toolCallId,
 				toolName,
 				details,
+				isError,
+				timestamp,
 			})),
 		).toEqual(
 			Array.from({ length: 5 }, (_, index) => ({
 				toolCallId: `call-${index + 1}`,
 				toolName: "custom_image",
 				details: { index: index + 1 },
+				isError: false,
+				timestamp: index + 1,
 			})),
 		);
 		expect(agent.state.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(9);
 	});
+
+	it.each([
+		{ limit: 2 as const, expected: ["image-4", "image-5"] },
+		{ limit: 0 as const, expected: [] },
+		{ limit: false as const, expected: ["image-1", "image-2", "image-3", "image-4", "image-5"] },
+	])("supports toolResultImageReplayLimit=$limit", async ({ limit, expected }) => {
+		const history: AgentMessage[] = Array.from({ length: 5 }, (_, index) => ({
+			role: "toolResult" as const,
+			toolCallId: `call-${index + 1}`,
+			toolName: "custom_image",
+			content: [{
+				type: "image" as const,
+				data: Buffer.from(`image-${index + 1}`).toString("base64"),
+				mimeType: "image/png",
+			}],
+			details: {},
+			isError: false,
+			timestamp: index + 1,
+		}));
+		let providerContext: Parameters<StreamFn>[1] | undefined;
+		const streamFn: StreamFn = (model, context) => {
+			providerContext = context;
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			stream.push({ type: "start", partial: message });
+			stream.push({ type: "done", reason: "stop", message });
+			stream.end(message);
+			return stream;
+		};
+		const agent = new CuaAgent({
+			browser,
+			client,
+			streamFn,
+			toolResultImageReplayLimit: limit,
+			initialState: { model: "anthropic:claude-haiku-4-5", messages: history },
+		});
+
+		await agent.prompt("next");
+
+		expect(
+			providerContext!.messages
+				.flatMap((message) => message.content)
+				.filter((block) => block.type === "image")
+				.map((block) => Buffer.from(block.data, "base64").toString()),
+		).toEqual(expected);
+		expect(agent.state.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
+	});
+
+	it("rejects invalid tool-result image replay limits", async () => {
+		for (const limit of [-1, 1.5, Number.NaN, Number.POSITIVE_INFINITY]) {
+			expect(() => new CuaAgent({
+				browser,
+				client,
+				toolResultImageReplayLimit: limit,
+				initialState: { model: "openai:gpt-5.5" },
+			})).toThrow(new TypeError("toolResultImageReplayLimit must be a finite non-negative integer or false"));
+
+			const services = await createHarnessServices();
+			expect(() => new CuaAgentHarness({
+				...services,
+				browser,
+				client,
+				model: "openai:gpt-5.5",
+				toolResultImageReplayLimit: limit,
+			})).toThrow(new TypeError("toolResultImageReplayLimit must be a finite non-negative integer or false"));
+		}
+	});
 });
 
 describe("CuaAgentHarness", () => {
+	it.each([
+		{ limit: 2 as const, expected: ["image-4", "image-5"] },
+		{ limit: 0 as const, expected: [] },
+		{ limit: false as const, expected: ["image-1", "image-2", "image-3", "image-4", "image-5"] },
+	])("supports toolResultImageReplayLimit=$limit", async ({ limit, expected }) => {
+		const services = await createHarnessServices();
+		for (let index = 1; index <= 5; index += 1) {
+			await services.session.appendMessage({
+				role: "toolResult",
+				toolCallId: `call-${index}`,
+				toolName: "custom_image",
+				content: [{ type: "image", data: Buffer.from(`image-${index}`).toString("base64"), mimeType: "image/png" }],
+				details: {},
+				isError: false,
+				timestamp: index,
+			});
+		}
+		let providerContext: Parameters<StreamFn>[1] | undefined;
+		const streamFn: StreamFn = (model, context) => {
+			providerContext = context;
+			const stream = createAssistantMessageEventStream();
+			const message = createAssistantMessage(model);
+			stream.push({ type: "start", partial: message });
+			stream.push({ type: "done", reason: "stop", message });
+			stream.end(message);
+			return stream;
+		};
+		const models = createCuaModels();
+		models.setProvider({
+			id: "anthropic",
+			name: "stateless test provider",
+			auth: { apiKey: { name: "test key", resolve: async () => ({ auth: { apiKey: "test-key" } }) } },
+			getModels: () => [],
+			stream: streamFn,
+			streamSimple: streamFn,
+		});
+		const harness = new CuaAgentHarness({
+			...services,
+			browser,
+			client,
+			model: "anthropic:claude-haiku-4-5",
+			models,
+			toolResultImageReplayLimit: limit,
+		});
+
+		await harness.prompt("next");
+
+		expect(
+			providerContext!.messages
+				.flatMap((message) => message.content)
+				.filter((block) => block.type === "image")
+				.map((block) => Buffer.from(block.data, "base64").toString()),
+		).toEqual(expected);
+		const stored = await services.session.buildContext();
+		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
+	});
+
+	it("uses last-result-wins context hook semantics", async () => {
+		const run = async (replace: boolean) => {
+			const services = await createHarnessServices();
+			for (let index = 1; index <= 5; index += 1) {
+				await services.session.appendMessage({
+					role: "toolResult",
+					toolCallId: `call-${index}`,
+					toolName: "custom_image",
+					content: [{ type: "image", data: `${index}`, mimeType: "image/png" }],
+					details: {},
+					isError: false,
+					timestamp: index,
+				});
+			}
+			let providerImageCount = -1;
+			let observedImageCount = -1;
+			const streamFn: StreamFn = (model, context) => {
+				providerImageCount = context.messages.flatMap((message) => message.content).filter((block) => block.type === "image").length;
+				const stream = createAssistantMessageEventStream();
+				const message = createAssistantMessage(model);
+				stream.push({ type: "start", partial: message });
+				stream.push({ type: "done", reason: "stop", message });
+				stream.end(message);
+				return stream;
+			};
+			const models = createCuaModels();
+			models.setProvider({
+				id: "anthropic",
+				name: "stateless test provider",
+				auth: { apiKey: { name: "test key", resolve: async () => ({ auth: { apiKey: "test-key" } }) } },
+				getModels: () => [],
+				stream: streamFn,
+				streamSimple: streamFn,
+			});
+			const harness = new CuaAgentHarness({ ...services, browser, client, model: "anthropic:claude-haiku-4-5", models });
+			harness.on("context", ({ messages }) => {
+				observedImageCount = messages.flatMap((message) => message.content).filter((block) => block.type === "image").length;
+				return replace ? { messages } : undefined;
+			});
+
+			await harness.prompt("next");
+			return { observedImageCount, providerImageCount };
+		};
+
+		await expect(run(false)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 4 });
+		await expect(run(true)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 5 });
+	});
+
 	it("bounds screenshots retained across long stateless CUA trajectories", async () => {
 		const contexts: Parameters<StreamFn>[1][] = [];
 		let providerCalls = 0;
@@ -549,32 +729,25 @@ describe("CuaAgentHarness", () => {
 			["screenshot-1", "screenshot-2", "screenshot-3"],
 			["screenshot-1", "screenshot-2", "screenshot-3", "screenshot-4"],
 			["screenshot-1", "screenshot-2", "screenshot-3", "screenshot-4"],
-			["screenshot-5"],
-			["screenshot-5"],
-			["screenshot-5", "screenshot-6"],
-			["screenshot-5", "screenshot-6"],
-			["screenshot-5", "screenshot-6", "screenshot-7"],
-			["screenshot-5", "screenshot-6", "screenshot-7"],
+			["screenshot-2", "screenshot-3", "screenshot-4", "screenshot-5"],
+			["screenshot-2", "screenshot-3", "screenshot-4", "screenshot-5"],
+			["screenshot-3", "screenshot-4", "screenshot-5", "screenshot-6"],
+			["screenshot-3", "screenshot-4", "screenshot-5", "screenshot-6"],
+			["screenshot-4", "screenshot-5", "screenshot-6", "screenshot-7"],
+			["screenshot-4", "screenshot-5", "screenshot-6", "screenshot-7"],
 			["screenshot-5", "screenshot-6", "screenshot-7", "screenshot-8"],
 			["screenshot-5", "screenshot-6", "screenshot-7", "screenshot-8"],
-			["screenshot-9"],
+			["screenshot-6", "screenshot-7", "screenshot-8", "screenshot-9"],
 		]);
 		expect(Math.max(...providerImages.map((images) => images.length))).toBeLessThanOrEqual(4);
 
-		// Appending within a chunk does not rewrite the already projected request prefix.
-		for (const [earlier, later] of [[9, 10], [9, 11], [11, 12], [11, 13], [13, 14], [13, 15]] as const) {
-			expect(contexts[later]!.messages.slice(0, contexts[earlier]!.messages.length)).toEqual(
-				contexts[earlier]!.messages,
-			);
-		}
-
 		const omissionCounts = contexts.map((context) =>
 			context.messages
 				.flatMap((message) => message.content)
 				.filter((block) => block.type === "text" && block.text === "[stale tool-result images omitted]").length,
 		);
-		expect(omissionCounts[9]).toBe(4);
-		expect(omissionCounts[17]).toBe(8);
+		expect(omissionCounts[9]).toBe(1);
+		expect(omissionCounts[17]).toBe(5);
 
 		const stored = await services.session.buildContext();
 		expect(
@@ -658,7 +831,7 @@ describe("CuaAgentHarness", () => {
 		)).toBe(true);
 		expect(
 			finalContext.messages.flatMap((message) => message.content).filter((block) => block.type === "image").map((block) => Buffer.from(block.data, "base64").toString()),
-		).toEqual(["builtin-4"]);
+		).toEqual(["builtin-1", "builtin-2", "builtin-3", "builtin-4"]);
 		const stored = await services.session.buildContext();
 		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
 	});
@@ -735,6 +908,8 @@ describe("CuaAgentHarness", () => {
 		});
 		expect(result?.content).toEqual([
 			{ type: "text", text: "[stale tool-result images omitted]" },
+			{ type: "image", data: Buffer.from("batch-3").toString("base64"), mimeType: "image/png" },
+			{ type: "image", data: Buffer.from("batch-4").toString("base64"), mimeType: "image/png" },
 			{ type: "text", text: "cursor_position(): 17,23" },
 			{ type: "image", data: Buffer.from("batch-5").toString("base64"), mimeType: "image/png" },
 			{ type: "image", data: Buffer.from("batch-6").toString("base64"), mimeType: "image/png" },
@@ -743,7 +918,7 @@ describe("CuaAgentHarness", () => {
 		expect(followUp.messages.some((message) =>
 			message.role === "assistant" && message.content.some((block) => block.type === "toolCall" && block.id === "batch-1"),
 		)).toBe(true);
-		expect(followUp.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(2);
+		expect(followUp.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(4);
 		const stored = await services.session.buildContext();
 		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(6);
 	});

From de2e34f2efe5ea568c33f382c0ab5aba6a8057cc Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Thu, 9 Jul 2026 22:46:12 +0000
Subject: [PATCH 3/6] Document image replay limit type

---
 packages/agent/src/agent.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index 14be1a9..9225fc9 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -40,6 +40,10 @@ type CuaRuntimeInput = CuaModelRef | Model<Api>;
 const DEFAULT_TOOL_RESULT_IMAGE_REPLAY_LIMIT = 4;
 const OMITTED_TOOL_RESULT_IMAGES = "[stale tool-result images omitted]";
 
+/**
+ * Maximum number of tool-result images replayed from local history, or `false`
+ * to keep the original message history unchanged.
+ */
 export type ToolResultImageReplayLimit = number | false;
 
 function resolveToolResultImageReplayLimit(limit: ToolResultImageReplayLimit | undefined): ToolResultImageReplayLimit {

From fc7060e551948ecbd20cceebaaedec183be9241e Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Fri, 10 Jul 2026 14:55:44 +0000
Subject: [PATCH 4/6] Apply response threading policy to all model calls

---
 packages/agent/src/agent.ts       | 16 +++++++++++++---
 packages/agent/test/agent.test.ts | 21 ++++++++++++++++++---
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index 62db8a0..44e7108 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -350,9 +350,19 @@ function resolveResponseThreading(responseThreading: boolean | undefined): boole
 	return responseThreading ?? true;
 }
 
+function disableResponseThreading<TOptions extends object | undefined>(
+	options: TOptions,
+): TOptions & { disableResponseThreading: true } {
+	return { ...options, disableResponseThreading: true };
+}
+
 function withoutResponseThreading(models: Models): Models {
+	const stream: Models["stream"] = (model, context, options) =>
+		models.stream(model, context, disableResponseThreading(options));
+	const complete: Models["complete"] = (model, context, options) =>
+		models.complete(model, context, disableResponseThreading(options));
 	const streamSimple: Models["streamSimple"] = (model, context, options) =>
-		models.streamSimple(model, context, { ...options, disableResponseThreading: true } as CuaSimpleStreamOptions);
+		models.streamSimple(model, context, disableResponseThreading(options));
 	return {
 		getProviders: () => models.getProviders(),
 		getProvider: (id) => models.getProvider(id),
@@ -360,8 +370,8 @@ function withoutResponseThreading(models: Models): Models {
 		getModel: (provider, id) => models.getModel(provider, id),
 		refresh: (provider) => models.refresh(provider),
 		getAuth: (model) => models.getAuth(model),
-		stream: (model, context, options) => models.stream(model, context, options),
-		complete: (model, context, options) => models.complete(model, context, options),
+		stream,
+		complete,
 		streamSimple,
 		completeSimple: (model, context, options) => streamSimple(model, context, options).result(),
 	};
diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts
index 220cd6a..f2bd2d9 100644
--- a/packages/agent/test/agent.test.ts
+++ b/packages/agent/test/agent.test.ts
@@ -718,7 +718,7 @@ describe("CuaAgentHarness", () => {
 		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
 	});
 
-	it.each([true, false])("passes responseThreading=$enabled through harness models", async (enabled) => {
+	it.each([true, false])("passes responseThreading=$enabled through all harness model methods", async (enabled) => {
 		const services = await createHarnessServices();
 		let streamOptions: CuaSimpleStreamOptions | undefined;
 		const streamFn: StreamFn = (model, _context, options) => {
@@ -748,9 +748,24 @@ describe("CuaAgentHarness", () => {
 			responseThreading: enabled,
 		});
 
-		await harness.prompt("next");
+		const expected = enabled ? undefined : true;
+		const context = { messages: [] };
+		const model = harness.getModel();
+
+		await harness.models.stream(model, context).result();
+		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		streamOptions = undefined;
+
+		await harness.models.complete(model, context);
+		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		streamOptions = undefined;
+
+		await harness.models.streamSimple(model, context).result();
+		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		streamOptions = undefined;
 
-		expect(streamOptions?.disableResponseThreading).toBe(enabled ? undefined : true);
+		await harness.models.completeSimple(model, context);
+		expect(streamOptions?.disableResponseThreading).toBe(expected);
 	});
 
 	it("uses last-result-wins context hook semantics", async () => {

From a2bf729993b551e64739ed7dc525f00fd1444b32 Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Fri, 10 Jul 2026 15:25:06 +0000
Subject: [PATCH 5/6] Compose harness context hooks with image limits

---
 packages/agent/README.md          | 12 ++++----
 packages/agent/src/agent.ts       | 46 +++++++++++++++++++++----------
 packages/agent/test/agent.test.ts | 39 +++++++++++++++++++++-----
 3 files changed, 70 insertions(+), 27 deletions(-)

diff --git a/packages/agent/README.md b/packages/agent/README.md
index a34772f..3fed6c5 100644
--- a/packages/agent/README.md
+++ b/packages/agent/README.md
@@ -137,13 +137,15 @@ pi provides several context-management layers with different lifetimes:
   `transformContext`, it runs first and the image limit is applied to its
   result.
 - `AgentHarness` exposes the same request-time stage through its `context`
-  hook. `CuaAgentHarness` registers the image limit there. pi's harness hooks
-  use last-result-wins semantics, so a later `context` handler that returns
-  `{ messages }` replaces the built-in projection rather than composing with
-  it.
+  hook. pi selects their final result first, then `CuaAgentHarness` applies
+  the image limit immediately before `Models.streamSimple`.
+  User handlers therefore compose with the built-in limit instead of replacing
+  it. Set `toolResultImageReplayLimit: false` if a handler should own image
+  filtering completely.
 - A harness `Session` can use `entryTransforms` and `entryProjectors` while
   building model context from stored session entries. These run before the
-  request-time `context` hook and do not rewrite stored entries.
+  request-time `context` hooks; the image limit runs afterward. None of these
+  projections rewrite stored entries.
 - `harness.compact()` is the persistent, model-generated option for reducing a
   long transcript. It writes a summary entry that future session contexts use.
   It is broader than the image limit and is not called automatically by this
diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index 44e7108..7a7eca6 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -60,7 +60,10 @@ function resolveToolResultImageReplayLimit(limit: ToolResultImageReplayLimit | u
 	return limit;
 }
 
-function projectToolResultImages(messages: AgentMessage[], limit: ToolResultImageReplayLimit): AgentMessage[] {
+function projectToolResultImages<TMessage extends AgentMessage>(
+	messages: TMessage[],
+	limit: ToolResultImageReplayLimit,
+): TMessage[] {
 	if (limit === false) return messages;
 
 	let imageCount = 0;
@@ -94,7 +97,7 @@ function projectToolResultImages(messages: AgentMessage[], limit: ToolResultImag
 				markerInserted = true;
 			}
 		}
-		return changed ? { ...message, content } : message;
+		return changed ? ({ ...message, content } as TMessage) : message;
 	});
 }
 
@@ -350,19 +353,34 @@ function resolveResponseThreading(responseThreading: boolean | undefined): boole
 	return responseThreading ?? true;
 }
 
-function disableResponseThreading<TOptions extends object | undefined>(
+function withResponseThreading<TOptions extends object | undefined>(
 	options: TOptions,
-): TOptions & { disableResponseThreading: true } {
-	return { ...options, disableResponseThreading: true };
+	enabled: boolean,
+): TOptions {
+	return enabled ? options : { ...options, disableResponseThreading: true };
 }
 
-function withoutResponseThreading(models: Models): Models {
+function withContextManagement(
+	models: Models,
+	imageReplayLimit: ToolResultImageReplayLimit,
+	responseThreading: boolean,
+): Models {
+	if (imageReplayLimit === false && responseThreading) return models;
+
 	const stream: Models["stream"] = (model, context, options) =>
-		models.stream(model, context, disableResponseThreading(options));
+		models.stream(model, context, withResponseThreading(options, responseThreading));
 	const complete: Models["complete"] = (model, context, options) =>
-		models.complete(model, context, disableResponseThreading(options));
+		models.complete(model, context, withResponseThreading(options, responseThreading));
 	const streamSimple: Models["streamSimple"] = (model, context, options) =>
-		models.streamSimple(model, context, disableResponseThreading(options));
+		models.streamSimple(
+			model,
+			imageReplayLimit === false
+				? context
+				: { ...context, messages: projectToolResultImages(context.messages, imageReplayLimit) },
+			withResponseThreading(options, responseThreading),
+		);
+	const completeSimple: Models["completeSimple"] = (model, context, options) =>
+		models.completeSimple(model, context, withResponseThreading(options, responseThreading));
 	return {
 		getProviders: () => models.getProviders(),
 		getProvider: (id) => models.getProvider(id),
@@ -373,7 +391,7 @@ function withoutResponseThreading(models: Models): Models {
 		stream,
 		complete,
 		streamSimple,
-		completeSimple: (model, context, options) => streamSimple(model, context, options).result(),
+		completeSimple,
 	};
 }
 
@@ -588,17 +606,16 @@ export class CuaAgentHarness<
 			onPayload,
 		});
 		const resolvedTools = runtime.tools();
-		const baseModels = models ?? cuaModels();
-		const contextModels = useResponseThreading ? baseModels : withoutResponseThreading(baseModels);
 		const retryingModels = withProviderRetryModels(
-			contextModels,
+			models ?? cuaModels(),
 			resolveProviderRetryPolicy(retry),
 		);
+		const contextModels = withContextManagement(retryingModels, imageReplayLimit, useResponseThreading);
 
 		super({
 			...harnessOptions,
 			model: runtime.model,
-			models: retryingModels,
+			models: contextModels,
 			tools: resolvedTools,
 			systemPrompt: systemPrompt ?? (() => runtime.systemPrompt),
 			activeToolNames: activeToolNames ?? resolvedTools.map((tool) => tool.name),
@@ -606,7 +623,6 @@ export class CuaAgentHarness<
 
 		this.runtime = runtime;
 		this.requestedActiveToolNames = activeToolNames;
-		this.on("context", ({ messages }) => ({ messages: projectToolResultImages(messages, imageReplayLimit) }));
 		this.on("before_provider_payload", async ({ model, payload }: { model: Model<Api>; payload: unknown }) => {
 			const onPayload = this.runtime.onPayload();
 			if (!onPayload) return { payload };
diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts
index f2bd2d9..72b5103 100644
--- a/packages/agent/test/agent.test.ts
+++ b/packages/agent/test/agent.test.ts
@@ -1,6 +1,11 @@
 import { describe, expect, it, vi } from "vitest";
 import { createAssistantMessageEventStream, type AssistantMessage } from "@earendil-works/pi-ai";
-import { createCuaModels, type CuaSimpleStreamOptions, resolveCuaRuntimeSpec } from "@onkernel/cua-ai";
+import {
+	type Context,
+	createCuaModels,
+	type CuaSimpleStreamOptions,
+	resolveCuaRuntimeSpec,
+} from "@onkernel/cua-ai";
 import type Kernel from "@onkernel/sdk";
 import {
 	Agent,
@@ -721,8 +726,10 @@ describe("CuaAgentHarness", () => {
 	it.each([true, false])("passes responseThreading=$enabled through all harness model methods", async (enabled) => {
 		const services = await createHarnessServices();
 		let streamOptions: CuaSimpleStreamOptions | undefined;
-		const streamFn: StreamFn = (model, _context, options) => {
+		let providerImageCount = -1;
+		const streamFn: StreamFn = (model, context, options) => {
 			streamOptions = options as CuaSimpleStreamOptions;
+			providerImageCount = context.messages.flatMap((message) => message.content).filter((block) => block.type === "image").length;
 			const stream = createAssistantMessageEventStream();
 			const message = createAssistantMessage(model);
 			stream.push({ type: "start", partial: message });
@@ -746,29 +753,43 @@ describe("CuaAgentHarness", () => {
 			model: "openai:gpt-5.5",
 			models,
 			responseThreading: enabled,
+			retry: { enabled: true, maxRetries: 1, baseDelayMs: 0 },
 		});
 
 		const expected = enabled ? undefined : true;
-		const context = { messages: [] };
+		const context: Context = {
+			messages: Array.from({ length: 5 }, (_, index) => ({
+				role: "toolResult" as const,
+				toolCallId: `call-${index}`,
+				toolName: "screenshot",
+				content: [{ type: "image" as const, data: `${index}`, mimeType: "image/png" }],
+				isError: false,
+				timestamp: index,
+			})),
+		};
 		const model = harness.getModel();
 
 		await harness.models.stream(model, context).result();
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		expect(providerImageCount).toBe(5);
 		streamOptions = undefined;
 
 		await harness.models.complete(model, context);
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		expect(providerImageCount).toBe(5);
 		streamOptions = undefined;
 
 		await harness.models.streamSimple(model, context).result();
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		expect(providerImageCount).toBe(4);
 		streamOptions = undefined;
 
 		await harness.models.completeSimple(model, context);
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
+		expect(providerImageCount).toBe(5);
 	});
 
-	it("uses last-result-wins context hook semantics", async () => {
+	it("applies image projection after user context hooks", async () => {
 		const run = async (replace: boolean) => {
 			const services = await createHarnessServices();
 			for (let index = 1; index <= 5; index += 1) {
@@ -783,9 +804,13 @@ describe("CuaAgentHarness", () => {
 				});
 			}
 			let providerImageCount = -1;
+			let providerOmissionCount = -1;
 			let observedImageCount = -1;
 			const streamFn: StreamFn = (model, context) => {
 				providerImageCount = context.messages.flatMap((message) => message.content).filter((block) => block.type === "image").length;
+				providerOmissionCount = context.messages
+					.flatMap((message) => message.content)
+					.filter((block) => block.type === "text" && block.text === "[stale tool-result images omitted]").length;
 				const stream = createAssistantMessageEventStream();
 				const message = createAssistantMessage(model);
 				stream.push({ type: "start", partial: message });
@@ -809,11 +834,11 @@ describe("CuaAgentHarness", () => {
 			});
 
 			await harness.prompt("next");
-			return { observedImageCount, providerImageCount };
+			return { observedImageCount, providerImageCount, providerOmissionCount };
 		};
 
-		await expect(run(false)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 4 });
-		await expect(run(true)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 5 });
+		await expect(run(false)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 4, providerOmissionCount: 1 });
+		await expect(run(true)).resolves.toEqual({ observedImageCount: 5, providerImageCount: 4, providerOmissionCount: 1 });
 	});
 
 	it("bounds screenshots retained across long stateless CUA trajectories", async () => {

From 51eb722739bf0044265cd06f30541188f5be9662 Mon Sep 17 00:00:00 2001
From: rgarcia <72655+rgarcia@users.noreply.github.com>
Date: Fri, 10 Jul 2026 15:30:54 +0000
Subject: [PATCH 6/6] Apply image limits to all model calls

---
 packages/agent/README.md          |  2 +-
 packages/agent/src/agent.ts       | 28 ++++++++++++++++++++++------
 packages/agent/test/agent.test.ts |  8 ++++----
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/packages/agent/README.md b/packages/agent/README.md
index 3fed6c5..bc8205b 100644
--- a/packages/agent/README.md
+++ b/packages/agent/README.md
@@ -138,7 +138,7 @@ pi provides several context-management layers with different lifetimes:
   result.
 - `AgentHarness` exposes the same request-time stage through its `context`
   hook. pi selects their final result first, then `CuaAgentHarness` applies
-  the image limit immediately before `Models.streamSimple`.
+  the image limit at the `Models` boundary before the provider call.
   User handlers therefore compose with the built-in limit instead of replacing
   it. Set `toolResultImageReplayLimit: false` if a handler should own image
   filtering completely.
diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index 7a7eca6..6b18861 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -12,6 +12,7 @@ import {
 } from "@earendil-works/pi-agent-core";
 import {
 	type Api,
+	type Context,
 	CUA_NAVIGATION_TOOL_NAME,
 	CUA_PLAYWRIGHT_TOOL_NAME,
 	cuaModels,
@@ -360,6 +361,11 @@ function withResponseThreading<TOptions extends object | undefined>(
 	return enabled ? options : { ...options, disableResponseThreading: true };
 }
 
+function projectModelContext(context: Context, imageReplayLimit: ToolResultImageReplayLimit): Context {
+	const messages = projectToolResultImages(context.messages, imageReplayLimit);
+	return messages === context.messages ? context : { ...context, messages };
+}
+
 function withContextManagement(
 	models: Models,
 	imageReplayLimit: ToolResultImageReplayLimit,
@@ -368,19 +374,29 @@ function withContextManagement(
 	if (imageReplayLimit === false && responseThreading) return models;
 
 	const stream: Models["stream"] = (model, context, options) =>
-		models.stream(model, context, withResponseThreading(options, responseThreading));
+		models.stream(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
 	const complete: Models["complete"] = (model, context, options) =>
-		models.complete(model, context, withResponseThreading(options, responseThreading));
+		models.complete(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
 	const streamSimple: Models["streamSimple"] = (model, context, options) =>
 		models.streamSimple(
 			model,
-			imageReplayLimit === false
-				? context
-				: { ...context, messages: projectToolResultImages(context.messages, imageReplayLimit) },
+			projectModelContext(context, imageReplayLimit),
 			withResponseThreading(options, responseThreading),
 		);
 	const completeSimple: Models["completeSimple"] = (model, context, options) =>
-		models.completeSimple(model, context, withResponseThreading(options, responseThreading));
+		models.completeSimple(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
 	return {
 		getProviders: () => models.getProviders(),
 		getProvider: (id) => models.getProvider(id),
diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts
index 72b5103..cc02565 100644
--- a/packages/agent/test/agent.test.ts
+++ b/packages/agent/test/agent.test.ts
@@ -723,7 +723,7 @@ describe("CuaAgentHarness", () => {
 		expect(stored.messages.flatMap((message) => message.content).filter((block) => block.type === "image")).toHaveLength(5);
 	});
 
-	it.each([true, false])("passes responseThreading=$enabled through all harness model methods", async (enabled) => {
+	it.each([true, false])("applies context management through all harness model methods", async (enabled) => {
 		const services = await createHarnessServices();
 		let streamOptions: CuaSimpleStreamOptions | undefined;
 		let providerImageCount = -1;
@@ -771,12 +771,12 @@ describe("CuaAgentHarness", () => {
 
 		await harness.models.stream(model, context).result();
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
-		expect(providerImageCount).toBe(5);
+		expect(providerImageCount).toBe(4);
 		streamOptions = undefined;
 
 		await harness.models.complete(model, context);
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
-		expect(providerImageCount).toBe(5);
+		expect(providerImageCount).toBe(4);
 		streamOptions = undefined;
 
 		await harness.models.streamSimple(model, context).result();
@@ -786,7 +786,7 @@ describe("CuaAgentHarness", () => {
 
 		await harness.models.completeSimple(model, context);
 		expect(streamOptions?.disableResponseThreading).toBe(expected);
-		expect(providerImageCount).toBe(5);
+		expect(providerImageCount).toBe(4);
 	});
 
 	it("applies image projection after user context hooks", async () => {