diff --git a/docs/.generated/plugin-sdk-api-baseline.sha256 b/docs/.generated/plugin-sdk-api-baseline.sha256 index 5a4165b3e93e..8364b5f3a613 100644 --- a/docs/.generated/plugin-sdk-api-baseline.sha256 +++ b/docs/.generated/plugin-sdk-api-baseline.sha256 @@ -1,2 +1,2 @@ -148fa180c89f7d4797c59b1d1779b51a25cf7123e874e64cce99024bd1cd771e plugin-sdk-api-baseline.json -ba44e7b8e6f7e6a50f3e8de5fe792199ce779c5eedb09928e8a6220e15316829 plugin-sdk-api-baseline.jsonl +5bab8bc957d81266c1c8cb760ff989e03ca8ce9c4057c990137f0ec2de2f003e plugin-sdk-api-baseline.json +c5f6cd4cb5919af53c6d59765aa133740f00ab6002c65761bc3eedabc6c3ba70 plugin-sdk-api-baseline.jsonl diff --git a/packages/agent-core/src/agent-loop.test.ts b/packages/agent-core/src/agent-loop.test.ts index 167e2f84c68f..bb084f8348a6 100644 --- a/packages/agent-core/src/agent-loop.test.ts +++ b/packages/agent-core/src/agent-loop.test.ts @@ -241,6 +241,109 @@ describe("agentLoop continuation guards", () => { role: "assistant", }); }); + + it("delivers a queued follow-up before continuing from a tool result", async () => { + let requestContext: Context | undefined; + const streamFn: StreamFn = (activeModel, context) => { + requestContext = context; + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "stop", + message: { + role: "assistant", + content: [{ type: "text", text: "done" }], + api: activeModel.api, + provider: activeModel.provider, + model: activeModel.id, + usage: TEST_USAGE, + stopReason: "stop", + timestamp: 3, + }, + }); + stream.end(); + }); + return stream; + }; + const agent = new Agent({ + initialState: { + model, + systemPrompt: "", + tools: [], + messages: [ + { + role: "toolResult", + toolCallId: "call-finish", + toolName: "finish", + content: [{ type: "text", text: "finished" }], + details: {}, + isError: false, + timestamp: 1, + }, + ], + }, + convertToLlm: (messages) => messages as Message[], + streamFn, + }); + agent.followUp({ role: "user", content: "queued after end", timestamp: 2 }); + + await agent.continue(); + + expect(requestContext?.messages.at(-1)).toMatchObject({ + role: "user", + content: "queued after end", + }); + }); + + it("keeps a queued follow-up behind a trailing user continuation", async () => { + const requestContexts: Context[] = []; + const streamFn: StreamFn = (activeModel, context) => { + requestContexts.push(context); + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "stop", + message: { + role: "assistant", + content: [{ type: "text", text: `answer ${requestContexts.length}` }], + api: activeModel.api, + provider: activeModel.provider, + model: activeModel.id, + usage: TEST_USAGE, + stopReason: "stop", + timestamp: requestContexts.length + 1, + }, + }); + stream.end(); + }); + return stream; + }; + const agent = new Agent({ + initialState: { + model, + systemPrompt: "", + tools: [], + messages: [{ role: "user", content: "retry this turn", timestamp: 1 }], + }, + convertToLlm: (messages) => messages as Message[], + streamFn, + }); + agent.followUp({ role: "user", content: "queued after retry", timestamp: 2 }); + + await agent.continue(); + + expect(requestContexts).toHaveLength(2); + expect(requestContexts[0]?.messages.at(-1)).toMatchObject({ + role: "user", + content: "retry this turn", + }); + expect(requestContexts[1]?.messages.at(-1)).toMatchObject({ + role: "user", + content: "queued after retry", + }); + }); }); describe("agentLoop streaming updates", () => { @@ -1053,6 +1156,104 @@ describe("agentLoop tool termination", () => { expect(events.at(-1)).toMatchObject({ type: "agent_end" }); }); + it("normalizes a tool result with missing content before the next model turn", async () => { + const contexts: Context[] = []; + let turn = 0; + const streamFn: StreamFn = (_activeModel, context) => { + contexts.push(context); + turn += 1; + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const message = + turn === 1 + ? makeAssistantMessage([ + { type: "toolCall", id: "call-empty", name: "empty", arguments: {} }, + ]) + : makeAssistantMessage([{ type: "text", text: "done" }]); + stream.push({ type: "done", reason: turn === 1 ? "toolUse" : "stop", message }); + stream.end(); + }); + return stream; + }; + const tool: AgentTool = { + name: "empty", + label: "empty", + description: "returns no display content", + parameters: Type.Object({}, { additionalProperties: false }), + execute: async () => ({ details: { ok: true } }) as AgentToolResult, + }; + + await collectEvents( + agentLoop( + [{ role: "user", content: "run", timestamp: 1 }], + { systemPrompt: "", messages: [], tools: [tool] }, + config, + undefined, + streamFn, + ), + ); + + expect(contexts).toHaveLength(2); + expect(contexts[1]?.messages).toContainEqual( + expect.objectContaining({ + role: "toolResult", + toolName: "empty", + content: [], + }), + ); + }); + + it("preserves extra tool result fields when an after hook patches the result", async () => { + const extra = { deliveryId: "delivery-1" }; + const originalResult = { + content: [{ type: "text" as const, text: "sent" }], + details: { phase: "original" }, + extra, + }; + const tool: AgentTool = { + name: "patched", + label: "patched", + description: "returns extended result metadata", + parameters: Type.Object({}, { additionalProperties: false }), + execute: async () => originalResult, + }; + const streamFn: StreamFn = () => { + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const message = makeAssistantMessage([ + { type: "toolCall", id: "call-patched", name: tool.name, arguments: {} }, + ]); + stream.push({ type: "done", reason: "toolUse", message }); + stream.end(); + }); + return stream; + }; + + const events = await collectEvents( + agentLoop( + [{ role: "user", content: "run", timestamp: 1 }], + { systemPrompt: "", messages: [], tools: [tool] }, + { + ...config, + afterToolCall: async () => ({ details: { phase: "patched" }, terminate: true }), + }, + undefined, + streamFn, + ), + ); + const endEvent = events.find( + (event): event is Extract => + event.type === "tool_execution_end", + ); + + expect(endEvent?.result).toMatchObject({ + content: originalResult.content, + details: { phase: "patched" }, + extra, + terminate: true, + }); + }); + it("marks policy-blocked tool calls as not executed", async () => { const executed: string[] = []; let turn = 0; @@ -1486,6 +1687,79 @@ describe("agentLoop tool termination", () => { }); }); +describe("Agent next-turn preparation", () => { + it("forwards completed-turn context and applies its update to the following request", async () => { + const nextModel = { ...model, id: "next-model" }; + const requests: Array<{ model: string; systemPrompt: string; tools: string[] }> = []; + let turn = 0; + const streamFn: StreamFn = (activeModel, context) => { + requests.push({ + model: activeModel.id, + systemPrompt: context.systemPrompt ?? "", + tools: context.tools?.map((tool) => tool.name) ?? [], + }); + turn += 1; + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const content: AssistantMessage["content"] = + turn === 1 + ? [{ type: "toolCall", id: "call-refresh", name: "refresh", arguments: {} }] + : [{ type: "text", text: "done" }]; + stream.push({ + type: "done", + reason: turn === 1 ? "toolUse" : "stop", + message: { + role: "assistant", + content, + api: activeModel.api, + provider: activeModel.provider, + model: activeModel.id, + usage: TEST_USAGE, + stopReason: turn === 1 ? "toolUse" : "stop", + timestamp: turn, + }, + }); + stream.end(); + }); + return stream; + }; + const tool: AgentTool = { + name: "refresh", + label: "refresh", + description: "refresh turn state", + parameters: Type.Object({}, { additionalProperties: false }), + execute: async () => ({ content: [{ type: "text", text: "refreshed" }], details: {} }), + }; + const prepareNextTurnWithContext = vi.fn(({ context }) => ({ + context: { ...context, systemPrompt: "refreshed prompt", tools: [] }, + model: nextModel, + })); + const prepareNextTurn = vi.fn(() => ({ + context: { systemPrompt: "legacy prompt", messages: [], tools: [tool] }, + })); + const agent = new Agent({ + initialState: { model, systemPrompt: "initial prompt", tools: [tool] }, + convertToLlm: (messages) => messages as Message[], + streamFn, + prepareNextTurn, + prepareNextTurnWithContext, + }); + + await agent.prompt("start"); + + expect(prepareNextTurnWithContext).toHaveBeenCalled(); + expect(prepareNextTurn).not.toHaveBeenCalled(); + expect(prepareNextTurnWithContext.mock.calls[0]?.[0]).toMatchObject({ + message: { role: "assistant", stopReason: "toolUse" }, + toolResults: [{ role: "toolResult", toolName: "refresh" }], + }); + expect(requests).toEqual([ + { model: model.id, systemPrompt: "initial prompt", tools: ["refresh"] }, + { model: nextModel.id, systemPrompt: "refreshed prompt", tools: [] }, + ]); + }); +}); + describe("agentLoop thinking state", () => { function makeAssistantMessage( activeModel: Model, diff --git a/packages/agent-core/src/agent-loop.ts b/packages/agent-core/src/agent-loop.ts index 37466082afac..c03258171c50 100644 --- a/packages/agent-core/src/agent-loop.ts +++ b/packages/agent-core/src/agent-loop.ts @@ -1063,6 +1063,7 @@ async function finalizeExecutedToolCall( ); if (afterResult) { result = { + ...result, content: afterResult.content ?? result.content, details: afterResult.details ?? result.details, terminate: afterResult.terminate ?? result.terminate, @@ -1112,7 +1113,7 @@ function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResul role: "toolResult", toolCallId: finalized.toolCall.id, toolName: finalized.toolCall.name, - content: finalized.result.content, + content: finalized.result.content ?? [], details: finalized.result.details, isError: finalized.isError, timestamp: Date.now(), diff --git a/packages/agent-core/src/agent.ts b/packages/agent-core/src/agent.ts index e8d32d163e09..93422b519083 100644 --- a/packages/agent-core/src/agent.ts +++ b/packages/agent-core/src/agent.ts @@ -29,6 +29,7 @@ import type { AgentTool, BeforeToolCallContext, BeforeToolCallResult, + PrepareNextTurnContext, QueueMode, StreamFn, ToolExecutionMode, @@ -133,6 +134,11 @@ export interface AgentOptions { prepareNextTurn?: ( signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; + /** Context-aware turn hook. Takes precedence over `prepareNextTurn` when both are provided. */ + prepareNextTurnWithContext?: ( + context: PrepareNextTurnContext, + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; /** Queue drain mode for steering messages injected before the next assistant response. */ steeringMode?: QueueMode; /** Queue drain mode for follow-up messages injected after the agent would otherwise stop. */ @@ -228,6 +234,10 @@ export class Agent { public prepareNextTurn?: ( signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; + public prepareNextTurnWithContext?: ( + context: PrepareNextTurnContext, + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; private activeRun?: ActiveRun; /** Session identifier forwarded to providers for cache-aware backends. */ public sessionId?: string; @@ -253,6 +263,7 @@ export class Agent { this.resolveDeferredTool = options.resolveDeferredTool; this.afterToolCall = options.afterToolCall; this.prepareNextTurn = options.prepareNextTurn; + this.prepareNextTurnWithContext = options.prepareNextTurnWithContext; this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); this.sessionId = options.sessionId; @@ -394,7 +405,7 @@ export class Agent { throw new Error("No messages to continue from"); } - if (lastMessage.role === "assistant") { + if (lastMessage.role === "assistant" || lastMessage.role === "toolResult") { const queuedSteering = this.steeringQueue.drain(); if (queuedSteering.length > 0) { await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true }); @@ -406,7 +417,9 @@ export class Agent { await this.runPromptMessages(queuedFollowUps); return; } + } + if (lastMessage.role === "assistant") { throw new TranscriptNotContinuableError(lastMessage.role); } @@ -487,9 +500,15 @@ export class Agent { beforeToolCall: this.beforeToolCall, resolveDeferredTool: this.resolveDeferredTool, afterToolCall: this.afterToolCall, - prepareNextTurn: this.prepareNextTurn - ? async () => await this.prepareNextTurn?.(this.signal) - : undefined, + prepareNextTurn: + this.prepareNextTurnWithContext || this.prepareNextTurn + ? async (context) => { + if (this.prepareNextTurnWithContext) { + return await this.prepareNextTurnWithContext(context, this.signal); + } + return await this.prepareNextTurn?.(this.signal); + } + : undefined, convertToLlm: this.convertToLlm, transformContext: this.transformContext, getApiKey: this.getApiKey, diff --git a/packages/agent-core/src/harness/compaction/compaction.test.ts b/packages/agent-core/src/harness/compaction/compaction.test.ts index e6ee99f55117..70f15e58c8dd 100644 --- a/packages/agent-core/src/harness/compaction/compaction.test.ts +++ b/packages/agent-core/src/harness/compaction/compaction.test.ts @@ -1,14 +1,69 @@ import { describe, expect, it, vi } from "vitest"; import { createAssistantMessageEventStream } from "../../llm.js"; -import type { AssistantMessage, Model, StreamFn } from "../../llm.js"; +import type { AssistantMessage, Model, StreamFn, Usage } from "../../llm.js"; +import type { AgentMessage } from "../../types.js"; +import type { SessionTreeEntry } from "../types.js"; import { calculateContextTokens, compact, estimateContextTokens, + findCutPoint, generateSummary, + getLastAssistantUsage, + prepareCompaction, } from "./compaction.js"; import { createFileOps } from "./utils.js"; +function createUsage(totalTokens: number): Usage { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + contextUsage: { state: "available", promptTokens: totalTokens, totalTokens }, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createAssistant(text: string, usage: Usage, timestamp: number): AssistantMessage { + return { + role: "assistant", + content: [{ type: "text", text }], + api: "test-api", + provider: "test-provider", + model: "test-model", + usage, + stopReason: "stop", + timestamp, + }; +} + +function createMessageEntry(message: AgentMessage, index: number): SessionTreeEntry { + return { + type: "message", + id: `entry-${index}`, + parentId: index === 0 ? null : `entry-${index - 1}`, + timestamp: new Date(message.timestamp).toISOString(), + message, + }; +} + +function createProjectedEntry( + type: "custom_message" | "branch_summary", + index: number, + content: string, +): SessionTreeEntry { + const common = { + id: `entry-${index}`, + parentId: index === 0 ? null : `entry-${index - 1}`, + timestamp: new Date(index + 1).toISOString(), + }; + return type === "custom_message" + ? { ...common, type, customType: "test", content, display: true } + : { ...common, type, fromId: common.parentId ?? common.id, summary: content }; +} + describe("calculateContextTokens", () => { it("prefers the final-iteration context snapshot over aggregate billing usage", () => { expect( @@ -121,6 +176,82 @@ describe("calculateContextTokens", () => { expect(estimate.tokens).toBeLessThan(927_907); expect(estimate.lastUsageIndex).toBe(0); }); + + it("ignores an all-zero terminal usage block", () => { + const validUsage = createUsage(20); + const messages: AgentMessage[] = [ + createAssistant("complete", validUsage, 1), + { role: "user", content: "continue", timestamp: 2 }, + createAssistant("partial", createUsage(0), 3), + ]; + const entries = messages.map(createMessageEntry); + + expect(getLastAssistantUsage(entries)).toBe(validUsage); + expect(estimateContextTokens(messages)).toMatchObject({ + usageTokens: 20, + lastUsageIndex: 0, + }); + expect(estimateContextTokens(messages).trailingTokens).toBeGreaterThan(0); + }); +}); + +describe("session-entry compaction budgeting", () => { + it.each(["custom_message", "branch_summary"] as const)( + "counts a %s entry that projects into model context", + (entryType) => { + const entries: SessionTreeEntry[] = [ + createMessageEntry({ role: "user", content: "hi", timestamp: 1 }, 0), + createMessageEntry(createAssistant("hello", createUsage(2), 2), 1), + createProjectedEntry(entryType, 2, "x".repeat(4_000)), + createMessageEntry(createAssistant("ok", createUsage(2), 4), 3), + ]; + + expect(findCutPoint(entries, 0, entries.length, 1)).toMatchObject({ + firstKeptEntryIndex: 3, + turnStartIndex: 2, + isSplitTurn: true, + }); + expect(findCutPoint(entries, 0, entries.length, 2)).toEqual({ + firstKeptEntryIndex: 2, + turnStartIndex: -1, + isSplitTurn: false, + }); + }, + ); + + it.each(["custom_message", "branch_summary"] as const)( + "does not rewind across adjacent %s entries", + (entryType) => { + const entries: SessionTreeEntry[] = [ + createMessageEntry({ role: "user", content: "hi", timestamp: 1 }, 0), + createMessageEntry(createAssistant("hello", createUsage(2), 2), 1), + createProjectedEntry(entryType, 2, "x".repeat(4_000)), + createProjectedEntry(entryType, 3, "y".repeat(4_000)), + createMessageEntry(createAssistant("ok", createUsage(2), 5), 4), + ]; + + expect(findCutPoint(entries, 0, entries.length, 2)).toEqual({ + firstKeptEntryIndex: 3, + turnStartIndex: -1, + isSplitTurn: false, + }); + }, + ); + + it("skips compaction when no history or turn prefix would be summarized", () => { + const entries = [ + createMessageEntry({ role: "user", content: "hello", timestamp: 1 }, 0), + createMessageEntry(createAssistant("done", createUsage(2), 2), 1), + ]; + + expect( + prepareCompaction(entries, { + enabled: true, + reserveTokens: 0, + keepRecentTokens: 10_000, + }), + ).toEqual({ ok: true, value: undefined }); + }); }); describe("generateSummary thinking options", () => { diff --git a/packages/agent-core/src/harness/compaction/compaction.ts b/packages/agent-core/src/harness/compaction/compaction.ts index f48ace5bea27..8a5f643c76e9 100644 --- a/packages/agent-core/src/harness/compaction/compaction.ts +++ b/packages/agent-core/src/harness/compaction/compaction.ts @@ -159,7 +159,8 @@ function getAssistantUsage(msg: AgentMessage): Usage | undefined { if ( assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && - assistantMsg.usage + assistantMsg.usage && + calculateContextTokens(assistantMsg.usage) > 0 ) { return assistantMsg.usage; } @@ -167,7 +168,7 @@ function getAssistantUsage(msg: AgentMessage): Usage | undefined { return undefined; } -/** Return usage from the last successful assistant message in session entries. */ +/** Return usage from the last valid assistant message in session entries. */ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined { for (const entry of entries.toReversed()) { if (entry.type === "message") { @@ -319,6 +320,43 @@ export function estimateTokens(message: AgentMessage): number { return 0; } +function isCutPointMessage(message: AgentMessage): boolean { + switch (message.role) { + case "user": + case "assistant": + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + return true; + case "toolResult": + return false; + } + + return false; +} + +function isTurnStartMessage(message: AgentMessage): boolean { + switch (message.role) { + case "user": + case "bashExecution": + case "custom": + case "branchSummary": + case "compactionSummary": + return true; + case "assistant": + case "toolResult": + return false; + } + + return false; +} + +function isTurnStartEntry(entry: SessionTreeEntry): boolean { + const message = getMessageFromEntryForCompaction(entry); + return message ? isTurnStartMessage(message) : false; +} + function findValidCutPoints( entries: SessionTreeEntry[], startIndex: number, @@ -330,35 +368,8 @@ function findValidCutPoints( if (!entry) { continue; } - switch (entry.type) { - case "message": { - const role = (entry.message as HarnessMessage).role; - switch (role) { - case "bashExecution": - case "custom": - case "branchSummary": - case "compactionSummary": - case "user": - case "assistant": - cutPoints.push(i); - break; - case "toolResult": - break; - } - break; - } - case "thinking_level_change": - case "model_change": - case "compaction": - case "branch_summary": - case "custom": - case "custom_message": - case "label": - case "session_info": - case "leaf": - break; - } - if (entry.type === "branch_summary" || entry.type === "custom_message") { + const message = getMessageFromEntryForCompaction(entry); + if (message && isCutPointMessage(message)) { cutPoints.push(i); } } @@ -376,15 +387,9 @@ export function findTurnStartIndex( if (!entry) { continue; } - if (entry.type === "branch_summary" || entry.type === "custom_message") { + if (isTurnStartEntry(entry)) { return i; } - if (entry.type === "message") { - const role = (entry.message as HarnessMessage).role; - if (role === "user" || role === "bashExecution") { - return i; - } - } } return -1; } @@ -420,10 +425,14 @@ export function findCutPoint( for (let i = endIndex - 1; i >= startIndex; i--) { const entry = entries[i]; - if (!entry || entry.type !== "message") { + if (!entry) { continue; } - const messageTokens = estimateTokens(entry.message); + const message = getMessageFromEntryForCompaction(entry); + if (!message) { + continue; + } + const messageTokens = estimateTokens(message); accumulatedTokens += messageTokens; if (accumulatedTokens >= keepRecentTokens) { const lastCutIndex = cutPoints.at(-1); @@ -448,7 +457,7 @@ export function findCutPoint( if (prevEntry.type === "compaction") { break; } - if (prevEntry.type === "message") { + if (getMessageFromEntryForCompaction(prevEntry)) { break; } cutIndex--; @@ -457,13 +466,13 @@ export function findCutPoint( if (!cutEntry) { throw new Error("compaction cut point does not reference a session entry"); } - const isUserMessage = cutEntry.type === "message" && cutEntry.message.role === "user"; - const turnStartIndex = isUserMessage ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); + const startsTurn = isTurnStartEntry(cutEntry); + const turnStartIndex = startsTurn ? -1 : findTurnStartIndex(entries, cutIndex, startIndex); return { firstKeptEntryIndex: cutIndex, turnStartIndex, - isSplitTurn: !isUserMessage && turnStartIndex !== -1, + isSplitTurn: !startsTurn && turnStartIndex !== -1, }; } @@ -757,6 +766,9 @@ export function prepareCompaction( } } } + if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) { + return ok(undefined); + } const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); if (cutPoint.isSplitTurn) { for (const msg of turnPrefixMessages) { diff --git a/src/agents/sessions/agent-session-base.ts b/src/agents/sessions/agent-session-base.ts index a678ac39e5a1..a16b00e7b98e 100644 --- a/src/agents/sessions/agent-session-base.ts +++ b/src/agents/sessions/agent-session-base.ts @@ -125,6 +125,7 @@ export abstract class AgentSessionBase { protected baseSystemPrompt = ""; protected baseSystemPromptOptions!: BuildSystemPromptOptions; protected exactBaseSystemPrompt: string | undefined; + protected systemPromptOverride: string | undefined; constructor(config: AgentSessionConfig) { this.agent = config.agent; @@ -285,9 +286,18 @@ export abstract class AgentSessionBase { // Track last assistant message for auto-compaction check protected lastAssistantMessage: AssistantMessage | undefined = undefined; + protected lastRunEndedForTurnHandoff = false; /** Internal handler for agent events - shared by subscribe and reconnect */ - protected handleAgentEvent = async (event: AgentEvent): Promise => { + protected handleAgentEvent = async (event: AgentEvent, signal?: AbortSignal): Promise => { + if (event.type === "agent_end") { + const reason: unknown = signal?.reason; + this.lastRunEndedForTurnHandoff = + signal?.aborted === true && + typeof reason === "object" && + reason !== null && + (reason as { turnHandoff?: unknown }).turnHandoff === true; + } if (this.eventMayWriteSession(event)) { await this.runWithSessionWriteLock(async () => await this.handleAgentEventUnlocked(event)); return; @@ -360,7 +370,9 @@ export abstract class AgentSessionBase { this.lastAssistantMessage = event.message; const assistantMsg = event.message; - if (assistantMsg.stopReason !== "error") { + // A length response may still need overflow recovery in checkCompaction(); + // retryCount is independent and resets for every non-error response below. + if (assistantMsg.stopReason !== "error" && assistantMsg.stopReason !== "length") { this.overflowRecoveryAttempted = false; } @@ -540,6 +552,21 @@ export abstract class AgentSessionBase { * Call this when completely done with the session. */ dispose(): void { + const abortOperations = [ + () => this.abortRetry(), + () => this.abortCompaction(), + () => this.abortBranchSummary(), + () => this.abortBash(), + () => this.agent.abort(), + ]; + for (const abortOperation of abortOperations) { + try { + abortOperation(); + } catch { + // One broken abort hook must not prevent the remaining work from being cancelled. + } + } + this.currentExtensionRunner.invalidate( "This extension ctx is stale after session replacement or reload. Do not use a captured api or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().", ); @@ -626,7 +653,7 @@ export abstract class AgentSessionBase { // Rebuild base system prompt with new tool set this.baseSystemPrompt = this.rebuildSystemPrompt(validToolNames); - this.agent.state.systemPrompt = this.baseSystemPrompt; + this.agent.state.systemPrompt = this.systemPromptOverride ?? this.baseSystemPrompt; } /** Set an exact base prompt owned by the current runtime. */ @@ -788,5 +815,8 @@ export abstract class AgentSessionBase { skipAbortedCheck?: boolean, ): Promise; abstract abortRetry(): void; + abstract abortCompaction(): void; + abstract abortBranchSummary(): void; + abstract abortBash(): void; protected abstract flushPendingBashMessages(): void; } diff --git a/src/agents/sessions/agent-session-compaction.ts b/src/agents/sessions/agent-session-compaction.ts index 431174033347..bd352fd69401 100644 --- a/src/agents/sessions/agent-session-compaction.ts +++ b/src/agents/sessions/agent-session-compaction.ts @@ -141,7 +141,14 @@ export abstract class AgentSessionCompaction extends AgentSessionInspection { } const pathEntries = this.sessionManager.getBranch(); - const preparation = unwrapCoreResult(prepareCompaction(pathEntries, options.settings)); + let preparation = unwrapCoreResult(prepareCompaction(pathEntries, options.settings)); + if (!preparation && isManual) { + // An explicit request should compact the smallest valid history instead of + // relying on the old empty-summary behavior when the whole session fits the keep budget. + preparation = unwrapCoreResult( + prepareCompaction(pathEntries, { ...options.settings, keepRecentTokens: 0 }), + ); + } if (!preparation) { if (isManual) { const lastEntry = pathEntries[pathEntries.length - 1]; @@ -265,8 +272,13 @@ export abstract class AgentSessionCompaction extends AgentSessionInspection { return false; } - // Case 1: Overflow - LLM returned context overflow error - if (sameModel && isContextOverflow(assistantMessage, contextWindow)) { + // Case 1: Overflow - an unsuccessful response needs compact-and-retry recovery. + // Successful high-usage responses fall through to threshold maintenance below. + if ( + sameModel && + (assistantMessage.stopReason === "error" || assistantMessage.stopReason === "length") && + isContextOverflow(assistantMessage, contextWindow) + ) { if (this.overflowRecoveryAttempted) { this.emit({ type: "compaction_end", @@ -281,8 +293,7 @@ export abstract class AgentSessionCompaction extends AgentSessionInspection { } this.overflowRecoveryAttempted = true; - // Remove the error message from agent state (it IS saved to session for history, - // but we don't want it in context for the retry) + // Keep the failed response in history, but exclude it from the retry context. const messages = this.agent.state.messages; if (messages.at(-1)?.role === "assistant") { this.agent.state.messages = messages.slice(0, -1); @@ -376,7 +387,10 @@ export abstract class AgentSessionCompaction extends AgentSessionInspection { if (willRetry) { const messages = this.agent.state.messages; const lastMsg = messages[messages.length - 1]; - if (lastMsg?.role === "assistant" && lastMsg.stopReason === "error") { + if ( + lastMsg?.role === "assistant" && + (lastMsg.stopReason === "error" || lastMsg.stopReason === "length") + ) { this.agent.state.messages = messages.slice(0, -1); } return true; diff --git a/src/agents/sessions/agent-session-extensions.ts b/src/agents/sessions/agent-session-extensions.ts index 7e91aaaba725..2a1b67685dbd 100644 --- a/src/agents/sessions/agent-session-extensions.ts +++ b/src/agents/sessions/agent-session-extensions.ts @@ -399,6 +399,8 @@ export abstract class AgentSessionExtensions extends AgentSessionCompaction { reason: "reload", }); await this.settingsManager.reload(); + this.agent.steeringMode = this.settingsManager.getSteeringMode(); + this.agent.followUpMode = this.settingsManager.getFollowUpMode(); resetApiProviders(defaultApiRegistry); await this.sessionResourceLoader.reload(); this.buildRuntime({ diff --git a/src/agents/sessions/agent-session-inspection.ts b/src/agents/sessions/agent-session-inspection.ts index fb2bd3729ab4..81123fb0fa4a 100644 --- a/src/agents/sessions/agent-session-inspection.ts +++ b/src/agents/sessions/agent-session-inspection.ts @@ -107,8 +107,8 @@ export abstract class AgentSessionInspection extends AgentSessionModels { if (contextTokens > 0) { hasPostCompactionUsage = true; estimateFromContent = false; + break; } - break; } } } diff --git a/src/agents/sessions/agent-session-loop-correctness.test.ts b/src/agents/sessions/agent-session-loop-correctness.test.ts new file mode 100644 index 000000000000..996c444a8564 --- /dev/null +++ b/src/agents/sessions/agent-session-loop-correctness.test.ts @@ -0,0 +1,669 @@ +import { + createAssistantMessageEventStream, + type AssistantMessage, + type Context, + type Model, + type SimpleStreamOptions, +} from "openclaw/plugin-sdk/llm"; +import { Type } from "typebox"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +const streamMocks = vi.hoisted(() => ({ + streamSimple: vi.fn(), +})); + +vi.mock("../../llm/stream.js", () => ({ + streamSimple: streamMocks.streamSimple, +})); + +import type { AgentTool } from "../runtime/index.js"; +import type { AgentSessionEvent } from "./agent-session-types.js"; +import { AgentSession } from "./agent-session.js"; +import { AuthStorage } from "./auth-storage.js"; +import { createExtensionRuntime } from "./extensions/loader.js"; +import type { LoadExtensionsResult, ToolDefinition } from "./extensions/types.js"; +import { ModelRegistry } from "./model-registry.js"; +import type { ResourceLoader } from "./resource-loader.js"; +import { createAgentSession } from "./sdk.js"; +import { SessionManager } from "./session-manager.js"; +import { SettingsManager } from "./settings-manager.js"; +import { createSyntheticSourceInfo } from "./source-info.js"; + +const testModel: Model = { + id: "test-model", + name: "Test Model", + api: "openai-responses", + provider: "test-provider", + baseUrl: "https://example.test", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 100, + maxTokens: 100, +}; + +const sessions: AgentSession[] = []; + +function createUsage(contextTokens: number) { + return { + input: contextTokens, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: contextTokens + 1, + contextUsage: { + state: "available" as const, + promptTokens: contextTokens, + totalTokens: contextTokens + 1, + }, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createAssistant( + activeModel: Model, + content: AssistantMessage["content"], + stopReason: AssistantMessage["stopReason"] = "stop", + contextTokens = 1, +): AssistantMessage { + return { + role: "assistant", + content, + api: activeModel.api, + provider: activeModel.provider, + model: activeModel.id, + usage: createUsage(contextTokens), + stopReason, + timestamp: Date.now(), + }; +} + +function createAssistantResultStream(message: AssistantMessage) { + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + if (message.stopReason === "error" || message.stopReason === "aborted") { + stream.push({ type: "error", reason: message.stopReason, error: message }); + } else { + stream.push({ type: "done", reason: message.stopReason, message }); + } + stream.end(); + }); + return stream; +} + +function createResourceLoader( + handlers: Map Promise>> = new Map(), +): ResourceLoader { + const extensionsResult: LoadExtensionsResult = { + extensions: + handlers.size > 0 + ? [ + { + path: "", + resolvedPath: "", + sourceInfo: createSyntheticSourceInfo("", { + source: "temporary", + }), + handlers, + tools: new Map(), + messageRenderers: new Map(), + commands: new Map(), + flags: new Map(), + shortcuts: new Map(), + }, + ] + : [], + errors: [], + runtime: createExtensionRuntime(), + }; + return { + getExtensions: () => extensionsResult, + getSkills: () => ({ skills: [], diagnostics: [] }), + getPrompts: () => ({ prompts: [], diagnostics: [] }), + getThemes: () => ({ themes: [], diagnostics: [] }), + getAgentsFiles: () => ({ agentsFiles: [] }), + getSystemPrompt: () => undefined, + getAppendSystemPrompt: () => [], + extendResources: () => {}, + reload: async () => {}, + }; +} + +function createCompactionHandlers() { + return new Map Promise>>([ + [ + "session_before_compact", + [ + async (event: unknown) => { + const preparation = ( + event as { + preparation: { firstKeptEntryId: string; tokensBefore: number }; + } + ).preparation; + return { + compaction: { + summary: "condensed history", + firstKeptEntryId: preparation.firstKeptEntryId, + tokensBefore: preparation.tokensBefore, + }, + }; + }, + ], + ], + ]); +} + +async function createTestSession( + options: { + model?: Model; + settingsManager?: SettingsManager; + sessionManager?: SessionManager; + resourceLoader?: ResourceLoader; + customTools?: ToolDefinition[]; + } = {}, +) { + const model = options.model ?? testModel; + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(model.provider, "test-api-key"); + const settingsManager = + options.settingsManager ?? + SettingsManager.inMemory({ + compaction: { enabled: false }, + retry: { enabled: false }, + }); + const sessionManager = options.sessionManager ?? SessionManager.inMemory(); + const result = await createAgentSession({ + model, + noTools: "builtin", + customTools: options.customTools, + resourceLoader: options.resourceLoader ?? createResourceLoader(), + sessionManager, + settingsManager, + modelRegistry: ModelRegistry.inMemory(authStorage), + }); + sessions.push(result.session); + return { ...result, settingsManager, sessionManager }; +} + +function appendHistory(sessionManager: SessionManager, assistant: AssistantMessage): void { + sessionManager.appendMessage({ role: "user", content: "old prompt", timestamp: Date.now() - 2 }); + sessionManager.appendMessage({ ...assistant, timestamp: Date.now() - 1 }); +} + +beforeEach(() => { + streamMocks.streamSimple.mockReset(); +}); + +afterEach(() => { + for (const session of sessions.splice(0)) { + session.dispose(); + } +}); + +describe("AgentSession loop correctness", () => { + it("manually compacts a completed turn smaller than the retained-token budget", async () => { + const sessionManager = SessionManager.inMemory(); + appendHistory( + sessionManager, + createAssistant(testModel, [{ type: "text", text: "short answer" }]), + ); + const settingsManager = SettingsManager.inMemory({ + compaction: { enabled: true, reserveTokens: 0, keepRecentTokens: 10_000 }, + retry: { enabled: false }, + }); + const { session } = await createTestSession({ + sessionManager, + settingsManager, + resourceLoader: createResourceLoader(createCompactionHandlers()), + }); + + const result = await session.compact(); + + expect(result.summary).toBe("condensed history"); + expect(sessionManager.getBranch().at(-1)).toMatchObject({ + type: "compaction", + summary: "condensed history", + }); + }); + + it("keeps a successful high-usage response and performs threshold maintenance without retry", async () => { + const settingsManager = SettingsManager.inMemory({ + compaction: { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }, + retry: { enabled: false }, + }); + const compactionEvents: AgentSessionEvent[] = []; + streamMocks.streamSimple.mockImplementation((activeModel: Model) => + createAssistantResultStream( + createAssistant(activeModel, [{ type: "text", text: "complete answer" }], "stop", 100), + ), + ); + const { session } = await createTestSession({ + settingsManager, + resourceLoader: createResourceLoader(createCompactionHandlers()), + }); + session.subscribe((event) => { + if (event.type === "compaction_end") { + compactionEvents.push(event); + } + }); + + await session.prompt("new prompt"); + + expect(streamMocks.streamSimple).toHaveBeenCalledOnce(); + expect(session.messages).toContainEqual( + expect.objectContaining({ + role: "assistant", + content: [{ type: "text", text: "complete answer" }], + }), + ); + expect(compactionEvents).toContainEqual( + expect.objectContaining({ type: "compaction_end", reason: "threshold", willRetry: false }), + ); + }); + + it("does not retry a high-usage turn terminated by a tool result", async () => { + const terminalTool: ToolDefinition = { + name: "finish", + label: "Finish", + description: "finishes the current run", + parameters: Type.Object({}), + execute: async () => ({ + content: [{ type: "text", text: "finished" }], + details: {}, + terminate: true, + }), + }; + const settingsManager = SettingsManager.inMemory({ + compaction: { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }, + retry: { enabled: false }, + }); + const compactionEvents: AgentSessionEvent[] = []; + streamMocks.streamSimple.mockImplementation((activeModel: Model) => + createAssistantResultStream( + createAssistant( + activeModel, + [{ type: "toolCall", id: "call-finish", name: "finish", arguments: {} }], + "toolUse", + 100, + ), + ), + ); + const { session } = await createTestSession({ + settingsManager, + resourceLoader: createResourceLoader(createCompactionHandlers()), + customTools: [terminalTool], + }); + session.subscribe((event) => { + if (event.type === "compaction_end") { + compactionEvents.push(event); + } + }); + + await session.prompt("finish now"); + + expect(streamMocks.streamSimple).toHaveBeenCalledOnce(); + expect(compactionEvents).toContainEqual( + expect.objectContaining({ type: "compaction_end", reason: "threshold", willRetry: false }), + ); + }); + + it("compacts and retries a high-usage length-truncated response", async () => { + const settingsManager = SettingsManager.inMemory({ + compaction: { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }, + retry: { enabled: false }, + }); + const compactionEvents: AgentSessionEvent[] = []; + let requestCount = 0; + streamMocks.streamSimple.mockImplementation((activeModel: Model) => { + requestCount += 1; + return createAssistantResultStream( + requestCount === 1 + ? { + ...createAssistant( + activeModel, + [{ type: "text", text: "truncated answer" }], + "length", + 100, + ), + usage: { ...createUsage(100), output: 0 }, + } + : createAssistant(activeModel, [{ type: "text", text: "complete retry" }]), + ); + }); + const { session } = await createTestSession({ + settingsManager, + resourceLoader: createResourceLoader(createCompactionHandlers()), + }); + session.subscribe((event) => { + if (event.type === "compaction_end") { + compactionEvents.push(event); + } + }); + + await session.prompt("long request"); + + expect(streamMocks.streamSimple).toHaveBeenCalledTimes(2); + expect(compactionEvents).toContainEqual( + expect.objectContaining({ type: "compaction_end", reason: "overflow", willRetry: true }), + ); + expect(session.getLastAssistantText()).toBe("complete retry"); + }); + + it("delivers a pending prompt immediately after pre-prompt compaction", async () => { + const sessionManager = SessionManager.inMemory(); + appendHistory( + sessionManager, + createAssistant(testModel, [{ type: "text", text: "old answer" }], "stop", 100), + ); + const settingsManager = SettingsManager.inMemory({ + compaction: { enabled: true, reserveTokens: 0, keepRecentTokens: 1 }, + retry: { enabled: false }, + }); + const requests: Context[] = []; + streamMocks.streamSimple.mockImplementation((activeModel: Model, context: Context) => { + requests.push(context); + return createAssistantResultStream( + createAssistant(activeModel, [{ type: "text", text: "new answer" }]), + ); + }); + const { session } = await createTestSession({ + sessionManager, + settingsManager, + resourceLoader: createResourceLoader(createCompactionHandlers()), + }); + const continueRun = vi.spyOn(session.agent, "continue"); + + await session.prompt("pending prompt"); + + expect(continueRun).not.toHaveBeenCalled(); + expect(requests).toHaveLength(1); + expect(JSON.stringify(requests[0]?.messages)).toContain("pending prompt"); + }); + + it("drains a follow-up queued by an agent-end handler", async () => { + const sessionRef: { current?: AgentSession } = {}; + let queued = false; + const handlers = new Map Promise>>([ + [ + "agent_end", + [ + async () => { + if (!queued) { + queued = true; + await sessionRef.current?.followUp("queued after end"); + } + return undefined; + }, + ], + ], + ]); + const requests: Context[] = []; + streamMocks.streamSimple.mockImplementation((activeModel: Model, context: Context) => { + requests.push(context); + return createAssistantResultStream( + createAssistant(activeModel, [{ type: "text", text: `answer ${requests.length}` }]), + ); + }); + const { session } = await createTestSession({ resourceLoader: createResourceLoader(handlers) }); + sessionRef.current = session; + + await session.prompt("initial prompt"); + + expect(requests).toHaveLength(2); + expect(JSON.stringify(requests[1]?.messages)).toContain("queued after end"); + expect(session.agent.hasQueuedMessages()).toBe(false); + }); + + it("leaves queued messages dormant after a turn handoff", async () => { + const sessionRef: { current?: AgentSession } = {}; + const yieldTool: ToolDefinition = { + name: "yield_turn", + label: "Yield turn", + description: "ends the current turn for an external handoff", + parameters: Type.Object({}), + execute: async () => { + const activeSession = sessionRef.current; + if (!activeSession) { + throw new Error("session not ready"); + } + activeSession.agent.steer({ + role: "custom", + customType: "test.turn-handoff", + content: "resume only for external delivery", + display: false, + timestamp: Date.now(), + }); + activeSession.agent.abort({ code: "turn_handoff", turnHandoff: true }); + return { content: [{ type: "text", text: "yielded" }], details: { yielded: true } }; + }, + }; + streamMocks.streamSimple.mockImplementation((activeModel: Model) => + createAssistantResultStream( + createAssistant( + activeModel, + [{ type: "toolCall", id: "call-yield", name: "yield_turn", arguments: {} }], + "toolUse", + ), + ), + ); + const { session } = await createTestSession({ customTools: [yieldTool] }); + sessionRef.current = session; + + await session.prompt("yield now"); + + expect(streamMocks.streamSimple).toHaveBeenCalledOnce(); + expect(session.agent.hasQueuedMessages()).toBe(true); + session.agent.clearAllQueues(); + }); + + it("applies session model, tool, and prompt changes on the following tool turn", async () => { + const nextModel = { ...testModel, id: "next-model" }; + const sessionRef: { current?: AgentSession } = {}; + const switchTool: ToolDefinition = { + name: "switch_state", + label: "Switch state", + description: "changes the next turn state", + parameters: Type.Object({}), + execute: async () => { + const activeSession = sessionRef.current; + if (!activeSession) { + throw new Error("session not ready"); + } + activeSession.setActiveToolsByName(["second_tool"]); + activeSession.agent.state.model = nextModel; + return { content: [{ type: "text", text: "switched" }], details: {} }; + }, + }; + const secondTool: ToolDefinition = { + name: "second_tool", + label: "Second tool", + description: "available after the switch", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "done" }], details: {} }), + }; + const handlers = new Map Promise>>([ + ["before_agent_start", [async () => ({ systemPrompt: "prompt override" })]], + ]); + const requests: Array<{ model: string; prompt: string; tools: string[] }> = []; + streamMocks.streamSimple.mockImplementation((activeModel: Model, context: Context) => { + requests.push({ + model: activeModel.id, + prompt: context.systemPrompt ?? "", + tools: context.tools?.map((tool) => tool.name) ?? [], + }); + const content: AssistantMessage["content"] = + requests.length === 1 + ? [{ type: "toolCall", id: "call-switch", name: "switch_state", arguments: {} }] + : [{ type: "text", text: "finished" }]; + return createAssistantResultStream( + createAssistant(activeModel, content, requests.length === 1 ? "toolUse" : "stop"), + ); + }); + const { session } = await createTestSession({ + resourceLoader: createResourceLoader(handlers), + customTools: [switchTool, secondTool], + }); + sessionRef.current = session; + session.setActiveToolsByName(["switch_state"]); + + await session.prompt("switch now"); + + expect(requests).toEqual([ + { model: testModel.id, prompt: "prompt override", tools: ["switch_state"] }, + { model: nextModel.id, prompt: "prompt override", tools: ["second_tool"] }, + ]); + }); + + it("preserves explicit updates from an existing next-turn hook", async () => { + const hookModel = { ...testModel, id: "hook-model" }; + const hookTool: AgentTool = { + name: "hook_tool", + label: "Hook tool", + description: "provided by the existing turn hook", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "done" }], details: {} }), + }; + const hookContext = { + systemPrompt: "hook prompt", + messages: [], + tools: [hookTool], + }; + let returnedUpdate = false; + const { session } = await createTestSession(); + session.agent.prepareNextTurn = () => { + if (returnedUpdate) { + return undefined; + } + returnedUpdate = true; + return { context: hookContext, model: hookModel, thinkingLevel: "high" }; + }; + const contextualHook = session.agent.prepareNextTurnWithContext; + if (!contextualHook) { + throw new Error("context-aware next-turn hook was not installed"); + } + const message = createAssistant(testModel, [{ type: "text", text: "turn complete" }]); + const newMessages = [message]; + + const firstUpdate = await contextualHook({ + message, + toolResults: [], + context: { systemPrompt: "loop prompt", messages: [], tools: [] }, + newMessages, + }); + const secondUpdate = await contextualHook({ + message, + toolResults: [], + context: firstUpdate?.context ?? hookContext, + newMessages, + }); + + for (const update of [firstUpdate, secondUpdate]) { + expect(update).toMatchObject({ + context: { + systemPrompt: "hook prompt", + tools: [expect.objectContaining({ name: "hook_tool" })], + }, + model: hookModel, + thinkingLevel: "high", + }); + } + }); + + it("preserves fields omitted by an existing next-turn context replacement", async () => { + const sessionTool: AgentTool = { + name: "session_tool", + label: "Session tool", + description: "available in session state", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "done" }], details: {} }), + }; + const initialHook = vi.fn(() => ({ + context: { systemPrompt: "stale prompt", messages: [], tools: [sessionTool] }, + })); + const replacementHook = vi.fn(() => ({ + context: { systemPrompt: "replacement prompt", messages: [] }, + })); + const { session } = await createTestSession({ customTools: [sessionTool] }); + session.setActiveToolsByName([sessionTool.name]); + session.agent.prepareNextTurn = initialHook; + session.agent.prepareNextTurn = replacementHook; + const message = createAssistant(testModel, [{ type: "text", text: "turn complete" }]); + const contextualHook = session.agent.prepareNextTurnWithContext; + if (!contextualHook) { + throw new Error("context-aware next-turn hook was not installed"); + } + + const update = await contextualHook({ + message, + toolResults: [], + context: { systemPrompt: "loop prompt", messages: [], tools: [sessionTool] }, + newMessages: [message], + }); + + expect(update?.context).toEqual({ systemPrompt: "replacement prompt", messages: [] }); + expect(replacementHook).toHaveBeenCalledOnce(); + expect(initialHook).not.toHaveBeenCalled(); + }); + + it("aborts in-flight work when disposed", async () => { + let providerSignal: AbortSignal | undefined; + streamMocks.streamSimple.mockImplementation( + (activeModel: Model, _context: Context, options?: SimpleStreamOptions) => { + providerSignal = options?.signal; + const stream = createAssistantMessageEventStream(); + options?.signal?.addEventListener( + "abort", + () => { + const message = createAssistant(activeModel, [], "aborted"); + stream.push({ type: "error", reason: "aborted", error: message }); + stream.end(); + }, + { once: true }, + ); + return stream; + }, + ); + const { session } = await createTestSession(); + const abortRetry = vi.spyOn(session, "abortRetry"); + const abortCompaction = vi.spyOn(session, "abortCompaction"); + const abortBranchSummary = vi.spyOn(session, "abortBranchSummary"); + const abortBash = vi.spyOn(session, "abortBash"); + const abortAgent = vi.spyOn(session.agent, "abort"); + abortRetry.mockImplementationOnce(() => { + throw new Error("retry abort failed"); + }); + const prompt = session.prompt("wait"); + await vi.waitFor(() => expect(providerSignal).toBeDefined()); + + session.dispose(); + await prompt; + + expect(providerSignal?.aborted).toBe(true); + expect(abortRetry).toHaveBeenCalledOnce(); + expect(abortCompaction).toHaveBeenCalledOnce(); + expect(abortBranchSummary).toHaveBeenCalledOnce(); + expect(abortBash).toHaveBeenCalledOnce(); + expect(abortAgent).toHaveBeenCalledOnce(); + }); + + it("resynchronizes queue modes when settings reload", async () => { + const settingsManager = SettingsManager.inMemory({ + steeringMode: "one-at-a-time", + followUpMode: "one-at-a-time", + compaction: { enabled: false }, + retry: { enabled: false }, + }); + const { session } = await createTestSession({ settingsManager }); + settingsManager.setSteeringMode("all"); + settingsManager.setFollowUpMode("all"); + await settingsManager.flush(); + + expect(session.agent.steeringMode).toBe("one-at-a-time"); + expect(session.agent.followUpMode).toBe("one-at-a-time"); + + await session.reload(); + + expect(session.agent.steeringMode).toBe("all"); + expect(session.agent.followUpMode).toBe("all"); + }); +}); diff --git a/src/agents/sessions/agent-session-prompting.ts b/src/agents/sessions/agent-session-prompting.ts index ea44ddf4755b..d34232b2f31d 100644 --- a/src/agents/sessions/agent-session-prompting.ts +++ b/src/agents/sessions/agent-session-prompting.ts @@ -26,6 +26,7 @@ export abstract class AgentSessionPrompting extends AgentSessionBase { await this.agent.continue(); } } finally { + this.systemPromptOverride = undefined; this.flushPendingBashMessages(); } } @@ -33,7 +34,10 @@ export abstract class AgentSessionPrompting extends AgentSessionBase { private async handlePostAgentRun(): Promise { const msg = this.lastAssistantMessage; this.lastAssistantMessage = undefined; - if (!msg) { + const endedForTurnHandoff = this.lastRunEndedForTurnHandoff; + this.lastRunEndedForTurnHandoff = false; + if (!msg || endedForTurnHandoff) { + // External delivery owns the next run after a deliberate turn handoff. return false; } @@ -51,7 +55,12 @@ export abstract class AgentSessionPrompting extends AgentSessionBase { this.retryCount = 0; } - return await this.checkCompaction(msg); + if (await this.checkCompaction(msg)) { + return true; + } + + // Messages queued by agent_end handlers arrive after the loop's final queue drain. + return this.agent.hasQueuedMessages(); } private createUserContent( @@ -149,17 +158,11 @@ export abstract class AgentSessionPrompting extends AgentSessionBase { throw new Error(formatNoApiKeyFoundMessage(this.model.provider)); } - // Check if we need to compact before sending (catches aborted responses) + // Check if we need to compact before sending (catches aborted responses). + // The pending user prompt below starts the next run; no intermediate continuation is needed. const lastAssistant = this.findLastAssistantMessage(); - if (lastAssistant && (await this.checkCompaction(lastAssistant, false))) { - try { - await this.agent.continue(); - while (await this.handlePostAgentRun()) { - await this.agent.continue(); - } - } finally { - this.flushPendingBashMessages(); - } + if (lastAssistant) { + await this.checkCompaction(lastAssistant, false); } // Build messages array (custom message if any, then user message) @@ -199,10 +202,12 @@ export abstract class AgentSessionPrompting extends AgentSessionBase { } } // Apply extension-modified system prompt, or reset to base - if (result?.systemPrompt) { + if (result?.systemPrompt !== undefined) { + this.systemPromptOverride = result.systemPrompt; this.agent.state.systemPrompt = result.systemPrompt; } else { // Ensure we're using the base prompt (in case previous turn had modifications) + this.systemPromptOverride = undefined; this.agent.state.systemPrompt = this.baseSystemPrompt; } } catch (error) { diff --git a/src/agents/sessions/agent-session.context-usage.test.ts b/src/agents/sessions/agent-session.context-usage.test.ts index 969eb2018cf6..19391d35a013 100644 --- a/src/agents/sessions/agent-session.context-usage.test.ts +++ b/src/agents/sessions/agent-session.context-usage.test.ts @@ -200,4 +200,77 @@ describe("AgentSession context usage", () => { expect(usage?.tokens).toBeGreaterThan(190_000); }); + + it("preserves an earlier exact post-compaction snapshot before zero usage", () => { + const exactUsage = { + input: 180_000, + output: 10_000, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 190_000, + contextUsage: { + state: "available" as const, + promptTokens: 180_000, + totalTokens: 190_000, + }, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + const zeroUsage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + contextUsage: { state: "available" as const, promptTokens: 0, totalTokens: 0 }, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + const messages = [ + { + role: "assistant", + content: [{ type: "text", text: "exact post-compaction answer" }], + stopReason: "stop", + usage: exactUsage, + }, + { role: "user", content: "small follow-up" }, + { + role: "assistant", + content: [], + stopReason: "stop", + usage: zeroUsage, + }, + ] as unknown as AgentMessage[]; + const branchEntries = [ + { + type: "compaction", + id: "compact-1", + parentId: null, + timestamp: "2026-07-05T00:00:00.000Z", + summary: "summary", + firstKeptEntryId: "assistant-exact", + tokensBefore: 120_000, + }, + { + type: "message", + id: "assistant-exact", + parentId: "compact-1", + timestamp: "2026-07-05T00:00:01.000Z", + message: messages[0], + }, + { + type: "message", + id: "assistant-zero", + parentId: "assistant-exact", + timestamp: "2026-07-05T00:00:02.000Z", + message: messages[2], + }, + ]; + + const usage = AgentSession.prototype.getContextUsage.call({ + model: { contextWindow: 200_000 }, + messages, + sessionManager: { getBranch: () => branchEntries }, + } as unknown as AgentSession); + + expect(usage?.tokens).toBeGreaterThan(190_000); + }); }); diff --git a/src/agents/sessions/agent-session.live.test.ts b/src/agents/sessions/agent-session.live.test.ts new file mode 100644 index 000000000000..0b327e31c64e --- /dev/null +++ b/src/agents/sessions/agent-session.live.test.ts @@ -0,0 +1,263 @@ +// Live end-to-end checks for AgentSession turns, compaction, and follow-up delivery. +import { mkdir, mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { Model } from "openclaw/plugin-sdk/llm"; +import { Type } from "typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import { getRuntimeConfig } from "../../config/config.js"; +import { discoverModels } from "../agent-model-discovery.js"; +import { isLiveTestEnabled } from "../live-test-helpers.js"; +import { ensureOpenClawModelsJson } from "../models-config.js"; +import type { AgentMessage } from "../runtime/index.js"; +import { AgentSession } from "./agent-session.js"; +import { AuthStorage } from "./auth-storage.js"; +import { createExtensionRuntime } from "./extensions/loader.js"; +import type { LoadExtensionsResult, ToolDefinition } from "./extensions/types.js"; +import type { ModelRegistry } from "./model-registry.js"; +import type { ResourceLoader } from "./resource-loader.js"; +import { createAgentSession } from "./sdk.js"; +import { SessionManager } from "./session-manager.js"; +import { SettingsManager } from "./settings-manager.js"; +import { createSyntheticSourceInfo } from "./source-info.js"; + +const API_KEY = process.env.ANTHROPIC_API_KEY?.trim() ?? ""; +const LIVE = isLiveTestEnabled() && API_KEY.length > 0; +const describeLive = LIVE ? describe : describe.skip; +const TEST_TIMEOUT_MS = 120_000; +const PROVIDER_TIMEOUT_MS = 30_000; +const DEFAULT_MODEL_ID = "claude-haiku-4-5"; + +const sessions: AgentSession[] = []; +const tempRoots: string[] = []; + +type ExtensionHandlers = Map Promise>>; + +function createResourceLoader(handlers: ExtensionHandlers = new Map()): ResourceLoader { + const extensionsResult: LoadExtensionsResult = { + extensions: + handlers.size > 0 + ? [ + { + path: "", + resolvedPath: "", + sourceInfo: createSyntheticSourceInfo("", { + source: "temporary", + }), + handlers, + tools: new Map(), + messageRenderers: new Map(), + commands: new Map(), + flags: new Map(), + shortcuts: new Map(), + }, + ] + : [], + errors: [], + runtime: createExtensionRuntime(), + }; + return { + getExtensions: () => extensionsResult, + getSkills: () => ({ skills: [], diagnostics: [] }), + getPrompts: () => ({ prompts: [], diagnostics: [] }), + getThemes: () => ({ themes: [], diagnostics: [] }), + getAgentsFiles: () => ({ agentsFiles: [] }), + getSystemPrompt: () => undefined, + getAppendSystemPrompt: () => [], + extendResources: () => {}, + reload: async () => {}, + }; +} + +async function resolveLiveModel( + agentDir: string, + authStorage: AuthStorage, +): Promise<{ model: Model; modelRegistry: ModelRegistry }> { + await ensureOpenClawModelsJson(getRuntimeConfig(), agentDir, { + providerDiscoveryProviderIds: ["anthropic"], + }); + const modelRegistry = discoverModels(authStorage, agentDir, { providerFilter: "anthropic" }); + const requestedModelId = + process.env.OPENCLAW_LIVE_AGENT_SESSION_MODEL?.trim() || DEFAULT_MODEL_ID; + const model = + modelRegistry.find("anthropic", requestedModelId) ?? + modelRegistry + .getAll() + .find((candidate) => candidate.provider === "anthropic" && /haiku/i.test(candidate.id)); + if (!model) { + throw new Error(`No Anthropic Haiku model found for ${requestedModelId}`); + } + return { + model: { ...model, maxTokens: Math.min(model.maxTokens, 128) }, + modelRegistry, + }; +} + +async function createLiveSession( + options: { + customTools?: ToolDefinition[]; + handlers?: ExtensionHandlers; + } = {}, +) { + const root = await mkdtemp(join(tmpdir(), "openclaw-agent-session-live-")); + tempRoots.push(root); + const cwd = join(root, "workspace"); + const agentDir = join(root, "agent"); + await mkdir(cwd, { recursive: true }); + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey("anthropic", API_KEY); + const { model, modelRegistry } = await resolveLiveModel(agentDir, authStorage); + const sessionManager = SessionManager.inMemory(); + const settingsManager = SettingsManager.inMemory({ + defaultThinkingLevel: "off", + compaction: { enabled: true, reserveTokens: 128, keepRecentTokens: 0 }, + retry: { + enabled: false, + provider: { timeoutMs: PROVIDER_TIMEOUT_MS, maxRetries: 0, maxRetryDelayMs: 0 }, + }, + }); + const { session } = await createAgentSession({ + cwd, + agentDir, + model, + thinkingLevel: "off", + noTools: "builtin", + customTools: options.customTools, + resourceLoader: createResourceLoader(options.handlers), + authStorage, + modelRegistry, + sessionManager, + settingsManager, + }); + sessions.push(session); + return { session, sessionManager }; +} + +function assistantText(message: AgentMessage): string { + if (message.role !== "assistant") { + return ""; + } + return message.content + .filter((block) => block.type === "text") + .map((block) => block.text) + .join("") + .trim(); +} + +afterEach(async () => { + for (const session of sessions.splice(0)) { + session.dispose(); + } + await Promise.all( + tempRoots.splice(0).map(async (root) => { + await rm(root, { recursive: true, force: true }); + }), + ); +}); + +describeLive("AgentSession live", () => { + it( + "completes a real tool turn", + async () => { + let toolExecutions = 0; + const echoParameters = Type.Object({ text: Type.String() }); + const echoTool: ToolDefinition = { + name: "live_echo", + label: "Live echo", + description: + "Return the supplied text unchanged. Use when the user explicitly requests it.", + parameters: echoParameters, + execute: async (_toolCallId, params) => { + toolExecutions += 1; + return { + content: [{ type: "text", text: params.text }], + details: {}, + }; + }, + }; + const { session } = await createLiveSession({ customTools: [echoTool] }); + + await session.prompt( + "Call live_echo exactly once with text OK. After its result, reply with exactly OK.", + ); + + const finalAssistant = session.messages.findLast( + (message) => message.role === "assistant" && message.stopReason === "stop", + ); + expect(toolExecutions).toBe(1); + expect(finalAssistant).toBeDefined(); + if (finalAssistant?.role !== "assistant") { + throw new Error("missing final assistant message"); + } + expect(finalAssistant.usage.output).toBeGreaterThan(0); + expect(assistantText(finalAssistant)).not.toBe(""); + }, + TEST_TIMEOUT_MS, + ); + + it( + "manually compacts a completed turn and remains usable", + async () => { + const { session, sessionManager } = await createLiveSession(); + await session.prompt("Reply with exactly OK."); + + const result = await session.compact(); + const compaction = sessionManager + .getBranch() + .findLast((entry) => entry.type === "compaction"); + + expect(result.summary.trim().length).toBeGreaterThan(0); + expect(compaction?.type).toBe("compaction"); + if (compaction?.type !== "compaction") { + throw new Error("missing compaction entry"); + } + expect(compaction.summary.trim().length).toBeGreaterThan(0); + + await session.prompt("Reply with exactly STILL_OK."); + expect((session.getLastAssistantText() ?? "").trim().length).toBeGreaterThan(0); + }, + TEST_TIMEOUT_MS, + ); + + it( + "drains a follow-up queued by an agent-end handler", + async () => { + const sessionRef: { current?: AgentSession } = {}; + let queued = false; + const handlers: ExtensionHandlers = new Map([ + [ + "agent_end", + [ + async () => { + if (!queued) { + queued = true; + await sessionRef.current?.followUp("Reply with exactly SECOND."); + } + return undefined; + }, + ], + ], + ]); + const { session } = await createLiveSession({ handlers }); + sessionRef.current = session; + + await session.prompt("Reply with exactly FIRST."); + + const assistants = session.messages.filter( + (message) => message.role === "assistant" && message.stopReason === "stop", + ); + const deliveredFollowUp = session.messages.some( + (message) => + message.role === "user" && + Array.isArray(message.content) && + message.content.some( + (block) => block.type === "text" && block.text === "Reply with exactly SECOND.", + ), + ); + expect(deliveredFollowUp).toBe(true); + expect(assistants).toHaveLength(2); + expect(assistants.every((message) => assistantText(message).length > 0)).toBe(true); + }, + TEST_TIMEOUT_MS, + ); +}); diff --git a/src/agents/sessions/sdk.ts b/src/agents/sessions/sdk.ts index 16927b0e2c38..772f99fa81b4 100644 --- a/src/agents/sessions/sdk.ts +++ b/src/agents/sessions/sdk.ts @@ -16,6 +16,7 @@ import { Agent, type AgentMessage, type AgentOptions, + type AgentTool, type ThinkingLevel, } from "../runtime/index.js"; import { AgentSession, type AgentSessionWriteLockRunner } from "./agent-session.js"; @@ -145,6 +146,66 @@ function getDefaultAgentDir(): string { return getAgentDir(); } +function createSessionPrepareNextTurnWithContext( + getAgent: () => Agent, +): NonNullable { + let activeRunMessages: AgentMessage[] | undefined; + let effectiveModel: Model | undefined; + let effectiveThinkingLevel: ThinkingLevel | undefined; + let lastSessionModel: Model | undefined; + let lastSessionThinkingLevel: ThinkingLevel | undefined; + let lastSessionPrompt: string | undefined; + let lastSessionTools: AgentTool[] = []; + const sameTools = (left: AgentTool[], right: AgentTool[]) => + left.length === right.length && left.every((tool, index) => tool === right[index]); + + return async (turn, signal) => { + const agent = getAgent(); + const firstTurnInRun = activeRunMessages !== turn.newMessages; + if (firstTurnInRun) { + activeRunMessages = turn.newMessages; + effectiveModel = agent.state.model; + effectiveThinkingLevel = agent.state.thinkingLevel; + } + + const previousSnapshot = await agent.prepareNextTurn?.(signal); + const sessionPrompt = agent.state.systemPrompt; + const sessionTools = agent.state.tools; + const sessionModelChanged = firstTurnInRun || agent.state.model !== lastSessionModel; + const sessionThinkingChanged = + firstTurnInRun || agent.state.thinkingLevel !== lastSessionThinkingLevel; + const sessionPromptChanged = firstTurnInRun || sessionPrompt !== lastSessionPrompt; + const sessionToolsChanged = firstTurnInRun || !sameTools(sessionTools, lastSessionTools); + + // Loop-only hook updates persist for the run; fresh session state wins only after it changes. + effectiveModel = + previousSnapshot?.model ?? (sessionModelChanged ? agent.state.model : effectiveModel); + effectiveThinkingLevel = + previousSnapshot?.thinkingLevel ?? + (sessionThinkingChanged ? agent.state.thinkingLevel : effectiveThinkingLevel); + + lastSessionModel = agent.state.model; + lastSessionThinkingLevel = agent.state.thinkingLevel; + lastSessionPrompt = sessionPrompt; + lastSessionTools = sessionTools.slice(); + + const nextContext = previousSnapshot?.context + ? { ...previousSnapshot.context } + : { + ...turn.context, + systemPrompt: sessionPromptChanged ? sessionPrompt : turn.context.systemPrompt, + tools: sessionToolsChanged ? sessionTools.slice() : turn.context.tools?.slice(), + }; + + return { + ...previousSnapshot, + context: nextContext, + model: effectiveModel, + thinkingLevel: effectiveThinkingLevel, + }; + }; +} + function getAttributionHeaders( model: Model, settingsManager: SettingsManager, @@ -438,6 +499,7 @@ export async function createAgentSession( return runner.emitContext(messages); }, resolveDeferredTool: options.resolveDeferredTool, + prepareNextTurnWithContext: createSessionPrepareNextTurnWithContext(() => agent), steeringMode: settingsManager.getSteeringMode(), followUpMode: settingsManager.getFollowUpMode(), transport: settingsManager.getTransport(),