@oh-my-pi/pi-agent-core 16.0.4 → 16.0.6

package/src/agent-loop.ts

@@ -3,12 +3,16 @@
  * Transforms to Message[] only at the LLM call boundary.
  */
 import {
-	type ApiKeyResolveContext,
 	type AssistantMessage,
 	type AssistantMessageEvent,
+	arkToWireSchema,
 	type Context,
 	EventStream,
+	isApiKeyResolver,
+	isArkSchema,
 	isZodSchema,
+	resolveApiKeyOnce,
+	seedApiKeyResolver,
 	streamSimple,
 	type ToolResultMessage,
 	type TSchema,
@@ -33,7 +37,7 @@ import {
 	signalListLabel,
 } from "@oh-my-pi/pi-ai/utils/harmony-leak";
 import { preferredDialect } from "@oh-my-pi/pi-catalog/identity";
-import { logger, sanitizeText } from "@oh-my-pi/pi-utils";
+import { sanitizeText } from "@oh-my-pi/pi-utils";
 import { type AgentRunCoverage, type AgentRunSummary, ToolCallBlockedError } from "./run-collector";
 import {
 	type AgentTelemetry,
@@ -111,6 +115,7 @@ function resolveOwnedDialectFromEnv(value: string | undefined): Dialect | undefi
 		case "qwen3":
 		case "gemini":
 		case "gemma":
+		case "minimax":
 			return value;
 		default:
 			return undefined;
@@ -563,6 +568,9 @@ export function normalizeTools(
 			if (isZodSchema(parameters)) {
 				const wired = zodToWireSchema(parameters);
 				parameters = injectIntentIntoSchema(wired, intentMode) as TSchema;
+			} else if (isArkSchema(parameters)) {
+				const wired = arkToWireSchema(parameters);
+				parameters = injectIntentIntoSchema(wired, intentMode) as TSchema;
 			} else {
 				parameters = injectIntentIntoSchema(parameters, intentMode) as TSchema;
 			}
@@ -636,6 +644,20 @@ interface StepCounter {
 	count: number;
 }
 
+function isDeadlineExceeded(deadline: number | undefined): boolean {
+	return deadline !== undefined && Date.now() >= deadline;
+}
+
+function endAgentStream(
+	stream: EventStream<AgentEvent, AgentMessage[]>,
+	newMessages: AgentMessage[],
+	telemetry: AgentTelemetry | undefined,
+	stepCount: number,
+): void {
+	stream.push(buildAgentEndEvent(newMessages, telemetry, stepCount));
+	stream.end(newMessages);
+}
+
 /**
  * Resolve aside entries at the moment the loop is about to inject them. Each entry
  * is either a ready {@link AgentMessage} or a sync thunk evaluated here so the
@@ -664,245 +686,287 @@ async function runLoopBody(
 	stepCounter: StepCounter,
 	streamFn?: StreamFn,
 ): Promise<void> {
-	let firstTurn = true;
-	// Check for steering messages at start (user may have typed while waiting).
-	// Skip when the run is already externally aborted — dequeuing would strand
-	// the messages in a run that is about to die.
-	let pendingMessages: AgentMessage[] = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
-	let harmonyRetryAttempt = 0;
-	let harmonyTruncateResumeCount = 0;
-	let pausedTurnContinuations = 0;
-
-	// Outer loop: continues when queued follow-up messages arrive after agent would stop
-	while (true) {
-		let hasMoreToolCalls = true;
-
-		// Inner loop: process tool calls and steering messages
-		while (hasMoreToolCalls || pendingMessages.length > 0) {
-			// Yield at the top of each iteration to prevent busy-wait when
-			// the agent loop is executing tool calls back-to-back.
-			await yieldIfDue();
-			if (!firstTurn) {
-				stream.push({ type: "turn_start" });
-			} else {
-				firstTurn = false;
-			}
+	let deadlineTimer: ReturnType<typeof setTimeout> | undefined;
+	if (config.deadline !== undefined) {
+		const deadlineAbortController = new AbortController();
+		const delay = config.deadline - Date.now();
+		if (delay <= 0) {
+			deadlineAbortController.abort("Deadline exceeded");
+		} else {
+			deadlineTimer = setTimeout(() => {
+				deadlineAbortController.abort("Deadline exceeded");
+			}, delay);
+		}
+		signal = signal ? AbortSignal.any([signal, deadlineAbortController.signal]) : deadlineAbortController.signal;
+	}
 
-			// Process pending messages (inject before next assistant response)
-			if (pendingMessages.length > 0) {
-				for (const message of pendingMessages) {
-					stream.push({ type: "message_start", message });
-					stream.push({ type: "message_end", message });
-					currentContext.messages.push(message);
-					newMessages.push(message);
+	try {
+		let firstTurn = true;
+		if (isDeadlineExceeded(config.deadline)) {
+			endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+			return;
+		}
+		// Check for steering messages at start (user may have typed while waiting).
+		// Skip when the run is already externally aborted — dequeuing would strand
+		// the messages in a run that is about to die.
+		let pendingMessages: AgentMessage[] = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+		let harmonyRetryAttempt = 0;
+		let harmonyTruncateResumeCount = 0;
+		let pausedTurnContinuations = 0;
+
+		// Outer loop: continues when queued follow-up messages arrive after agent would stop
+		while (true) {
+			let hasMoreToolCalls = true;
+
+			// Inner loop: process tool calls and steering messages
+			while (hasMoreToolCalls || pendingMessages.length > 0) {
+				if (isDeadlineExceeded(config.deadline)) {
+					endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+					return;
+				}
+				// Yield at the top of each iteration to prevent busy-wait when
+				// the agent loop is executing tool calls back-to-back.
+				await yieldIfDue();
+				if (!firstTurn) {
+					stream.push({ type: "turn_start" });
+				} else {
+					firstTurn = false;
 				}
-				pendingMessages = [];
-			}
 
-			// Refresh prompt/tool context from live state before each model call
-			if (config.syncContextBeforeModelCall) {
-				await config.syncContextBeforeModelCall(currentContext);
-			}
+				// Process pending messages (inject before next assistant response)
+				if (pendingMessages.length > 0) {
+					for (const message of pendingMessages) {
+						stream.push({ type: "message_start", message });
+						stream.push({ type: "message_end", message });
+						currentContext.messages.push(message);
+						newMessages.push(message);
+					}
+					pendingMessages = [];
+				}
 
-			// Stream assistant response
-			let recovered: HarmonyRecoveredToolCall | undefined;
-			let message: AssistantMessage;
-			try {
-				message = await streamAssistantResponse(
-					currentContext,
-					config,
-					signal,
-					stream,
-					telemetry,
-					invokeAgentSpan,
-					stepCounter,
-					streamFn,
-					harmonyRetryAttempt,
-				);
-				harmonyRetryAttempt = 0;
-				harmonyTruncateResumeCount = 0;
-			} catch (err) {
-				if (!(err instanceof HarmonyLeakInterruption)) throw err;
-				if (err.recovered) {
-					if (harmonyTruncateResumeCount >= 2) {
-						await emitHarmonyAudit(config, err, "escalated", harmonyRetryAttempt);
-						throw new Error(
-							`GPT-5 Harmony leak recurred after truncate-and-resume recovery (${signalListLabel(err.detection.signals)}).`,
-						);
+				// Refresh prompt/tool context from live state before each model call
+				if (config.syncContextBeforeModelCall) {
+					await config.syncContextBeforeModelCall(currentContext);
+				}
+
+				// Stream assistant response
+				let recovered: HarmonyRecoveredToolCall | undefined;
+				let message: AssistantMessage;
+				try {
+					message = await streamAssistantResponse(
+						currentContext,
+						config,
+						signal,
+						stream,
+						telemetry,
+						invokeAgentSpan,
+						stepCounter,
+						streamFn,
+						harmonyRetryAttempt,
+					);
+					harmonyRetryAttempt = 0;
+					harmonyTruncateResumeCount = 0;
+				} catch (err) {
+					if (!(err instanceof HarmonyLeakInterruption)) throw err;
+					if (err.recovered) {
+						if (harmonyTruncateResumeCount >= 2) {
+							await emitHarmonyAudit(config, err, "escalated", harmonyRetryAttempt);
+							throw new Error(
+								`GPT-5 Harmony leak recurred after truncate-and-resume recovery (${signalListLabel(err.detection.signals)}).`,
+							);
+						}
+						harmonyTruncateResumeCount++;
+						recovered = err.recovered;
+						message = recovered.message;
+						await emitHarmonyAudit(config, err, "truncate_resume", harmonyRetryAttempt);
+					} else {
+						if (harmonyRetryAttempt >= 2) {
+							await emitHarmonyAudit(config, err, "escalated", harmonyRetryAttempt);
+							throw new Error(
+								`GPT-5 Harmony leak persisted after ${harmonyRetryAttempt} retries (${signalListLabel(err.detection.signals)}).`,
+							);
+						}
+						await emitHarmonyAudit(config, err, "abort_retry", harmonyRetryAttempt);
+						harmonyRetryAttempt++;
+						continue;
 					}
-					harmonyTruncateResumeCount++;
-					recovered = err.recovered;
-					message = recovered.message;
-					await emitHarmonyAudit(config, err, "truncate_resume", harmonyRetryAttempt);
-				} else {
-					if (harmonyRetryAttempt >= 2) {
-						await emitHarmonyAudit(config, err, "escalated", harmonyRetryAttempt);
-						throw new Error(
-							`GPT-5 Harmony leak persisted after ${harmonyRetryAttempt} retries (${signalListLabel(err.detection.signals)}).`,
-						);
+				}
+				if (recovered) {
+					message = snapshotAssistantMessage(message);
+					currentContext.messages.push(message);
+					stream.push({ type: "message_start", message: snapshotAssistantMessage(message) });
+					stream.push({ type: "message_end", message: snapshotAssistantMessage(message) });
+				}
+				newMessages.push(message);
+
+				if (message.stopReason === "error" || message.stopReason === "aborted") {
+					// Create placeholder tool results for any tool calls in the aborted message
+					// This maintains the tool_use/tool_result pairing that the API requires
+					type ToolCallContent = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
+					const toolCalls = message.content.filter((c): c is ToolCallContent => c.type === "toolCall");
+					const toolResults: ToolResultMessage[] = [];
+					for (const toolCall of toolCalls) {
+						const result = createAbortedToolResult(toolCall, stream, message.stopReason, message.errorMessage);
+						currentContext.messages.push(result);
+						newMessages.push(result);
+						toolResults.push(result);
+						// The placeholder result above keeps the API's tool_use/tool_result
+						// pairing intact, but no execute_tool span is started for these
+						// calls. Mirror the run-collector entry directly so the run
+						// summary's tool counters and `coverage.toolsInvoked` reflect
+						// what the user actually saw on the wire.
+						recordSkippedTool(telemetry, {
+							toolCallId: toolCall.id,
+							toolName: toolCall.name,
+							status: message.stopReason === "aborted" ? "aborted" : "error",
+						});
 					}
-					await emitHarmonyAudit(config, err, "abort_retry", harmonyRetryAttempt);
-					harmonyRetryAttempt++;
-					continue;
+					await emitTurnEnd(stream, currentContext, message, toolResults, config, signal);
+
+					stream.push(buildAgentEndEvent(newMessages, telemetry, stepCounter.count));
+					stream.end(newMessages);
+					return;
 				}
-			}
-			if (recovered) {
-				message = snapshotAssistantMessage(message);
-				currentContext.messages.push(message);
-				stream.push({ type: "message_start", message: snapshotAssistantMessage(message) });
-				stream.push({ type: "message_end", message: snapshotAssistantMessage(message) });
-			}
-			newMessages.push(message);
 
-			if (message.stopReason === "error" || message.stopReason === "aborted") {
-				// Create placeholder tool results for any tool calls in the aborted message
-				// This maintains the tool_use/tool_result pairing that the API requires
+				// Run tools whenever the turn carries tool_use blocks AND was not truncated.
+				// `stop_reason` is provider metadata that never goes back on the wire, so it
+				// does not gate continuation validity: replaying a tool_use turn with the
+				// tool_results appended is accepted whether the turn ended on `tool_use` or
+				// `end_turn` (adaptive/interleaved-thinking Opus routinely emits tool calls
+				// under `end_turn`; verified against the live Anthropic API). The only
+				// continuation hazard is a thinking block carrying a stale/invalid signature,
+				// which `transformMessages` already neutralizes — it strips the signature on
+				// non-`toolUse` turns and the encoder downgrades the unsigned block to text,
+				// which the API accepts. So treat `stop` (end_turn/pause_turn) the same as
+				// `toolUse`. `length` (max_tokens) is the one reason we must NOT run: the
+				// trailing tool_use may be truncated with incomplete arguments — those calls
+				// are abandoned below. (`error`/`aborted` already returned above.)
 				type ToolCallContent = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
 				const toolCalls = message.content.filter((c): c is ToolCallContent => c.type === "toolCall");
-				const toolResults: ToolResultMessage[] = [];
-				for (const toolCall of toolCalls) {
-					const result = createAbortedToolResult(toolCall, stream, message.stopReason, message.errorMessage);
-					currentContext.messages.push(result);
-					newMessages.push(result);
-					toolResults.push(result);
-					// The placeholder result above keeps the API's tool_use/tool_result
-					// pairing intact, but no execute_tool span is started for these
-					// calls. Mirror the run-collector entry directly so the run
-					// summary's tool counters and `coverage.toolsInvoked` reflect
-					// what the user actually saw on the wire.
-					recordSkippedTool(telemetry, {
-						toolCallId: toolCall.id,
-						toolName: toolCall.name,
-						status: message.stopReason === "aborted" ? "aborted" : "error",
-					});
+				const runnableStop = message.stopReason === "toolUse" || message.stopReason === "stop";
+				hasMoreToolCalls = runnableStop && toolCalls.length > 0;
+
+				const deadlinePassed = isDeadlineExceeded(config.deadline);
+				if (hasMoreToolCalls && deadlinePassed) {
+					hasMoreToolCalls = false;
 				}
-				await emitTurnEnd(stream, currentContext, message, toolResults, config, signal);
 
-				stream.push(buildAgentEndEvent(newMessages, telemetry, stepCounter.count));
-				stream.end(newMessages);
-				return;
-			}
+				const toolResults: ToolResultMessage[] = [];
+				if (hasMoreToolCalls) {
+					const executionResult = await executeToolCalls(
+						currentContext,
+						message,
+						signal,
+						stream,
+						config,
+						telemetry,
+						invokeAgentSpan,
+					);
 
-			// Run tools whenever the turn carries tool_use blocks AND was not truncated.
-			// `stop_reason` is provider metadata that never goes back on the wire, so it
-			// does not gate continuation validity: replaying a tool_use turn with the
-			// tool_results appended is accepted whether the turn ended on `tool_use` or
-			// `end_turn` (adaptive/interleaved-thinking Opus routinely emits tool calls
-			// under `end_turn`; verified against the live Anthropic API). The only
-			// continuation hazard is a thinking block carrying a stale/invalid signature,
-			// which `transformMessages` already neutralizes — it strips the signature on
-			// non-`toolUse` turns and the encoder downgrades the unsigned block to text,
-			// which the API accepts. So treat `stop` (end_turn/pause_turn) the same as
-			// `toolUse`. `length` (max_tokens) is the one reason we must NOT run: the
-			// trailing tool_use may be truncated with incomplete arguments — those calls
-			// are abandoned below. (`error`/`aborted` already returned above.)
-			type ToolCallContent = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
-			const toolCalls = message.content.filter((c): c is ToolCallContent => c.type === "toolCall");
-			const runnableStop = message.stopReason === "toolUse" || message.stopReason === "stop";
-			hasMoreToolCalls = runnableStop && toolCalls.length > 0;
-
-			const toolResults: ToolResultMessage[] = [];
-			if (hasMoreToolCalls) {
-				const executionResult = await executeToolCalls(
-					currentContext,
-					message,
-					signal,
-					stream,
-					config,
-					telemetry,
-					invokeAgentSpan,
-				);
+					toolResults.push(...executionResult.toolResults);
 
-				toolResults.push(...executionResult.toolResults);
+					for (const result of toolResults) {
+						currentContext.messages.push(result);
+						newMessages.push(result);
+					}
+				} else if (toolCalls.length > 0) {
+					// Turn ended on a non-runnable reason (`length` truncation) or deadline was exceeded
+					// but left toolCall blocks behind. pair each with a placeholder result.
+					const skipReason = deadlinePassed ? "aborted" : message.stopReason === "length" ? "length" : "skipped";
+					const skipErrMsg = deadlinePassed ? "Deadline exceeded" : undefined;
+					for (const toolCall of toolCalls) {
+						const result = createAbortedToolResult(toolCall, stream, skipReason, skipErrMsg);
+						currentContext.messages.push(result);
+						newMessages.push(result);
+						toolResults.push(result);
+						recordSkippedTool(telemetry, {
+							toolCallId: toolCall.id,
+							toolName: toolCall.name,
+							status: deadlinePassed ? "aborted" : "skipped",
+						});
+					}
+					if (message.stopReason === "length" && toolResults.length > 0 && !deadlinePassed) {
+						hasMoreToolCalls = true;
+					}
+				}
 
-				for (const result of toolResults) {
-					currentContext.messages.push(result);
-					newMessages.push(result);
+				if (toolCalls.length > 0) {
+					pausedTurnContinuations = 0;
+				} else if (
+					!hasMoreToolCalls &&
+					message.stopReason === "stop" &&
+					message.stopDetails?.type === "pause_turn" &&
+					pausedTurnContinuations < MAX_PAUSED_TURN_CONTINUATIONS
+				) {
+					// Non-terminal stop: the provider ended the response but not the turn
+					// (e.g. Codex `end_turn: false` on a commentary-only progress update).
+					// Re-sample with the assistant message replayed so the model keeps
+					// working; the next round folds steering/asides in like any other
+					// mid-work turn.
+					pausedTurnContinuations++;
+					hasMoreToolCalls = true;
 				}
-			} else if (toolCalls.length > 0) {
-				// Turn ended on a non-runnable reason (`length` truncation) but left
-				// toolCall blocks behind. The trailing call's arguments may be incomplete,
-				// so don't execute or continue — pair each with a placeholder result to keep
-				// the tool_use/tool_result contract valid for any later request that
-				// replays this turn. When the truncation was `length`, surface an actionable
-				// hint so the model doesn't loop by re-emitting the same oversized payload
-				// (e.g. 1000+ line `write` content blowing past the model's output cap).
-				const skipReason = message.stopReason === "length" ? "length" : "skipped";
-				for (const toolCall of toolCalls) {
-					const result = createAbortedToolResult(toolCall, stream, skipReason);
-					currentContext.messages.push(result);
-					newMessages.push(result);
-					toolResults.push(result);
-					recordSkippedTool(telemetry, {
-						toolCallId: toolCall.id,
-						toolName: toolCall.name,
-						status: "skipped",
-					});
+
+				await emitTurnEnd(stream, currentContext, message, toolResults, config, signal);
+
+				if (isDeadlineExceeded(config.deadline)) {
+					endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+					return;
 				}
-				if (message.stopReason === "length" && toolResults.length > 0) {
-					hasMoreToolCalls = true;
+				// On external abort (user interrupt), leave the steering queue intact: the
+				// session aborts then continues, delivering the queue into a fresh run.
+				// Draining it here would inject the messages right before a model call that
+				// instantly aborts — message lands in history, agent never responds. The
+				// mid-batch interrupt poll only peeks (hasSteeringMessages), so the queue
+				// still owns every message until this dequeue.
+				const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+				if (hasMoreToolCalls) {
+					// Mid-work: fold any non-interrupting asides into the next turn alongside steering.
+					const asides = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.());
+					pendingMessages = asides.length > 0 ? [...steering, ...asides] : steering;
+				} else {
+					// Stop boundary: only steering (live user input) forces another turn here. Leave
+					// asides for the outer drain below so a passive aside can't trigger an extra model
+					// turn ahead of a queued follow-up — the outer drain batches asides + follow-ups together.
+					pendingMessages = steering;
 				}
 			}
 
-			if (toolCalls.length > 0) {
-				pausedTurnContinuations = 0;
-			} else if (
-				!hasMoreToolCalls &&
-				message.stopReason === "stop" &&
-				message.stopDetails?.type === "pause_turn" &&
-				pausedTurnContinuations < MAX_PAUSED_TURN_CONTINUATIONS
-			) {
-				// Non-terminal stop: the provider ended the response but not the turn
-				// (e.g. Codex `end_turn: false` on a commentary-only progress update).
-				// Re-sample with the assistant message replayed so the model keeps
-				// working; the next round folds steering/asides in like any other
-				// mid-work turn.
-				pausedTurnContinuations++;
-				hasMoreToolCalls = true;
+			if (isDeadlineExceeded(config.deadline)) {
+				endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+				return;
 			}
 
-			await emitTurnEnd(stream, currentContext, message, toolResults, config, signal);
-
-			// On external abort (user interrupt), leave the steering queue intact: the
-			// session aborts then continues, delivering the queue into a fresh run.
-			// Draining it here would inject the messages right before a model call that
-			// instantly aborts — message lands in history, agent never responds. The
-			// mid-batch interrupt poll only peeks (hasSteeringMessages), so the queue
-			// still owns every message until this dequeue.
-			const steering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
-			if (hasMoreToolCalls) {
-				// Mid-work: fold any non-interrupting asides into the next turn alongside steering.
-				const asides = resolveAsides(await config.getAsideMessages?.());
-				pendingMessages = asides.length > 0 ? [...steering, ...asides] : steering;
-			} else {
-				// Stop boundary: only steering (live user input) forces another turn here. Leave
-				// asides for the outer drain below so a passive aside can't trigger an extra model
-				// turn ahead of a queued follow-up — the outer drain batches asides + follow-ups together.
-				pendingMessages = steering;
+			// Agent would stop here. Drain non-interrupting asides + follow-up messages.
+			await config.onBeforeYield?.();
+
+			if (isDeadlineExceeded(config.deadline)) {
+				endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+				return;
+			}
+			// Skip queue drains when externally aborted (same stranding hazard as above).
+			// Re-poll steering too: a steer can land between the stop-boundary dequeue
+			// above and this yield point (e.g. queued while onBeforeYield ran). Without
+			// this poll it would strand in the queue until the next manual prompt.
+			const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
+			const asideMessages = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.());
+			const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.()) || [];
+			if (lateSteering.length > 0 || asideMessages.length > 0 || followUpMessages.length > 0) {
+				// Set as pending so the inner loop processes them before stopping.
+				pendingMessages = [...lateSteering, ...asideMessages, ...followUpMessages];
+				continue;
 			}
-		}
 
-		// Agent would stop here. Drain non-interrupting asides + follow-up messages.
-		await config.onBeforeYield?.();
-		// Skip queue drains when externally aborted (same stranding hazard as above).
-		// Re-poll steering too: a steer can land between the stop-boundary dequeue
-		// above and this yield point (e.g. queued while onBeforeYield ran). Without
-		// this poll it would strand in the queue until the next manual prompt.
-		const lateSteering = signal?.aborted ? [] : (await config.getSteeringMessages?.()) || [];
-		const asideMessages = signal?.aborted ? [] : resolveAsides(await config.getAsideMessages?.());
-		const followUpMessages = signal?.aborted ? [] : (await config.getFollowUpMessages?.()) || [];
-		if (lateSteering.length > 0 || asideMessages.length > 0 || followUpMessages.length > 0) {
-			// Set as pending so the inner loop processes them before stopping.
-			pendingMessages = [...lateSteering, ...asideMessages, ...followUpMessages];
-			continue;
+			// No more messages, exit
+			break;
 		}
 
-		// No more messages, exit
-		break;
+		endAgentStream(stream, newMessages, telemetry, stepCounter.count);
+	} finally {
+		if (deadlineTimer) {
+			clearTimeout(deadlineTimer);
+		}
 	}
-
-	stream.push(buildAgentEndEvent(newMessages, telemetry, stepCounter.count));
-	stream.end(newMessages);
 }
 
 async function emitHarmonyAudit(
@@ -985,17 +1049,6 @@ async function streamAssistantResponse(
 
 	const streamFunction = streamFn || streamSimple;
 
-	// Resolve API key (important for expiring tokens) — do this before resolving
-	// metadata so that the session-sticky credential recorded by getApiKey is
-	// visible to metadataResolver (e.g. for the correct account_uuid in metadata.user_id).
-	const staticApiKey = typeof config.apiKey === "string" ? config.apiKey : undefined;
-	const resolvedApiKey =
-		(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || staticApiKey;
-
-	// Re-resolve metadata after credential selection so the per-request value
-	// reflects the credential actually used, not the snapshot from AgentLoopConfig construction.
-	const resolvedMetadata = config.metadataResolver ? config.metadataResolver(config.model.provider) : config.metadata;
-
 	const dynamicToolChoice = config.getToolChoice?.();
 	const dynamicReasoning = config.getReasoning?.();
 	const dynamicDisableReasoning = config.getDisableReasoning?.();
@@ -1006,7 +1059,6 @@ async function streamAssistantResponse(
 			? AbortSignal.any([signal, harmonyAbortController.signal])
 			: harmonyAbortController.signal
 		: signal;
-	const repetitionAbortController = new AbortController();
 	// Owned tool calling: aborted by the stream wrapper when the model starts
 	// fabricating a `<tool_response>`, so the provider stops generating the rest of
 	// the hallucinated turn. Merged into the provider signal ONLY (not
@@ -1015,10 +1067,20 @@ async function streamAssistantResponse(
 	const promptToolAbortController = ownedDialect ? new AbortController() : undefined;
 	const providerAbortSignals: AbortSignal[] = [];
 	if (requestSignal) providerAbortSignals.push(requestSignal);
-	providerAbortSignals.push(repetitionAbortController.signal);
 	if (promptToolAbortController) providerAbortSignals.push(promptToolAbortController.signal);
 	const finalRequestSignal =
-		providerAbortSignals.length === 1 ? providerAbortSignals[0]! : AbortSignal.any(providerAbortSignals);
+		providerAbortSignals.length === 0
+			? undefined
+			: providerAbortSignals.length === 1
+				? providerAbortSignals[0]!
+				: AbortSignal.any(providerAbortSignals);
+	const requestApiKey = (config.getApiKey ? await config.getApiKey(config.model) : undefined) ?? config.apiKey;
+	const resolvedApiKey = await resolveApiKeyOnce(requestApiKey, finalRequestSignal);
+	const apiKey = isApiKeyResolver(requestApiKey) ? seedApiKeyResolver(resolvedApiKey, requestApiKey) : requestApiKey;
+
+	// Re-resolve metadata after credential selection so the per-request value
+	// reflects the credential actually used, not the snapshot from AgentLoopConfig construction.
+	const resolvedMetadata = config.metadataResolver ? config.metadataResolver(config.model.provider) : config.metadata;
 	const effectiveTemperature =
 		harmonyRetryAttempt > 0 && config.temperature !== undefined ? config.temperature + 0.05 : config.temperature;
 	// Owned tool calling sends no native tools, so any tool_choice would error.
@@ -1069,19 +1131,7 @@ async function streamAssistantResponse(
 		return await runInActiveSpan(chatSpan, async () => {
 			let response = await streamFunction(config.model, llmContext, {
 				...config,
-				// Hand streamSimple a resolver so its central auth-retry policy can
-				// re-resolve on 401 / usage-limit: the initial step reuses the key
-				// already resolved above (which set the session-sticky credential
-				// feeding metadataResolver), and retry steps forward the a/b/c ctx
-				// to config.getApiKey (force-refresh, then rotate). With no
-				// getApiKey hook the caller's own apiKey (string or resolver) flows
-				// through unchanged.
-				apiKey: config.getApiKey
-					? (ctx: ApiKeyResolveContext) =>
-							ctx.error === undefined
-								? resolvedApiKey
-								: Promise.resolve(config.getApiKey!(config.model.provider, ctx))
-					: config.apiKey,
+				apiKey,
 				metadata: resolvedMetadata,
 				toolChoice: effectiveToolChoice,
 				reasoning: effectiveReasoning,
@@ -1130,56 +1180,6 @@ async function streamAssistantResponse(
 				return aborted;
 			};
 
-			const finishRepetitionStream = async (
-				kind: "text" | "thinking",
-				pattern: string,
-				count: number,
-			): Promise<AssistantMessage> => {
-				repetitionAbortController.abort();
-				try {
-					const cleanup = responseIterator.return?.();
-					if (cleanup) void cleanup.catch(() => {});
-				} catch {
-					// ignore
-				}
-				if (partialMessage) {
-					truncateRepetition(partialMessage, kind, pattern);
-					partialMessage.stopReason = "error";
-					partialMessage.errorMessage = `Repetition loop detected: assistant repeated "${pattern.trim()}" ${count} times consecutively.`;
-				}
-				const finalMsg = snapshotAssistantMessage(
-					partialMessage ?? {
-						role: "assistant",
-						content: [],
-						api: config.model.api,
-						provider: config.model.provider,
-						model: config.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: "error",
-						errorMessage: `Repetition loop detected.`,
-						timestamp: Date.now(),
-					},
-				);
-				if (addedPartial) {
-					context.messages[context.messages.length - 1] = finalMsg;
-				} else {
-					context.messages.push(finalMsg);
-				}
-				if (!addedPartial) {
-					stream.push({ type: "message_start", message: snapshotAssistantMessage(finalMsg) });
-				}
-				stream.push({ type: "message_end", message: snapshotAssistantMessage(finalMsg) });
-				await finishChat(finalMsg);
-				return finalMsg;
-			};
-
 			// Set up a single abort race: register the abort listener once for the whole
 			// stream and reuse the same race promise for every iterator.next() instead of
 			// allocating Promise.withResolvers and add/removeEventListener per event.
@@ -1196,14 +1196,6 @@ async function streamAssistantResponse(
 				detachAbortListener = () => requestSignal.removeEventListener("abort", onAbort);
 			}
 
-			// Rolling tail of streamed text/thinking used for repetition-loop detection.
-			// Bounded to REPETITION_WINDOW chars and reset when the active block kind
-			// switches (text <-> thinking) so detection stays O(1) per delta and never
-			// miscounts a repeated unit across a thinking/answer boundary.
-			let repetitionTail = "";
-			let repetitionKind: "text" | "thinking" | undefined;
-			const isGeminiModel = config.model.provider.includes("google") || config.model.provider.includes("gemini");
-
 			try {
 				while (true) {
 					let next: IteratorResult<AssistantMessageEvent>;
@@ -1239,6 +1231,12 @@ async function streamAssistantResponse(
 							}
 						}
 						finalMessage = snapshotAssistantMessage(finalMessage);
+						// Expand inline macros (and any other registered rewrite) on the
+						// finalized message before it reaches the context, the UI, or tool
+						// dispatch — so a single mutation is the source of truth for all three.
+						if (config.transformAssistantMessage) {
+							await config.transformAssistantMessage(finalMessage, requestSignal);
+						}
 						if (addedPartial) {
 							context.messages[context.messages.length - 1] = finalMessage;
 						} else {
@@ -1262,9 +1260,19 @@ async function streamAssistantResponse(
 					switch (event.type) {
 						case "start":
 							partialMessage = event.partial;
-							context.messages.push(partialMessage);
-							addedPartial = true;
-							stream.push({ type: "message_start", message: snapshotAssistantMessage(partialMessage) });
+							if (addedPartial) {
+								context.messages[context.messages.length - 1] = partialMessage;
+								completedToolCallIds.clear();
+								stream.push({
+									type: "message_update",
+									assistantMessageEvent: snapshotAssistantMessageEvent(event),
+									message: snapshotAssistantMessage(partialMessage),
+								});
+							} else {
+								context.messages.push(partialMessage);
+								addedPartial = true;
+								stream.push({ type: "message_start", message: snapshotAssistantMessage(partialMessage) });
+							}
 							break;
 
 						case "text_start":
@@ -1288,27 +1296,6 @@ async function streamAssistantResponse(
 									assistantMessageEvent: snapshotAssistantMessageEvent(event),
 									message: snapshotAssistantMessage(partialMessage),
 								});
-
-								if (isGeminiModel && (event.type === "text_delta" || event.type === "thinking_delta")) {
-									const kind = event.type === "text_delta" ? "text" : "thinking";
-									if (repetitionKind !== kind) {
-										repetitionKind = kind;
-										repetitionTail = "";
-									}
-									repetitionTail += event.delta;
-									if (repetitionTail.length > REPETITION_WINDOW) {
-										repetitionTail = repetitionTail.slice(-REPETITION_WINDOW);
-									}
-									const repetition = detectRepetition(repetitionTail);
-									if (repetition) {
-										const [pattern, count] = repetition;
-										logger.warn("Repetition loop detected during assistant stream, aborting.", {
-											pattern,
-											count,
-										});
-										return await finishRepetitionStream(kind, pattern, count);
-									}
-								}
 							}
 							break;
 					}
@@ -1944,97 +1931,3 @@ function createSkippedToolResult(): AgentToolResult<any> {
 		details: {},
 	};
 }
-
-const REPETITION_WINDOW = 250;
-const REPETITION_MIN_REPEATED_CHARS = 180;
-
-function detectRepetition(text: string): [pattern: string, count: number] | null {
-	if (text.length < REPETITION_MIN_REPEATED_CHARS) return null;
-
-	const windowSize = Math.min(text.length, REPETITION_WINDOW);
-	const searchSpace = text.slice(-windowSize);
-
-	for (let len = 2; len <= 60; len++) {
-		if (searchSpace.length < len * 4) continue;
-
-		const pattern = searchSpace.slice(-len);
-		// Only treat a repeated unit as a pathological loop when it carries real
-		// linguistic content (a letter or a pictographic emoji). Runs made purely of
-		// digits, whitespace or punctuation are legitimate in tabular / hex / numeric
-		// output (e.g. "00 00 00", "0, 0, 0", "| -- | -- |") and must not trip.
-		if (!/[\p{L}\p{Extended_Pictographic}]/u.test(pattern)) continue;
-
-		let count = 0;
-		let pos = searchSpace.length;
-		while (pos >= len) {
-			const chunk = searchSpace.slice(pos - len, pos);
-			if (chunk === pattern) {
-				count++;
-				pos -= len;
-			} else {
-				break;
-			}
-		}
-
-		if (count >= 4 && len * count >= REPETITION_MIN_REPEATED_CHARS) {
-			return [pattern, count];
-		}
-	}
-	return null;
-}
-
-function truncateRepetition(message: AssistantMessage, kind: "text" | "thinking", pattern: string): void {
-	// A repetition loop streams into a single growing block (real providers) or a run
-	// of same-kind blocks (some transports), always at the tail of the message. Gather
-	// that trailing contiguous run and collapse its repeated copies down to one, so the
-	// committed transcript keeps a representative sample instead of the full runaway.
-	const matches = (block: AssistantContentBlock): boolean =>
-		kind === "text" ? block.type === "text" : block.type === "thinking";
-	const readBlock = (block: AssistantContentBlock): string =>
-		block.type === "text" ? block.text : block.type === "thinking" ? block.thinking : "";
-	const clearThinkingReplayAnchors = (block: AssistantContentBlock): void => {
-		if (block.type !== "thinking") return;
-		block.thinkingSignature = undefined;
-		block.itemId = undefined;
-	};
-	const writeBlock = (block: AssistantContentBlock, value: string): void => {
-		if (block.type === "text") {
-			block.text = value;
-		} else if (block.type === "thinking") {
-			block.thinking = value;
-			clearThinkingReplayAnchors(block);
-		}
-	};
-
-	const trailing: AssistantContentBlock[] = [];
-	for (let i = message.content.length - 1; i >= 0; i--) {
-		const block = message.content[i];
-		if (!matches(block)) break;
-		trailing.unshift(block);
-	}
-	if (trailing.length === 0) return;
-	if (kind === "thinking") {
-		for (const block of trailing) clearThinkingReplayAnchors(block);
-	}
-
-	let joined = "";
-	for (const block of trailing) joined += readBlock(block);
-
-	let kept = joined;
-	while (kept.length >= pattern.length * 2 && kept.slice(kept.length - pattern.length * 2) === pattern + pattern) {
-		kept = kept.slice(0, kept.length - pattern.length);
-	}
-
-	let remainingToRemove = joined.length - kept.length;
-	for (let i = trailing.length - 1; i >= 0 && remainingToRemove > 0; i--) {
-		const block = trailing[i];
-		const value = readBlock(block);
-		if (value.length <= remainingToRemove) {
-			remainingToRemove -= value.length;
-			writeBlock(block, "");
-		} else {
-			writeBlock(block, value.slice(0, value.length - remainingToRemove));
-			remainingToRemove = 0;
-		}
-	}
-}