Wait on turns

Author: benbrandt

src/CodexAcpClient.ts

@@ -325,6 +325,10 @@ export class CodexAcpClient {
         await this.codexClient.threadGoalClear({threadId: sessionId});
     }
 
+    async awaitTurnCompleted(params: { threadId: string, turnId: string }): Promise<TurnCompletedNotification> {
+        return await this.codexClient.awaitTurnCompleted(params.threadId, params.turnId);
+    }
+
     async awaitMcpServerStartup(serverNames: Array<string>, afterVersion: number): Promise<McpStartupResult> {
         return await this.codexClient.awaitMcpServerStartup(serverNames, afterVersion);
     }

src/CodexAcpServer.ts

@@ -15,6 +15,7 @@ import type {
     ReasoningEffortOption,
     Thread,
     ThreadItem,
+    TurnCompletedNotification,
     UserInput
 } from "./app-server/v2";
 import type {RateLimitsMap} from "./RateLimitsMap";
@@ -1279,6 +1280,30 @@ export class CodexAcpServer {
         return turnId;
     }
 
+    private async awaitCurrentCommandTurnCompletion(
+        sessionState: SessionState,
+        activePrompt: ActivePrompt,
+    ): Promise<TurnCompletedNotification | null | undefined> {
+        const turnId = sessionState.currentTurnId;
+        if (!turnId) {
+            return undefined;
+        }
+
+        const turnCompleted = await Promise.race([
+            this.runWithProcessCheck(() => this.codexAcpClient.awaitTurnCompleted({
+                threadId: sessionState.sessionId,
+                turnId,
+            })),
+            activePrompt.closeSignal,
+        ]);
+        if (turnCompleted === null) {
+            return null;
+        }
+
+        await this.codexAcpClient.waitForSessionNotifications(sessionState.sessionId);
+        return turnCompleted;
+    }
+
     async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
         logger.log("Prompt received", {
             sessionId: params.sessionId,
@@ -1306,7 +1331,18 @@ export class CodexAcpServer {
             if (commandResult.handled) {
                 logger.log("Prompt handled by a command");
                 await this.codexAcpClient.waitForSessionNotifications(params.sessionId);
-                if (commandResult.turnCompleted?.turn.status === "interrupted") {
+                let turnCompleted: TurnCompletedNotification | null | undefined = commandResult.turnCompleted;
+                if (!turnCompleted && commandResult.waitForCurrentTurnCompletion) {
+                    turnCompleted = await this.awaitCurrentCommandTurnCompletion(sessionState, activePrompt);
+                    if (turnCompleted === null) {
+                        return {
+                            stopReason: "cancelled",
+                            usage: this.buildPromptUsage(sessionState.lastTokenUsage),
+                            _meta: this.buildQuotaMeta(sessionState),
+                        };
+                    }
+                }
+                if (turnCompleted?.turn.status === "interrupted") {
                     if (!this.sessionIsClosing(params.sessionId) && this.sessions.has(params.sessionId)) {
                         await this.connection.notify(acp.methods.client.session.update, {
                             sessionId: params.sessionId,

src/CodexCommands.ts

@@ -15,7 +15,7 @@ type ParsedSlashCommand = {
 
 export type CommandHandleResult =
     | { handled: false }
-    | { handled: true, turnCompleted?: TurnCompletedNotification };
+    | { handled: true, turnCompleted?: TurnCompletedNotification, waitForCurrentTurnCompletion?: boolean };
 
 export class CodexCommands {
     private readonly connection: AcpClientConnection;
@@ -157,8 +157,7 @@ export class CodexCommands {
                 return { handled: true };
             }
             case "goal": {
-                await this.runGoalCommand(sessionId, command.rest);
-                return { handled: true };
+                return await this.runGoalCommand(sessionId, command.rest);
             }
             case "review": {
                 const target = this.buildReviewTarget(command.rest);
@@ -260,23 +259,23 @@ export class CodexCommands {
         ));
     }
 
-    private async runGoalCommand(sessionId: string, rest: string): Promise<void> {
+    private async runGoalCommand(sessionId: string, rest: string): Promise<CommandHandleResult> {
         const argument = rest.trim();
         if (argument.length === 0) {
             await this.sendCommandUsageMessage("goal", "[<objective>|clear|edit|pause|resume]", sessionId);
-            return;
+            return { handled: true };
         }
 
         switch (argument.toLowerCase()) {
             case "pause":
                 await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionId, "paused"));
-                return;
+                return { handled: true };
             case "resume":
                 await this.runWithProcessCheck(() => this.codexAcpClient.setGoalStatus(sessionId, "active"));
-                return;
+                return { handled: true, waitForCurrentTurnCompletion: true };
             case "clear":
                 await this.runWithProcessCheck(() => this.codexAcpClient.clearGoal(sessionId));
-                return;
+                return { handled: true };
         }
 
         if (argument.length > 4000) {
@@ -288,10 +287,11 @@ export class CodexCommands {
                     text: 'Command "/goal" requires goal text of at most 4000 characters.'
                 }
             });
-            return;
+            return { handled: true };
         }
 
         await this.runWithProcessCheck(() => this.codexAcpClient.setGoal(sessionId, argument));
+        return { handled: true, waitForCurrentTurnCompletion: true };
     }
 
     private buildReviewTarget(instructions: string): ReviewTarget {

src/__tests__/CodexACPAgent/CodexAcpClient.test.ts

@@ -1141,6 +1141,54 @@ describe('ACP server test', { timeout: 40_000 }, () => {
         expect(turnStartSpy).not.toHaveBeenCalled();
     });
 
+    it('waits for goal slash command turn completion', async () => {
+        const { mockFixture } = setupPromptFixture();
+        vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalSet")
+            .mockImplementation(async () => {
+                mockFixture.sendServerNotification({
+                    method: "turn/started",
+                    params: {
+                        threadId: "session-id",
+                        turn: createTurn("goal-turn-id", "inProgress"),
+                    },
+                });
+                return { goal: createThreadGoal() };
+            });
+        let completeGoal: (value: TurnCompletedNotification) => void = () => {};
+        const goalCompletedPromise = new Promise<TurnCompletedNotification>((resolve) => {
+            completeGoal = resolve;
+        });
+        vi.spyOn(mockFixture.getCodexAppServerClient(), "awaitTurnCompleted")
+            .mockReturnValue(goalCompletedPromise);
+
+        let promptResolved = false;
+        const promptPromise = mockFixture.getCodexAcpAgent().prompt({
+            sessionId: "session-id",
+            prompt: [{ type: "text", text: "/goal Ship the migration and keep tests green" }],
+        }).then((response) => {
+            promptResolved = true;
+            return response;
+        });
+
+        await vi.waitFor(() => {
+            expect(mockFixture.getCodexAppServerClient().awaitTurnCompleted).toHaveBeenCalledWith(
+                "session-id",
+                "goal-turn-id",
+            );
+        });
+        await Promise.resolve();
+        expect(promptResolved).toBe(false);
+
+        completeGoal({
+            threadId: "session-id",
+            turn: createTurn("goal-turn-id", "completed"),
+        });
+        await expect(promptPromise).resolves.toEqual(expect.objectContaining({
+            stopReason: "end_turn",
+        }));
+        expect(promptResolved).toBe(true);
+    });
+
     it('reports missing goal slash command input', async () => {
         const { mockFixture } = setupPromptFixture();
         const goalSetSpy = vi.spyOn(mockFixture.getCodexAppServerClient(), "threadGoalSet")