diff --git a/packages/agent/docs/agent-harness.md b/packages/agent/docs/agent-harness.md new file mode 100644 index 00000000..8455e63f --- /dev/null +++ b/packages/agent/docs/agent-harness.md @@ -0,0 +1,136 @@ +# AgentHarness lifecycle + +`AgentHarness` is the orchestration layer above the low-level `Agent`. It owns session persistence, runtime configuration, resource resolution, operation locking, and extension-facing mutation semantics. + +## State model + +The harness separates state into four categories. + +### Harness config + +Harness config is the latest runtime configuration set by the application or extensions: + +- model +- thinking level +- tools +- active tool names +- resources or resource provider +- system prompt or system prompt provider + +Getters return harness config. They do not return the snapshot used by an in-flight provider request. + +Setters update harness config immediately, including while a turn is in flight. Changes affect the next turn snapshot, not the currently running provider request. + +### Turn snapshot + +A turn snapshot is the concrete state used for one LLM turn. It is created by `createTurnState()` and contains: + +- persisted session messages +- resolved resources +- resolved system prompt +- model +- thinking level +- active tools + +Static option values are used directly. Provider callbacks are invoked once per `createTurnState()` call. All logic for that turn uses the same snapshot. + +### Session + +The session contains persisted entries only. Session reads return persisted state and do not include queued writes. + +### Pending session writes + +Session writes requested while an operation is active are queued as pending session writes. Pending writes are visible through an explicit pending-writes API, not through normal session reads. + +Pending session writes are always persisted. They are flushed at save points, at operation settlement, and in failure cleanup. + +## Operation phases + +The harness has an explicit phase: + +```ts +type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry"; +``` + +Structural operations require `phase === "idle"` and synchronously set the phase before the first `await`: + +- `prompt` +- `skill` +- `promptFromTemplate` +- `compact` +- `navigateTree` + +Starting another structural operation while the harness is not idle throws. + +The following operations are allowed during a turn where appropriate: + +- `steer` +- `followUp` +- `nextTurn` +- `abort` +- runtime config setters +- session facade writes + +## Turn execution + +`prompt`, `skill`, and `promptFromTemplate` follow the same flow: + +1. Assert idle and set phase to `"turn"`. +2. Create a turn snapshot with `createTurnState()`. +3. Derive invocation text from that snapshot. +4. Execute the turn with `executeTurn()`. + +`skill` and `promptFromTemplate` resolve their resource from the same snapshot that is passed to the turn. They do not resolve resources separately. + +`nextTurn` queues arbitrary `AgentMessage`s for the next user-initiated turn. Queued messages are inserted before the new user message. + +## Save points + +A save point occurs after an assistant turn and its tool-result messages have completed. + +At a save point the harness: + +1. flushes pending session writes after the agent-emitted messages for that turn +2. creates a fresh turn snapshot if the low-level loop may continue +3. applies the fresh context/model/reasoning state before the next provider request + +This lets model, thinking level, tool, resource, and system prompt changes made during a turn affect the next turn in the same run, while never mutating an in-flight provider request. The loop callbacks are not recreated at save points. + +No state refresh is needed on `agent_end` except flushing leftover pending session writes and clearing the operation phase. + +## Session facade + +Extensions and callbacks interact with a harness-scoped session facade rather than the raw session. + +Reads delegate to persisted session state. Writes behave as follows: + +- idle: persist immediately +- busy: enqueue as pending session writes + +The facade exposes pending writes explicitly for diagnostics and UI: + +```ts +getPendingWrites(): readonly PendingSessionWrite[] +``` + +Agent-emitted messages are persisted on `message_end` to preserve transcript ordering. Pending extension/session writes flush after those messages at save points. + +## Extension context + +Event payloads describe what is happening. Harness getters describe latest config for future snapshots. + +Event contexts expose the harness and session facade. Events that belong to a turn also expose the immutable turn snapshot used for that turn. Extensions may update harness config at any time; updates affect the next snapshot. + +## Abort + +Abort is allowed during a turn. It aborts the low-level run and clears low-level steering/follow-up queues. + +Abort does not discard pending session writes. Pending writes flush at the next save point if reached, at `agent_end`, or in operation failure cleanup. + +## Compaction and tree navigation + +Compaction and tree navigation are structural session mutations. + +They are allowed only while idle and are not queued. They operate on persisted session state. The next prompt creates a fresh turn snapshot. + +Branch summary generation is part of the tree navigation operation. diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 958aad15..7226082a 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -153,13 +153,15 @@ function createAgentStream(): EventStream { * Main loop logic shared by agentLoop and agentLoopContinue. */ async function runLoop( - currentContext: AgentContext, + initialContext: AgentContext, newMessages: AgentMessage[], - config: AgentLoopConfig, + initialConfig: AgentLoopConfig, signal: AbortSignal | undefined, emit: AgentEventSink, streamFn?: StreamFn, ): Promise { + let currentContext = initialContext; + let config = initialConfig; let firstTurn = true; // Check for steering messages at start (user may have typed while waiting) let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; @@ -215,6 +217,27 @@ async function runLoop( await emit({ type: "turn_end", message, toolResults }); + const nextTurnContext = { + message, + toolResults, + context: currentContext, + newMessages, + }; + const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext); + if (nextTurnSnapshot) { + currentContext = nextTurnSnapshot.context ?? currentContext; + config = { + ...config, + model: nextTurnSnapshot.model ?? config.model, + reasoning: + nextTurnSnapshot.thinkingLevel === undefined + ? config.reasoning + : nextTurnSnapshot.thinkingLevel === "off" + ? undefined + : nextTurnSnapshot.thinkingLevel, + }; + } + if ( await config.shouldStopAfterTurn?.({ message, diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 12aa5911..6eafd030 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -15,6 +15,7 @@ import type { AgentContext, AgentEvent, AgentLoopConfig, + AgentLoopTurnUpdate, AgentMessage, AgentState, AgentTool, @@ -52,7 +53,7 @@ const DEFAULT_MODEL = { maxTokens: 0, } satisfies Model; -type QueueMode = "all" | "one-at-a-time"; +export type QueueMode = "all" | "one-at-a-time"; type MutableAgentState = Omit & { isStreaming: boolean; @@ -101,6 +102,9 @@ export interface AgentOptions { onResponse?: SimpleStreamOptions["onResponse"]; beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; + prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; steeringMode?: QueueMode; followUpMode?: QueueMode; sessionId?: string; @@ -175,6 +179,9 @@ export class Agent { context: AfterToolCallContext, signal?: AbortSignal, ) => Promise; + public prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; private activeRun?: ActiveRun; /** Session identifier forwarded to providers for cache-aware backends. */ public sessionId?: string; @@ -197,6 +204,7 @@ export class Agent { this.onResponse = options.onResponse; this.beforeToolCall = options.beforeToolCall; this.afterToolCall = options.afterToolCall; + this.prepareNextTurn = options.prepareNextTurn; this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); this.sessionId = options.sessionId; @@ -421,6 +429,7 @@ export class Agent { toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, + prepareNextTurn: this.prepareNextTurn ? async () => await this.prepareNextTurn?.(this.signal) : undefined, convertToLlm: this.convertToLlm, transformContext: this.transformContext, getApiKey: this.getApiKey, diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index 5d1d516d..dd7bf90b 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -1,6 +1,5 @@ -import { randomUUID } from "node:crypto"; -import type { AssistantMessage, ImageContent, Model } from "@earendil-works/pi-ai"; -import { Agent } from "../agent.js"; +import type { AssistantMessage, ImageContent, Model, UserMessage } from "@earendil-works/pi-ai"; +import { Agent, type QueueMode } from "../agent.js"; import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../types.js"; import { collectEntriesForBranchSummary, generateBranchSummary } from "./compaction/branch-summarization.js"; import { compact, DEFAULT_COMPACTION_SETTINGS, prepareCompaction } from "./compaction/compaction.js"; @@ -8,20 +7,20 @@ import { formatPromptTemplateInvocation } from "./prompt-templates.js"; import { formatSkillInvocation } from "./skills.js"; import type { AbortResult, - AgentHarnessContext, - AgentHarnessConversationState, AgentHarnessEvent, AgentHarnessEventResultMap, - AgentHarnessOperationState, AgentHarnessOptions, AgentHarnessOwnEvent, + AgentHarnessPhase, AgentHarnessResources, + AgentHarnessTurnState, ExecutionEnv, NavigateTreeResult, + PendingSessionWrite, Session, } from "./types.js"; -function createUserMessage(text: string, images?: ImageContent[]): AgentMessage { +function createUserMessage(text: string, images?: ImageContent[]): UserMessage { const content: Array<{ type: "text"; text: string } | ImageContent> = [{ type: "text", text }]; if (images) content.push(...images); return { role: "user", content, timestamp: Date.now() }; @@ -30,19 +29,21 @@ function createUserMessage(text: string, images?: ImageContent[]): AgentMessage export class AgentHarness { readonly agent: Agent; readonly env: ExecutionEnv; - readonly conversation: AgentHarnessConversationState; - readonly operation: AgentHarnessOperationState; - private session: Session; - private resourcesInput?: AgentHarnessOptions["resources"]; + private model: Model; + private thinkingLevel: ThinkingLevel; + private activeToolNames: string[]; + private nextTurnQueue: AgentMessage[] = []; + private phase: AgentHarnessPhase = "idle"; + private steerQueue: UserMessage[] = []; + private followUpQueue: UserMessage[] = []; + private pendingSessionWrites: PendingSessionWrite[] = []; + private resources?: AgentHarnessOptions["resources"]; private systemPrompt: AgentHarnessOptions["systemPrompt"]; - private requestAuth?: AgentHarnessOptions["requestAuth"]; - private toolRegistry = new Map(); + private getApiKeyAndHeaders?: AgentHarnessOptions["getApiKeyAndHeaders"]; + private tools = new Map(); private listeners = new Set<(event: AgentHarnessEvent, signal?: AbortSignal) => Promise | void>(); - private hooks = new Map< - keyof AgentHarnessEventResultMap, - Set<(event: any, ctx: AgentHarnessContext) => Promise | any> - >(); + private hooks = new Map Promise | any>>(); constructor(options: AgentHarnessOptions) { this.agent = new Agent({ @@ -51,119 +52,79 @@ export class AgentHarness { thinkingLevel: options.thinkingLevel, tools: options.tools ?? [], }, + steeringMode: options.steeringMode, + followUpMode: options.followUpMode, }); this.env = options.env; this.session = options.session; - this.resourcesInput = options.resources; + this.resources = options.resources; this.systemPrompt = options.systemPrompt; - this.requestAuth = options.requestAuth; + this.getApiKeyAndHeaders = options.getApiKeyAndHeaders; for (const tool of this.agent.state.tools) { - this.toolRegistry.set(tool.name, tool); + this.tools.set(tool.name, tool); } - this.conversation = { - session: options.session, - model: options.model, - thinkingLevel: options.thinkingLevel ?? this.agent.state.thinkingLevel, - activeToolNames: options.activeToolNames ?? this.agent.state.tools.map((tool) => tool.name), - nextTurnQueue: [], - }; - this.agent.state.model = this.conversation.model; - this.agent.state.thinkingLevel = this.conversation.thinkingLevel; + this.model = options.model; + this.thinkingLevel = options.thinkingLevel ?? this.agent.state.thinkingLevel; + this.activeToolNames = options.activeToolNames ?? this.agent.state.tools.map((tool) => tool.name); + this.agent.state.model = this.model; + this.agent.state.thinkingLevel = this.thinkingLevel; this.agent.getApiKey = async (provider) => { - const model = this.conversation.model; - if (!this.requestAuth || model.provider !== provider) return undefined; - return (await this.requestAuth(model))?.apiKey; + const model = this.model; + if (!this.getApiKeyAndHeaders || model.provider !== provider) return undefined; + return (await this.getApiKeyAndHeaders(model))?.apiKey; }; - this.operation = { - idle: true, - abortRequested: false, - steerQueue: [], - followUpQueue: [], - pendingMutations: { - appendMessages: [], - }, - }; - this.agent.transformContext = async (messages, signal) => { - const result = await this.emitHook("context", { type: "context", messages: [...messages] }, signal); + this.agent.transformContext = async (messages) => { + const result = await this.emitHook({ type: "context", messages: [...messages] }); return result?.messages ?? messages; }; - this.agent.beforeToolCall = async ({ toolCall, args }, signal) => { - const result = await this.emitHook( - "tool_call", - { - type: "tool_call", - toolCallId: toolCall.id, - toolName: toolCall.name, - input: args as Record, - }, - signal, - ); + this.agent.beforeToolCall = async ({ toolCall, args }) => { + const result = await this.emitHook({ + type: "tool_call", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + }); return result ? { block: result.block, reason: result.reason } : undefined; }; - this.agent.afterToolCall = async ({ toolCall, args, result, isError }, signal) => { - const patch = await this.emitHook( - "tool_result", - { - type: "tool_result", - toolCallId: toolCall.id, - toolName: toolCall.name, - input: args as Record, - content: result.content, - details: result.details, - isError, - }, - signal, - ); + this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => { + const patch = await this.emitHook({ + type: "tool_result", + toolCallId: toolCall.id, + toolName: toolCall.name, + input: args as Record, + content: result.content, + details: result.details, + isError, + }); return patch ? { content: patch.content, details: patch.details, isError: patch.isError, terminate: patch.terminate } : undefined; }; this.agent.onPayload = async (payload) => { - const result = await this.emitHook("before_provider_request", { type: "before_provider_request", payload }); + const result = await this.emitHook({ type: "before_provider_request", payload }); return result?.payload ?? payload; }; this.agent.onResponse = async (response) => { const headers = { ...(response.headers as Record) }; await this.emitOwn({ type: "after_provider_response", status: response.status, headers }, this.agent.signal); }; + this.agent.prepareNextTurn = async () => { + await this.flushPendingSessionWrites(); + const turnState = await this.createTurnState(); + this.applyTurnState(turnState); + return { + context: { + systemPrompt: turnState.systemPrompt, + messages: turnState.messages.slice(), + tools: turnState.activeTools.slice(), + }, + model: turnState.model, + thinkingLevel: turnState.thinkingLevel, + }; + }; this.agent.subscribe(async (event, signal) => { await this.handleAgentEvent(event, signal); }); - void this.syncFromTree(); - } - - private createContext(signal?: AbortSignal): AgentHarnessContext { - return { - env: this.env, - conversation: this.conversation, - operation: this.operation, - abortSignal: signal, - }; - } - - private async resolveResources(signal?: AbortSignal): Promise { - return typeof this.resourcesInput === "function" - ? await this.resourcesInput(this.createContext(signal)) - : (this.resourcesInput ?? {}); - } - - private getActiveTools(): AgentTool[] { - return this.conversation.activeToolNames - .map((name) => this.toolRegistry.get(name)) - .filter((tool): tool is AgentTool => tool !== undefined); - } - - private async resolveSystemPrompt(resources: AgentHarnessResources): Promise { - if (!this.systemPrompt) return "You are a helpful assistant."; - if (typeof this.systemPrompt === "string") return this.systemPrompt; - return await this.systemPrompt({ - env: this.env, - session: this.session, - model: this.conversation.model, - thinkingLevel: this.conversation.thinkingLevel, - activeTools: this.getActiveTools(), - resources, - }); } private async emitOwn(event: AgentHarnessOwnEvent, signal?: AbortSignal): Promise { @@ -179,15 +140,13 @@ export class AgentHarness { } private async emitHook( - type: TType, event: Extract, - signal?: AbortSignal, ): Promise { - const handlers = this.hooks.get(type); + const handlers = this.hooks.get(event.type as TType); if (!handlers || handlers.size === 0) return undefined; let lastResult: AgentHarnessEventResultMap[TType] | undefined; for (const handler of handlers) { - const result = await handler(event, this.createContext(signal)); + const result = await handler(event); if (result !== undefined) { lastResult = result; } @@ -198,69 +157,89 @@ export class AgentHarness { private async emitQueueUpdate(): Promise { await this.emitOwn({ type: "queue_update", - steer: [...this.operation.steerQueue], - followUp: [...this.operation.followUpQueue], - nextTurn: [...this.conversation.nextTurnQueue], + steer: [...this.steerQueue], + followUp: [...this.followUpQueue], + nextTurn: [...this.nextTurnQueue], }); } - private async syncFromTree(): Promise { + private async createTurnState(): Promise { const context = await this.session.buildContext(); - this.agent.state.messages = context.messages; - if (context.model && this.conversation.model) { - // leave active model untouched; harness-level model is source of truth + const resources = typeof this.resources === "function" ? await this.resources() : (this.resources ?? {}); + const tools = [...this.tools.values()]; + const activeTools = this.activeToolNames + .map((name) => this.tools.get(name)) + .filter((tool): tool is AgentTool => tool !== undefined); + let systemPrompt = "You are a helpful assistant."; + if (typeof this.systemPrompt === "string") { + systemPrompt = this.systemPrompt; + } else if (this.systemPrompt) { + systemPrompt = await this.systemPrompt({ + env: this.env, + session: this.session, + model: this.model, + thinkingLevel: this.thinkingLevel, + activeTools, + resources, + }); } - this.agent.state.systemPrompt = await this.resolveSystemPrompt(await this.resolveResources()); + return { + messages: context.messages, + resources, + systemPrompt, + model: this.model, + thinkingLevel: this.thinkingLevel, + tools, + activeTools, + }; } - private async applyPendingMutations(): Promise { - for (const message of this.operation.pendingMutations.appendMessages) { - await this.session.appendMessage(message); - } - this.operation.pendingMutations.appendMessages = []; + private applyTurnState(turnState: AgentHarnessTurnState): void { + this.agent.state.messages = turnState.messages; + this.agent.state.systemPrompt = turnState.systemPrompt; + this.agent.state.model = turnState.model; + this.agent.state.thinkingLevel = turnState.thinkingLevel; + this.agent.state.tools = turnState.activeTools; + } - if (this.operation.pendingMutations.model) { - const model = this.operation.pendingMutations.model; - const previousModel = this.conversation.model; - this.conversation.model = model; - this.agent.state.model = model; - await this.session.appendModelChange(model.provider, model.id); - await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); - this.operation.pendingMutations.model = undefined; - } + private validateToolNames(toolNames: string[]): void { + const missing = toolNames.filter((name) => !this.tools.has(name)); + if (missing.length > 0) throw new Error(`Unknown tool(s): ${missing.join(", ")}`); + } - if (this.operation.pendingMutations.thinkingLevel !== undefined) { - const level = this.operation.pendingMutations.thinkingLevel; - const previousLevel = this.conversation.thinkingLevel; - this.conversation.thinkingLevel = level; - this.agent.state.thinkingLevel = level; - await this.session.appendThinkingLevelChange(level); - await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); - this.operation.pendingMutations.thinkingLevel = undefined; + private async flushPendingSessionWrites(): Promise { + const writes = this.pendingSessionWrites; + this.pendingSessionWrites = []; + for (const write of writes) { + if (write.type === "message") { + await this.session.appendMessage(write.message); + } else if (write.type === "model_change") { + await this.session.appendModelChange(write.provider, write.modelId); + } else if (write.type === "thinking_level_change") { + await this.session.appendThinkingLevelChange(write.thinkingLevel); + } else if (write.type === "custom") { + await this.session.appendCustomEntry(write.customType, write.data); + } else if (write.type === "custom_message") { + await this.session.appendCustomMessageEntry(write.customType, write.content, write.display, write.details); + } else if (write.type === "label") { + await this.session.appendLabel(write.targetId, write.label); + } else if (write.type === "session_info") { + await this.session.appendSessionName(write.name ?? ""); + } } - - if (this.operation.pendingMutations.activeToolNames) { - this.conversation.activeToolNames = [...this.operation.pendingMutations.activeToolNames]; - this.agent.state.tools = this.conversation.activeToolNames - .map((name) => this.toolRegistry.get(name)) - .filter((tool): tool is (typeof this.agent.state.tools)[number] => tool !== undefined); - this.operation.pendingMutations.activeToolNames = undefined; - } - - await this.syncFromTree(); } private async handleAgentEvent(event: AgentEvent, signal?: AbortSignal): Promise { await this.emitAny(event, signal); - if (event.type === "message_start") { - const steerIndex = this.operation.steerQueue.indexOf(event.message); + if (event.type === "message_start" && event.message.role === "user") { + const steerIndex = this.steerQueue.indexOf(event.message); if (steerIndex !== -1) { - this.operation.steerQueue.splice(steerIndex, 1); + this.steerQueue.splice(steerIndex, 1); await this.emitQueueUpdate(); } else { - const followUpIndex = this.operation.followUpQueue.indexOf(event.message); + const followUpIndex = this.followUpQueue.indexOf(event.message); if (followUpIndex !== -1) { - this.operation.followUpQueue.splice(followUpIndex, 1); + this.followUpQueue.splice(followUpIndex, 1); await this.emitQueueUpdate(); } } @@ -269,55 +248,47 @@ export class AgentHarness { await this.session.appendMessage(event.message); } if (event.type === "turn_end") { - const hadPendingMutations = - this.operation.pendingMutations.appendMessages.length > 0 || - this.operation.pendingMutations.model !== undefined || - this.operation.pendingMutations.thinkingLevel !== undefined || - this.operation.pendingMutations.activeToolNames !== undefined; - await this.emitOwn( - { type: "save_point", liveOperationId: this.operation.liveOperationId ?? "unknown", hadPendingMutations }, - signal, - ); - if (hadPendingMutations) { - await this.applyPendingMutations(); - } + const hadPendingMutations = this.pendingSessionWrites.length > 0; + await this.flushPendingSessionWrites(); + await this.emitOwn({ + type: "save_point", + hadPendingMutations, + }); } if (event.type === "agent_end") { - this.operation.idle = true; - this.operation.liveOperationId = undefined; - this.operation.abortRequested = false; - await this.syncFromTree(); - await this.emitOwn({ type: "settled", nextTurnCount: this.conversation.nextTurnQueue.length }, signal); + await this.flushPendingSessionWrites(); + this.phase = "idle"; + await this.emitOwn({ type: "settled", nextTurnCount: this.nextTurnQueue.length }, signal); } } - async prompt(text: string, options?: { images?: ImageContent[] }): Promise { - if (!this.operation.idle) throw new Error("AgentHarness is busy"); + private async executeTurn( + turnState: AgentHarnessTurnState, + text: string, + options?: { images?: ImageContent[] }, + ): Promise { + this.applyTurnState(turnState); const beforeLength = this.agent.state.messages.length; - this.operation.idle = false; - this.operation.liveOperationId = randomUUID(); - const resources = await this.resolveResources(this.agent.signal); let messages: AgentMessage[] = [createUserMessage(text, options?.images)]; - if (this.conversation.nextTurnQueue.length > 0) { - messages = [messages[0]!, ...this.conversation.nextTurnQueue]; - this.conversation.nextTurnQueue = []; + if (this.nextTurnQueue.length > 0) { + messages = [...this.nextTurnQueue, messages[0]!]; + this.nextTurnQueue = []; await this.emitQueueUpdate(); } - this.agent.state.systemPrompt = await this.resolveSystemPrompt(resources); - const beforeResult = await this.emitHook( - "before_agent_start", - { - type: "before_agent_start", - prompt: text, - images: options?.images, - systemPrompt: this.agent.state.systemPrompt, - resources, - }, - this.agent.signal, - ); + const beforeResult = await this.emitHook({ + type: "before_agent_start", + prompt: text, + images: options?.images, + systemPrompt: turnState.systemPrompt, + resources: turnState.resources, + }); if (beforeResult?.messages) messages = [...beforeResult.messages, ...messages]; if (beforeResult?.systemPrompt) this.agent.state.systemPrompt = beforeResult.systemPrompt; - await this.agent.prompt(messages); + try { + await this.agent.prompt(messages); + } finally { + await this.flushPendingSessionWrites(); + } let response: AssistantMessage | undefined; const newMessages = this.agent.state.messages.slice(beforeLength); for (let i = newMessages.length - 1; i >= 0; i--) { @@ -331,81 +302,98 @@ export class AgentHarness { return response; } + async prompt(text: string, options?: { images?: ImageContent[] }): Promise { + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + return await this.executeTurn(turnState, text, options); + } catch (error) { + this.phase = "idle"; + throw error; + } + } + async skill(name: string, additionalInstructions?: string): Promise { - const resources = await this.resolveResources(); - const skill = (resources.skills ?? []).find((candidate) => candidate.name === name); - if (!skill) throw new Error(`Unknown skill: ${name}`); - return await this.prompt(formatSkillInvocation(skill, additionalInstructions)); + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + const skill = (turnState.resources.skills ?? []).find((candidate) => candidate.name === name); + if (!skill) throw new Error(`Unknown skill: ${name}`); + return await this.executeTurn(turnState, formatSkillInvocation(skill, additionalInstructions)); + } catch (error) { + this.phase = "idle"; + throw error; + } } async promptFromTemplate(name: string, args: string[] = []): Promise { - const resources = await this.resolveResources(); - const template = (resources.promptTemplates ?? []).find((candidate) => candidate.name === name); - if (!template) throw new Error(`Unknown prompt template: ${name}`); - return await this.prompt(formatPromptTemplateInvocation(template, args)); + if (this.phase !== "idle") throw new Error("AgentHarness is busy"); + this.phase = "turn"; + try { + const turnState = await this.createTurnState(); + const template = (turnState.resources.promptTemplates ?? []).find((candidate) => candidate.name === name); + if (!template) throw new Error(`Unknown prompt template: ${name}`); + return await this.executeTurn(turnState, formatPromptTemplateInvocation(template, args)); + } catch (error) { + this.phase = "idle"; + throw error; + } } - steer(message: AgentMessage): void { - if (this.operation.idle) throw new Error("Cannot steer while idle"); - this.operation.steerQueue.push(message); + steer(text: string, options?: { images?: ImageContent[] }): void { + if (this.phase === "idle") throw new Error("Cannot steer while idle"); + const message = createUserMessage(text, options?.images); + this.steerQueue.push(message); this.agent.steer(message); void this.emitQueueUpdate(); } - followUp(message: AgentMessage): void { - if (this.operation.idle) throw new Error("Cannot follow up while idle"); - this.operation.followUpQueue.push(message); + followUp(text: string, options?: { images?: ImageContent[] }): void { + if (this.phase === "idle") throw new Error("Cannot follow up while idle"); + const message = createUserMessage(text, options?.images); + this.followUpQueue.push(message); this.agent.followUp(message); void this.emitQueueUpdate(); } - nextTurn(message: AgentMessage): void { - this.conversation.nextTurnQueue.push(message); + nextTurn(text: string, options?: { images?: ImageContent[] }): void { + this.nextTurnQueue.push(createUserMessage(text, options?.images)); void this.emitQueueUpdate(); } async appendMessage(message: AgentMessage): Promise { - if (this.operation.idle) { + if (this.phase === "idle") { await this.session.appendMessage(message); - await this.syncFromTree(); } else { - this.operation.pendingMutations.appendMessages.push(message); + this.pendingSessionWrites.push({ type: "message", message }); } } - async shell( - command: string, - options?: { - cwd?: string; - env?: Record; - timeout?: number; - signal?: AbortSignal; - onStdout?: (chunk: string) => void; - onStderr?: (chunk: string) => void; - }, - ): Promise<{ stdout: string; stderr: string; exitCode: number }> { - return await this.env.exec(command, options); - } - async compact( customInstructions?: string, ): Promise<{ summary: string; firstKeptEntryId: string; tokensBefore: number; details?: unknown }> { - if (!this.operation.idle) throw new Error("compact() requires idle harness"); - const model = this.conversation.model; + if (this.phase !== "idle") throw new Error("compact() requires idle harness"); + this.phase = "compaction"; + const model = this.model; if (!model) throw new Error("No model set for compaction"); - const auth = await this.requestAuth?.(model); + const auth = await this.getApiKeyAndHeaders?.(model); if (!auth) throw new Error("No auth available for compaction"); const branchEntries = await this.session.getBranch(); const preparation = prepareCompaction(branchEntries, DEFAULT_COMPACTION_SETTINGS); if (!preparation) throw new Error("Nothing to compact"); - const hookResult = await this.emitHook("session_before_compact", { + const hookResult = await this.emitHook({ type: "session_before_compact", preparation, branchEntries, customInstructions, signal: new AbortController().signal, }); - if (hookResult?.cancel) throw new Error("Compaction cancelled"); + if (hookResult?.cancel) { + this.phase = "idle"; + throw new Error("Compaction cancelled"); + } const provided = hookResult?.compaction; const result = provided ?? @@ -416,7 +404,7 @@ export class AgentHarness { auth.headers, customInstructions, undefined, - this.conversation.thinkingLevel, + this.thinkingLevel, )); const entryId = await this.session.appendCompaction( result.summary, @@ -426,10 +414,10 @@ export class AgentHarness { provided !== undefined, ); const entry = await this.session.getEntry(entryId); - await this.syncFromTree(); if (entry?.type === "compaction") { await this.emitOwn({ type: "session_compact", compactionEntry: entry, fromHook: provided !== undefined }); } + this.phase = "idle"; return result; } @@ -437,9 +425,13 @@ export class AgentHarness { targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, ): Promise { - if (!this.operation.idle) throw new Error("navigateTree() requires idle harness"); + if (this.phase !== "idle") throw new Error("navigateTree() requires idle harness"); + this.phase = "branch_summary"; const oldLeafId = await this.session.getLeafId(); - if (oldLeafId === targetId) return { cancelled: false }; + if (oldLeafId === targetId) { + this.phase = "idle"; + return { cancelled: false }; + } const targetEntry = await this.session.getEntry(targetId); if (!targetEntry) throw new Error(`Entry ${targetId} not found`); const { entries, commonAncestorId } = await collectEntriesForBranchSummary(this.session, oldLeafId, targetId); @@ -454,19 +446,22 @@ export class AgentHarness { label: options?.label, }; const signal = new AbortController().signal; - const hookResult = await this.emitHook("session_before_tree", { + const hookResult = await this.emitHook({ type: "session_before_tree", preparation, signal, }); - if (hookResult?.cancel) return { cancelled: true }; + if (hookResult?.cancel) { + this.phase = "idle"; + return { cancelled: true }; + } let summaryEntry: any | undefined; let summaryText: string | undefined = hookResult?.summary?.summary; let summaryDetails: unknown = hookResult?.summary?.details; if (!summaryText && options?.summarize && entries.length > 0) { - const model = this.conversation.model; + const model = this.model; if (!model) throw new Error("No model set for branch summary"); - const auth = await this.requestAuth?.(model); + const auth = await this.getApiKeyAndHeaders?.(model); if (!auth) throw new Error("No auth available for branch summary"); const branchSummary = await generateBranchSummary(entries, { model, @@ -476,7 +471,10 @@ export class AgentHarness { customInstructions: hookResult?.customInstructions ?? options?.customInstructions, replaceInstructions: hookResult?.replaceInstructions ?? options?.replaceInstructions, }); - if (branchSummary.aborted) return { cancelled: true }; + if (branchSummary.aborted) { + this.phase = "idle"; + return { cancelled: true }; + } if (branchSummary.error) throw new Error(branchSummary.error); summaryText = branchSummary.summary; summaryDetails = { @@ -521,7 +519,6 @@ export class AgentHarness { if (summaryId) { summaryEntry = await this.session.getEntry(summaryId); } - await this.syncFromTree(); await this.emitOwn({ type: "session_tree", newLeafId: await this.session.getLeafId(), @@ -529,48 +526,80 @@ export class AgentHarness { summaryEntry, fromHook: hookResult?.summary !== undefined, }); + this.phase = "idle"; return { cancelled: false, editorText, summaryEntry }; } async setModel(model: Model): Promise { - if (this.operation.idle) { - const previousModel = this.conversation.model; - this.conversation.model = model; + const previousModel = this.model; + this.model = model; + if (this.phase === "idle") { this.agent.state.model = model; await this.session.appendModelChange(model.provider, model.id); - await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); } else { - this.operation.pendingMutations.model = model; + this.pendingSessionWrites.push({ type: "model_change", provider: model.provider, modelId: model.id }); } + await this.emitOwn({ type: "model_select", model, previousModel, source: "set" }); } async setThinkingLevel(level: ThinkingLevel): Promise { - if (this.operation.idle) { - const previousLevel = this.conversation.thinkingLevel; - this.conversation.thinkingLevel = level; + const previousLevel = this.thinkingLevel; + this.thinkingLevel = level; + if (this.phase === "idle") { this.agent.state.thinkingLevel = level; await this.session.appendThinkingLevelChange(level); - await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); } else { - this.operation.pendingMutations.thinkingLevel = level; + this.pendingSessionWrites.push({ type: "thinking_level_change", thinkingLevel: level }); } + await this.emitOwn({ type: "thinking_level_select", level, previousLevel }); } async setActiveTools(toolNames: string[]): Promise { - if (this.operation.idle) { - this.conversation.activeToolNames = [...toolNames]; - this.agent.state.tools = this.getActiveTools(); + this.validateToolNames(toolNames); + this.activeToolNames = [...toolNames]; + if (this.phase === "idle") { + this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!); + } + } + + get steeringMode(): QueueMode { + return this.agent.steeringMode; + } + + set steeringMode(mode: QueueMode) { + this.agent.steeringMode = mode; + } + + get followUpMode(): QueueMode { + return this.agent.followUpMode; + } + + set followUpMode(mode: QueueMode) { + this.agent.followUpMode = mode; + } + + async setResources(resources: AgentHarnessResources): Promise { + this.resources = resources; + } + + async setTools(tools: AgentTool[], activeToolNames?: string[]): Promise { + this.tools = new Map(tools.map((tool) => [tool.name, tool])); + if (activeToolNames) { + this.validateToolNames(activeToolNames); + this.activeToolNames = [...activeToolNames]; } else { - this.operation.pendingMutations.activeToolNames = [...toolNames]; + this.validateToolNames(this.activeToolNames); + } + if (this.phase === "idle") { + this.agent.state.tools = this.activeToolNames.map((name) => this.tools.get(name)!); } } async abort(): Promise { - this.operation.abortRequested = true; - const clearedSteer = [...this.operation.steerQueue]; - const clearedFollowUp = [...this.operation.followUpQueue]; - this.operation.steerQueue = []; - this.operation.followUpQueue = []; + const clearedSteer = [...this.steerQueue]; + const clearedFollowUp = [...this.followUpQueue]; + this.steerQueue = []; + this.followUpQueue = []; this.agent.clearAllQueues(); await this.emitQueueUpdate(); this.agent.abort(); @@ -591,8 +620,7 @@ export class AgentHarness { on( type: TType, handler: ( - event: Extract, - ctx: AgentHarnessContext, + event: Extract, ) => Promise | AgentHarnessEventResultMap[TType], ): () => void { let handlers = this.hooks.get(type); diff --git a/packages/agent/src/harness/types.ts b/packages/agent/src/harness/types.ts index 5c768d78..4c368ce0 100644 --- a/packages/agent/src/harness/types.ts +++ b/packages/agent/src/harness/types.ts @@ -1,4 +1,5 @@ import type { ImageContent, Model, TextContent } from "@earendil-works/pi-ai"; +import type { QueueMode } from "../agent.js"; import type { AgentEvent, AgentMessage, AgentTool, ThinkingLevel } from "../index.js"; import type { Session } from "./session/session.js"; @@ -286,43 +287,22 @@ export interface JsonlSessionListOptions { export interface JsonlSessionRepoApi extends SessionRepo {} -export interface AgentHarnessPendingMutations { - appendMessages: AgentMessage[]; - model?: Model; - thinkingLevel?: ThinkingLevel; - activeToolNames?: string[]; -} +export type AgentHarnessPhase = "idle" | "turn" | "compaction" | "branch_summary" | "retry"; -export interface AgentHarnessConversationState { - session: Session; +export type PendingSessionWrite = SessionTreeEntry extends infer TEntry + ? TEntry extends SessionTreeEntry + ? Omit + : never + : never; + +export interface AgentHarnessTurnState { + messages: AgentMessage[]; + resources: AgentHarnessResources; + systemPrompt: string; model: Model; thinkingLevel: ThinkingLevel; - activeToolNames: string[]; - nextTurnQueue: AgentMessage[]; -} - -export interface AgentHarnessOperationState { - idle: boolean; - liveOperationId?: string; - abortRequested: boolean; - steerQueue: AgentMessage[]; - followUpQueue: AgentMessage[]; - pendingMutations: AgentHarnessPendingMutations; -} - -export interface SavePointSnapshot { - messages: AgentMessage[]; - model: Model | undefined; - thinkingLevel: ThinkingLevel; - activeToolNames: string[]; - systemPrompt: string; -} - -export interface AgentHarnessContext { - env: ExecutionEnv; - conversation: AgentHarnessConversationState; - operation: AgentHarnessOperationState; - abortSignal?: AbortSignal; + tools: AgentTool[]; + activeTools: AgentTool[]; } export interface QueueUpdateEvent { @@ -334,7 +314,6 @@ export interface QueueUpdateEvent { export interface SavePointEvent { type: "save_point"; - liveOperationId: string; hadPendingMutations: boolean; } @@ -586,9 +565,7 @@ export interface AgentHarnessOptions { env: ExecutionEnv; session: Session; tools?: AgentTool[]; - resources?: - | AgentHarnessResources - | ((context: AgentHarnessContext) => AgentHarnessResources | Promise); + resources?: AgentHarnessResources | (() => AgentHarnessResources | Promise); systemPrompt?: | string | ((context: { @@ -599,10 +576,14 @@ export interface AgentHarnessOptions { activeTools: AgentTool[]; resources: AgentHarnessResources; }) => string | Promise); - requestAuth?: (model: Model) => Promise<{ apiKey: string; headers?: Record } | undefined>; + getApiKeyAndHeaders?: ( + model: Model, + ) => Promise<{ apiKey: string; headers?: Record } | undefined>; model: Model; thinkingLevel?: ThinkingLevel; activeToolNames?: string[]; + steeringMode?: QueueMode; + followUpMode?: QueueMode; } export type { AgentHarness } from "./agent-harness.js"; diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 24b12e83..285c1b02 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -112,6 +112,18 @@ export interface ShouldStopAfterTurnContext { newMessages: AgentMessage[]; } +/** Replacement runtime state used by the agent loop before starting another provider request. */ +export interface AgentLoopTurnUpdate { + /** Context for the next provider request. */ + context?: AgentContext; + /** Model for the next provider request. */ + model?: Model; + /** Thinking level for the next provider request. */ + thinkingLevel?: ThinkingLevel; +} + +export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {} + export interface AgentLoopConfig extends SimpleStreamOptions { model: Model; @@ -187,6 +199,15 @@ export interface AgentLoopConfig extends SimpleStreamOptions { */ shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise; + /** + * Called after `turn_end` and before the loop decides whether another provider request should start. + * Return replacement context/model/thinking state to affect the next turn in this run. + * Return undefined to keep using the current context/config. + */ + prepareNextTurn?: ( + context: PrepareNextTurnContext, + ) => AgentLoopTurnUpdate | undefined | Promise; + /** * Returns steering messages to inject into the conversation mid-run. * diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 26c4fcf7..261b42fc 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -894,6 +894,79 @@ describe("agentLoop with AgentMessage", () => { expect(parallelObserved).toBe(true); }); + it("should use prepareNextTurn snapshot before continuing", async () => { + const toolSchema = Type.Object({ value: Type.String() }); + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + }; + }, + }; + const context: AgentContext = { + systemPrompt: "first prompt", + messages: [], + tools: [tool], + }; + let convertedSecondTurnSystemPrompt = ""; + let prepared = false; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + prepareNextTurn: async ({ context: currentContext }) => { + if (prepared) return undefined; + prepared = true; + return { + context: { + systemPrompt: "second prompt", + messages: currentContext.messages.slice(), + tools: currentContext.tools, + }, + }; + }, + }; + + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, (_model, ctx) => { + llmCalls++; + if (llmCalls === 2) { + convertedSecondTurnSystemPrompt = ctx.systemPrompt ?? ""; + } + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + if (llmCalls === 1) { + mockStream.push({ + type: "done", + reason: "toolUse", + message: createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ), + }); + } else { + mockStream.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "done" }]), + }); + } + }); + return mockStream; + }); + + for await (const _event of stream) { + // consume + } + + expect(llmCalls).toBe(2); + expect(convertedSecondTurnSystemPrompt).toBe("second prompt"); + }); + it("should stop after the current turn when shouldStopAfterTurn returns true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; diff --git a/packages/agent/test/harness/factory.test.ts b/packages/agent/test/harness/factory.test.ts index cac0e882..d8863482 100644 --- a/packages/agent/test/harness/factory.test.ts +++ b/packages/agent/test/harness/factory.test.ts @@ -18,10 +18,21 @@ describe("harness factories", () => { const session = createSession(new InMemorySessionStorage()); const env = new NodeExecutionEnv({ cwd: process.cwd() }); const initialModel = getModel("anthropic", "claude-sonnet-4-5"); - const harness = createAgentHarness({ env, session, model: initialModel, systemPrompt: "You are helpful." }); + const harness = createAgentHarness({ + env, + session, + model: initialModel, + systemPrompt: "You are helpful.", + steeringMode: "all", + followUpMode: "all", + }); expect(harness.env).toBe(env); - expect(harness.conversation.session).toBe(session); - expect(harness.conversation.model).toBe(initialModel); expect(harness.agent.state.model).toBe(initialModel); + expect(harness.steeringMode).toBe("all"); + expect(harness.followUpMode).toBe("all"); + harness.steeringMode = "one-at-a-time"; + harness.followUpMode = "one-at-a-time"; + expect(harness.agent.steeringMode).toBe("one-at-a-time"); + expect(harness.agent.followUpMode).toBe("one-at-a-time"); }); });