Skip to content

Commit 908bb0f

Browse files
steipetesteipete
committed
test: clear timeout compaction broad matchers
1 parent f83b831 commit 908bb0f

1 file changed

Lines changed: 130 additions & 92 deletions

File tree

‎src/agents/pi-embedded-runner/run.timeout-triggered-compaction.test.ts‎

Lines changed: 130 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,87 @@ const useTwoAuthProfiles = () => {
2626
}));
2727
};
2828

29+
type CompactRuntimeContext = {
30+
promptCache?: {
31+
retention?: string;
32+
lastCallUsage?: {
33+
input?: number;
34+
cacheRead?: number;
35+
};
36+
observation?: {
37+
broke?: boolean;
38+
cacheRead?: number;
39+
};
40+
lastCacheTouchAt?: number;
41+
};
42+
trigger?: string;
43+
attempt?: number;
44+
maxAttempts?: number;
45+
messageChannel?: string;
46+
messageProvider?: string;
47+
agentAccountId?: string;
48+
currentChannelId?: string;
49+
currentThreadTs?: string;
50+
currentMessageId?: string;
51+
senderId?: string;
52+
senderIsOwner?: boolean;
53+
authProfileId?: string;
54+
};
55+
56+
type CompactParams = {
57+
sessionId?: string;
58+
sessionFile?: string;
59+
tokenBudget?: number;
60+
force?: boolean;
61+
compactionTarget?: string;
62+
runtimeContext?: CompactRuntimeContext;
63+
};
64+
65+
type AttemptParams = {
66+
sessionId?: string;
67+
sessionFile?: string;
68+
authProfileId?: string;
69+
};
70+
71+
type HookEvent = {
72+
messageCount?: number;
73+
compactedCount?: number;
74+
tokenCount?: number;
75+
sessionFile?: string;
76+
};
77+
78+
type HookContext = {
79+
sessionKey?: string;
80+
};
81+
82+
function compactCallAt(index: number): CompactParams {
83+
const call = mockedCompactDirect.mock.calls[index];
84+
if (!call) {
85+
throw new Error(`expected compact call ${index + 1}`);
86+
}
87+
return call[0] as CompactParams;
88+
}
89+
90+
function attemptCallAt(index: number): AttemptParams {
91+
const call = mockedRunEmbeddedAttempt.mock.calls[index];
92+
if (!call) {
93+
throw new Error(`expected embedded attempt call ${index + 1}`);
94+
}
95+
return call[0] as AttemptParams;
96+
}
97+
98+
function hookCallAt(index: number, kind: "before" | "after"): [HookEvent, HookContext] {
99+
const mock =
100+
kind === "before"
101+
? mockedGlobalHookRunner.runBeforeCompaction
102+
: mockedGlobalHookRunner.runAfterCompaction;
103+
const call = mock.mock.calls[index];
104+
if (!call) {
105+
throw new Error(`expected ${kind} compaction hook call ${index + 1}`);
106+
}
107+
return call as unknown as [HookEvent, HookContext];
108+
}
109+
29110
describe("timeout-triggered compaction", () => {
30111
beforeAll(async () => {
31112
({ runEmbeddedPiAgent } = await loadRunOverflowCompactionHarness());
@@ -72,32 +153,21 @@ describe("timeout-triggered compaction", () => {
72153
const result = await runEmbeddedPiAgent(overflowBaseRunParams);
73154

74155
expect(mockedCompactDirect).toHaveBeenCalledTimes(1);
75-
expect(mockedCompactDirect).toHaveBeenCalledWith(
76-
expect.objectContaining({
77-
sessionId: "test-session",
78-
sessionFile: "/tmp/session.json",
79-
tokenBudget: 200000,
80-
force: true,
81-
compactionTarget: "budget",
82-
runtimeContext: expect.objectContaining({
83-
promptCache: expect.objectContaining({
84-
retention: "short",
85-
lastCallUsage: expect.objectContaining({
86-
input: 150000,
87-
cacheRead: 32000,
88-
}),
89-
observation: expect.objectContaining({
90-
broke: false,
91-
cacheRead: 32000,
92-
}),
93-
lastCacheTouchAt: 1_700_000_000_000,
94-
}),
95-
trigger: "timeout_recovery",
96-
attempt: 1,
97-
maxAttempts: 2,
98-
}),
99-
}),
100-
);
156+
const compactParams = compactCallAt(0);
157+
expect(compactParams.sessionId).toBe("test-session");
158+
expect(compactParams.sessionFile).toBe("/tmp/session.json");
159+
expect(compactParams.tokenBudget).toBe(200000);
160+
expect(compactParams.force).toBe(true);
161+
expect(compactParams.compactionTarget).toBe("budget");
162+
expect(compactParams.runtimeContext?.promptCache?.retention).toBe("short");
163+
expect(compactParams.runtimeContext?.promptCache?.lastCallUsage?.input).toBe(150000);
164+
expect(compactParams.runtimeContext?.promptCache?.lastCallUsage?.cacheRead).toBe(32000);
165+
expect(compactParams.runtimeContext?.promptCache?.observation?.broke).toBe(false);
166+
expect(compactParams.runtimeContext?.promptCache?.observation?.cacheRead).toBe(32000);
167+
expect(compactParams.runtimeContext?.promptCache?.lastCacheTouchAt).toBe(1_700_000_000_000);
168+
expect(compactParams.runtimeContext?.trigger).toBe("timeout_recovery");
169+
expect(compactParams.runtimeContext?.attempt).toBe(1);
170+
expect(compactParams.runtimeContext?.maxAttempts).toBe(2);
101171
expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2);
102172
expect(result.meta.error).toBeUndefined();
103173
expect(result.meta.agentMeta?.compactionTokensAfter).toBe(80_000);
@@ -136,13 +206,9 @@ describe("timeout-triggered compaction", () => {
136206

137207
// Verify the loop continued (retry happened)
138208
expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2);
139-
expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith(
140-
2,
141-
expect.objectContaining({
142-
sessionId: "timeout-rotated-session",
143-
sessionFile: "/tmp/timeout-rotated-session.json",
144-
}),
145-
);
209+
const retryParams = attemptCallAt(1);
210+
expect(retryParams.sessionId).toBe("timeout-rotated-session");
211+
expect(retryParams.sessionFile).toBe("/tmp/timeout-rotated-session.json");
146212
expect(mockedRunPostCompactionSideEffects).not.toHaveBeenCalled();
147213
expect(result.meta.error).toBeUndefined();
148214
});
@@ -177,20 +243,16 @@ describe("timeout-triggered compaction", () => {
177243
senderIsOwner: true,
178244
});
179245

180-
expect(mockedCompactDirect).toHaveBeenCalledWith(
181-
expect.objectContaining({
182-
runtimeContext: expect.objectContaining({
183-
messageChannel: "slack",
184-
messageProvider: "slack",
185-
agentAccountId: "acct-1",
186-
currentChannelId: "channel-1",
187-
currentThreadTs: "thread-1",
188-
currentMessageId: "message-1",
189-
senderId: "sender-1",
190-
senderIsOwner: true,
191-
}),
192-
}),
193-
);
246+
expect(mockedCompactDirect).toHaveBeenCalledTimes(1);
247+
const compactParams = compactCallAt(0);
248+
expect(compactParams.runtimeContext?.messageChannel).toBe("slack");
249+
expect(compactParams.runtimeContext?.messageProvider).toBe("slack");
250+
expect(compactParams.runtimeContext?.agentAccountId).toBe("acct-1");
251+
expect(compactParams.runtimeContext?.currentChannelId).toBe("channel-1");
252+
expect(compactParams.runtimeContext?.currentThreadTs).toBe("thread-1");
253+
expect(compactParams.runtimeContext?.currentMessageId).toBe("message-1");
254+
expect(compactParams.runtimeContext?.senderId).toBe("sender-1");
255+
expect(compactParams.runtimeContext?.senderIsOwner).toBe(true);
194256
});
195257

196258
it("falls through to normal handling when timeout compaction fails", async () => {
@@ -448,23 +510,17 @@ describe("timeout-triggered compaction", () => {
448510

449511
await runEmbeddedPiAgent(overflowBaseRunParams);
450512

451-
expect(mockedGlobalHookRunner.runBeforeCompaction).toHaveBeenCalledWith(
452-
{ messageCount: -1, sessionFile: "/tmp/session.json" },
453-
expect.objectContaining({
454-
sessionKey: "test-key",
455-
}),
456-
);
457-
expect(mockedGlobalHookRunner.runAfterCompaction).toHaveBeenCalledWith(
458-
{
459-
messageCount: -1,
460-
compactedCount: -1,
461-
tokenCount: 70,
462-
sessionFile: "/tmp/session.json",
463-
},
464-
expect.objectContaining({
465-
sessionKey: "test-key",
466-
}),
467-
);
513+
const [beforeEvent, beforeContext] = hookCallAt(0, "before");
514+
expect(beforeEvent).toEqual({ messageCount: -1, sessionFile: "/tmp/session.json" });
515+
expect(beforeContext.sessionKey).toBe("test-key");
516+
const [afterEvent, afterContext] = hookCallAt(0, "after");
517+
expect(afterEvent).toEqual({
518+
messageCount: -1,
519+
compactedCount: -1,
520+
tokenCount: 70,
521+
sessionFile: "/tmp/session.json",
522+
});
523+
expect(afterContext.sessionKey).toBe("test-key");
468524
expect(mockedRunPostCompactionSideEffects).toHaveBeenCalledTimes(1);
469525
});
470526

@@ -506,26 +562,14 @@ describe("timeout-triggered compaction", () => {
506562
const result = await runEmbeddedPiAgent(overflowBaseRunParams);
507563

508564
expect(mockedCompactDirect).toHaveBeenCalledTimes(2);
509-
expect(mockedCompactDirect).toHaveBeenNthCalledWith(
510-
1,
511-
expect.objectContaining({
512-
runtimeContext: expect.objectContaining({
513-
authProfileId: "profile-a",
514-
attempt: 1,
515-
maxAttempts: 2,
516-
}),
517-
}),
518-
);
519-
expect(mockedCompactDirect).toHaveBeenNthCalledWith(
520-
2,
521-
expect.objectContaining({
522-
runtimeContext: expect.objectContaining({
523-
authProfileId: "profile-b",
524-
attempt: 2,
525-
maxAttempts: 2,
526-
}),
527-
}),
528-
);
565+
const firstCompact = compactCallAt(0);
566+
expect(firstCompact.runtimeContext?.authProfileId).toBe("profile-a");
567+
expect(firstCompact.runtimeContext?.attempt).toBe(1);
568+
expect(firstCompact.runtimeContext?.maxAttempts).toBe(2);
569+
const secondCompact = compactCallAt(1);
570+
expect(secondCompact.runtimeContext?.authProfileId).toBe("profile-b");
571+
expect(secondCompact.runtimeContext?.attempt).toBe(2);
572+
expect(secondCompact.runtimeContext?.maxAttempts).toBe(2);
529573
expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2);
530574
expect(result.payloads?.[0]?.isError).toBe(true);
531575
expect(result.payloads?.[0]?.text).toContain("timed out");
@@ -562,14 +606,8 @@ describe("timeout-triggered compaction", () => {
562606

563607
expect(mockedCompactDirect).toHaveBeenCalledTimes(2);
564608
expect(mockedRunEmbeddedAttempt).toHaveBeenCalledTimes(2);
565-
expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith(
566-
1,
567-
expect.objectContaining({ authProfileId: "profile-a" }),
568-
);
569-
expect(mockedRunEmbeddedAttempt).toHaveBeenNthCalledWith(
570-
2,
571-
expect.objectContaining({ authProfileId: "profile-b" }),
572-
);
609+
expect(attemptCallAt(0).authProfileId).toBe("profile-a");
610+
expect(attemptCallAt(1).authProfileId).toBe("profile-b");
573611
expect(result.payloads?.[0]?.isError).toBe(true);
574612
expect(result.payloads?.[0]?.text).toContain("timed out");
575613
});

0 commit comments

Comments
 (0)