diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index e722fa9f5..bb3e89b62 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,30 @@ ## [Unreleased] +### Breaking Changes + +- Removed `AgentOptions.shouldStopAfterTurn` and `AgentLoopConfig.shouldStopAfterTurn`. Use `finishTurn` and return `{ action: "end" }` to stop after the completed turn: + + ```ts + // Before + shouldStopAfterTurn: async (turn, signal) => await shouldStop(turn, signal), + + // After + finishTurn: async (turn, signal) => { + // shouldStopAfterTurn previously ran only for normal responses. + if (turn.message.stopReason === "error" || turn.message.stopReason === "aborted") return; + return (await shouldStop(turn, signal)) ? { action: "end" } : undefined; + }, + ``` + + `finishTurn` runs after the assistant and all tool results are finalized but before `turn_end`; its decision is applied after `turn_end`. It also runs for error and aborted responses, whose decisions are ignored because those responses remain hard exits. The guard in the migration preserves the old hook's normal-response-only invocation, including avoiding predicate side effects on hard exits. Returning `{ action: "end" }` leaves steering and follow-up queues untouched and skips `prepareNextTurn`. + +### Added + +- Added `prepareRequest`, which runs before every provider request, including the first. For example, return `{ context: { ...context, messages: persistedMessages } }` to install canonical context after already-selected input is emitted without introducing another queue poll. +- Added `finishTurn`, which runs after assistant/tool-result finalization and before `turn_end` for normal, error, and aborted responses. Return `{ action: "end" }` to end a normal run after `turn_end`, or `undefined` to preserve normal scheduling. `{ action: "continue" }` ensures one next provider request: existing tool-result, steering, or follow-up scheduling can satisfy that request without adding another one; otherwise the loop makes one context-only request. Error and aborted responses remain hard exits. +- Added `Agent.peekQueuedMessages()` to preview the next queue-selected batch without consuming it. + ### Fixed - Fixed harness reads misclassifying text files beginning with `GIF` as images ([#9755](https://github.com/earendil-works/pi/issues/9755)). diff --git a/packages/agent/README.md b/packages/agent/README.md index b107e2fff..aa43845c6 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -127,38 +127,68 @@ The `beforeToolCall` hook runs after `tool_execution_start` and validated argume Tools, blocked `beforeToolCall` results, and `afterToolCall` overrides can return `terminate: true` to hint that the automatic follow-up LLM call should be skipped. The loop only stops early when every finalized tool result in that batch sets `terminate: true`. Mixed batches continue normally. -The `Agent` class accepts `shouldStopAfterTurn` in `AgentOptions`. Low-level loop callers can set the same hook in `AgentLoopConfig`: - -```typescript -const stream = agentLoop( - prompts, - context, - { - model, - convertToLlm, - shouldStopAfterTurn: async ({ message, toolResults, context, newMessages }) => { - return shouldCompactBeforeNextTurn(context.messages); - }, - }, - undefined, - models.streamSimple.bind(models), -); -``` - -`shouldStopAfterTurn` runs after `turn_end` is emitted and after the assistant response and any tool executions have completed normally. If it returns `true`, the loop emits `agent_end` and exits before polling steering or follow-up queues, and before starting another LLM call. It does not abort the provider stream, does not cancel running tools, and does not alter the assistant message stop reason. The `AgentOptions` callback also receives the active run's `AbortSignal` as its second argument. - When you use the `Agent` class, assistant `message_end` processing is treated as a barrier before tool preflight begins. That means `beforeToolCall` sees agent state that already includes the assistant message that requested the tool call. -### continue() Event Sequence +### Request preparation and turn finalization -`continue()` resumes from existing context without adding a new message. Use it for retries after errors. +`prepareRequest` runs immediately before every conversational provider request, including the first. Use it to install canonical persisted context after pending input has been emitted: ```typescript -// After an error, retry from current state -await agent.continue(); +agent.prepareRequest = async ({ context }) => ({ + context: { ...context, messages: await session.loadModelContext() }, +}); ``` -The last message in context must be `user` or `toolResult` (not `assistant`). +`prepareRequest` does not poll queues. Steering queued while it runs waits for the next normal steering poll. + +`finishTurn` runs after the assistant and all tool results are finalized, but before `turn_end`. It runs for normal, error, and aborted responses: + +```typescript +agent.finishTurn = async ({ message }) => { + if (message.stopReason === "error" || message.stopReason === "aborted") return; + if (shouldEndRun(message)) return { action: "end" }; + return needsAnotherResponse(message) ? { action: "continue" } : undefined; +}; +``` + +Returning `undefined` preserves normal scheduling. `{ action: "end" }` stops immediately after `turn_end`, before polling steering or follow-up queues or preparing another request. On a normal response, `{ action: "continue" }` ensures one next provider request. If tool results, steering, or a follow-up already cause that request, they satisfy the decision and no additional request is made; otherwise the loop makes one context-only request. Error and aborted responses remain hard exits, so their decisions are ignored. `finishTurn` runs again after the next request, so returning `{ action: "continue" }` unconditionally creates an endless loop. + +To migrate from the removed `shouldStopAfterTurn`, return `{ action: "end" }`. Guard error and aborted responses to preserve the old hook's normal-response-only invocation, especially when the predicate has side effects or assumes a successful response: + +```typescript +finishTurn: async (turn, signal) => { + if (turn.message.stopReason === "error" || turn.message.stopReason === "aborted") return; + return (await shouldStop(turn, signal)) ? { action: "end" } : undefined; +}, +``` + +Each provider turn follows this lifecycle: + +```text +selected input events +→ prepareRequest +→ provider response +→ tool results +→ finishTurn +→ turn_end +→ existing continuation scheduling or agent_end +``` + +### continue() and queued input + +`continue()` retains its existing queue behavior. Empty and system-only transcripts reject without consuming queues. A non-assistant tail continues from existing context: steering is polled at startup, while follow-up input waits until the response naturally stops. + +```typescript +agent.followUp({ role: "user", content: "After the retry", timestamp: Date.now() }); +await agent.continue(); // The first request retries the existing user/toolResult tail. +``` + +An assistant tail cannot be sent directly, so `continue()` falls back to one queued steering batch, then one queued follow-up batch. Queue mode still controls whether that selected batch contains one message or all messages: + +```typescript +agent.steer({ role: "user", content: "Continue from here", timestamp: Date.now() }); +await agent.continue(); // Uses the queued message only because the tail is assistant. +``` ### Event Types @@ -233,9 +263,16 @@ const agent = new Agent({ } }, - // Stop gracefully after a completed turn, before queued messages are polled. - shouldStopAfterTurn: async ({ context }, signal) => { - return shouldCompactBeforeNextTurn(context.messages, signal); + // Rebuild finalized context immediately before every provider request. + prepareRequest: async ({ context }, signal) => { + return { context: { ...context, messages: await loadCanonicalMessages(signal) } }; + }, + + // Finalize a completed turn before turn_end is emitted. + // `continue` ensures one next request; existing tool/queue scheduling can satisfy it. + // `end` ends this run after turn_end without polling queues. + finishTurn: async ({ message, toolResults }, signal) => { + return shouldContinue(message, toolResults) ? { action: "continue" } : undefined; }, // Custom thinking budgets for token-based providers @@ -298,7 +335,7 @@ await agent.prompt("What's in this image?", [ // AgentMessage directly await agent.prompt({ role: "user", content: "Hello", timestamp: Date.now() }); -// Continue from current context (last message must be user or toolResult) +// Continue existing non-assistant input; an assistant tail may use queued input as fallback await agent.continue(); ``` @@ -311,9 +348,13 @@ agent.state.tools = [myTool]; agent.toolExecution = "sequential"; agent.beforeToolCall = async ({ toolCall }) => undefined; agent.afterToolCall = async ({ toolCall, result }) => undefined; -agent.shouldStopAfterTurn = async ({ context }) => shouldCompactBeforeNextTurn(context.messages); +agent.prepareRequest = async ({ context }) => ({ + context: { ...context, messages: await loadCanonicalMessages() }, +}); +agent.finishTurn = async () => undefined; agent.state.messages = newMessages; // top-level array is copied agent.state.messages.push(message); +const nextQueuedMessages = agent.peekQueuedMessages(); // respects queue modes; does not consume agent.reset(); ``` diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 735225f9a..e4fd515df 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -170,6 +170,7 @@ async function runLoop( let currentContext = initialContext; let config = initialConfig; let lastCompletedTurn: PrepareNextTurnContext | undefined; + let explicitContinuation = false; // Check for steering messages at start (user may have typed while waiting) let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; @@ -214,11 +215,40 @@ async function runLoop( } pendingMessages = []; + const requestUpdate = await config.prepareRequest?.( + { + context: currentContext, + model: config.model, + thinkingLevel: config.reasoning ?? "off", + }, + signal, + ); + if (requestUpdate) { + currentContext = requestUpdate.context ?? currentContext; + config = { + ...config, + model: requestUpdate.model ?? config.model, + reasoning: + requestUpdate.thinkingLevel === undefined + ? config.reasoning + : requestUpdate.thinkingLevel === "off" + ? undefined + : requestUpdate.thinkingLevel, + }; + } + // Stream assistant response const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFunction); newMessages.push(message); if (message.stopReason === "error" || message.stopReason === "aborted") { + lastCompletedTurn = { + message, + toolResults: [], + context: currentContext, + newMessages, + }; + await config.finishTurn?.(lastCompletedTurn, signal); await emit({ type: "turn_end", message, toolResults: [] }); await emit({ type: "agent_end", messages: newMessages }); return; @@ -246,31 +276,42 @@ async function runLoop( } } - await emit({ type: "turn_end", message, toolResults }); - lastCompletedTurn = { message, toolResults, context: currentContext, newMessages, }; + const decision = await config.finishTurn?.(lastCompletedTurn, signal); + await emit({ type: "turn_end", message, toolResults }); - if (await config.shouldStopAfterTurn?.(lastCompletedTurn)) { + if (decision?.action === "end") { await emit({ type: "agent_end", messages: newMessages }); return; } + explicitContinuation = decision?.action === "continue"; pendingMessages = (await config.getSteeringMessages?.()) || []; + if (hasMoreToolCalls || pendingMessages.length > 0) { + explicitContinuation = false; + } } // Agent would stop here. Check for follow-up messages. const followUpMessages = (await config.getFollowUpMessages?.()) || []; if (followUpMessages.length > 0) { // Set as pending so inner loop processes them + explicitContinuation = false; pendingMessages = followUpMessages; continue; } + // No natural request was selected, so fulfill the continuation decision with one context-only turn. + if (explicitContinuation) { + explicitContinuation = false; + continue; + } + // No more messages, exit break; } diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 78d287e47..92af2c791 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -25,9 +25,10 @@ import type { AgentTool, BeforeToolCallContext, BeforeToolCallResult, + FinishTurn, PrepareNextTurnContext, + PrepareRequest, QueueMode, - ShouldStopAfterTurnContext, StreamFn, ToolExecutionMode, } from "./types.ts"; @@ -120,7 +121,8 @@ export interface AgentOptions { onResponse?: SimpleStreamOptions["onResponse"]; beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; - shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext, signal?: AbortSignal) => boolean | Promise; + finishTurn?: FinishTurn; + prepareRequest?: PrepareRequest; prepareNextTurn?: ( signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; @@ -153,19 +155,16 @@ class PendingMessageQueue { return this.messages.length > 0; } - drain(): AgentMessage[] { - if (this.mode === "all") { - const drained = this.messages.slice(); - this.messages = []; - return drained; - } - + peek(): AgentMessage[] { + if (this.mode === "all") return this.messages.slice(); const first = this.messages[0]; - if (!first) { - return []; - } - this.messages = this.messages.slice(1); - return [first]; + return first ? [first] : []; + } + + drain(): AgentMessage[] { + const drained = this.peek(); + this.messages = this.messages.slice(drained.length); + return drained; } clear(): void { @@ -205,10 +204,8 @@ export class Agent { context: AfterToolCallContext, signal?: AbortSignal, ) => Promise; - public shouldStopAfterTurn?: ( - context: ShouldStopAfterTurnContext, - signal?: AbortSignal, - ) => boolean | Promise; + public finishTurn?: FinishTurn; + public prepareRequest?: PrepareRequest; public prepareNextTurn?: ( signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; @@ -240,7 +237,8 @@ export class Agent { this.onResponse = runtimeOptions.onResponse; this.beforeToolCall = runtimeOptions.beforeToolCall; this.afterToolCall = runtimeOptions.afterToolCall; - this.shouldStopAfterTurn = runtimeOptions.shouldStopAfterTurn; + this.finishTurn = runtimeOptions.finishTurn; + this.prepareRequest = runtimeOptions.prepareRequest; this.prepareNextTurn = runtimeOptions.prepareNextTurn; this.prepareNextTurnWithContext = runtimeOptions.prepareNextTurnWithContext; this.steeringQueue = new PendingMessageQueue(runtimeOptions.steeringMode ?? "one-at-a-time"); @@ -325,6 +323,12 @@ export class Agent { return this.steeringQueue.hasItems() || this.followUpQueue.hasItems(); } + /** Preview the messages selected for the next turn without consuming them. */ + peekQueuedMessages(): AgentMessage[] { + const steering = this.steeringQueue.peek(); + return steering.length > 0 ? steering : this.followUpQueue.peek(); + } + /** Active abort signal for the current run, if any. */ get signal(): AbortSignal | undefined { return this.activeRun?.abortController.signal; @@ -459,7 +463,6 @@ export class Agent { private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig { let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; - const shouldStopAfterTurn = this.shouldStopAfterTurn; return { model: this._state.model, reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, @@ -472,9 +475,8 @@ export class Agent { toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, - shouldStopAfterTurn: shouldStopAfterTurn - ? async (context) => await shouldStopAfterTurn(context, this.signal) - : undefined, + finishTurn: this.finishTurn, + prepareRequest: this.prepareRequest, prepareNextTurn: this.prepareNextTurnWithContext || this.prepareNextTurn ? async (context) => { diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index fb718b635..8da843f56 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -127,11 +127,11 @@ export interface AfterToolCallContext { context: AgentContext; } -/** Context passed to `shouldStopAfterTurn`. */ -export interface ShouldStopAfterTurnContext { +/** Context passed to completed-turn callbacks. */ +export interface AgentTurnContext { /** The assistant message that completed the turn. */ message: AssistantMessage; - /** Tool result messages passed to the preceding `turn_end` event. */ + /** Tool result messages emitted for the completed turn. */ toolResults: ToolResultMessage[]; /** Current agent context after the turn's assistant message and tool results have been appended. */ context: AgentContext; @@ -139,6 +139,20 @@ export interface ShouldStopAfterTurnContext { newMessages: AgentMessage[]; } +/** Decision returned by {@link FinishTurn}. Returning undefined preserves normal scheduling. */ +export type AgentTurnDecision = { action: "continue" } | { action: "end" }; + +/** + * Called after a completed assistant turn and all of its tool-result messages, but before `turn_end`. + * On a normal turn, `{ action: "continue" }` ensures one next provider request. Tool-result, steering, or + * follow-up scheduling can satisfy that request and adds no extra request; otherwise the loop continues once + * with the current context. Error and aborted responses remain hard exits. + */ +export type FinishTurn = ( + turn: AgentTurnContext, + signal?: AbortSignal, +) => AgentTurnDecision | void | Promise | Promise; + /** Replacement runtime state used by the agent loop before starting another provider request. */ export interface AgentLoopTurnUpdate { /** Context for the next provider request. */ @@ -151,7 +165,26 @@ export interface AgentLoopTurnUpdate { thinkingLevel?: ThinkingLevel; } -export interface PrepareNextTurnContext extends ShouldStopAfterTurnContext {} +/** Runtime state available immediately before a conversational provider request. */ +export interface PrepareRequestContext { + context: AgentContext; + model: Model; + thinkingLevel: ThinkingLevel; +} + +/** Replacement runtime state for the provider request being prepared. */ +export type AgentRequestUpdate = Omit; + +/** + * Called immediately before every conversational provider request, including the first. + * Pending messages have already been appended and emitted when this callback runs. + */ +export type PrepareRequest = ( + request: PrepareRequestContext, + signal?: AbortSignal, +) => AgentRequestUpdate | void | Promise | Promise; + +export interface PrepareNextTurnContext extends AgentTurnContext {} export interface AgentLoopConfig extends SimpleStreamOptions { model: Model; @@ -217,17 +250,21 @@ export interface AgentLoopConfig extends SimpleStreamOptions { getApiKey?: (provider: string) => Promise | string | undefined; /** - * Called after each turn fully completes and `turn_end` has been emitted. - * - * If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues, - * without starting another LLM call. The current assistant response and any tool executions finish normally. - * This callback sees the completed-turn context and runs before `prepareNextTurn`. - * - * Use this to request a graceful stop after the current turn, e.g. before context gets too full. - * - * Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence. + * Called after the assistant message and all tool-result messages have been emitted, immediately before `turn_end`. + * `{ action: "end" }` ends the run without polling queues or preparing another request. + * On a normal turn, `{ action: "continue" }` ensures one next provider request. Tool-result, steering, or + * follow-up scheduling can satisfy that request and adds no extra request; otherwise the loop continues once + * with the current context. Returning undefined preserves normal scheduling. Error and aborted responses remain + * hard exits. */ - shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise; + finishTurn?: FinishTurn; + + /** + * Called immediately before every conversational provider request, including the first. + * Pending messages have already been appended. The returned context, model, and thinking level + * replace the runtime values for this and later requests in the run. This hook does not poll queues. + */ + prepareRequest?: PrepareRequest; /** * Called after `turn_end` when the loop will continue, immediately before the next turn starts. @@ -241,7 +278,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { /** * Returns steering messages to inject into the conversation mid-run. * - * Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first. + * Called after the current assistant turn finishes executing its tool calls, unless `finishTurn` ends the run. * If messages are returned, they are added to the context before the next LLM call. * Tool calls from the current assistant message are not skipped. * diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 862b2a6ef..018b0f7d4 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -8,7 +8,7 @@ import { } from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; -import { agentLoop, agentLoopContinue } from "../src/agent-loop.ts"; +import { agentLoop, agentLoopContinue, runAgentLoop } from "../src/agent-loop.ts"; import { setDefaultStreamFn } from "../src/index.ts"; import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.ts"; @@ -1057,6 +1057,401 @@ describe("agentLoop with AgentMessage", () => { expect(parallelObserved).toBe(true); }); + it("runs finishTurn after tool-result messages and before turn_end", 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: params.value }], + details: { value: params.value }, + terminate: true, + }; + }, + }; + const ordering: string[] = []; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: ({ context, toolResults }) => { + ordering.push("finishTurn"); + expect(toolResults).toHaveLength(1); + expect(context.messages.at(-1)?.role).toBe("toolResult"); + }, + }; + + await runAgentLoop( + [createUserMessage("echo")], + { messages: [], tools: [tool] }, + config, + (event) => { + if (event.type === "message_end") ordering.push(`message_end:${event.message.role}`); + if (event.type === "turn_end") ordering.push("turn_end"); + }, + undefined, + () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ), + }); + }); + return stream; + }, + ); + + expect(ordering.slice(-3)).toEqual(["message_end:toolResult", "finishTurn", "turn_end"]); + }); + + it.each(["error", "aborted"] as const)( + "runs finishTurn for a %s assistant before turn_end without changing the hard exit", + async (reason) => { + const ordering: string[] = []; + let providerCalls = 0; + let steeringPolls = 0; + let followUpPolls = 0; + await runAgentLoop( + [createUserMessage("run")], + { messages: [], tools: [] }, + { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: ({ message }) => { + expect(message.stopReason).toBe(reason); + ordering.push("finishTurn"); + return { action: "continue" }; + }, + getSteeringMessages: async () => { + steeringPolls++; + return []; + }, + getFollowUpMessages: async () => { + followUpPolls++; + return [createUserMessage("queued")]; + }, + }, + (event) => { + if (event.type === "turn_end") ordering.push("turn_end"); + }, + undefined, + () => { + providerCalls++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "error", + reason, + error: { + ...createAssistantMessage([], reason), + errorMessage: reason, + }, + }); + }); + return stream; + }, + ); + + expect(ordering).toEqual(["finishTurn", "turn_end"]); + expect(providerCalls).toBe(1); + expect(steeringPolls).toBe(1); + expect(followUpPolls).toBe(0); + }, + ); + + it("action:end skips queue polling and next-turn preparation", async () => { + const toolSchema = Type.Object({}); + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop tool", + parameters: toolSchema, + async execute() { + return { content: [{ type: "text", text: "done" }], details: undefined }; + }, + }; + let providerCalls = 0; + let steeringPolls = 0; + let followUpPolls = 0; + let prepareNextTurnCalls = 0; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: () => ({ action: "end" }), + prepareNextTurn: () => { + prepareNextTurnCalls++; + return undefined; + }, + getSteeringMessages: async () => { + steeringPolls++; + return []; + }, + getFollowUpMessages: async () => { + followUpPolls++; + return [createUserMessage("queued")]; + }, + }; + + const stream = agentLoop([createUserMessage("run")], { messages: [], tools: [tool] }, config, undefined, () => { + providerCalls++; + const response = new MockAssistantStream(); + queueMicrotask(() => { + response.push({ + type: "done", + reason: "toolUse", + message: createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "noop", arguments: {} }], + "toolUse", + ), + }); + }); + return response; + }); + await stream.result(); + + expect(providerCalls).toBe(1); + expect(steeringPolls).toBe(1); + expect(followUpPolls).toBe(0); + expect(prepareNextTurnCalls).toBe(0); + }); + + it("makes exactly one context-only request when no natural request satisfies continuation", async () => { + let providerCalls = 0; + let finishCalls = 0; + const stream = agentLoop( + [createUserMessage("run")], + { messages: [], tools: [] }, + { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: () => { + finishCalls++; + return finishCalls === 1 ? { action: "continue" } : undefined; + }, + }, + undefined, + () => { + providerCalls++; + const response = new MockAssistantStream(); + queueMicrotask(() => { + response.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: `response ${providerCalls}` }]), + }); + }); + return response; + }, + ); + await stream.result(); + + expect(providerCalls).toBe(2); + expect(finishCalls).toBe(2); + }); + + it("lets a natural tool-result request satisfy continuation", async () => { + const toolSchema = Type.Object({}); + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop tool", + parameters: toolSchema, + async execute() { + return { content: [{ type: "text", text: "done" }], details: undefined }; + }, + }; + let providerCalls = 0; + let finishCalls = 0; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: () => { + finishCalls++; + return finishCalls === 1 ? { action: "continue" } : undefined; + }, + }; + + const stream = agentLoop([createUserMessage("run")], { messages: [], tools: [tool] }, config, undefined, () => { + providerCalls++; + const response = new MockAssistantStream(); + queueMicrotask(() => { + const message = + providerCalls === 1 + ? createAssistantMessage([{ type: "toolCall", id: "tool-1", name: "noop", arguments: {} }], "toolUse") + : createAssistantMessage([{ type: "text", text: "done" }]); + response.push({ type: "done", reason: providerCalls === 1 ? "toolUse" : "stop", message }); + }); + return response; + }); + await stream.result(); + + expect(providerCalls).toBe(2); + expect(finishCalls).toBe(2); + }); + + it.each(["steering", "follow-up"] as const)("lets a natural %s request satisfy continuation", async (queueKind) => { + const queuedMessage = createUserMessage(queueKind); + let providerCalls = 0; + let finishCalls = 0; + let steeringPolls = 0; + let followUpDelivered = false; + const secondRequestUsers: string[] = []; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + finishTurn: () => { + finishCalls++; + return finishCalls === 1 ? { action: "continue" } : undefined; + }, + getSteeringMessages: async () => { + steeringPolls++; + return queueKind === "steering" && steeringPolls === 2 ? [queuedMessage] : []; + }, + getFollowUpMessages: async () => { + if (queueKind !== "follow-up" || followUpDelivered) return []; + followUpDelivered = true; + return [queuedMessage]; + }, + }; + + const stream = agentLoop( + [createUserMessage("run")], + { messages: [], tools: [] }, + config, + undefined, + (_model, context) => { + providerCalls++; + if (providerCalls === 2) { + secondRequestUsers.push( + ...context.messages.flatMap((message) => + message.role === "user" && typeof message.content === "string" ? [message.content] : [], + ), + ); + } + const response = new MockAssistantStream(); + queueMicrotask(() => { + response.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "done" }]), + }); + }); + return response; + }, + ); + await stream.result(); + + expect(providerCalls).toBe(2); + expect(finishCalls).toBe(2); + expect(secondRequestUsers).toContain(queueKind); + }); + + it("prepares the initial request after pending messages and can replace request state", async () => { + const replacementModel = { ...createModel(), id: "replacement", name: "replacement" }; + const canonicalMessage = createUserMessage("canonical projection"); + const steeringMessage = createUserMessage("steering"); + const completedMessages: AgentMessage[] = []; + let steeringDelivered = false; + let prepareCalls = 0; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + getSteeringMessages: async () => { + if (steeringDelivered) return []; + steeringDelivered = true; + return [steeringMessage]; + }, + prepareRequest: ({ context }) => { + prepareCalls++; + expect(completedMessages).toContain(steeringMessage); + expect(context.messages).toContain(steeringMessage); + return { + context: { ...context, messages: [canonicalMessage] }, + model: replacementModel, + thinkingLevel: "high", + }; + }, + }; + + await runAgentLoop( + [createUserMessage("prompt")], + { messages: [], tools: [] }, + config, + (event) => { + if (event.type === "message_end") completedMessages.push(event.message); + }, + undefined, + (model, context, options) => { + expect(model).toBe(replacementModel); + expect(context.messages).toEqual([canonicalMessage]); + expect(options?.reasoning).toBe("high"); + const response = new MockAssistantStream(); + queueMicrotask(() => { + response.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "done" }]), + }); + }); + return response; + }, + ); + + expect(prepareCalls).toBe(1); + }); + + it("does not poll steering after prepareRequest", async () => { + const queued: AgentMessage[] = []; + const lateSteering = createUserMessage("late steering"); + const requestIncludedSteering: boolean[] = []; + let requestPreparations = 0; + let steeringPolls = 0; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + getSteeringMessages: async () => { + steeringPolls++; + return queued.splice(0); + }, + prepareRequest: () => { + requestPreparations++; + if (requestPreparations === 1) queued.push(lateSteering); + }, + }; + + const stream = agentLoop( + [createUserMessage("run")], + { messages: [], tools: [] }, + config, + undefined, + (_model, context) => { + requestIncludedSteering.push(context.messages.includes(lateSteering)); + const response = new MockAssistantStream(); + queueMicrotask(() => { + response.push({ + type: "done", + reason: "stop", + message: createAssistantMessage([{ type: "text", text: "done" }]), + }); + }); + return response; + }, + ); + await stream.result(); + + expect(requestIncludedSteering).toEqual([false, true]); + expect(requestPreparations).toBe(2); + // Startup, post-turn delivery, then the final natural-stop check. + expect(steeringPolls).toBe(3); + }); + it("should use prepareNextTurn snapshot before continuing", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool = { @@ -1134,7 +1529,60 @@ describe("agentLoop with AgentMessage", () => { expect(convertedSecondTurnHasUpdate).toBe(true); }); - it("should stop after the current turn when shouldStopAfterTurn returns true", async () => { + it("picks up steering queued during prepareNextTurn before the next request", async () => { + const toolSchema = Type.Object({}); + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop tool", + parameters: toolSchema, + async execute() { + return { content: [{ type: "text", text: "done" }], details: undefined }; + }, + }; + const queued: AgentMessage[] = []; + const lateSteering = createUserMessage("late steering"); + let providerCalls = 0; + let secondRequestIncludedSteering = false; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + prepareNextTurn: () => { + queued.push(lateSteering); + return undefined; + }, + getSteeringMessages: async () => queued.splice(0), + }; + + const stream = agentLoop( + [createUserMessage("run")], + { messages: [], tools: [tool] }, + config, + undefined, + (_model, context) => { + providerCalls++; + if (providerCalls === 2) secondRequestIncludedSteering = context.messages.includes(lateSteering); + const response = new MockAssistantStream(); + queueMicrotask(() => { + const message = + providerCalls === 1 + ? createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "noop", arguments: {} }], + "toolUse", + ) + : createAssistantMessage([{ type: "text", text: "done" }]); + response.push({ type: "done", reason: providerCalls === 1 ? "toolUse" : "stop", message }); + }); + return response; + }, + ); + await stream.result(); + + expect(providerCalls).toBe(2); + expect(secondRequestIncludedSteering).toBe(true); + }); + + it("action:end receives finalized turn context and stops before queue polling", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; const tool: AgentTool = { @@ -1163,6 +1611,12 @@ describe("agentLoop with AgentMessage", () => { const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, + finishTurn: async ({ message, toolResults, context }) => { + expect(message.role).toBe("assistant"); + callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId); + callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role); + return { action: "end" }; + }, getSteeringMessages: async () => { steeringPolls++; return []; @@ -1171,12 +1625,6 @@ describe("agentLoop with AgentMessage", () => { followUpPolls++; return [createUserMessage("follow up should stay queued")]; }, - shouldStopAfterTurn: async ({ message, toolResults, context }) => { - expect(message.role).toBe("assistant"); - callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId); - callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role); - return true; - }, }; let llmCalls = 0; diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 94b14de50..bed5dd427 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -5,6 +5,7 @@ import { getCurrentSystemMessage, getModel, toToolDeclaration, + type UserMessage, } from "@earendil-works/pi-ai/compat"; import { Type } from "typebox"; import { describe, expect, it } from "vitest"; @@ -31,6 +32,10 @@ class MockAssistantStream extends EventStream { expect(agent.state.messages[agent.state.messages.length - 1].role).toBe("assistant"); }); - it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => { - let responseCount = 0; + it.each([ + { mode: "one-at-a-time" as const, expectedRequests: 2 }, + { mode: "all" as const, expectedRequests: 1 }, + ])("continue() keeps $mode steering semantics for assistant-tail fallback", async ({ mode, expectedRequests }) => { + const requests: string[][] = []; const agent = new Agent({ - streamFn: () => { + steeringMode: mode, + streamFn: (_model, context) => { + requests.push( + context.messages.flatMap((message) => + message.role === "user" && typeof message.content === "string" ? [message.content] : [], + ), + ); const stream = new MockAssistantStream(); - responseCount++; queueMicrotask(() => { - stream.push({ - type: "done", - reason: "stop", - message: createAssistantMessage(`Processed ${responseCount}`), - }); + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") }); }); return stream; }, }); - - agent.state.messages = [ - { - role: "user", - content: [{ type: "text", text: "Initial" }], - timestamp: Date.now() - 10, - }, - createAssistantMessage("Initial response"), - ]; - - agent.steer({ - role: "user", - content: [{ type: "text", text: "Steering 1" }], - timestamp: Date.now(), - }); - agent.steer({ - role: "user", - content: [{ type: "text", text: "Steering 2" }], - timestamp: Date.now() + 1, - }); + agent.state.messages = [createUserMessage("Initial"), createAssistantMessage("Initial response")]; + agent.steer(createUserMessage("Steering 1")); + agent.steer(createUserMessage("Steering 2")); await expect(agent.continue()).resolves.toBeUndefined(); - const recentMessages = agent.state.messages.slice(-4); - expect(recentMessages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]); - expect(responseCount).toBe(2); + expect(requests).toHaveLength(expectedRequests); + expect(requests[0]).toContain("Steering 1"); + if (mode === "one-at-a-time") { + expect(requests[0]).not.toContain("Steering 2"); + expect(requests[1]).toContain("Steering 2"); + } else { + expect(requests[0]).toContain("Steering 2"); + } }); it("keeps legacy prepareNextTurn signal callback behavior", async () => { @@ -915,7 +912,7 @@ describe("Agent", () => { expect(sawAbortSignal).toBe(true); }); - it("forwards shouldStopAfterTurn through AgentOptions", async () => { + it("forwards finishTurn through AgentOptions with the active abort signal", async () => { const schema = Type.Object({}); const tool: AgentTool = { name: "noop", @@ -929,10 +926,10 @@ describe("Agent", () => { let callbackContextRoles: string[] = []; const agent = new Agent({ initialState: { tools: [tool] }, - shouldStopAfterTurn: (context, signal) => { + finishTurn: (context, signal) => { sawAbortSignal = signal instanceof AbortSignal; callbackContextRoles = context.context.messages.map((message) => message.role); - return true; + return { action: "end" }; }, streamFn: () => { requestCount++; @@ -959,6 +956,216 @@ describe("Agent", () => { expect(callbackContextRoles).toEqual(["system", "user", "assistant", "toolResult"]); }); + it.each([ + { name: "empty", messages: [] }, + { name: "system-only", messages: [{ role: "system" as const, content: "system only", timestamp: 1 }] }, + ])("rejects a queued continuation from $name context without draining queues", async ({ messages }) => { + const agent = new Agent({ initialState: { messages }, streamFn: unusedStreamFunction }); + const steering = createUserMessage("steering"); + const followUp = createUserMessage("follow-up"); + agent.steer(steering); + agent.followUp(followUp); + + await expect(agent.continue()).rejects.toThrow("No messages to continue from"); + expect(agent.peekQueuedMessages()).toEqual([steering]); + agent.clearSteeringQueue(); + expect(agent.peekQueuedMessages()).toEqual([followUp]); + }); + + it.each([ + { + name: "user", + messages: [createUserMessage("existing user")], + }, + { + name: "toolResult", + messages: [ + createUserMessage("existing user"), + createAssistantToolUseMessage([{ type: "toolCall", id: "call-1", name: "noop", arguments: {} }]), + { + role: "toolResult" as const, + toolCallId: "call-1", + toolName: "noop", + content: [{ type: "text" as const, text: "done" }], + isError: false, + timestamp: 1, + }, + ], + }, + ])("defers follow-up input on the first continuation request from a $name tail", async ({ messages }) => { + const requests: string[][] = []; + const agent = new Agent({ + initialState: { messages }, + streamFn: (_model, context) => { + requests.push( + context.messages.flatMap((message) => + message.role === "user" && typeof message.content === "string" ? [message.content] : [], + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), + ); + return stream; + }, + }); + agent.followUp(createUserMessage("follow-up")); + + await agent.continue(); + + expect(requests).toHaveLength(2); + expect(requests[0]).not.toContain("follow-up"); + expect(requests[1]).toContain("follow-up"); + }); + + it.each([ + { mode: "one-at-a-time" as const, expectedRequests: 2 }, + { mode: "all" as const, expectedRequests: 1 }, + ])("polls $mode steering at continuation startup", async ({ mode, expectedRequests }) => { + const requests: string[][] = []; + const agent = new Agent({ + initialState: { messages: [createUserMessage("existing")] }, + steeringMode: mode, + streamFn: (_model, context) => { + requests.push( + context.messages.flatMap((message) => + message.role === "user" && typeof message.content === "string" ? [message.content] : [], + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), + ); + return stream; + }, + }); + agent.steer(createUserMessage("first")); + agent.steer(createUserMessage("second")); + + await agent.continue(); + + expect(requests).toHaveLength(expectedRequests); + expect(requests[0]).toContain("first"); + if (mode === "one-at-a-time") { + expect(requests[0]).not.toContain("second"); + expect(requests[1]).toContain("second"); + } else { + expect(requests[0]).toContain("second"); + } + }); + + it("keeps steering ahead of follow-up from a non-assistant continuation tail", async () => { + const requests: string[][] = []; + const agent = new Agent({ + initialState: { messages: [createUserMessage("existing")] }, + streamFn: (_model, context) => { + requests.push( + context.messages.flatMap((message) => + message.role === "user" && typeof message.content === "string" ? [message.content] : [], + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), + ); + return stream; + }, + }); + agent.steer(createUserMessage("steering")); + agent.followUp(createUserMessage("follow-up")); + + await agent.continue(); + + expect(requests).toHaveLength(2); + expect(requests[0]).toContain("steering"); + expect(requests[0]).not.toContain("follow-up"); + expect(requests[1]).toContain("follow-up"); + }); + + it.each(["error", "aborted"] as const)( + "keeps queues on a %s response even when finishTurn requests continuation", + async (stopReason) => { + const queuedDuringResponse = createUserMessage("steering"); + const followUp = createUserMessage("follow-up"); + const agent = new Agent({ + finishTurn: () => ({ action: "continue" }), + streamFn: () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "error", + reason: stopReason, + error: { + ...createAssistantMessage(stopReason), + stopReason, + errorMessage: stopReason, + }, + }); + }); + return stream; + }, + }); + agent.followUp(followUp); + agent.subscribe((event) => { + if (event.type === "message_end" && event.message.role === "assistant") { + agent.steer(queuedDuringResponse); + } + }); + + await agent.prompt("start"); + + expect(agent.peekQueuedMessages()).toEqual([queuedDuringResponse]); + agent.clearSteeringQueue(); + expect(agent.peekQueuedMessages()).toEqual([followUp]); + }, + ); + + it("keeps queues when finishTurn ends the run", async () => { + const queuedDuringResponse = createUserMessage("steering"); + const followUp = createUserMessage("follow-up"); + const agent = new Agent({ + finishTurn: () => ({ action: "end" }), + streamFn: () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }), + ); + return stream; + }, + }); + agent.followUp(followUp); + agent.subscribe((event) => { + if (event.type === "message_end" && event.message.role === "assistant") { + agent.steer(queuedDuringResponse); + } + }); + + await agent.prompt("start"); + + expect(agent.peekQueuedMessages()).toEqual([queuedDuringResponse]); + agent.clearSteeringQueue(); + expect(agent.peekQueuedMessages()).toEqual([followUp]); + }); + + it("previews the next selected queued messages without consuming them", () => { + const agent = new Agent({ + steeringMode: "one-at-a-time", + followUpMode: "all", + streamFn: () => new MockAssistantStream(), + }); + const first = createUserMessage("first steering"); + const second = createUserMessage("second steering"); + const followUp = createUserMessage("follow-up"); + agent.steer(first); + agent.steer(second); + agent.followUp(followUp); + + expect(agent.peekQueuedMessages()).toEqual([first]); + expect(agent.peekQueuedMessages()).toEqual([first]); + agent.clearSteeringQueue(); + expect(agent.peekQueuedMessages()).toEqual([followUp]); + }); + it("forwards sessionId to streamFunction options", async () => { let receivedSessionId: string | undefined; const agent = new Agent({ diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 57024c8e3..e16245f9e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,12 +6,27 @@ - **Per-model image input limits** — Configure cache-safe image resizing per model for attachments, `read`, and tool-result images. See [Image Input Limits](docs/models.md#image-input-limits). +### Breaking Changes + +- Removed the inherited `shouldStopAfterTurn` agent option. Use `finishTurn` and return `{ action: "end" }` instead. `finishTurn` runs before `turn_end` but applies the decision afterward, and it also receives error and aborted responses; migrate normal-response predicates by returning `undefined` for those hard exits. See the `@earendil-works/pi-agent-core` changelog for a complete before-and-after example. +- Added `ContextEditEntry` to the exported `SessionEntry` union. TypeScript consumers with exhaustive entry switches must handle `context_edit`; use `replacement: null` for omission and a content replacement otherwise. +- Made `SessionManager` canonical for `AgentSession` provider context. Assigning `session.agent.state.messages` no longer replaces future request history; restore with `SessionManager.inMemory(cwd, { id }, entries)`, navigate with `session.navigateTree()`, or append through `session.sessionManager` and call `session.refreshContext()`. +- Expanded `TurnEndEvent` with required boundary fields and added `AgentBeforeSettleEvent` to the exported `ExtensionEvent` union. Consumers constructing events or exhaustively switching on `ExtensionEvent` must handle the new shapes. `ExtensionRunner.emit()` no longer accepts `turn_end`; host integrations dispatch actionable boundaries with `emitBoundary(baseEvent, buildContext)`. +- Deferred runs requested from `agent_settled` handlers until all settled handlers finish. Handlers still observe `ctx.isIdle() === true`, but no longer see a reentrant `agent_start` during the same notification dispatch. + ### Added +- Added append-only model-context edits. For example, `sessionManager.appendContextEdit(entryId, null)` omits one message from future provider context without changing raw history, usage, or UI history. +- Added actionable `turn_end` and `agent_before_settle` extension boundaries. Return `{ entries: [...event.entries, draft], continue: true }` to persist structural entries in order and ensure one next provider request without changing steering or follow-up scheduling. +- Added retain-none compaction input: `sessionManager.appendCompaction(summary, null, tokensBefore)` stores the compaction's own ID as its kept boundary. - Added per-model image resize profiles through `inputLimits.images.resize` in `models.json`, applied to file attachments, image reads, and tool-result images ([#9631](https://github.com/earendil-works/pi/issues/9631)). ### Fixed +- Fixed string context-edit replacements producing invalid assistant and tool-result message content instead of text blocks. +- Fixed context-invisible boundary metadata and replacement edits causing newly appended or replaced input to be summarized before its first provider request. +- Fixed edited-context accounting both discarding valid assistant usage captured after the latest context edit and reusing that usage after a later compaction made it stale. +- Fixed selected error retries and final length/overflow recovery retaining abandoned model attempts in future provider context; post-run recovery omissions are now persisted without hiding raw transcript history or changing queue scheduling. - Fixed `/bug` allowing uploads in offline mode while preserving local zip exports. - Fixed idle prompt-cache warming rebuilding expired caches when its timer or an extension decision is delayed. - Improved crash diagnostics with hints identifying loaded extensions that appear in the stack trace. diff --git a/packages/coding-agent/README.md b/packages/coding-agent/README.md index ae0f18175..6466b91ac 100644 --- a/packages/coding-agent/README.md +++ b/packages/coding-agent/README.md @@ -282,6 +282,15 @@ Long sessions can exhaust context windows. Compaction summarizes older messages Compaction is lossy. The full history remains in the JSONL file; use `/tree` to revisit. Customize compaction behavior via [extensions](#extensions). See [docs/compaction.md](docs/compaction.md) for internals. +Session model context is projected from append-only history. Extensions can append a `context_edit` to omit or replace an earlier message only for future model requests; raw history and usage remain unchanged: + +```typescript +const assistantId = sessionManager.appendMessage(partialAssistant); +sessionManager.appendContextEdit(assistantId, null); // Hidden from model context, retained in JSONL. +``` + +A retain-none compaction uses `appendCompaction(summary, null, tokensBefore)` to make the exact summary the new context root while preserving prior raw entries. + --- ## Settings @@ -383,6 +392,33 @@ export default function (pi: ExtensionAPI) { } ``` +`turn_end` and `agent_before_settle` are actionable persistence boundaries. Handlers can append structural entries and request one continuation. Later handlers see earlier proposals: + +```typescript +let replacedResponse = false; +pi.on("turn_end", (event) => { + if (replacedResponse || event.outcome !== "completed" || event.toolResults.length > 0) return; + replacedResponse = true; + return { + entries: [ + ...event.entries, + { type: "context_edit", targetId: event.messageEntryId, replacement: null }, + { + type: "custom_message", + customType: "replacement-instruction", + content: "Answer again using the persisted user request.", + display: false, + }, + ], + continue: true, + }; +}); +``` + +Continuation is one-shot per boundary result, not per registered handler. It ensures one next provider request: tool-result, steering, or follow-up scheduling can satisfy that request without adding another one; otherwise Pi makes one context-only request. Guard handlers like the example above because an unconditional `continue: true` is evaluated again after the next response and can loop indefinitely. Error and aborted responses remain hard exits. + +Use `agent_before_settle` for final actions after retries, compaction, and queued input are exhausted. See [docs/extensions.md](docs/extensions.md#agent_start--agent_end--agent_before_settle--agent_settled). + The default export can also be `async`. pi waits for async extension factories before startup continues, which is useful for one-time initialization such as fetching remote model lists before calling `pi.registerProvider()`. **What's possible:** diff --git a/packages/coding-agent/docs/compaction.md b/packages/coding-agent/docs/compaction.md index cdc94931e..861dcf0a6 100644 --- a/packages/coding-agent/docs/compaction.md +++ b/packages/coding-agent/docs/compaction.md @@ -34,14 +34,16 @@ contextTokens > contextWindow - reserveTokens By default, `reserveTokens` is 16384 tokens (configurable in `~/.pi/agent/settings.json` or `/.pi/settings.json`). This leaves room for the LLM's response. -During a multi-turn agent run, Pi checks this threshold after tools finish and their results are appended, before starting the next assistant response. If the threshold is crossed, Pi compacts inside the same agent run and resumes with the summary and retained messages. It skips this between-turn check when the completed tool batch terminates the run and no queued message requires another response. Pi also checks the threshold before a new user prompt and after a low-level agent run ends. +During a multi-turn agent run, Pi checks the canonical projected context after tools finish and their results are appended, before starting the next assistant response. If the threshold is crossed, Pi compacts during `prepareNextTurn`, then performs the existing catch-up steering poll before `turn_start`. It skips this between-turn check when the completed tool batch terminates the run and no queued message requires another response. Pi also checks before a new user prompt and performs final-attempt overflow recovery after the low-level run ends. + +A provider context-overflow error or an early final `stopReason: "length"` can select one compact-and-retry recovery attempt. Length responses with tool calls retain their synthetic failed tool results and follow the ordinary tool/queue scheduler rather than forcing the run to end. You can also trigger manually with `/compact [instructions]`, where optional instructions focus the summary. ### How It Works -1. **Find cut point**: Walk backwards from newest message, accumulating token estimates until `keepRecentTokens` (default 20k, configurable in `~/.pi/agent/settings.json` or `/.pi/settings.json`) is reached -2. **Extract messages**: Collect messages from the previous kept boundary (or session start) up to the cut point +1. **Find cut point**: Walk backwards through the finalized session projection, accumulating token estimates until `keepRecentTokens` (default 20k, configurable in `~/.pi/agent/settings.json` or `/.pi/settings.json`) is reached +2. **Extract messages**: Collect projected messages from the previous kept boundary (or session start) up to the cut point 3. **Generate summary**: Call LLM to summarize with structured format, passing the previous summary as iterative context when present 4. **Append entry**: Save `CompactionEntry` with summary and `firstKeptEntryId` 5. **Rebuilds context**: Session rebuilds the context for the next request, using summary + messages from `firstKeptEntryId` onwards @@ -78,7 +80,22 @@ What the LLM sees: prompt from cmp messages from firstKeptEntryId ``` -On repeated compactions, the summarized span starts at the previous compaction's kept boundary (`firstKeptEntryId`), not at the compaction entry itself, falling back to the entry after the previous compaction if that kept entry cannot be found in the path. This preserves messages that survived the earlier compaction by including them in the next summarization pass as well. Pi also recalculates `tokensBefore` from the rebuilt session context before writing the new `CompactionEntry`, so the token count reflects the actual pre-compaction context being replaced. +On repeated compactions, the summarized span starts at the previous compaction's kept boundary (`firstKeptEntryId`), not at the compaction entry itself, falling back to the entry after the previous compaction if that kept entry cannot be found in the path. A retain-none compaction records its own ID as `firstKeptEntryId`; repeated compaction starts after that entry. This preserves messages that survived the earlier compaction by including them in the next summarization pass as well. Pi also recalculates `tokensBefore` from the rebuilt, context-edited session projection before writing the new `CompactionEntry`, so the token count reflects the actual pre-compaction context being replaced. Omitted raw entries remain stored but do not affect cut selection, summaries, checkpoints, or token estimates. + +### Overflow and Length Recovery Ordering + +Recovery preserves the existing lifecycle and queue order. The completed attempt remains visible to `turn_end` and `agent_end`; post-run recovery then repairs persisted model context before a fresh retry: + +```text +persist final assistant response +→ extension/public turn_end +→ extension/public agent_end +→ append context_edit omissions for the selected attempt +→ for overflow/length: run session_before_compact and append compaction on success +→ start the retry as a fresh run +``` + +If recovery compaction fails or is cancelled, Pi keeps the omission edits, appends no compaction, and schedules no internal retry. Existing queued work remains governed by ordinary steering and follow-up rules. `agent_before_settle` sees the repaired projection after recovery processing. Raw transcript history, exports, billing totals, and history-search extensions can still inspect the omitted attempt. ### Split Turns @@ -118,6 +135,8 @@ Valid cut points are: Never cut at tool results (they must stay with their tool call). +Preparation advances the kept boundary into a context-invisible suffix only when that suffix contains an omitted assistant attempt and no unomitted context-producing entries. Recovery `context_edit` omissions satisfy this rule; intrinsically context-invisible metadata may coexist with them. Metadata alone and newly appended custom messages do not move the cut. A replacement edit affecting the candidate input or summarized prefix also blocks advancement because the omitted assistant answered the pre-edit input; replacements of suffix entries that are ultimately omitted remain safe. This allows an over-budget recovered input to be summarized while retaining the edits that keep the abandoned attempt omitted, without making bookkeeping change whether new model input is preserved verbatim. + ### CompactionEntry Structure Defined in [`session-manager.ts`](https://github.com/earendil-works/pi/blob/main/packages/coding-agent/src/core/session-manager.ts): diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index b45163108..d58304c3b 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -307,9 +307,13 @@ user sends prompt ──────────────────── │ │ └─► tool_execution_end │ │ │ │ │ │ │ └─► turn_end │ │ + │ └─► threshold compaction before a naturally required next turn │ │ ├─► agent_end │ - └─► agent_settled (no retry/compaction/follow-up left) │ + ├─► retry backoff or final-attempt recovery (when selected) + │ └─► fresh agent_start on successful recovery │ + ├─► agent_before_settle (can append entries and continue)│ + └─► agent_settled (final, notification only) │ │ user sends another prompt ◄────────────────────────────────┘ @@ -567,9 +571,9 @@ The `systemPromptOptions` field gives extensions access to the same structured d Inside `before_agent_start`, `event.systemPrompt` and `ctx.getSystemPrompt()` both reflect the chained system prompt as of the current handler. Later `before_agent_start` handlers can still modify it again. -#### agent_start / agent_end / agent_settled +#### agent_start / agent_end / agent_before_settle / agent_settled -`agent_start` fires when a low-level agent run begins. `agent_end` fires when that run ends, but Pi may still auto-retry, auto-compact and retry, or continue with queued follow-up messages. Use `agent_settled` for status integrations that need to know Pi will not continue running automatically. +`agent_start` fires when a low-level agent run begins. `agent_end` fires when that run ends, but Pi may still auto-retry, auto-compact and retry, or continue with queued follow-up messages. `agent_before_settle` is the final actionable boundary: it can append session entries and request one continuation. `agent_settled` is final and notification-only; use it for status integrations that need to know Pi will not continue running automatically. ```typescript pi.on("agent_start", async (_event, ctx) => {}); @@ -578,11 +582,28 @@ pi.on("agent_end", async (event, ctx) => { // event.messages - messages from this low-level run }); +let addedReviewReminder = false; +pi.on("agent_before_settle", async (event, ctx) => { + if (addedReviewReminder) return; + addedReviewReminder = true; + return { + entries: [...event.entries, { + type: "custom_message", + customType: "review-reminder", + content: "Review the final diff before replying.", + display: false, + }], + continue: true, + }; +}); + pi.on("agent_settled", async (_event, ctx) => { - // ctx.isIdle() is true here unless another extension started a new run. + // ctx.isIdle() is true; runs requested here start after all settled handlers finish. }); ``` +If the run is aborted while `agent_before_settle` handlers are running, valid returned entries are still committed, but requested continuation is suppressed. Work requested from `agent_settled` is deferred until every settled handler completes, so notification dispatch is non-reentrant. + #### ui_prompt_start / ui_prompt_end Notification-only lifecycle events for blocking user-facing extension UI prompts. They fire around `ctx.ui.select()`, `ctx.ui.confirm()`, `ctx.ui.input()`, `ctx.ui.editor()`, and `ctx.ui.custom()` so host/status integrations can report "waiting for user" instead of just "running". @@ -610,11 +631,34 @@ pi.on("turn_start", async (event, ctx) => { // event.turnIndex, event.timestamp }); +let replacedResponse = false; pi.on("turn_end", async (event, ctx) => { // event.turnIndex, event.message, event.toolResults + // event.entries contains the structural entries proposed so far. + if (replacedResponse || event.outcome !== "completed" || event.toolResults.length > 0) return; + replacedResponse = true; + return { + entries: [ + ...event.entries, + { type: "context_edit", targetId: event.messageEntryId, replacement: null }, + { + type: "custom_message", + customType: "replacement-instruction", + content: "Answer again using the persisted user request.", + display: false, + }, + ], + continue: true, + }; }); ``` +`turn_end` runs after the assistant and tool-result messages have been persisted and before the low-level `turn_end` event. Retry backoff and final-attempt recovery still happen after `agent_end`, preserving their existing lifecycle and queue ordering; `agent_before_settle` sees the repaired projection after that work completes. Boundary handlers run in extension load and registration order. Each handler sees prior proposals in `event.entries` and sees `event.context` rebuilt from them. Returning `entries` or `continue` replaces only that field; omitted fields preserve the current proposal. Allowed draft entry types are `custom`, `custom_message`, `context_edit`, and `compaction`. The complete proposal is validated before it is appended in list order after all handlers finish; a handler error is reported and later handlers still run. Validation prevents partially applied semantic errors, but persistence is not transactional. + +`continue: true` ensures one next provider request for that boundary invocation. If tool results, steering, or a follow-up already cause that request, they satisfy the decision and no additional request is made; otherwise Pi makes one context-only request. Error and aborted responses remain hard exits. `continue: false` never suppresses natural work. Guard continuation conditions: an unconditional `continue: true` is evaluated again after the next response and can create an endless loop. A `custom_message` draft contributes a user-role model message but is extension-authored: it does not run human input hooks, slash commands, skills, or prompt templates. + +Host integrations that construct `TurnEndEvent` values must now provide `messageEntryId`, `toolResultEntryIds`, `outcome`, `entries`, `continue`, and `context`. `ExtensionEvent` exhaustive switches must also handle `agent_before_settle`. `ExtensionRunner.emit()` excludes actionable turn boundaries; dispatch `turn_end` and `agent_before_settle` through `emitBoundary(baseEvent, buildContext)` so handlers receive chained previews. Other dedicated runner methods still return results for events such as `session_before_*`. + #### message_start / message_update / message_end Fired for message lifecycle updates. diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index aaa415317..a830fe682 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -91,6 +91,8 @@ interface AgentSession { // State access agent: Agent; + sessionManager: SessionManager; + refreshContext(): void; model: Model | undefined; thinkingLevel: ThinkingLevel; messages: AgentMessage[]; @@ -251,8 +253,8 @@ const state = session.agent.state; // state.streamingMessage?: AgentMessage - current partial assistant message // state.errorMessage?: string - latest assistant error -// Replace messages (useful for branching or restoration) -session.agent.state.messages = messages; // copies the top-level array +// Model-visible messages are projected from session.sessionManager. +// agent.state.messages is a refreshed inspection cache; do not assign it for restoration. // Replace tools session.agent.state.tools = tools; // copies the top-level array @@ -261,6 +263,15 @@ session.agent.state.tools = tools; // copies the top-level array await session.agent.waitForIdle(); ``` +Provider requests use `session.sessionManager` as the canonical finalized context. Assigning `session.agent.state.messages` does not replace persisted context and may be overwritten at the next request boundary. Restore externally stored history when constructing the session instead: + +```typescript +const restoredManager = SessionManager.inMemory(process.cwd(), { id: sessionId }, entries); +const { session } = await createAgentSession({ sessionManager: restoredManager }); +``` + +For an existing session, use `session.navigateTree(entryId)` to move its active branch. Use `session.sessionManager.appendMessage(...)` plus `session.refreshContext()` only when intentionally appending externally managed entries. + ### Events Subscribe to events to receive streaming output and lifecycle notifications. diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index c3811c636..522cf60d7 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -274,7 +274,7 @@ Created when context is compacted. Stores a summary of earlier messages and a co {"type":"compaction","id":"f6g7h8i9","parentId":"e5f6g7h8","timestamp":"2024-12-03T14:10:00.000Z","summary":"User discussed X, Y, Z...","firstKeptEntryId":"c3d4e5f6","tokensBefore":50000,"systemMessage":{"role":"system","content":"You are a coding assistant.","toolsAdded":[],"timestamp":1733235000000}} ``` -`firstKeptEntryId` is required. It identifies the first entry retained from before the compaction entry. When rebuilding context, Pi replaces older summarized entries with the compaction summary and keeps the range beginning at this entry. +`firstKeptEntryId` is required. It identifies the first entry retained from before the compaction entry. When rebuilding context, Pi replaces older summarized entries with the compaction summary and keeps the range beginning at this entry. A retain-none compaction stores its own ID in this field, so no preceding entries are retained. Optional fields: - `systemMessage`: The replayed prompt sections and tool declarations at the compaction boundary; it becomes the leading system message of the compacted context, and system messages among the kept entries are dropped in its favor. It is absent on older session entries. @@ -282,6 +282,16 @@ Optional fields: - `details`: Implementation-specific data (e.g., `{ readFiles: string[], modifiedFiles: string[] }` for default, or custom data for extensions) - `fromHook`: `true` if generated by an extension, `false`/`undefined` if pi-generated (legacy field name) +### ContextEditEntry + +Append-only edit of one earlier context-producing entry. It changes only future model context; the target entry and its metadata remain unchanged in raw history, UI, exports, and session accounting. + +```json +{"type":"context_edit","id":"g6h7i8j9","parentId":"f6g7h8i9","timestamp":"2024-12-03T14:11:00.000Z","targetId":"c3d4e5f6","replacement":null} +``` + +Targets may be user, assistant, tool-result, or custom-message entries. `replacement: null` omits the target from model context. A non-null `replacement` replaces only the target message content. String replacements for assistant and tool-result entries are normalized to one text block because those roles require content arrays. If several edits target the same entry, the latest edit on the active branch wins. Edits are branch-relative: navigating to a point before the edit reveals the target's original contribution again. + ### BranchSummaryEntry Created when switching branches via `/tree` with an LLM generated summary of the left branch up to the common ancestor. Captures context from the abandoned path. @@ -366,7 +376,9 @@ Entries normally form one tree, but navigation APIs can create multiple roots: - Includes entries after the compaction entry 3. Preserves non-message entries in the selected range so interactive mode can render them -`buildSessionContext()` builds on that entry list to produce the message list for the LLM: +`buildSessionProjection()` then applies the latest `context_edit` for each selected target. It returns the model-visible messages together with their source entries. Omitted targets produce no message; replacements retain the source entry's role and metadata while changing only content. The raw selected entries are not modified. + +`buildSessionContext()` builds on that projection to produce the message list for the LLM: 1. Extracts current model and thinking level settings from the full path 2. Converts selected entries to messages: @@ -374,6 +386,7 @@ Entries normally form one tree, but navigation APIs can create multiple roots: - `compaction` -> complete system checkpoint followed by `compactionSummary` - `branch_summary` -> `branchSummary` - `custom_message` -> `CustomMessage` + - `context_edit` -> no context message of its own - `usage` and `custom` -> no context message The compaction summary replaces entries before `firstKeptEntryId`. Pre-compaction system messages are folded into the complete checkpoint rather than replayed from the retained range. Retained non-system entries and all entries after the compaction remain available to the LLM. diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 861c00f75..851258b45 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -253,7 +253,7 @@ export class AgentSessionRuntime { ); if (options?.setup) { await options.setup(this.session.sessionManager); - this.session.agent.state.messages = this.session.sessionManager.buildSessionContext().messages; + this.session.refreshContext(); } await this.finishSessionReplacement(options?.withSession); return { cancelled: false }; diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 9845d9b0e..338d0e0d0 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -34,6 +34,7 @@ import type { ProviderHeaders, SystemMessage, TextContent, + ToolResultMessage, Usage, } from "@earendil-works/pi-ai/compat"; import { @@ -64,6 +65,7 @@ import { collectEntriesForBranchSummary, compact, estimateContextTokens, + estimateProjectedContextTokens, estimateTokens, generateBranchSummary, prepareCompaction, @@ -73,6 +75,8 @@ import { DEFAULT_THINKING_LEVEL, THINKING_LEVEL_OPTIONS } from "./defaults.ts"; import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.ts"; import { createToolHtmlRenderer } from "./export-html/tool-renderer.ts"; import { + type AgentActivityOutcome, + type BoundaryContextPreview, type ContextUsage, type ExtensionCommandContextActions, type ExtensionErrorListener, @@ -86,6 +90,7 @@ import { type ReplacedSessionContext, type SessionBeforeCompactResult, type SessionBeforeTreeResult, + type SessionBoundaryDraft, type SessionCompactFailedEvent, type SessionStartEvent, type ShutdownHandler, @@ -95,19 +100,24 @@ import { type ToolExecutionUpdateEvent, type ToolInfo, type TreePreparation, - type TurnEndEvent, type TurnStartEvent, wrapRegisteredTools, } from "./extensions/index.ts"; import { emitSessionShutdownEvent } from "./extensions/runner.ts"; -import type { BashExecutionMessage, CustomMessage } from "./messages.ts"; +import { type BashExecutionMessage, type CustomMessage, convertToLlm } from "./messages.ts"; import { ModelRegistry } from "./model-registry.ts"; import type { ModelRuntime } from "./model-runtime.ts"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts"; import { exportSessionToJsonl } from "./session-export.ts"; -import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts"; -import { getLatestCompactionEntry } from "./session-manager.ts"; +import { + type BranchSummaryEntry, + type CompactionEntry, + type ContextEditEntry, + getLatestCompactionEntry, + type SessionEntry, + SessionManager, +} from "./session-manager.ts"; import type { CacheWarmingMode, SettingsManager } from "./settings-manager.ts"; import type { SlashCommandInfo } from "./slash-commands.ts"; import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts"; @@ -358,6 +368,15 @@ export class AgentSession { // Extension system private _extensionRunner!: ExtensionRunner; private _turnIndex = 0; + private readonly _entryIdsByMessage = new WeakMap(); + private readonly _boundaryDispatchedMessages = new WeakSet(); + private _lastAssistantMessage: AssistantMessage | undefined; + private _lastAssistantToolResults: AgentMessage[] = []; + private _lastActivityOutcome: AgentActivityOutcome = "completed"; + private _isBeforeSettle = false; + private _abortDuringBeforeSettle = false; + private _isEmittingAgentSettled = false; + private readonly _deferredSettledActions: Array<() => Promise> = []; private _resourceLoader: ResourceLoader; private _customTools: ToolDefinition[]; @@ -415,6 +434,8 @@ export class AgentSession { this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent); this._installAgentToolHooks(); this._installAgentNextTurnRefresh(); + this._installAgentRequestProjection(); + this._installAgentBoundaryHooks(); this._installAgentForcedPromptProjection(); this._buildRuntime({ @@ -566,19 +587,100 @@ export class AgentSession { private async _compactBeforeNextAssistantResponse(context: AgentContext): Promise { const model = this.model; const settings = this.settingsManager.getCompactionSettings(model); + const projection = this.sessionManager.buildSessionProjection(); if ( !model || model.contextWindow <= 0 || - !shouldCompact(estimateContextTokens(context.messages).tokens, model.contextWindow, settings) + !shouldCompact( + estimateProjectedContextTokens(projection, this.sessionManager.getBranch()).tokens, + model.contextWindow, + settings, + ) ) { - return context; + return { ...context, messages: projection.messages }; } await this._runAutoCompaction("threshold", false); - return { - ...context, - messages: this.agent.state.messages.slice(), + return { ...context, messages: this.sessionManager.buildSessionProjection().messages }; + } + + private _installAgentRequestProjection(): void { + const previousPrepareRequest = this.agent.prepareRequest; + this.agent.prepareRequest = async (request, signal) => { + const canonicalContext = { + ...request.context, + messages: this.sessionManager.buildSessionProjection().messages, + // Messages declare the provider-visible loadout; context.tools keeps executable implementations. + tools: this.agent.state.tools.slice(), + }; + const previous = await previousPrepareRequest?.( + { + ...request, + context: canonicalContext, + model: this.agent.state.model, + thinkingLevel: this.agent.state.thinkingLevel, + }, + signal, + ); + return { + ...previous, + context: previous?.context ?? canonicalContext, + model: previous?.model ?? this.agent.state.model, + thinkingLevel: previous?.thinkingLevel ?? this.agent.state.thinkingLevel, + }; + }; + } + + private async _dispatchTurnEndBoundary( + message: AssistantMessage, + toolResults: ToolResultMessage[], + ): Promise { + this._lastActivityOutcome = + message.stopReason === "aborted" ? "aborted" : message.stopReason === "error" ? "error" : "completed"; + const messageEntryId = this._findPersistedMessageEntryId(message); + if (!this._extensionRunner.hasHandlers("turn_end")) return false; + if (!messageEntryId) { + this._extensionRunner.emitError({ + extensionPath: "", + event: "turn_end", + error: "turn_end could not resolve the persisted assistant entry ID", + }); + return false; + } + const toolResultEntryIds = toolResults.flatMap((result) => { + const entryId = this._findPersistedMessageEntryId(result); + return entryId ? [entryId] : []; + }); + const boundary = await this._extensionRunner.emitBoundary( + { + type: "turn_end", + turnIndex: this._turnIndex, + message, + toolResults, + messageEntryId, + toolResultEntryIds, + outcome: this._lastActivityOutcome, + }, + (entries) => this._buildBoundaryContext(entries, "turn_end"), + ); + this._commitBoundaryDrafts(boundary.entries); + if (boundary.continue && !this._buildBoundaryContext([], "turn_end").canContinue) { + this._reportInvalidBoundaryContinuation("turn_end"); + return false; + } + return boundary.continue; + } + + private _installAgentBoundaryHooks(): void { + const previousFinishTurn = this.agent.finishTurn; + this.agent.finishTurn = async (turn, signal) => { + this._boundaryDispatchedMessages.add(turn.message); + const extensionContinue = await this._dispatchTurnEndBoundary(turn.message, turn.toolResults); + const previousDecision = await previousFinishTurn?.(turn, signal); + if (previousDecision?.action === "end") return previousDecision; + if (extensionContinue || previousDecision?.action === "continue") return { action: "continue" }; + return undefined; }; } @@ -589,7 +691,10 @@ export class AgentSession { ? async (_turn: PrepareNextTurnContext, signal?: AbortSignal) => await this.agent.prepareNextTurn?.(signal) : undefined); this.agent.prepareNextTurnWithContext = async (turn, signal) => { - const context = await this._compactBeforeNextAssistantResponse(turn.context); + const context = await this._compactBeforeNextAssistantResponse({ + ...turn.context, + messages: this.sessionManager.buildSessionProjection().messages, + }); const previousSnapshot = await previousPrepareNextTurnWithContext?.({ ...turn, context }, signal); const nextContext = previousSnapshot?.context ?? context; const runOptions = this._runSystemPromptOptions ?? this._baseSystemPromptOptions; @@ -622,6 +727,106 @@ export class AgentSession { // Event Subscription // ========================================================================= + private _refreshFinalizedContext(): void { + const projection = this.sessionManager.buildSessionProjection(); + for (const entry of projection.entries) { + for (const message of entry.messages) this._entryIdsByMessage.set(message, entry.sourceEntry.id); + } + this.agent.state.messages = projection.messages; + } + + private _applyBoundaryDrafts(manager: SessionManager, drafts: SessionBoundaryDraft[]): SessionEntry[] { + const appended: SessionEntry[] = []; + for (const draft of drafts) { + let entryId: string; + switch (draft.type) { + case "custom": + entryId = manager.appendCustomEntry(draft.customType, draft.data); + break; + case "custom_message": + entryId = manager.appendCustomMessageEntry( + draft.customType, + draft.content, + draft.display, + draft.details, + ); + break; + case "context_edit": + entryId = manager.appendContextEdit(draft.targetId, draft.replacement); + break; + case "compaction": { + const tokensBefore = estimateProjectedContextTokens( + manager.buildSessionProjection(), + manager.getBranch(), + ).tokens; + entryId = manager.appendCompaction( + draft.summary, + draft.firstKeptEntryId, + tokensBefore, + draft.details, + true, + draft.usage, + ); + break; + } + } + const entry = manager.getEntry(entryId); + if (entry) appended.push(entry); + } + return appended; + } + + private _createBoundaryPreviewManager(drafts: SessionBoundaryDraft[]): SessionManager { + const header = this.sessionManager.getHeader(); + if (!header) throw new Error("Session header is missing"); + const manager = SessionManager.inMemory(this._cwd, undefined, [header, ...this.sessionManager.getBranch()]); + this._applyBoundaryDrafts(manager, drafts); + return manager; + } + + private _getPendingBoundaryMessages(): AgentMessage[] { + return [...this.agent.peekQueuedMessages(), ...this._pendingCustomMessages]; + } + + private _buildBoundaryContext( + drafts: SessionBoundaryDraft[], + boundary: "turn_end" | "agent_before_settle", + ): BoundaryContextPreview { + const projection = this._createBoundaryPreviewManager(drafts).buildSessionProjection(); + const pendingMessages = this._getPendingBoundaryMessages(); + const llmMessages = convertToLlm(projection.messages); + const finalRole = llmMessages[llmMessages.length - 1]?.role; + const hasNonSystemContext = llmMessages.some((message) => message.role !== "system"); + const contextCanContinue = hasNonSystemContext && finalRole !== "assistant"; + const pendingCustomContext = this._pendingCustomMessages.length > 0; + return { + contextEntries: projection.entries, + contextMessages: projection.messages, + llmMessages, + pendingMessages, + canContinue: + contextCanContinue || + pendingCustomContext || + (boundary === "turn_end" + ? this.agent.hasQueuedMessages() + : finalRole === "assistant" && this.agent.hasQueuedMessages()), + }; + } + + private _commitBoundaryDrafts(drafts: SessionBoundaryDraft[]): void { + const appended = this._applyBoundaryDrafts(this.sessionManager, drafts); + this._refreshFinalizedContext(); + for (const entry of appended) this._emit({ type: "entry_appended", entry }); + } + + private _reportInvalidBoundaryContinuation(event: "turn_end" | "agent_before_settle"): void { + this._extensionRunner.emitError({ + extensionPath: "", + event, + error: `${event} requested continuation without runnable model context`, + }); + } + /** Emit an event to all listeners */ private _emit(event: AgentSessionEvent): void { for (const l of this._eventListeners) { @@ -665,16 +870,25 @@ export class AgentSession { private async _emitAgentSettled(): Promise { this._cacheWarmer?.onAgentSettled(); this._isAgentRunActive = false; + this._isEmittingAgentSettled = true; try { await this._extensionRunner.emit({ type: "agent_settled" }); this._emit({ type: "agent_settled" }); } finally { - this._resolveIdleWaitIfIdle(); + this._isEmittingAgentSettled = false; } - } - // Track last assistant message for auto-compaction check - private _lastAssistantMessage: AssistantMessage | undefined = undefined; + const deferred = this._deferredSettledActions.splice(0); + if (deferred.length > 0) { + try { + for (const action of deferred) await action(); + } finally { + this._resolveIdleWaitIfIdle(); + } + return; + } + this._resolveIdleWaitIfIdle(); + } /** Internal handler for agent events - shared by subscribe and reconnect */ private _handleAgentEvent = async (event: AgentEvent): Promise => { @@ -700,18 +914,17 @@ export class AgentSession { } } - // Emit to extensions first + // Emit to extensions first, then notify public listeners. await this._emitExtensionEvent(event); - - // Notify all listeners this._emit(event.type === "agent_end" ? { ...event, willRetry: this._willRetryAfterAgentEnd(event) } : event); // Handle session persistence if (event.type === "message_end") { + let entryId: string | undefined; // Check if this is a custom message from extensions if (event.message.role === "custom") { // Persist as CustomMessageEntry - this.sessionManager.appendCustomMessageEntry( + entryId = this.sessionManager.appendCustomMessageEntry( event.message.customType, event.message.content, event.message.display, @@ -724,15 +937,14 @@ export class AgentSession { event.message.role === "toolResult" ) { // Regular LLM message - persist as SessionMessageEntry - this.sessionManager.appendMessage(event.message); + entryId = this.sessionManager.appendMessage(event.message); } + if (entryId) this._entryIdsByMessage.set(event.message, entryId); // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere - // Track assistant message for auto-compaction (checked on agent_end) if (event.message.role === "assistant") { - this._lastAssistantMessage = event.message; - const assistantMsg = event.message as AssistantMessage; + this._lastAssistantMessage = assistantMsg; if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "length") { this._overflowRecoveryAttempted = false; } @@ -756,6 +968,7 @@ export class AgentSession { // extension and listener dispatch above also picks up messages that turn_end // handlers queued. if (event.type === "turn_end") { + this._lastAssistantToolResults = event.toolResults; this._flushPendingCustomMessages(); } }; @@ -776,6 +989,47 @@ export class AgentSession { return false; } + private _findPersistedMessageEntryId(message: AgentMessage): string | undefined { + const mapped = this._entryIdsByMessage.get(message); + if (mapped) return mapped; + for (const entry of [...this.sessionManager.getBranch()].reverse()) { + if (entry.type === "message" && entry.message === message) return entry.id; + } + + const messageIndex = this.agent.state.messages.indexOf(message); + if (messageIndex < 0) return undefined; + const projection = this.sessionManager.buildSessionProjection(); + let projectedIndex = 0; + for (const entry of projection.entries) { + for (let i = 0; i < entry.messages.length; i++) { + if (projectedIndex === messageIndex) { + this._entryIdsByMessage.set(message, entry.sourceEntry.id); + return entry.sourceEntry.id; + } + projectedIndex++; + } + } + return undefined; + } + + private _omitRecoveryAttempt(message: AssistantMessage, toolResults: AgentMessage[] = []): void { + const targets = [message, ...toolResults]; + const targetIds = targets.map((target) => this._findPersistedMessageEntryId(target)); + const unresolvedProjectedTarget = targets.some( + (target, index) => targetIds[index] === undefined && this.agent.state.messages.includes(target), + ); + if (unresolvedProjectedTarget) { + throw new Error("Cannot persist recovery omission because a projected message has no source entry"); + } + for (const targetId of targetIds) { + if (!targetId) continue; + const editId = this.sessionManager.appendContextEdit(targetId, null); + const entry = this.sessionManager.getEntry(editId); + if (entry) this._emit({ type: "entry_appended", entry }); + } + this._refreshFinalizedContext(); + } + /** Find the last assistant message in agent state (including aborted ones) */ private _findLastAssistantMessage(): AssistantMessage | undefined { const messages = this.agent.state.messages; @@ -819,13 +1073,9 @@ export class AgentSession { }; await this._extensionRunner.emit(extensionEvent); } else if (event.type === "turn_end") { - const extensionEvent: TurnEndEvent = { - type: "turn_end", - turnIndex: this._turnIndex, - message: event.message, - toolResults: event.toolResults, - }; - await this._extensionRunner.emit(extensionEvent); + if (event.message.role === "assistant" && !this._boundaryDispatchedMessages.delete(event.message)) { + await this._dispatchTurnEndBoundary(event.message, event.toolResults); + } this._turnIndex++; } else if (event.type === "message_start") { const extensionEvent: MessageStartEvent = { @@ -944,6 +1194,11 @@ export class AgentSession { // Read-only State Access // ========================================================================= + /** Refresh the public finalized transcript from the canonical session projection. */ + refreshContext(): void { + this._refreshFinalizedContext(); + } + /** Full agent state */ get state(): AgentState { return this.agent.state; @@ -1215,7 +1470,13 @@ export class AgentSession { this._isAgentRunActive = true; try { await this.agent.prompt(messages); - while (await this._handlePostAgentRun()) { + while (!this._agentRunAbortRequested) { + if (await this._handlePostAgentRun()) { + if (this._agentRunAbortRequested) break; + await this.agent.continue(); + continue; + } + if (this._agentRunAbortRequested || !(await this._runBeforeSettleBoundary())) break; if (this._agentRunAbortRequested) break; await this.agent.continue(); } @@ -1229,17 +1490,17 @@ export class AgentSession { } private async _handlePostAgentRun(): Promise { - const msg = this._lastAssistantMessage; + const message = this._lastAssistantMessage; + const toolResults = this._lastAssistantToolResults; this._lastAssistantMessage = undefined; + this._lastAssistantToolResults = []; if (this._agentRunAbortRequested) { this._finishCancelledRetry(); return false; } - if (!msg) { - return false; - } + if (!message) return this.agent.hasQueuedMessages(); - if (this._isRetryableError(msg) && (await this._prepareRetry(msg))) { + if (this._isRetryableError(message) && (await this._prepareRetry(message))) { if (this._agentRunAbortRequested) this._finishCancelledRetry(); return !this._agentRunAbortRequested; } @@ -1248,25 +1509,49 @@ export class AgentSession { return false; } - if (msg.stopReason === "error" && this._retryAttempt > 0) { + if (message.stopReason === "error" && this._retryAttempt > 0) { this._emit({ type: "auto_retry_end", success: false, attempt: this._retryAttempt, - finalError: msg.errorMessage, + finalError: message.errorMessage, }); this._retryAttempt = 0; } - if (await this._checkCompaction(msg)) { + if (await this._checkCompaction(message, true, toolResults)) { return !this._agentRunAbortRequested; } - // The agent loop drains both queues before emitting agent_end. Any messages - // here were queued by agent_end extension handlers and need a continuation. + // The low-level loop drains both queues before agent_end. Messages queued by + // agent_end handlers require a fresh run before pre-settlement handlers fire. return !this._agentRunAbortRequested && this.agent.hasQueuedMessages(); } + private async _runBeforeSettleBoundary(): Promise { + if (!this._extensionRunner.hasHandlers("agent_before_settle")) return this.agent.hasQueuedMessages(); + this._isBeforeSettle = true; + this._abortDuringBeforeSettle = false; + try { + const result = await this._extensionRunner.emitBoundary( + { type: "agent_before_settle", outcome: this._lastActivityOutcome }, + (entries) => this._buildBoundaryContext(entries, "agent_before_settle"), + ); + this._commitBoundaryDrafts(result.entries); + this._flushPendingCustomMessages(); + const finalContext = this._buildBoundaryContext([], "agent_before_settle"); + if (this._abortDuringBeforeSettle) return false; + const shouldContinue = result.continue || this.agent.hasQueuedMessages(); + if (shouldContinue && !finalContext.canContinue) { + if (result.continue) this._reportInvalidBoundaryContinuation("agent_before_settle"); + return false; + } + return shouldContinue; + } finally { + this._isBeforeSettle = false; + } + } + private async _runInputHandlers( text: string, images: ImageContent[] | undefined, @@ -1319,6 +1604,10 @@ export class AgentSession { * @throws Error if no model selected or no API key available (when not streaming) */ async prompt(text: string, options?: PromptOptions): Promise { + if (this._isEmittingAgentSettled) { + this._deferredSettledActions.push(async () => await this.prompt(text, options)); + return; + } const expandPromptTemplates = options?.expandPromptTemplates ?? true; const preflightResult = options?.preflightResult; let messages: AgentMessage[] | undefined; @@ -1611,11 +1900,7 @@ export class AgentSession { if (images) { content.push(...images); } - this.agent.followUp({ - role: "user", - content, - timestamp: Date.now(), - }); + this.agent.followUp({ role: "user", content, timestamp: Date.now() }); } /** @@ -1668,6 +1953,10 @@ export class AgentSession { this.agent.steer(appMessage); } } else if (options?.triggerTurn) { + if (this._isEmittingAgentSettled) { + this._deferredSettledActions.push(async () => await this._runAgentPrompt(appMessage)); + return; + } await this._runAgentPrompt(appMessage); } else if (this.isStreaming) { // Appending now would put the message between an assistant tool call and its @@ -1681,13 +1970,13 @@ export class AgentSession { } private _appendCustomMessage(appMessage: CustomMessage): void { - this.agent.state.messages.push(appMessage); this.sessionManager.appendCustomMessageEntry( appMessage.customType, appMessage.content, appMessage.display, appMessage.details, ); + this._refreshFinalizedContext(); this._emit({ type: "message_start", message: appMessage }); this._emit({ type: "message_end", message: appMessage }); } @@ -1790,6 +2079,7 @@ export class AgentSession { this.abortRetry(); this.abortCompaction(); this.abortBranchSummary(); + if (this._isBeforeSettle) this._abortDuringBeforeSettle = true; this.agent.abort(); await this.waitForIdle(); } @@ -2208,9 +2498,8 @@ export class AgentSession { this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); - const sessionContext = this.sessionManager.buildSessionContext(); - this.agent.state.messages = sessionContext.messages; - const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); + this._refreshFinalizedContext(); + const estimatedTokensAfter = estimateMessagesTokens(this.sessionManager.buildSessionProjection().messages); // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as @@ -2307,7 +2596,11 @@ export class AgentSession { * @param skipAbortedCheck If false, include aborted messages (for pre-prompt check). Default: true * @returns Whether the post-run loop should call `agent.continue()` for overflow recovery or queued messages */ - private async _checkCompaction(assistantMessage: AssistantMessage, skipAbortedCheck = true): Promise { + private async _checkCompaction( + assistantMessage: AssistantMessage, + skipAbortedCheck = true, + toolResults: AgentMessage[] = [], + ): Promise { const settings = this.settingsManager.getCompactionSettings(this.model); if (!settings.enabled) return false; @@ -2336,8 +2629,36 @@ export class AgentSession { // Automatic cases 1 and 2: context overflow. // A length stop is recoverable when output ended below the model's original desired limit, // independent of the configured context size or any context-clamped provider request limit. - const contextOverflow = sameModel && isContextOverflow(assistantMessage, contextWindow); - const recoverableLength = sameModel && isRecoverableLength(assistantMessage, this.model?.maxTokens ?? 0); + const currentProjection = this.sessionManager.buildSessionProjection(); + const assistantEntryId = this._findPersistedMessageEntryId(assistantMessage); + const assistantIsProjected = + assistantEntryId === undefined || + currentProjection.entries.some( + (entry) => + entry.sourceEntry.id === assistantEntryId && + entry.messages.some((message) => message.role === "assistant"), + ); + const branch = this.sessionManager.getBranch(); + const assistantIndex = assistantEntryId ? branch.findIndex((entry) => entry.id === assistantEntryId) : -1; + const entriesAfterAssistant = assistantIndex >= 0 ? branch.slice(assistantIndex + 1) : []; + const hasPostAssistantContextEdit = entriesAfterAssistant.some((entry) => entry.type === "context_edit"); + const latestAssistantEdit = entriesAfterAssistant + .filter( + (entry): entry is ContextEditEntry => entry.type === "context_edit" && entry.targetId === assistantEntryId, + ) + .at(-1); + const assistantRetainedForExplicitRecovery = + assistantEntryId === undefined || + (!entriesAfterAssistant.some((entry) => entry.type === "compaction") && + latestAssistantEdit?.replacement !== null); + const assistantUsageMatchesProjection = assistantIsProjected && !hasPostAssistantContextEdit; + const explicitOverflow = assistantMessage.stopReason === "error" && isContextOverflow(assistantMessage); + const contextOverflow = + sameModel && + ((explicitOverflow && assistantRetainedForExplicitRecovery) || + (assistantUsageMatchesProjection && isContextOverflow(assistantMessage, contextWindow))); + const recoverableLength = + sameModel && assistantIsProjected && isRecoverableLength(assistantMessage, this.model?.maxTokens ?? 0); if (contextOverflow || recoverableLength) { const willRetry = assistantMessage.stopReason !== "stop"; @@ -2369,13 +2690,9 @@ export class AgentSession { return false; } - // Case 1: remove the failed or truncated message from agent state, compact, and - // retry once. The message remains in session history but is excluded from retry context. + // Persistently omit the selected final attempt before post-run recovery compaction. this._overflowRecoveryAttempted = true; - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.state.messages = messages.slice(0, -1); - } + this._omitRecoveryAttempt(assistantMessage, toolResults); return await this._runAutoCompaction("overflow", willRetry); } @@ -2384,8 +2701,12 @@ export class AgentSession { // This ensures sessions that hit persistent API errors (e.g. 529) or malformed zero-usage // responses can still compact and do not reset context accounting. let contextTokens: number; + const projection = currentProjection; + const hasContextEdits = projection.entries.some((entry) => entry.sourceEntry.type === "context_edit"); const directContextTokens = assistantMessage.usage ? calculateContextTokens(assistantMessage.usage) : 0; - if (assistantMessage.stopReason === "error" || directContextTokens === 0) { + if (hasContextEdits) { + contextTokens = estimateProjectedContextTokens(projection, branch).tokens; + } else if (assistantMessage.stopReason === "error" || directContextTokens === 0) { const messages = this.agent.state.messages; const estimate = estimateContextTokens(messages); // Without provider usage, estimate.tokens is the pure message-size estimate. @@ -2516,9 +2837,8 @@ export class AgentSession { this.sessionManager.appendCompaction(summary, firstKeptEntryId, tokensBefore, details, fromExtension, usage); const newEntries = this.sessionManager.getEntries(); - const sessionContext = this.sessionManager.buildSessionContext(); - this.agent.state.messages = sessionContext.messages; - const estimatedTokensAfter = estimateMessagesTokens(sessionContext.messages); + this._refreshFinalizedContext(); + const estimatedTokensAfter = estimateMessagesTokens(this.sessionManager.buildSessionProjection().messages); // Get the saved compaction entry for the extension event const savedCompactionEntry = newEntries.find((e) => e.type === "compaction" && e.summary === summary) as @@ -2545,18 +2865,7 @@ export class AgentSession { }; this._emit({ type: "compaction_end", reason, result, aborted: false, willRetry }); - if (willRetry) { - const messages = this.agent.state.messages; - const lastMsg = messages[messages.length - 1]; - // The overflow response was persisted on message_end before _checkCompaction() removed it - // from agent state. Rebuilding state from the new compaction can restore that kept entry, - // leaving an assistant as the final message. agent.continue() rejects that state, so remove - // the retriable error or truncated-length response again before continuing the interrupted turn. - if (lastMsg?.role === "assistant" && (lastMsg.stopReason === "error" || lastMsg.stopReason === "length")) { - this.agent.state.messages = messages.slice(0, -1); - } - return true; - } + if (willRetry) return true; // Auto-compaction can complete while follow-up/steering/custom messages are waiting. // Continue once so queued messages are delivered. @@ -3091,11 +3400,8 @@ export class AgentSession { errorMessage: message.errorMessage || "Unknown error", }); - // Remove error message from agent state (keep in session for history) - const messages = this.agent.state.messages; - if (messages.length > 0 && messages[messages.length - 1].role === "assistant") { - this.agent.state.messages = messages.slice(0, -1); - } + // Keep the failed attempt in raw history while durably omitting it from model projection. + this._omitRecoveryAttempt(message); // Wait with exponential backoff (abortable) this._retryAbortController = new AbortController(); @@ -3205,11 +3511,8 @@ export class AgentSession { // Queue for later - will be flushed on agent_end this._pendingBashMessages.push(bashMessage); } else { - // Add to agent state immediately - this.agent.state.messages.push(bashMessage); - - // Save to session this.sessionManager.appendMessage(bashMessage); + this._refreshFinalizedContext(); } } @@ -3240,14 +3543,10 @@ export class AgentSession { if (this._pendingBashMessages.length === 0) return; for (const bashMessage of this._pendingBashMessages) { - // Add to agent state - this.agent.state.messages.push(bashMessage); - - // Save to session this.sessionManager.appendMessage(bashMessage); } - this._pendingBashMessages = []; + this._refreshFinalizedContext(); } // ========================================================================= @@ -3455,9 +3754,8 @@ export class AgentSession { this.sessionManager.appendLabelChange(targetId, label); } - // Update agent state - const sessionContext = this.sessionManager.buildSessionContext(); - this.agent.state.messages = sessionContext.messages; + // Update finalized context from the canonical session projection. + this._refreshFinalizedContext(); this._restoreToolsFromTranscript(); // Emit session_tree event @@ -3567,33 +3865,32 @@ export class AgentSession { // After compaction, the last assistant usage reflects pre-compaction context size. // We can only trust usage from an assistant that responded after the latest compaction. // If no such assistant exists, context token count is unknown until the next LLM response. - const branchEntries = this.sessionManager.getBranch(); - const latestCompaction = getLatestCompactionEntry(branchEntries); + const projection = this.sessionManager.buildSessionProjection(); + const branch = this.sessionManager.getBranch(); + const latestCompaction = getLatestCompactionEntry(branch); if (latestCompaction) { - // Check if there's a valid assistant usage after the compaction boundary - const compactionIndex = branchEntries.lastIndexOf(latestCompaction); - let hasPostCompactionUsage = false; - for (let i = branchEntries.length - 1; i > compactionIndex; i--) { - const entry = branchEntries[i]; - if (entry.type === "message" && entry.message.role === "assistant") { - const assistant = entry.message; - if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { - const contextTokens = calculateContextTokens(assistant.usage); - if (contextTokens > 0) { - hasPostCompactionUsage = true; - break; - } - } - } - } - - if (!hasPostCompactionUsage) { - return { tokens: null, contextWindow, percent: null }; - } + const projectedAssistants = new Set( + projection.entries.flatMap((entry) => + entry.messages.some( + (message) => + message.role === "assistant" && + message.stopReason !== "aborted" && + message.stopReason !== "error" && + calculateContextTokens(message.usage) > 0, + ) + ? [entry.sourceEntry.id] + : [], + ), + ); + const compactionIndex = branch.findIndex((entry) => entry.id === latestCompaction.id); + const hasPostCompactionUsage = branch + .slice(compactionIndex + 1) + .some((entry) => projectedAssistants.has(entry.id)); + if (!hasPostCompactionUsage) return { tokens: null, contextWindow, percent: null }; } - const estimate = estimateContextTokens(this.messages); + const estimate = estimateProjectedContextTokens(projection, branch); const percent = (estimate.tokens / contextWindow) * 100; return { diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 32a5ed643..ea45ea6e0 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -8,6 +8,7 @@ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; import { contentText, + getCurrentSystemMessage, normalizeContext, type RetryCallbacks, type RetryPolicy, @@ -18,15 +19,18 @@ import type { AssistantMessage, Model, SimpleStreamOptions, + SystemMessage, TranscriptContext, Usage, } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm } from "../messages.ts"; import { - buildSessionContext, + buildSessionProjection, type CompactionEntry, + type ProjectedSessionEntry, type SessionEntry, + type SessionProjection, sessionEntryToContextMessages, } from "../session-manager.ts"; import { @@ -90,13 +94,10 @@ function extractFileOperations( * Extract AgentMessage from an entry if it produces one. * Returns undefined for entries that don't contribute to LLM context. */ -function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | undefined { - if (entry.type === "compaction") { - return undefined; - } +function getMessagesFromProjectedEntryForCompaction(entry: ProjectedSessionEntry): AgentMessage[] { + if (entry.sourceEntry.type === "compaction") return []; // System messages are prompt state, not conversation; the compaction entry carries their replay. - const message = sessionEntryToContextMessages(entry)[0]; - return message?.role === "system" ? undefined : message; + return entry.messages.filter((message) => message.role !== "system"); } /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ @@ -244,6 +245,44 @@ export function estimateContextTokens(messages: AgentMessage[]): ContextUsageEst }; } +/** Estimate projected context without trusting usage captured before a later edit or compaction. */ +export function estimateProjectedContextTokens( + projection: SessionProjection, + branchEntries: SessionEntry[], +): ContextUsageEstimate { + const estimate = estimateContextTokens(projection.messages); + if (estimate.lastUsageIndex !== null) { + let projectedMessageIndex = 0; + let usageEntryId: string | undefined; + for (const entry of projection.entries) { + const nextMessageIndex = projectedMessageIndex + entry.messages.length; + if (estimate.lastUsageIndex < nextMessageIndex) { + usageEntryId = entry.sourceEntry.id; + break; + } + projectedMessageIndex = nextMessageIndex; + } + + const usageEntryIndex = usageEntryId ? branchEntries.findIndex((entry) => entry.id === usageEntryId) : -1; + let latestInvalidatingEntryIndex = -1; + for (let i = branchEntries.length - 1; i >= 0; i--) { + const entry = branchEntries[i]; + if (entry.type === "context_edit" || entry.type === "compaction") { + latestInvalidatingEntryIndex = i; + break; + } + } + if (usageEntryIndex > latestInvalidatingEntryIndex) return estimate; + } + + const currentSystem = getCurrentSystemMessage(projection.messages); + let tokens = currentSystem ? estimateTokens(currentSystem) : 0; + for (const message of projection.messages) { + if (message.role !== "system") tokens += estimateTokens(message); + } + return { tokens, usageTokens: 0, trailingTokens: tokens, lastUsageIndex: null }; +} + /** * Check if compaction should trigger based on context usage. */ @@ -282,6 +321,17 @@ export function estimateTokens(message: AgentMessage): number { let chars = 0; switch (message.role) { + case "system": { + const system = message as SystemMessage; + chars = estimateTextAndImageContentChars(system.content); + if (system.sections) { + for (const section of Object.values(system.sections)) { + if (section) chars += section.length; + } + } + if (system.toolsAdded) chars += JSON.stringify(system.toolsAdded).length; + return Math.ceil(chars / 4); + } case "user": { chars = estimateTextAndImageContentChars( (message as { content: string | Array<{ type: string; text?: string }> }).content, @@ -759,6 +809,88 @@ export interface CompactionPreparation { settings: CompactionSettings; } +function isProjectedTurnStart(entry: ProjectedSessionEntry): boolean { + if (entry.sourceEntry.type === "compaction") return false; + return entry.messages.some(isTurnStartMessage); +} + +function findProjectedTurnStartIndex(entries: ProjectedSessionEntry[], entryIndex: number, startIndex: number): number { + for (let i = entryIndex; i >= startIndex; i--) { + if (isProjectedTurnStart(entries[i])) return i; + } + return -1; +} + +function findProjectedCutPoint( + entries: ProjectedSessionEntry[], + startIndex: number, + endIndex: number, + keepRecentTokens: number, +): CutPointResult { + const cutPoints: number[] = []; + for (let i = startIndex; i < endIndex; i++) { + const entry = entries[i]; + if (entry.sourceEntry.type !== "compaction" && entry.messages.some(isCutPointMessage)) cutPoints.push(i); + } + if (cutPoints.length === 0) { + return { firstKeptEntryIndex: startIndex, turnStartIndex: -1, isSplitTurn: false }; + } + + let accumulatedTokens = 0; + let exceededBudget = false; + let cutIndex = cutPoints[0]; + for (let i = endIndex - 1; i >= startIndex; i--) { + const messageTokens = entries[i].messages.reduce((sum, message) => sum + estimateTokens(message), 0); + if (messageTokens === 0) continue; + accumulatedTokens += messageTokens; + if (accumulatedTokens >= keepRecentTokens) { + exceededBudget = true; + cutIndex = cutPoints.find((candidate) => candidate >= i) ?? cutPoints[cutPoints.length - 1]; + break; + } + } + + // A recovery attempt and its omission edits are context-invisible after the last + // visible input. Advance only for a closed suffix containing an omitted assistant + // attempt; arbitrary metadata must not move the cut past unsent input. + const suffix = entries.slice(cutIndex + 1, endIndex); + const isIntrinsicallyVisible = (entry: ProjectedSessionEntry): boolean => + entry.sourceEntry.type !== "context_edit" && sessionEntryToContextMessages(entry.sourceEntry).length > 0; + const isOmitted = (entry: ProjectedSessionEntry): boolean => + isIntrinsicallyVisible(entry) && entry.messages.length === 0; + const omittedSuffixIds = new Set(suffix.filter(isOmitted).map((entry) => entry.sourceEntry.id)); + const hasExternalReplacement = suffix.some( + (entry) => + entry.sourceEntry.type === "context_edit" && + entry.sourceEntry.replacement !== null && + !omittedSuffixIds.has(entry.sourceEntry.targetId), + ); + const isRecoveryOmissionSuffix = + exceededBudget && + !hasExternalReplacement && + suffix.some( + (entry) => + entry.sourceEntry.type === "message" && entry.sourceEntry.message.role === "assistant" && isOmitted(entry), + ) && + suffix.every( + (entry) => entry.sourceEntry.type !== "compaction" && (!isIntrinsicallyVisible(entry) || isOmitted(entry)), + ); + if (isRecoveryOmissionSuffix) cutIndex++; + + while (cutIndex > startIndex) { + const previous = entries[cutIndex - 1]; + if (previous.sourceEntry.type === "compaction" || previous.messages.length > 0) break; + cutIndex--; + } + const startsTurn = isProjectedTurnStart(entries[cutIndex]); + const turnStartIndex = startsTurn ? -1 : findProjectedTurnStartIndex(entries, cutIndex, startIndex); + return { + firstKeptEntryIndex: cutIndex, + turnStartIndex, + isSplitTurn: !startsTurn && turnStartIndex !== -1, + }; +} + export function prepareCompaction( pathEntries: SessionEntry[], settings: CompactionSettings, @@ -767,59 +899,44 @@ export function prepareCompaction( return undefined; } - let prevCompactionIndex = -1; - for (let i = pathEntries.length - 1; i >= 0; i--) { - if (pathEntries[i].type === "compaction") { - prevCompactionIndex = i; - break; - } - } + const projection = buildSessionProjection(pathEntries); + const projectedEntries = projection.entries; + const sourceEntries = projectedEntries.map((entry) => entry.sourceEntry); + // The newest compaction is projected first. Older compaction entries can still + // occur in its retained raw range, but their projected contribution is empty. + const prevCompactionIndex = projectedEntries.findIndex( + (entry) => entry.sourceEntry.type === "compaction" && entry.messages.length > 0, + ); let previousSummary: string | undefined; let boundaryStart = 0; if (prevCompactionIndex >= 0) { - const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; - previousSummary = prevCompaction.summary; - const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); - boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; + previousSummary = (projectedEntries[prevCompactionIndex].sourceEntry as CompactionEntry).summary; + // The canonical projection has already selected the previous compaction's retained tail. + boundaryStart = prevCompactionIndex + 1; } - const boundaryEnd = pathEntries.length; + const boundaryEnd = projectedEntries.length; + const tokensBefore = estimateProjectedContextTokens(projection, pathEntries).tokens; + const cutPoint = findProjectedCutPoint(projectedEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens); - const tokensBefore = estimateContextTokens(buildSessionContext(pathEntries).messages).tokens; - - const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, settings.keepRecentTokens); - - // Get UUID of first kept entry - const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex]; - if (!firstKeptEntry?.id) { - return undefined; // Session needs migration - } + const firstKeptEntry = projectedEntries[cutPoint.firstKeptEntryIndex]?.sourceEntry; + if (!firstKeptEntry?.id) return undefined; const firstKeptEntryId = firstKeptEntry.id; - const historyEnd = cutPoint.isSplitTurn ? cutPoint.turnStartIndex : cutPoint.firstKeptEntryIndex; - // Messages to summarize (will be discarded after summary) - const messagesToSummarize: AgentMessage[] = []; - for (let i = boundaryStart; i < historyEnd; i++) { - const msg = getMessageFromEntryForCompaction(pathEntries[i]); - if (msg) messagesToSummarize.push(msg); - } + const messagesToSummarize = projectedEntries + .slice(boundaryStart, historyEnd) + .flatMap(getMessagesFromProjectedEntryForCompaction); + const turnPrefixMessages = cutPoint.isSplitTurn + ? projectedEntries + .slice(cutPoint.turnStartIndex, cutPoint.firstKeptEntryIndex) + .flatMap(getMessagesFromProjectedEntryForCompaction) + : []; - // Messages for turn prefix summary (if splitting a turn) - const turnPrefixMessages: AgentMessage[] = []; - if (cutPoint.isSplitTurn) { - for (let i = cutPoint.turnStartIndex; i < cutPoint.firstKeptEntryIndex; i++) { - const msg = getMessageFromEntryForCompaction(pathEntries[i]); - if (msg) turnPrefixMessages.push(msg); - } - } + if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) return undefined; - if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) { - return undefined; - } - - // Extract file operations from messages and previous compaction - const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); + // Extract file operations from edited model-visible messages and the previous compaction. + const fileOps = extractFileOperations(messagesToSummarize, sourceEntries, prevCompactionIndex); // Also extract file ops from turn prefix if splitting if (cutPoint.isSplitTurn) { diff --git a/packages/coding-agent/src/core/export-html/template.js b/packages/coding-agent/src/core/export-html/template.js index cfdd16120..e4578bfcc 100644 --- a/packages/coding-agent/src/core/export-html/template.js +++ b/packages/coding-agent/src/core/export-html/template.js @@ -357,6 +357,9 @@ case 'thinking_level_change': parts.push('thinking', entry.thinkingLevel); break; + case 'context_edit': + parts.push('context edit', entry.replacement === null ? 'omit' : 'replace', entry.targetId); + break; } return parts.join(' ').toLowerCase(); @@ -385,7 +388,7 @@ } // Apply filter mode - const isSettingsEntry = ['label', 'custom', 'model_change', 'thinking_level_change'].includes(entry.type); + const isSettingsEntry = ['label', 'custom', 'context_edit', 'model_change', 'thinking_level_change'].includes(entry.type); let passesFilter = true; switch (filterMode) { @@ -696,6 +699,8 @@ return labelHtml + `[model: ${escapeHtml(entry.modelId)}]`; case 'thinking_level_change': return labelHtml + `[thinking: ${escapeHtml(entry.thinkingLevel)}]`; + case 'context_edit': + return labelHtml + `[context ${entry.replacement === null ? 'omit' : 'replace'}: ${escapeHtml(entry.targetId)}]`; default: return labelHtml + `[${escapeHtml(entry.type)}]`; } diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index 35d255817..57b2e77bf 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -21,6 +21,9 @@ export type { export { ExtensionRunner } from "./runner.ts"; export type { AfterProviderResponseEvent, + AgentActivityOutcome, + AgentBeforeSettleEvent, + AgentBeforeSettleEventResult, AgentEndEvent, AgentSettledEvent, AgentStartEvent, @@ -39,16 +42,23 @@ export type { BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, + BoundaryContextPreview, + BoundaryResult, + BoundaryState, BuildSystemPromptOptions, CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult, + CompactionEntryDraft, // Context CompactOptions, + ContextEditEntryDraft, // Events - Agent ContextEvent, // Event Results ContextEventResult, ContextUsage, + CustomEntryDraft, + CustomMessageEntryDraft, CustomToolCallEvent, CustomToolResultEvent, EditorFactory, @@ -139,6 +149,7 @@ export type { SessionBeforeSwitchResult, SessionBeforeTreeEvent, SessionBeforeTreeResult, + SessionBoundaryDraft, SessionCompactEvent, SessionCompactFailedEvent, SessionEvent, @@ -170,6 +181,7 @@ export type { ToolResultEventResult, TreePreparation, TurnEndEvent, + TurnEndEventResult, TurnStartEvent, UIPromptEndEvent, UIPromptKind, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 3ae663fba..44576a1bb 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -19,10 +19,13 @@ import { normalizeBuildSystemPromptOptions, } from "../system-prompt.ts"; import type { + AgentBeforeSettleEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, BeforeProviderHeadersEvent, BeforeProviderRequestEvent, + BoundaryContextPreview, + BoundaryResult, CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult, CompactOptions, @@ -65,11 +68,13 @@ import type { SessionBeforeForkResult, SessionBeforeSwitchResult, SessionBeforeTreeResult, + SessionBoundaryDraft, SessionShutdownEvent, ToolCallEvent, ToolCallEventResult, ToolResultEvent, ToolResultEventResult, + TurnEndEvent, UIPromptKind, UserBashEvent, UserBashEventResult, @@ -171,6 +176,8 @@ type RunnerEmitEvent = Exclude< | MessageEndEvent | ResourcesDiscoverEvent | InputEvent + | TurnEndEvent + | AgentBeforeSettleEvent >; type SessionBeforeEvent = Extract< @@ -196,6 +203,17 @@ type RunnerEmitResult = TEvent extends { type: " export type ExtensionErrorListener = (error: ExtensionError) => void; +type BoundaryBaseEvent = + | Omit + | Omit; + +interface BoundaryDispatchResult { + entries: SessionBoundaryDraft[]; + continue: boolean; + context: BoundaryContextPreview; + valid: boolean; +} + export type NewSessionHandler = (options?: { parentSession?: string; setup?: (sessionManager: SessionManager) => Promise; @@ -877,6 +895,57 @@ export class ExtensionRunner { return context; } + async emitBoundary( + baseEvent: BoundaryBaseEvent, + buildContext: (entries: SessionBoundaryDraft[]) => BoundaryContextPreview | Promise, + ): Promise { + const ctx = this.createContext(); + let entries: SessionBoundaryDraft[] = []; + let shouldContinue = false; + let context = await buildContext(entries); + let valid = true; + + for (const { ext, handlers } of snapshotEventHandlers(this.extensions, baseEvent.type)) { + for (const handler of handlers) { + const event = { + ...baseEvent, + entries, + continue: shouldContinue, + context, + } as TurnEndEvent | AgentBeforeSettleEvent; + try { + const handlerResult = (await handler(event, ctx)) as BoundaryResult | undefined; + if (handlerResult?.entries !== undefined) entries = handlerResult.entries; + if (handlerResult?.continue !== undefined) shouldContinue = handlerResult.continue; + } catch (err) { + this.emitError({ + extensionPath: ext.path, + event: baseEvent.type, + error: err instanceof Error ? err.message : String(err), + stack: err instanceof Error ? err.stack : undefined, + }); + } + + try { + context = await buildContext(entries); + valid = true; + } catch (err) { + valid = false; + this.emitError({ + extensionPath: ext.path, + event: baseEvent.type, + error: `Invalid boundary entries: ${err instanceof Error ? err.message : String(err)}`, + stack: err instanceof Error ? err.stack : undefined, + }); + } + } + } + + return valid + ? { entries, continue: shouldContinue, context, valid: true } + : { entries: [], continue: false, context, valid: false }; + } + private isSessionBeforeEvent(event: RunnerEmitEvent): event is SessionBeforeEvent { return ( event.type === "session_before_switch" || diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index c470a36c9..d10078701 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -21,6 +21,7 @@ import type { AssistantMessageEventStream, ConstrainedSamplingConfig, ImageContent, + Message, Model, OAuthCredentials, OAuthLoginCallbacks, @@ -59,7 +60,9 @@ import type { ScopedModel } from "../model-resolver.ts"; import type { BranchSummaryEntry, CompactionEntry, + ContextEditEntry, CustomEntry, + ProjectedSessionEntry, ReadonlySessionManager, SessionEntry, SessionManager, @@ -738,6 +741,68 @@ export interface AgentEndEvent { messages: AgentMessage[]; } +export type AgentActivityOutcome = "completed" | "aborted" | "error"; + +export interface CustomEntryDraft { + type: "custom"; + customType: string; + data?: unknown; +} + +export interface CustomMessageEntryDraft { + type: "custom_message"; + customType: string; + content: string | (TextContent | ImageContent)[]; + display: boolean; + details?: unknown; +} + +export interface ContextEditEntryDraft { + type: "context_edit"; + targetId: string; + replacement: ContextEditEntry["replacement"]; +} + +export interface CompactionEntryDraft { + type: "compaction"; + summary: string; + /** Null creates a self-retaining compaction that keeps no preceding entries. */ + firstKeptEntryId: string | null; + details?: unknown; + usage?: Usage; +} + +export type SessionBoundaryDraft = + | CustomEntryDraft + | CustomMessageEntryDraft + | ContextEditEntryDraft + | CompactionEntryDraft; + +export interface BoundaryContextPreview { + contextEntries: ProjectedSessionEntry[]; + contextMessages: AgentMessage[]; + llmMessages: Message[]; + pendingMessages: AgentMessage[]; + canContinue: boolean; +} + +export interface BoundaryState { + entries: SessionBoundaryDraft[]; + continue: boolean; + context: BoundaryContextPreview; + outcome: AgentActivityOutcome; +} + +export interface BoundaryResult { + entries?: SessionBoundaryDraft[]; + continue?: boolean; +} + +/** Fired before final settlement. May append entries and ensure one next provider request. */ +export interface AgentBeforeSettleEvent extends BoundaryState { + type: "agent_before_settle"; +} + /** Fired after an agent run has fully settled and no automatic retry, compaction, or queued continuation will run. */ export interface AgentSettledEvent { type: "agent_settled"; @@ -769,11 +834,13 @@ export interface TurnStartEvent { } /** Fired at the end of each turn */ -export interface TurnEndEvent { +export interface TurnEndEvent extends BoundaryState { type: "turn_end"; turnIndex: number; message: AgentMessage; toolResults: ToolResultMessage[]; + messageEntryId: string; + toolResultEntryIds: string[]; } /** Fired when a message starts (user, assistant, or toolResult) */ @@ -1096,6 +1163,7 @@ export type ExtensionEvent = | BeforeAgentStartEvent | AgentStartEvent | AgentEndEvent + | AgentBeforeSettleEvent | AgentSettledEvent | UIPromptStartEvent | UIPromptEndEvent @@ -1122,6 +1190,9 @@ export interface ContextEventResult { messages?: AgentMessage[]; } +export type TurnEndEventResult = BoundaryResult; +export type AgentBeforeSettleEventResult = BoundaryResult; + export type BeforeProviderRequestEventResult = unknown; export type { CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult } from "../cache-warmer.ts"; @@ -1307,11 +1378,15 @@ export interface ExtensionAPI { ): () => void; on(event: "agent_start", handler: ExtensionHandler): () => void; on(event: "agent_end", handler: ExtensionHandler): () => void; + on( + event: "agent_before_settle", + handler: ExtensionHandler, + ): () => void; on(event: "agent_settled", handler: ExtensionHandler): () => void; on(event: "ui_prompt_start", handler: ExtensionHandler): () => void; on(event: "ui_prompt_end", handler: ExtensionHandler): () => void; on(event: "turn_start", handler: ExtensionHandler): () => void; - on(event: "turn_end", handler: ExtensionHandler): () => void; + on(event: "turn_end", handler: ExtensionHandler): () => void; on(event: "message_start", handler: ExtensionHandler): () => void; on(event: "message_update", handler: ExtensionHandler): () => void; on(event: "message_end", handler: ExtensionHandler): () => void; diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index aceb1e13e..178ed7105 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -32,6 +32,9 @@ export { createEventBus, type EventBus, type EventBusController } from "./event- export { areExperimentalFeaturesEnabled } from "./experimental.ts"; // Extensions system export { + type AgentActivityOutcome, + type AgentBeforeSettleEvent, + type AgentBeforeSettleEventResult, type AgentEndEvent, type AgentSettledEvent, type AgentStartEvent, @@ -39,10 +42,17 @@ export { type AgentToolUpdateCallback, type BeforeAgentStartEvent, type BeforeAgentStartEventResult, + type BoundaryContextPreview, + type BoundaryResult, + type BoundaryState, type BuildSystemPromptOptions, type CacheWarmingDecisionEvent, type CacheWarmingDecisionEventResult, + type CompactionEntryDraft, + type ContextEditEntryDraft, type ContextEvent, + type CustomEntryDraft, + type CustomMessageEntryDraft, defineTool, discoverAndLoadExtensions, type ExecOptions, @@ -68,6 +78,7 @@ export { type SessionBeforeForkEvent, type SessionBeforeSwitchEvent, type SessionBeforeTreeEvent, + type SessionBoundaryDraft, type SessionCompactEvent, type SessionShutdownEvent, type SessionStartEvent, @@ -78,6 +89,7 @@ export { type ToolRenderResultOptions, type ToolResultEvent, type TurnEndEvent, + type TurnEndEventResult, type TurnStartEvent, type WorkingIndicatorOptions, } from "./extensions/index.ts"; diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 98a45ed14..2ab852514 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -264,8 +264,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} options.tools ?? (options.noTools ? [] : (configuredDefaultToolNames ?? defaultActiveToolNames)) ).filter((name) => !excludedToolNameSet?.has(name)); - let agent: Agent; - // Create convertToLlm wrapper that filters images if blockImages is enabled (defense-in-depth) const convertToLlmWithBlockImages = (messages: AgentMessage[]): Message[] => { const converted = convertToLlm(messages); @@ -365,12 +363,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }); }; - agent = new Agent({ + const agent = new Agent({ initialState: { systemPrompt: "", model, thinkingLevel, tools: [], + messages: existingSession.messages, }, convertToLlm: convertToLlmWithBlockImages, streamFn: async (model, context, options) => { @@ -400,9 +399,8 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} maxRetryDelayMs: settingsManager.getProviderRetrySettings().maxRetryDelayMs, }); - // Restore messages if session has existing data + // Restore missing settings metadata for older sessions. if (hasExistingSession) { - agent.state.messages = existingSession.messages; if (!hasThinkingEntry) { sessionManager.appendThinkingLevelChange(thinkingLevel); } diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 6d507d8fe..fd87a2d88 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -1,11 +1,14 @@ import type { AgentMessage } from "@earendil-works/pi-agent-core"; import { + type AssistantMessage, getCurrentSystemMessage, type ImageContent, type Message, type SystemMessage, type TextContent, + type ToolResultMessage, type Usage, + type UserMessage, uuidv7, } from "@earendil-works/pi-ai"; import { randomUUID } from "crypto"; @@ -161,6 +164,21 @@ export interface CustomMessageEntry extends SessionEntryBase { display: boolean; } +/** Content that an append-only context edit may replace without changing message metadata. */ +export type ContextEditableContent = + | UserMessage["content"] + | AssistantMessage["content"] + | ToolResultMessage["content"] + | CustomMessage["content"]; + +/** Append-only change to one earlier entry's contribution to model context. */ +export interface ContextEditEntry extends SessionEntryBase { + type: "context_edit"; + targetId: string; + /** Null omits the target from model context. A value replaces only its content. */ + replacement: { content: ContextEditableContent } | null; +} + /** Session entry - has id/parentId for tree structure (returned by "read" methods in SessionManager) */ export type SessionEntry = | SessionMessageEntry @@ -171,6 +189,7 @@ export type SessionEntry = | BranchSummaryEntry | CustomEntry | CustomMessageEntry + | ContextEditEntry | LabelEntry | SessionInfoEntry; @@ -187,6 +206,20 @@ export interface SessionTreeNode { labelTimestamp?: string; } +export interface ProjectedSessionEntry { + /** Raw append-only entry that owns this projected contribution. */ + sourceEntry: SessionEntry; + /** Model-visible messages after context edits. Empty for state-only entries and omissions. */ + messages: AgentMessage[]; +} + +export interface SessionProjection { + entries: ProjectedSessionEntry[]; + messages: AgentMessage[]; + thinkingLevel: string; + model: { provider: string; modelId: string } | null; +} + export interface SessionContext { messages: AgentMessage[]; thinkingLevel: string; @@ -221,6 +254,7 @@ export type ReadonlySessionManager = Pick< | "getLabel" | "getBranch" | "buildContextEntries" + | "buildSessionProjection" | "getHeader" | "getEntries" | "getTree" @@ -482,14 +516,69 @@ export function buildContextEntries( * If leafId is provided, walks from that entry to root. * Handles compaction and branch summaries along the path. */ +function projectContextEntry(entry: SessionEntry, edit: ContextEditEntry | undefined): AgentMessage[] { + const messages = sessionEntryToContextMessages(entry); + if (!edit) return messages; + const replacement = edit.replacement; + if (replacement === null) return []; + + return messages.map((message) => { + if ( + message.role !== "user" && + message.role !== "assistant" && + message.role !== "toolResult" && + message.role !== "custom" + ) { + return message; + } + const content = + (message.role === "assistant" || message.role === "toolResult") && typeof replacement.content === "string" + ? [{ type: "text" as const, text: replacement.content }] + : replacement.content; + return { ...message, content } as AgentMessage; + }); +} + +/** Build provenance-preserving, compaction-aware model context. */ +export function buildSessionProjection( + entries: SessionEntry[], + leafId?: string | null, + byId?: Map, +): SessionProjection { + const path = buildSessionPath(entries, leafId, byId); + const { thinkingLevel, model } = getSessionContextSettings(path); + const contextEntries = buildContextEntries(entries, leafId, byId); + const edits = new Map(); + for (const entry of contextEntries) { + if (entry.type === "context_edit") edits.set(entry.targetId, entry); + } + const projectedEntries = contextEntries.map( + (sourceEntry, index): ProjectedSessionEntry => ({ + sourceEntry, + // buildContextEntries() may retain an older compaction entry because its + // raw ID lies inside the newest retained range. Only the newest compaction + // at index zero contributes a checkpoint and summary. + messages: + sourceEntry.type === "compaction" && index > 0 + ? [] + : projectContextEntry(sourceEntry, edits.get(sourceEntry.id)), + }), + ); + return { + entries: projectedEntries, + messages: projectedEntries.flatMap((entry) => entry.messages), + thinkingLevel, + model, + }; +} + +/** Build the finalized model context from the canonical session projection. */ export function buildSessionContext( entries: SessionEntry[], leafId?: string | null, byId?: Map, ): SessionContext { - const path = buildSessionPath(entries, leafId, byId); - const { thinkingLevel, model } = getSessionContextSettings(path); - const messages = buildContextEntries(entries, leafId, byId).flatMap(sessionEntryToContextMessages); + const { messages, thinkingLevel, model } = buildSessionProjection(entries, leafId, byId); return { messages, thinkingLevel, model }; } @@ -1169,21 +1258,22 @@ export class SessionManager { /** Append a compaction summary as child of current leaf, then advance leaf. Returns entry id. */ appendCompaction( summary: string, - firstKeptEntryId: string, + firstKeptEntryId: string | null, tokensBefore: number, details?: T, fromHook?: boolean, usage?: Usage, ): string { const timestamp = new Date().toISOString(); - const systemMessage = getCurrentSystemMessage(this.buildSessionContext().messages); + const systemMessage = getCurrentSystemMessage(this.buildSessionProjection().messages); + const id = generateId(this.byId); const entry: CompactionEntry = { type: "compaction", - id: generateId(this.byId), + id, parentId: this.leafId, timestamp, summary, - firstKeptEntryId, + firstKeptEntryId: firstKeptEntryId ?? id, tokensBefore, details, usage, @@ -1264,6 +1354,47 @@ export class SessionManager { return entry.id; } + /** Append a branch-local edit to an earlier model-visible entry. */ + appendContextEdit(targetId: string, replacement: ContextEditEntry["replacement"]): string { + if ( + replacement !== null && + (typeof replacement !== "object" || + !("content" in replacement) || + (typeof replacement.content !== "string" && !Array.isArray(replacement.content))) + ) { + throw new Error("Context edit replacement must be null or contain string/array content"); + } + const target = this.byId.get(targetId); + if (!target) throw new Error(`Entry ${targetId} not found`); + if (!this.getBranch().some((entry) => entry.id === targetId)) { + throw new Error(`Entry ${targetId} is not on the active branch`); + } + const editable = + target.type === "custom_message" || + (target.type === "message" && + (target.message.role === "user" || + target.message.role === "assistant" || + target.message.role === "toolResult")); + if (!editable) throw new Error(`Entry ${targetId} does not contribute editable model content`); + const targetRole = target.type === "message" ? target.message.role : "custom"; + const normalizedReplacement = + replacement !== null && + (targetRole === "assistant" || targetRole === "toolResult") && + typeof replacement.content === "string" + ? { content: [{ type: "text" as const, text: replacement.content }] } + : replacement; + const entry: ContextEditEntry = { + type: "context_edit", + id: generateId(this.byId), + parentId: this.leafId, + timestamp: new Date().toISOString(), + targetId, + replacement: normalizedReplacement, + }; + this._appendEntry(entry); + return entry.id; + } + // ========================================================================= // Tree Traversal // ========================================================================= @@ -1357,8 +1488,13 @@ export class SessionManager { * Build the session context (what gets sent to the LLM). * Uses tree traversal from current leaf. */ + buildSessionProjection(): SessionProjection { + return buildSessionProjection(this.getEntries(), this.leafId, this.byId); + } + buildSessionContext(): SessionContext { - return buildSessionContext(this.getEntries(), this.leafId, this.byId); + const { messages, thinkingLevel, model } = this.buildSessionProjection(); + return { messages, thinkingLevel, model }; } /** @@ -1514,7 +1650,10 @@ export class SessionManager { ? { ...entry, parentId: pathParentId, - firstKeptEntryId: replacementByLabelId.get(entry.firstKeptEntryId) ?? entry.firstKeptEntryId, + firstKeptEntryId: + entry.firstKeptEntryId === entry.id + ? entry.id + : (replacementByLabelId.get(entry.firstKeptEntryId) ?? entry.firstKeptEntryId), } : { ...entry, parentId: pathParentId }, ); diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index d4f2009ec..86907ebba 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -53,6 +53,9 @@ export { createEventBus, type EventBus, type EventBusController } from "./core/e // Extension system export type { AfterProviderResponseEvent, + AgentActivityOutcome, + AgentBeforeSettleEvent, + AgentBeforeSettleEventResult, AgentEndEvent, AgentSettledEvent, AgentStartEvent, @@ -66,13 +69,20 @@ export type { BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, + BoundaryContextPreview, + BoundaryResult, + BoundaryState, BuildSystemPromptOptions, CacheWarmingDecisionEvent, CacheWarmingDecisionEventResult, + CompactionEntryDraft, CompactOptions, + ContextEditEntryDraft, ContextEvent, ContextEventResult, ContextUsage, + CustomEntryDraft, + CustomMessageEntryDraft, CustomToolCallEvent, EditToolCallEvent, EntryRenderer, @@ -138,6 +148,7 @@ export type { SessionBeforeSwitchResult, SessionBeforeTreeEvent, SessionBeforeTreeResult, + SessionBoundaryDraft, SessionCompactEvent, SessionCompactFailedEvent, SessionInfoChangedEvent, @@ -161,6 +172,7 @@ export type { ToolResultEvent, ToolResultEventResult, TurnEndEvent, + TurnEndEventResult, TurnStartEvent, UIPromptEndEvent, UIPromptKind, @@ -251,7 +263,10 @@ export { type BranchSummaryEntry, buildContextEntries, buildSessionContext, + buildSessionProjection, type CompactionEntry, + type ContextEditableContent, + type ContextEditEntry, CURRENT_SESSION_VERSION, type CustomEntry, type CustomMessageEntry, @@ -260,6 +275,7 @@ export { type ModelChangeEntry, migrateSessionEntries, type NewSessionOptions, + type ProjectedSessionEntry, parseSessionEntries, type SessionContext, type SessionEntry, @@ -269,6 +285,7 @@ export { type SessionInfoEntry, SessionManager, type SessionMessageEntry, + type SessionProjection, type SessionTreeNode, sessionEntryToContextMessages, type ThinkingLevelChangeEntry, diff --git a/packages/coding-agent/src/modes/interactive/components/tree-selector.ts b/packages/coding-agent/src/modes/interactive/components/tree-selector.ts index 336e19c52..577bfa289 100644 --- a/packages/coding-agent/src/modes/interactive/components/tree-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/tree-selector.ts @@ -358,6 +358,7 @@ class TreeList implements Component { // Entry types hidden in default view (settings/bookkeeping) const isSettingsEntry = entry.type === "label" || + entry.type === "context_edit" || entry.type === "custom" || entry.type === "model_change" || entry.type === "thinking_level_change" || @@ -607,6 +608,9 @@ class TreeList implements Component { case "custom": parts.push("custom", entry.customType); break; + case "context_edit": + parts.push("context edit", entry.replacement === null ? "omit" : "replace", entry.targetId); + break; case "label": parts.push("label", entry.label ?? ""); break; @@ -837,6 +841,9 @@ class TreeList implements Component { case "custom": result = theme.fg("dim", `[custom: ${entry.customType}]`); break; + case "context_edit": + result = theme.fg("dim", `[context ${entry.replacement === null ? "omit" : "replace"}: ${entry.targetId}]`); + break; case "label": result = theme.fg("dim", `[label: ${entry.label ?? "(cleared)"}]`); break; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 96cae384f..968f74140 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -92,7 +92,7 @@ import type { import { FooterDataProvider, type ReadonlyFooterDataProvider } from "../../core/footer-data-provider.ts"; import { configureHttpDispatcher, formatHttpIdleTimeoutMs } from "../../core/http-dispatcher.ts"; import { type AppKeybinding, KeybindingsManager } from "../../core/keybindings.ts"; -import { createCompactionSummaryMessage } from "../../core/messages.ts"; +import { createCompactionSummaryMessage, createCustomMessage } from "../../core/messages.ts"; import { defaultModelPerProvider, findExactModelReferenceMatch, @@ -465,6 +465,7 @@ export class InteractiveMode { // Streaming message tracking private streamingComponent: AssistantMessageComponent | undefined = undefined; + private readonly entriesRenderedByBoundaryCompaction = new Set(); private streamingMessage: AssistantMessage | undefined = undefined; // Tool execution tracking: toolCallId -> component @@ -3318,12 +3319,47 @@ export class InteractiveMode { break; case "entry_appended": + if (this.entriesRenderedByBoundaryCompaction.delete(event.entry.id)) break; if (event.entry.type === "custom") { this.addCustomEntryToChat(event.entry); this.ui.requestRender(); } else if (event.entry.type === "usage" && event.entry.kind === "cache_warm") { this.addCacheWarmingUsage(event.entry); this.ui.requestRender(); + } else if (event.entry.type === "custom_message" && event.entry.display) { + this.addMessageToChat( + createCustomMessage( + event.entry.customType, + event.entry.content, + event.entry.display, + event.entry.details, + event.entry.timestamp, + ), + ); + this.ui.requestRender(); + } else if (event.entry.type === "compaction") { + const entries = this.sessionManager.buildContextEntries(); + if (entries[0]?.id !== event.entry.id) break; + this.chatContainer.clear(); + const branch = this.sessionManager.getBranch(); + const compactionIndex = branch.findIndex((entry) => entry.id === event.entry.id); + const entriesAfterCompaction = new Set(branch.slice(compactionIndex + 1).map((entry) => entry.id)); + const retainedEntries = entries.slice(1); + this.renderSessionEntries(retainedEntries.filter((entry) => !entriesAfterCompaction.has(entry.id))); + this.addMessageToChat( + createCompactionSummaryMessage(event.entry.summary, event.entry.tokensBefore, event.entry.timestamp), + ); + if (event.entry.usage) { + this.addCompactionCostNotice({ + type: "compaction_cost", + kind: "compaction", + usage: event.entry.usage, + }); + } + this.renderSessionEntries(retainedEntries.filter((entry) => entriesAfterCompaction.has(entry.id))); + for (const entryId of entriesAfterCompaction) this.entriesRenderedByBoundaryCompaction.add(entryId); + this.footer.invalidate(); + this.ui.requestRender(); } break; diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index e77b5fd08..2660fb4df 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -887,6 +887,149 @@ describe("ExtensionRunner", () => { }); }); + describe("boundary chaining", () => { + it("chains shared draft proposals and preserves omitted result fields", async () => { + const runtime = createExtensionRuntime(); + const eventBus = createEventBus(); + const observations: Array<{ entries: number; continuation: boolean; preview: number }> = []; + const first = await loadExtensionFromFactory( + (pi) => { + pi.on("agent_before_settle", (event) => { + observations.push({ + entries: event.entries.length, + continuation: event.continue, + preview: event.context.contextEntries.length, + }); + event.entries.push({ type: "custom", customType: "first", data: 1 }); + return { continue: true }; + }); + }, + tempDir, + eventBus, + runtime, + "", + ); + const second = await loadExtensionFromFactory( + (pi) => { + pi.on("agent_before_settle", (event) => { + observations.push({ + entries: event.entries.length, + continuation: event.continue, + preview: event.context.contextEntries.length, + }); + return { entries: [] }; + }); + }, + tempDir, + eventBus, + runtime, + "", + ); + const runner = new ExtensionRunner([first, second], runtime, tempDir, sessionManager, modelRegistry); + + const result = await runner.emitBoundary({ type: "agent_before_settle", outcome: "completed" }, (entries) => ({ + contextEntries: entries.map((entry, index) => ({ + sourceEntry: { + type: "custom", + id: `draft-${index}`, + parentId: null, + timestamp: "", + customType: entry.type, + }, + messages: [], + })), + contextMessages: [], + llmMessages: [], + pendingMessages: [], + canContinue: false, + })); + + expect(observations).toEqual([ + { entries: 0, continuation: false, preview: 0 }, + { entries: 1, continuation: true, preview: 1 }, + ]); + expect(result.entries).toEqual([]); + expect(result.continue).toBe(true); + }); + + it("reports invalid boundary previews and lets later handlers repair the proposal", async () => { + const runtime = createExtensionRuntime(); + let secondRan = false; + const first = await loadExtensionFromFactory( + (pi) => { + pi.on("agent_before_settle", () => ({ + entries: [{ type: "context_edit", targetId: "missing", replacement: null }], + })); + }, + tempDir, + createEventBus(), + runtime, + "", + ); + const second = await loadExtensionFromFactory( + (pi) => { + pi.on("agent_before_settle", (event) => { + secondRan = true; + expect(event.entries).toHaveLength(1); + return { entries: [] }; + }); + }, + tempDir, + createEventBus(), + runtime, + "", + ); + const runner = new ExtensionRunner([first, second], runtime, tempDir, sessionManager, modelRegistry); + const errors: string[] = []; + runner.onError((error) => errors.push(error.error)); + + const result = await runner.emitBoundary({ type: "agent_before_settle", outcome: "completed" }, (entries) => { + if (entries.some((entry) => entry.type === "context_edit")) throw new Error("Entry missing not found"); + return { + contextEntries: [], + contextMessages: [], + llmMessages: [], + pendingMessages: [], + canContinue: false, + }; + }); + + expect(secondRan).toBe(true); + expect(errors).toContain("Invalid boundary entries: Entry missing not found"); + expect(result.entries).toEqual([]); + expect(result.valid).toBe(true); + }); + + it("keeps shared mutations made before a handler throws", async () => { + const runtime = createExtensionRuntime(); + const extension = await loadExtensionFromFactory( + (pi) => { + pi.on("agent_before_settle", (event) => { + event.entries.push({ type: "custom", customType: "kept" }); + throw new Error("boundary failed"); + }); + }, + tempDir, + createEventBus(), + runtime, + ); + const runner = new ExtensionRunner([extension], runtime, tempDir, sessionManager, modelRegistry); + const errors: string[] = []; + runner.onError((error) => errors.push(error.error)); + + const result = await runner.emitBoundary({ type: "agent_before_settle", outcome: "completed" }, () => ({ + contextEntries: [], + contextMessages: [], + llmMessages: [], + pendingMessages: [], + canContinue: false, + })); + + expect(result.entries).toMatchObject([{ type: "custom", customType: "kept" }]); + expect(errors).toEqual(["boundary failed"]); + }); + }); + describe("tool_result chaining", () => { it("chains content modifications across handlers", async () => { const extCode1 = ` diff --git a/packages/coding-agent/test/session-context-edit.test.ts b/packages/coding-agent/test/session-context-edit.test.ts new file mode 100644 index 000000000..28ed37b86 --- /dev/null +++ b/packages/coding-agent/test/session-context-edit.test.ts @@ -0,0 +1,364 @@ +import type { AssistantMessage, ToolResultMessage } from "@earendil-works/pi-ai"; +import { describe, expect, it } from "vitest"; +import { + DEFAULT_COMPACTION_SETTINGS, + estimateProjectedContextTokens, + prepareCompaction, +} from "../src/core/compaction/index.ts"; +import { SessionManager } from "../src/core/session-manager.ts"; + +function assistant(text: string): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "faux", + provider: "faux", + model: "faux", + usage: { + input: 10, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 11, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +function text(message: { content: string | Array<{ type: string; text?: string }> }): string { + return typeof message.content === "string" + ? message.content + : message.content.flatMap((part) => (part.type === "text" ? [part.text ?? ""] : [])).join(""); +} + +describe("session context edits", () => { + it("omits a target only from model projection", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "request", timestamp: Date.now() }); + const assistantId = session.appendMessage(assistant("partial")); + const result: ToolResultMessage = { + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "raw output" }], + details: { path: "large.txt" }, + isError: true, + timestamp: Date.now(), + }; + const resultId = session.appendMessage(result); + session.appendContextEdit(assistantId, null); + session.appendContextEdit(resultId, null); + + expect(session.getBranch().filter((entry) => entry.type === "message")).toHaveLength(3); + expect(session.buildSessionProjection().messages.map((message) => message.role)).toEqual(["user"]); + expect((session.getEntry(resultId) as { message: ToolResultMessage }).message).toBe(result); + }); + + it("replaces only content and lets the latest edit win", () => { + const session = SessionManager.inMemory(); + const targetId = session.appendMessage(assistant("original")); + session.appendContextEdit(targetId, { content: [{ type: "text", text: "first" }] }); + session.appendContextEdit(targetId, null); + session.appendContextEdit(targetId, { content: [{ type: "text", text: "restored" }] }); + + const projected = session.buildSessionProjection().messages[0]; + expect(projected.role).toBe("assistant"); + if (projected.role !== "assistant") throw new Error("expected assistant"); + expect(text(projected)).toBe("restored"); + expect(projected.usage.totalTokens).toBe(11); + expect(text((session.getEntry(targetId) as { message: AssistantMessage }).message)).toBe("original"); + }); + + it("normalizes string replacements for array-only assistant and tool-result roles", () => { + const session = SessionManager.inMemory(); + const assistantId = session.appendMessage(assistant("original")); + const resultId = session.appendMessage({ + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "original result" }], + isError: false, + timestamp: Date.now(), + }); + const assistantEditId = session.appendContextEdit(assistantId, { content: "assistant replacement" }); + const resultEditId = session.appendContextEdit(resultId, { content: "result replacement" }); + + expect(session.getEntry(assistantEditId)).toMatchObject({ + replacement: { content: [{ type: "text", text: "assistant replacement" }] }, + }); + expect(session.getEntry(resultEditId)).toMatchObject({ + replacement: { content: [{ type: "text", text: "result replacement" }] }, + }); + const projected = session.buildSessionProjection().messages; + expect(projected[0]).toMatchObject({ + role: "assistant", + content: [{ type: "text", text: "assistant replacement" }], + }); + expect(projected[1]).toMatchObject({ + role: "toolResult", + content: [{ type: "text", text: "result replacement" }], + }); + }); + + it("normalizes imported string replacements while projecting array-only roles", () => { + const session = SessionManager.inMemory(); + const assistantId = session.appendMessage(assistant("original")); + const editId = session.appendContextEdit(assistantId, null); + const edit = session.getEntry(editId); + if (edit?.type !== "context_edit") throw new Error("expected context edit"); + edit.replacement = { content: "imported replacement" }; + + expect(session.buildSessionProjection().messages[0]).toMatchObject({ + role: "assistant", + content: [{ type: "text", text: "imported replacement" }], + }); + }); + + it("keeps edits branch-relative", () => { + const session = SessionManager.inMemory(); + const targetId = session.appendMessage({ role: "user", content: "original", timestamp: Date.now() }); + session.appendContextEdit(targetId, { content: "edited" }); + expect(text(session.buildSessionProjection().messages[0] as { content: string })).toBe("edited"); + + session.branch(targetId); + expect(text(session.buildSessionProjection().messages[0] as { content: string })).toBe("original"); + }); + + it("uses a self-referencing compaction to retain no preceding entries", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "discarded", timestamp: Date.now() }); + const compactionId = session.appendCompaction("exact handoff", null, 100); + session.appendMessage({ role: "user", content: "after", timestamp: Date.now() }); + + const compaction = session.getEntry(compactionId); + expect(compaction).toMatchObject({ type: "compaction", firstKeptEntryId: compactionId }); + expect(session.buildSessionProjection().messages.map((message) => message.role)).toEqual([ + "compactionSummary", + "user", + ]); + expect( + session + .buildSessionProjection() + .messages.map((message) => ("summary" in message ? message.summary : text(message as { content: string }))), + ).toEqual(["exact handoff", "after"]); + }); + + it("applies post-compaction edits to retained pre-compaction entries", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "summarized", timestamp: Date.now() }); + const retainedId = session.appendMessage({ role: "user", content: "original retained", timestamp: Date.now() }); + session.appendCompaction("summary", retainedId, 100); + session.appendContextEdit(retainedId, { content: "edited retained" }); + + expect( + session + .buildSessionProjection() + .messages.map((message) => ("summary" in message ? message.summary : text(message as { content: string }))), + ).toEqual(["summary", "edited retained"]); + }); + + it("uses only the newest summary when a repeated compaction retains entries before the older compaction", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "summarized first", timestamp: Date.now() }); + const retainedId = session.appendMessage({ role: "user", content: "retained", timestamp: Date.now() }); + session.appendCompaction("first summary", retainedId, 100); + session.appendMessage(assistant("after first compaction")); + session.appendCompaction("second summary", retainedId, 80); + session.appendMessage({ role: "user", content: "new tail ".repeat(100), timestamp: Date.now() }); + + const summaries = session + .buildSessionProjection() + .messages.flatMap((message) => ("summary" in message ? [message.summary] : [])); + expect(summaries).toEqual(["second summary"]); + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation?.previousSummary).toBe("second summary"); + }); + + it("supports repeated retain-none compactions", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "discarded", timestamp: Date.now() }); + session.appendCompaction("first handoff", null, 100); + session.appendMessage({ role: "user", content: "also discarded", timestamp: Date.now() }); + const secondId = session.appendCompaction("second handoff", null, 50); + + expect(session.getEntry(secondId)).toMatchObject({ firstKeptEntryId: secondId }); + expect( + session.buildSessionProjection().messages.map((message) => ("summary" in message ? message.summary : "")), + ).toEqual(["second handoff"]); + }); + + it("does not trust pre-edit assistant usage for projected context estimates", () => { + const session = SessionManager.inMemory(); + const largeUserId = session.appendMessage({ + role: "user", + content: "discarded input ".repeat(2_000), + timestamp: Date.now(), + }); + const response = assistant("small answer"); + response.usage = { ...response.usage, input: 10_000, totalTokens: 10_001 }; + const assistantId = session.appendMessage(response); + session.appendContextEdit(largeUserId, null); + + const editedEstimate = estimateProjectedContextTokens(session.buildSessionProjection(), session.getBranch()); + expect(editedEstimate.usageTokens).toBe(0); + expect(editedEstimate.tokens).toBeLessThan(100); + + session.appendContextEdit(assistantId, null); + expect(estimateProjectedContextTokens(session.buildSessionProjection(), session.getBranch()).tokens).toBe(0); + }); + + it("uses assistant usage captured after the latest context edit", () => { + const session = SessionManager.inMemory(); + const userId = session.appendMessage({ role: "user", content: "original", timestamp: Date.now() }); + session.appendContextEdit(userId, { content: "edited" }); + const response = assistant("answer"); + response.usage = { ...response.usage, input: 4_000, output: 100, totalTokens: 4_100 }; + session.appendMessage(response); + session.appendMessage({ role: "user", content: "next", timestamp: Date.now() }); + + const estimate = estimateProjectedContextTokens(session.buildSessionProjection(), session.getBranch()); + expect(estimate.usageTokens).toBe(4_100); + expect(estimate.trailingTokens).toBe(1); + expect(estimate.tokens).toBe(4_101); + }); + + it("does not reuse post-edit assistant usage after a later compaction", () => { + const session = SessionManager.inMemory(); + const userId = session.appendMessage({ role: "user", content: "small input", timestamp: Date.now() }); + session.appendContextEdit(userId, { content: "edited input" }); + const response = assistant("answer"); + response.usage = { ...response.usage, input: 50_000, output: 1, totalTokens: 50_001 }; + session.appendMessage(response); + session.appendCompaction("small summary", userId, 50_001); + + const estimate = estimateProjectedContextTokens(session.buildSessionProjection(), session.getBranch()); + expect(estimate.usageTokens).toBe(0); + expect(estimate.tokens).toBeLessThan(100); + }); + + it("includes effective system and tool context in edited estimates", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ + role: "system", + content: "system prompt ".repeat(3_000), + toolsAdded: [ + { + name: "example", + description: "tool declaration ".repeat(100), + parameters: { type: "object", properties: {} }, + }, + ], + timestamp: Date.now(), + }); + const userId = session.appendMessage({ role: "user", content: "ask", timestamp: Date.now() }); + session.appendMessage(assistant("done")); + session.appendContextEdit(userId, { content: "ask" }); + + expect( + estimateProjectedContextTokens(session.buildSessionProjection(), session.getBranch()).tokens, + ).toBeGreaterThan(10_000); + }); + + it("does not advance past a boundary replacement of the candidate input", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "old request", timestamp: Date.now() }); + session.appendMessage(assistant("old answer")); + const replacedUserId = session.appendMessage({ + role: "user", + content: "original input", + timestamp: Date.now(), + }); + const assistantId = session.appendMessage(assistant("answered original input")); + session.appendContextEdit(replacedUserId, { content: "NEW-INSTRUCTION ".repeat(100) }); + session.appendContextEdit(assistantId, null); + session.appendCustomEntry("bookkeeping", { source: "test" }); + + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation?.firstKeptEntryId).toBe(replacedUserId); + expect(JSON.stringify(preparation?.messagesToSummarize)).not.toContain("NEW-INSTRUCTION"); + expect(JSON.stringify(preparation?.turnPrefixMessages)).not.toContain("NEW-INSTRUCTION"); + }); + + it("does not let metadata move the cut past unsent boundary input", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ role: "user", content: "old request", timestamp: Date.now() }); + session.appendMessage(assistant("old answer")); + const instructionId = session.appendCustomMessageEntry("next-work", "UNSENT-INSTRUCTION ".repeat(100), false); + session.appendCustomEntry("bookkeeping", { source: "test" }); + + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation?.firstKeptEntryId).toBe(instructionId); + expect(JSON.stringify(preparation?.messagesToSummarize)).not.toContain("UNSENT-INSTRUCTION"); + expect(JSON.stringify(preparation?.turnPrefixMessages)).not.toContain("UNSENT-INSTRUCTION"); + }); + + it("does not treat an omitted custom message as a recovery attempt", () => { + const session = SessionManager.inMemory(); + session.appendMessage({ + role: "user", + content: "unanswered input ".repeat(100), + timestamp: Date.now(), + }); + const customId = session.appendCustomMessageEntry("temporary", "temporary context", false); + session.appendContextEdit(customId, null); + + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation).toBeUndefined(); + }); + + it("advances past input for an omitted assistant recovery suffix with metadata", () => { + const session = SessionManager.inMemory(); + const userId = session.appendMessage({ + role: "user", + content: "recovery input ".repeat(100), + timestamp: Date.now(), + }); + const attemptId = session.appendMessage(assistant("failed attempt")); + session.appendContextEdit(attemptId, null); + session.appendCustomEntry("bookkeeping", { source: "test" }); + + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation?.firstKeptEntryId).toBe(attemptId); + expect(preparation?.turnPrefixMessages).toEqual([ + expect.objectContaining({ role: "user", content: expect.stringContaining("recovery input") }), + ]); + expect(preparation?.messagesToSummarize).not.toContainEqual( + expect.objectContaining({ role: "user", content: expect.stringContaining("recovery input") }), + ); + expect(userId).not.toBe(attemptId); + }); + + it("prepares compaction from edited model content", () => { + const session = SessionManager.inMemory(); + const omittedId = session.appendMessage({ role: "user", content: "OMIT-ME ".repeat(100), timestamp: Date.now() }); + session.appendMessage(assistant("old answer ".repeat(100))); + session.appendContextEdit(omittedId, null); + session.appendMessage({ role: "user", content: "keep", timestamp: Date.now() }); + session.appendMessage(assistant("suffix")); + + const preparation = prepareCompaction(session.getBranch(), { + ...DEFAULT_COMPACTION_SETTINGS, + keepRecentTokens: 1, + }); + expect(preparation).toBeDefined(); + expect(JSON.stringify(preparation?.messagesToSummarize)).not.toContain("OMIT-ME"); + expect(JSON.stringify(preparation?.turnPrefixMessages)).not.toContain("OMIT-ME"); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-boundaries.test.ts b/packages/coding-agent/test/suite/agent-session-boundaries.test.ts new file mode 100644 index 000000000..5115efb74 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-boundaries.test.ts @@ -0,0 +1,1062 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createHarness, getMessageText, type Harness } from "./harness.ts"; + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve = () => {}; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describe("AgentSession actionable boundaries", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) harnesses.pop()?.cleanup(); + }); + + it("commits a retain-none turn_end compaction and explicitly continues once", async () => { + let handled = false; + const observedIds: string[] = []; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event) => { + observedIds.push(event.messageEntryId); + if (handled) return; + handled = true; + return { + entries: [ + { + type: "compaction", + summary: "exact handoff", + firstKeptEntryId: null, + details: { source: "test" }, + }, + ], + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("discarded response"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("continued from handoff"); + }, + ]); + + await harness.session.prompt("discarded prompt"); + + const compaction = harness.sessionManager.getEntries().find((entry) => entry.type === "compaction"); + expect(compaction).toMatchObject({ type: "compaction", summary: "exact handoff" }); + if (compaction?.type !== "compaction") throw new Error("expected compaction"); + expect(compaction.firstKeptEntryId).toBe(compaction.id); + expect(requests).toHaveLength(1); + expect(requests[0]).toContain("exact handoff"); + expect(requests[0]).not.toContain("discarded prompt"); + expect(requests[0]).not.toContain("discarded response"); + expect(observedIds).toHaveLength(2); + expect(harness.eventsOfType("agent_settled")).toHaveLength(1); + }); + + it.each(["steering", "follow-up", "both"] as const)( + "preserves %s queue scheduling around a turn_end handoff", + async (queueKind) => { + let handled = false; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("turn_end", () => { + if (handled) return; + handled = true; + if (queueKind === "steering" || queueKind === "both") { + pi.sendUserMessage("queued steering", { deliverAs: "steer" }); + } + if (queueKind === "follow-up" || queueKind === "both") { + pi.sendUserMessage("queued follow-up", { deliverAs: "followUp" }); + } + return { + entries: [{ type: "compaction", summary: "exact handoff", firstKeptEntryId: null }], + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("third"); + }, + ]); + + await harness.session.prompt("start"); + + expect(requests[0]).toContain("exact handoff"); + if (queueKind === "steering") { + expect(harness.faux.state.callCount).toBe(2); + expect(requests[0]).toContain("queued steering"); + expect(requests[0]).not.toContain("queued follow-up"); + } else if (queueKind === "follow-up") { + expect(harness.faux.state.callCount).toBe(2); + expect(requests[0]).toContain("queued follow-up"); + } else { + expect(harness.faux.state.callCount).toBe(3); + expect(requests[0]).toContain("queued steering"); + expect(requests[0]).not.toContain("queued follow-up"); + expect(requests[1]).toContain("queued follow-up"); + } + }, + ); + + it("keeps a boundary replacement verbatim through threshold compaction", async () => { + let handled = false; + const requests: string[] = []; + const instruction = "EXACT-REPLACEMENT-INSTRUCTION ".repeat(100); + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 2_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event, ctx) => { + if (handled) return; + handled = true; + const user = [...ctx.sessionManager.getBranch()] + .reverse() + .find((entry) => entry.type === "message" && entry.message.role === "user"); + if (!user) throw new Error("missing user entry"); + return { + entries: [ + { type: "context_edit", targetId: user.id, replacement: { content: instruction } }, + { type: "context_edit", targetId: event.messageEntryId, replacement: null }, + { type: "custom", customType: "bookkeeping", data: { source: "test" } }, + ], + continue: true, + }; + }); + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "older history summary", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.sessionManager.appendMessage({ role: "user", content: "older input", timestamp: Date.now() - 2 }); + harness.sessionManager.appendMessage(fauxAssistantMessage("older answer", { timestamp: Date.now() - 1 })); + harness.session.refreshContext(); + harness.setResponses([ + fauxAssistantMessage("answered original input"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("answered replacement"); + }, + ]); + + await harness.session.prompt("original input"); + + expect(harness.eventsOfType("compaction_start").length).toBeGreaterThan(0); + expect(requests).toHaveLength(1); + expect(requests[0]).toContain("EXACT-REPLACEMENT-INSTRUCTION"); + }); + + it("keeps boundary input verbatim through threshold compaction when metadata follows it", async () => { + let handled = false; + const requests: string[] = []; + const instruction = "EXACT-UNSENT-INSTRUCTION ".repeat(100); + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 2_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("turn_end", () => { + if (handled) return; + handled = true; + return { + entries: [ + { + type: "custom_message", + customType: "next-work", + content: instruction, + display: false, + }, + { type: "custom", customType: "bookkeeping", data: { source: "test" } }, + ], + continue: true, + }; + }); + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "older history summary", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + ]); + + await harness.session.prompt("old input ".repeat(500)); + + expect(harness.eventsOfType("compaction_start").length).toBeGreaterThan(0); + expect(requests).toHaveLength(1); + expect(requests[0]).toContain("EXACT-UNSENT-INSTRUCTION"); + }); + + it("refreshes canonical context before publishing boundary entry notifications", async () => { + let handled = false; + const snapshots: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("turn_end", () => { + if (handled) return; + handled = true; + return { + entries: [ + { type: "custom", customType: "metadata", data: true }, + { + type: "custom_message", + customType: "visible-context", + content: "committed context", + display: true, + }, + ], + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.session.subscribe((event) => { + if (event.type === "entry_appended") snapshots.push(JSON.stringify(harness.session.messages)); + }); + harness.setResponses([fauxAssistantMessage("done")]); + + await harness.session.prompt("start"); + + expect(snapshots).toHaveLength(2); + expect(snapshots.every((snapshot) => snapshot.includes("committed context"))).toBe(true); + }); + + it("continues from an agent_before_settle custom message before final settlement", async () => { + let requested = false; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_before_settle", () => { + if (requested) return; + requested = true; + return { + entries: [ + { + type: "custom_message", + customType: "test-continuation", + content: "continue now", + display: false, + }, + ], + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + ]); + + await harness.session.prompt("start"); + + expect(requests[0]).toContain("continue now"); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ type: "custom_message", customType: "test-continuation", display: false }), + ); + expect(harness.eventsOfType("agent_start")).toHaveLength(2); + expect(harness.eventsOfType("agent_settled")).toHaveLength(1); + }); + + it("persists custom context queued by agent_end before pre-settlement continuation", async () => { + let firstRun = true; + let continued = false; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_end", () => { + if (!firstRun) return; + firstRun = false; + pi.sendMessage( + { customType: "agent-end-context", content: "queued after agent end", display: false }, + { triggerTurn: false }, + ); + }); + pi.on("agent_before_settle", (event) => { + if (continued) return; + continued = true; + expect(JSON.stringify(event.context.pendingMessages)).toContain("queued after agent end"); + expect(JSON.stringify(event.context.contextMessages)).not.toContain("queued after agent end"); + return { continue: true }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + ]); + + await harness.session.prompt("start"); + + expect(requests[0]).toContain("queued after agent end"); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ type: "custom_message", customType: "agent-end-context" }), + ); + }); + + it("keeps a pre-settlement follow-up deferred until the explicit continuation would stop", async () => { + let handled = false; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_before_settle", () => { + if (handled) return; + handled = true; + pi.sendUserMessage("queued follow-up", { deliverAs: "followUp" }); + return { + entries: [ + { + type: "custom_message", + customType: "boundary", + content: "boundary context", + display: false, + }, + ], + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("follow-up response"); + }, + ]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(3); + expect(requests[0]).toContain("boundary context"); + expect(requests[0]).not.toContain("queued follow-up"); + expect(requests[1]).toContain("queued follow-up"); + }); + + it("defers runs started by agent_settled handlers until every settled handler completes", async () => { + let triggered = false; + const lifecycle: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_start", () => { + lifecycle.push("start"); + }); + pi.on("agent_settled", (_event, ctx) => { + lifecycle.push(`settled-first:${ctx.isIdle()}`); + if (triggered) return; + triggered = true; + pi.sendMessage( + { customType: "settled-trigger", content: "start later", display: false }, + { triggerTurn: true }, + ); + }); + pi.on("agent_settled", (_event, ctx) => { + lifecycle.push(`settled-second:${ctx.isIdle()}`); + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]); + + await harness.session.prompt("start"); + + expect(lifecycle).toEqual([ + "start", + "settled-first:true", + "settled-second:true", + "start", + "settled-first:true", + "settled-second:true", + ]); + }); + + it("does not let an invalid explicit continuation suppress natural tool continuation", async () => { + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "done" }], details: {} }), + }; + const harness = await createHarness({ + tools: [tool], + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event, ctx) => { + const user = [...ctx.sessionManager.getBranch()] + .reverse() + .find((entry) => entry.type === "message" && entry.message.role === "user"); + if (!user) throw new Error("missing user entry"); + return { + entries: [user.id, event.messageEntryId, ...event.toolResultEntryIds].map((targetId) => ({ + type: "context_edit" as const, + targetId, + replacement: null, + })), + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("noop", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("must not run"), + ]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(2); + }); + + it("dispatches actionable turn_end for synthetic run failures", async () => { + let turnEnds = 0; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event) => { + turnEnds++; + expect(event.outcome).toBe("error"); + return { entries: [{ type: "custom", customType: "failure-boundary", data: true }] }; + }); + }, + ], + }); + harnesses.push(harness); + harness.session.agent.prepareRequest = () => { + throw new Error("request preparation failed"); + }; + + await harness.session.prompt("start"); + + expect(turnEnds).toBe(1); + expect(harness.faux.state.callCount).toBe(0); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ type: "custom", customType: "failure-boundary" }), + ); + }); + + it("does not compact from usage belonging to a boundary-omitted assistant", async () => { + let handled = false; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 10_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 300 } }, + extensionFactories: [ + (pi) => { + pi.on("message_end", (event) => { + if (event.message.role !== "assistant") return; + return { + message: { + ...event.message, + usage: { ...event.message.usage, input: 9_800, output: 1, totalTokens: 9_801 }, + }, + }; + }); + pi.on("turn_end", (event) => { + if (handled) return; + handled = true; + return { + entries: [{ type: "context_edit", targetId: event.messageEntryId, replacement: null }], + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("short response")]); + + await harness.session.prompt("small prompt"); + + expect(harness.eventsOfType("compaction_start")).toEqual([]); + expect(harness.session.getContextUsage()?.tokens).toBeLessThan(2_000); + }); + + it("does not trigger successful-response overflow from usage captured before a boundary edit", async () => { + let handled = false; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 5_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("message_end", (event) => { + if (event.message.role !== "assistant") return; + return { + message: { + ...event.message, + usage: { ...event.message.usage, input: 5_100, output: 1, totalTokens: 5_101 }, + }, + }; + }); + pi.on("turn_end", (_event, ctx) => { + if (handled) return; + handled = true; + const user = [...ctx.sessionManager.getBranch()] + .reverse() + .find((entry) => entry.type === "message" && entry.message.role === "user"); + if (!user) throw new Error("missing user entry"); + return { entries: [{ type: "context_edit", targetId: user.id, replacement: null }] }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("done")]); + + await harness.session.prompt("large input that is later omitted"); + + expect(harness.eventsOfType("compaction_start")).toEqual([]); + expect(harness.session.getContextUsage()?.tokens).toBeLessThan(2_000); + }); + + it("does not trigger threshold compaction from post-edit usage captured before a later compaction", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 10_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + }); + harnesses.push(harness); + const userId = harness.sessionManager.appendMessage({ + role: "user", + content: "small input", + timestamp: Date.now() - 3, + }); + harness.sessionManager.appendContextEdit(userId, { content: "edited input" }); + const response = fauxAssistantMessage("answer", { timestamp: Date.now() - 2 }); + response.usage = { ...response.usage, input: 50_000, output: 1, totalTokens: 50_001 }; + harness.sessionManager.appendMessage(response); + harness.sessionManager.appendCompaction("small summary", userId, 50_001); + harness.session.refreshContext(); + const runAutoCompaction = vi.spyOn( + harness.session as unknown as { + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; + }, + "_runAutoCompaction", + ); + const checkCompaction = ( + harness.session as unknown as { + _checkCompaction: (message: ReturnType) => Promise; + } + )._checkCompaction.bind(harness.session); + const error = fauxAssistantMessage("", { + stopReason: "error", + errorMessage: "invalid_api_key", + timestamp: Date.now() + 1_000, + }); + + await checkCompaction(error); + + expect(runAutoCompaction).not.toHaveBeenCalled(); + }); + + it("does not treat retained pre-compaction assistant usage as post-compaction usage", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const retained = fauxAssistantMessage("retained"); + retained.usage = { ...retained.usage, input: 10_000, totalTokens: 10_001 }; + const retainedId = harness.sessionManager.appendMessage(retained); + harness.sessionManager.appendCompaction("summary", retainedId, 10_001); + harness.session.refreshContext(); + + expect(harness.session.getContextUsage()?.tokens).toBeNull(); + }); + + it("persists custom context sent during pre-settlement before continuing", async () => { + let handled = false; + const requests: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_before_settle", () => { + if (handled) return; + handled = true; + pi.sendMessage( + { customType: "pending-boundary", content: "persist before continue", display: false }, + { triggerTurn: false }, + ); + return { continue: true }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first"), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("second"); + }, + ]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(2); + expect(requests[0]).toContain("persist before continue"); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ type: "custom_message", customType: "pending-boundary" }), + ); + }); + + it("does not consume queued input when pre-settlement drafts leave system-only context", async () => { + let handled = false; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_before_settle", (_event, ctx) => { + if (handled) return; + handled = true; + pi.sendUserMessage("still queued", { deliverAs: "followUp" }); + const targets = ctx.sessionManager + .getBranch() + .flatMap((entry) => + entry.type === "message" && + (entry.message.role === "user" || + entry.message.role === "assistant" || + entry.message.role === "toolResult") + ? [entry.id] + : [], + ); + return { + entries: targets.map((targetId) => ({ + type: "context_edit" as const, + targetId, + replacement: null, + })), + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("must not run")]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(1); + expect(harness.session.pendingMessageCount).toBe(1); + }); + + it("commits pre-settlement drafts but suppresses continuation when aborted during the hook", async () => { + const started = deferred(); + const release = deferred(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("agent_before_settle", async () => { + started.resolve(); + await release.promise; + return { + entries: [{ type: "custom", customType: "committed-after-abort", data: true }], + continue: true, + }; + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("must not run")]); + + const prompt = harness.session.prompt("start"); + await started.promise; + const abort = harness.session.abort(); + release.resolve(); + await Promise.all([prompt, abort]); + + expect(harness.faux.state.callCount).toBe(1); + expect(harness.sessionManager.getEntries()).toContainEqual( + expect.objectContaining({ type: "custom", customType: "committed-after-abort", data: true }), + ); + expect(harness.eventsOfType("agent_settled")).toHaveLength(1); + }); +}); + +describe("durable length recovery", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) harnesses.pop()?.cleanup(); + }); + + it("keeps truncated tool attempts in context for the natural next turn", async () => { + let executed = false; + const requests: string[] = []; + const tool: AgentTool = { + name: "unsafe_truncated_tool", + label: "Unsafe truncated tool", + description: "Must not execute from a length response", + parameters: Type.Object({ value: Type.String() }), + execute: async () => { + executed = true; + return { content: [{ type: "text", text: "executed" }], details: {} }; + }, + }; + const harness = await createHarness({ tools: [tool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("unsafe_truncated_tool", { value: "partial" }), { + stopReason: "length", + }), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("completed natural continuation"); + }, + ]); + + await harness.session.prompt("start"); + + expect(executed).toBe(false); + expect(harness.faux.state.callCount).toBe(2); + expect(requests[0]).toContain("may be truncated"); + expect(harness.sessionManager.getEntries().some((entry) => entry.type === "context_edit")).toBe(false); + }); + + it("resets length recovery after a successful intermediate assistant turn", async () => { + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "done" }], details: {} }), + }; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 1000, maxTokens: 100 }], + settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } }, + tools: [tool], + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "recovered input", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first partial", { stopReason: "length" }), + fauxAssistantMessage(fauxToolCall("noop", {}), { stopReason: "toolUse" }), + () => fauxAssistantMessage("second partial", { stopReason: "length", timestamp: Date.now() + 1_000 }), + () => fauxAssistantMessage("completed second recovery", { timestamp: Date.now() + 2_000 }), + ]); + + await harness.session.prompt("x".repeat(5000)); + + const omittedIds = harness.sessionManager + .getEntries() + .filter((entry) => entry.type === "context_edit") + .map((entry) => entry.targetId); + const lengthResponses = harness.sessionManager + .getEntries() + .filter( + (entry) => + entry.type === "message" && entry.message.role === "assistant" && entry.message.stopReason === "length", + ); + expect(lengthResponses).toHaveLength(2); + expect(omittedIds).toEqual(expect.arrayContaining(lengthResponses.map((entry) => entry.id))); + expect(harness.faux.state.callCount).toBe(3); + }); + + it("gives a distinct queued follow-up its own length-recovery budget", async () => { + let queued = false; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 1000, maxTokens: 100 }], + settings: { compaction: { keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "recovered input", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + })); + pi.on("agent_end", (event) => { + if (queued || !event.messages.some((message) => getMessageText(message) === "first recovered")) + return; + queued = true; + pi.sendUserMessage("distinct follow-up", { deliverAs: "followUp" }); + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("first partial", { stopReason: "length" }), + fauxAssistantMessage("first recovered"), + () => fauxAssistantMessage("follow-up partial", { stopReason: "length", timestamp: Date.now() + 1_000 }), + () => fauxAssistantMessage("follow-up recovered", { timestamp: Date.now() + 2_000 }), + ]); + + await harness.session.prompt("x".repeat(5000)); + + const lengthIds = harness.sessionManager + .getEntries() + .flatMap((entry) => + entry.type === "message" && entry.message.role === "assistant" && entry.message.stopReason === "length" + ? [entry.id] + : [], + ); + const omitted = harness.sessionManager + .getEntries() + .flatMap((entry) => (entry.type === "context_edit" ? [entry.targetId] : [])); + expect(lengthIds).toHaveLength(2); + expect(omitted).toEqual(expect.arrayContaining(lengthIds)); + expect(harness.faux.state.callCount).toBe(4); + }); + + it("finishes retry bookkeeping when a retry receives a nonretryable error", async () => { + const harness = await createHarness({ + settings: { retry: { enabled: true, maxRetries: 2, baseDelayMs: 1 } }, + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "invalid_api_key" }), + ]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(2); + expect(harness.eventsOfType("auto_retry_end")).toContainEqual( + expect.objectContaining({ success: false, attempt: 1, finalError: "invalid_api_key" }), + ); + }); + + it("omits a recoverable projected replacement by its source entry ID", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 1_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + harness.sessionManager.appendMessage({ role: "user", content: "x".repeat(5_000), timestamp: Date.now() - 2 }); + const partial = fauxAssistantMessage("original partial", { + stopReason: "length", + timestamp: Date.now() - 1, + }); + const partialId = harness.sessionManager.appendMessage(partial); + harness.sessionManager.appendContextEdit(partialId, { + content: [{ type: "text", text: "edited partial" }], + }); + harness.session.refreshContext(); + harness.setResponses([fauxAssistantMessage("new answer")]); + + await harness.session.prompt("next prompt"); + + const edits = harness.sessionManager.getEntries().filter((entry) => entry.type === "context_edit"); + expect(edits.filter((entry) => entry.targetId === partialId).at(-1)?.replacement).toBeNull(); + expect( + harness.sessionManager + .buildSessionProjection() + .messages.some((message) => getMessageText(message) === "edited partial"), + ).toBe(false); + }); + + it("recovers an explicit overflow error after a retained boundary replacement", async () => { + let replaced = false; + let overflowId: string | undefined; + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 1_000, maxTokens: 100 }], + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 0 } }, + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event) => { + if (replaced || event.outcome !== "error") return; + replaced = true; + overflowId = event.messageEntryId; + return { + entries: [ + { + type: "context_edit", + targetId: event.messageEntryId, + replacement: { content: [{ type: "text", text: "retained error" }] }, + }, + ], + }; + }); + pi.on("session_before_compact", (event) => ({ + compaction: { + summary: "recovered overflow", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("retained error", { stopReason: "error", errorMessage: "prompt is too long" }), + fauxAssistantMessage("recovered"), + ]); + + await harness.session.prompt("x".repeat(5_000)); + + expect(harness.faux.state.callCount).toBe(2); + expect(overflowId).toBeDefined(); + const edits = harness.sessionManager.getEntries().filter((entry) => entry.type === "context_edit"); + expect(edits.filter((entry) => entry.targetId === overflowId).at(-1)?.replacement).toBeNull(); + }); + + it("keeps follow-up work behind an automatic error retry", async () => { + let queued = false; + const requests: string[] = []; + const lifecycle: string[] = []; + const harness = await createHarness({ + settings: { retry: { enabled: true, maxRetries: 2, baseDelayMs: 1 } }, + extensionFactories: [ + (pi) => { + pi.on("turn_end", (event) => { + if (queued || event.outcome !== "error") return; + queued = true; + pi.sendUserMessage("queued follow-up", { deliverAs: "followUp" }); + }); + }, + ], + }); + harnesses.push(harness); + harness.session.subscribe((event) => { + if (event.type === "agent_end" || event.type === "auto_retry_start") lifecycle.push(event.type); + }); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("retry recovered"); + }, + (context) => { + requests.push(JSON.stringify(context.messages)); + return fauxAssistantMessage("follow-up completed"); + }, + ]); + + await harness.session.prompt("start"); + + expect(harness.faux.state.callCount).toBe(3); + expect(requests[0]).not.toContain("queued follow-up"); + expect(requests[1]).toContain("queued follow-up"); + expect(lifecycle.slice(0, 2)).toEqual(["agent_end", "auto_retry_start"]); + }); + + it("marks the exhausted retry run as final", async () => { + const harness = await createHarness({ + settings: { retry: { enabled: true, maxRetries: 1, baseDelayMs: 1 } }, + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + ]); + + await harness.session.prompt("start"); + + expect(harness.eventsOfType("agent_end").map((event) => event.willRetry)).toEqual([true, false]); + expect(harness.eventsOfType("auto_retry_end")).toContainEqual( + expect.objectContaining({ success: false, attempt: 1 }), + ); + }); + + it("keeps omissions and does not retry when recovery compaction fails", async () => { + const harness = await createHarness({ + models: [{ id: "faux-1", contextWindow: 1000, maxTokens: 100 }], + settings: { + compaction: { keepRecentTokens: 1, reserveTokens: 0 }, + retry: { enabled: false, maxRetries: 0, baseDelayMs: 1 }, + }, + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("partial response", { stopReason: "length" }), + fauxAssistantMessage("summary failed", { stopReason: "error", errorMessage: "summary failed" }), + fauxAssistantMessage("must not retry"), + ]); + + await harness.session.prompt("x".repeat(5000)); + + const entries = harness.sessionManager.getEntries(); + expect(entries.some((entry) => entry.type === "context_edit")).toBe(true); + expect(entries.some((entry) => entry.type === "compaction")).toBe(false); + expect( + entries.some((entry) => entry.type === "message" && getMessageText(entry.message) === "partial response"), + ).toBe(true); + expect( + harness.sessionManager + .buildSessionProjection() + .messages.some((message) => getMessageText(message) === "partial response"), + ).toBe(false); + expect(harness.faux.state.callCount).toBe(2); + }); +}); diff --git a/packages/coding-agent/test/tree-selector.test.ts b/packages/coding-agent/test/tree-selector.test.ts index 9dc750d5d..579511b33 100644 --- a/packages/coding-agent/test/tree-selector.test.ts +++ b/packages/coding-agent/test/tree-selector.test.ts @@ -156,6 +156,43 @@ describe("TreeSelectorComponent", () => { expect(list.getSelectedNode()?.entry.id).toBe("user-2"); }); + test("hides context edits by default and labels them in all mode", () => { + const entries: SessionEntry[] = [ + userMessage("user-1", null, "hello"), + assistantMessage("asst-1", "user-1", "hi"), + { + type: "context_edit", + id: "edit-1", + parentId: "asst-1", + timestamp: new Date().toISOString(), + targetId: "asst-1", + replacement: null, + }, + ]; + const tree = buildTree(entries); + const defaultSelector = new TreeSelectorComponent( + tree, + "edit-1", + 24, + () => {}, + () => {}, + ); + expect(defaultSelector.getTreeList().getSelectedNode()?.entry.id).toBe("asst-1"); + + const allSelector = new TreeSelectorComponent( + tree, + "edit-1", + 24, + () => {}, + () => {}, + undefined, + undefined, + "all", + ); + const rendered = allSelector.getTreeList().render(200).map(stripVTControlCharacters).join("\n"); + expect(rendered).toContain("[context omit: asst-1]"); + }); + test("focuses nearest visible ancestor when currentLeafId is a thinking_level_change entry", () => { // Similar structure with thinking_level_change instead of model_change const entries = [