test: stabilize e2e approval and cancellation tests

src/__tests__/CodexACPAgent/e2e/acp-e2e-shell-approval.test.ts

@@ -2,6 +2,7 @@ import fs from "node:fs";
 import path from "node:path";
 import {afterEach, beforeEach, expect, it, onTestFinished, vi} from "vitest";
 import {AgentMode} from "../../../AgentMode";
+import {ASYNC_TASK_STOP_METHOD} from "../../../async-tasks/AsyncTaskExtension";
 import {ApprovalOptionId} from "../../../permissions/option-ids";
 import {
     createAuthenticatedFixture,
@@ -18,6 +19,9 @@ import {
 
 const FIRST_FILE_NAME = "approval-first.txt";
 const SECOND_FILE_NAME = "approval-second.txt";
+// Without this hint the model may run the command sandboxed, hit the read-only
+// filesystem and end the turn without ever requesting approval.
+const ESCALATE = "It writes outside the sandbox, so request escalated permissions on the first attempt.";
 
 describeE2E("E2E shell approval tests", () => {
     let fixture: SpawnedAgentFixture;
@@ -43,8 +47,8 @@ describeE2E("E2E shell approval tests", () => {
         expectedStopReasons: ["end_turn" | "cancelled", "end_turn" | "cancelled"],
     ): Promise<void> {
         const prompts = [
-            `Use your shell tool to run exactly \`${command}\`.`,
-            `Use your shell tool to run exactly the same command again: \`${command}\`.`,
+            `Use your shell tool to run exactly \`${command}\`. ${ESCALATE}`,
+            `Use your shell tool to run exactly the same command again: \`${command}\`. ${ESCALATE}`,
         ];
         for (const [index, text] of prompts.entries()) {
             const response = await fixture.connection.prompt({
@@ -105,7 +109,7 @@ describeE2E("E2E read-only mode shell permission tests", () => {
 
     it("requests permission for a command that writes inside the workspace", async () => {
         fixture.setPermissionResponder(createPermissionResponder("execute", ApprovalOptionId.AllowOnce));
-        const sessionId = await writeToFile(fixture, path.join(fixture.workspaceDir, generateFileNameForTest()));
+        const sessionId = await writeToFile(fixture, path.join(fixture.workspaceDir, generateFileNameForTest()), ESCALATE);
 
         expect(fixture.readPermissionRequests(sessionId, "execute").length).toBeGreaterThanOrEqual(1);
         expect(fixture.readPermissionRequests(sessionId, "edit")).toHaveLength(0);
@@ -141,7 +145,7 @@ describeE2E("E2E read-only mode shell permission tests", () => {
     it("requests permission for a command that writes outside the workspace", async () => {
         const dir = createDirOutsideWorkspace(fixture);
         fixture.setPermissionResponder(createPermissionResponder("execute", ApprovalOptionId.AllowOnce));
-        const sessionId = await writeToFile(fixture, path.join(dir, generateFileNameForTest()));
+        const sessionId = await writeToFile(fixture, path.join(dir, generateFileNameForTest()), ESCALATE);
         expectPermissionRequests(fixture, sessionId, {execute: 1, edit: 0});
     });
 });
@@ -198,8 +202,10 @@ describeE2E("E2E shell cancellation tests", () => {
         }
     }
 
-    it("cancels a running shell command", async () => {
-        fixture = await createAuthenticatedFixture(AgentMode.AgentFullAccess);
+    // Codex keeps a running command as a background terminal when the turn is interrupted,
+    // so cancel ends the turn and the async task stop request ends the process.
+    it("cancels the turn of a running shell command and stops it as an async task", async () => {
+        fixture = await createAuthenticatedFixture(AgentMode.AgentFullAccess, undefined, ["asyncTasks"]);
         const sessionId = (await fixture.createSession()).sessionId;
         const pidFilePath = path.join(fixture.workspaceDir, "cancel-command.pid");
         const command = `/bin/sh -c 'echo $$ > "${pidFilePath}"; exec sleep 100'`;
@@ -215,25 +221,45 @@ describeE2E("E2E shell cancellation tests", () => {
             expect(parsed).toBeGreaterThan(0);
             return parsed;
         }, {timeout: 10_000});
+        onTestFinished(() => {
+            if (isProcessRunning(pid)) process.kill(pid, "SIGKILL");
+        });
         expect(isProcessRunning(pid)).toBe(true);
+        // The process can start before Codex tracks it in the turn, so wait until the client sees it.
+        const toolCallId = await vi.waitFor(() => {
+            const [toolCall] = fixture?.readToolCalls(sessionId, "execute") ?? [];
+            expect(toolCall).toBeDefined();
+            return toolCall!.toolCallId;
+        }, {timeout: 10_000});
         await fixture.connection.cancel({sessionId});
 
         expect((await promptResponse).stopReason).toBe("cancelled");
+        expect(isProcessRunning(pid)).toBe(true);
+
+        // The SDK client rejects the `async_task_spawned` update it has no schema for, but a root
+        // command's async task id is its tool call id. Stop fails until the adapter publishes the task.
+        await vi.waitFor(async () => {
+            const response = await fixture?.connection.extMethod(ASYNC_TASK_STOP_METHOD, {
+                sessionId,
+                asyncTaskId: toolCallId,
+            });
+            expect(response).toEqual({stopped: true});
+        }, {timeout: 10_000});
         await vi.waitFor(() => {
             expect(isProcessRunning(pid)).toBe(false);
         }, {timeout: 5_000});
     });
 });
 
-async function writeToFile(fixture: SpawnedAgentFixture, filePath: string): Promise<string> {
+async function writeToFile(fixture: SpawnedAgentFixture, filePath: string, hint?: string): Promise<string> {
     const content = "hello from e2e";
     const command = `printf '${content}' > '${filePath}'`;
     const sessionId = (await fixture.createSession()).sessionId;
     const response = await fixture.connection.prompt({
         sessionId,
         prompt: [{
             type: "text",
-            text: `Use your shell tool to run exactly \`${command}\`. Do not modify files any other way.`,
+            text: `Use your shell tool to run exactly \`${command}\`. Do not modify files any other way.${hint ? ` ${hint}` : ""}`,
         }],
     });
     expectEndTurn(response);

src/__tests__/CodexACPAgent/e2e/acp-e2e-test-utils.ts

@@ -47,7 +47,11 @@ export function expectNoPermissionRequests(fixture: SpawnedAgentFixture, session
     expectPermissionRequests(fixture, sessionId, { edit: 0, execute: 0 });
 }
 
-export async function createAuthenticatedFixture(initialMode?: AgentMode, mcpServers?: acp.McpServerStdio[]): Promise<SpawnedAgentFixture> {
+export async function createAuthenticatedFixture(
+    initialMode?: AgentMode,
+    mcpServers?: acp.McpServerStdio[],
+    airCapabilities?: string[],
+): Promise<SpawnedAgentFixture> {
     const apiKey = requireLiveApiKey();
     const extraEnv = initialMode ? {INITIAL_AGENT_MODE: initialMode.id} : undefined;
     return await createSpawnedFixture(async (connection, authMethods) => {
@@ -68,7 +72,7 @@ export async function createAuthenticatedFixture(initialMode?: AgentMode, mcpSer
         if (authenticationStatus["type"] !== "api-key") {
             throw new Error(`Unexpected authentication status: ${JSON.stringify(authenticationStatus)}`);
         }
-    }, extraEnv, mcpServers);
+    }, extraEnv, mcpServers, airCapabilities);
 }
 
 export async function createGatewayFixture(
@@ -97,7 +101,7 @@ export async function createGatewayFixture(
     });
 }
 
-function buildClientCapabilities(): acp.ClientCapabilities {
+function buildClientCapabilities(airCapabilities?: string[]): acp.ClientCapabilities {
     return {
         fs: {
             readTextFile: true,
@@ -111,6 +115,7 @@ function buildClientCapabilities(): acp.ClientCapabilities {
         },
         _meta: {
             "terminal-auth": true,
+            ...(airCapabilities ? {jetbrains: {air: {version: 1, capabilities: airCapabilities}}} : {}),
         },
     };
 }
@@ -121,11 +126,12 @@ async function createSpawnedFixture(
     authenticate: Authenticator,
     extraEnv?: NodeJS.ProcessEnv,
     mcpServers?: acp.McpServerStdio[],
+    airCapabilities?: string[],
 ): Promise<SpawnedAgentFixture> {
     return await createSpawnedAgentFixture(async (connection) => {
         const initializeResponse = await connection.initialize({
             protocolVersion: acp.PROTOCOL_VERSION,
-            clientCapabilities: buildClientCapabilities(),
+            clientCapabilities: buildClientCapabilities(airCapabilities),
             clientInfo: {
                 name: "vitest",
                 version: "1.0.0",

src/__tests__/CodexACPAgent/e2e/spawned-agent-fixture.ts

@@ -9,7 +9,7 @@ import {removeDirectoryWithRetry, writeCodexHomeConfig} from "../../acp-test-uti
 import type {PermissionResponder} from "./permission-responders";
 import type {LegacyNewSessionResponse} from "../../../AcpExtensions";
 
-export const DEFAULT_TEST_MODEL_ID = ModelId.create("gpt-5.2", "none");
+export const DEFAULT_TEST_MODEL_ID = ModelId.create("gpt-6-luna", "low");
 export const OTHER_TEST_MODEL_ID = ModelId.create("gpt-5.5", "low");
 
 export interface TestSkill {
@@ -38,6 +38,7 @@ export interface SpawnedAgentFixture {
         sessionId: string,
         toolCallKind: acp.ToolKind,
     ): acp.RequestPermissionRequest[];
+    readToolCalls(sessionId: string, toolCallKind: acp.ToolKind): acp.ToolCall[];
     dispose(): Promise<void>;
 }
 
@@ -107,6 +108,7 @@ class RecordingClient implements acp.Client {
     private readonly textBySessionId = new Map<string, string>();
     private readonly availableCommandsBySessionId = new Map<string, acp.AvailableCommand[]>();
     private readonly permissionRequestsBySessionId = new Map<string, acp.RequestPermissionRequest[]>();
+    private readonly toolCallsBySessionId = new Map<string, acp.ToolCall[]>();
     private permissionResponder: PermissionResponder = () => ({
         outcome: {outcome: "cancelled"},
     });
@@ -131,6 +133,13 @@ class RecordingClient implements acp.Client {
             return;
         }
 
+        if (params.update.sessionUpdate === "tool_call") {
+            const toolCalls = this.toolCallsBySessionId.get(params.sessionId) ?? [];
+            toolCalls.push(params.update);
+            this.toolCallsBySessionId.set(params.sessionId, toolCalls);
+            return;
+        }
+
         if (params.update.sessionUpdate !== "agent_message_chunk" || params.update.content.type !== "text") {
             return;
         }
@@ -154,6 +163,11 @@ class RecordingClient implements acp.Client {
         const requests = this.permissionRequestsBySessionId.get(sessionId) ?? [];
         return requests.filter((request) => request.toolCall.kind === toolCallKind);
     }
+
+    readToolCalls(sessionId: string, toolCallKind: acp.ToolKind): acp.ToolCall[] {
+        const toolCalls = this.toolCallsBySessionId.get(sessionId) ?? [];
+        return toolCalls.filter((toolCall) => toolCall.kind === toolCallKind);
+    }
 }
 
 class SpawnedAgentFixtureImpl implements SpawnedAgentFixture {
@@ -261,6 +275,10 @@ class SpawnedAgentFixtureImpl implements SpawnedAgentFixture {
         return this.client.readPermissionRequests(sessionId, toolCallKind);
     }
 
+    readToolCalls(sessionId: string, toolCallKind: acp.ToolKind): acp.ToolCall[] {
+        return this.client.readToolCalls(sessionId, toolCallKind);
+    }
+
     async dispose(): Promise<void> {
         if (this.disposed) {
             return;