Skip to content

Commit f4163f0

Browse files
benbrandtbenbrandt
committed
Handle ACP request cancellation
ACP SDK v0.29.0 added support for the draft of request cancellation.
1 parent b93225f commit f4163f0

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.0",
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
@@ -292,7 +292,7 @@ export class CodexAcpClient {
292292
async runReview(
293293
sessionId: string,
294294
target: ReviewTarget,
295-
onTurnStarted?: (turnId: string) => void,
295+
onTurnStarted?: (turnId: string, threadId: string) => void,
296296
): Promise<TurnCompletedNotification> {
297297
return await this.codexClient.runReview({
298298
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
}
@@ -1146,17 +1150,33 @@ export class CodexAcpServer {
11461150
const closeSignal = new Promise<null>((resolve) => {
11471151
resolveCloseSignal = resolve;
11481152
});
1153+
let resolveCancelSignal: (value: null) => void = () => {};
1154+
const cancelSignal = new Promise<null>((resolve) => {
1155+
resolveCancelSignal = resolve;
1156+
});
1157+
const abortController = new AbortController();
11491158

11501159
let completed = false;
11511160
let closeRequested = false;
11521161
const activePrompt: ActivePrompt = {
11531162
completion,
11541163
closeSignal,
1164+
cancelSignal,
1165+
signal: abortController.signal,
1166+
currentTurn: null,
1167+
requestCancel: () => {
1168+
if (abortController.signal.aborted) {
1169+
return;
1170+
}
1171+
abortController.abort();
1172+
resolveCancelSignal(null);
1173+
},
11551174
requestClose: () => {
11561175
if (closeRequested) {
11571176
return;
11581177
}
11591178
closeRequested = true;
1179+
activePrompt.requestCancel();
11601180
resolveCloseSignal(null);
11611181
},
11621182
complete: () => {
@@ -1175,6 +1195,48 @@ export class CodexAcpServer {
11751195
return activePrompt;
11761196
}
11771197

1198+
private cancelBeforeTurnStarted(activePrompt: ActivePrompt): Promise<null> {
1199+
return activePrompt.cancelSignal.then(() => {
1200+
if (activePrompt.currentTurn === null) {
1201+
return null;
1202+
}
1203+
return new Promise<null>(() => {});
1204+
});
1205+
}
1206+
1207+
private observePromptRequestCancellation(
1208+
signal: AbortSignal | undefined,
1209+
sessionState: SessionState,
1210+
activePrompt: ActivePrompt,
1211+
): () => void {
1212+
if (!signal) {
1213+
return () => {};
1214+
}
1215+
1216+
const onAbort = () => {
1217+
if (this.activePrompts.get(sessionState.sessionId) !== activePrompt) {
1218+
return;
1219+
}
1220+
logger.log("Prompt request cancelled", {sessionId: sessionState.sessionId});
1221+
activePrompt.requestCancel();
1222+
const turn = activePrompt.currentTurn;
1223+
if (!turn) {
1224+
return;
1225+
}
1226+
void this.interruptPromptTurn(turn, "Cancel").catch((err) => {
1227+
logger.error("Prompt request cancellation failed to interrupt turn", err);
1228+
});
1229+
};
1230+
1231+
if (signal.aborted) {
1232+
onAbort();
1233+
return () => {};
1234+
}
1235+
1236+
signal.addEventListener("abort", onAbort, {once: true});
1237+
return () => signal.removeEventListener("abort", onAbort);
1238+
}
1239+
11781240
private createPendingTurnStart(): PendingTurnStart {
11791241
let resolve: (turnId: string | null) => void = () => {};
11801242
const promise = new Promise<string | null>((innerResolve) => {
@@ -1183,26 +1245,39 @@ export class CodexAcpServer {
11831245
return {promise, resolve};
11841246
}
11851247

1186-
private interruptLateStartedTurn(sessionId: string, turnId: string): void {
1248+
private async interruptPromptTurn(
1249+
turn: { threadId: string, turnId: string },
1250+
requestName: "Cancel" | "Close",
1251+
): Promise<void> {
11871252
this.codexAcpClient.markTurnStale({
1188-
threadId: sessionId,
1189-
turnId,
1253+
threadId: turn.threadId,
1254+
turnId: turn.turnId,
11901255
});
1191-
void this.runWithProcessCheck(() => this.codexAcpClient.turnInterrupt({
1192-
threadId: sessionId,
1193-
turnId,
1194-
})).catch((err) => {
1195-
logger.error(`Close - late turnInterrupt failed`, err);
1196-
}).finally(() => {
1256+
try {
1257+
await this.runWithProcessCheck(() => this.codexAcpClient.turnInterrupt({
1258+
threadId: turn.threadId,
1259+
turnId: turn.turnId,
1260+
}));
1261+
logger.log(`${requestName} - turnInterrupt succeeded`, {
1262+
sessionId: turn.threadId,
1263+
currentTurnId: turn.turnId,
1264+
});
1265+
} catch (err) {
1266+
logger.error(`${requestName} - turnInterrupt failed`, err);
1267+
} finally {
11971268
this.codexAcpClient.resolveTurnInterrupted({
1198-
threadId: sessionId,
1199-
turnId,
1269+
threadId: turn.threadId,
1270+
turnId: turn.turnId,
12001271
});
1201-
});
1272+
}
1273+
}
1274+
1275+
private interruptLateStartedTurn(turn: { threadId: string, turnId: string }): void {
1276+
void this.interruptPromptTurn(turn, "Close");
12021277
}
12031278

1204-
private promptIsClosedOrStale(sessionId: string, activePrompt: ActivePrompt): boolean {
1205-
return this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId);
1279+
private promptShouldStop(sessionId: string, activePrompt: ActivePrompt): boolean {
1280+
return activePrompt.signal.aborted || this.activePrompts.get(sessionId) !== activePrompt || this.sessionIsClosing(sessionId);
12061281
}
12071282

12081283
private async interruptSessionTurn(
@@ -1272,7 +1347,7 @@ export class CodexAcpServer {
12721347
return turnId;
12731348
}
12741349

1275-
async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
1350+
async prompt(params: acp.PromptRequest, signal?: AbortSignal): Promise<acp.PromptResponse> {
12761351
logger.log("Prompt received", {
12771352
sessionId: params.sessionId,
12781353
prompt: params.prompt,
@@ -1282,11 +1357,12 @@ export class CodexAcpServer {
12821357
sessionState.lastTokenUsage = null;
12831358
const activePrompt = this.trackActivePrompt(params.sessionId);
12841359
let pendingTurnStart: PendingTurnStart | null = null;
1360+
const disposePromptRequestCancellation = this.observePromptRequestCancellation(signal, sessionState, activePrompt);
12851361

12861362
try {
12871363
const eventHandler = new CodexEventHandler(this.connection, sessionState);
1288-
const approvalHandler = new CodexApprovalHandler(this.connection, sessionState);
1289-
const elicitationHandler = new CodexElicitationHandler(this.connection, sessionState);
1364+
const approvalHandler = new CodexApprovalHandler(this.connection, sessionState, activePrompt.signal);
1365+
const elicitationHandler = new CodexElicitationHandler(this.connection, sessionState, activePrompt.signal);
12901366
await this.codexAcpClient.subscribeToSessionEvents(params.sessionId,
12911367
(event) => {
12921368
elicitationHandler.handleNotification(event);
@@ -1295,28 +1371,40 @@ export class CodexAcpServer {
12951371
approvalHandler,
12961372
elicitationHandler);
12971373

1298-
const commandResult = await this.availableCommands.tryHandleCommand(params.prompt, sessionState);
1374+
if (activePrompt.signal.aborted) {
1375+
return this.cancelledPromptResponse(sessionState);
1376+
}
1377+
1378+
const commandPromise = this.availableCommands.tryHandleCommand(params.prompt, sessionState, {
1379+
onTurnStarted: (turnId, threadId) => {
1380+
const turn = {threadId, turnId};
1381+
activePrompt.currentTurn = turn;
1382+
if (this.promptShouldStop(params.sessionId, activePrompt)) {
1383+
this.interruptLateStartedTurn(turn);
1384+
return;
1385+
}
1386+
sessionState.currentTurnId = turnId;
1387+
},
1388+
});
1389+
void commandPromise.catch((err) => {
1390+
if (this.activePrompts.get(params.sessionId) !== activePrompt) {
1391+
logger.error(`Command for cancelled prompt ${params.sessionId} failed after prompt returned`, err);
1392+
}
1393+
});
1394+
const commandResult = await Promise.race([
1395+
commandPromise,
1396+
activePrompt.closeSignal,
1397+
this.cancelBeforeTurnStarted(activePrompt),
1398+
]);
1399+
if (commandResult === null) {
1400+
return this.cancelledPromptResponse(sessionState);
1401+
}
12991402
if (commandResult.handled) {
13001403
logger.log("Prompt handled by a command");
13011404
await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
13021405
if (commandResult.turnCompleted?.turn.status === "interrupted") {
1303-
if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
1304-
await this.connection.notify(acp.methods.client.session.update, {
1305-
sessionId: params.sessionId,
1306-
update: {
1307-
sessionUpdate: "agent_message_chunk",
1308-
content: {
1309-
type: "text",
1310-
text: "*Conversation interrupted*"
1311-
}
1312-
}
1313-
});
1314-
}
1315-
return {
1316-
stopReason: "cancelled",
1317-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1318-
_meta: this.buildQuotaMeta(sessionState),
1319-
};
1406+
await this.notifyConversationInterrupted(params.sessionId);
1407+
return this.cancelledPromptResponse(sessionState);
13201408
}
13211409
const error = eventHandler.getFailure()
13221410
if (error) {
@@ -1331,11 +1419,7 @@ export class CodexAcpServer {
13311419
}
13321420

13331421
if (this.sessionIsClosing(params.sessionId)) {
1334-
return {
1335-
stopReason: "cancelled",
1336-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1337-
_meta: this.buildQuotaMeta(sessionState),
1338-
};
1422+
return this.cancelledPromptResponse(sessionState);
13391423
}
13401424

13411425
const modelId = ModelId.fromString(sessionState.currentModelId);
@@ -1370,54 +1454,37 @@ export class CodexAcpServer {
13701454
sessionState.cwd,
13711455
sessionState.additionalDirectories,
13721456
(turnId) => {
1373-
if (this.promptIsClosedOrStale(params.sessionId, activePrompt)) {
1374-
this.interruptLateStartedTurn(params.sessionId, turnId);
1457+
const turn = {threadId: params.sessionId, turnId};
1458+
activePrompt.currentTurn = turn;
1459+
if (this.promptShouldStop(params.sessionId, activePrompt)) {
1460+
this.interruptLateStartedTurn(turn);
13751461
return;
13761462
}
13771463
sessionState.currentTurnId = turnId;
13781464
pendingTurnStart?.resolve(turnId);
13791465
},
1380-
() => this.promptIsClosedOrStale(params.sessionId, activePrompt),
1466+
() => this.promptShouldStop(params.sessionId, activePrompt),
13811467
));
13821468
void sendPromptPromise.catch((err) => {
13831469
if (this.activePrompts.get(params.sessionId) !== activePrompt) {
1384-
logger.error(`Prompt for closed session ${params.sessionId} failed after close`, err);
1470+
logger.error(`Prompt for cancelled session ${params.sessionId} failed after prompt returned`, err);
13851471
}
13861472
});
13871473
const turnCompleted = await Promise.race([
13881474
sendPromptPromise,
13891475
activePrompt.closeSignal,
1476+
this.cancelBeforeTurnStarted(activePrompt),
13901477
]);
13911478

13921479
if (turnCompleted === null) {
1393-
return {
1394-
stopReason: "cancelled",
1395-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1396-
_meta: this.buildQuotaMeta(sessionState),
1397-
};
1480+
return this.cancelledPromptResponse(sessionState);
13981481
}
13991482

14001483
await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
14011484

1402-
// Check if turn was interrupted (cancelled)
14031485
if (turnCompleted.turn.status === "interrupted") {
1404-
if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
1405-
await this.connection.notify(acp.methods.client.session.update, {
1406-
sessionId: params.sessionId,
1407-
update: {
1408-
sessionUpdate: "agent_message_chunk",
1409-
content: {
1410-
type: "text",
1411-
text: "*Conversation interrupted*"
1412-
}
1413-
}
1414-
});
1415-
}
1416-
return {
1417-
stopReason: "cancelled",
1418-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1419-
_meta: this.buildQuotaMeta(sessionState),
1420-
};
1486+
await this.notifyConversationInterrupted(params.sessionId);
1487+
return this.cancelledPromptResponse(sessionState);
14211488
}
14221489

14231490
const error = eventHandler.getFailure()
@@ -1436,6 +1503,7 @@ export class CodexAcpServer {
14361503
throw err;
14371504
} finally {
14381505
logger.log("Prompt completed", {sessionId: params.sessionId});
1506+
disposePromptRequestCancellation();
14391507
sessionState.currentTurnId = null;
14401508
if (pendingTurnStart !== null && this.pendingTurnStarts.get(params.sessionId) === pendingTurnStart) {
14411509
this.pendingTurnStarts.delete(params.sessionId);
@@ -1445,6 +1513,30 @@ export class CodexAcpServer {
14451513
}
14461514
}
14471515

1516+
private cancelledPromptResponse(sessionState: SessionState): acp.PromptResponse {
1517+
return {
1518+
stopReason: "cancelled",
1519+
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1520+
_meta: this.buildQuotaMeta(sessionState),
1521+
};
1522+
}
1523+
1524+
private async notifyConversationInterrupted(sessionId: string): Promise<void> {
1525+
if (this.sessionIsClosing(sessionId) || !this.sessions.has(sessionId)) {
1526+
return;
1527+
}
1528+
await this.connection.notify(acp.methods.client.session.update, {
1529+
sessionId,
1530+
update: {
1531+
sessionUpdate: "agent_message_chunk",
1532+
content: {
1533+
type: "text",
1534+
text: "*Conversation interrupted*"
1535+
}
1536+
}
1537+
});
1538+
}
1539+
14481540
private buildQuotaMeta(sessionState: SessionState): { quota: QuotaMeta } {
14491541
const lastTokenUsage = sessionState.lastTokenUsage;
14501542

0 commit comments

Comments
 (0)