mirror of
https://github.com/openclaw/openclaw.git
synced 2026-08-12 21:53:00 -06:00
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:
committed by
GitHub
parent
da69daeb72
commit
cbaf112355
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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,
|
||||
);
|
||||
});
|
||||
@@ -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(),
|
||||
|
||||
Reference in New Issue
Block a user