@sayknow-cli/agent-core 0.3.16 → 0.4.0

package/src/agent.ts

@@ -20,6 +20,7 @@ import {
 	type ToolChoice,
 	type ToolResultMessage,
 } from "@sayknow-cli/ai";
+import { extractHttpStatusFromError } from "@sayknow-cli/utils";
 import { agentLoop, agentLoopContinue } from "./agent-loop";
 import type { AppendOnlyContextManager } from "./append-only-context";
 import type { HarmonyAuditEvent } from "./harmony-leak";
@@ -32,6 +33,12 @@ import type {
 	AgentState,
 	AgentTool,
 	AgentToolContext,
+	ManagedAttemptContinuation,
+	ManagedAttemptContinuationOwnership,
+	ManagedAttemptDecision,
+	ManagedAttemptOutcome,
+	ManagedLogicalRunId,
+	RunTerminalRequest,
 	StreamFn,
 	ToolCallContext,
 } from "./types";
@@ -88,6 +95,13 @@ function refreshToolChoiceForActiveTools(
 	return tools.some(tool => tool.name === toolName) ? toolChoice : undefined;
 }
 
+export class ManagedCursorInvariantError extends Error {
+	constructor(message: string = "Managed Cursor attempt received a provider-side tool result") {
+		super(message);
+		this.name = "ManagedCursorInvariantError";
+	}
+}
+
 export class AgentBusyError extends Error {
 	constructor(
 		message: string = "Agent is already processing. Use steer() or followUp() to queue messages, or wait for completion.",
@@ -276,6 +290,16 @@ export interface AgentOptions {
 
 export interface AgentPromptOptions {
 	toolChoice?: ToolChoice;
+	/** Disable transport replay; fallback accounting is owned by the caller. */
+	fallbackManaged?: boolean;
+	/** Called synchronously after this invocation claims the agent run, before asynchronous provider work. */
+	onRunAccepted?: () => void;
+	/** Called once immediately before every managed upstream request. */
+	nextFallbackAttempt?: AgentLoopConfig["nextFallbackAttempt"];
+	/** Called after a managed upstream request is accepted and committed. */
+	onManagedAttemptAccepted?: AgentLoopConfig["onManagedAttemptAccepted"];
+	/** Receives a discarded managed attempt without exposing assistant lifecycle events. */
+	onManagedAttemptOutcome?: AgentLoopConfig["onManagedAttemptOutcome"];
 }
 
 /** Buffered Cursor tool result with text position at time of call */
@@ -284,6 +308,11 @@ interface CursorToolResultEntry {
 	textLengthAtCall: number;
 }
 
+export type AgentQueueSnapshot = {
+	steering: AgentMessage[];
+	followUp: AgentMessage[];
+};
+
 export class Agent {
 	#state: AgentState = {
 		systemPrompt: [],
@@ -332,6 +361,8 @@ export class Agent {
 	#resolveRunningPrompt?: () => void;
 	#runSequence = 0;
 	#activeRunId?: number;
+	#continuationGeneration = 0;
+	#activeFallbackManaged = false;
 	#kimiApiFormat?: "openai" | "anthropic";
 	#preferWebsockets?: boolean;
 	#transformToolCallArguments?: (args: Record<string, unknown>, toolName: string) => Record<string, unknown>;
@@ -345,6 +376,7 @@ export class Agent {
 	#onHarmonyLeak?: (event: HarmonyAuditEvent) => void | Promise<void>;
 	#onBeforeYield?: () => Promise<void> | void;
 	#shouldPause?: AgentLoopConfig["shouldPause"];
+	#maintainContext?: AgentLoopConfig["maintainContext"];
 	#telemetry?: AgentLoopConfig["telemetry"];
 	#appendOnlyContext?: AppendOnlyContextManager;
 
@@ -354,6 +386,8 @@ export class Agent {
 
 	/** Buffered Cursor tool results with text length at time of call (for correct ordering) */
 	#cursorToolResultBuffer: CursorToolResultEntry[] = [];
+	#terminalizedLogicalRunIds = new Set<ManagedLogicalRunId>();
+	#managedLogicalRunOwner?: ManagedLogicalRunId;
 
 	streamFn: StreamFn;
 	getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
@@ -680,6 +714,10 @@ export class Agent {
 		this.#shouldPause = fn;
 	}
 
+	setMaintainContext(fn: AgentLoopConfig["maintainContext"] | undefined): void {
+		this.#maintainContext = fn;
+	}
+
 	emitExternalEvent(event: AgentEvent) {
 		switch (event.type) {
 			case "message_start":
@@ -830,7 +868,7 @@ export class Agent {
 		this.#contextRevision++;
 	}
 
-	setModel(m: Model) {
+	setModel(m: Model | undefined) {
 		this.#state.model = m;
 		this.#contextRevision++;
 	}
@@ -976,6 +1014,20 @@ export class Agent {
 		this.#followUpQueue = [...messages, ...this.#followUpQueue];
 	}
 
+	/** Snapshot both executable queues as one atomic session-level view. */
+	snapshotQueues(): AgentQueueSnapshot {
+		return {
+			steering: this.#steeringQueue.slice(),
+			followUp: this.#followUpQueue.slice(),
+		};
+	}
+
+	/** Replace both executable queues with a prior snapshot. */
+	restoreQueues(snapshot: AgentQueueSnapshot): void {
+		this.#steeringQueue = snapshot.steering.slice();
+		this.#followUpQueue = snapshot.followUp.slice();
+	}
+
 	#dequeueSteeringMessages(): AgentMessage[] {
 		if (this.#steeringMode === "one-at-a-time") {
 			if (this.#steeringQueue.length > 0) {
@@ -1088,23 +1140,30 @@ export class Agent {
 	 * #runLoop guards every state mutation with a run id.
 	 */
 	forceAbort(reason = "Force aborted"): boolean {
-		const hadActiveRun = this.#runningPrompt !== undefined || this.#state.isStreaming;
+		const runId = this.#activeRunId;
+		const managedLogicalRunId = this.#managedLogicalRunOwner;
+		const hadActiveRun = runId !== undefined && (this.#runningPrompt !== undefined || this.#state.isStreaming);
 		if (!hadActiveRun) return false;
 
 		this.#abortController?.abort(reason);
-		this.#activeRunId = undefined;
+		this.#continuationGeneration++;
 		this.#state.isStreaming = false;
 		this.#state.streamMessage = null;
 		this.#state.pendingToolCalls = new Set<string>();
 		this.#abortController = undefined;
 		this.#cursorToolResultBuffer = [];
+		this.#managedLogicalRunOwner = undefined;
 
 		const resolve = this.#resolveRunningPrompt;
 		this.#runningPrompt = undefined;
 		this.#resolveRunningPrompt = undefined;
+		this.#activeRunId = undefined;
 		resolve?.();
-
-		this.#emit({ type: "agent_end", messages: [] });
+		if (this.#activeFallbackManaged) {
+			this.requestRunTerminal(managedLogicalRunId ?? runId, { stopReason: "cancelled" });
+		} else {
+			this.#finalizeRun(runId, { type: "agent_end", messages: [] });
+		}
 		return true;
 	}
 
@@ -1112,12 +1171,56 @@ export class Agent {
 		return this.#runningPrompt ?? Promise.resolve();
 	}
 
+	/** The active per-attempt run identifier. */
+	get activeRunId(): number | undefined {
+		return this.#activeRunId;
+	}
+
+	/**
+	 * Stable identifier for the active managed logical run, shared by every retry
+	 * attempt. Pass this value to requestRunTerminal(); never retain activeRunId
+	 * for managed terminal completion.
+	 */
+	get currentManagedLogicalRunId(): ManagedLogicalRunId | undefined {
+		return this.#managedLogicalRunOwner;
+	}
+
+	/**
+	 * Request terminal completion through the single logical-run keyed finalizer.
+	 *
+	 * For managed runs, logicalRunId must be currentManagedLogicalRunId from any
+	 * attempt in the retry chain. Non-managed runs use their activeRunId. Terminal
+	 * requests with messages emit a committed message_start/message_end lifecycle
+	 * for each diagnostic before agent_end. Requests without messages (such as
+	 * cancellation) emit only agent_end.
+	 */
+	requestRunTerminal(logicalRunId: ManagedLogicalRunId, request: RunTerminalRequest): boolean {
+		if (this.#terminalizedLogicalRunIds.has(logicalRunId)) return false;
+		this.#finalizeRun(
+			logicalRunId,
+			{
+				type: "agent_end",
+				messages: request.messages ?? [],
+				...(request.stopReason === "cancelled" ? { stopReason: "cancelled" as const } : {}),
+			},
+			() => {
+				for (const message of request.messages ?? []) {
+					this.#emit({ type: "message_start", message });
+					this.appendMessage(message);
+					this.#emit({ type: "message_end", message });
+				}
+			},
+		);
+		return true;
+	}
+
 	reset() {
 		this.#state.messages = [];
 		this.#contextRevision++;
 		this.#state.isStreaming = false;
 		this.#state.streamMessage = null;
 		this.#state.pendingToolCalls = new Set<string>();
+		this.#managedLogicalRunOwner = undefined;
 		this.#state.error = undefined;
 		this.#steeringQueue = [];
 		this.#followUpQueue = [];
@@ -1170,6 +1273,10 @@ export class Agent {
 		}
 
 		assertUserImagePlaceholdersHavePayload(msgs);
+		if (this.#managedLogicalRunOwner !== undefined) {
+			this.requestRunTerminal(this.#managedLogicalRunOwner, { stopReason: "cancelled" });
+			this.#managedLogicalRunOwner = undefined;
+		}
 
 		await this.#runLoop(msgs, promptOptions);
 	}
@@ -1177,7 +1284,7 @@ export class Agent {
 	/**
 	 * Continue from current context (used for retries and resuming queued messages).
 	 */
-	async continue() {
+	async continue(options?: AgentPromptOptions) {
 		if (this.#state.isStreaming) {
 			throw new AgentBusyError();
 		}
@@ -1189,13 +1296,13 @@ export class Agent {
 		if (messages[messages.length - 1].role === "assistant") {
 			const queuedSteering = this.#dequeueSteeringMessages();
 			if (queuedSteering.length > 0) {
-				await this.#runLoop(queuedSteering, { skipInitialSteeringPoll: true });
+				await this.#runLoop(queuedSteering, { ...options, skipInitialSteeringPoll: true });
 				return;
 			}
 
 			const queuedFollowUp = this.#dequeueFollowUpMessages();
 			if (queuedFollowUp.length > 0) {
-				await this.#runLoop(queuedFollowUp);
+				await this.#runLoop(queuedFollowUp, options);
 				return;
 			}
 
@@ -1206,7 +1313,7 @@ export class Agent {
 			throw new Error("No messages to continue from");
 		}
 
-		await this.#runLoop(undefined);
+		await this.#runLoop(undefined, options);
 	}
 
 	/**
@@ -1225,18 +1332,41 @@ export class Agent {
 		this.#resolveRunningPrompt = resolve;
 
 		const runId = ++this.#runSequence;
+		const continuationGeneration = ++this.#continuationGeneration;
 		this.#activeRunId = runId;
 		const abortController = new AbortController();
 		this.#abortController = abortController;
 		this.#state.isStreaming = true;
 		this.#state.streamMessage = null;
 		this.#state.error = undefined;
-
-		// Clear Cursor tool result buffer at start of each run
+		options?.onRunAccepted?.();
+
+		const fallbackManaged = options?.fallbackManaged === true;
+		const managedLogicalRunOwner = fallbackManaged ? (this.#managedLogicalRunOwner ?? runId) : undefined;
+		const startsManagedLogicalRun = fallbackManaged && this.#managedLogicalRunOwner === undefined;
+		if (startsManagedLogicalRun) {
+			this.#managedLogicalRunOwner = managedLogicalRunOwner;
+			this.#emit({ type: "agent_start" });
+		}
+		if (fallbackManaged && this.#cursorToolResultBuffer.length > 0) {
+			const error = new ManagedCursorInvariantError(
+				"Managed Cursor attempt started with buffered provider-side tool results",
+			);
+			this.#state.isStreaming = false;
+			this.#abortController = undefined;
+			this.#activeRunId = undefined;
+			this.#runningPrompt = undefined;
+			this.#resolveRunningPrompt = undefined;
+			resolve();
+			this.requestRunTerminal(managedLogicalRunOwner ?? runId, { stopReason: "error" });
+			this.#managedLogicalRunOwner = undefined;
+			throw error;
+		}
+		// Each run gets a fresh buffer only after managed stale-state validation.
 		this.#cursorToolResultBuffer = [];
+		this.#activeFallbackManaged = fallbackManaged;
 
 		const reasoning = this.#state.thinkingLevel;
-
 		const context: AgentContext = {
 			systemPrompt: this.#state.systemPrompt,
 			messages: this.#state.messages.slice(),
@@ -1244,7 +1374,7 @@ export class Agent {
 		};
 
 		const cursorOnToolResult =
-			this.#cursorExecHandlers || this.#cursorOnToolResult
+			!fallbackManaged && (this.#cursorExecHandlers || this.#cursorOnToolResult)
 				? async (message: ToolResultMessage) => {
 						let finalMessage = message;
 						if (this.#activeRunId !== runId) {
@@ -1261,7 +1391,6 @@ export class Agent {
 								}
 							} catch {}
 						}
-						// Buffer tool result with current text length for correct ordering later.
 						// Cursor executes tools server-side during streaming, so the assistant message
 						// already incorporates results. We buffer here and emit in correct order
 						// when the assistant message ends.
@@ -1273,7 +1402,10 @@ export class Agent {
 
 		const getToolChoice = () =>
 			this.#getToolChoice?.() ?? refreshToolChoiceForActiveTools(options?.toolChoice, this.#state.tools);
-		const cursorExecHandlers = this.#cursorExecHandlersForRun(runId);
+		const cursorExecHandlers = fallbackManaged ? undefined : this.#cursorExecHandlersForRun(runId);
+		let managedDecision: ManagedAttemptDecision | undefined;
+		let managedOutcome: ManagedAttemptOutcome | undefined;
+		let maintenanceInterrupted = false;
 
 		const config: AgentLoopConfig = {
 			model,
@@ -1296,6 +1428,21 @@ export class Agent {
 			maxRetryDelayMs: this.#maxRetryDelayMs,
 			requestMaxRetries: this.#requestMaxRetries,
 			streamMaxRetries: this.#streamMaxRetries,
+			...(fallbackManaged
+				? {
+						fallbackManaged: true,
+						nextFallbackAttempt: options?.nextFallbackAttempt,
+						onManagedAttemptAccepted: options?.onManagedAttemptAccepted,
+						onManagedAttemptOutcome: async outcome => {
+							managedOutcome = outcome;
+							managedDecision = (await options?.onManagedAttemptOutcome?.(outcome)) ?? {
+								type: "terminal",
+								terminal: { stopReason: outcome.type === "run_terminal" ? outcome.reason : "error" },
+							};
+							return managedDecision;
+						},
+					}
+				: {}),
 			kimiApiFormat: this.#kimiApiFormat,
 			preferWebsockets: this.#preferWebsockets,
 			convertToLlm: this.#convertToLlm,
@@ -1314,8 +1461,8 @@ export class Agent {
 				context.systemPrompt = this.#state.systemPrompt;
 				context.tools = this.#state.tools;
 			},
-			cursorExecHandlers,
-			cursorOnToolResult,
+			...(cursorExecHandlers ? { cursorExecHandlers } : {}),
+			...(cursorOnToolResult ? { cursorOnToolResult } : {}),
 			transformToolCallArguments: this.#transformToolCallArguments,
 			intentTracing: this.#intentTracing,
 			appendOnlyContext: this.#appendOnlyContext,
@@ -1384,6 +1531,12 @@ export class Agent {
 				if (this.#activeRunId !== runId) return false;
 				return this.#shouldPause?.() === true;
 			},
+			maintainContext: this.#maintainContext
+				? async (context, lifecycle) => {
+						if (this.#activeRunId !== runId) return "not-needed";
+						return (await this.#maintainContext?.(context, lifecycle)) ?? "not-needed";
+					}
+				: undefined,
 			telemetry: this.#telemetry,
 		};
 
@@ -1391,8 +1544,8 @@ export class Agent {
 
 		try {
 			const stream = messages
-				? agentLoop(messages, context, config, abortController.signal, this.streamFn)
-				: agentLoopContinue(context, config, abortController.signal, this.streamFn);
+				? agentLoop(messages, context, config, abortController.signal, this.streamFn, !fallbackManaged)
+				: agentLoopContinue(context, config, abortController.signal, this.streamFn, !fallbackManaged);
 
 			for await (const event of stream) {
 				if (this.#activeRunId !== runId) {
@@ -1412,6 +1565,9 @@ export class Agent {
 						break;
 
 					case "message_end":
+						if (fallbackManaged && this.#cursorToolResultBuffer.length > 0) {
+							throw new ManagedCursorInvariantError();
+						}
 						partial = null;
 						// Check if this is an assistant message with buffered Cursor tool results.
 						// If so, split the message to emit tool results at the correct position.
@@ -1444,9 +1600,18 @@ export class Agent {
 						break;
 
 					case "agent_end":
+						if (fallbackManaged && managedOutcome) {
+							continue;
+						}
 						this.#state.isStreaming = false;
 						this.#state.streamMessage = null;
-						break;
+						if (event.stopReason === "maintenance") {
+							maintenanceInterrupted = true;
+							this.#emit(event);
+							continue;
+						}
+						this.#finalizeRun(managedLogicalRunOwner ?? runId, event);
+						continue;
 				}
 
 				// Emit to listeners
@@ -1456,6 +1621,15 @@ export class Agent {
 			if (this.#activeRunId !== runId) {
 				return;
 			}
+			if (managedOutcome) {
+				if (managedDecision?.type === "terminal") {
+					this.requestRunTerminal(managedLogicalRunOwner ?? runId, managedDecision.terminal);
+				} else if (managedOutcome.type === "run_terminal") {
+					this.requestRunTerminal(managedLogicalRunOwner ?? runId, { stopReason: managedOutcome.reason });
+				} else if (managedDecision?.type !== "retry" && managedDecision?.type !== "maintenance") {
+					this.#finalizeRun(managedLogicalRunOwner ?? runId);
+				}
+			}
 
 			// Handle any remaining partial message
 			if (partial && partial.role === "assistant" && Array.isArray(partial.content) && partial.content.length > 0) {
@@ -1494,23 +1668,73 @@ export class Agent {
 				},
 				stopReason: abortController.signal.aborted ? "aborted" : "error",
 				errorMessage: err?.message || String(err),
+				errorStatus: extractHttpStatusFromError({ status: err?.errorStatus }) ?? extractHttpStatusFromError(err),
 				timestamp: Date.now(),
 			} as AgentMessage;
 
-			this.appendMessage(errorMsg);
 			this.#state.error = err?.message || String(err);
-			this.#emit({ type: "agent_end", messages: [errorMsg] });
+			this.requestRunTerminal(managedLogicalRunOwner ?? runId, {
+				stopReason: abortController.signal.aborted ? "cancelled" : "error",
+				messages: [errorMsg],
+			});
 		} finally {
+			let continuation: ManagedAttemptContinuation | undefined;
+			if (
+				managedOutcome?.type !== "run_terminal" &&
+				(managedDecision?.type === "retry" || managedDecision?.type === "maintenance")
+			) {
+				continuation = managedDecision.continuation;
+			}
+			const ownership: ManagedAttemptContinuationOwnership = {
+				runId,
+				logicalRunId: managedLogicalRunOwner ?? runId,
+				generation: continuationGeneration,
+				isCurrent: () => this.#continuationGeneration === continuationGeneration && this.#activeRunId === undefined,
+			};
 			if (this.#activeRunId === runId) {
 				this.#state.isStreaming = false;
 				this.#state.streamMessage = null;
 				this.#state.pendingToolCalls = new Set<string>();
 				this.#abortController = undefined;
 				this.#activeRunId = undefined;
+				this.#activeFallbackManaged = false;
 				this.#resolveRunningPrompt?.();
 				this.#runningPrompt = undefined;
 				this.#resolveRunningPrompt = undefined;
 			}
+			if (
+				fallbackManaged &&
+				!continuation &&
+				!maintenanceInterrupted &&
+				this.#managedLogicalRunOwner === managedLogicalRunOwner
+			) {
+				this.#managedLogicalRunOwner = undefined;
+			}
+			if (continuation && ownership.isCurrent()) {
+				try {
+					await continuation(ownership);
+					if (
+						managedDecision?.type === "maintenance" &&
+						this.#terminalizedLogicalRunIds.has(managedLogicalRunOwner ?? runId) &&
+						this.#managedLogicalRunOwner === managedLogicalRunOwner
+					) {
+						this.#managedLogicalRunOwner = undefined;
+					}
+					if (
+						managedDecision?.type !== "maintenance" &&
+						this.#activeRunId === undefined &&
+						this.#managedLogicalRunOwner === managedLogicalRunOwner
+					) {
+						this.#managedLogicalRunOwner = undefined;
+					}
+				} catch (err) {
+					if (ownership.isCurrent()) {
+						this.#state.error = err instanceof Error ? err.message : String(err);
+						this.requestRunTerminal(managedLogicalRunOwner ?? runId, { stopReason: "error" });
+						if (this.#managedLogicalRunOwner === managedLogicalRunOwner) this.#managedLogicalRunOwner = undefined;
+					}
+				}
+			}
 		}
 	}
 
@@ -1521,6 +1745,20 @@ export class Agent {
 	}
 
 	/** Calculate total text length from an assistant message's content blocks */
+	#finalizeRun(
+		logicalRunId: ManagedLogicalRunId,
+		event?: Extract<AgentEvent, { type: "agent_end" }>,
+		beforeEvent?: () => void,
+	): void {
+		if (this.#terminalizedLogicalRunIds.has(logicalRunId)) return;
+		this.#terminalizedLogicalRunIds.add(logicalRunId);
+		if (this.#terminalizedLogicalRunIds.size > 256) {
+			this.#terminalizedLogicalRunIds.delete(this.#terminalizedLogicalRunIds.values().next().value!);
+		}
+		beforeEvent?.();
+		if (event) this.#emit(event);
+	}
+
 	#getAssistantTextLength(message: AgentMessage | null): number {
 		if (message?.role !== "assistant" || !Array.isArray(message.content)) {
 			return 0;