Skip to content

Commit b46d043

Browse files
benbrandtbenbrandt
committed
Better handle the async current turns
1 parent 625c5b7 commit b46d043

14 files changed

Lines changed: 449 additions & 121 deletions

‎src/CodexAcpClient.ts‎

Lines changed: 48 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,12 +31,15 @@ import type {
3131
SkillsListResponse,
3232
Thread,
3333
ThreadSourceKind,
34+
ThreadUnsubscribeResponse,
3435
TurnCompletedNotification,
3536
UserInput,
3637
} from "./app-server/v2";
3738
import packageJson from "../package.json";
3839
import type {AuthenticationLogoutResponse, AuthenticationStatusResponse} from "./AcpExtensions";
3940

41+
const CLOSE_TURN_COMPLETION_TIMEOUT_MS = 10_000;
42+
4043
/**
4144
* API for accessing the Codex App Server using ACP requests.
4245
* Converts ACP requests into corresponding app-server operations.
@@ -382,17 +385,34 @@ export class CodexAcpClient {
382385
}
383386

384387
async closeSession(sessionId: string, currentTurnId: string | null): Promise<void> {
385-
if (currentTurnId && !this.codexClient.hasTurnCompleted(sessionId, currentTurnId)) {
386-
await this.codexClient.turnInterrupt({
387-
threadId: sessionId,
388-
turnId: currentTurnId,
389-
});
390-
await this.codexClient.awaitTurnCompleted(sessionId, currentTurnId);
391-
}
388+
try {
389+
if (currentTurnId && !this.codexClient.hasTurnCompleted(sessionId, currentTurnId)) {
390+
try {
391+
await this.codexClient.turnInterrupt({
392+
threadId: sessionId,
393+
turnId: currentTurnId,
394+
});
395+
await this.waitForTurnCompletionDuringClose(sessionId, currentTurnId);
396+
} catch (err) {
397+
logger.error(`Failed to interrupt active turn while closing session ${sessionId}`, err);
398+
}
399+
}
392400

393-
await this.codexClient.threadUnsubscribe({threadId: sessionId});
394-
this.unsubscribeFromSessionEvents(sessionId);
395-
this.codexClient.clearThreadState(sessionId);
401+
try {
402+
const response = await this.codexClient.threadUnsubscribe({threadId: sessionId}) as ThreadUnsubscribeResponse | undefined;
403+
if (response?.status && response.status !== "unsubscribed") {
404+
logger.log("Thread unsubscribe completed without an active subscription", {
405+
sessionId,
406+
status: response.status,
407+
});
408+
}
409+
} catch (err) {
410+
logger.error(`Failed to unsubscribe thread while closing session ${sessionId}`, err);
411+
}
412+
} finally {
413+
this.unsubscribeFromSessionEvents(sessionId);
414+
this.codexClient.clearThreadState(sessionId);
415+
}
396416
}
397417

398418
async startPrompt(
@@ -428,6 +448,24 @@ export class CodexAcpClient {
428448
return await this.codexClient.awaitTurnCompleted(sessionId, turnId);
429449
}
430450

451+
private async waitForTurnCompletionDuringClose(sessionId: string, turnId: string): Promise<void> {
452+
let timeout: ReturnType<typeof setTimeout> | null = null;
453+
try {
454+
await Promise.race([
455+
this.codexClient.awaitTurnCompleted(sessionId, turnId),
456+
new Promise<never>((_, reject) => {
457+
timeout = setTimeout(() => {
458+
reject(new Error(`Timed out waiting for turn ${turnId} to complete during close`));
459+
}, CLOSE_TURN_COMPLETION_TIMEOUT_MS);
460+
}),
461+
]);
462+
} finally {
463+
if (timeout) {
464+
clearTimeout(timeout);
465+
}
466+
}
467+
}
468+
431469
hasTurnCompleted(sessionId: string, turnId: string): boolean {
432470
return this.codexClient.hasTurnCompleted(sessionId, turnId);
433471
}

‎src/CodexAcpServer.ts‎

Lines changed: 73 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ export class CodexAcpServer implements acp.Agent {
7070
private readonly sessions: Map<string, SessionState>;
7171
private readonly pendingMcpStartupSessions: Map<string, PendingMcpStartupSession>;
7272
private readonly closingSessions: Set<string>;
73+
private readonly closeSessionPromises: Map<string, Promise<void>>;
7374

7475
constructor(
7576
connection: acp.AgentSideConnection,
@@ -80,6 +81,7 @@ export class CodexAcpServer implements acp.Agent {
8081
this.sessions = new Map();
8182
this.pendingMcpStartupSessions = new Map();
8283
this.closingSessions = new Set();
84+
this.closeSessionPromises = new Map();
8385
this.connection = connection;
8486
this.codexAcpClient = codexAcpClient;
8587
this.defaultAuthRequest = defaultAuthRequest ?? null;
@@ -246,24 +248,40 @@ export class CodexAcpServer implements acp.Agent {
246248

247249
async closeSession(params: acp.CloseSessionRequest): Promise<acp.CloseSessionResponse> {
248250
logger.log("Closing session...", {sessionId: params.sessionId});
249-
const sessionState = this.getSessionState(params.sessionId);
250-
this.closingSessions.add(params.sessionId);
251+
const existingClose = this.closeSessionPromises.get(params.sessionId);
252+
if (existingClose) {
253+
await existingClose;
254+
logger.log("Session close already completed", {sessionId: params.sessionId});
255+
return {};
256+
}
251257

258+
const closePromise = this.closeSessionInternal(params.sessionId);
259+
this.closeSessionPromises.set(params.sessionId, closePromise);
252260
try {
253-
const activeTurnId = await this.resolveActiveTurnId(sessionState);
254-
await this.runWithProcessCheck(() =>
255-
this.codexAcpClient.closeSession(params.sessionId, activeTurnId)
256-
);
257-
this.forgetSession(params.sessionId);
258-
} catch (err) {
259-
this.closingSessions.delete(params.sessionId);
260-
throw err;
261+
await closePromise;
262+
} finally {
263+
this.closeSessionPromises.delete(params.sessionId);
261264
}
262265

263266
logger.log("Session closed", {sessionId: params.sessionId});
264267
return {};
265268
}
266269

270+
private async closeSessionInternal(sessionId: string): Promise<void> {
271+
logger.log("Closing session internals...", {sessionId});
272+
const sessionState = this.getSessionState(sessionId);
273+
this.closingSessions.add(sessionId);
274+
275+
try {
276+
const activeTurnId = await this.resolveActiveTurnId(sessionState);
277+
await this.runWithProcessCheck(() =>
278+
this.codexAcpClient.closeSession(sessionId, activeTurnId)
279+
);
280+
} finally {
281+
this.forgetSession(sessionId);
282+
}
283+
}
284+
267285
async newSession(
268286
params: acp.NewSessionRequest,
269287
): Promise<acp.NewSessionResponse> {
@@ -305,6 +323,9 @@ export class CodexAcpServer implements acp.Agent {
305323
});
306324
const sessionState = this.sessions.get(_params.sessionId);
307325
if (!sessionState) throw new Error(`Session ${_params.sessionId} not found`);
326+
if (this.closingSessions.has(_params.sessionId)) {
327+
throw RequestError.invalidRequest("Session is closing");
328+
}
308329

309330
const newMode = AgentMode.find(_params.modeId);
310331
if (!newMode) {
@@ -321,6 +342,9 @@ export class CodexAcpServer implements acp.Agent {
321342
});
322343
const sessionState = this.sessions.get(params.sessionId);
323344
if (!sessionState) throw new Error(`Session ${params.sessionId} not found`);
345+
if (this.closingSessions.has(params.sessionId)) {
346+
throw RequestError.invalidRequest("Session is closing");
347+
}
324348

325349
const requestedModelId= ModelId.fromString(params.modelId);
326350
const requestedModelName = requestedModelId.model;
@@ -685,10 +709,8 @@ export class CodexAcpServer implements acp.Agent {
685709
}
686710
}
687711

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-
}
712+
private isSessionClosedOrClosing(sessionId: string): boolean {
713+
return this.closingSessions.has(sessionId) || !this.sessions.has(sessionId);
692714
}
693715

694716
private forgetSession(sessionId: string): void {
@@ -781,8 +803,9 @@ export class CodexAcpServer implements acp.Agent {
781803
prompt: params.prompt,
782804
});
783805
const sessionState = this.getSessionState(params.sessionId);
784-
const requireTrackedSession = this.sessions.has(params.sessionId);
785-
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
806+
if (this.isSessionClosedOrClosing(params.sessionId)) {
807+
return this.buildCancelledPromptResponse(sessionState);
808+
}
786809
sessionState.currentTurnId = null;
787810
sessionState.lastTokenUsage = null;
788811

@@ -800,18 +823,25 @@ export class CodexAcpServer implements acp.Agent {
800823
},
801824
approvalHandler,
802825
elicitationHandler);
803-
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
826+
if (this.isSessionClosedOrClosing(params.sessionId)) {
827+
return this.buildCancelledPromptResponse(sessionState);
828+
}
804829

805830
if (await this.availableCommands.tryHandle(params.prompt, sessionState)) {
806831
logger.log("Prompt handled by a command");
832+
if (this.isSessionClosedOrClosing(params.sessionId)) {
833+
return this.buildCancelledPromptResponse(sessionState);
834+
}
807835
return {
808836
stopReason: "end_turn",
809837
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
810838
_meta: this.buildQuotaMeta(sessionState),
811839
};
812840
}
813841

814-
this.ensureSessionIsActive(params.sessionId, requireTrackedSession);
842+
if (this.isSessionClosedOrClosing(params.sessionId)) {
843+
return this.buildCancelledPromptResponse(sessionState);
844+
}
815845

816846
const modelId = ModelId.fromString(sessionState.currentModelId);
817847
const modelLacksReasoning = sessionState.supportedReasoningEfforts.length > 0
@@ -832,8 +862,20 @@ export class CodexAcpServer implements acp.Agent {
832862
const turnIdPromise = this.runWithProcessCheck(
833863
() => this.codexAcpClient.startPrompt(params, agentMode, modelId, disableSummary, sessionState.cwd)
834864
);
865+
turnIdPromise.catch(() => {
866+
// The prompt path awaits this promise, but close also reads currentTurnId.
867+
// Attach a handler immediately so future readers cannot create an unhandled rejection.
868+
});
835869
sessionState.currentTurnId = turnIdPromise;
836-
const turnId = await turnIdPromise;
870+
let turnId: string;
871+
try {
872+
turnId = await turnIdPromise;
873+
} catch (err) {
874+
if (this.isSessionClosedOrClosing(params.sessionId)) {
875+
return this.buildCancelledPromptResponse(sessionState);
876+
}
877+
throw err;
878+
}
837879
sessionState.currentTurnId = turnId;
838880

839881
const turnCompleted = await this.runWithProcessCheck(
@@ -854,11 +896,7 @@ export class CodexAcpServer implements acp.Agent {
854896
}
855897
});
856898
}
857-
return {
858-
stopReason: "cancelled",
859-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
860-
_meta: this.buildQuotaMeta(sessionState),
861-
};
899+
return this.buildCancelledPromptResponse(sessionState);
862900
}
863901

864902
const error = eventHandler.getFailure()
@@ -873,6 +911,9 @@ export class CodexAcpServer implements acp.Agent {
873911
_meta: this.buildQuotaMeta(sessionState),
874912
};
875913
} catch (err) {
914+
if (this.isSessionClosedOrClosing(params.sessionId)) {
915+
return this.buildCancelledPromptResponse(sessionState);
916+
}
876917
logger.error(`Prompt for session ${params.sessionId} failed`, err);
877918
throw err;
878919
} finally {
@@ -907,6 +948,14 @@ export class CodexAcpServer implements acp.Agent {
907948
return toPromptUsage(lastTokenUsage);
908949
}
909950

951+
private buildCancelledPromptResponse(sessionState: SessionState): acp.PromptResponse {
952+
return {
953+
stopReason: "cancelled",
954+
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
955+
_meta: this.buildQuotaMeta(sessionState),
956+
};
957+
}
958+
910959
private async runWithProcessCheck<T>(operation: () => Promise<T>): Promise<T> {
911960
try {
912961
return await operation();

0 commit comments

Comments
 (0)