Skip to content

Commit 48cd521

Browse files
benbrandtbenbrandt
committed
Wait for goal turn
1 parent e9bc724 commit 48cd521

5 files changed

Lines changed: 168 additions & 65 deletions

File tree

‎src/CodexAcpClient.ts‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -306,12 +306,16 @@ export class CodexAcpClient {
306306
await this.codexClient.runCompact({threadId: sessionId});
307307
}
308308

309-
async setGoal(sessionId: string, objective: string): Promise<void> {
310-
await this.codexClient.threadGoalSet({
309+
async setGoal(
310+
sessionId: string,
311+
objective: string,
312+
onTurnStarted?: (turnId: string) => void,
313+
): Promise<TurnCompletedNotification> {
314+
return await this.codexClient.runGoalSet({
311315
threadId: sessionId,
312316
objective,
313317
status: "active",
314-
});
318+
}, onTurnStarted);
315319
}
316320

317321
async setGoalStatus(sessionId: string, status: ThreadGoalStatus): Promise<void> {
@@ -321,12 +325,18 @@ export class CodexAcpClient {
321325
});
322326
}
323327

324-
async clearGoal(sessionId: string): Promise<void> {
325-
await this.codexClient.threadGoalClear({threadId: sessionId});
328+
async resumeGoal(
329+
sessionId: string,
330+
onTurnStarted?: (turnId: string) => void,
331+
): Promise<TurnCompletedNotification> {
332+
return await this.codexClient.runGoalSet({
333+
threadId: sessionId,
334+
status: "active",
335+
}, onTurnStarted);
326336
}
327337

328-
async awaitTurnCompleted(params: { threadId: string, turnId: string }): Promise<TurnCompletedNotification> {
329-
return await this.codexClient.awaitTurnCompleted(params.threadId, params.turnId);
338+
async clearGoal(sessionId: string): Promise<void> {
339+
await this.codexClient.threadGoalClear({threadId: sessionId});
330340
}
331341

332342
async awaitMcpServerStartup(serverNames: Array<string>, afterVersion: number): Promise<McpStartupResult> {

‎src/CodexAcpServer.ts‎

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1280,30 +1280,6 @@ export class CodexAcpServer {
12801280
return turnId;
12811281
}
12821282

1283-
private async awaitCurrentCommandTurnCompletion(
1284-
sessionState: SessionState,
1285-
activePrompt: ActivePrompt,
1286-
): Promise<TurnCompletedNotification | null | undefined> {
1287-
const turnId = sessionState.currentTurnId;
1288-
if (!turnId) {
1289-
return undefined;
1290-
}
1291-
1292-
const turnCompleted = await Promise.race([
1293-
this.runWithProcessCheck(() => this.codexAcpClient.awaitTurnCompleted({
1294-
threadId: sessionState.sessionId,
1295-
turnId,
1296-
})),
1297-
activePrompt.closeSignal,
1298-
]);
1299-
if (turnCompleted === null) {
1300-
return null;
1301-
}
1302-
1303-
await this.codexAcpClient.waitForSessionNotifications(sessionState.sessionId);
1304-
return turnCompleted;
1305-
}
1306-
13071283
async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
13081284
logger.log("Prompt received", {
13091285
sessionId: params.sessionId,
@@ -1331,17 +1307,7 @@ export class CodexAcpServer {
13311307
if (commandResult.handled) {
13321308
logger.log("Prompt handled by a command");
13331309
await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
1334-
let turnCompleted: TurnCompletedNotification | null | undefined = commandResult.turnCompleted;
1335-
if (!turnCompleted && commandResult.waitForCurrentTurnCompletion) {
1336-
turnCompleted = await this.awaitCurrentCommandTurnCompletion(sessionState, activePrompt);
1337-
if (turnCompleted === null) {
1338-
return {
1339-
stopReason: "cancelled",
1340-
usage: this.buildPromptUsage(sessionState.lastTokenUsage),
1341-
_meta: this.buildQuotaMeta(sessionState),
1342-
};
1343-
}
1344-
}
1310+
const turnCompleted = commandResult.turnCompleted;
13451311
if (turnCompleted?.turn.status === "interrupted") {
13461312
if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
13471313
await this.connection.notify(acp.methods.client.session.update, {

‎src/CodexAppServerClient.ts‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ export class CodexAppServerClient {
119119
private readonly pendingTurnCompletionResolvers = new Map<string, Map<string, (event: TurnCompletedNotification) => void>>();
120120
private readonly pendingCompactionCompletionResolvers = new Map<string, Set<(event: CompactionCompletedNotification) => void>>();
121121
private readonly turnCompletionCaptures = new Map<string, Set<(event: TurnCompletedNotification) => void>>();
122+
private readonly turnRoutingCaptures = new Map<string, Set<(turnId: string) => void>>();
122123
private readonly staleTurnIds = new Map<string, Set<string>>();
123124

124125
constructor(connection: MessageConnection) {
@@ -151,6 +152,7 @@ export class CodexAppServerClient {
151152
}
152153
return;
153154
}
155+
this.recordTurnRouting(routing);
154156
this.notify(serverNotification);
155157
for (const callback of this.codexEventHandlers) {
156158
callback({ eventType: "notification", ...serverNotification });
@@ -266,6 +268,47 @@ export class CodexAppServerClient {
266268
}
267269
}
268270

271+
async runGoalSet(params: ThreadGoalSetParams, onTurnStarted?: (turnId: string) => void): Promise<TurnCompletedNotification> {
272+
const capturedCompletions: Array<TurnCompletedNotification> = [];
273+
const releaseCompletionCapture = this.captureTurnCompletions(params.threadId, (event) => {
274+
capturedCompletions.push(event);
275+
});
276+
let goalTurnId: string | null = null;
277+
let resolveGoalTurnStarted: () => void = () => {};
278+
const goalTurnStarted = new Promise<void>((resolve) => {
279+
resolveGoalTurnStarted = resolve;
280+
});
281+
const releaseRoutingCapture = this.captureTurnRoutings(params.threadId, (turnId) => {
282+
if (goalTurnId !== null) {
283+
return;
284+
}
285+
goalTurnId = turnId;
286+
onTurnStarted?.(turnId);
287+
resolveGoalTurnStarted();
288+
});
289+
290+
try {
291+
await this.threadGoalSet(params);
292+
if (goalTurnId === null) {
293+
await goalTurnStarted;
294+
}
295+
const turnId = goalTurnId;
296+
if (turnId === null) {
297+
throw new Error("Goal command did not start a turn");
298+
}
299+
const earlyCompletion = capturedCompletions.find(event => event.turn.id === turnId);
300+
releaseCompletionCapture();
301+
releaseRoutingCapture();
302+
if (earlyCompletion) {
303+
return earlyCompletion;
304+
}
305+
return await this.awaitTurnCompleted(params.threadId, turnId);
306+
} finally {
307+
releaseCompletionCapture();
308+
releaseRoutingCapture();
309+
}
310+
}
311+
269312
async runCompact(params: ThreadCompactStartParams): Promise<CompactionCompletedNotification> {
270313
const compactionCompleted = this.awaitCompactionCompleted(params.threadId);
271314
await this.threadCompactStart(params);
@@ -478,6 +521,19 @@ export class CodexAppServerClient {
478521
}
479522
}
480523

524+
private recordTurnRouting(routing: { threadId: string | null, turnId: string | null }): void {
525+
if (routing.threadId === null || routing.turnId === null) {
526+
return;
527+
}
528+
const captures = this.turnRoutingCaptures.get(routing.threadId);
529+
if (!captures) {
530+
return;
531+
}
532+
for (const capture of captures) {
533+
capture(routing.turnId);
534+
}
535+
}
536+
481537
private isStaleTurn(threadId: string | null, turnId: string | null): boolean {
482538
if (threadId === null || turnId === null) {
483539
return false;
@@ -523,6 +579,23 @@ export class CodexAppServerClient {
523579
};
524580
}
525581

582+
private captureTurnRoutings(threadId: string, capture: (turnId: string) => void): () => void {
583+
const captures = this.turnRoutingCaptures.get(threadId) ?? new Set<(turnId: string) => void>();
584+
captures.add(capture);
585+
this.turnRoutingCaptures.set(threadId, captures);
586+
let released = false;
587+
return () => {
588+
if (released) {
589+
return;
590+
}
591+
released = true;
592+
captures.delete(capture);
593+
if (captures.size === 0) {
594+
this.turnRoutingCaptures.delete(threadId);
595+
}
596+
};
597+
}
598+
526599
private resolveMcpServerStartupResolvers(): void {
527600
const pendingResolvers: Array<McpServerStartupResolver> = [];
528601
for (const resolver of this.mcpServerStartupResolvers) {

‎src/CodexCommands.ts‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ type ParsedSlashCommand = {
1515

1616
export type CommandHandleResult =
1717
| { handled: false }
18-
| { handled: true, turnCompleted?: TurnCompletedNotification, waitForCurrentTurnCompletion?: boolean };
18+
| { handled: true, turnCompleted?: TurnCompletedNotification };
1919

2020
export class CodexCommands {
2121
private readonly connection: AcpClientConnection;
@@ -157,7 +157,7 @@ export class CodexCommands {
157157
return { handled: true };
158158
}
159159
case "goal": {
160-
return await this.runGoalCommand(sessionId, command.rest);
160+
return await this.runGoalCommand(sessionState, command.rest);
161161
}
162162
case "review": {
163163
const target = this.buildReviewTarget(command.rest);
@@ -259,7 +259,8 @@ export class CodexCommands {
259259
));
260260
}
261261

262-
private async runGoalCommand(sessionId: string, rest: string): Promise<CommandHandleResult> {
262+
private async runGoalCommand(sessionState: SessionState, rest: string): Promise<CommandHandleResult> {
263+
const sessionId = sessionState.sessionId;
263264
const argument = rest.trim();
264265
if (argument.length === 0) {
265266
await this.sendCommandUsageMessage("goal", "[<objective>|clear|edit|pause|resume]", sessionId);
@@ -271,8 +272,15 @@ export class CodexCommands {
271272
await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionId, "paused"));
272273
return { handled: true };
273274
case "resume":
274-
await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionId, "active"));
275-
return { handled: true, waitForCurrentTurnCompletion: true };
275+
return {
276+
handled: true,
277+
turnCompleted: await this.runWithProcessCheck(() => this.codexAcpClient.resumeGoal(
278+
sessionId,
279+
(turnId) => {
280+
sessionState.currentTurnId = turnId;
281+
},
282+
)),
283+
};
276284
case "clear":
277285
await this.runWithProcessCheck(() => this.codexAcpClient.clearGoal(sessionId));
278286
return { handled: true };
@@ -290,8 +298,16 @@ export class CodexCommands {
290298
return { handled: true };
291299
}
292300

293-
await this.runWithProcessCheck(() => this.codexAcpClient.setGoal(sessionId, argument));
294-
return { handled: true, waitForCurrentTurnCompletion: true };
301+
return {
302+
handled: true,
303+
turnCompleted: await this.runWithProcessCheck(() => this.codexAcpClient.setGoal(
304+
sessionId,
305+
argument,
306+
(turnId) => {
307+
sessionState.currentTurnId = turnId;
308+
},
309+
)),
310+
};
295311
}
296312

297313
private buildReviewTarget(instructions: string): ReviewTarget {

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

Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1102,7 +1102,12 @@ describe('ACP server test', { timeout: 40_000 }, () => {
11021102

11031103
it('handles goal slash commands through Codex app server', async () => {
11041104
const { mockFixture, turnStartSpy } = setupPromptFixture();
1105-
const goalSetSpy = vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalSet")
1105+
const goalRunSpy = vi.spyOn(mockFixture.getCodexAppServerClient(), "runGoalSet")
1106+
.mockResolvedValue({
1107+
threadId: "session-id",
1108+
turn: createTurn("goal-turn-id", "completed"),
1109+
});
1110+
const goalStatusSpy = vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalSet")
11061111
.mockResolvedValue({ goal: createThreadGoal() });
11071112
const goalClearSpy = vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalClear")
11081113
.mockResolvedValue({ cleared: true });
@@ -1124,36 +1129,27 @@ describe('ACP server test', { timeout: 40_000 }, () => {
11241129
prompt: [{ type: "text", text: "/goal clear" }],
11251130
});
11261131

1127-
expect(goalSetSpy).toHaveBeenNthCalledWith(1, {
1132+
expect(goalRunSpy).toHaveBeenNthCalledWith(1, {
11281133
threadId: "session-id",
11291134
objective: "Ship the migration and keep tests green",
11301135
status: "active",
1131-
});
1132-
expect(goalSetSpy).toHaveBeenNthCalledWith(2, {
1136+
}, expect.any(Function));
1137+
expect(goalStatusSpy).toHaveBeenCalledWith({
11331138
threadId: "session-id",
11341139
status: "paused",
11351140
});
1136-
expect(goalSetSpy).toHaveBeenNthCalledWith(3, {
1141+
expect(goalRunSpy).toHaveBeenNthCalledWith(2, {
11371142
threadId: "session-id",
11381143
status: "active",
1139-
});
1144+
}, expect.any(Function));
11401145
expect(goalClearSpy).toHaveBeenCalledWith({ threadId: "session-id" });
11411146
expect(turnStartSpy).not.toHaveBeenCalled();
11421147
});
11431148

11441149
it('waits for goal slash command turn completion', async () => {
11451150
const { mockFixture } = setupPromptFixture();
11461151
vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalSet")
1147-
.mockImplementation(async () => {
1148-
mockFixture.sendServerNotification({
1149-
method: "turn/started",
1150-
params: {
1151-
threadId: "session-id",
1152-
turn: createTurn("goal-turn-id", "inProgress"),
1153-
},
1154-
});
1155-
return { goal: createThreadGoal() };
1156-
});
1152+
.mockResolvedValue({ goal: createThreadGoal() });
11571153
let completeGoal: (value: TurnCompletedNotification) => void = () => {};
11581154
const goalCompletedPromise = new Promise<TurnCompletedNotification>((resolve) => {
11591155
completeGoal = resolve;
@@ -1170,6 +1166,48 @@ describe('ACP server test', { timeout: 40_000 }, () => {
11701166
return response;
11711167
});
11721168

1169+
await vi.waitFor(() => {
1170+
expect(mockFixture.getCodexAppServerClient().threadGoalSet).toHaveBeenCalledWith({
1171+
threadId: "session-id",
1172+
objective: "Ship the migration and keep tests green",
1173+
status: "active",
1174+
});
1175+
});
1176+
await Promise.resolve();
1177+
expect(promptResolved).toBe(false);
1178+
expect(mockFixture.getCodexAppServerClient().awaitTurnCompleted).not.toHaveBeenCalled();
1179+
1180+
mockFixture.sendServerNotification({
1181+
method: "thread/goal/updated",
1182+
params: {
1183+
threadId: "session-id",
1184+
turnId: null,
1185+
goal: createThreadGoal(),
1186+
},
1187+
});
1188+
mockFixture.sendServerNotification({
1189+
method: "thread/status/changed",
1190+
params: {
1191+
threadId: "session-id",
1192+
status: {
1193+
type: "active",
1194+
activeFlags: [],
1195+
},
1196+
},
1197+
});
1198+
await Promise.resolve();
1199+
expect(promptResolved).toBe(false);
1200+
expect(mockFixture.getCodexAppServerClient().awaitTurnCompleted).not.toHaveBeenCalled();
1201+
1202+
mockFixture.sendServerNotification({
1203+
method: "item/agentMessage/delta",
1204+
params: {
1205+
threadId: "session-id",
1206+
turnId: "goal-turn-id",
1207+
itemId: "goal-message-id",
1208+
delta: "I",
1209+
},
1210+
});
11731211
await vi.waitFor(() => {
11741212
expect(mockFixture.getCodexAppServerClient().awaitTurnCompleted).toHaveBeenCalledWith(
11751213
"session-id",

0 commit comments

Comments
 (0)