Skip to content

Commit 5e09113

Browse files
vincentkocvincentkoc
committed
refactor: share selected global session test setup
1 parent bff66a3 commit 5e09113

3 files changed

Lines changed: 109 additions & 132 deletions

File tree

‎src/gateway/server.sessions.compaction.test.ts‎

Lines changed: 88 additions & 103 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import fs from "node:fs/promises";
22
import os from "node:os";
33
import path from "node:path";
44
import { expect, test, vi } from "vitest";
5+
import type { SessionCompactionCheckpoint } from "../config/sessions.js";
56
import { withEnvAsync } from "../test-utils/env.js";
67
import {
78
embeddedRunMock,
@@ -21,42 +22,90 @@ import {
2122
directSessionReq,
2223
} from "./test/server-sessions.test-helpers.js";
2324

24-
const { createSessionStoreDir, openClient } = setupGatewaySessionsTestHarness();
25+
const { createSessionStoreDir, createSelectedGlobalSessionStore, openClient } =
26+
setupGatewaySessionsTestHarness();
27+
28+
type CheckpointFixture = Awaited<ReturnType<typeof createCheckpointFixture>>;
29+
30+
function compactionCheckpointEntry(
31+
fixture: CheckpointFixture,
32+
options: {
33+
checkpointId: string;
34+
sessionKey: string;
35+
createdAt: number;
36+
reason: SessionCompactionCheckpoint["reason"];
37+
summary: string;
38+
tokensBefore?: number;
39+
tokensAfter?: number;
40+
},
41+
) {
42+
return {
43+
checkpointId: options.checkpointId,
44+
sessionKey: options.sessionKey,
45+
sessionId: fixture.sessionId,
46+
createdAt: options.createdAt,
47+
reason: options.reason,
48+
summary: options.summary,
49+
...(options.tokensBefore === undefined ? {} : { tokensBefore: options.tokensBefore }),
50+
...(options.tokensAfter === undefined ? {} : { tokensAfter: options.tokensAfter }),
51+
firstKeptEntryId: fixture.preCompactionLeafId,
52+
preCompaction: {
53+
sessionId: fixture.sessionId,
54+
leafId: fixture.preCompactionLeafId,
55+
},
56+
postCompaction: {
57+
sessionId: fixture.sessionId,
58+
sessionFile: fixture.sessionFile,
59+
leafId: fixture.postCompactionLeafId,
60+
entryId: fixture.postCompactionLeafId,
61+
},
62+
};
63+
}
64+
65+
function isCompactOperationEvent(message: unknown, phase: "start" | "end") {
66+
const candidate = message as {
67+
event?: unknown;
68+
payload?: { operation?: unknown; phase?: unknown };
69+
type?: unknown;
70+
};
71+
return (
72+
candidate.type === "event" &&
73+
candidate.event === "session.operation" &&
74+
candidate.payload?.operation === "compact" &&
75+
candidate.payload?.phase === phase
76+
);
77+
}
78+
79+
function expectMainCompactionResult(
80+
compacted: { ok?: boolean; payload?: { compacted?: boolean; key?: string } | null },
81+
expectedCompacted: boolean,
82+
) {
83+
expect(compacted.ok).toBe(true);
84+
expect(compacted.payload?.key).toBe("agent:main:main");
85+
expect(compacted.payload?.compacted).toBe(expectedCompacted);
86+
}
2587

2688
test("sessions.compaction.* lists checkpoints and branches or restores from compacted transcripts", async () => {
2789
const { dir, storePath } = await createSessionStoreDir();
2890
const fixture = await createCheckpointFixture(dir, { legacyPreCompactionSnapshot: false });
2991
expect((await fs.readdir(dir)).some((file) => file.includes(".checkpoint."))).toBe(false);
3092
const checkpointEntryCount = fixture.session.getEntries().length;
3193
const checkpointCreatedAt = Date.now();
94+
const checkpointEntry = compactionCheckpointEntry(fixture, {
95+
checkpointId: "checkpoint-1",
96+
sessionKey: "agent:main:main",
97+
createdAt: checkpointCreatedAt,
98+
reason: "manual",
99+
summary: "checkpoint summary",
100+
tokensBefore: 123,
101+
tokensAfter: 45,
102+
});
32103
const { SessionManager } = await getSessionManagerModule();
33104
await writeSessionStore({
34105
entries: {
35106
main: sessionStoreEntry(fixture.sessionId, {
36107
sessionFile: fixture.sessionFile,
37-
compactionCheckpoints: [
38-
{
39-
checkpointId: "checkpoint-1",
40-
sessionKey: "agent:main:main",
41-
sessionId: fixture.sessionId,
42-
createdAt: checkpointCreatedAt,
43-
reason: "manual",
44-
tokensBefore: 123,
45-
tokensAfter: 45,
46-
summary: "checkpoint summary",
47-
firstKeptEntryId: fixture.preCompactionLeafId,
48-
preCompaction: {
49-
sessionId: fixture.sessionId,
50-
leafId: fixture.preCompactionLeafId,
51-
},
52-
postCompaction: {
53-
sessionId: fixture.sessionId,
54-
sessionFile: fixture.sessionFile,
55-
leafId: fixture.postCompactionLeafId,
56-
entryId: fixture.postCompactionLeafId,
57-
},
58-
},
59-
],
108+
compactionCheckpoints: [checkpointEntry],
60109
}),
61110
},
62111
});
@@ -101,27 +150,7 @@ test("sessions.compaction.* lists checkpoints and branches or restores from comp
101150
expect(listedCheckpoints.ok).toBe(true);
102151
expect(listedCheckpoints.payload?.key).toBe("agent:main:main");
103152
expect(listedCheckpoints.payload?.checkpoints).toHaveLength(1);
104-
expect(listedCheckpoints.payload?.checkpoints[0]).toEqual({
105-
checkpointId: "checkpoint-1",
106-
sessionKey: "agent:main:main",
107-
sessionId: fixture.sessionId,
108-
createdAt: checkpointCreatedAt,
109-
reason: "manual",
110-
summary: "checkpoint summary",
111-
tokensBefore: 123,
112-
tokensAfter: 45,
113-
firstKeptEntryId: fixture.preCompactionLeafId,
114-
preCompaction: {
115-
sessionId: fixture.sessionId,
116-
leafId: fixture.preCompactionLeafId,
117-
},
118-
postCompaction: {
119-
sessionId: fixture.sessionId,
120-
sessionFile: fixture.sessionFile,
121-
leafId: fixture.postCompactionLeafId,
122-
entryId: fixture.postCompactionLeafId,
123-
},
124-
});
153+
expect(listedCheckpoints.payload?.checkpoints[0]).toEqual(checkpointEntry);
125154

126155
const checkpoint = await rpcReq<{
127156
ok: true;
@@ -277,20 +306,21 @@ test("sessions.compaction.* lists checkpoints and branches or restores from comp
277306
});
278307

279308
test("sessions.compaction.* scopes selected global checkpoints to the requested agent", async () => {
280-
const { dir } = await createSessionStoreDir();
281-
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
282-
testState.sessionStorePath = storeTemplate;
283-
testState.sessionConfig = { scope: "global" };
284-
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
285-
const mainStorePath = storeTemplate.replace("{agentId}", "main");
286-
const workStorePath = storeTemplate.replace("{agentId}", "work");
309+
const { mainStorePath, workStorePath } = await createSelectedGlobalSessionStore();
287310
const workDir = path.dirname(workStorePath);
288311
await fs.mkdir(path.dirname(mainStorePath), { recursive: true });
289312
await fs.mkdir(workDir, { recursive: true });
290313
const mainSessionFile = path.join(path.dirname(mainStorePath), "sess-main-global.jsonl");
291314
await fs.writeFile(mainSessionFile, `${JSON.stringify({ role: "user", content: "main" })}\n`);
292315
const fixture = await createCheckpointFixture(workDir, { legacyPreCompactionSnapshot: false });
293316
const checkpointCreatedAt = Date.now();
317+
const checkpointEntry = compactionCheckpointEntry(fixture, {
318+
checkpointId: "checkpoint-work",
319+
sessionKey: "global",
320+
createdAt: checkpointCreatedAt,
321+
reason: "manual",
322+
summary: "work checkpoint",
323+
});
294324
await fs.writeFile(
295325
mainStorePath,
296326
JSON.stringify(
@@ -305,27 +335,7 @@ test("sessions.compaction.* scopes selected global checkpoints to the requested
305335
{
306336
global: sessionStoreEntry(fixture.sessionId, {
307337
sessionFile: fixture.sessionFile,
308-
compactionCheckpoints: [
309-
{
310-
checkpointId: "checkpoint-work",
311-
sessionKey: "global",
312-
sessionId: fixture.sessionId,
313-
createdAt: checkpointCreatedAt,
314-
reason: "manual",
315-
summary: "work checkpoint",
316-
firstKeptEntryId: fixture.preCompactionLeafId,
317-
preCompaction: {
318-
sessionId: fixture.sessionId,
319-
leafId: fixture.preCompactionLeafId,
320-
},
321-
postCompaction: {
322-
sessionId: fixture.sessionId,
323-
sessionFile: fixture.sessionFile,
324-
leafId: fixture.postCompactionLeafId,
325-
entryId: fixture.postCompactionLeafId,
326-
},
327-
},
328-
],
338+
compactionCheckpoints: [checkpointEntry],
329339
}),
330340
},
331341
null,
@@ -410,22 +420,8 @@ test("sessions.compact without maxLines runs embedded manual compaction for chec
410420

411421
const { ws } = await openClient();
412422
await rpcReq(ws, "sessions.subscribe", {});
413-
const startEventPromise = onceMessage(
414-
ws,
415-
(message) =>
416-
message.type === "event" &&
417-
message.event === "session.operation" &&
418-
(message.payload as { operation?: unknown; phase?: unknown })?.operation === "compact" &&
419-
(message.payload as { operation?: unknown; phase?: unknown })?.phase === "start",
420-
);
421-
const endEventPromise = onceMessage(
422-
ws,
423-
(message) =>
424-
message.type === "event" &&
425-
message.event === "session.operation" &&
426-
(message.payload as { operation?: unknown; phase?: unknown })?.operation === "compact" &&
427-
(message.payload as { operation?: unknown; phase?: unknown })?.phase === "end",
428-
);
423+
const startEventPromise = onceMessage(ws, (message) => isCompactOperationEvent(message, "start"));
424+
const endEventPromise = onceMessage(ws, (message) => isCompactOperationEvent(message, "end"));
429425
const compacted = await rpcReq<{
430426
ok: true;
431427
key: string;
@@ -435,9 +431,7 @@ test("sessions.compact without maxLines runs embedded manual compaction for chec
435431
key: "main",
436432
});
437433

438-
expect(compacted.ok).toBe(true);
439-
expect(compacted.payload?.key).toBe("agent:main:main");
440-
expect(compacted.payload?.compacted).toBe(true);
434+
expectMainCompactionResult(compacted, true);
441435
const startEvent = await startEventPromise;
442436
const endEvent = await endEventPromise;
443437
const startPayload = startEvent.payload as {
@@ -568,14 +562,7 @@ test("sessions.compact treats Codex native compaction start as pending, not comp
568562

569563
const { ws } = await openClient();
570564
await rpcReq(ws, "sessions.subscribe", {});
571-
const endEventPromise = onceMessage(
572-
ws,
573-
(message) =>
574-
message.type === "event" &&
575-
message.event === "session.operation" &&
576-
(message.payload as { operation?: unknown; phase?: unknown })?.operation === "compact" &&
577-
(message.payload as { operation?: unknown; phase?: unknown })?.phase === "end",
578-
);
565+
const endEventPromise = onceMessage(ws, (message) => isCompactOperationEvent(message, "end"));
579566

580567
const compacted = await rpcReq<{
581568
ok: true;
@@ -586,9 +573,7 @@ test("sessions.compact treats Codex native compaction start as pending, not comp
586573
key: "main",
587574
});
588575

589-
expect(compacted.ok).toBe(true);
590-
expect(compacted.payload?.key).toBe("agent:main:main");
591-
expect(compacted.payload?.compacted).toBe(false);
576+
expectMainCompactionResult(compacted, false);
592577
expect(compacted.payload?.result?.details).toMatchObject({
593578
backend: "codex-app-server",
594579
threadId: "thread-1",

‎src/gateway/server.sessions.create.test.ts‎

Lines changed: 6 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@ import {
1010
sessionLifecycleHookMocks,
1111
} from "./test/server-sessions.test-helpers.js";
1212

13-
const { createSessionStoreDir, openClient } = setupGatewaySessionsTestHarness();
13+
const { createSessionStoreDir, createSelectedGlobalSessionStore, openClient } =
14+
setupGatewaySessionsTestHarness();
1415

1516
function requireNonEmptyString(value: string | undefined, label: string): string {
1617
if (!value) {
@@ -302,13 +303,7 @@ test("sessions.create preserves global and unknown sentinel keys", async () => {
302303
});
303304

304305
test("sessions.create stores selected global sessions in the requested agent store", async () => {
305-
const { dir } = await createSessionStoreDir();
306-
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
307-
const mainStorePath = storeTemplate.replace("{agentId}", "main");
308-
const workStorePath = storeTemplate.replace("{agentId}", "work");
309-
testState.sessionStorePath = storeTemplate;
310-
testState.sessionConfig = { scope: "global" };
311-
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
306+
const { mainStorePath, workStorePath } = await createSelectedGlobalSessionStore();
312307
const broadcastToConnIds = vi.fn();
313308

314309
const created = await directSessionReq<{
@@ -350,13 +345,7 @@ test("sessions.create stores selected global sessions in the requested agent sto
350345
});
351346

352347
test("sessions.create loads selected global parent from the requested agent store", async () => {
353-
const { dir } = await createSessionStoreDir();
354-
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
355-
const mainStorePath = storeTemplate.replace("{agentId}", "main");
356-
const workStorePath = storeTemplate.replace("{agentId}", "work");
357-
testState.sessionStorePath = storeTemplate;
358-
testState.sessionConfig = { scope: "global" };
359-
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
348+
const { mainStorePath, workStorePath } = await createSelectedGlobalSessionStore();
360349
try {
361350
await writeSessionStore({
362351
storePath: mainStorePath,
@@ -430,10 +419,7 @@ test("sessions.create loads selected global parent from the requested agent stor
430419
});
431420

432421
test("sessions.get reads selected global messages from the requested agent store", async () => {
433-
const { dir } = await createSessionStoreDir();
434-
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
435-
const mainStorePath = storeTemplate.replace("{agentId}", "main");
436-
const workStorePath = storeTemplate.replace("{agentId}", "work");
422+
const { mainStorePath, workStorePath } = await createSelectedGlobalSessionStore();
437423
const mainTranscriptPath = path.join(path.dirname(mainStorePath), "sess-main-global.jsonl");
438424
const workTranscriptPath = path.join(path.dirname(workStorePath), "sess-work-global.jsonl");
439425
await fs.mkdir(path.dirname(mainTranscriptPath), { recursive: true });
@@ -448,9 +434,6 @@ test("sessions.get reads selected global messages from the requested agent store
448434
`${JSON.stringify({ type: "message", id: "work-msg", message: { role: "user", content: "work global" } })}\n`,
449435
"utf-8",
450436
);
451-
testState.sessionStorePath = storeTemplate;
452-
testState.sessionConfig = { scope: "global" };
453-
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
454437
try {
455438
await writeSessionStore({
456439
storePath: mainStorePath,
@@ -487,11 +470,7 @@ test("sessions.get reads selected global messages from the requested agent store
487470
});
488471

489472
test("sessions.create sends selected global initial tasks to the requested agent", async () => {
490-
const { dir } = await createSessionStoreDir();
491-
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
492-
testState.sessionStorePath = storeTemplate;
493-
testState.sessionConfig = { scope: "global" };
494-
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
473+
const { mainStorePath, workStorePath } = await createSelectedGlobalSessionStore();
495474
const { ws } = await openClient();
496475

497476
const created = await rpcReq<{
@@ -510,8 +489,6 @@ test("sessions.create sends selected global initial tasks to the requested agent
510489
const runId = requireNonEmptyString(created.payload?.runId, "selected global run id");
511490
const wait = await rpcReq(ws, "agent.wait", { runId, timeoutMs: 1_000 });
512491
expect(wait.ok).toBe(true);
513-
const workStorePath = storeTemplate.replace("{agentId}", "work");
514-
const mainStorePath = storeTemplate.replace("{agentId}", "main");
515492
const workStore = JSON.parse(await fs.readFile(workStorePath, "utf-8")) as Record<
516493
string,
517494
{ sessionFile?: string }

‎src/gateway/test/server-sessions.test-helpers.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,20 @@ export function setupGatewaySessionsTestHarness() {
320320
return { dir, storePath };
321321
}
322322

323+
async function createSelectedGlobalSessionStore() {
324+
const { dir } = await createSessionStoreDir();
325+
const storeTemplate = path.join(dir, "{agentId}", "sessions.json");
326+
testState.sessionStorePath = storeTemplate;
327+
testState.sessionConfig = { scope: "global" };
328+
testState.agentsConfig = { list: [{ id: "main", default: true }, { id: "work" }] };
329+
return {
330+
dir,
331+
storeTemplate,
332+
mainStorePath: storeTemplate.replace("{agentId}", "main"),
333+
workStorePath: storeTemplate.replace("{agentId}", "work"),
334+
};
335+
}
336+
323337
async function seedActiveMainSession() {
324338
const { dir, storePath } = await createSessionStoreDir();
325339
await writeSingleLineSession(dir, "sess-main", "hello");
@@ -333,6 +347,7 @@ export function setupGatewaySessionsTestHarness() {
333347

334348
return {
335349
createSessionStoreDir,
350+
createSelectedGlobalSessionStore,
336351
getHarness: requireHarness,
337352
openClient,
338353
seedActiveMainSession,

0 commit comments

Comments
 (0)