diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 171d75917..2e07f9698 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -8,6 +8,8 @@ ### Fixed +- Skipped next-turn preparation when every tool result terminates the current batch and no queued message requires another provider turn, retained queued input when next-turn preparation fails, and stopped before provider continuation when next-turn preparation aborts. + ### Removed ## [2026.7.13] - 2026-07-13 diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 08501bbf8..da828de52 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -16,6 +16,7 @@ import type { AgentContext, AgentEvent, AgentLoopConfig, + AgentLoopTurnUpdate, AgentMessage, AgentTool, AgentToolCall, @@ -182,6 +183,17 @@ async function runLoop( let firstTurn = true; // Check for steering messages at start (user may have typed while waiting) let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || []; + let drainedTerminatingQueue: "steering" | "followUp" | undefined; + const refreshTerminatingQueueDrain = async (): Promise => { + if (!drainedTerminatingQueue || !config.restorePendingMessages) return; + await config.restorePendingMessages(drainedTerminatingQueue, pendingMessages); + pendingMessages = (await config.getSteeringMessages?.()) || []; + drainedTerminatingQueue = pendingMessages.length > 0 ? "steering" : undefined; + if (pendingMessages.length === 0) { + pendingMessages = (await config.getFollowUpMessages?.()) || []; + drainedTerminatingQueue = pendingMessages.length > 0 ? "followUp" : undefined; + } + }; // Outer loop: continues when queued follow-up messages arrive after agent would stop while (true) { @@ -194,6 +206,14 @@ async function runLoop( } else { firstTurn = false; } + if (drainedTerminatingQueue) { + await refreshTerminatingQueueDrain(); + if (pendingMessages.length === 0) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + drainedTerminatingQueue = undefined; + } // Process pending messages (inject before next assistant response) if (pendingMessages.length > 0) { @@ -221,6 +241,7 @@ async function runLoop( const toolResults: ToolResultMessage[] = []; hasMoreToolCalls = false; + let toolBatchTerminated = false; if (toolCalls.length > 0) { // A "length" stop means the output was cut off by the token limit, so // every tool call in the message may carry truncated arguments. Fail @@ -230,6 +251,7 @@ async function runLoop( ? await failToolCallsFromTruncatedMessage(toolCalls, emit) : await executeToolCalls(currentContext, message, config, signal, emit); toolResults.push(...executedToolBatch.messages); + toolBatchTerminated = executedToolBatch.terminate; hasMoreToolCalls = !executedToolBatch.terminate; for (const result of toolResults) { @@ -250,7 +272,36 @@ async function runLoop( context: currentContext, newMessages, }; - const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext); + if (toolBatchTerminated) { + if (await config.shouldStopAfterTurn?.(nextTurnContext)) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + + pendingMessages = (await config.getSteeringMessages?.()) || []; + if (pendingMessages.length > 0) { + drainedTerminatingQueue = "steering"; + } + if (pendingMessages.length === 0) { + pendingMessages = (await config.getFollowUpMessages?.()) || []; + if (pendingMessages.length > 0) { + drainedTerminatingQueue = "followUp"; + } + } + if (pendingMessages.length === 0) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + } + let nextTurnSnapshot: AgentLoopTurnUpdate | undefined; + try { + nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext); + } catch (error) { + if (drainedTerminatingQueue) { + await config.restorePendingMessages?.(drainedTerminatingQueue, pendingMessages); + } + throw error; + } if (nextTurnSnapshot) { currentContext = nextTurnSnapshot.context ?? currentContext; config = { @@ -264,20 +315,37 @@ async function runLoop( : nextTurnSnapshot.thinkingLevel, }; } + if (signal?.aborted) { + if (drainedTerminatingQueue) { + await config.restorePendingMessages?.(drainedTerminatingQueue, pendingMessages); + } + await emit({ type: "agent_end", messages: newMessages }); + return; + } + if (drainedTerminatingQueue) { + await refreshTerminatingQueueDrain(); + if (pendingMessages.length === 0) { + await emit({ type: "agent_end", messages: newMessages }); + return; + } + } if ( - await config.shouldStopAfterTurn?.({ + !toolBatchTerminated && + (await config.shouldStopAfterTurn?.({ message, toolResults, context: currentContext, newMessages, - }) + })) ) { await emit({ type: "agent_end", messages: newMessages }); return; } - pendingMessages = (await config.getSteeringMessages?.()) || []; + if (!toolBatchTerminated) { + pendingMessages = (await config.getSteeringMessages?.()) || []; + } } // Agent would stop here. Check for follow-up messages. diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 434882238..18043a1a7 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -123,6 +123,7 @@ export interface AgentOptions { class PendingMessageQueue { private messages: AgentMessage[] = []; + private clearGeneration = 0; public mode: QueueMode; constructor(mode: QueueMode) { @@ -137,6 +138,10 @@ class PendingMessageQueue { return this.messages.length > 0; } + getClearGeneration(): number { + return this.clearGeneration; + } + drain(): AgentMessage[] { if (this.mode === "all") { const drained = this.messages.slice(); @@ -152,8 +157,13 @@ class PendingMessageQueue { return [first]; } + prepend(messages: AgentMessage[]): void { + this.messages = [...messages, ...this.messages]; + } + clear(): void { this.messages = []; + this.clearGeneration++; } } @@ -434,6 +444,8 @@ export class Agent { private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig { let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; + let steeringQueueGeneration = this.steeringQueue.getClearGeneration(); + let followUpQueueGeneration = this.followUpQueue.getClearGeneration(); return { model: this._state.model, reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel, @@ -464,9 +476,22 @@ export class Agent { skipInitialSteeringPoll = false; return []; } + steeringQueueGeneration = this.steeringQueue.getClearGeneration(); return this.steeringQueue.drain(); }, - getFollowUpMessages: async () => this.followUpQueue.drain(), + getFollowUpMessages: async () => { + followUpQueueGeneration = this.followUpQueue.getClearGeneration(); + return this.followUpQueue.drain(); + }, + restorePendingMessages: (queue, messages) => { + if (queue === "steering") { + if (this.steeringQueue.getClearGeneration() !== steeringQueueGeneration) return; + this.steeringQueue.prepend(messages); + return; + } + if (this.followUpQueue.getClearGeneration() !== followUpQueueGeneration) return; + this.followUpQueue.prepend(messages); + }, }; } diff --git a/packages/agent/src/harness/agent-harness.ts b/packages/agent/src/harness/agent-harness.ts index cd58a8caa..9d2dda493 100644 --- a/packages/agent/src/harness/agent-harness.ts +++ b/packages/agent/src/harness/agent-harness.ts @@ -418,6 +418,8 @@ export class AgentHarness< setTurnState: (turnState: AgentHarnessTurnState) => void, ): AgentLoopConfig { const turnState = getTurnState(); + let steeringQueueAtDrain = this.steerQueue; + let followUpQueueAtDrain = this.followUpQueue; return { model: turnState.model, reasoning: turnState.thinkingLevel === "off" ? undefined : turnState.thinkingLevel, @@ -459,8 +461,27 @@ export class AgentHarness< thinkingLevel: nextTurnState.thinkingLevel, }; }, - getSteeringMessages: async () => this.drainQueuedMessages(this.steerQueue, this.steeringQueueMode), - getFollowUpMessages: async () => this.drainQueuedMessages(this.followUpQueue, this.followUpQueueMode), + getSteeringMessages: async () => { + steeringQueueAtDrain = this.steerQueue; + return this.drainQueuedMessages(steeringQueueAtDrain, this.steeringQueueMode); + }, + getFollowUpMessages: async () => { + followUpQueueAtDrain = this.followUpQueue; + return this.drainQueuedMessages(followUpQueueAtDrain, this.followUpQueueMode); + }, + restorePendingMessages: async (queue, messages) => { + const target = queue === "steering" ? this.steerQueue : this.followUpQueue; + const drainedQueue = queue === "steering" ? steeringQueueAtDrain : followUpQueueAtDrain; + if (target !== drainedQueue) return; + for (let index = messages.length - 1; index >= 0; index--) { + const message = messages[index]; + if (message?.role !== "user") { + throw new AgentHarnessError("invalid_state", "Only user messages can be restored to input queues"); + } + target.unshift(message); + } + await this.emitQueueUpdate(); + }, }; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 40027efa6..cdf55cf10 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -249,6 +249,12 @@ export interface AgentLoopConfig extends SimpleStreamOptions { */ getFollowUpMessages?: () => Promise; + /** + * Restores messages previously returned by a queue callback when next-turn preparation cannot continue. + * Restored messages must be placed before messages queued after the drain. + */ + restorePendingMessages?: (queue: "steering" | "followUp", messages: AgentMessage[]) => Promise | void; + /** * Tool execution mode. * - "sequential": execute tool calls one by one diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 8c1f8a327..77b34bfb0 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -1446,6 +1446,62 @@ describe("agentLoop with AgentMessage", () => { expect(convertedSecondTurnSystemPrompt).toBe("second prompt"); }); + it("should stop before provider call 2 when prepareNextTurn aborts", async () => { + // given + const toolSchema = Type.Object({ value: Type.String() }); + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + }; + }, + }; + const controller = new AbortController(); + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + prepareNextTurn: async () => { + controller.abort(); + return undefined; + }, + }; + let llmCalls = 0; + const stream = agentLoop([createUserMessage("echo something")], context, config, controller.signal, () => { + llmCalls++; + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + if (llmCalls === 1) { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + return; + } + const message = createAssistantMessage([{ type: "text", text: "should not run" }]); + mockStream.push({ type: "done", reason: "stop", message }); + }); + return mockStream; + }); + + // when + const { events } = await collectAgentEvents(stream); + + // then + expect(llmCalls).toBe(1); + expect(events.filter((event) => event.type === "agent_end")).toHaveLength(1); + }); + it("should stop after the current turn when shouldStopAfterTurn returns true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; @@ -1568,6 +1624,7 @@ describe("agentLoop with AgentMessage", () => { const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, + prepareNextTurn: vi.fn(async () => undefined), }; let llmCalls = 0; @@ -1591,10 +1648,81 @@ describe("agentLoop with AgentMessage", () => { const messages = await stream.result(); expect(llmCalls).toBe(1); + expect(config.prepareNextTurn).not.toHaveBeenCalled(); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); expect(events.filter((event) => event.type === "turn_end")).toHaveLength(1); }); + it("should continue a terminating tool batch when steering is queued", async () => { + // given + const toolSchema = Type.Object({ value: Type.String() }); + const tool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo tool", + parameters: toolSchema, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `echoed: ${params.value}` }], + details: { value: params.value }, + terminate: true, + }; + }, + }; + const context: AgentContext = { + systemPrompt: "", + messages: [], + tools: [tool], + }; + const queuedMessage = createUserMessage("continue after termination"); + let steeringPolls = 0; + let toolTurnPreparations = 0; + const config: AgentLoopConfig = { + model: createModel(), + convertToLlm: identityConverter, + getSteeringMessages: async () => { + steeringPolls++; + return steeringPolls === 2 ? [queuedMessage] : []; + }, + prepareNextTurn: async ({ toolResults }) => { + if (toolResults.length > 0) toolTurnPreparations++; + return undefined; + }, + }; + let llmCalls = 0; + let queuedMessageReachedCall2 = false; + const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, (_model, ctx) => { + llmCalls++; + if (llmCalls === 2) { + queuedMessageReachedCall2 = ctx.messages.includes(queuedMessage); + } + const mockStream = new MockAssistantStream(); + queueMicrotask(() => { + if (llmCalls === 1) { + const message = createAssistantMessage( + [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], + "toolUse", + ); + mockStream.push({ type: "done", reason: "toolUse", message }); + return; + } + const message = createAssistantMessage([{ type: "text", text: "done" }]); + mockStream.push({ type: "done", reason: "stop", message }); + }); + return mockStream; + }); + + // when + for await (const _event of stream) { + // consume + } + + // then + expect(llmCalls).toBe(2); + expect(toolTurnPreparations).toBe(1); + expect(queuedMessageReachedCall2).toBe(true); + }); + it("should continue after parallel tool calls when not all tool results terminate", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool = { diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 052e61c67..fac0f2be7 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -746,6 +746,477 @@ describe("Agent", () => { expect(sawAbortSignal).toBe(true); }); + it.each([ + "steering", + "followUp", + ] as const)("retains queued %s input when next-turn preparation fails after a terminating tool", async (queue) => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let requestCount = 0; + const agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => { + throw new Error("required preparation failed"); + }, + streamFn: () => { + requestCount++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + }); + return stream; + }, + }); + const queuedMessage: AgentMessage = { + role: "user", + content: [{ type: "text", text: "queued safety instruction" }], + timestamp: Date.now(), + }; + + // when + const prompt = agent.prompt("run the terminating tool"); + await toolStarted.promise; + if (queue === "steering") { + agent.steer(queuedMessage); + } else { + agent.followUp(queuedMessage); + } + releaseTool.resolve(); + await prompt; + + // then + expect(requestCount).toBe(1); + expect(agent.hasQueuedMessages()).toBe(true); + expect(agent.state.messages).not.toContain(queuedMessage); + expect(agent.state.errorMessage).toContain("required preparation failed"); + }); + + it.each([ + "steering", + "followUp", + ] as const)("retains queued %s input when next-turn preparation aborts after a terminating tool", async (queue) => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let providerCalls = 0; + let agent: Agent; + agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => { + agent.abort(); + return undefined; + }, + streamFn: () => { + providerCalls++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + }); + return stream; + }, + }); + const queuedMessage: AgentMessage = { + role: "user", + content: [{ type: "text", text: "queued safety instruction" }], + timestamp: Date.now(), + }; + + // when + const prompt = agent.prompt("run the terminating tool"); + await toolStarted.promise; + if (queue === "steering") { + agent.steer(queuedMessage); + } else { + agent.followUp(queuedMessage); + } + releaseTool.resolve(); + await prompt; + + // then + expect(providerCalls).toBe(1); + expect(agent.hasQueuedMessages()).toBe(true); + expect(agent.state.messages).not.toContain(queuedMessage); + expect(agent.state.errorMessage).toBeUndefined(); + }); + + it.each([ + "steering", + "followUp", + ] as const)("clears queued %s input when next-turn preparation clears it before aborting", async (queue) => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let providerCalls = 0; + let agent: Agent; + agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => { + if (queue === "steering") { + agent.clearSteeringQueue(); + } else { + agent.clearFollowUpQueue(); + } + agent.abort(); + return undefined; + }, + streamFn: () => { + providerCalls++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + }); + return stream; + }, + }); + const queuedMessage: AgentMessage = { + role: "user", + content: [{ type: "text", text: "queued safety instruction" }], + timestamp: Date.now(), + }; + + // when + const prompt = agent.prompt("run the terminating tool"); + await toolStarted.promise; + if (queue === "steering") { + agent.steer(queuedMessage); + } else { + agent.followUp(queuedMessage); + } + releaseTool.resolve(); + await prompt; + + // then + expect(providerCalls).toBe(1); + expect(agent.hasQueuedMessages()).toBe(false); + expect(agent.state.messages).not.toContain(queuedMessage); + expect(agent.state.errorMessage).toBeUndefined(); + }); + + it.each([ + ["steering", false], + ["steering", true], + ["followUp", false], + ["followUp", true], + ] as const)("does not deliver cleared %s input after successful next-turn preparation (replacement: %s)", async (queue, withReplacement) => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let providerCalls = 0; + const providerUserTexts: string[][] = []; + let preparedTerminatingTurn = false; + const replacementMessage: AgentMessage = { + role: "user", + content: [{ type: "text", text: "replacement instruction" }], + timestamp: Date.now(), + }; + let agent: Agent; + agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => { + if (preparedTerminatingTurn) return undefined; + preparedTerminatingTurn = true; + if (queue === "steering") { + agent.clearSteeringQueue(); + if (withReplacement) agent.steer(replacementMessage); + } else { + agent.clearFollowUpQueue(); + if (withReplacement) agent.followUp(replacementMessage); + } + return undefined; + }, + streamFn: (_model, context) => { + providerCalls++; + providerUserTexts.push( + context.messages + .filter((message) => message.role === "user") + .map((message) => + typeof message.content === "string" + ? message.content + : message.content + .filter((content) => content.type === "text") + .map((content) => content.text) + .join("\n"), + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (providerCalls === 1) { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + return; + } + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); + }); + return stream; + }, + }); + const withdrawnMessage: AgentMessage = { + role: "user", + content: [{ type: "text", text: "withdrawn safety instruction" }], + timestamp: Date.now(), + }; + + // when + const prompt = agent.prompt("run the terminating tool"); + await toolStarted.promise; + if (queue === "steering") { + agent.steer(withdrawnMessage); + } else { + agent.followUp(withdrawnMessage); + } + releaseTool.resolve(); + await prompt; + + // then + expect(providerCalls).toBe(withReplacement ? 2 : 1); + expect(providerUserTexts.flat()).not.toContain("withdrawn safety instruction"); + expect(agent.state.messages).not.toContain(withdrawnMessage); + if (withReplacement) { + expect(providerUserTexts.flat()).toEqual([ + "run the terminating tool", + "run the terminating tool", + "replacement instruction", + ]); + } + }); + + it("delivers steering queued during terminating-turn preparation before an older follow-up", async () => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const preparationStarted = createDeferred(); + const releasePreparation = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let providerCalls = 0; + const providerUserTexts: string[][] = []; + const agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => { + preparationStarted.resolve(); + await releasePreparation.promise; + return undefined; + }, + streamFn: (_model, context) => { + providerCalls++; + providerUserTexts.push( + context.messages + .filter((message) => message.role === "user") + .flatMap((message) => + typeof message.content === "string" + ? [message.content] + : message.content.filter((content) => content.type === "text").map((content) => content.text), + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (providerCalls === 1) { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + return; + } + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); + }); + return stream; + }, + }); + + // when + const prompt = agent.prompt("start"); + await toolStarted.promise; + agent.followUp({ role: "user", content: [{ type: "text", text: "old follow-up" }], timestamp: Date.now() }); + releaseTool.resolve(); + await preparationStarted.promise; + agent.steer({ role: "user", content: [{ type: "text", text: "urgent steering" }], timestamp: Date.now() }); + releasePreparation.resolve(); + await prompt; + + // then + expect(providerCalls).toBe(3); + expect(providerUserTexts[1]).toEqual(["start", "urgent steering"]); + expect(providerUserTexts[2]).toEqual(["start", "urgent steering", "old follow-up"]); + }); + + it.each([ + ["steering", false], + ["steering", true], + ["followUp", false], + ["followUp", true], + ] as const)("honors %s clear during terminating continuation turn_start (replacement: %s)", async (queue, replace) => { + // given + const toolStarted = createDeferred(); + const releaseTool = createDeferred(); + const schema = Type.Object({}); + const tool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: schema, + execute: async () => { + toolStarted.resolve(); + await releaseTool.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + let providerCalls = 0; + const providerUserTexts: string[][] = []; + const agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurnWithContext: async () => undefined, + streamFn: (_model, context) => { + providerCalls++; + providerUserTexts.push( + context.messages + .filter((message) => message.role === "user") + .flatMap((message) => + typeof message.content === "string" + ? [message.content] + : message.content.filter((content) => content.type === "text").map((content) => content.text), + ), + ); + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (providerCalls === 1) { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "terminating", arguments: {} }, + ]), + }); + return; + } + stream.push({ type: "done", reason: "stop", message: createAssistantMessage("done") }); + }); + return stream; + }, + }); + let turnStarts = 0; + agent.subscribe((event) => { + if (event.type !== "turn_start" || ++turnStarts !== 2) return; + if (queue === "steering") { + agent.clearSteeringQueue(); + if (replace) { + agent.steer({ role: "user", content: [{ type: "text", text: "replacement" }], timestamp: Date.now() }); + } + } else { + agent.clearFollowUpQueue(); + if (replace) { + agent.followUp({ + role: "user", + content: [{ type: "text", text: "replacement" }], + timestamp: Date.now(), + }); + } + } + }); + + // when + const prompt = agent.prompt("start"); + await toolStarted.promise; + const withdrawn = { + role: "user" as const, + content: [{ type: "text" as const, text: "withdrawn" }], + timestamp: Date.now(), + }; + if (queue === "steering") agent.steer(withdrawn); + else agent.followUp(withdrawn); + releaseTool.resolve(); + await prompt; + + // then + expect(providerCalls).toBe(replace ? 2 : 1); + expect(providerUserTexts.flat()).not.toContain("withdrawn"); + if (replace) expect(providerUserTexts[1]).toEqual(["start", "replacement"]); + }); + it("forwards sessionId to streamFn options", async () => { let receivedSessionId: string | undefined; const agent = new Agent({ diff --git a/packages/agent/test/harness/agent-harness.test.ts b/packages/agent/test/harness/agent-harness.test.ts index d13eca84a..80218585a 100644 --- a/packages/agent/test/harness/agent-harness.test.ts +++ b/packages/agent/test/harness/agent-harness.test.ts @@ -7,6 +7,7 @@ import { type RegisterFauxProviderOptions, } from "@earendil-works/pi-ai"; import { getModel } from "@earendil-works/pi-ai/compat"; +import { Type } from "typebox"; import { describe, expect, it } from "vitest"; import { AgentHarness } from "../../src/harness/agent-harness.ts"; import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts"; @@ -216,6 +217,74 @@ describe("AgentHarness", () => { expect(secondRequestText).toEqual(["first", "next", "second"]); }); + it.each([ + "steering", + "followUp", + ] as const)("does not restore drained %s input after abort replaces the queue during next-turn preparation", async (queue) => { + // given + const registration = newFaux(); + const toolStarted = deferred(); + const toolRelease = deferred(); + const preparationStarted = deferred(); + const preparationRelease = deferred(); + const terminatingTool: AgentTool = { + name: "terminating", + label: "Terminating", + description: "Terminates after release", + parameters: Type.Object({}), + execute: async () => { + toolStarted.resolve(); + await toolRelease.promise; + return { content: [{ type: "text", text: "done" }], details: {}, terminate: true }; + }, + }; + const session = new Session(new InMemorySessionStorage()); + const originalBuildContext = session.buildContext.bind(session); + let buildContextCalls = 0; + session.buildContext = async () => { + buildContextCalls++; + if (buildContextCalls === 2) { + preparationStarted.resolve(); + await preparationRelease.promise; + } + return originalBuildContext(); + }; + const secondRequestText: string[] = []; + registration.setResponses([ + fauxAssistantMessage(fauxToolCall("terminating", {}), { stopReason: "toolUse" }), + (context) => { + secondRequestText.push(...textFromUserMessages(context.messages)); + return fauxAssistantMessage("second"); + }, + ]); + const harness = new AgentHarness({ + models, + env: new NodeExecutionEnv({ cwd: process.cwd() }), + session, + model: registration.getModel(), + tools: [terminatingTool], + }); + const promptPromise = harness.prompt("first"); + await toolStarted.promise; + if (queue === "steering") { + await harness.steer("withdrawn instruction"); + } else { + await harness.followUp("withdrawn instruction"); + } + toolRelease.resolve(); + await preparationStarted.promise; + + // when + const abortPromise = harness.abort(); + preparationRelease.resolve(); + await Promise.all([promptPromise, abortPromise]); + await harness.prompt("second"); + + // then + expect(registration.state.callCount).toBe(2); + expect(secondRequestText).toEqual(["first", "second"]); + }); + it("drains follow-up messages one at a time after the agent would otherwise stop", async () => { const registration = newFaux(); const userCounts: number[] = []; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e1521afaf..4389463c0 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -8,6 +8,7 @@ ### Fixed +- Compacted context after tool results cross the threshold before issuing the next provider request, applied provider-bound context transforms to current and persisted summary content, and stopped safely when required compaction could not complete. - Fixed built-in web search discovery so model aliases sharing a provider endpoint no longer multiply serial fallback attempts, aborting a search stops pending native authentication discovery, and dotted private host spellings are rejected before auth or fetch ([upstream #5](https://github.com/code-yeongyu/pi-websearch/pull/5), [upstream #6](https://github.com/code-yeongyu/pi-websearch/pull/6)). ### Removed diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 31b18981b..69ae3c558 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -272,6 +272,13 @@ class CompactionRejectedError extends Error { } } +class RequiredCompactionError extends Error { + constructor() { + super("Context remains above the compaction threshold because compaction did not complete"); + this.name = "RequiredCompactionError"; + } +} + export interface ExtensionBindings { uiContext?: ExtensionUIContext; mode?: ExtensionMode; @@ -409,6 +416,7 @@ export class AgentSession { private _retryPromise: Promise | undefined = undefined; private _retryResolve: (() => void) | undefined = undefined; private _userAbortPromise: Promise | undefined = undefined; + private _suppressQueuedContinuationAfterUserAbort = false; // Bash execution state private _bashAbortController: AbortController | undefined = undefined; @@ -614,8 +622,28 @@ export class AgentSession { ? async (_turn: PrepareNextTurnContext, signal?: AbortSignal) => await this.agent.prepareNextTurn?.(signal) : undefined); this.agent.prepareNextTurnWithContext = async (turn, signal) => { - const previousSnapshot = await previousPrepareNextTurnWithContext?.(turn, signal); - const previousContext = previousSnapshot?.context ?? turn.context; + let messages = turn.context.messages; + if (turn.toolResults.length > 0) { + await this._agentEventQueue; + const compacted = await this._checkCompaction(turn.message, true, "threshold"); + if (compacted) { + messages = this.agent.state.messages.slice(); + } else { + const settings = this.settingsManager.getCompactionSettings(); + const contextTokens = estimateContextTokens( + filterContextExcludedMessages(this.sessionManager.buildSessionContext().messages), + ).tokens; + if (settings.enabled && this.model && shouldCompact(contextTokens, this.model.contextWindow, settings)) { + throw new RequiredCompactionError(); + } + } + } + const postCompactionTurn = { + ...turn, + context: { ...turn.context, messages }, + }; + const previousSnapshot = await previousPrepareNextTurnWithContext?.(postCompactionTurn, signal); + const previousContext = previousSnapshot?.context ?? postCompactionTurn.context; return { ...previousSnapshot, @@ -911,8 +939,15 @@ export class AgentSession { // Check auto-retry and auto-compaction after agent completes let launchedContinuation = false; + const userAbortSuppressedQueuedContinuation = + event.type === "agent_end" && this._suppressQueuedContinuationAfterUserAbort; + if (userAbortSuppressedQueuedContinuation) { + this._suppressQueuedContinuationAfterUserAbort = false; + } const allowsQueuedContinuation = - event.type === "agent_end" ? this._agentEndAllowsQueuedContinuation(event.messages) : false; + event.type === "agent_end" && !userAbortSuppressedQueuedContinuation + ? this._agentEndAllowsQueuedContinuation(event.messages) + : false; if (event.type === "agent_end" && this._lastAssistantMessage) { const msg = this._lastAssistantMessage; this._lastAssistantMessage = undefined; @@ -1578,7 +1613,16 @@ export class AgentSession { // The user's new prompt is sent below, so do not call agent.continue() here. const lastAssistant = this._findLastAssistantMessage(); if (lastAssistant) { - await this._checkCompaction(lastAssistant, false, "pre_prompt"); + const compacted = await this._checkCompaction(lastAssistant, false, "pre_prompt"); + if (!compacted) { + const settings = this.settingsManager.getCompactionSettings(); + const contextTokens = estimateContextTokens( + filterContextExcludedMessages(this.sessionManager.buildSessionContext().messages), + ).tokens; + if (settings.enabled && this.model && shouldCompact(contextTokens, this.model.contextWindow, settings)) { + throw new RequiredCompactionError(); + } + } } // Build messages array (custom message if any, then user message) @@ -2016,6 +2060,9 @@ export class AgentSession { await this._userAbortPromise; return; } + if (this.isStreaming) { + this._suppressQueuedContinuationAfterUserAbort = true; + } const abortPromise = (async () => { this.agent.abort(); @@ -2490,6 +2537,7 @@ export class AgentSession { this.thinkingLevel, this.agent.streamFn, env, + this.agent.transformContext, ); } } @@ -2569,7 +2617,7 @@ export class AgentSession { [...pathEntries, simulatedCompactionEntry], simulatedCompactionEntry.id, ).messages; - const contextTokens = estimateContextTokens(filterContextExcludedMessages(simulatedMessages)).tokens; + const contextTokens = estimateMessagesTokens(filterContextExcludedMessages(simulatedMessages)); const settings = this.settingsManager.getCompactionSettings(); return contextTokens > this.model.contextWindow - settings.reserveTokens; } @@ -2622,7 +2670,7 @@ export class AgentSession { private async _checkCompaction( assistantMessage: AssistantMessage, skipAbortedCheck = true, - requestReason?: "pre_prompt", + inlineReason?: "pre_prompt" | "threshold", ): Promise { const settings = this.settingsManager.getCompactionSettings(); if (!settings.enabled) return false; @@ -2676,7 +2724,7 @@ export class AgentSession { willRetry: false, errorMessage, }); - if (requestReason === "pre_prompt") { + if (inlineReason === "pre_prompt") { throw new Error(errorMessage); } return false; @@ -2690,12 +2738,11 @@ export class AgentSession { this.agent.state.messages = messages.slice(0, -1); this._incrementMessageRevision(); } - if (requestReason) { - await this._runPrePromptCompaction(assistantMessage, skipAbortedCheck, "overflow", willRetry); + if (inlineReason) { + return await this._runPrePromptCompaction(assistantMessage, skipAbortedCheck, "overflow", willRetry); } else { return await this._runAutoCompaction("overflow", willRetry); } - return false; } // Case 2: Threshold - context is getting large @@ -2703,29 +2750,34 @@ 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 directContextTokens = assistantMessage.usage ? calculateContextTokens(assistantMessage.usage) : 0; - if (assistantMessage.stopReason === "error" || directContextTokens === 0) { - const messages = filterContextExcludedMessages(this.agent.state.messages); - const estimate = estimateContextTokens(messages); - if (estimate.lastUsageIndex === null) return false; // No usage data at all - // Verify the usage source is post-compaction. Kept pre-compaction messages - // have stale usage reflecting the old (larger) context and would falsely - // trigger compaction right after one just finished. - const usageMsg = messages[estimate.lastUsageIndex]; - if ( - compactionEntry && - usageMsg.role === "assistant" && - (usageMsg as AssistantMessage).timestamp <= new Date(compactionEntry.timestamp).getTime() - ) { - return false; - } - contextTokens = estimate.tokens; + if (inlineReason) { + const messages = filterContextExcludedMessages(this.sessionManager.buildSessionContext().messages); + contextTokens = estimateContextTokens(messages).tokens; } else { - contextTokens = directContextTokens; + const directContextTokens = assistantMessage.usage ? calculateContextTokens(assistantMessage.usage) : 0; + if (assistantMessage.stopReason !== "error" && directContextTokens !== 0) { + contextTokens = directContextTokens; + } else { + const messages = filterContextExcludedMessages(this.agent.state.messages); + const estimate = estimateContextTokens(messages); + if (estimate.lastUsageIndex === null) return false; // No usage data at all + // Verify the usage source is post-compaction. Kept pre-compaction messages + // have stale usage reflecting the old (larger) context and would falsely + // trigger compaction right after one just finished. + const usageMsg = messages[estimate.lastUsageIndex]; + if ( + compactionEntry && + usageMsg.role === "assistant" && + (usageMsg as AssistantMessage).timestamp <= new Date(compactionEntry.timestamp).getTime() + ) { + return false; + } + contextTokens = estimate.tokens; + } } if (shouldCompact(contextTokens, contextWindow, settings)) { - if (requestReason) { - await this._runPrePromptCompaction(assistantMessage, skipAbortedCheck); + if (inlineReason) { + return await this._runPrePromptCompaction(assistantMessage, skipAbortedCheck, inlineReason); } else { return await this._runAutoCompaction("threshold", false); } @@ -2736,9 +2788,9 @@ export class AgentSession { private async _runPrePromptCompaction( lastAssistantMessage: AssistantMessage, skipAbortedCheck: boolean, - reason: "pre_prompt" | "overflow" = "pre_prompt", + reason: "pre_prompt" | "overflow" | "threshold" = "pre_prompt", willRetry = false, - ): Promise { + ): Promise { this._emit({ type: "compaction_start", reason }); this._compactionAbortController = new AbortController(); @@ -2752,6 +2804,7 @@ export class AgentSession { if (!execution.accepted && isContextOverflow(lastAssistantMessage, this.model?.contextWindow ?? 0)) { this._overflowRecoveryAttempted = false; } + return execution.accepted; } catch (error) { if (isContextOverflow(lastAssistantMessage, this.model?.contextWindow ?? 0)) { this._overflowRecoveryAttempted = false; @@ -2767,6 +2820,7 @@ export class AgentSession { willRetry, errorMessage: aborted ? undefined : `Pre-prompt compaction failed: ${errorMessage}`, }); + return false; } finally { this._compactionAbortController = undefined; } diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index e9c130785..ee42ba232 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -573,6 +573,41 @@ async function completeSummarization( return stream.result(); } +async function transformSummarySource( + currentMessages: AgentMessage[], + previousSummary: string | undefined, + transformContext: ((messages: AgentMessage[], signal?: AbortSignal) => Promise) | undefined, + signal: AbortSignal | undefined, +): Promise<{ readonly messages: AgentMessage[]; readonly previousSummary: string | undefined }> { + if (!transformContext) return { messages: currentMessages, previousSummary }; + if (!previousSummary) { + return { messages: await transformContext(currentMessages, signal), previousSummary: undefined }; + } + + const timestamps = new Set(currentMessages.map((message) => message.timestamp)); + let summaryTimestamp = -1; + while (timestamps.has(summaryTimestamp)) { + summaryTimestamp--; + } + const transformed = await transformContext( + [ + { + role: "user", + content: [{ type: "text", text: previousSummary }], + timestamp: summaryTimestamp, + }, + ...currentMessages, + ], + signal, + ); + const transformedSummary = transformed.filter((message) => message.timestamp === summaryTimestamp); + return { + messages: transformed.filter((message) => message.timestamp !== summaryTimestamp), + previousSummary: + transformedSummary.length > 0 ? serializeConversation(convertToLlm(transformedSummary)) : undefined, + }; +} + /** * Generate a summary of the conversation using the LLM. * If previousSummary is provided, uses the update prompt to merge. @@ -590,27 +625,31 @@ export async function generateSummary( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise, ): Promise { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, ); + const transformedSource = await transformSummarySource(currentMessages, previousSummary, transformContext, signal); + const providerPreviousSummary = transformedSource.previousSummary; + // Use update prompt if we have a previous summary, otherwise initial prompt - let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + let basePrompt = providerPreviousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; if (customInstructions) { basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; } // Serialize conversation to text so model doesn't try to continue it // Convert to LLM messages first (handles custom types like bashExecution, custom, etc.) - const llmMessages = convertToLlm(currentMessages); + const llmMessages = convertToLlm(transformedSource.messages); const conversationText = serializeConversation(llmMessages); // Build the prompt with conversation wrapped in tags let promptText = `\n${conversationText}\n\n\n`; - if (previousSummary) { - promptText += `\n${previousSummary}\n\n\n`; + if (providerPreviousSummary) { + promptText += `\n${providerPreviousSummary}\n\n\n`; } promptText += basePrompt; @@ -793,6 +832,7 @@ export async function compact( thinkingLevel?: ThinkingLevel, streamFn?: StreamFn, env?: Record, + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise, ): Promise { const { firstKeptEntryId, @@ -824,6 +864,7 @@ export async function compact( thinkingLevel, streamFn, env, + transformContext, ) : "No prior history."; const turnPrefixResult = await generateTurnPrefixSummary( @@ -837,6 +878,7 @@ export async function compact( extraBody, thinkingLevel, streamFn, + transformContext, ); // Merge into single summary summary = `${historyResult}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult}`; @@ -855,6 +897,7 @@ export async function compact( thinkingLevel, streamFn, env, + transformContext, ); } @@ -888,12 +931,14 @@ async function generateTurnPrefixSummary( extraBody?: Record, thinkingLevel?: ThinkingLevel, streamFn?: SummarizationStreamFn, + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise, ): Promise { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, ); // Smaller budget for turn prefix - const llmMessages = convertToLlm(messages); + const providerMessages = transformContext ? await transformContext(messages, signal) : messages; + const llmMessages = convertToLlm(providerMessages); const conversationText = serializeConversation(llmMessages); const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; const summarizationMessages = [ diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index 97a6da341..bafe98b5d 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -1,6 +1,13 @@ -import { type AssistantMessage, createAssistantMessageEventStream, fauxAssistantMessage } from "@earendil-works/pi-ai"; +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import { + type AssistantMessage, + createAssistantMessageEventStream, + fauxAssistantMessage, + fauxToolCall, +} from "@earendil-works/pi-ai"; +import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { estimateTokens } from "../../src/core/compaction/index.ts"; +import { estimateContextTokens, estimateTokens, generateSummary } from "../../src/core/compaction/index.ts"; import type { ExtensionAPI } from "../../src/core/extensions/index.ts"; import { createHarness, type Harness } from "./harness.ts"; @@ -159,6 +166,32 @@ function seedCompactableSession(harness: Harness): void { harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; } +function createResultToolExtension(toolResult: string, compactionSummary?: string, terminate = false) { + return (pi: ExtensionAPI): void => { + pi.registerTool({ + name: "large_result", + label: "Large Result", + description: "Return text for next-turn compaction tests", + parameters: Type.Object({}), + execute: async () => ({ + content: [{ type: "text", text: toolResult }], + details: {}, + terminate, + }), + }); + if (compactionSummary !== undefined) { + pi.on("session_before_compact", async (event) => ({ + compaction: { + summary: compactionSummary, + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + } + }; +} + describe("AgentSession compaction characterization", () => { const harnesses: Harness[] = []; @@ -296,6 +329,497 @@ describe("AgentSession compaction characterization", () => { expect(harness.eventsOfType("compaction_end")).toHaveLength(0); }); + it("compacts trailing tool results before provider call 2 when they push context over the threshold", async () => { + // given + const contextWindow = 5_000; + const reserveTokens = 1_000; + const threshold = contextWindow - reserveTokens; + const largeToolResult = "tool output ".repeat(300); + const harness = await createHarness({ + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [createResultToolExtension(largeToolResult, "tool result threshold summary")], + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "prior context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + text: "prior response", + stopReason: "stop", + totalTokens: 700, + timestamp: seedTimestamp + 1_000, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + let call2Context = ""; + let compactionEndsAtCall2 = 0; + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), + (context) => { + call2Context = JSON.stringify(context.messages); + compactionEndsAtCall2 = harness.eventsOfType("compaction_end").length; + return fauxAssistantMessage("done after compaction"); + }, + ]); + + // when + await harness.session.prompt("run the large result tool"); + + // then + const persistedMessages = harness.sessionManager + .getEntries() + .filter((entry) => entry.type === "message") + .map((entry) => entry.message); + const toolCallResponse = persistedMessages.find( + (message) => message.role === "assistant" && message.stopReason === "toolUse", + ); + const toolResult = persistedMessages.find((message) => message.role === "toolResult"); + if (toolCallResponse?.role !== "assistant" || toolResult?.role !== "toolResult") { + throw new Error("Expected a successful assistant tool call and its appended tool result"); + } + const assembledContext = estimateContextTokens([toolCallResponse, toolResult]); + expect(toolCallResponse.usage.totalTokens).toBeLessThan(threshold); + expect(assembledContext.tokens).toBeGreaterThan(threshold); + expect(compactionEndsAtCall2).toBe(1); + expect(call2Context).toContain("tool result threshold summary"); + expect(call2Context).toContain(largeToolResult); + expect(call2Context).not.toContain("prior context "); + expect(harness.faux.state.callCount).toBe(2); + expect(harness.eventsOfType("compaction_start").map((event) => event.reason)).toEqual(["threshold"]); + expect(harness.eventsOfType("compaction_end")).toHaveLength(1); + expect(harness.eventsOfType("compaction_end")[0]).toMatchObject({ + reason: "threshold", + aborted: false, + willRetry: false, + accepted: true, + }); + }); + + it("compacts a terminating tool result before the next separate user prompt", async () => { + // given + const contextWindow = 5_000; + const reserveTokens = 1_000; + const threshold = contextWindow - reserveTokens; + const largeToolResult = "terminating tool output ".repeat(300); + const compactionSummary = "terminating tool result summary"; + let call2Context = ""; + let compactionEndsAtCall2 = 0; + const harness = await createHarness({ + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [createResultToolExtension(largeToolResult, compactionSummary, true)], + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "prior terminating context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + text: "prior response", + stopReason: "stop", + totalTokens: 700, + timestamp: seedTimestamp + 1_000, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + const terminatingToolCall = Object.assign(createAssistant(harness, { stopReason: "toolUse", totalTokens: 700 }), { + content: [fauxToolCall("large_result", {})], + }); + harness.setResponses([ + terminatingToolCall, + (context) => { + call2Context = JSON.stringify(context.messages); + compactionEndsAtCall2 = harness.eventsOfType("compaction_end").length; + return createAssistant(harness, { + text: "done after pre-prompt compaction", + stopReason: "stop", + totalTokens: 1_000, + }); + }, + ]); + + await harness.session.prompt("run the terminating result tool"); + const toolResult = harness.session.messages.find((message) => message.role === "toolResult"); + if (toolResult?.role !== "toolResult") { + throw new Error("Expected the terminating tool result in the persisted session context"); + } + const persistedContext = estimateContextTokens(harness.sessionManager.buildSessionContext().messages); + expect(terminatingToolCall.usage.totalTokens).toBeLessThan(threshold); + expect(persistedContext.tokens).toBeGreaterThan(threshold); + expect(harness.faux.state.callCount).toBe(1); + expect(harness.eventsOfType("compaction_start")).toHaveLength(0); + + // when + await harness.session.prompt("continue after the terminating tool"); + + // then + expect(compactionEndsAtCall2).toBe(1); + expect(call2Context).toContain(compactionSummary); + expect(call2Context).toContain(largeToolResult); + expect(call2Context).not.toContain("prior terminating context "); + expect(harness.faux.state.callCount).toBe(2); + expect(harness.eventsOfType("compaction_start").filter((event) => event.reason === "pre_prompt")).toHaveLength(1); + }); + + it("stops before a separate prompt provider call when required pre-prompt compaction is cancelled", async () => { + // given + const contextWindow = 5_000; + const reserveTokens = 1_000; + const largeToolResult = "terminating oversized output ".repeat(350); + const harness = await createHarness({ + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [ + createResultToolExtension(largeToolResult, undefined, true), + (pi) => { + pi.on("session_before_compact", async () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "prior separate prompt context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + text: "prior response", + stopReason: "stop", + totalTokens: 700, + timestamp: seedTimestamp + 1_000, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + harness.setResponses([ + Object.assign(createAssistant(harness, { stopReason: "toolUse", totalTokens: 700 }), { + content: [fauxToolCall("large_result", {})], + }), + () => { + throw new Error("provider call 2 must not be reached"); + }, + ]); + await harness.session.prompt("run the terminating result tool"); + expect(estimateContextTokens(harness.sessionManager.buildSessionContext().messages).tokens).toBeGreaterThan( + contextWindow - reserveTokens, + ); + + // when / then + await expect(harness.session.prompt("continue after cancelled compaction")).rejects.toThrow( + "Context remains above the compaction threshold because compaction did not complete", + ); + expect(harness.faux.state.callCount).toBe(1); + expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toHaveLength(0); + expect(harness.eventsOfType("compaction_end")).toContainEqual( + expect.objectContaining({ reason: "pre_prompt", accepted: false, rejectionCause: "cancelled-by-extension" }), + ); + }); + + it("preserves a constructor-supplied next-turn callback after rebuilding compacted context", async () => { + // given + const largeToolResult = "callback tool output ".repeat(300); + const compactionSummary = "callback-preserved summary"; + const callbackContexts: string[] = []; + const prepareNextTurnWithContext = vi.fn(async (turn) => { + callbackContexts.push(JSON.stringify(turn.context.messages)); + return undefined; + }); + const harness = await createHarness({ + settings: { + compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: 1_000 }, + retry: { enabled: false }, + }, + models: [{ id: "faux-1", contextWindow: 5_000 }], + extensionFactories: [createResultToolExtension(largeToolResult, compactionSummary)], + prepareNextTurnWithContext, + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "callback prior context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + text: "prior response", + stopReason: "stop", + totalTokens: 700, + timestamp: seedTimestamp + 1_000, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done after callback", { stopReason: "error", errorMessage: "test complete" }), + ]); + + // when + await harness.session.prompt("run the callback result tool"); + + // then + expect(prepareNextTurnWithContext).toHaveBeenCalledTimes(1); + expect(callbackContexts[0]).toContain(compactionSummary); + expect(callbackContexts[0]).toContain(largeToolResult); + expect(callbackContexts[0]).not.toContain("callback prior context "); + }); + + it("applies the provider context transform to inline compaction summarization", async () => { + // given + const sensitiveToolOutput = "SENSITIVE_TOOL_OUTPUT"; + const permittedToolOutput = "PERMITTED_TOOL_OUTPUT"; + const contextWindow = 5_000; + const reserveTokens = 1_000; + const largeToolResult = "tool output ".repeat(300); + const priorToolResult = `${sensitiveToolOutput}\n${permittedToolOutput}`; + const priorToolCall = fauxToolCall("prior_tool", {}); + const compactionSummary = "provider-generated callback summary"; + const harness = await createHarness({ + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [ + createResultToolExtension(largeToolResult), + (pi) => { + pi.on("context", async (event) => ({ + messages: event.messages.map((message) => { + if (message.role !== "toolResult") return message; + return { + ...message, + content: message.content.map((content) => + content.type === "text" + ? { ...content, text: content.text.replaceAll(sensitiveToolOutput, "") } + : content, + ), + }; + }), + })); + }, + ], + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "prior context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + Object.assign( + createAssistant(harness, { + stopReason: "toolUse", + totalTokens: 700, + timestamp: seedTimestamp + 500, + }), + { content: [priorToolCall] }, + ), + ); + harness.sessionManager.appendMessage({ + role: "toolResult", + toolCallId: priorToolCall.id, + toolName: "prior_tool", + content: [{ type: "text", text: priorToolResult }], + details: {}, + isError: false, + timestamp: seedTimestamp + 1_000, + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), + fauxAssistantMessage(compactionSummary), + fauxAssistantMessage("provider-generated turn prefix summary"), + fauxAssistantMessage("done after callback-filtered compaction"), + ]); + + // when + await harness.session.prompt("run the sensitive result tool"); + + // then + const callLog = harness.faux.getCallLog(); + const historyCompactionContext = JSON.stringify(callLog[1]?.context.messages ?? []); + const turnPrefixCompactionContext = JSON.stringify(callLog[2]?.context.messages ?? []); + const continuationProviderContext = JSON.stringify(callLog[3]?.context.messages ?? []); + expect(historyCompactionContext).toContain(permittedToolOutput); + expect(historyCompactionContext).not.toContain(sensitiveToolOutput); + expect(turnPrefixCompactionContext).not.toContain(sensitiveToolOutput); + expect(continuationProviderContext).toContain(compactionSummary); + expect(continuationProviderContext).not.toContain(sensitiveToolOutput); + expect(harness.faux.state.callCount).toBe(4); + expect(harness.eventsOfType("compaction_end")).toContainEqual( + expect.objectContaining({ reason: "threshold", accepted: true }), + ); + }); + + it("applies the provider context transform to a persisted previous summary", async () => { + // given + const sensitiveSummary = "SENSITIVE_PREVIOUS_SUMMARY"; + const injectedDirective = "INJECTED_SAFETY_DIRECTIVE"; + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("updated summary")]); + const transformContext = vi.fn(async (messages: AgentMessage[]) => [ + ...messages.map((message) => { + if (message.role !== "user" || typeof message.content === "string") return message; + return { + ...message, + content: message.content.map((content) => + content.type === "text" + ? { ...content, text: content.text.replaceAll(sensitiveSummary, "") } + : content, + ), + }; + }), + { + role: "user" as const, + content: [{ type: "text" as const, text: injectedDirective }], + timestamp: Date.now(), + }, + ]); + + // when + await generateSummary( + [ + { + role: "user", + content: [{ type: "text", text: "new conversation content" }], + timestamp: Date.now(), + }, + ], + harness.getModel(), + 1_000, + "faux-key", + undefined, + undefined, + undefined, + sensitiveSummary, + undefined, + undefined, + undefined, + undefined, + transformContext, + ); + + // then + const providerContext = JSON.stringify(harness.faux.getCallLog()[0]?.context.messages ?? []); + expect(transformContext).toHaveBeenCalledTimes(1); + expect(providerContext).not.toContain(sensitiveSummary); + expect(providerContext).toContain("new conversation content"); + expect(providerContext.split(injectedDirective)).toHaveLength(2); + expect(harness.faux.state.callCount).toBe(1); + }); + + it("stops before provider call 2 when required inline threshold compaction is cancelled", async () => { + // given + const contextWindow = 5_000; + const reserveTokens = 1_000; + const largeToolResult = "tool output ".repeat(300); + const harness = await createHarness({ + settings: { compaction: { enabled: true, keepRecentTokens: 1, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [ + createResultToolExtension(largeToolResult), + (pi) => { + pi.on("session_before_compact", async () => ({ cancel: true })); + }, + ], + }); + harnesses.push(harness); + const seedTimestamp = Date.now() - 2_000; + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "prior context ".repeat(220) }], + timestamp: seedTimestamp, + }); + harness.sessionManager.appendMessage( + createAssistant(harness, { + text: "prior response", + stopReason: "stop", + totalTokens: 700, + timestamp: seedTimestamp + 1_000, + }), + ); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + harness.setResponses([fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" })]); + + // when + await harness.session.prompt("run the large result tool"); + + // then + expect(harness.faux.state.callCount).toBe(1); + expect(harness.session.agent.state.errorMessage).toContain( + "Context remains above the compaction threshold because compaction did not complete", + ); + expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toHaveLength(0); + expect(harness.eventsOfType("compaction_end")).toContainEqual( + expect.objectContaining({ + reason: "threshold", + accepted: false, + rejectionCause: "cancelled-by-extension", + }), + ); + }); + + it("continues to provider call 2 without compaction when trailing tool results stay below the threshold", async () => { + // given + const harness = await createHarness({ + settings: { compaction: { enabled: true, reserveTokens: 1_000 } }, + models: [{ id: "faux-1", contextWindow: 5_000 }], + extensionFactories: [createResultToolExtension("small tool output")], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done without compaction"), + ]); + + // when + await harness.session.prompt("run the small result tool"); + + // then + expect(harness.faux.state.callCount).toBe(2); + expect(harness.eventsOfType("compaction_start")).toHaveLength(0); + expect(harness.eventsOfType("compaction_end")).toHaveLength(0); + }); + + it("continues to provider call 2 without compaction when compaction is disabled", async () => { + // given + const contextWindow = 5_000; + const reserveTokens = 1_000; + const threshold = contextWindow - reserveTokens; + const largeToolResult = "tool output ".repeat(2_000); + const harness = await createHarness({ + settings: { compaction: { enabled: false, reserveTokens } }, + models: [{ id: "faux-1", contextWindow }], + extensionFactories: [createResultToolExtension(largeToolResult)], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("large_result", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done with compaction disabled"), + ]); + + // when + await harness.session.prompt("run the large result tool"); + + // then + const toolResult = harness.session.messages.find((message) => message.role === "toolResult"); + if (toolResult?.role !== "toolResult") { + throw new Error("Expected the large tool result in the session context"); + } + expect(estimateTokens(toolResult)).toBeGreaterThan(threshold); + expect(harness.faux.state.callCount).toBe(2); + expect(harness.eventsOfType("compaction_start")).toHaveLength(0); + expect(harness.eventsOfType("compaction_end")).toHaveLength(0); + }); + it("cancels in-progress manual compaction when abortCompaction is called", async () => { const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } }, diff --git a/packages/coding-agent/test/suite/compaction-race.test.ts b/packages/coding-agent/test/suite/compaction-race.test.ts index ae0827110..3157ae6e3 100644 --- a/packages/coding-agent/test/suite/compaction-race.test.ts +++ b/packages/coding-agent/test/suite/compaction-race.test.ts @@ -144,7 +144,12 @@ describe("AgentSession compaction race handling", () => { } expect(harness.faux.state.callCount).toBe(2); - expect(getUserTexts(harness)).toEqual([initialPrompt, "prompt during compaction"]); + const secondCall = harness.faux.getCallLog()[1]; + expect( + secondCall?.context.messages.some((message) => hasCompactionSummary(message, "slow threshold summary")), + ).toBe(true); + expect(secondCall?.context.messages.some((message) => contextHasText(message, initialPrompt))).toBe(false); + expect(getUserTexts(harness)).toEqual(["prompt during compaction"]); }); it("given overflow compaction has a queued follow-up when a new prompt arrives, when recovery finishes, then queued messages drain before the new prompt", async () => { @@ -233,6 +238,11 @@ describe("AgentSession compaction race handling", () => { // then expect(harness.faux.state.callCount).toBe(2); - expect(getUserTexts(harness)).toEqual([initialPrompt, "after compaction prompt"]); + const secondCall = harness.faux.getCallLog()[1]; + expect( + secondCall?.context.messages.some((message) => hasCompactionSummary(message, "normal threshold summary")), + ).toBe(true); + expect(secondCall?.context.messages.some((message) => contextHasText(message, initialPrompt))).toBe(false); + expect(getUserTexts(harness)).toEqual(["after compaction prompt"]); }); }); diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index 9407297f3..74d8279a9 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -5,7 +5,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import type { AgentMessage, AgentTool } from "@earendil-works/pi-agent-core"; +import type { AgentMessage, AgentOptions, AgentTool } from "@earendil-works/pi-agent-core"; import { Agent } from "@earendil-works/pi-agent-core"; import type { FauxModelDefinition, @@ -75,6 +75,7 @@ export interface HarnessOptions { withConfiguredAuth?: boolean; upstreamModelId?: string; onPayload?: (payload: unknown) => void; + prepareNextTurnWithContext?: AgentOptions["prepareNextTurnWithContext"]; persistSession?: boolean; autoTitleSessions?: boolean; } @@ -181,6 +182,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise; resolve: () => void } { + let resolve: (() => void) | undefined; + const promise = new Promise((deferredResolve) => { + resolve = deferredResolve; + }); + return { promise, resolve: () => resolve?.() }; +} + +async function createBlockedTerminatingRun(): Promise<{ + readonly harness: Harness; + readonly promptPromise: Promise; + readonly releaseTool: () => void; + readonly preparationStarted: Promise; + readonly releasePreparation: () => void; +}> { + const toolStarted = createDeferred(); + const toolRelease = createDeferred(); + const preparationStarted = createDeferred(); + const preparationRelease = createDeferred(); + const terminatingTool: AgentTool = { + name: "terminate_after_wait", + label: "Terminate after wait", + description: "Wait, then terminate the current run", + parameters: Type.Object({}), + execute: async () => { + toolStarted.resolve(); + await toolRelease.promise; + return { + content: [{ type: "text", text: "terminated" }], + details: {}, + terminate: true, + }; + }, + }; + const harness = await createHarness({ + tools: [terminatingTool], + prepareNextTurnWithContext: async () => { + preparationStarted.resolve(); + await preparationRelease.promise; + return undefined; + }, + }); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("terminate_after_wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("replacement handled"), + ]); + const promptPromise = harness.session.prompt("start"); + await toolStarted.promise; + return { + harness, + promptPromise, + releaseTool: toolRelease.resolve, + preparationStarted: preparationStarted.promise, + releasePreparation: preparationRelease.resolve, + }; +} + +async function queueMessage(harness: Harness, queue: QueueName, text: string): Promise { + if (queue === "steering") { + await harness.session.steer(text); + return; + } + await harness.session.followUp(text); +} + +describe("terminating-tool queue ownership", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it.each([ + "steering", + "followUp", + ] as const)("retains %s input without restarting provider work when public abort interrupts preparation", async (queue) => { + // given + const run = await createBlockedTerminatingRun(); + harnesses.push(run.harness); + await queueMessage(run.harness, queue, "retain after abort"); + run.releaseTool(); + await run.preparationStarted; + + // when + const abortPromise = run.harness.session.abort(); + run.releasePreparation(); + await Promise.all([run.promptPromise, abortPromise]); + + // then + expect(run.harness.faux.state.callCount).toBe(1); + expect(getUserTexts(run.harness)).toEqual(["start"]); + expect(run.harness.session.pendingMessageCount).toBe(1); + expect(run.harness.session.agent.hasQueuedMessages()).toBe(true); + }); + + it.each([ + ["steering", false], + ["steering", true], + ["followUp", false], + ["followUp", true], + ] as const)("does not deliver cleared %s input after successful preparation (replacement: %s)", async (queue, replace) => { + // given + const run = await createBlockedTerminatingRun(); + harnesses.push(run.harness); + await queueMessage(run.harness, queue, "withdrawn safety instruction"); + run.releaseTool(); + await run.preparationStarted; + + // when + const cleared = run.harness.session.clearQueue(); + if (replace) { + await queueMessage(run.harness, queue, "replacement instruction"); + } + run.releasePreparation(); + await run.promptPromise; + + // then + expect(cleared).toEqual( + queue === "steering" + ? { steering: ["withdrawn safety instruction"], followUp: [] } + : { steering: [], followUp: ["withdrawn safety instruction"] }, + ); + expect(run.harness.faux.state.callCount).toBe(replace ? 2 : 1); + expect(getUserTexts(run.harness)).toEqual(replace ? ["start", "replacement instruction"] : ["start"]); + expect(run.harness.session.pendingMessageCount).toBe(0); + expect(run.harness.session.agent.hasQueuedMessages()).toBe(false); + }); +});