Configure CUA image context and response threading

packages/agent/README.md

@@ -99,6 +99,68 @@ 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 included in each model request
+- `responseThreading` to choose whether OpenAI and Tzafon continue from
+  provider-stored response state or receive the full request context each time
+
+### Context management
+
+Both constructors accept the same request-context options:
+
+```ts
+const agent = new CuaAgent({
+  // ...browser, client, and initialState
+  toolResultImageReplayLimit: 4,
+  responseThreading: true,
+});
+```
+
+`toolResultImageReplayLimit` defaults to `4`. A non-negative integer keeps that
+many of the newest image blocks from tool results in each model request. `0`
+removes all tool-result images from the request, while `false` disables this
+filter and includes every tool-result image. Text, metadata, and message order
+are preserved; a changed tool result gets one
+`[stale tool-result images omitted]` marker where its first removed image was.
+The limit applies across built-in tools, custom tools, and multi-image results.
+
+This is a request-time view of the transcript, not a history mutation.
+`CuaAgent.state.messages`, harness session entries, and
+`session.buildContext()` retain the original images.
+
+#### How this uses pi context primitives
+
+pi provides several context-management layers with different lifetimes:
+
+- `Agent.transformContext` changes the messages passed to one model request.
+  `CuaAgent` uses this hook for the image limit. If you provide your own
+  `transformContext`, it runs first and the image limit is applied to its
+  result.
+- `AgentHarness` exposes the same request-time stage through its `context`
+  hook. pi selects their final result first, then `CuaAgentHarness` applies
+  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.
+- A harness `Session` can use `entryTransforms` and `entryProjectors` while
+  building model context from stored session entries. These run before the
+  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
+  package.
+
+`responseThreading` defaults to `true`. For OpenAI and Tzafon Responses APIs,
+that means each request can continue from `previous_response_id`; the provider
+keeps earlier response state and CUA sends only messages added since that
+response. Because the provider retains that state, removing an older image from
+the current request view does not remove it from an existing provider thread.
+
+Set `responseThreading: false` when every request should be built entirely from
+the current pi context instead. In that mode no `previous_response_id` is sent,
+and `toolResultImageReplayLimit` determines which tool-result images are
+included in the full request. Other providers do not use this option.
 
 If auth callbacks are omitted, both classes default to CUA env var conventions:
 - OpenAI: `OPENAI_API_KEY`

packages/agent/src/agent.ts

@@ -2,6 +2,7 @@ import {
 	Agent,
 	AgentHarness,
 	type AgentHarnessOptions,
+	type AgentMessage,
 	type AgentOptions,
 	type AgentState,
 	type AgentTool,
@@ -11,6 +12,7 @@ import {
 } from "@earendil-works/pi-agent-core";
 import {
 	type Api,
+	type Context,
 	CUA_NAVIGATION_TOOL_NAME,
 	CUA_PLAYWRIGHT_TOOL_NAME,
 	cuaModels,
@@ -42,6 +44,64 @@ import { InternalComputerTranslator, type KernelBrowser } from "./translator/tra
  */
 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 included in the request-time message
+ * projection, or `false` to leave image blocks unchanged.
+ */
+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<TMessage extends AgentMessage>(
+	messages: TMessage[],
+	limit: ToolResultImageReplayLimit,
+): TMessage[] {
+	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 <= limit) return messages;
+
+	const firstRetainedImage = Math.max(0, imageCount - 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 } as TMessage) : message;
+	});
+}
+
 /**
  * Agent state exposed by {@link CuaAgent}.
  *
@@ -84,6 +144,10 @@ 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 included from message history per provider request. Defaults to 4; false disables projection. */
+	toolResultImageReplayLimit?: ToolResultImageReplayLimit;
+	/** Chain OpenAI and Tzafon requests through provider-stored response state. Defaults to true. */
+	responseThreading?: boolean;
 	/** Optional CUA-level retries around each provider request. Disabled by default. */
 	retry?: CuaRetryOptions;
 };
@@ -124,6 +188,10 @@ export type CuaAgentHarnessOptions<
 	playwright?: boolean;
 	/** Optional payload hook composed after the provider-specific CUA payload hook. */
 	onPayload?: SimpleStreamOptions["onPayload"];
+	/** Maximum tool-result images included from message history per provider request. Defaults to 4; false disables projection. */
+	toolResultImageReplayLimit?: ToolResultImageReplayLimit;
+	/** Chain OpenAI and Tzafon requests through provider-stored response state. Defaults to true. */
+	responseThreading?: boolean;
 	/** Optional CUA-level retries around each provider request. Disabled by default. */
 	retry?: CuaRetryOptions;
 };
@@ -279,6 +347,70 @@ class CuaRuntimeController {
 /** Default stream path: the shared CUA `Models` collection. */
 const defaultCuaStream: StreamFn = (model, context, options) => cuaModels().streamSimple(model, context, options);
 
+function resolveResponseThreading(responseThreading: boolean | undefined): boolean {
+	if (responseThreading !== undefined && typeof responseThreading !== "boolean") {
+		throw new TypeError("responseThreading must be a boolean");
+	}
+	return responseThreading ?? true;
+}
+
+function withResponseThreading<TOptions extends object | undefined>(
+	options: TOptions,
+	enabled: boolean,
+): TOptions {
+	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,
+	responseThreading: boolean,
+): Models {
+	if (imageReplayLimit === false && responseThreading) return models;
+
+	const stream: Models["stream"] = (model, context, options) =>
+		models.stream(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
+	const complete: Models["complete"] = (model, context, options) =>
+		models.complete(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
+	const streamSimple: Models["streamSimple"] = (model, context, options) =>
+		models.streamSimple(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
+	const completeSimple: Models["completeSimple"] = (model, context, options) =>
+		models.completeSimple(
+			model,
+			projectModelContext(context, imageReplayLimit),
+			withResponseThreading(options, responseThreading),
+		);
+	return {
+		getProviders: () => models.getProviders(),
+		getProvider: (id) => models.getProvider(id),
+		getModels: (provider) => models.getModels(provider),
+		getModel: (provider, id) => models.getModel(provider, id),
+		refresh: (provider) => models.refresh(provider),
+		getAuth: (model) => models.getAuth(model),
+		stream,
+		complete,
+		streamSimple,
+		completeSimple,
+	};
+}
+
 /**
  * Pi `Agent` configured for Kernel browser computer use.
  *
@@ -302,13 +434,18 @@ export class CuaAgent extends Agent {
 			onPayload,
 			streamFn,
 			prepareNextTurn,
+			transformContext,
 			extraTools,
 			mode,
 			nativeTool,
 			playwright,
+			toolResultImageReplayLimit,
+			responseThreading,
 			retry,
 			...agentOptions
 		} = options;
+		const imageReplayLimit = resolveToolResultImageReplayLimit(toolResultImageReplayLimit);
+		const useResponseThreading = resolveResponseThreading(responseThreading);
 		const runtime = new CuaRuntimeController({
 			browser,
 			client,
@@ -328,6 +465,7 @@ export class CuaAgent extends Agent {
 				...streamOptions,
 				onPayload: runtime.onPayload(),
 				keepToolNames: runtime.keepToolNames(),
+				disableResponseThreading: !useResponseThreading,
 			};
 			return retryingStream(model, context, optionsWithCuaRuntime);
 		};
@@ -336,6 +474,11 @@ export class CuaAgent extends Agent {
 			...agentOptions,
 			getApiKey: agentOptions.getApiKey ?? getCuaEnvApiKey,
 			streamFn: wrappedStreamFn,
+			transformContext: async (messages, signal) =>
+				projectToolResultImages(
+					transformContext ? await transformContext(messages, signal) : messages,
+					imageReplayLimit,
+				),
 			initialState: {
 				...initialState,
 				model: runtime.model,
@@ -461,9 +604,13 @@ export class CuaAgentHarness<
 			systemPrompt,
 			onPayload,
 			activeToolNames,
+			toolResultImageReplayLimit,
+			responseThreading,
 			retry,
 			...harnessOptions
 		} = options;
+		const imageReplayLimit = resolveToolResultImageReplayLimit(toolResultImageReplayLimit);
+		const useResponseThreading = resolveResponseThreading(responseThreading);
 		const runtime = new CuaRuntimeController({
 			browser,
 			client,
@@ -479,11 +626,12 @@ export class CuaAgentHarness<
 			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),

packages/agent/src/index.ts

@@ -16,5 +16,10 @@ export type {
 	PlaywrightDetails,
 } from "./tools";
 export { CuaAgent, CuaAgentHarness } from "./agent";
-export type { CuaAgentHarnessOptions, CuaAgentOptions, CuaAgentState } from "./agent";
+export type {
+	CuaAgentHarnessOptions,
+	CuaAgentOptions,
+	CuaAgentState,
+	ToolResultImageReplayLimit,
+} from "./agent";
 export type { CuaRetryOptions } from "./provider-retry";