Skip to content

Commit fbfb2c4

Browse files
benbrandtbenbrandt
committed
Cleanups
1 parent 843654a commit fbfb2c4

6 files changed

Lines changed: 154 additions & 31 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,6 +387,7 @@ export class CodexAcpClient {
387387

388388
await this.codexClient.threadUnsubscribe({threadId: sessionId});
389389
this.unsubscribeFromSessionEvents(sessionId);
390+
this.codexClient.clearThreadState(sessionId);
390391
}
391392

392393
async startPrompt(

‎src/CodexAcpServer.ts‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -685,6 +685,12 @@ export class CodexAcpServer implements acp.Agent {
685685
}
686686
}
687687

688+
private ensureSessionIsActive(sessionId: string, requireTrackedSession: boolean = false): void {
689+
if (this.closingSessions.has(sessionId) || (requireTrackedSession && !this.sessions.has(sessionId))) {
690+
throw RequestError.invalidRequest("Session is closing");
691+
}
692+
}
693+
688694
private forgetSession(sessionId: string): void {
689695
this.sessions.delete(sessionId);
690696
this.pendingMcpStartupSessions.delete(sessionId);
@@ -721,20 +727,23 @@ export class CodexAcpServer implements acp.Agent {
721727
return;
722728
}
723729

730+
const requestedServers = pendingStartup.requestedServers;
731+
const afterVersion = pendingStartup.afterVersion;
732+
724733
try {
725734
const mcpStartup = await this.runWithProcessCheck(() =>
726735
this.codexAcpClient.awaitMcpServerStartup(
727-
Array.from(pendingStartup.requestedServers),
728-
pendingStartup.afterVersion,
736+
Array.from(requestedServers),
737+
afterVersion,
729738
)
730739
);
731740
const sessionState = this.sessions.get(sessionId);
732-
const pendingStartup = this.pendingMcpStartupSessions.get(sessionId);
733-
if (sessionState && pendingStartup) {
741+
const currentPendingStartup = this.pendingMcpStartupSessions.get(sessionId);
742+
if (sessionState && currentPendingStartup) {
734743
sessionState.sessionMcpServers = mcpStartup.ready.filter(serverName =>
735-
pendingStartup.requestedServers.has(serverName)
744+
currentPendingStartup.requestedServers.has(serverName)
736745
);
737-
await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.requestedServers);
746+
await this.publishMcpStartupStatus(sessionId, mcpStartup, currentPendingStartup.requestedServers);
738747
} else {
739748
logger.log("Skipping MCP startup status for closed session", {sessionId});
740749
}
@@ -772,9 +781,8 @@ await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.request
772781
prompt: params.prompt,
773782
});
774783
const sessionState = this.getSessionState(params.sessionId);
775-
if (this.closingSessions.has(params.sessionId)) {
776-
throw RequestError.invalidRequest("Session is closing");
777-
}
784+
const requireTrackedSession = this.sessions.has(params.sessionId);
785+
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
778786
sessionState.pendingTurnId = null;
779787
sessionState.currentTurnId = null;
780788
sessionState.lastTokenUsage = null;
@@ -793,6 +801,7 @@ await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.request
793801
},
794802
approvalHandler,
795803
elicitationHandler);
804+
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
796805

797806
if (await this.availableCommands.tryHandle(params.prompt, sessionState)) {
798807
logger.log("Prompt handled by a command");
@@ -803,6 +812,8 @@ await this.publishMcpStartupStatus(sessionId, mcpStartup, pendingStartup.request
803812
};
804813
}
805814

815+
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
816+
806817
const modelId = ModelId.fromString(sessionState.currentModelId);
807818
const modelLacksReasoning = sessionState.supportedReasoningEfforts.length > 0
808819
&& sessionState.supportedReasoningEfforts.every(e => e.reasoningEffort === "none");

‎src/CodexAppServerClient.ts‎

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -103,19 +103,18 @@ export class CodexAppServerClient {
103103
this.connection = connection;
104104
this.connection.onUnhandledNotification((data) => {
105105
const serverNotification = data as ServerNotification;
106-
if (isMcpServerStatusUpdatedNotification(serverNotification)) {
107-
this.mcpServerStartupVersion += 1;
108-
this.mcpServerStartupStates.set(serverNotification.params.name, {
109-
status: serverNotification.params.status,
110-
error: serverNotification.params.error,
111-
version: this.mcpServerStartupVersion,
112-
});
113-
this.resolveMcpServerStartupResolvers();
114-
}
106+
if (isMcpServerStatusUpdatedNotification(serverNotification)) {
107+
this.mcpServerStartupVersion += 1;
108+
this.mcpServerStartupStates.set(serverNotification.params.name, {
109+
status: serverNotification.params.status,
110+
error: serverNotification.params.error,
111+
version: this.mcpServerStartupVersion,
112+
});
113+
this.resolveMcpServerStartupResolvers();
114+
}
115115
if (isTurnCompletedNotification(serverNotification)) {
116116
this.lastTurnCompletedByThread.set(serverNotification.params.threadId, serverNotification.params);
117117
this.resolveTurnCompleted(serverNotification.params);
118-
}
119118
}
120119
this.notify(serverNotification);
121120
for (const callback of this.codexEventHandlers) {
@@ -282,6 +281,11 @@ export class CodexAppServerClient {
282281
this.notificationHandlers.delete(sessionId);
283282
}
284283

284+
clearThreadState(threadId: string): void {
285+
this.turnCompletedResolvers.delete(threadId);
286+
this.lastTurnCompletedByThread.delete(threadId);
287+
}
288+
285289
private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
286290
onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
287291
this.codexEventHandlers.push(callback);
@@ -418,12 +422,6 @@ type TurnCompletedResolver = {
418422
resolve: (event: TurnCompletedNotification) => void;
419423
};
420424

421-
type McpStartupCompleteNotification = {
422-
method: "codex/event/mcp_startup_complete",
423-
params: {
424-
msg: McpStartupCompleteEvent & { type: "mcp_startup_complete" }
425-
}
426-
};
427425

428426
function isMcpServerStatusUpdatedNotification(notification: ServerNotification): notification is {
429427
method: "mcpServer/startupStatus/updated";

‎src/CodexCommands.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,7 +205,7 @@ export class CodexCommands {
205205
}
206206

207207
private async updateSession(sessionId: string, update: UpdateSessionEvent): Promise<void> {
208-
if (this.isSessionClosing(sessionId)) {
208+
if (this.isSessionClosing(sessionId) || !this.hasTrackedSession(sessionId)) {
209209
return;
210210
}
211211
const session = new ACPSessionConnection(this.connection, sessionId);

‎src/__tests__/CodexACPAgent/CodexAcpClient.test.ts‎

Lines changed: 115 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -649,6 +649,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
649649

650650
const sessionState: SessionState = createTestSessionState();
651651
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
652+
// @ts-expect-error seeding private session store for slash-command test
653+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
652654

653655
await codexAcpAgent.prompt({ sessionId: "session-id", prompt: [{ type: "text", text: "/status" }] });
654656
await expect(mockFixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/command-status.json");
@@ -661,6 +663,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
661663
const sessionState: SessionState = createTestSessionState();
662664

663665
const logoutSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "logout").mockResolvedValue({});
666+
// @ts-expect-error seeding private session store for slash-command test
667+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
664668

665669
// @ts-expect-error - exercising private helper
666670
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "logout", input: null }, sessionState);
@@ -687,6 +691,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
687691
}]
688692
};
689693
const skillsSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "listSkills").mockResolvedValue(skillsResponse);
694+
// @ts-expect-error seeding private session store for slash-command test
695+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
690696

691697
// @ts-expect-error - exercising private helper
692698
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "skills", input: null }, sessionState);
@@ -722,6 +728,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
722728
nextCursor: null
723729
};
724730
const mcpSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "listMcpServers").mockResolvedValue(mcpResponse);
731+
// @ts-expect-error seeding private session store for slash-command test
732+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
725733

726734
// @ts-expect-error - exercising private helper
727735
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "mcp", input: null }, sessionState);
@@ -811,7 +819,10 @@ describe('ACP server test', { timeout: 40_000 }, () => {
811819
durationMs: null,
812820
}
813821
});
814-
vi.spyOn(mockFixture.getCodexAcpAgent(), "getSessionState").mockReturnValue(sessionState);
822+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
823+
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
824+
// @ts-expect-error seeding private session store for prompt test
825+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
815826
return { mockFixture, sessionState, turnStartSpy };
816827
}
817828

@@ -1055,6 +1066,24 @@ describe('ACP server test', { timeout: 40_000 }, () => {
10551066
expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(true);
10561067
// @ts-expect-error verifying test-only access to private maps
10571068
expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(true);
1069+
// @ts-expect-error verifying test-only access to private maps
1070+
codexAppServerClient.turnCompletedResolvers.set("session-close", [{
1071+
turnId: "other-turn",
1072+
resolve: vi.fn(),
1073+
}]);
1074+
// @ts-expect-error verifying test-only access to private maps
1075+
codexAppServerClient.lastTurnCompletedByThread.set("session-close", {
1076+
threadId: "session-close",
1077+
turn: {
1078+
id: "completed-turn",
1079+
items: [],
1080+
status: "completed",
1081+
error: null,
1082+
startedAt: null,
1083+
completedAt: null,
1084+
durationMs: null,
1085+
},
1086+
});
10581087

10591088
await codexAcpClient.closeSession("session-close", "turn-close");
10601089

@@ -1079,6 +1108,10 @@ describe('ACP server test', { timeout: 40_000 }, () => {
10791108
expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(false);
10801109
// @ts-expect-error verifying test-only access to private maps
10811110
expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(false);
1111+
// @ts-expect-error verifying test-only access to private maps
1112+
expect(codexAppServerClient.turnCompletedResolvers.has("session-close")).toBe(false);
1113+
// @ts-expect-error verifying test-only access to private maps
1114+
expect(codexAppServerClient.lastTurnCompletedByThread.has("session-close")).toBe(false);
10821115
});
10831116

10841117
it('waits for the matching turn completion before resolving a prompt', async () => {
@@ -1212,6 +1245,45 @@ describe('ACP server test', { timeout: 40_000 }, () => {
12121245
expect(closeSpy).toHaveBeenCalledWith(sessionState.sessionId, "turn-pending");
12131246
});
12141247

1248+
it('rejects a prompt if the session closes before turn start begins', async () => {
1249+
const mockFixture = createCodexMockTestFixture();
1250+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
1251+
const codexAcpClient = mockFixture.getCodexAcpClient();
1252+
const sessionState = createTestSessionState({
1253+
sessionId: "session-close",
1254+
});
1255+
1256+
let resolveTryHandle!: (handled: boolean) => void;
1257+
const tryHandlePromise = new Promise<boolean>((resolve) => {
1258+
resolveTryHandle = resolve;
1259+
});
1260+
1261+
// @ts-expect-error seeding private session store for focused close-session test
1262+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
1263+
// @ts-expect-error exercising private helper through the availableCommands collaborator
1264+
const tryHandleSpy = vi.spyOn(codexAcpAgent.availableCommands, "tryHandle").mockReturnValue(tryHandlePromise);
1265+
const startPromptSpy = vi.spyOn(codexAcpClient, "startPrompt");
1266+
vi.spyOn(codexAcpClient, "closeSession").mockResolvedValue();
1267+
1268+
const promptPromise = codexAcpAgent.prompt({
1269+
sessionId: sessionState.sessionId,
1270+
prompt: [{ type: "text", text: "normal prompt" }],
1271+
});
1272+
1273+
await vi.waitFor(() => {
1274+
expect(tryHandleSpy).toHaveBeenCalled();
1275+
});
1276+
1277+
await expect(
1278+
codexAcpAgent.unstable_closeSession({sessionId: sessionState.sessionId})
1279+
).resolves.toEqual({});
1280+
1281+
resolveTryHandle(false);
1282+
1283+
await expect(promptPromise).rejects.toThrow("Invalid request");
1284+
expect(startPromptSpy).not.toHaveBeenCalled();
1285+
});
1286+
12151287
it('skips late MCP startup updates after a session is closed', async () => {
12161288
const mockFixture = createCodexMockTestFixture();
12171289
const codexAcpAgent = mockFixture.getCodexAcpAgent();
@@ -1224,7 +1296,7 @@ describe('ACP server test', { timeout: 40_000 }, () => {
12241296
const startupPromise = new Promise((resolve) => {
12251297
resolveStartup = resolve;
12261298
});
1227-
vi.spyOn(codexAcpClient, "awaitMcpStartupResult").mockReturnValue(startupPromise as Promise<any>);
1299+
vi.spyOn(codexAcpClient, "awaitMcpServerStartup").mockReturnValue(startupPromise as Promise<any>);
12281300

12291301
// @ts-expect-error seeding private session store for focused close-session test
12301302
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
@@ -1288,4 +1360,45 @@ describe('ACP server test', { timeout: 40_000 }, () => {
12881360

12891361
expect(mockFixture.getAcpConnectionEvents([])).toEqual([]);
12901362
});
1363+
1364+
it('skips late slash command updates after a session is closed', async () => {
1365+
const mockFixture = createCodexMockTestFixture();
1366+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
1367+
const codexAcpClient = mockFixture.getCodexAcpClient();
1368+
const sessionState = createTestSessionState({
1369+
sessionId: "session-close",
1370+
});
1371+
1372+
let resolveSkills!: (response: SkillsListResponse) => void;
1373+
const skillsPromise = new Promise<SkillsListResponse>((resolve) => {
1374+
resolveSkills = resolve;
1375+
});
1376+
vi.spyOn(codexAcpClient, "listSkills").mockReturnValue(skillsPromise);
1377+
1378+
// @ts-expect-error seeding private session store for focused close-session test
1379+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
1380+
1381+
// @ts-expect-error exercising private helper through the availableCommands collaborator
1382+
const commandPromise = codexAcpAgent.availableCommands.handleCommand({ name: "skills", input: null }, sessionState);
1383+
// @ts-expect-error exercising private helper to simulate close-session cleanup
1384+
codexAcpAgent.forgetSession(sessionState.sessionId);
1385+
1386+
resolveSkills({
1387+
data: [{
1388+
cwd: "/workspace",
1389+
skills: [{
1390+
name: "build",
1391+
description: "Build the project",
1392+
shortDescription: "Build",
1393+
path: "/workspace",
1394+
scope: "user",
1395+
enabled: true,
1396+
}],
1397+
errors: [],
1398+
}],
1399+
});
1400+
1401+
await expect(commandPromise).resolves.toBe(true);
1402+
expect(mockFixture.getAcpConnectionEvents([])).toEqual([]);
1403+
});
12911404
});

‎src/__tests__/acp-test-utils.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -345,8 +345,8 @@ export async function setupPromptAndSendNotifications(
345345

346346
for (const notification of notifications) {
347347
const routedNotification = (() => {
348-
const params = notification.params as { threadId?: unknown };
349-
if (typeof params.threadId !== "string") {
348+
const params = notification.params;
349+
if ("threadId" in params && typeof params.threadId !== "string") {
350350
return notification;
351351
}
352352
return {
@@ -355,7 +355,7 @@ export async function setupPromptAndSendNotifications(
355355
...notification.params,
356356
threadId: sessionId,
357357
},
358-
} satisfies ServerNotification;
358+
} as ServerNotification;
359359
})();
360360

361361
fixture.sendServerNotification(routedNotification);

0 commit comments

Comments
 (0)