fix: compaction usage guards and agent-loop turn correctness (#109709)

* fix(agents): compaction usage guards and agent-loop turn correctness

* fix(agents): manual compaction replans with zero retention for small sessions

* ci: retrigger pull-request validation

* ci: refresh pull-request validation

* test(agents): live end-to-end agent session coverage

* fix(agents): keep yielded runs settled so external delivery can claim the session

* docs(agents): clarify length retry bookkeeping
This commit is contained in:
Peter Steinberger
2026-07-17 14:13:29 +01:00
committed by GitHub
parent da69daeb72
commit cbaf112355
15 changed files with 1630 additions and 75 deletions
@@ -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
+274
View File
@@ -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<unknown>,
};
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<AgentEvent, { type: "tool_execution_end" }> =>
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,
+2 -1
View File
@@ -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(),
+23 -4
View File
@@ -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> | AgentLoopTurnUpdate | undefined;
/** Context-aware turn hook. Takes precedence over `prepareNextTurn` when both are provided. */
prepareNextTurnWithContext?: (
context: PrepareNextTurnContext,
signal?: AbortSignal,
) => Promise<AgentLoopTurnUpdate | undefined> | 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> | AgentLoopTurnUpdate | undefined;
public prepareNextTurnWithContext?: (
context: PrepareNextTurnContext,
signal?: AbortSignal,
) => Promise<AgentLoopTurnUpdate | undefined> | 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,
@@ -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", () => {
@@ -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) {
+33 -3
View File
@@ -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<void> => {
protected handleAgentEvent = async (event: AgentEvent, signal?: AbortSignal): Promise<void> => {
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<boolean>;
abstract abortRetry(): void;
abstract abortCompaction(): void;
abstract abortBranchSummary(): void;
abstract abortBash(): void;
protected abstract flushPendingBashMessages(): void;
}
@@ -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;
@@ -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({
@@ -107,8 +107,8 @@ export abstract class AgentSessionInspection extends AgentSessionModels {
if (contextTokens > 0) {
hasPostCompactionUsage = true;
estimateFromContent = false;
break;
}
break;
}
}
}
@@ -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<string, Array<(...args: unknown[]) => Promise<unknown>>> = new Map(),
): ResourceLoader {
const extensionsResult: LoadExtensionsResult = {
extensions:
handlers.size > 0
? [
{
path: "<test-extension>",
resolvedPath: "<test-extension>",
sourceInfo: createSyntheticSourceInfo("<test-extension>", {
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<string, Array<(...args: unknown[]) => Promise<unknown>>>([
[
"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<string, Array<(...args: unknown[]) => Promise<unknown>>>([
[
"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<string, Array<(...args: unknown[]) => Promise<unknown>>>([
["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");
});
});
+18 -13
View File
@@ -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<boolean> {
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) {
@@ -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);
});
});
@@ -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<string, Array<(...args: unknown[]) => Promise<unknown>>>;
function createResourceLoader(handlers: ExtensionHandlers = new Map()): ResourceLoader {
const extensionsResult: LoadExtensionsResult = {
extensions:
handlers.size > 0
? [
{
path: "<live-test-extension>",
resolvedPath: "<live-test-extension>",
sourceInfo: createSyntheticSourceInfo("<live-test-extension>", {
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<typeof echoParameters> = {
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,
);
});
+62
View File
@@ -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<AgentOptions["prepareNextTurnWithContext"]> {
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(),