fix(agent): simplify state API and update consumers fixes #2633

packages/agent/CHANGELOG.md

@@ -2,6 +2,33 @@
 
 ## [Unreleased]
 
+### Breaking Changes
+
+- `AgentState` has been reshaped:
+  - `streamMessage` was renamed to `streamingMessage`
+  - `error` was renamed to `errorMessage`
+  - `isStreaming`, `streamingMessage`, `pendingToolCalls`, and `errorMessage` are now readonly in the public API
+  - `pendingToolCalls` is now typed as `ReadonlySet<string>`
+  - `tools` and `messages` are now accessor properties, and assigning either field copies the provided top-level array instead of preserving array identity
+- `AgentOptions.initialState` no longer accepts runtime-owned fields. Remove `isStreaming`, `streamingMessage`, `pendingToolCalls`, and `errorMessage` from `initialState` values.
+- Removed `Agent` mutator methods in favor of direct property access:
+  - `agent.setSystemPrompt(value)` -> `agent.state.systemPrompt = value`
+  - `agent.setModel(model)` -> `agent.state.model = model`
+  - `agent.setThinkingLevel(level)` -> `agent.state.thinkingLevel = level`
+  - `agent.setTools(tools)` -> `agent.state.tools = tools`
+  - `agent.replaceMessages(messages)` -> `agent.state.messages = messages`
+  - `agent.appendMessage(message)` -> `agent.state.messages.push(message)`
+  - `agent.clearMessages()` -> `agent.state.messages = []`
+  - `agent.setToolExecution(mode)` -> `agent.toolExecution = mode`
+  - `agent.setBeforeToolCall(fn)` -> `agent.beforeToolCall = fn`
+  - `agent.setAfterToolCall(fn)` -> `agent.afterToolCall = fn`
+  - `agent.setTransport(transport)` -> `agent.transport = transport`
+- Removed queue mode getter/setter methods in favor of properties:
+  - `agent.setSteeringMode(mode)` -> `agent.steeringMode = mode`
+  - `agent.getSteeringMode()` -> `agent.steeringMode`
+  - `agent.setFollowUpMode(mode)` -> `agent.followUpMode = mode`
+  - `agent.getFollowUpMode()` -> `agent.followUpMode`
+
 ## [0.64.0] - 2026-03-29
 
 ### Added

packages/agent/README.md

@@ -204,14 +204,18 @@ interface AgentState {
   thinkingLevel: ThinkingLevel;
   tools: AgentTool<any>[];
   messages: AgentMessage[];
-  isStreaming: boolean;
+  readonly isStreaming: boolean;
-  streamMessage: AgentMessage | null;  // Current partial during streaming
+  readonly streamingMessage?: AgentMessage;
-  pendingToolCalls: Set<string>;
+  readonly pendingToolCalls: ReadonlySet<string>;
-  error?: string;
+  readonly errorMessage?: string;
 }
 ```
 
-Access via `agent.state`. During streaming, `streamMessage` contains the partial assistant message.
+Access state via `agent.state`.
+
+Assigning `agent.state.tools = [...]` or `agent.state.messages = [...]` copies the top-level array before storing it. Mutating the returned array mutates the current agent state.
+
+During streaming, `agent.state.streamingMessage` contains the current partial assistant message.
 
 ## Methods
 
@@ -236,17 +240,16 @@ await agent.continue();
 ### State Management
 
 ```typescript
-agent.setSystemPrompt("New prompt");
+agent.state.systemPrompt = "New prompt";
-agent.setModel(getModel("openai", "gpt-4o"));
+agent.state.model = getModel("openai", "gpt-4o");
-agent.setThinkingLevel("medium");
+agent.state.thinkingLevel = "medium";
-agent.setTools([myTool]);
+agent.state.tools = [myTool];
-agent.setToolExecution("sequential");
+agent.toolExecution = "sequential";
-agent.setBeforeToolCall(async ({ toolCall }) => undefined);
+agent.beforeToolCall = async ({ toolCall }) => undefined;
-agent.setAfterToolCall(async ({ toolCall, result }) => undefined);
+agent.afterToolCall = async ({ toolCall, result }) => undefined;
-agent.replaceMessages(newMessages);
+agent.state.messages = newMessages; // top-level array is copied
-agent.appendMessage(message);
+agent.state.messages.push(message);
-agent.clearMessages();
+agent.reset();
-agent.reset();  // Clear everything
 ```
 
 ### Session and Thinking Budgets
@@ -283,8 +286,8 @@ unsubscribe();
 Steering messages let you interrupt the agent while tools are running. Follow-up messages let you queue work after the agent would otherwise stop.
 
 ```typescript
-agent.setSteeringMode("one-at-a-time");
+agent.steeringMode = "one-at-a-time";
-agent.setFollowUpMode("one-at-a-time");
+agent.followUpMode = "one-at-a-time";
 
 // While agent is running tools
 agent.steer({
@@ -300,8 +303,8 @@ agent.followUp({
   timestamp: Date.now(),
 });
 
-const steeringMode = agent.getSteeringMode();
+const steeringMode = agent.steeringMode;
-const followUpMode = agent.getFollowUpMode();
+const followUpMode = agent.followUpMode;
 
 agent.clearSteeringQueue();
 agent.clearFollowUpQueue();
@@ -370,7 +373,7 @@ const readFileTool: AgentTool = {
   },
 };
 
-agent.setTools([readFileTool]);
+agent.state.tools = [readFileTool];
 ```
 
 ### Error Handling

packages/agent/src/agent.ts

@@ -1,10 +1,4 @@
-/**
- * Agent class that uses the agent-loop directly.
- * No transport abstraction - calls streamSimple via the loop.
- */
-
 import {
-	getModel,
 	type ImageContent,
 	type Message,
 	type Model,
@@ -27,453 +21,466 @@ import type {
 	BeforeToolCallContext,
 	BeforeToolCallResult,
 	StreamFn,
-	ThinkingLevel,
 	ToolExecutionMode,
 } from "./types.js";
 
-/**
- * Default convertToLlm: Keep only LLM-compatible messages, convert attachments.
- */
 function defaultConvertToLlm(messages: AgentMessage[]): Message[] {
-	return messages.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult");
+	return messages.filter(
+		(message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult",
+	);
 }
 
-export interface AgentOptions {
+const EMPTY_USAGE = {
-	initialState?: Partial<AgentState>;
+	input: 0,
+	output: 0,
+	cacheRead: 0,
+	cacheWrite: 0,
+	totalTokens: 0,
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+};
 
-	/**
+const DEFAULT_MODEL = {
-	 * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.
+	id: "unknown",
-	 * Default filters to user/assistant/toolResult and converts attachments.
+	name: "unknown",
-	 */
+	api: "unknown",
-	convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
+	provider: "unknown",
+	baseUrl: "",
+	reasoning: false,
+	input: [],
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+	contextWindow: 0,
+	maxTokens: 0,
+} satisfies Model<any>;
 
-	/**
+type QueueMode = "all" | "one-at-a-time";
-	 * Optional transform applied to context before convertToLlm.
-	 * Use for context pruning, injecting external context, etc.
-	 */
-	transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
 
-	/**
+type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & {
-	 * Steering mode: "all" = send all steering messages at once, "one-at-a-time" = one per turn
+	isStreaming: boolean;
-	 */
+	streamingMessage?: AgentMessage;
-	steeringMode?: "all" | "one-at-a-time";
+	pendingToolCalls: Set<string>;
+	errorMessage?: string;
+};
 
-	/**
+function createMutableAgentState(
-	 * Follow-up mode: "all" = send all follow-up messages at once, "one-at-a-time" = one per turn
+	initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>,
-	 */
+): MutableAgentState {
-	followUpMode?: "all" | "one-at-a-time";
+	let tools = initialState?.tools?.slice() ?? [];
+	let messages = initialState?.messages?.slice() ?? [];
 
-	/**
+	return {
-	 * Custom stream function (for proxy backends, etc.). Default uses streamSimple.
+		systemPrompt: initialState?.systemPrompt ?? "",
-	 */
+		model: initialState?.model ?? DEFAULT_MODEL,
-	streamFn?: StreamFn;
+		thinkingLevel: initialState?.thinkingLevel ?? "off",
-
+		get tools() {
-	/**
+			return tools;
-	 * Optional session identifier forwarded to LLM providers.
+		},
-	 * Used by providers that support session-based caching (e.g., OpenAI Codex).
+		set tools(nextTools: AgentTool<any>[]) {
-	 */
+			tools = nextTools.slice();
-	sessionId?: string;
+		},
-
+		get messages() {
-	/**
+			return messages;
-	 * Resolves an API key dynamically for each LLM call.
+		},
-	 * Useful for expiring tokens (e.g., GitHub Copilot OAuth).
+		set messages(nextMessages: AgentMessage[]) {
-	 */
+			messages = nextMessages.slice();
-	getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
+		},
-
-	/**
-	 * Inspect or replace provider payloads before they are sent.
-	 */
-	onPayload?: SimpleStreamOptions["onPayload"];
-
-	/**
-	 * Custom token budgets for thinking levels (token-based providers only).
-	 */
-	thinkingBudgets?: ThinkingBudgets;
-
-	/**
-	 * Preferred transport for providers that support multiple transports.
-	 */
-	transport?: Transport;
-
-	/**
-	 * Maximum delay in milliseconds to wait for a retry when the server requests a long wait.
-	 * If the server's requested delay exceeds this value, the request fails immediately,
-	 * allowing higher-level retry logic to handle it with user visibility.
-	 * Default: 60000 (60 seconds). Set to 0 to disable the cap.
-	 */
-	maxRetryDelayMs?: number;
-
-	/** Tool execution mode. Default: "parallel" */
-	toolExecution?: ToolExecutionMode;
-
-	/** Called before a tool is executed, after arguments have been validated. */
-	beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
-
-	/** Called after a tool finishes executing, before final tool events are emitted. */
-	afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
-}
-
-export class Agent {
-	private _state: AgentState = {
-		systemPrompt: "",
-		model: getModel("google", "gemini-2.5-flash-lite-preview-06-17"),
-		thinkingLevel: "off",
-		tools: [],
-		messages: [],
 		isStreaming: false,
-		streamMessage: null,
+		streamingMessage: undefined,
 		pendingToolCalls: new Set<string>(),
-		error: undefined,
+		errorMessage: undefined,
 	};
+}
 
-	private listeners = new Set<(e: AgentEvent) => void>();
+/** Options for constructing an {@link Agent}. */
-	private abortController?: AbortController;
+export interface AgentOptions {
-	private convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
+	initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>;
-	private transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
+	convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
-	private steeringQueue: AgentMessage[] = [];
+	transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
-	private followUpQueue: AgentMessage[] = [];
+	streamFn?: StreamFn;
-	private steeringMode: "all" | "one-at-a-time";
+	getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
-	private followUpMode: "all" | "one-at-a-time";
+	onPayload?: SimpleStreamOptions["onPayload"];
+	beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
+	afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
+	steeringMode?: QueueMode;
+	followUpMode?: QueueMode;
+	sessionId?: string;
+	thinkingBudgets?: ThinkingBudgets;
+	transport?: Transport;
+	maxRetryDelayMs?: number;
+	toolExecution?: ToolExecutionMode;
+}
+
+class PendingMessageQueue {
+	private messages: AgentMessage[] = [];
+
+	constructor(public mode: QueueMode) {}
+
+	enqueue(message: AgentMessage): void {
+		this.messages.push(message);
+	}
+
+	hasItems(): boolean {
+		return this.messages.length > 0;
+	}
+
+	drain(): AgentMessage[] {
+		if (this.mode === "all") {
+			const drained = this.messages.slice();
+			this.messages = [];
+			return drained;
+		}
+
+		const first = this.messages[0];
+		if (!first) {
+			return [];
+		}
+		this.messages = this.messages.slice(1);
+		return [first];
+	}
+
+	clear(): void {
+		this.messages = [];
+	}
+}
+
+type ActiveRun = {
+	promise: Promise<void>;
+	resolve: () => void;
+	abortController: AbortController;
+};
+
+/**
+ * Stateful wrapper around the low-level agent loop.
+ *
+ * `Agent` owns the current transcript, emits lifecycle events, executes tools,
+ * and exposes queueing APIs for steering and follow-up messages.
+ */
+export class Agent {
+	private _state: MutableAgentState;
+	private readonly listeners = new Set<(event: AgentEvent) => void>();
+	private readonly steeringQueue: PendingMessageQueue;
+	private readonly followUpQueue: PendingMessageQueue;
+
+	public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
+	public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
 	public streamFn: StreamFn;
-	private _sessionId?: string;
 	public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
-	private _onPayload?: SimpleStreamOptions["onPayload"];
+	public onPayload?: SimpleStreamOptions["onPayload"];
-	private runningPrompt?: Promise<void>;
+	public beforeToolCall?: (
-	private resolveRunningPrompt?: () => void;
-	private _thinkingBudgets?: ThinkingBudgets;
-	private _transport: Transport;
-	private _maxRetryDelayMs?: number;
-	private _toolExecution: ToolExecutionMode;
-	private _beforeToolCall?: (
 		context: BeforeToolCallContext,
 		signal?: AbortSignal,
 	) => Promise<BeforeToolCallResult | undefined>;
-	private _afterToolCall?: (
+	public afterToolCall?: (
 		context: AfterToolCallContext,
 		signal?: AbortSignal,
 	) => Promise<AfterToolCallResult | undefined>;
+	private activeRun?: ActiveRun;
+	/** Session identifier forwarded to providers for cache-aware backends. */
+	public sessionId?: string;
+	/** Optional per-level thinking token budgets forwarded to the stream function. */
+	public thinkingBudgets?: ThinkingBudgets;
+	/** Preferred transport forwarded to the stream function. */
+	public transport: Transport;
+	/** Optional cap for provider-requested retry delays. */
+	public maxRetryDelayMs?: number;
+	/** Tool execution strategy for assistant messages that contain multiple tool calls. */
+	public toolExecution: ToolExecutionMode;
 
-	constructor(opts: AgentOptions = {}) {
+	constructor(options: AgentOptions = {}) {
-		this._state = { ...this._state, ...opts.initialState };
+		this._state = createMutableAgentState(options.initialState);
-		this.convertToLlm = opts.convertToLlm || defaultConvertToLlm;
+		this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm;
-		this.transformContext = opts.transformContext;
+		this.transformContext = options.transformContext;
-		this.steeringMode = opts.steeringMode || "one-at-a-time";
+		this.streamFn = options.streamFn ?? streamSimple;
-		this.followUpMode = opts.followUpMode || "one-at-a-time";
+		this.getApiKey = options.getApiKey;
-		this.streamFn = opts.streamFn || streamSimple;
+		this.onPayload = options.onPayload;
-		this._sessionId = opts.sessionId;
+		this.beforeToolCall = options.beforeToolCall;
-		this.getApiKey = opts.getApiKey;
+		this.afterToolCall = options.afterToolCall;
-		this._onPayload = opts.onPayload;
+		this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time");
-		this._thinkingBudgets = opts.thinkingBudgets;
+		this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time");
-		this._transport = opts.transport ?? "sse";
+		this.sessionId = options.sessionId;
-		this._maxRetryDelayMs = opts.maxRetryDelayMs;
+		this.thinkingBudgets = options.thinkingBudgets;
-		this._toolExecution = opts.toolExecution ?? "parallel";
+		this.transport = options.transport ?? "sse";
-		this._beforeToolCall = opts.beforeToolCall;
+		this.maxRetryDelayMs = options.maxRetryDelayMs;
-		this._afterToolCall = opts.afterToolCall;
+		this.toolExecution = options.toolExecution ?? "parallel";
+	}
+
+	/** Subscribe to agent lifecycle events. Returns an unsubscribe function. */
+	subscribe(listener: (event: AgentEvent) => void): () => void {
+		this.listeners.add(listener);
+		return () => this.listeners.delete(listener);
 	}
 
 	/**
-	 * Get the current session ID used for provider caching.
+	 * Current agent state.
+	 *
+	 * Assigning `state.tools` or `state.messages` copies the provided top-level array.
 	 */
-	get sessionId(): string | undefined {
-		return this._sessionId;
-	}
-
-	/**
-	 * Set the session ID for provider caching.
-	 * Call this when switching sessions (new session, branch, resume).
-	 */
-	set sessionId(value: string | undefined) {
-		this._sessionId = value;
-	}
-
-	/**
-	 * Get the current thinking budgets.
-	 */
-	get thinkingBudgets(): ThinkingBudgets | undefined {
-		return this._thinkingBudgets;
-	}
-
-	/**
-	 * Set custom thinking budgets for token-based providers.
-	 */
-	set thinkingBudgets(value: ThinkingBudgets | undefined) {
-		this._thinkingBudgets = value;
-	}
-
-	/**
-	 * Get the current preferred transport.
-	 */
-	get transport(): Transport {
-		return this._transport;
-	}
-
-	/**
-	 * Set the preferred transport.
-	 */
-	setTransport(value: Transport) {
-		this._transport = value;
-	}
-
-	/**
-	 * Get the current max retry delay in milliseconds.
-	 */
-	get maxRetryDelayMs(): number | undefined {
-		return this._maxRetryDelayMs;
-	}
-
-	/**
-	 * Set the maximum delay to wait for server-requested retries.
-	 * Set to 0 to disable the cap.
-	 */
-	set maxRetryDelayMs(value: number | undefined) {
-		this._maxRetryDelayMs = value;
-	}
-
-	get toolExecution(): ToolExecutionMode {
-		return this._toolExecution;
-	}
-
-	setToolExecution(value: ToolExecutionMode) {
-		this._toolExecution = value;
-	}
-
-	setBeforeToolCall(
-		value:
-			| ((context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>)
-			| undefined,
-	) {
-		this._beforeToolCall = value;
-	}
-
-	setAfterToolCall(
-		value:
-			| ((context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>)
-			| undefined,
-	) {
-		this._afterToolCall = value;
-	}
-
 	get state(): AgentState {
 		return this._state;
 	}
 
-	subscribe(fn: (e: AgentEvent) => void): () => void {
+	/** Controls how queued steering messages are drained. */
-		this.listeners.add(fn);
+	set steeringMode(mode: QueueMode) {
-		return () => this.listeners.delete(fn);
+		this.steeringQueue.mode = mode;
 	}
 
-	// State mutators
+	get steeringMode(): QueueMode {
-	setSystemPrompt(v: string) {
+		return this.steeringQueue.mode;
-		this._state.systemPrompt = v;
 	}
 
-	setModel(m: Model<any>) {
+	/** Controls how queued follow-up messages are drained. */
-		this._state.model = m;
+	set followUpMode(mode: QueueMode) {
+		this.followUpQueue.mode = mode;
 	}
 
-	setThinkingLevel(l: ThinkingLevel) {
+	get followUpMode(): QueueMode {
-		this._state.thinkingLevel = l;
+		return this.followUpQueue.mode;
 	}
 
-	setSteeringMode(mode: "all" | "one-at-a-time") {
+	/** Queue a message to be injected after the current assistant turn finishes. */
-		this.steeringMode = mode;
+	steer(message: AgentMessage): void {
+		this.steeringQueue.enqueue(message);
 	}
 
-	getSteeringMode(): "all" | "one-at-a-time" {
+	/** Queue a message to run only after the agent would otherwise stop. */
-		return this.steeringMode;
+	followUp(message: AgentMessage): void {
+		this.followUpQueue.enqueue(message);
 	}
 
-	setFollowUpMode(mode: "all" | "one-at-a-time") {
+	/** Remove all queued steering messages. */
-		this.followUpMode = mode;
+	clearSteeringQueue(): void {
+		this.steeringQueue.clear();
 	}
 
-	getFollowUpMode(): "all" | "one-at-a-time" {
+	/** Remove all queued follow-up messages. */
-		return this.followUpMode;
+	clearFollowUpQueue(): void {
+		this.followUpQueue.clear();
 	}
 
-	setTools(t: AgentTool<any>[]) {
+	/** Remove all queued steering and follow-up messages. */
-		this._state.tools = t;
+	clearAllQueues(): void {
-	}
+		this.clearSteeringQueue();
-
+		this.clearFollowUpQueue();
-	replaceMessages(ms: AgentMessage[]) {
-		this._state.messages = ms.slice();
-	}
-
-	appendMessage(m: AgentMessage) {
-		this._state.messages = [...this._state.messages, m];
-	}
-
-	/**
-	 * Queue a steering message while the agent is running.
-	 * Delivered after the current assistant turn finishes executing its tool calls,
-	 * before the next LLM call.
-	 */
-	steer(m: AgentMessage) {
-		this.steeringQueue.push(m);
-	}
-
-	/**
-	 * Queue a follow-up message to be processed after the agent finishes.
-	 * Delivered only when agent has no more tool calls or steering messages.
-	 */
-	followUp(m: AgentMessage) {
-		this.followUpQueue.push(m);
-	}
-
-	clearSteeringQueue() {
-		this.steeringQueue = [];
-	}
-
-	clearFollowUpQueue() {
-		this.followUpQueue = [];
-	}
-
-	clearAllQueues() {
-		this.steeringQueue = [];
-		this.followUpQueue = [];
 	}
 
+	/** Returns true when either queue still contains pending messages. */
 	hasQueuedMessages(): boolean {
-		return this.steeringQueue.length > 0 || this.followUpQueue.length > 0;
+		return this.steeringQueue.hasItems() || this.followUpQueue.hasItems();
 	}
 
-	private dequeueSteeringMessages(): AgentMessage[] {
+	/** Active abort signal for the current run, if any. */
-		if (this.steeringMode === "one-at-a-time") {
-			if (this.steeringQueue.length > 0) {
-				const first = this.steeringQueue[0];
-				this.steeringQueue = this.steeringQueue.slice(1);
-				return [first];
-			}
-			return [];
-		}
-
-		const steering = this.steeringQueue.slice();
-		this.steeringQueue = [];
-		return steering;
-	}
-
-	private dequeueFollowUpMessages(): AgentMessage[] {
-		if (this.followUpMode === "one-at-a-time") {
-			if (this.followUpQueue.length > 0) {
-				const first = this.followUpQueue[0];
-				this.followUpQueue = this.followUpQueue.slice(1);
-				return [first];
-			}
-			return [];
-		}
-
-		const followUp = this.followUpQueue.slice();
-		this.followUpQueue = [];
-		return followUp;
-	}
-
-	clearMessages() {
-		this._state.messages = [];
-	}
-
-	/** The current abort signal, or undefined when the agent is not streaming. */
 	get signal(): AbortSignal | undefined {
-		return this.abortController?.signal;
+		return this.activeRun?.abortController.signal;
 	}
 
-	abort() {
+	/** Abort the current run, if one is active. */
-		this.abortController?.abort();
+	abort(): void {
+		this.activeRun?.abortController.abort();
 	}
 
+	/** Resolve when the current run has finished. */
 	waitForIdle(): Promise<void> {
-		return this.runningPrompt ?? Promise.resolve();
+		return this.activeRun?.promise ?? Promise.resolve();
 	}
 
-	reset() {
+	/** Clear transcript state, runtime state, and queued messages. */
+	reset(): void {
 		this._state.messages = [];
 		this._state.isStreaming = false;
-		this._state.streamMessage = null;
+		this._state.streamingMessage = undefined;
 		this._state.pendingToolCalls = new Set<string>();
-		this._state.error = undefined;
+		this._state.errorMessage = undefined;
-		this.steeringQueue = [];
+		this.clearFollowUpQueue();
-		this.followUpQueue = [];
+		this.clearSteeringQueue();
 	}
 
-	/** Send a prompt with an AgentMessage */
+	/** Start a new prompt from text, a single message, or a batch of messages. */
 	async prompt(message: AgentMessage | AgentMessage[]): Promise<void>;
 	async prompt(input: string, images?: ImageContent[]): Promise<void>;
-	async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]) {
+	async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise<void> {
-		if (this._state.isStreaming) {
+		if (this.activeRun) {
 			throw new Error(
 				"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
 			);
 		}
-
+		const messages = this.normalizePromptInput(input, images);
-		const model = this._state.model;
+		await this.runPromptMessages(messages);
-		if (!model) throw new Error("No model configured");
-
-		let msgs: AgentMessage[];
-
-		if (Array.isArray(input)) {
-			msgs = input;
-		} else if (typeof input === "string") {
-			const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
-			if (images && images.length > 0) {
-				content.push(...images);
-			}
-			msgs = [
-				{
-					role: "user",
-					content,
-					timestamp: Date.now(),
-				},
-			];
-		} else {
-			msgs = [input];
-		}
-
-		await this._runLoop(msgs);
 	}
 
-	/**
+	/** Continue from the current transcript. The last message must be a user or tool-result message. */
-	 * Continue from current context (used for retries and resuming queued messages).
+	async continue(): Promise<void> {
-	 */
+		if (this.activeRun) {
-	async continue() {
-		if (this._state.isStreaming) {
 			throw new Error("Agent is already processing. Wait for completion before continuing.");
 		}
 
-		const messages = this._state.messages;
+		const lastMessage = this._state.messages[this._state.messages.length - 1];
-		if (messages.length === 0) {
+		if (!lastMessage) {
 			throw new Error("No messages to continue from");
 		}
-		if (messages[messages.length - 1].role === "assistant") {
+
-			const queuedSteering = this.dequeueSteeringMessages();
+		if (lastMessage.role === "assistant") {
+			const queuedSteering = this.steeringQueue.drain();
 			if (queuedSteering.length > 0) {
-				await this._runLoop(queuedSteering, { skipInitialSteeringPoll: true });
+				await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true });
 				return;
 			}
 
-			const queuedFollowUp = this.dequeueFollowUpMessages();
+			const queuedFollowUps = this.followUpQueue.drain();
-			if (queuedFollowUp.length > 0) {
+			if (queuedFollowUps.length > 0) {
-				await this._runLoop(queuedFollowUp);
+				await this.runPromptMessages(queuedFollowUps);
 				return;
 			}
 
 			throw new Error("Cannot continue from message role: assistant");
 		}
 
-		await this._runLoop(undefined);
+		await this.runContinuation();
 	}
 
-	private _processLoopEvent(event: AgentEvent): void {
+	private normalizePromptInput(
+		input: string | AgentMessage | AgentMessage[],
+		images?: ImageContent[],
+	): AgentMessage[] {
+		if (Array.isArray(input)) {
+			return input;
+		}
+
+		if (typeof input !== "string") {
+			return [input];
+		}
+
+		const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
+		if (images && images.length > 0) {
+			content.push(...images);
+		}
+		return [{ role: "user", content, timestamp: Date.now() }];
+	}
+
+	private async runPromptMessages(
+		messages: AgentMessage[],
+		options: { skipInitialSteeringPoll?: boolean } = {},
+	): Promise<void> {
+		await this.runWithLifecycle(async (signal) => {
+			await runAgentLoop(
+				messages,
+				this.createContextSnapshot(),
+				this.createLoopConfig(options),
+				(event) => this.processEvents(event),
+				signal,
+				this.streamFn,
+			);
+		});
+	}
+
+	private async runContinuation(): Promise<void> {
+		await this.runWithLifecycle(async (signal) => {
+			await runAgentLoopContinue(
+				this.createContextSnapshot(),
+				this.createLoopConfig(),
+				(event) => this.processEvents(event),
+				signal,
+				this.streamFn,
+			);
+		});
+	}
+
+	private createContextSnapshot(): AgentContext {
+		return {
+			systemPrompt: this._state.systemPrompt,
+			messages: this._state.messages.slice(),
+			tools: this._state.tools.slice(),
+		};
+	}
+
+	private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig {
+		let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true;
+		return {
+			model: this._state.model,
+			reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel,
+			sessionId: this.sessionId,
+			onPayload: this.onPayload,
+			transport: this.transport,
+			thinkingBudgets: this.thinkingBudgets,
+			maxRetryDelayMs: this.maxRetryDelayMs,
+			toolExecution: this.toolExecution,
+			beforeToolCall: this.beforeToolCall,
+			afterToolCall: this.afterToolCall,
+			convertToLlm: this.convertToLlm,
+			transformContext: this.transformContext,
+			getApiKey: this.getApiKey,
+			getSteeringMessages: async () => {
+				if (skipInitialSteeringPoll) {
+					skipInitialSteeringPoll = false;
+					return [];
+				}
+				return this.steeringQueue.drain();
+			},
+			getFollowUpMessages: async () => this.followUpQueue.drain(),
+		};
+	}
+
+	private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> {
+		if (this.activeRun) {
+			throw new Error("Agent is already processing.");
+		}
+
+		const abortController = new AbortController();
+		let resolvePromise = () => {};
+		const promise = new Promise<void>((resolve) => {
+			resolvePromise = resolve;
+		});
+		this.activeRun = { promise, resolve: resolvePromise, abortController };
+
+		this._state.isStreaming = true;
+		this._state.streamingMessage = undefined;
+		this._state.errorMessage = undefined;
+
+		try {
+			await executor(abortController.signal);
+		} catch (error) {
+			this.handleRunFailure(error, abortController.signal.aborted);
+		} finally {
+			this.finishRun();
+		}
+	}
+
+	private handleRunFailure(error: unknown, aborted: boolean): void {
+		const failureMessage = {
+			role: "assistant",
+			content: [{ type: "text", text: "" }],
+			api: this._state.model.api,
+			provider: this._state.model.provider,
+			model: this._state.model.id,
+			usage: EMPTY_USAGE,
+			stopReason: aborted ? "aborted" : "error",
+			errorMessage: error instanceof Error ? error.message : String(error),
+			timestamp: Date.now(),
+		} satisfies AgentMessage;
+		this._state.messages.push(failureMessage);
+		this._state.errorMessage = failureMessage.errorMessage;
+		this.processEvents({ type: "agent_end", messages: [failureMessage] });
+	}
+
+	private finishRun(): void {
+		this._state.isStreaming = false;
+		this._state.streamingMessage = undefined;
+		this._state.pendingToolCalls = new Set<string>();
+		this.activeRun?.resolve();
+		this.activeRun = undefined;
+	}
+
+	private processEvents(event: AgentEvent): void {
 		switch (event.type) {
 			case "message_start":
-				this._state.streamMessage = event.message;
+				this._state.streamingMessage = event.message;
 				break;
 
 			case "message_update":
-				this._state.streamMessage = event.message;
+				this._state.streamingMessage = event.message;
 				break;
 
 			case "message_end":
-				this._state.streamMessage = null;
+				this._state.streamingMessage = undefined;
-				this.appendMessage(event.message);
+				this._state.messages.push(event.message);
 				break;
 
 			case "tool_execution_start": {
@@ -491,128 +498,19 @@ export class Agent {
 			}
 
 			case "turn_end":
-				if (event.message.role === "assistant" && (event.message as any).errorMessage) {
+				if (event.message.role === "assistant" && event.message.errorMessage) {
-					this._state.error = (event.message as any).errorMessage;
+					this._state.errorMessage = event.message.errorMessage;
 				}
 				break;
 
 			case "agent_end":
 				this._state.isStreaming = false;
-				this._state.streamMessage = null;
+				this._state.streamingMessage = undefined;
 				break;
 		}
 
-		this.emit(event);
-	}
-
-	/**
-	 * Run the agent loop.
-	 * If messages are provided, starts a new conversation turn with those messages.
-	 * Otherwise, continues from existing context.
-	 */
-	private async _runLoop(messages?: AgentMessage[], options?: { skipInitialSteeringPoll?: boolean }) {
-		const model = this._state.model;
-		if (!model) throw new Error("No model configured");
-
-		this.runningPrompt = new Promise<void>((resolve) => {
-			this.resolveRunningPrompt = resolve;
-		});
-
-		this.abortController = new AbortController();
-		this._state.isStreaming = true;
-		this._state.streamMessage = null;
-		this._state.error = undefined;
-
-		const reasoning = this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel;
-
-		const context: AgentContext = {
-			systemPrompt: this._state.systemPrompt,
-			messages: this._state.messages.slice(),
-			tools: this._state.tools,
-		};
-
-		let skipInitialSteeringPoll = options?.skipInitialSteeringPoll === true;
-
-		const config: AgentLoopConfig = {
-			model,
-			reasoning,
-			sessionId: this._sessionId,
-			onPayload: this._onPayload,
-			transport: this._transport,
-			thinkingBudgets: this._thinkingBudgets,
-			maxRetryDelayMs: this._maxRetryDelayMs,
-			toolExecution: this._toolExecution,
-			beforeToolCall: this._beforeToolCall,
-			afterToolCall: this._afterToolCall,
-			convertToLlm: this.convertToLlm,
-			transformContext: this.transformContext,
-			getApiKey: this.getApiKey,
-			getSteeringMessages: async () => {
-				if (skipInitialSteeringPoll) {
-					skipInitialSteeringPoll = false;
-					return [];
-				}
-				return this.dequeueSteeringMessages();
-			},
-			getFollowUpMessages: async () => this.dequeueFollowUpMessages(),
-		};
-
-		try {
-			if (messages) {
-				await runAgentLoop(
-					messages,
-					context,
-					config,
-					async (event) => this._processLoopEvent(event),
-					this.abortController.signal,
-					this.streamFn,
-				);
-			} else {
-				await runAgentLoopContinue(
-					context,
-					config,
-					async (event) => this._processLoopEvent(event),
-					this.abortController.signal,
-					this.streamFn,
-				);
-			}
-		} catch (err: any) {
-			const errorMsg: AgentMessage = {
-				role: "assistant",
-				content: [{ type: "text", text: "" }],
-				api: model.api,
-				provider: model.provider,
-				model: model.id,
-				usage: {
-					input: 0,
-					output: 0,
-					cacheRead: 0,
-					cacheWrite: 0,
-					totalTokens: 0,
-					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-				},
-				stopReason: this.abortController?.signal.aborted ? "aborted" : "error",
-				errorMessage: err?.message || String(err),
-				timestamp: Date.now(),
-			} as AgentMessage;
-
-			this.appendMessage(errorMsg);
-			this._state.error = err?.message || String(err);
-			this.emit({ type: "agent_end", messages: [errorMsg] });
-		} finally {
-			this._state.isStreaming = false;
-			this._state.streamMessage = null;
-			this._state.pendingToolCalls = new Set<string>();
-			this.abortController = undefined;
-			this.resolveRunningPrompt?.();
-			this.runningPrompt = undefined;
-			this.resolveRunningPrompt = undefined;
-		}
-	}
-
-	private emit(e: AgentEvent) {
 		for (const listener of this.listeners) {
-			listener(e);
+			listener(event);
 		}
 	}
 }

packages/agent/src/types.ts

@@ -245,37 +245,55 @@ export interface CustomAgentMessages {
 export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];
 
 /**
- * Agent state containing all configuration and conversation data.
+ * Public agent state.
+ *
+ * `tools` and `messages` use accessor properties so implementations can copy
+ * assigned arrays before storing them.
  */
 export interface AgentState {
+	/** System prompt sent with each model request. */
 	systemPrompt: string;
+	/** Active model used for future turns. */
 	model: Model<any>;
+	/** Requested reasoning level for future turns. */
 	thinkingLevel: ThinkingLevel;
-	tools: AgentTool<any>[];
+	/** Available tools. Assigning a new array copies the top-level array. */
-	messages: AgentMessage[]; // Can include attachments + custom message types
+	set tools(tools: AgentTool<any>[]);
-	isStreaming: boolean;
+	get tools(): AgentTool<any>[];
-	streamMessage: AgentMessage | null;
+	/** Conversation transcript. Assigning a new array copies the top-level array. */
-	pendingToolCalls: Set<string>;
+	set messages(messages: AgentMessage[]);
-	error?: string;
+	get messages(): AgentMessage[];
+	/** True while the agent is processing a prompt or continuation. */
+	readonly isStreaming: boolean;
+	/** Partial assistant message for the current streamed response, if any. */
+	readonly streamingMessage?: AgentMessage;
+	/** Tool call ids currently executing. */
+	readonly pendingToolCalls: ReadonlySet<string>;
+	/** Error message from the most recent failed or aborted assistant turn, if any. */
+	readonly errorMessage?: string;
 }
 
+/** Final or partial result produced by a tool. */
 export interface AgentToolResult<T> {
-	// Content blocks supporting text and images
+	/** Text or image content returned to the model. */
 	content: (TextContent | ImageContent)[];
-	// Details to be displayed in a UI or logged
+	/** Arbitrary structured details for logs or UI rendering. */
 	details: T;
 }
 
-// Callback for streaming tool execution updates
+/** Callback used by tools to stream partial execution updates. */
 export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
 
-// AgentTool extends Tool but adds argument preparation and execution hooks
+/** Tool definition used by the agent runtime. */
 export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
-	// A human-readable label for the tool to be displayed in UI
+	/** Human-readable label for UI display. */
 	label: string;
-	// Optional compatibility shim to prepare raw tool call arguments before schema validation.
+	/**
-	// Must return an object conforming to TParameters.
+	 * Optional compatibility shim for raw tool-call arguments before schema validation.
+	 * Must return an object that matches `TParameters`.
+	 */
 	prepareArguments?: (args: unknown) => Static<TParameters>;
+	/** Execute the tool call. Throw on failure instead of encoding errors in `content`. */
 	execute: (
 		toolCallId: string,
 		params: Static<TParameters>,
@@ -284,10 +302,13 @@ export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any
 	) => Promise<AgentToolResult<TDetails>>;
 }
 
-// AgentContext is like Context but uses AgentTool
+/** Context snapshot passed into the low-level agent loop. */
 export interface AgentContext {
+	/** System prompt included with the request. */
 	systemPrompt: string;
+	/** Transcript visible to the model. */
 	messages: AgentMessage[];
+	/** Tools available for this run. */
 	tools?: AgentTool<any>[];
 }
 

packages/agent/test/agent.test.ts

@@ -47,9 +47,9 @@ describe("Agent", () => {
 		expect(agent.state.tools).toEqual([]);
 		expect(agent.state.messages).toEqual([]);
 		expect(agent.state.isStreaming).toBe(false);
-		expect(agent.state.streamMessage).toBe(null);
+		expect(agent.state.streamingMessage).toBe(undefined);
 		expect(agent.state.pendingToolCalls).toEqual(new Set());
-		expect(agent.state.error).toBeUndefined();
+		expect(agent.state.errorMessage).toBeUndefined();
 	});
 
 	it("should create an agent instance with custom initial state", () => {
@@ -79,13 +79,13 @@ describe("Agent", () => {
 		expect(eventCount).toBe(0);
 
 		// State mutators don't emit events
-		agent.setSystemPrompt("Test prompt");
+		agent.state.systemPrompt = "Test prompt";
 		expect(eventCount).toBe(0);
 		expect(agent.state.systemPrompt).toBe("Test prompt");
 
 		// Unsubscribe should work
 		unsubscribe();
-		agent.setSystemPrompt("Another prompt");
+		agent.state.systemPrompt = "Another prompt";
 		expect(eventCount).toBe(0); // Should not increase
 	});
 
@@ -93,37 +93,38 @@ describe("Agent", () => {
 		const agent = new Agent();
 
 		// Test setSystemPrompt
-		agent.setSystemPrompt("Custom prompt");
+		agent.state.systemPrompt = "Custom prompt";
 		expect(agent.state.systemPrompt).toBe("Custom prompt");
 
 		// Test setModel
 		const newModel = getModel("google", "gemini-2.5-flash");
-		agent.setModel(newModel);
+		agent.state.model = newModel;
 		expect(agent.state.model).toBe(newModel);
 
 		// Test setThinkingLevel
-		agent.setThinkingLevel("high");
+		agent.state.thinkingLevel = "high";
 		expect(agent.state.thinkingLevel).toBe("high");
 
 		// Test setTools
 		const tools = [{ name: "test", description: "test tool" } as any];
-		agent.setTools(tools);
+		agent.state.tools = tools;
-		expect(agent.state.tools).toBe(tools);
+		expect(agent.state.tools).toEqual(tools);
+		expect(agent.state.tools).not.toBe(tools); // Should be a copy
 
 		// Test replaceMessages
 		const messages = [{ role: "user" as const, content: "Hello", timestamp: Date.now() }];
-		agent.replaceMessages(messages);
+		agent.state.messages = messages;
 		expect(agent.state.messages).toEqual(messages);
 		expect(agent.state.messages).not.toBe(messages); // Should be a copy
 
 		// Test appendMessage
 		const newMessage = { role: "assistant" as const, content: [{ type: "text" as const, text: "Hi" }] };
-		agent.appendMessage(newMessage as any);
+		agent.state.messages.push(newMessage as any);
 		expect(agent.state.messages).toHaveLength(2);
 		expect(agent.state.messages[1]).toBe(newMessage);
 
 		// Test clearMessages
-		agent.clearMessages();
+		agent.state.messages = [];
 		expect(agent.state.messages).toEqual([]);
 	});
 
@@ -241,14 +242,14 @@ describe("Agent", () => {
 			},
 		});
 
-		agent.replaceMessages([
+		agent.state.messages = [
 			{
 				role: "user",
 				content: [{ type: "text", text: "Initial" }],
 				timestamp: Date.now() - 10,
 			},
 			createAssistantMessage("Initial response"),
-		]);
+		];
 
 		agent.followUp({
 			role: "user",
@@ -285,14 +286,14 @@ describe("Agent", () => {
 			},
 		});
 
-		agent.replaceMessages([
+		agent.state.messages = [
 			{
 				role: "user",
 				content: [{ type: "text", text: "Initial" }],
 				timestamp: Date.now() - 10,
 			},
 			createAssistantMessage("Initial response"),
-		]);
+		];
 
 		agent.steer({
 			role: "user",

packages/agent/test/bedrock-models.test.ts

@@ -1,287 +0,0 @@
-/**
- * A test suite to ensure Amazon Bedrock models work correctly with the agent loop.
- *
- * Some Bedrock models don't support all features (e.g., reasoning signatures).
- * This test suite verifies that the agent loop works with various Bedrock models.
- *
- * This test suite is not enabled by default unless AWS credentials and
- * `BEDROCK_EXTENSIVE_MODEL_TEST` environment variables are set.
- *
- * You can run this test suite with:
- * ```bash
- * $ AWS_REGION=us-east-1 BEDROCK_EXTENSIVE_MODEL_TEST=1 AWS_PROFILE=pi npm test -- ./test/bedrock-models.test.ts
- * ```
- *
- * ## Known Issues by Category
- *
- * 1. **Inference Profile Required**: Some models require an inference profile ARN instead of on-demand.
- * 2. **Invalid Model ID**: Model identifiers that don't exist in the current region.
- * 3. **Max Tokens Exceeded**: Model's maxTokens in our config exceeds the actual limit.
- * 4. **No Reasoning in User Messages**: Model rejects reasoning content when replayed in conversation.
- * 5. **Invalid Signature Format**: Model validates signature format (Anthropic newer models).
- */
-
-import type { AssistantMessage } from "@mariozechner/pi-ai";
-import { getModels } from "@mariozechner/pi-ai";
-import { describe, expect, it } from "vitest";
-import { Agent } from "../src/index.js";
-import { hasBedrockCredentials } from "./bedrock-utils.js";
-
-// =============================================================================
-// Known Issue Categories
-// =============================================================================
-
-/** Models that require inference profile ARN (not available on-demand in us-east-1) */
-const REQUIRES_INFERENCE_PROFILE = new Set([
-	"anthropic.claude-3-5-haiku-20241022-v1:0",
-	"anthropic.claude-3-5-sonnet-20241022-v2:0",
-	"anthropic.claude-3-opus-20240229-v1:0",
-	"meta.llama3-1-70b-instruct-v1:0",
-	"meta.llama3-1-8b-instruct-v1:0",
-]);
-
-/** Models with invalid identifiers (not available in us-east-1 or don't exist) */
-const INVALID_MODEL_ID = new Set([
-	"deepseek.v3-v1:0",
-	"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
-	"eu.anthropic.claude-opus-4-5-20251101-v1:0",
-	"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
-	"qwen.qwen3-235b-a22b-2507-v1:0",
-	"qwen.qwen3-coder-480b-a35b-v1:0",
-]);
-
-/** Models where our maxTokens config exceeds the model's actual limit */
-const MAX_TOKENS_EXCEEDED = new Set([
-	"us.meta.llama4-maverick-17b-instruct-v1:0",
-	"us.meta.llama4-scout-17b-instruct-v1:0",
-]);
-
-/**
- * Models that reject reasoning content in user messages (when replaying conversation).
- * These work for multi-turn but fail when synthetic thinking is injected.
- */
-const NO_REASONING_IN_USER_MESSAGES = new Set([
-	// Mistral models
-	"mistral.ministral-3-14b-instruct",
-	"mistral.ministral-3-8b-instruct",
-	"mistral.mistral-large-2402-v1:0",
-	"mistral.voxtral-mini-3b-2507",
-	"mistral.voxtral-small-24b-2507",
-	// Nvidia models
-	"nvidia.nemotron-nano-12b-v2",
-	"nvidia.nemotron-nano-9b-v2",
-	// Qwen models
-	"qwen.qwen3-coder-30b-a3b-v1:0",
-	// Amazon Nova models
-	"us.amazon.nova-lite-v1:0",
-	"us.amazon.nova-micro-v1:0",
-	"us.amazon.nova-premier-v1:0",
-	"us.amazon.nova-pro-v1:0",
-	// Meta Llama models
-	"us.meta.llama3-2-11b-instruct-v1:0",
-	"us.meta.llama3-2-1b-instruct-v1:0",
-	"us.meta.llama3-2-3b-instruct-v1:0",
-	"us.meta.llama3-2-90b-instruct-v1:0",
-	"us.meta.llama3-3-70b-instruct-v1:0",
-	// DeepSeek
-	"us.deepseek.r1-v1:0",
-	// Older Anthropic models
-	"anthropic.claude-3-5-sonnet-20240620-v1:0",
-	"anthropic.claude-3-haiku-20240307-v1:0",
-	"anthropic.claude-3-sonnet-20240229-v1:0",
-	// Cohere models
-	"cohere.command-r-plus-v1:0",
-	"cohere.command-r-v1:0",
-	// Google models
-	"google.gemma-3-27b-it",
-	"google.gemma-3-4b-it",
-	// Non-Anthropic models that don't support signatures (now handled by omitting signature)
-	// but still reject reasoning content in user messages
-	"global.amazon.nova-2-lite-v1:0",
-	"minimax.minimax-m2",
-	"moonshot.kimi-k2-thinking",
-	"openai.gpt-oss-120b-1:0",
-	"openai.gpt-oss-20b-1:0",
-	"openai.gpt-oss-safeguard-120b",
-	"openai.gpt-oss-safeguard-20b",
-	"qwen.qwen3-32b-v1:0",
-	"qwen.qwen3-next-80b-a3b",
-	"qwen.qwen3-vl-235b-a22b",
-]);
-
-/**
- * Models that validate signature format (Anthropic newer models).
- * These work for multi-turn but fail when synthetic/invalid signature is injected.
- */
-const VALIDATES_SIGNATURE_FORMAT = new Set([
-	"global.anthropic.claude-haiku-4-5-20251001-v1:0",
-	"global.anthropic.claude-opus-4-5-20251101-v1:0",
-	"global.anthropic.claude-sonnet-4-20250514-v1:0",
-	"global.anthropic.claude-sonnet-4-5-20250929-v1:0",
-	"us.anthropic.claude-3-7-sonnet-20250219-v1:0",
-	"us.anthropic.claude-opus-4-1-20250805-v1:0",
-	"us.anthropic.claude-opus-4-20250514-v1:0",
-]);
-
-/**
- * DeepSeek R1 fails multi-turn because it rejects reasoning in the replayed assistant message.
- */
-const REJECTS_REASONING_ON_REPLAY = new Set(["us.deepseek.r1-v1:0"]);
-
-// =============================================================================
-// Helper Functions
-// =============================================================================
-
-function isModelUnavailable(modelId: string): boolean {
-	return REQUIRES_INFERENCE_PROFILE.has(modelId) || INVALID_MODEL_ID.has(modelId) || MAX_TOKENS_EXCEEDED.has(modelId);
-}
-
-function failsMultiTurnWithThinking(modelId: string): boolean {
-	return REJECTS_REASONING_ON_REPLAY.has(modelId);
-}
-
-function failsSyntheticSignature(modelId: string): boolean {
-	return NO_REASONING_IN_USER_MESSAGES.has(modelId) || VALIDATES_SIGNATURE_FORMAT.has(modelId);
-}
-
-// =============================================================================
-// Tests
-// =============================================================================
-
-describe("Amazon Bedrock Models - Agent Loop", () => {
-	const shouldRunExtensiveTests = hasBedrockCredentials() && process.env.BEDROCK_EXTENSIVE_MODEL_TEST;
-
-	// Get all Amazon Bedrock models
-	const allBedrockModels = getModels("amazon-bedrock");
-
-	if (shouldRunExtensiveTests) {
-		for (const model of allBedrockModels) {
-			const modelId = model.id;
-
-			describe(`Model: ${modelId}`, () => {
-				// Skip entirely unavailable models
-				const unavailable = isModelUnavailable(modelId);
-
-				it.skipIf(unavailable)("should handle basic text prompt", { timeout: 60_000 }, async () => {
-					const agent = new Agent({
-						initialState: {
-							systemPrompt: "You are a helpful assistant. Be extremely concise.",
-							model,
-							thinkingLevel: "off",
-							tools: [],
-						},
-					});
-
-					await agent.prompt("Reply with exactly: 'OK'");
-
-					if (agent.state.error) {
-						throw new Error(`Basic prompt error: ${agent.state.error}`);
-					}
-
-					expect(agent.state.isStreaming).toBe(false);
-					expect(agent.state.messages.length).toBe(2);
-
-					const assistantMessage = agent.state.messages[1];
-					if (assistantMessage.role !== "assistant") throw new Error("Expected assistant message");
-
-					console.log(`${modelId}: OK`);
-				});
-
-				// Skip if model is unavailable or known to fail multi-turn with thinking
-				const skipMultiTurn = unavailable || failsMultiTurnWithThinking(modelId);
-
-				it.skipIf(skipMultiTurn)(
-					"should handle multi-turn conversation with thinking content in history",
-					{ timeout: 120_000 },
-					async () => {
-						const agent = new Agent({
-							initialState: {
-								systemPrompt: "You are a helpful assistant. Be extremely concise.",
-								model,
-								thinkingLevel: "medium",
-								tools: [],
-							},
-						});
-
-						// First turn
-						await agent.prompt("My name is Alice.");
-
-						if (agent.state.error) {
-							throw new Error(`First turn error: ${agent.state.error}`);
-						}
-
-						// Second turn - this should replay the first assistant message which may contain thinking
-						await agent.prompt("What is my name?");
-
-						if (agent.state.error) {
-							throw new Error(`Second turn error: ${agent.state.error}`);
-						}
-
-						expect(agent.state.messages.length).toBe(4);
-						console.log(`${modelId}: multi-turn OK`);
-					},
-				);
-
-				// Skip if model is unavailable or known to fail synthetic signature
-				const skipSynthetic = unavailable || failsSyntheticSignature(modelId);
-
-				it.skipIf(skipSynthetic)(
-					"should handle conversation with synthetic thinking signature in history",
-					{ timeout: 60_000 },
-					async () => {
-						const agent = new Agent({
-							initialState: {
-								systemPrompt: "You are a helpful assistant. Be extremely concise.",
-								model,
-								thinkingLevel: "off",
-								tools: [],
-							},
-						});
-
-						// Inject a message with a thinking block that has a signature
-						const syntheticAssistantMessage: AssistantMessage = {
-							role: "assistant",
-							content: [
-								{
-									type: "thinking",
-									thinking: "I need to remember the user's name.",
-									thinkingSignature: "synthetic-signature-123",
-								},
-								{ type: "text", text: "Nice to meet you, Alice!" },
-							],
-							api: "bedrock-converse-stream",
-							provider: "amazon-bedrock",
-							model: modelId,
-							usage: {
-								input: 10,
-								output: 20,
-								cacheRead: 0,
-								cacheWrite: 0,
-								totalTokens: 30,
-								cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
-							},
-							stopReason: "stop",
-							timestamp: Date.now(),
-						};
-
-						agent.replaceMessages([
-							{ role: "user", content: "My name is Alice.", timestamp: Date.now() },
-							syntheticAssistantMessage,
-						]);
-
-						await agent.prompt("What is my name?");
-
-						if (agent.state.error) {
-							throw new Error(`Synthetic signature error: ${agent.state.error}`);
-						}
-
-						expect(agent.state.messages.length).toBe(4);
-						console.log(`${modelId}: synthetic signature OK`);
-					},
-				);
-			});
-		}
-	} else {
-		it.skip("skipped - set AWS credentials and BEDROCK_EXTENSIVE_MODEL_TEST=1 to run", () => {});
-	}
-});

packages/agent/test/bedrock-utils.ts

@@ -1,18 +0,0 @@
-/**
- * Utility functions for Amazon Bedrock tests
- */
-
-/**
- * Check if any valid AWS credentials are configured for Bedrock.
- * Returns true if any of the following are set:
- * - AWS_PROFILE (named profile from ~/.aws/credentials)
- * - AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY (IAM keys)
- * - AWS_BEARER_TOKEN_BEDROCK (Bedrock API key)
- */
-export function hasBedrockCredentials(): boolean {
-	return !!(
-		process.env.AWS_PROFILE ||
-		(process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) ||
-		process.env.AWS_BEARER_TOKEN_BEDROCK
-	);
-}

packages/agent/test/e2e.test.ts

@@ -1,13 +1,41 @@
-import type { AssistantMessage, Model, ToolResultMessage, UserMessage } from "@mariozechner/pi-ai";
+import {
-import { getModel } from "@mariozechner/pi-ai";
+	type AssistantMessage,
-import { describe, expect, it } from "vitest";
+	type FauxProviderRegistration,
-import { Agent } from "../src/index.js";
+	fauxAssistantMessage,
-import { hasBedrockCredentials } from "./bedrock-utils.js";
+	fauxText,
+	fauxThinking,
+	fauxToolCall,
+	type Model,
+	registerFauxProvider,
+	type ToolResultMessage,
+	type UserMessage,
+} from "@mariozechner/pi-ai";
+import { afterEach, describe, expect, it } from "vitest";
+import { Agent, type AgentEvent } from "../src/index.js";
 import { calculateTool } from "./utils/calculate.js";
 
-delete process.env.ANTHROPIC_OAUTH_TOKEN;
+const registrations: FauxProviderRegistration[] = [];
 
-async function basicPrompt(model: Model<any>) {
+function createFauxRegistration(options: Parameters<typeof registerFauxProvider>[0] = {}): FauxProviderRegistration {
+	const registration = registerFauxProvider(options);
+	registrations.push(registration);
+	return registration;
+}
+
+function getTextContent(message: AssistantMessage | ToolResultMessage): string {
+	return message.content
+		.filter((block) => block.type === "text")
+		.map((block) => block.text)
+		.join("\n");
+}
+
+afterEach(() => {
+	while (registrations.length > 0) {
+		registrations.pop()?.unregister();
+	}
+});
+
+async function basicPrompt(model: Model<string>) {
 	const agent = new Agent({
 		initialState: {
 			systemPrompt: "You are a helpful assistant. Keep your responses concise.",
@@ -26,15 +54,10 @@ async function basicPrompt(model: Model<any>) {
 
 	const assistantMessage = agent.state.messages[1];
 	if (assistantMessage.role !== "assistant") throw new Error("Expected assistant message");
-	expect(assistantMessage.content.length).toBeGreaterThan(0);
+	expect(getTextContent(assistantMessage)).toContain("4");
-
-	const textContent = assistantMessage.content.find((c) => c.type === "text");
-	expect(textContent).toBeDefined();
-	if (textContent?.type !== "text") throw new Error("Expected text content");
-	expect(textContent.text).toContain("4");
 }
 
-async function toolExecution(model: Model<any>) {
+async function toolExecution(model: Model<string>) {
 	const agent = new Agent({
 		initialState: {
 			systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.",
@@ -44,52 +67,49 @@ async function toolExecution(model: Model<any>) {
 		},
 	});
 
+	const pendingToolCallsDuringEvents: Array<{ type: AgentEvent["type"]; ids: string[] }> = [];
+	agent.subscribe((event) => {
+		if (event.type === "tool_execution_start" || event.type === "tool_execution_end") {
+			pendingToolCallsDuringEvents.push({
+				type: event.type,
+				ids: [...agent.state.pendingToolCalls],
+			});
+		}
+	});
+
 	await agent.prompt("Calculate 123 * 456 using the calculator tool.");
 
 	expect(agent.state.isStreaming).toBe(false);
-	expect(agent.state.messages.length).toBeGreaterThanOrEqual(3);
+	expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
-
+	const toolResultMsg = agent.state.messages.find((message) => message.role === "toolResult");
-	const toolResultMsg = agent.state.messages.find((m) => m.role === "toolResult");
 	expect(toolResultMsg).toBeDefined();
 	if (toolResultMsg?.role !== "toolResult") throw new Error("Expected tool result message");
-	const textContent =
+	expect(getTextContent(toolResultMsg)).toContain("123 * 456 = 56088");
-		toolResultMsg.content
-			?.filter((c) => c.type === "text")
-			.map((c: any) => c.text)
-			.join("\n") || "";
-	expect(textContent).toBeDefined();
-
-	const expectedResult = 123 * 456;
-	expect(textContent).toContain(String(expectedResult));
 
 	const finalMessage = agent.state.messages[agent.state.messages.length - 1];
 	if (finalMessage.role !== "assistant") throw new Error("Expected final assistant message");
-	const finalText = finalMessage.content.find((c) => c.type === "text");
+	expect(getTextContent(finalMessage)).toContain("56088");
-	expect(finalText).toBeDefined();
+	expect(agent.state.pendingToolCalls.size).toBe(0);
-	if (finalText?.type !== "text") throw new Error("Expected text content");
+	expect(pendingToolCallsDuringEvents).toEqual([
-	// Check for number with or without comma formatting
+		{ type: "tool_execution_start", ids: ["calc-1"] },
-	const hasNumber =
+		{ type: "tool_execution_end", ids: [] },
-		finalText.text.includes(String(expectedResult)) ||
+	]);
-		finalText.text.includes("56,088") ||
-		finalText.text.includes("56088");
-	expect(hasNumber).toBe(true);
 }
 
-async function abortExecution(model: Model<any>) {
+async function abortExecution(model: Model<string>) {
 	const agent = new Agent({
 		initialState: {
 			systemPrompt: "You are a helpful assistant.",
 			model,
 			thinkingLevel: "off",
-			tools: [calculateTool],
+			tools: [],
 		},
 	});
 
-	const promptPromise = agent.prompt("Calculate 100 * 200, then 300 * 400, then sum the results.");
+	const promptPromise = agent.prompt("Count slowly from 1 to 20.");
-
 	setTimeout(() => {
 		agent.abort();
-	}, 100);
+	}, 30);
 
 	await promptPromise;
 
@@ -100,11 +120,10 @@ async function abortExecution(model: Model<any>) {
 	if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
 	expect(lastMessage.stopReason).toBe("aborted");
 	expect(lastMessage.errorMessage).toBeDefined();
-	expect(agent.state.error).toBeDefined();
+	expect(agent.state.errorMessage).toBe(lastMessage.errorMessage);
-	expect(agent.state.error).toBe(lastMessage.errorMessage);
 }
 
-async function stateUpdates(model: Model<any>) {
+async function stateUpdates(model: Model<string>) {
 	const agent = new Agent({
 		initialState: {
 			systemPrompt: "You are a helpful assistant.",
@@ -114,29 +133,29 @@ async function stateUpdates(model: Model<any>) {
 		},
 	});
 
-	const events: Array<string> = [];
+	const events: AgentEvent["type"][] = [];
-
 	agent.subscribe((event) => {
 		events.push(event.type);
 	});
 
 	await agent.prompt("Count from 1 to 5.");
 
-	// Should have received lifecycle events
 	expect(events).toContain("agent_start");
-	expect(events).toContain("agent_end");
+	expect(events).toContain("turn_start");
 	expect(events).toContain("message_start");
+	expect(events).toContain("message_update");
 	expect(events).toContain("message_end");
-	// May have message_update events during streaming
+	expect(events).toContain("turn_end");
-	const hasMessageUpdates = events.some((e) => e === "message_update");
+	expect(events).toContain("agent_end");
-	expect(hasMessageUpdates).toBe(true);
+	expect(events.indexOf("agent_start")).toBeLessThan(events.indexOf("message_start"));
+	expect(events.indexOf("message_start")).toBeLessThan(events.indexOf("message_end"));
+	expect(events.indexOf("message_end")).toBeLessThan(events.lastIndexOf("agent_end"));
 
-	// Check final state
 	expect(agent.state.isStreaming).toBe(false);
-	expect(agent.state.messages.length).toBe(2); // User message + assistant response
+	expect(agent.state.messages.length).toBe(2);
 }
 
-async function multiTurnConversation(model: Model<any>) {
+async function multiTurnConversation(model: Model<string>) {
 	const agent = new Agent({
 		initialState: {
 			systemPrompt: "You are a helpful assistant.",
@@ -154,232 +173,120 @@ async function multiTurnConversation(model: Model<any>) {
 
 	const lastMessage = agent.state.messages[3];
 	if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
-	const lastText = lastMessage.content.find((c) => c.type === "text");
+	expect(getTextContent(lastMessage).toLowerCase()).toContain("alice");
-	if (lastText?.type !== "text") throw new Error("Expected text content");
-	expect(lastText.text.toLowerCase()).toContain("alice");
 }
 
-describe("Agent E2E Tests", () => {
+describe("Agent integration with faux provider", () => {
-	describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider (gemini-2.5-flash)", () => {
+	it("handles a basic text prompt", async () => {
-		const model = getModel("google", "gemini-2.5-flash");
+		const faux = createFauxRegistration();
-
+		faux.setResponses([fauxAssistantMessage("4")]);
-		it("should handle basic text prompt", async () => {
+		await basicPrompt(faux.getModel());
-			await basicPrompt(model);
-		});
-
-		it("should execute tools correctly", async () => {
-			await toolExecution(model);
-		});
-
-		it("should handle abort during execution", async () => {
-			await abortExecution(model);
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
 	});
 
-	describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Provider (gpt-4o-mini)", () => {
+	it("executes tools and tracks pending tool calls", async () => {
-		const model = getModel("openai", "gpt-4o-mini");
+		const faux = createFauxRegistration();
-
+		faux.setResponses([
-		it("should handle basic text prompt", async () => {
+			fauxAssistantMessage(
-			await basicPrompt(model);
+				[
-		});
+					fauxText("Let me calculate that."),
-
+					fauxToolCall("calculate", { expression: "123 * 456" }, { id: "calc-1" }),
-		it("should execute tools correctly", async () => {
+				],
-			await toolExecution(model);
+				{ stopReason: "toolUse" },
-		});
+			),
-
+			fauxAssistantMessage("The result is 56088."),
-		it("should handle abort during execution", async () => {
+		]);
-			await abortExecution(model);
+		await toolExecution(faux.getModel());
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
 	});
 
-	describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider (claude-haiku-4-5)", () => {
+	it("handles abort during streaming", async () => {
-		const model = getModel("anthropic", "claude-haiku-4-5");
+		const faux = createFauxRegistration({
-
+			tokensPerSecond: 20,
-		it("should handle basic text prompt", async () => {
+			tokenSize: { min: 2, max: 2 },
-			await basicPrompt(model);
-		});
-
-		it("should execute tools correctly", async () => {
-			await toolExecution(model);
-		});
-
-		it("should handle abort during execution", async () => {
-			await abortExecution(model);
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
 		});
+		faux.setResponses([
+			fauxAssistantMessage(
+				"one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen",
+			),
+		]);
+		await abortExecution(faux.getModel());
 	});
 
-	describe.skipIf(!process.env.XAI_API_KEY)("xAI Provider (grok-3)", () => {
+	it("emits lifecycle updates while streaming", async () => {
-		const model = getModel("xai", "grok-3");
+		const faux = createFauxRegistration({ tokenSize: { min: 1, max: 1 } });
-
+		faux.setResponses([fauxAssistantMessage("1 2 3 4 5")]);
-		it("should handle basic text prompt", async () => {
+		await stateUpdates(faux.getModel());
-			await basicPrompt(model);
-		});
-
-		it("should execute tools correctly", async () => {
-			await toolExecution(model);
-		});
-
-		it("should handle abort during execution", async () => {
-			await abortExecution(model);
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
 	});
 
-	describe.skipIf(!process.env.GROQ_API_KEY)("Groq Provider (openai/gpt-oss-20b)", () => {
+	it("maintains context across multiple turns", async () => {
-		const model = getModel("groq", "openai/gpt-oss-20b");
+		const faux = createFauxRegistration();
-
+		faux.setResponses([
-		it("should handle basic text prompt", async () => {
+			fauxAssistantMessage("Nice to meet you, Alice."),
-			await basicPrompt(model);
+			(context) => {
-		});
+				const hasAlice = context.messages.some((message) => {
-
+					if (message.role !== "user") return false;
-		it("should execute tools correctly", async () => {
+					if (typeof message.content === "string") return message.content.includes("Alice");
-			await toolExecution(model);
+					return message.content.some((block) => block.type === "text" && block.text.includes("Alice"));
-		});
+				});
-
+				return fauxAssistantMessage(hasAlice ? "Your name is Alice." : "I do not know your name.");
-		it("should handle abort during execution", async () => {
+			},
-			await abortExecution(model);
+		]);
-		});
+		await multiTurnConversation(faux.getModel());
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
 	});
 
-	/*describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras Provider (gpt-oss-120b)", () => {
+	it("preserves thinking content blocks", async () => {
-		const model = getModel("cerebras", "gpt-oss-120b");
+		const faux = createFauxRegistration({ models: [{ id: "faux-reasoning", reasoning: true }] });
+		faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]);
 
-		it("should handle basic text prompt", async () => {
+		const agent = new Agent({
-			await basicPrompt(model);
+			initialState: {
+				systemPrompt: "You are a helpful assistant.",
+				model: faux.getModel(),
+				thinkingLevel: "low",
+				tools: [],
+			},
 		});
 
-		it("should execute tools correctly", async () => {
+		await agent.prompt("What is 2+2?");
-			await toolExecution(model);
-		});
 
-		it("should handle abort during execution", async () => {
+		const assistantMessage = agent.state.messages[1];
-			await abortExecution(model);
+		if (assistantMessage?.role !== "assistant") throw new Error("Expected assistant message");
-		});
+		expect(assistantMessage.content).toEqual([
-
+			{ type: "thinking", thinking: "step by step" },
-		it("should emit state updates during streaming", async () => {
+			{ type: "text", text: "4" },
-			await stateUpdates(model);
+		]);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
-	});*/
-
-	describe.skipIf(!process.env.ZAI_API_KEY)("zAI Provider (glm-4.5-air)", () => {
-		const model = getModel("zai", "glm-4.5-air");
-
-		it("should handle basic text prompt", async () => {
-			await basicPrompt(model);
-		});
-
-		it("should execute tools correctly", async () => {
-			await toolExecution(model);
-		});
-
-		it("should handle abort during execution", async () => {
-			await abortExecution(model);
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
-	});
-
-	describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock Provider (claude-sonnet-4-5)", () => {
-		const model = getModel("amazon-bedrock", "global.anthropic.claude-sonnet-4-5-20250929-v1:0");
-
-		it("should handle basic text prompt", async () => {
-			await basicPrompt(model);
-		});
-
-		it("should execute tools correctly", async () => {
-			await toolExecution(model);
-		});
-
-		it("should handle abort during execution", async () => {
-			await abortExecution(model);
-		});
-
-		it("should emit state updates during streaming", async () => {
-			await stateUpdates(model);
-		});
-
-		it("should maintain context across multiple turns", async () => {
-			await multiTurnConversation(model);
-		});
 	});
 });
 
-describe("Agent.continue()", () => {
+describe("Agent.continue() with faux provider", () => {
 	describe("validation", () => {
-		it("should throw when no messages in context", async () => {
+		it("throws when no messages in context", async () => {
+			const faux = createFauxRegistration();
 			const agent = new Agent({
 				initialState: {
 					systemPrompt: "Test",
-					model: getModel("openai", "gpt-5.4"),
+					model: faux.getModel(),
 				},
 			});
 
 			await expect(agent.continue()).rejects.toThrow("No messages to continue from");
 		});
 
-		it("should throw when last message is assistant", async () => {
+		it("throws when last message is assistant", async () => {
+			const faux = createFauxRegistration();
+			const model = faux.getModel();
 			const agent = new Agent({
 				initialState: {
 					systemPrompt: "Test",
-					model: getModel("openai", "gpt-5.4"),
+					model,
 				},
 			});
 
 			const assistantMessage: AssistantMessage = {
 				role: "assistant",
 				content: [{ type: "text", text: "Hello" }],
-				api: "openai-responses",
+				api: model.api,
-				provider: "openai",
+				provider: model.provider,
-				model: "gpt-5.4",
+				model: model.id,
 				usage: {
 					input: 0,
 					output: 0,
@@ -391,34 +298,32 @@ describe("Agent.continue()", () => {
 				stopReason: "stop",
 				timestamp: Date.now(),
 			};
-			agent.replaceMessages([assistantMessage]);
+			agent.state.messages = [assistantMessage];
 
 			await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant");
 		});
 	});
 
-	describe.skipIf(!process.env.OPENAI_API_KEY)("continue from user message", () => {
+	describe("continue from user message", () => {
-		const model = getModel("openai", "gpt-5.4");
+		it("continues and gets a response when last message is user", async () => {
-
+			const faux = createFauxRegistration();
-		it("should continue and get response when last message is user", async () => {
+			faux.setResponses([fauxAssistantMessage("HELLO WORLD")]);
 			const agent = new Agent({
 				initialState: {
 					systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
-					model,
+					model: faux.getModel(),
 					thinkingLevel: "off",
 					tools: [],
 				},
 			});
 
-			// Manually add a user message without calling prompt()
 			const userMessage: UserMessage = {
 				role: "user",
 				content: [{ type: "text", text: "Say exactly: HELLO WORLD" }],
 				timestamp: Date.now(),
 			};
-			agent.replaceMessages([userMessage]);
+			agent.state.messages = [userMessage];
 
-			// Continue from the user message
 			await agent.continue();
 
 			expect(agent.state.isStreaming).toBe(false);
@@ -426,19 +331,17 @@ describe("Agent.continue()", () => {
 			expect(agent.state.messages[0].role).toBe("user");
 			expect(agent.state.messages[1].role).toBe("assistant");
 
-			const assistantMsg = agent.state.messages[1] as AssistantMessage;
+			const assistantMsg = agent.state.messages[1];
-			const textContent = assistantMsg.content.find((c) => c.type === "text");
+			if (assistantMsg.role !== "assistant") throw new Error("Expected assistant message");
-			expect(textContent).toBeDefined();
+			expect(getTextContent(assistantMsg).toUpperCase()).toContain("HELLO WORLD");
-			if (textContent?.type === "text") {
-				expect(textContent.text.toUpperCase()).toContain("HELLO WORLD");
-			}
 		});
 	});
 
-	describe.skipIf(!process.env.OPENAI_API_KEY)("continue from tool result", () => {
+	describe("continue from tool result", () => {
-		const model = getModel("openai", "gpt-5.4");
+		it("continues and processes tool results", async () => {
-
+			const faux = createFauxRegistration();
-		it("should continue and process tool results", async () => {
+			const model = faux.getModel();
+			faux.setResponses([fauxAssistantMessage("The answer is 8.")]);
 			const agent = new Agent({
 				initialState: {
 					systemPrompt:
@@ -449,7 +352,6 @@ describe("Agent.continue()", () => {
 				},
 			});
 
-			// Set up a conversation state as if tool was just executed
 			const userMessage: UserMessage = {
 				role: "user",
 				content: [{ type: "text", text: "What is 5 + 3?" }],
@@ -462,9 +364,9 @@ describe("Agent.continue()", () => {
 					{ type: "text", text: "Let me calculate that." },
 					{ type: "toolCall", id: "calc-1", name: "calculate", arguments: { expression: "5 + 3" } },
 				],
-				api: "anthropic-messages",
+				api: model.api,
-				provider: "anthropic",
+				provider: model.provider,
-				model: "claude-haiku-4-5",
+				model: model.id,
 				usage: {
 					input: 0,
 					output: 0,
@@ -486,26 +388,17 @@ describe("Agent.continue()", () => {
 				timestamp: Date.now(),
 			};
 
-			agent.replaceMessages([userMessage, assistantMessage, toolResult]);
+			agent.state.messages = [userMessage, assistantMessage, toolResult];
 
-			// Continue from the tool result
 			await agent.continue();
 
 			expect(agent.state.isStreaming).toBe(false);
-			// Should have added an assistant response
 			expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
 
 			const lastMessage = agent.state.messages[agent.state.messages.length - 1];
 			expect(lastMessage.role).toBe("assistant");
-
+			if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
-			if (lastMessage.role === "assistant") {
+			expect(getTextContent(lastMessage)).toContain("8");
-				const textContent = lastMessage.content
-					.filter((c) => c.type === "text")
-					.map((c) => (c as { type: "text"; text: string }).text)
-					.join(" ");
-				// Should mention 8 in the response
-				expect(textContent).toMatch(/8/);
-			}
 		});
 	});
 });

packages/ai/src/models.generated.ts

@@ -2304,6 +2304,23 @@ export const MODELS = {
 			contextWindow: 400000,
 			maxTokens: 128000,
 		} satisfies Model<"azure-openai-responses">,
+		"gpt-5.3-chat-latest": {
+			id: "gpt-5.3-chat-latest",
+			name: "GPT-5.3 Chat (latest)",
+			api: "azure-openai-responses",
+			provider: "azure-openai-responses",
+			baseUrl: "",
+			reasoning: false,
+			input: ["text", "image"],
+			cost: {
+				input: 1.75,
+				output: 14,
+				cacheRead: 0.175,
+				cacheWrite: 0,
+			},
+			contextWindow: 128000,
+			maxTokens: 16384,
+		} satisfies Model<"azure-openai-responses">,
 		"gpt-5.3-codex": {
 			id: "gpt-5.3-codex",
 			name: "GPT-5.3 Codex",
@@ -5029,22 +5046,39 @@ export const MODELS = {
 			contextWindow: 128000,
 			maxTokens: 16384,
 		} satisfies Model<"mistral-conversations">,
+		"mistral-small-2603": {
+			id: "mistral-small-2603",
+			name: "Mistral Small 4",
+			api: "mistral-conversations",
+			provider: "mistral",
+			baseUrl: "https://api.mistral.ai",
+			reasoning: true,
+			input: ["text", "image"],
+			cost: {
+				input: 0.15,
+				output: 0.6,
+				cacheRead: 0,
+				cacheWrite: 0,
+			},
+			contextWindow: 256000,
+			maxTokens: 256000,
+		} satisfies Model<"mistral-conversations">,
 		"mistral-small-latest": {
 			id: "mistral-small-latest",
 			name: "Mistral Small (latest)",
 			api: "mistral-conversations",
 			provider: "mistral",
 			baseUrl: "https://api.mistral.ai",
-			reasoning: false,
+			reasoning: true,
 			input: ["text", "image"],
 			cost: {
-				input: 0.1,
+				input: 0.15,
-				output: 0.3,
+				output: 0.6,
 				cacheRead: 0,
 				cacheWrite: 0,
 			},
-			contextWindow: 128000,
+			contextWindow: 256000,
-			maxTokens: 16384,
+			maxTokens: 256000,
 		} satisfies Model<"mistral-conversations">,
 		"open-mistral-7b": {
 			id: "open-mistral-7b",
@@ -5575,6 +5609,23 @@ export const MODELS = {
 			contextWindow: 400000,
 			maxTokens: 128000,
 		} satisfies Model<"openai-responses">,
+		"gpt-5.3-chat-latest": {
+			id: "gpt-5.3-chat-latest",
+			name: "GPT-5.3 Chat (latest)",
+			api: "openai-responses",
+			provider: "openai",
+			baseUrl: "https://api.openai.com/v1",
+			reasoning: false,
+			input: ["text", "image"],
+			cost: {
+				input: 1.75,
+				output: 14,
+				cacheRead: 0.175,
+				cacheWrite: 0,
+			},
+			contextWindow: 128000,
+			maxTokens: 16384,
+		} satisfies Model<"openai-responses">,
 		"gpt-5.3-codex": {
 			id: "gpt-5.3-codex",
 			name: "GPT-5.3 Codex",
@@ -10779,12 +10830,12 @@ export const MODELS = {
 			reasoning: false,
 			input: ["text"],
 			cost: {
-				input: 0.071,
+				input: 0.22,
-				output: 0.463,
+				output: 0.88,
-				cacheRead: 0,
+				cacheRead: 0.11,
 				cacheWrite: 0,
 			},
-			contextWindow: 40960,
+			contextWindow: 32768,
 			maxTokens: 16384,
 		} satisfies Model<"anthropic-messages">,
 		"alibaba/qwen-3-30b": {
@@ -10813,13 +10864,13 @@ export const MODELS = {
 			reasoning: true,
 			input: ["text"],
 			cost: {
-				input: 0.29,
+				input: 0.16,
-				output: 0.59,
+				output: 0.64,
-				cacheRead: 0.145,
+				cacheRead: 0,
 				cacheWrite: 0,
 			},
-			contextWindow: 131072,
+			contextWindow: 128000,
-			maxTokens: 40960,
+			maxTokens: 8192,
 		} satisfies Model<"anthropic-messages">,
 		"alibaba/qwen3-235b-a22b-thinking": {
 			id: "alibaba/qwen3-235b-a22b-thinking",
@@ -10847,13 +10898,13 @@ export const MODELS = {
 			reasoning: false,
 			input: ["text"],
 			cost: {
-				input: 0.39999999999999997,
+				input: 1.5,
-				output: 1.5999999999999999,
+				output: 7.5,
-				cacheRead: 0.022,
+				cacheRead: 0.3,
 				cacheWrite: 0,
 			},
 			contextWindow: 262144,
-			maxTokens: 66536,
+			maxTokens: 65536,
 		} satisfies Model<"anthropic-messages">,
 		"alibaba/qwen3-coder-30b-a3b": {
 			id: "alibaba/qwen3-coder-30b-a3b",
@@ -11323,13 +11374,13 @@ export const MODELS = {
 			reasoning: true,
 			input: ["text"],
 			cost: {
-				input: 0.5,
+				input: 0.56,
-				output: 1.5,
+				output: 1.68,
-				cacheRead: 0,
+				cacheRead: 0.28,
 				cacheWrite: 0,
 			},
 			contextWindow: 163840,
-			maxTokens: 16384,
+			maxTokens: 8192,
 		} satisfies Model<"anthropic-messages">,
 		"deepseek/deepseek-v3.1-terminus": {
 			id: "deepseek/deepseek-v3.1-terminus",
@@ -11601,7 +11652,7 @@ export const MODELS = {
 				cacheWrite: 0,
 			},
 			contextWindow: 128000,
-			maxTokens: 100000,
+			maxTokens: 8192,
 		} satisfies Model<"anthropic-messages">,
 		"meituan/longcat-flash-thinking": {
 			id: "meituan/longcat-flash-thinking",
@@ -11646,13 +11697,13 @@ export const MODELS = {
 			reasoning: false,
 			input: ["text"],
 			cost: {
-				input: 0.09999999999999999,
+				input: 0.22,
-				output: 0.09999999999999999,
+				output: 0.22,
-				cacheRead: 0.09999999999999999,
+				cacheRead: 0,
 				cacheWrite: 0,
 			},
 			contextWindow: 128000,
-			maxTokens: 16384,
+			maxTokens: 8192,
 		} satisfies Model<"anthropic-messages">,
 		"meta/llama-3.2-11b": {
 			id: "meta/llama-3.2-11b",
@@ -11714,12 +11765,12 @@ export const MODELS = {
 			reasoning: false,
 			input: ["text", "image"],
 			cost: {
-				input: 0.24,
+				input: 0.35,
-				output: 0.9700000000000001,
+				output: 1.15,
 				cacheRead: 0,
 				cacheWrite: 0,
 			},
-			contextWindow: 128000,
+			contextWindow: 524288,
 			maxTokens: 8192,
 		} satisfies Model<"anthropic-messages">,
 		"meta/llama-4-scout": {
@@ -12623,6 +12674,23 @@ export const MODELS = {
 			contextWindow: 1050000,
 			maxTokens: 128000,
 		} satisfies Model<"anthropic-messages">,
+		"openai/gpt-oss-120b": {
+			id: "openai/gpt-oss-120b",
+			name: "gpt-oss-120b",
+			api: "anthropic-messages",
+			provider: "vercel-ai-gateway",
+			baseUrl: "https://ai-gateway.vercel.sh",
+			reasoning: true,
+			input: ["text"],
+			cost: {
+				input: 0.15,
+				output: 0.6,
+				cacheRead: 0,
+				cacheWrite: 0,
+			},
+			contextWindow: 128000,
+			maxTokens: 8192,
+		} satisfies Model<"anthropic-messages">,
 		"openai/gpt-oss-20b": {
 			id: "openai/gpt-oss-20b",
 			name: "gpt-oss-20b",
@@ -13108,9 +13176,9 @@ export const MODELS = {
 			reasoning: true,
 			input: ["text"],
 			cost: {
-				input: 0.09999999999999999,
+				input: 0.09,
-				output: 0.3,
+				output: 0.29,
-				cacheRead: 0.02,
+				cacheRead: 0.045,
 				cacheWrite: 0,
 			},
 			contextWindow: 262144,
@@ -13246,7 +13314,7 @@ export const MODELS = {
 			cost: {
 				input: 0.6,
 				output: 2.2,
-				cacheRead: 0,
+				cacheRead: 0.11,
 				cacheWrite: 0,
 			},
 			contextWindow: 200000,
@@ -13286,6 +13354,23 @@ export const MODELS = {
 			contextWindow: 200000,
 			maxTokens: 128000,
 		} satisfies Model<"anthropic-messages">,
+		"zai/glm-5": {
+			id: "zai/glm-5",
+			name: "GLM 5",
+			api: "anthropic-messages",
+			provider: "vercel-ai-gateway",
+			baseUrl: "https://ai-gateway.vercel.sh",
+			reasoning: true,
+			input: ["text"],
+			cost: {
+				input: 1,
+				output: 3.1999999999999997,
+				cacheRead: 0.19999999999999998,
+				cacheWrite: 0,
+			},
+			contextWindow: 202800,
+			maxTokens: 131100,
+		} satisfies Model<"anthropic-messages">,
 		"zai/glm-5-turbo": {
 			id: "zai/glm-5-turbo",
 			name: "GLM 5 Turbo",

packages/coding-agent/docs/sdk.md

@@ -171,10 +171,15 @@ const state = session.agent.state;
 // state.model: Model - current model
 // state.thinkingLevel: ThinkingLevel - current thinking level
 // state.systemPrompt: string - system prompt
-// state.tools: Tool[] - available tools
+// state.tools: AgentTool[] - available tools
+// state.streamingMessage?: AgentMessage - current partial assistant message
+// state.errorMessage?: string - latest assistant error
 
-// Replace messages (useful for branching, restoration)
+// Replace messages (useful for branching or restoration)
-session.agent.replaceMessages(messages);
+session.agent.state.messages = messages; // copies the top-level array
+
+// Replace tools
+session.agent.state.tools = tools; // copies the top-level array
 
 // Wait for agent to finish processing
 await session.agent.waitForIdle();

packages/coding-agent/docs/tree.md

@@ -141,7 +141,7 @@ Flow:
 4. Fire `session_before_tree` event (hook can cancel or provide summary)
 5. Run default summarizer if needed
 6. Switch leaf via `branch()` or `branchWithSummary()`
-7. Update agent: `agent.replaceMessages(sessionManager.buildSessionContext().messages)`
+7. Update agent: `agent.state.messages = sessionManager.buildSessionContext().messages`
 8. Fire `session_tree` event
 9. Notify custom tools via session event
 10. Return result with `editorText` if user message was selected

packages/coding-agent/src/core/agent-session.ts

@@ -358,7 +358,7 @@ export class AgentSession {
 	 * happens here instead of in wrappers.
 	 */
 	private _installAgentToolHooks(): void {
-		this.agent.setBeforeToolCall(async ({ toolCall, args }) => {
+		this.agent.beforeToolCall = async ({ toolCall, args }) => {
 			const runner = this._extensionRunner;
 			if (!runner?.hasHandlers("tool_call")) {
 				return undefined;
@@ -379,9 +379,9 @@ export class AgentSession {
 				}
 				throw new Error(`Extension failed, blocking execution: ${String(err)}`);
 			}
-		});
+		};
 
-		this.agent.setAfterToolCall(async ({ toolCall, args, result, isError }) => {
+		this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
 			const runner = this._extensionRunner;
 			if (!runner?.hasHandlers("tool_result")) {
 				return undefined;
@@ -405,7 +405,7 @@ export class AgentSession {
 				content: hookResult.content,
 				details: hookResult.details,
 			};
-		});
+		};
 	}
 
 	// =========================================================================
@@ -790,11 +790,11 @@ export class AgentSession {
 				validToolNames.push(name);
 			}
 		}
-		this.agent.setTools(tools);
+		this.agent.state.tools = tools;
 
 		// Rebuild base system prompt with new tool set
 		this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);
-		this.agent.setSystemPrompt(this._baseSystemPrompt);
+		this.agent.state.systemPrompt = this._baseSystemPrompt;
 	}
 
 	/** Whether compaction or branch summarization is currently running */
@@ -813,12 +813,12 @@ export class AgentSession {
 
 	/** Current steering mode */
 	get steeringMode(): "all" | "one-at-a-time" {
-		return this.agent.getSteeringMode();
+		return this.agent.steeringMode;
 	}
 
 	/** Current follow-up mode */
 	get followUpMode(): "all" | "one-at-a-time" {
-		return this.agent.getFollowUpMode();
+		return this.agent.followUpMode;
 	}
 
 	/** Current session file path, or undefined if sessions are disabled */
@@ -1051,10 +1051,10 @@ export class AgentSession {
 			}
 			// Apply extension-modified system prompt, or reset to base
 			if (result?.systemPrompt) {
-				this.agent.setSystemPrompt(result.systemPrompt);
+				this.agent.state.systemPrompt = result.systemPrompt;
 			} else {
 				// Ensure we're using the base prompt (in case previous turn had modifications)
-				this.agent.setSystemPrompt(this._baseSystemPrompt);
+				this.agent.state.systemPrompt = this._baseSystemPrompt;
 			}
 		}
 
@@ -1251,7 +1251,7 @@ export class AgentSession {
 		} else if (options?.triggerTurn) {
 			await this.agent.prompt(appMessage);
 		} else {
-			this.agent.appendMessage(appMessage);
+			this.agent.state.messages.push(appMessage);
 			this.sessionManager.appendCustomMessageEntry(
 				message.customType,
 				message.content,
@@ -1388,7 +1388,7 @@ export class AgentSession {
 			await options.setup(this.sessionManager);
 			// Sync agent state with session manager after setup
 			const sessionContext = this.sessionManager.buildSessionContext();
-			this.agent.replaceMessages(sessionContext.messages);
+			this.agent.state.messages = sessionContext.messages;
 		}
 
 		this._reconnectToAgent();
@@ -1437,7 +1437,7 @@ export class AgentSession {
 
 		const previousModel = this.model;
 		const thinkingLevel = this._getThinkingLevelForModelSwitch();
-		this.agent.setModel(model);
+		this.agent.state.model = model;
 		this.sessionManager.appendModelChange(model.provider, model.id);
 		this.settingsManager.setDefaultModelAndProvider(model.provider, model.id);
 
@@ -1478,7 +1478,7 @@ export class AgentSession {
 		const thinkingLevel = this._getThinkingLevelForModelSwitch(next.thinkingLevel);
 
 		// Apply model
-		this.agent.setModel(next.model);
+		this.agent.state.model = next.model;
 		this.sessionManager.appendModelChange(next.model.provider, next.model.id);
 		this.settingsManager.setDefaultModelAndProvider(next.model.provider, next.model.id);
 
@@ -1506,7 +1506,7 @@ export class AgentSession {
 		const nextModel = availableModels[nextIndex];
 
 		const thinkingLevel = this._getThinkingLevelForModelSwitch();
-		this.agent.setModel(nextModel);
+		this.agent.state.model = nextModel;
 		this.sessionManager.appendModelChange(nextModel.provider, nextModel.id);
 		this.settingsManager.setDefaultModelAndProvider(nextModel.provider, nextModel.id);
 
@@ -1534,7 +1534,7 @@ export class AgentSession {
 		// Only persist if actually changing
 		const isChanging = effectiveLevel !== this.agent.state.thinkingLevel;
 
-		this.agent.setThinkingLevel(effectiveLevel);
+		this.agent.state.thinkingLevel = effectiveLevel;
 
 		if (isChanging) {
 			this.sessionManager.appendThinkingLevelChange(effectiveLevel);
@@ -1620,7 +1620,7 @@ export class AgentSession {
 	 * Saves to settings.
 	 */
 	setSteeringMode(mode: "all" | "one-at-a-time"): void {
-		this.agent.setSteeringMode(mode);
+		this.agent.steeringMode = mode;
 		this.settingsManager.setSteeringMode(mode);
 	}
 
@@ -1629,7 +1629,7 @@ export class AgentSession {
 	 * Saves to settings.
 	 */
 	setFollowUpMode(mode: "all" | "one-at-a-time"): void {
-		this.agent.setFollowUpMode(mode);
+		this.agent.followUpMode = mode;
 		this.settingsManager.setFollowUpMode(mode);
 	}
 
@@ -1724,7 +1724,7 @@ export class AgentSession {
 			this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
-			this.agent.replaceMessages(sessionContext.messages);
+			this.agent.state.messages = sessionContext.messages;
 
 			// Get the saved compaction entry for the extension event
 			const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
@@ -1843,7 +1843,7 @@ export class AgentSession {
 			// but we don't want it in context for the retry)
 			const messages = this.agent.state.messages;
 			if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
-				this.agent.replaceMessages(messages.slice(0, -1));
+				this.agent.state.messages = messages.slice(0, -1);
 			}
 			await this._runAutoCompaction("overflow", true);
 			return;
@@ -1995,7 +1995,7 @@ export class AgentSession {
 			this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension);
 			const newEntries = this.sessionManager.getEntries();
 			const sessionContext = this.sessionManager.buildSessionContext();
-			this.agent.replaceMessages(sessionContext.messages);
+			this.agent.state.messages = sessionContext.messages;
 
 			// Get the saved compaction entry for the extension event
 			const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as
@@ -2022,7 +2022,7 @@ export class AgentSession {
 				const messages = this.agent.state.messages;
 				const lastMsg = messages[messages.length - 1];
 				if (lastMsg?.role === "assistant" && (lastMsg as AssistantMessage).stopReason === "error") {
-					this.agent.replaceMessages(messages.slice(0, -1));
+					this.agent.state.messages = messages.slice(0, -1);
 				}
 
 				setTimeout(() => {
@@ -2108,7 +2108,7 @@ export class AgentSession {
 
 		this._resourceLoader.extendResources(extensionPaths);
 		this._baseSystemPrompt = this._rebuildSystemPrompt(this.getActiveToolNames());
-		this.agent.setSystemPrompt(this._baseSystemPrompt);
+		this.agent.state.systemPrompt = this._baseSystemPrompt;
 	}
 
 	private buildExtensionResourcePaths(entries: Array<{ path: string; extensionPath: string }>): Array<{
@@ -2160,7 +2160,7 @@ export class AgentSession {
 			return;
 		}
 
-		this.agent.setModel(refreshedModel);
+		this.agent.state.model = refreshedModel;
 	}
 
 	private _bindExtensionCore(runner: ExtensionRunner): void {
@@ -2500,7 +2500,7 @@ export class AgentSession {
 		// Remove error message from agent state (keep in session for history)
 		const messages = this.agent.state.messages;
 		if (messages.length > 0 && messages[messages.length - 1].role === "assistant") {
-			this.agent.replaceMessages(messages.slice(0, -1));
+			this.agent.state.messages = messages.slice(0, -1);
 		}
 
 		// Wait with exponential backoff (abortable)
@@ -2633,7 +2633,7 @@ export class AgentSession {
 			this._pendingBashMessages.push(bashMessage);
 		} else {
 			// Add to agent state immediately
-			this.agent.appendMessage(bashMessage);
+			this.agent.state.messages.push(bashMessage);
 
 			// Save to session
 			this.sessionManager.appendMessage(bashMessage);
@@ -2666,7 +2666,7 @@ export class AgentSession {
 
 		for (const bashMessage of this._pendingBashMessages) {
 			// Add to agent state
-			this.agent.appendMessage(bashMessage);
+			this.agent.state.messages.push(bashMessage);
 
 			// Save to session
 			this.sessionManager.appendMessage(bashMessage);
@@ -2725,7 +2725,7 @@ export class AgentSession {
 
 		// Emit session event to custom tools
 
-		this.agent.replaceMessages(sessionContext.messages);
+		this.agent.state.messages = sessionContext.messages;
 
 		// Restore model if saved
 		if (sessionContext.model) {
@@ -2735,7 +2735,7 @@ export class AgentSession {
 				(m) => m.provider === sessionContext.model!.provider && m.id === sessionContext.model!.modelId,
 			);
 			if (match) {
-				this.agent.setModel(match);
+				this.agent.state.model = match;
 				await this._emitModelSelect(match, previousModel, "restore");
 			}
 		}
@@ -2751,7 +2751,7 @@ export class AgentSession {
 			const effectiveLevel = availableLevels.includes(defaultThinkingLevel)
 				? defaultThinkingLevel
 				: this._clampThinkingLevel(defaultThinkingLevel, availableLevels);
-			this.agent.setThinkingLevel(effectiveLevel);
+			this.agent.state.thinkingLevel = effectiveLevel;
 			this.sessionManager.appendThinkingLevelChange(effectiveLevel);
 		}
 
@@ -2824,7 +2824,7 @@ export class AgentSession {
 		// Emit session event to custom tools (with reason "fork")
 
 		if (!skipConversationRestore) {
-			this.agent.replaceMessages(sessionContext.messages);
+			this.agent.state.messages = sessionContext.messages;
 		}
 
 		return { selectedText, cancelled: false };
@@ -3006,7 +3006,7 @@ export class AgentSession {
 
 		// Update agent state
 		const sessionContext = this.sessionManager.buildSessionContext();
-		this.agent.replaceMessages(sessionContext.messages);
+		this.agent.state.messages = sessionContext.messages;
 
 		// Emit session_tree event
 		if (this._extensionRunner) {

packages/coding-agent/src/core/sdk.ts

@@ -326,7 +326,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 
 	// Restore messages if session has existing data
 	if (hasExistingSession) {
-		agent.replaceMessages(existingSession.messages);
+		agent.state.messages = existingSession.messages;
 		if (!hasThinkingEntry) {
 			sessionManager.appendThinkingLevelChange(thinkingLevel);
 		}

packages/coding-agent/src/modes/interactive/interactive-mode.ts

@@ -3319,7 +3319,7 @@ export class InteractiveMode {
 					},
 					onTransportChange: (transport) => {
 						this.settingsManager.setTransport(transport);
-						this.session.agent.setTransport(transport);
+						this.session.agent.transport = transport;
 					},
 					onThinkingLevelChange: (level) => {
 						this.session.setThinkingLevel(level);

packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts

@@ -278,12 +278,12 @@ describe("AgentSession auto-compaction queue resume", () => {
 		};
 
 		// Put both messages into agent state so estimateContextTokens can find the successful one
-		session.agent.replaceMessages([
+		session.agent.state.messages = [
 			{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 },
 			successfulAssistant,
 			{ role: "user", content: [{ type: "text", text: "another prompt" }], timestamp: Date.now() + 500 },
 			errorAssistant,
-		]);
+		];
 
 		const runAutoCompactionSpy = vi
 			.spyOn(
@@ -328,10 +328,10 @@ describe("AgentSession auto-compaction queue resume", () => {
 			timestamp: Date.now(),
 		};
 
-		session.agent.replaceMessages([
+		session.agent.state.messages = [
 			{ role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 },
 			errorAssistant,
-		]);
+		];
 
 		const runAutoCompactionSpy = vi
 			.spyOn(
@@ -407,12 +407,12 @@ describe("AgentSession auto-compaction queue resume", () => {
 		};
 
 		// Agent state has the kept assistant (pre-compaction) and the error (post-compaction)
-		session.agent.replaceMessages([
+		session.agent.state.messages = [
 			{ role: "user", content: [{ type: "text", text: "kept user msg" }], timestamp: preCompactionTimestamp - 1000 },
 			keptAssistant,
 			{ role: "user", content: [{ type: "text", text: "new prompt" }], timestamp: Date.now() - 500 },
 			errorAssistant,
-		]);
+		];
 
 		const runAutoCompactionSpy = vi
 			.spyOn(

packages/coding-agent/test/agent-session-stats.test.ts

@@ -74,7 +74,7 @@ function createSession() {
 }
 
 function syncAgentMessages(session: AgentSession, sessionManager: SessionManager): void {
-	session.agent.replaceMessages(sessionManager.buildSessionContext().messages);
+	session.agent.state.messages = sessionManager.buildSessionContext().messages;
 }
 
 describe("AgentSession.getSessionStats", () => {

packages/mom/src/agent.ts

@@ -446,7 +446,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
 	// Load existing messages
 	const loadedSession = sessionManager.buildSessionContext();
 	if (loadedSession.messages.length > 0) {
-		agent.replaceMessages(loadedSession.messages);
+		agent.state.messages = loadedSession.messages;
 		log.logInfo(`[${channelId}] Loaded ${loadedSession.messages.length} messages from context.jsonl`);
 	}
 
@@ -656,7 +656,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
 			// This picks up any messages synced above
 			const reloadedSession = sessionManager.buildSessionContext();
 			if (reloadedSession.messages.length > 0) {
-				agent.replaceMessages(reloadedSession.messages);
+				agent.state.messages = reloadedSession.messages;
 				log.logInfo(`[${channelId}] Reloaded ${reloadedSession.messages.length} messages from context`);
 			}
 
@@ -672,7 +672,7 @@ function createRunner(sandboxConfig: SandboxConfig, channelId: string, channelDi
 				ctx.users,
 				skills,
 			);
-			session.agent.setSystemPrompt(systemPrompt);
+			session.agent.state.systemPrompt = systemPrompt;
 
 			// Set up file upload function
 			setUploadFunction(async (filePath: string, title?: string) => {

packages/web-ui/README.md

@@ -202,10 +202,10 @@ await agent.prompt({ role: 'user-with-attachments', content: 'Check this', attac
 
 // Control
 agent.abort();
-agent.setModel(newModel);
+agent.state.model = newModel;
-agent.setThinkingLevel('medium');
+agent.state.thinkingLevel = 'medium';
-agent.setTools([...]);
+agent.state.tools = [...];
-agent.queueMessage(customMessage);
+agent.followUp(customMessage);
 ```
 
 ## Message Types
@@ -312,7 +312,7 @@ replTool.runtimeProvidersFactory = () => [
   new ArtifactsRuntimeProvider(artifactsPanel, agent, true), // read-write
 ];
 
-agent.setTools([replTool]);
+agent.state.tools = [replTool];
 ```
 
 ### Extract Document
@@ -325,7 +325,7 @@ import { createExtractDocumentTool } from '@mariozechner/pi-web-ui';
 const extractTool = createExtractDocumentTool();
 extractTool.corsProxyUrl = 'https://corsproxy.io/?';
 
-agent.setTools([extractTool]);
+agent.state.tools = [extractTool];
 ```
 
 ### Artifacts Tool
@@ -337,7 +337,7 @@ const artifactsPanel = new ArtifactsPanel();
 artifactsPanel.agent = agent;
 
 // The tool is available as artifactsPanel.tool
-agent.setTools([artifactsPanel.tool]);
+agent.state.tools = [artifactsPanel.tool];
 ```
 
 ### Custom Tool Renderers
@@ -551,7 +551,7 @@ const success = await ApiKeyPromptDialog.prompt('anthropic');
 import { ModelSelector } from '@mariozechner/pi-web-ui';
 
 ModelSelector.open(currentModel, (selectedModel) => {
-  agent.setModel(selectedModel);
+  agent.state.model = selectedModel;
 });
 ```
 

packages/web-ui/src/ChatPanel.ts

@@ -141,7 +141,7 @@ export class ChatPanel extends LitElement {
 		const additionalTools =
 			config?.toolsFactory?.(agent, this.agentInterface, this.artifactsPanel, runtimeProvidersFactory) || [];
 		const tools = [this.artifactsPanel.tool, ...additionalTools];
-		this.agent.setTools(tools);
+		this.agent.state.tools = tools;
 
 		// Reconstruct artifacts from existing messages
 		// Temporarily disable the onArtifactsChange callback to prevent auto-opening on load

packages/web-ui/src/components/AgentInterface.ts

@@ -376,13 +376,15 @@ export class AgentInterface extends LitElement {
 								if (this.onModelSelect) {
 									this.onModelSelect();
 								} else {
-									ModelSelector.open(state.model, (model) => session.setModel(model));
+									ModelSelector.open(state.model, (model) => {
+										session.state.model = model;
+									});
 								}
 							}}
 							.onThinkingChange=${
 								this.enableThinkingSelector
 									? (level: "off" | "minimal" | "low" | "medium" | "high") => {
-											session.setThinkingLevel(level);
+											session.state.thinkingLevel = level;
 										}
 									: undefined
 							}

packages/web-ui/src/components/MessageList.ts

@@ -11,7 +11,7 @@ import { renderMessage } from "./message-renderer-registry.js";
 export class MessageList extends LitElement {
 	@property({ type: Array }) messages: AgentMessage[] = [];
 	@property({ type: Array }) tools: AgentTool[] = [];
-	@property({ type: Object }) pendingToolCalls?: Set<string>;
+	@property({ type: Object }) pendingToolCalls?: ReadonlySet<string>;
 	@property({ type: Boolean }) isStreaming: boolean = false;
 	@property({ attribute: false }) onCostClick?: () => void;
 

packages/web-ui/src/components/Messages.ts

@@ -85,7 +85,7 @@ export class UserMessage extends LitElement {
 export class AssistantMessage extends LitElement {
 	@property({ type: Object }) message!: AssistantMessageType;
 	@property({ type: Array }) tools?: AgentTool<any>[];
-	@property({ type: Object }) pendingToolCalls?: Set<string>;
+	@property({ type: Object }) pendingToolCalls?: ReadonlySet<string>;
 	@property({ type: Boolean }) hideToolCalls = false;
 	@property({ type: Object }) toolResultsById?: Map<string, ToolResultMessageType>;
 	@property({ type: Boolean }) isStreaming: boolean = false;

packages/web-ui/src/components/StreamingMessageContainer.ts

@@ -6,7 +6,7 @@ import { property, state } from "lit/decorators.js";
 export class StreamingMessageContainer extends LitElement {
 	@property({ type: Array }) tools: AgentTool[] = [];
 	@property({ type: Boolean }) isStreaming = false;
-	@property({ type: Object }) pendingToolCalls?: Set<string>;
+	@property({ type: Object }) pendingToolCalls?: ReadonlySet<string>;
 	@property({ type: Object }) toolResultsById?: Map<string, ToolResultMessage>;
 	@property({ attribute: false }) onCostClick?: () => void;
 

packages/web-ui/src/components/sandbox/ArtifactsRuntimeProvider.ts

@@ -1,3 +1,4 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
 import {
 	ARTIFACTS_RUNTIME_PROVIDER_DESCRIPTION_RO,
 	ARTIFACTS_RUNTIME_PROVIDER_DESCRIPTION_RW,
@@ -13,7 +14,7 @@ interface ArtifactsPanelLike {
 }
 
 interface AgentLike {
-	appendMessage(message: any): void;
+	state: { messages: AgentMessage[] };
 }
 
 /**
@@ -171,7 +172,7 @@ export class ArtifactsRuntimeProvider implements SandboxRuntimeProvider {
 							filename,
 							content,
 						});
-						this.agent?.appendMessage({
+						this.agent?.state.messages.push({
 							role: "artifact",
 							action,
 							filename,
@@ -192,7 +193,7 @@ export class ArtifactsRuntimeProvider implements SandboxRuntimeProvider {
 							command: "delete",
 							filename,
 						});
-						this.agent?.appendMessage({
+						this.agent?.state.messages.push({
 							role: "artifact",
 							action: "delete",
 							filename,