Skip to content

Commit 3486dad

Browse files
benbrandtbenbrandt
committed
Cleanups
1 parent c39706a commit 3486dad

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
@@ -392,6 +392,7 @@ export class CodexAcpClient {
392392

393393
await this.codexClient.threadUnsubscribe({threadId: sessionId});
394394
this.unsubscribeFromSessionEvents(sessionId);
395+
this.codexClient.clearThreadState(sessionId);
395396
}
396397

397398
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
@@ -658,6 +658,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
658658

659659
const sessionState: SessionState = createTestSessionState();
660660
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
661+
// @ts-expect-error seeding private session store for slash-command test
662+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
661663

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

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

674678
// @ts-expect-error - exercising private helper
675679
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "logout", input: null }, sessionState);
@@ -696,6 +700,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
696700
}]
697701
};
698702
const skillsSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "listSkills").mockResolvedValue(skillsResponse);
703+
// @ts-expect-error seeding private session store for slash-command test
704+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
699705

700706
// @ts-expect-error - exercising private helper
701707
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "skills", input: null }, sessionState);
@@ -731,6 +737,8 @@ describe('ACP server test', { timeout: 40_000 }, () => {
731737
nextCursor: null
732738
};
733739
const mcpSpy = vi.spyOn(mockFixture.getCodexAcpClient(), "listMcpServers").mockResolvedValue(mcpResponse);
740+
// @ts-expect-error seeding private session store for slash-command test
741+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
734742

735743
// @ts-expect-error - exercising private helper
736744
const handled = await codexAcpAgent.availableCommands.handleCommand({ name: "mcp", input: null }, sessionState);
@@ -820,7 +828,10 @@ describe('ACP server test', { timeout: 40_000 }, () => {
820828
durationMs: null,
821829
}
822830
});
823-
vi.spyOn(mockFixture.getCodexAcpAgent(), "getSessionState").mockReturnValue(sessionState);
831+
const codexAcpAgent = mockFixture.getCodexAcpAgent();
832+
vi.spyOn(codexAcpAgent, "getSessionState").mockReturnValue(sessionState);
833+
// @ts-expect-error seeding private session store for prompt test
834+
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
824835
return { mockFixture, sessionState, turnStartSpy };
825836
}
826837

@@ -1064,6 +1075,24 @@ describe('ACP server test', { timeout: 40_000 }, () => {
10641075
expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(true);
10651076
// @ts-expect-error verifying test-only access to private maps
10661077
expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(true);
1078+
// @ts-expect-error verifying test-only access to private maps
1079+
codexAppServerClient.turnCompletedResolvers.set("session-close", [{
1080+
turnId: "other-turn",
1081+
resolve: vi.fn(),
1082+
}]);
1083+
// @ts-expect-error verifying test-only access to private maps
1084+
codexAppServerClient.lastTurnCompletedByThread.set("session-close", {
1085+
threadId: "session-close",
1086+
turn: {
1087+
id: "completed-turn",
1088+
items: [],
1089+
status: "completed",
1090+
error: null,
1091+
startedAt: null,
1092+
completedAt: null,
1093+
durationMs: null,
1094+
},
1095+
});
10671096

10681097
await codexAcpClient.closeSession("session-close", "turn-close");
10691098

@@ -1088,6 +1117,10 @@ describe('ACP server test', { timeout: 40_000 }, () => {
10881117
expect(codexAppServerClient.approvalHandlers.has("session-close")).toBe(false);
10891118
// @ts-expect-error verifying test-only access to private maps
10901119
expect(codexAppServerClient.elicitationHandlers.has("session-close")).toBe(false);
1120+
// @ts-expect-error verifying test-only access to private maps
1121+
expect(codexAppServerClient.turnCompletedResolvers.has("session-close")).toBe(false);
1122+
// @ts-expect-error verifying test-only access to private maps
1123+
expect(codexAppServerClient.lastTurnCompletedByThread.has("session-close")).toBe(false);
10911124
});
10921125

10931126
it('waits for the matching turn completion before resolving a prompt', async () => {
@@ -1221,6 +1254,45 @@ describe('ACP server test', { timeout: 40_000 }, () => {
12211254
expect(closeSpy).toHaveBeenCalledWith(sessionState.sessionId, "turn-pending");
12221255
});
12231256

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

12381310
// @ts-expect-error seeding private session store for focused close-session test
12391311
codexAcpAgent.sessions.set(sessionState.sessionId, sessionState);
@@ -1297,4 +1369,45 @@ describe('ACP server test', { timeout: 40_000 }, () => {
12971369

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

‎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)