File size: 3,660 Bytes
e249c6d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | import type { Model } from "openclaw/plugin-sdk/llm";
import { describe, expect, it, vi } from "vitest";
import { subscribeEmbeddedAgentSession } from "../embedded-agent-subscribe.js";
import { guardSessionManager } from "../session-tool-result-guard-wrapper.js";
import {
createAssistant,
createAssistantResultStream,
createAutoCompactionSettings,
createOverflowAssistant,
createTestSession,
registerAgentSessionLoopTestLifecycle,
streamMocks,
testModel,
} from "./agent-session-loop-correctness.test-support.js";
import {
createCompactionHandlers,
createResourceLoader,
} from "./agent-session-loop-resource-loader.test-support.js";
import { SessionManager } from "./session-manager.js";
registerAgentSessionLoopTestLifecycle();
describe("AgentSession compaction correlation", () => {
it.each(["manual", "automatic"])(
"correlates repeated %s compactions and tokens through the subscriber",
async (mode) => {
const runId = "run-tokens-after";
const sessionManager = guardSessionManager(SessionManager.inMemory(), { runId });
sessionManager.appendMessage({ role: "user", content: "old prompt", timestamp: 1 });
sessionManager.appendMessage({
...createAssistant(testModel, [{ type: "text", text: "retained answer" }]),
timestamp: 2,
});
let requests = 0;
streamMocks.streamSimple.mockImplementation((activeModel: Model) =>
createAssistantResultStream(
++requests % 2 === 1
? createOverflowAssistant(activeModel)
: createAssistant(activeModel, [{ type: "text", text: "complete retry" }]),
),
);
const { session } = await createTestSession({
sessionManager,
settingsManager: createAutoCompactionSettings(),
resourceLoader: createResourceLoader(createCompactionHandlers()),
});
const onAgentEvent = vi.fn();
const subscription = subscribeEmbeddedAgentSession({ session, runId, onAgentEvent });
try {
for (const prompt of ["first long request", "second long request"]) {
if (mode === "manual") {
sessionManager.appendMessage({ role: "user", content: prompt, timestamp: 3 });
sessionManager.appendMessage(
createAssistant(testModel, [{ type: "text", text: "answer" }]),
);
await session.compact();
} else {
await session.prompt(prompt);
}
}
const compactions = sessionManager
.getBranch()
.filter((entry) => entry.type === "compaction");
expect(compactions).toHaveLength(2);
const itemIds = compactions.map(({ __openclaw: metadata }) => metadata?.itemId);
expect(itemIds).toEqual([expect.any(String), expect.any(String)]);
expect(new Set(itemIds).size).toBe(2);
expect(compactions.map(({ __openclaw: metadata }) => metadata?.runId)).toEqual([
runId,
runId,
]);
const events = onAgentEvent.mock.calls
.map(([event]) => event)
.filter((event) => event.stream === "compaction");
expect(events.map(({ data }) => ({ phase: data.phase, itemId: data.itemId }))).toEqual(
itemIds.flatMap((itemId) => [
{ phase: "start", itemId },
{ phase: "end", itemId },
]),
);
expect(subscription.getCompactionCount()).toBe(2);
expect(subscription.getLastCompactionTokensAfter()).toEqual(expect.any(Number));
expect(subscription.getLastCompactionTokensAfter()).toBeGreaterThan(0);
} finally {
subscription.unsubscribe();
}
},
);
});
|