Skip to content

Commit 3be096f

Browse files
xtmqslapoguzov
xtmq
authored and
slapoguzov
committed
Support turn cancellation
1 parent 7fec958 commit 3be096f

5 files changed

Lines changed: 91 additions & 30 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type {
99
SetDefaultModelParams,
1010
SetDefaultModelResponse
1111
} from "./app-server";
12+
import type {TurnCompletedNotification } from "./app-server/v2";
1213
import type {JsonValue} from "./app-server/serde_json/JsonValue";
1314
import type {Model} from "./app-server/v2";
1415
import {ModelId} from "./ModelId";
@@ -101,7 +102,10 @@ export class CodexAcpClient {
101102
};
102103
}
103104

104-
async sendPrompt(request: acp.PromptRequest, eventHandler: (result: ServerNotification) => void): Promise<void> {
105+
async sendPrompt(
106+
request: acp.PromptRequest,
107+
eventHandler: (result: ServerNotification) => void
108+
): Promise<TurnCompletedNotification> {
105109
this.codexClient.onServerNotification(request.sessionId, eventHandler);
106110

107111
const input = request.prompt.filter(b => b.type === "text")
@@ -119,13 +123,22 @@ export class CodexAcpClient {
119123
model: null,
120124
});
121125

122-
await this.codexClient.awaitTurnCompleted();
126+
// Wait for turn completion
127+
// If turnInterrupt() was called, Codex will send turn/completed event with status "interrupted"
128+
return await this.codexClient.awaitTurnCompleted();
123129
}
124130

125131
async setModel(params: SetDefaultModelParams): Promise<SetDefaultModelResponse> {
126132
return this.codexClient.setModelRequest(params);
127133
}
128134

135+
async turnInterrupt(params: { threadId: string, turnId: string }): Promise<void> {
136+
await this.codexClient.turnInterrupt({
137+
threadId: params.threadId,
138+
turnId: params.turnId
139+
});
140+
}
141+
129142
private async fetchAvailableModels(): Promise<Model[]> {
130143
const models: Model[] = [];
131144
let cursor: string | null = null;

‎src/CodexAcpServer.ts‎

Lines changed: 39 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import {ModelId} from "./ModelId";
1010

1111
export interface SessionState {
1212
sessionMetadata: SessionMetadata;
13-
pendingPrompt: AbortController | null;
13+
currentTurnId: string | null;
1414
}
1515

1616
export class CodexAcpServer implements acp.Agent {
@@ -62,7 +62,7 @@ export class CodexAcpServer implements acp.Agent {
6262
const {sessionId, currentModelId, models} = sessionMetadata;
6363
this.sessions.set(sessionId, {
6464
sessionMetadata: sessionMetadata,
65-
pendingPrompt: null
65+
currentTurnId: null
6666
});
6767

6868
const availableModels = this.buildAvailableModels(models);
@@ -153,25 +153,28 @@ export class CodexAcpServer implements acp.Agent {
153153
async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
154154
const sessionState = this.getSessionState(params.sessionId);
155155

156-
sessionState.pendingPrompt?.abort();
157-
sessionState.pendingPrompt = new AbortController();
156+
sessionState.currentTurnId = null;
158157

159158
try {
160159
const messageHandler = new CodexEventHandler(this.connection, sessionState);
161-
await this.runWithProcessCheck(() => this.codexAcpClient.sendPrompt(params, (event) => messageHandler.handleNotification(event)));
162-
} catch (err) {
163-
if (sessionState.pendingPrompt.signal.aborted) {
164-
return {stopReason: "cancelled"};
160+
const turnCompleted = await this.runWithProcessCheck(() => this.codexAcpClient.sendPrompt(params, (event) => messageHandler.handleNotification(event)));
161+
162+
// Check if turn was interrupted (cancelled)
163+
if (turnCompleted.turn.status === "interrupted") {
164+
return {
165+
stopReason: "cancelled",
166+
};
165167
}
166168

169+
return {
170+
stopReason: "end_turn",
171+
};
172+
} catch (err) {
173+
console.error(`Prompt for session ${params.sessionId} failed:`, err);
167174
throw err;
175+
} finally {
176+
sessionState.currentTurnId = null;
168177
}
169-
170-
sessionState.pendingPrompt = null;
171-
172-
return {
173-
stopReason: "end_turn",
174-
};
175178
}
176179

177180
private async runWithProcessCheck<T>(operation: () => Promise<T>): Promise<T> {
@@ -191,7 +194,27 @@ export class CodexAcpServer implements acp.Agent {
191194
}
192195

193196
async cancel(params: acp.CancelNotification): Promise<void> {
194-
//TODO not supported yet
195-
this.sessions.get(params.sessionId)?.pendingPrompt?.abort();
197+
const sessionState = this.sessions.get(params.sessionId);
198+
if (!sessionState) {
199+
console.info(`Can not cancel: session ${params.sessionId} not found`);
200+
return;
201+
}
202+
203+
if (!sessionState.currentTurnId) {
204+
console.info(`Can not cancel: session ${params.sessionId} has no current turn`);
205+
return;
206+
}
207+
208+
console.info(`Cancel session ${params.sessionId}, currentTurnId: ${sessionState.currentTurnId}...`);
209+
try {
210+
// After turnInterrupt(), Codex will send turn/completed event, which will naturally complete awaitTurnCompleted()
211+
await this.codexAcpClient.turnInterrupt({
212+
threadId: params.sessionId,
213+
turnId: sessionState.currentTurnId
214+
});
215+
console.log(`Cancel - turnInterrupt succeeded`);
216+
} catch (err) {
217+
console.error(`Cancel - turnInterrupt failed:`, err);
218+
}
196219
}
197220
}

‎src/CodexAppServerClient.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ import type {
1313
ThreadStartParams,
1414
ThreadStartResponse,
1515
TurnCompletedNotification,
16+
TurnInterruptParams,
17+
TurnInterruptResponse,
1618
TurnStartParams,
1719
TurnStartResponse
1820
} from "./app-server/v2";
@@ -45,6 +47,10 @@ export class CodexAppServerClient {
4547
return await this.sendRequest({ method: "turn/start", params: params });
4648
}
4749

50+
async turnInterrupt(params: TurnInterruptParams): Promise<TurnInterruptResponse> {
51+
return await this.sendRequest({ method: "turn/interrupt", params: params });
52+
}
53+
4854
async threadStart(params: ThreadStartParams): Promise<ThreadStartResponse> {
4955
return await this.sendRequest({ method: "thread/start", params: params });
5056
}

‎src/CodexEventHandler.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,16 @@ export class CodexEventHandler {
5353
return await this.updatePlan(notification.params);
5454
case "error":
5555
return await this.createErrorEvent(notification.params);
56+
case "turn/started":
57+
this.sessionState.currentTurnId = notification.params.turn.id;
58+
return null;
59+
case "turn/completed":
60+
this.sessionState.currentTurnId = null;
61+
return null;
5662
case "item/reasoning/summaryTextDelta": //TODO streaming reasoning?
5763
case "item/reasoning/summaryPartAdded":
5864
//skipped events
5965
case "item/reasoning/textDelta": //for raw output
60-
case "turn/started":
61-
case "turn/completed":
6266
case "turn/diff/updated":
6367
case "item/commandExecution/outputDelta":
6468
case "item/fileChange/outputDelta":

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

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -77,10 +77,15 @@ describe('ACP server test', { timeout: 40_000 }, () => {
7777

7878
const codexAcpAgent = fixture.getCodexAcpAgent();
7979

80-
fixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
81-
fixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
80+
fixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue({
81+
turn: { id: "turn-id", items: [], status: "inProgress", error: null }
82+
});
83+
fixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue({
84+
threadId: "id",
85+
turn: { id: "turn-id", items: [], status: "completed", error: null }
86+
});
8287
const sessionState: SessionState = {
83-
pendingPrompt: null,
88+
currentTurnId: null,
8489
sessionMetadata: {
8590
sessionId: "id",
8691
currentModelId: "model-id",
@@ -99,11 +104,16 @@ describe('ACP server test', { timeout: 40_000 }, () => {
99104
const mockFixture = createCodexMockTestFixture();
100105
const codexAcpAgent = mockFixture.getCodexAcpAgent();
101106

102-
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
103-
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
107+
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue({
108+
turn: { id: "turn-id", items: [], status: "inProgress", error: null }
109+
});
110+
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue({
111+
threadId: "id",
112+
turn: { id: "turn-id", items: [], status: "completed", error: null }
113+
});
104114

105115
const sessionState: SessionState = {
106-
pendingPrompt: null,
116+
currentTurnId: null,
107117
sessionMetadata: {
108118
sessionId: "id",
109119
currentModelId: "model-id",
@@ -143,19 +153,24 @@ describe('ACP server test', { timeout: 40_000 }, () => {
143153
const mockFixture = createCodexMockTestFixture();
144154
const codexAcpAgent = mockFixture.getCodexAcpAgent();
145155

146-
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
147-
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
156+
mockFixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue({
157+
turn: { id: "turn-id", items: [], status: "inProgress", error: null }
158+
});
159+
mockFixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue({
160+
threadId: "id",
161+
turn: { id: "turn-id", items: [], status: "completed", error: null }
162+
});
148163

149164
const sessionState1: SessionState = {
150-
pendingPrompt: null,
165+
currentTurnId: null,
151166
sessionMetadata: {
152167
sessionId: "session-1",
153168
currentModelId: "model-id",
154169
models: [],
155170
}
156171
};
157172
const sessionState2: SessionState = {
158-
pendingPrompt: null,
173+
currentTurnId: null,
159174
sessionMetadata: {
160175
sessionId: "session-2",
161176
currentModelId: "model-id",

0 commit comments

Comments
 (0)