Skip to content

Commit 3ecdaed

Browse files
benbrandtbenbrandt
committed
Handle ACP request cancellation
ACP SDK v0.29.0 added support for the draft of request cancellation.
1 parent 72c9a10 commit 3ecdaed

10 files changed

Lines changed: 441 additions & 93 deletions

‎package-lock.json‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@
6161
"vitest": "^4.0.10"
6262
},
6363
"dependencies": {
64-
"@agentclientprotocol/sdk": "^0.28.1",
64+
"@agentclientprotocol/sdk": "^0.29.0",
6565
"@openai/codex": "^0.142.2",
6666
"diff": "^8.0.3",
6767
"open": "^11.0.0",

‎src/CodexAcpClient.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ export class CodexAcpClient {
309309
async runReview(
310310
sessionId: string,
311311
target: ReviewTarget,
312-
onTurnStarted?: (turnId: string) => void,
312+
onTurnStarted?: (turnId: string, threadId: string) => void,
313313
): Promise<TurnCompletedNotification> {
314314
return await this.codexClient.runReview({
315315
threadId: sessionId,

‎src/CodexAcpServer.ts‎

Lines changed: 159 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,10 @@ interface PendingTurnStart {
100100
interface ActivePrompt {
101101
completion: Promise<void>;
102102
closeSignal: Promise<null>;
103+
cancelSignal: Promise<null>;
104+
signal: AbortSignal;
105+
currentTurn: { threadId: string, turnId: string } | null;
106+
requestCancel: () => void;
103107
requestClose: () => void;
104108
complete: () => void;
105109
}
@@ -1116,17 +1120,33 @@ export class CodexAcpServer {
11161120
const closeSignal = new Promise<null>((resolve) => {
11171121
resolveCloseSignal = resolve;
11181122
});
1123+
let resolveCancelSignal: (value: null) => void = () => {};
1124+
const cancelSignal = new Promise<null>((resolve) => {
1125+
resolveCancelSignal = resolve;
1126+
});
1127+
const abortController = new AbortController();
11191128

11201129
let completed = false;
11211130
let closeRequested = false;
11221131
const activePrompt: ActivePrompt = {
11231132
completion,
11241133
closeSignal,
1134+
cancelSignal,
1135+
signal: abortController.signal,
1136+
currentTurn: null,
1137+
requestCancel: () => {
1138+
if (abortController.signal.aborted) {
1139+
return;
1140+
}
1141+
abortController.abort();
1142+
resolveCancelSignal(null);
1143+
},
11251144
requestClose: () => {
11261145
if (closeRequested) {
11271146
return;
11281147
}
11291148
closeRequested = true;
1149+
activePrompt.requestCancel();
11301150
resolveCloseSignal(null);
11311151
},
11321152
complete: () => {
@@ -1145,6 +1165,48 @@ export class CodexAcpServer {
11451165
return activePrompt;
11461166
}
11471167

1168+
private cancelBeforeTurnStarted(activePrompt: ActivePrompt): Promise<null> {
1169+
return activePrompt.cancelSignal.then(() => {
1170+
if (activePrompt.currentTurn === null) {
1171+
return null;
1172+
}
1173+
return new Promise<null>(() => {});
1174+
});
1175+
}
1176+
1177+
private observePromptRequestCancellation(
1178+
signal: AbortSignal | undefined,
1179+
sessionState: SessionState,
1180+
activePrompt: ActivePrompt,
1181+
): () => void {
1182+
if (!signal) {
1183+
return () => {};
1184+
}
1185+
1186+
const onAbort = () => {
1187+
if (this.activePrompts.get(sessionState.sessionId) !== activePrompt) {
1188+
return;
1189+
}
1190+
logger.log("Prompt request cancelled", {sessionId: sessionState.sessionId});
1191+
activePrompt.requestCancel();
1192+
const turn = activePrompt.currentTurn;
1193+
if (!turn) {
1194+
return;
1195+
}
1196+
void this.interruptPromptTurn(turn, "Cancel").catch((err) => {
1197+
logger.error("Prompt request cancellation failed to interrupt turn", err);
1198+
});
1199+
};
1200+
1201+
if (signal.aborted) {
1202+
onAbort();
1203+
return () => {};
1204+
}
1205+
1206+
signal.addEventListener("abort", onAbort, {once: true});
1207+
return () => signal.removeEventListener("abort", onAbort);
1208+
}
1209+
11481210
private createPendingTurnStart(): PendingTurnStart {
11491211
let resolve: (turnId: string | null) => void = () => {};
11501212
const promise = new Promise<string | null>((innerResolve) => {
@@ -1153,26 +1215,39 @@ export class CodexAcpServer {
11531215
return {promise, resolve};
11541216
}
11551217

1156-
private interruptLateStartedTurn(sessionId: string, turnId: string): void {
1218+
private async interruptPromptTurn(
1219+
turn: { threadId: string, turnId: string },
1220+
requestName: "Cancel" | "Close",
1221+
): Promise<void> {
11571222
this.codexAcpClient.markTurnStale({
1158-
threadId: sessionId,
1159-
turnId,
1223+
threadId: turn.threadId,
1224+
turnId: turn.turnId,
11601225
});
1161-
void this.runWithProcessCheck(() => this.codexAcpClient.turnInterrupt({
1162-
threadId: sessionId,
1163-
turnId,
1164-
})).catch((err) => {
1165-
logger.error(`Close - late turnInterrupt failed`, err);
1166-
}).finally(() => {
1226+
try {
1227+
await this.runWithProcessCheck(() => this.codexAcpClient.turnInterrupt({
1228+
threadId: turn.threadId,
1229+
turnId: turn.turnId,
1230+
}));
1231+
logger.log(`${requestName} - turnInterrupt succeeded`, {
1232+
sessionId: turn.threadId,
1233+
currentTurnId: turn.turnId,
1234+
});
1235+
} catch (err) {
1236+
logger.error(`${requestName} - turnInterrupt failed`, err);
1237+
} finally {
11671238
this.codexAcpClient.resolveTurnInterrupted({
1168-
threadId: sessionId,
1169-
turnId,
1239+
threadId: turn.threadId,
1240+
turnId: turn.turnId,
11701241
});
1171-
});
1242+
}
1243+
}
1244+
1245+
private interruptLateStartedTurn(turn: { threadId: string, turnId: string }): void {
1246+
void this.interruptPromptTurn(turn, "Close");
11721247
}
11731248

1174-
private promptIsClosedOrStale(sessionId: string, activePrompt: ActivePrompt): boolean {
1175-
return this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId);
1249+
private promptShouldStop(sessionId: string, activePrompt: ActivePrompt): boolean {
1250+
return activePrompt.signal.aborted || this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId);
11761251
}
11771252

11781253
private async interruptSessionTurn(
@@ -1242,7 +1317,7 @@ export class CodexAcpServer {
12421317
return turnId;
12431318
}
12441319

1245-
async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
1320+
async prompt(params: acp.PromptRequest, signal?: AbortSignal): Promise<acp.PromptResponse> {
12461321
logger.log("Prompt received", {
12471322
sessionId: params.sessionId,
12481323
prompt: params.prompt,
@@ -1252,11 +1327,12 @@ export class CodexAcpServer {
12521327
sessionState.lastTokenUsage = null;
12531328
const activePrompt = this.trackActivePrompt(params.sessionId);
12541329
let pendingTurnStart: PendingTurnStart | null = null;
1330+
const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt);
12551331

12561332
try {
12571333
const eventHandler = new CodexEventHandler(this.connection, sessionState);
1258-
const approvalHandler = new CodexApprovalHandler(this.connection, sessionState);
1259-
const elicitationHandler = new CodexElicitationHandler(this.connection, sessionState);
1334+
const approvalHandler = new CodexApprovalHandler(this.connection, sessionState, activePrompt.signal);
1335+
const elicitationHandler = new CodexElicitationHandler(this.connection, sessionState, activePrompt.signal);
12601336
await this.codexAcpClient.subscribeToSessionEvents(params.sessionId,
12611337
(event) => {
12621338
elicitationHandler.handleNotification(event);
@@ -1265,28 +1341,40 @@ export class CodexAcpServer {
12651341
approvalHandler,
12661342
elicitationHandler);
12671343

1268-
const commandResult = await this.availableCommands.tryHandleCommand(params.prompt, sessionState);
1344+
if (activePrompt.signal.aborted) {
1345+
return this.cancelledPromptResponse(sessionState);
1346+
}
1347+
1348+
const commandPromise = this.availableCommands.tryHandleCommand(params.prompt, sessionState, {
1349+
onTurnStarted: (turnId, threadId) => {
1350+
const turn = {threadId, turnId};
1351+
activePrompt.currentTurn = turn;
1352+
if (this.promptShouldStop(params.sessionId, activePrompt)) {
1353+
this.interruptLateStartedTurn(turn);
1354+
return;
1355+
}
1356+
sessionState.currentTurnId = turnId;
1357+
},
1358+
});
1359+
void commandPromise.catch((err) => {
1360+
if (this.activePrompts.get(params.sessionId) !== activePrompt) {
1361+
logger.error(`Command for cancelled prompt ${params.sessionId} failed after prompt returned`, err);
1362+
}
1363+
});
1364+
const commandResult = await Promise.race([
1365+
commandPromise,
1366+
activePrompt.closeSignal,
1367+
this.cancelBeforeTurnStarted(activePrompt),
1368+
]);
1369+
if (commandResult === null) {
1370+
return this.cancelledPromptResponse(sessionState);
1371+
}
12691372
if (commandResult.handled) {
12701373
logger.log("Prompt handled by a command");
12711374
await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
12721375
if (commandResult.turnCompleted?.turn.status === "interrupted") {
1273-
if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
1274-
await this.connection.notify(acp.methods.client.session.update, {
1275-
sessionId: params.sessionId,
1276-
update: {
1277-
sessionUpdate: "agent_message_chunk",
1278-
content: {
1279-
type: "text",
1280-
text: "*Conversation interrupted*"
1281-
}
1282-
}
1283-
});
1284-
}
1285-
return {
1286-
stopReason: "cancelled",
1287-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1288-
_meta: this.buildQuotaMeta(sessionState),
1289-
};
1376+
await this.notifyConversationInterrupted(params.sessionId);
1377+
return this.cancelledPromptResponse(sessionState);
12901378
}
12911379
const error = eventHandler.getFailure();
12921380
if (error) {
@@ -1301,11 +1389,7 @@ export class CodexAcpServer {
13011389
}
13021390

13031391
if (this.sessionIsClosing(params.sessionId)) {
1304-
return {
1305-
stopReason: "cancelled",
1306-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1307-
_meta: this.buildQuotaMeta(sessionState),
1308-
};
1392+
return this.cancelledPromptResponse(sessionState);
13091393
}
13101394

13111395
const modelId = ModelId.fromString(sessionState.currentModelId);
@@ -1340,54 +1424,37 @@ export class CodexAcpServer {
13401424
sessionState.cwd,
13411425
sessionState.additionalDirectories,
13421426
(turnId) => {
1343-
if (this.promptIsClosedOrStale(params.sessionId, activePrompt)) {
1344-
this.interruptLateStartedTurn(params.sessionId, turnId);
1427+
const turn = {threadId: params.sessionId, turnId};
1428+
activePrompt.currentTurn = turn;
1429+
if (this.promptShouldStop(params.sessionId, activePrompt)) {
1430+
this.interruptLateStartedTurn(turn);
13451431
return;
13461432
}
13471433
sessionState.currentTurnId = turnId;
13481434
pendingTurnStart?.resolve(turnId);
13491435
},
1350-
() => this.promptIsClosedOrStale(params.sessionId, activePrompt),
1436+
() => this.promptShouldStop(params.sessionId, activePrompt),
13511437
));
13521438
void sendPromptPromise.catch((err) => {
13531439
if (this.activePrompts.get(params.sessionId) !== activePrompt) {
1354-
logger.error(`Prompt for closed session ${params.sessionId} failed after close`, err);
1440+
logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err);
13551441
}
13561442
});
13571443
const turnCompleted = await Promise.race([
13581444
sendPromptPromise,
13591445
activePrompt.closeSignal,
1446+
this.cancelBeforeTurnStarted(activePrompt),
13601447
]);
13611448

13621449
if (turnCompleted === null) {
1363-
return {
1364-
stopReason: "cancelled",
1365-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1366-
_meta: this.buildQuotaMeta(sessionState),
1367-
};
1450+
return this.cancelledPromptResponse(sessionState);
13681451
}
13691452

13701453
await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
13711454

1372-
// Check if turn was interrupted (cancelled)
13731455
if (turnCompleted.turn.status === "interrupted") {
1374-
if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
1375-
await this.connection.notify(acp.methods.client.session.update, {
1376-
sessionId: params.sessionId,
1377-
update: {
1378-
sessionUpdate: "agent_message_chunk",
1379-
content: {
1380-
type: "text",
1381-
text: "*Conversation interrupted*"
1382-
}
1383-
}
1384-
});
1385-
}
1386-
return {
1387-
stopReason: "cancelled",
1388-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1389-
_meta: this.buildQuotaMeta(sessionState),
1390-
};
1456+
await this.notifyConversationInterrupted(params.sessionId);
1457+
return this.cancelledPromptResponse(sessionState);
13911458
}
13921459

13931460
const error = eventHandler.getFailure();
@@ -1406,6 +1473,7 @@ export class CodexAcpServer {
14061473
throw err;
14071474
} finally {
14081475
logger.log("Prompt completed", {sessionId: params.sessionId});
1476+
disposePromptRequestCancellation();
14091477
sessionState.currentTurnId = null;
14101478
if (pendingTurnStart !== null && this.pendingTurnStarts.get(params.sessionId) === pendingTurnStart) {
14111479
this.pendingTurnStarts.delete(params.sessionId);
@@ -1415,6 +1483,30 @@ export class CodexAcpServer {
14151483
}
14161484
}
14171485

1486+
private cancelledPromptResponse(sessionState: SessionState): acp.PromptResponse {
1487+
return {
1488+
stopReason: "cancelled",
1489+
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1490+
_meta: this.buildQuotaMeta(sessionState),
1491+
};
1492+
}
1493+
1494+
private async notifyConversationInterrupted(sessionId: string): Promise<void> {
1495+
if (this.sessionIsClosing(sessionId) || !this.sessions.has(sessionId)) {
1496+
return;
1497+
}
1498+
await this.connection.notify(acp.methods.client.session.update, {
1499+
sessionId,
1500+
update: {
1501+
sessionUpdate: "agent_message_chunk",
1502+
content: {
1503+
type: "text",
1504+
text: "*Conversation interrupted*"
1505+
}
1506+
}
1507+
});
1508+
}
1509+
14181510
private buildQuotaMeta(sessionState: SessionState): { quota: QuotaMeta } {
14191511
const lastTokenUsage = sessionState.lastTokenUsage;
14201512

0 commit comments

Comments
 (0)