@vellumai/assistant 0.8.11 → 0.8.12-staging.2

package/src/acp/__tests__/agent-process.test.ts

@@ -6,9 +6,9 @@
  * connection is stubbed directly so no child process is spawned.
  */
 
-import { describe, expect, test } from "bun:test";
+import { describe, expect, mock, test } from "bun:test";
 
-import type { InitializeResponse } from "@agentclientprotocol/sdk";
+import type { AuthMethod, InitializeResponse } from "@agentclientprotocol/sdk";
 
 import { AcpAgentProcess } from "../agent-process.js";
 
@@ -159,3 +159,316 @@ describe("AcpAgentProcess loadSession/resumeSession", () => {
     ).rejects.toThrow('ACP agent "test-agent" is not spawned');
   });
 });
+
+describe("AcpAgentProcess auth_required retry", () => {
+  // spawnedEnv is injected the way spawn() builds it ({ ...process.env,
+  // ...config.env }), so the advertised var names use a VELLUM_TEST_ prefix
+  // guaranteed absent from the test process env. The structure mirrors
+  // codex-acp's real advertisement: an agent-driven browser login (no `type`,
+  // must never be auto-selected) followed by two env_var methods, in that
+  // order.
+  const CODEX_VAR = "VELLUM_TEST_FAKE_CODEX_API_KEY";
+  const OPENAI_VAR = "VELLUM_TEST_FAKE_OPENAI_API_KEY";
+
+  const codexLikeAuthMethods: AuthMethod[] = [
+    { id: "chatgpt", name: "Login with ChatGPT" },
+    {
+      type: "env_var",
+      id: "codex-api-key",
+      name: "Use CODEX_API_KEY",
+      vars: [{ name: CODEX_VAR }],
+    },
+    {
+      type: "env_var",
+      id: "openai-api-key",
+      name: "Use OPENAI_API_KEY",
+      vars: [{ name: OPENAI_VAR }],
+    },
+  ];
+
+  const authRequiredError = {
+    code: -32000,
+    message: "Authentication required",
+  };
+
+  /**
+   * Creates a process whose connection is stubbed with mocks, then runs
+   * initialize() so the advertised authMethods are captured.
+   */
+  async function setupAuthProcess(options: {
+    env?: Record<string, string>;
+    authMethods?: AuthMethod[];
+    /** Rejections to throw from successive newSession calls before succeeding. */
+    newSessionRejections?: unknown[];
+    loadSessionRejections?: unknown[];
+    resumeSessionRejections?: unknown[];
+    promptRejections?: unknown[];
+  }) {
+    const proc = new AcpAgentProcess(
+      "codex",
+      { command: "echo", args: [], env: options.env },
+      () => {
+        throw new Error("client factory should not be called in this test");
+      },
+    );
+
+    function failThenSucceed(rejections: unknown[], result: unknown) {
+      let calls = 0;
+      return mock(() => {
+        const rejection = rejections[calls];
+        calls += 1;
+        return rejection != null
+          ? Promise.reject(rejection)
+          : Promise.resolve(result);
+      });
+    }
+
+    const newSession = failThenSucceed(options.newSessionRejections ?? [], {
+      sessionId: "session-1",
+    });
+    const loadSession = failThenSucceed(
+      options.loadSessionRejections ?? [],
+      {},
+    );
+    const resumeSession = failThenSucceed(
+      options.resumeSessionRejections ?? [],
+      {},
+    );
+    const prompt = failThenSucceed(options.promptRejections ?? [], {
+      stopReason: "end_turn",
+    });
+    const authenticate = mock(() => Promise.resolve({}));
+
+    const internals = proc as unknown as {
+      connection: unknown;
+      spawnedEnv: NodeJS.ProcessEnv;
+    };
+    internals.connection = {
+      initialize: () =>
+        Promise.resolve({
+          protocolVersion: 1,
+          authMethods: options.authMethods ?? codexLikeAuthMethods,
+        }),
+      newSession,
+      loadSession,
+      resumeSession,
+      prompt,
+      authenticate,
+    };
+    internals.spawnedEnv = { ...process.env, ...options.env };
+
+    await proc.initialize();
+
+    return {
+      proc,
+      newSession,
+      loadSession,
+      resumeSession,
+      prompt,
+      authenticate,
+      internals,
+    };
+  }
+
+  test("createSession does not authenticate when newSession succeeds", async () => {
+    const { proc, newSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+    });
+
+    const sessionId = await proc.createSession("/tmp/project");
+
+    expect(sessionId).toBe("session-1");
+    expect(newSession).toHaveBeenCalledTimes(1);
+    expect(authenticate).not.toHaveBeenCalled();
+  });
+
+  test("auth_required with OPENAI-style key authenticates and retries once", async () => {
+    const { proc, newSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      newSessionRejections: [authRequiredError],
+    });
+
+    const sessionId = await proc.createSession("/tmp/project");
+
+    expect(sessionId).toBe("session-1");
+    expect(authenticate).toHaveBeenCalledTimes(1);
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+    expect(newSession).toHaveBeenCalledTimes(2);
+  });
+
+  test("auth_required with both keys present selects the earlier advertised method", async () => {
+    // Both methods are satisfiable, so this pins advertised-order selection
+    // rather than only-satisfiable selection.
+    const { proc, authenticate } = await setupAuthProcess({
+      env: { [CODEX_VAR]: "sk-codex", [OPENAI_VAR]: "sk-openai" },
+      newSessionRejections: [authRequiredError],
+    });
+
+    await proc.createSession("/tmp/project");
+
+    expect(authenticate).toHaveBeenCalledTimes(1);
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "codex-api-key" });
+  });
+
+  test("auth_required with no satisfiable env_var method throws an actionable error", async () => {
+    const { proc, newSession, authenticate } = await setupAuthProcess({
+      env: {},
+      newSessionRejections: [authRequiredError],
+    });
+
+    const promise = proc.createSession("/tmp/project");
+
+    await expect(promise).rejects.toThrow(
+      'ACP agent "codex" requires authentication',
+    );
+    await expect(promise).rejects.toThrow("Login with ChatGPT");
+    await expect(promise).rejects.toThrow(CODEX_VAR);
+    await expect(promise).rejects.toThrow(OPENAI_VAR);
+    expect(authenticate).not.toHaveBeenCalled();
+    expect(newSession).toHaveBeenCalledTimes(1);
+  });
+
+  test("a second auth_required after authenticating propagates without a retry loop", async () => {
+    const { proc, newSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      newSessionRejections: [authRequiredError, authRequiredError],
+    });
+
+    await expect(proc.createSession("/tmp/project")).rejects.toMatchObject({
+      code: -32000,
+    });
+    expect(authenticate).toHaveBeenCalledTimes(1);
+    expect(newSession).toHaveBeenCalledTimes(2);
+  });
+
+  test("an absent optional var does not block satisfiability", async () => {
+    const { proc, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      authMethods: [
+        {
+          type: "env_var",
+          id: "openai-api-key",
+          name: "Use OPENAI_API_KEY",
+          vars: [
+            { name: OPENAI_VAR },
+            { name: "VELLUM_TEST_FAKE_OPTIONAL_VAR", optional: true },
+          ],
+        },
+      ],
+      newSessionRejections: [authRequiredError],
+    });
+
+    await proc.createSession("/tmp/project");
+
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+  });
+
+  test("loadSession authenticates and retries on auth_required", async () => {
+    const { proc, loadSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      loadSessionRejections: [authRequiredError],
+    });
+
+    await proc.loadSession("session-1", "/tmp/project");
+
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+    expect(loadSession).toHaveBeenCalledTimes(2);
+  });
+
+  test("resumeSession authenticates and retries on auth_required", async () => {
+    const { proc, resumeSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      resumeSessionRejections: [authRequiredError],
+    });
+
+    await proc.resumeSession("session-1", "/tmp/project");
+
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+    expect(resumeSession).toHaveBeenCalledTimes(2);
+  });
+
+  test("prompt authenticates and retries on auth_required, returning the response", async () => {
+    const { proc, prompt, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      promptRejections: [authRequiredError],
+    });
+
+    const response = await proc.prompt("session-1", "hello");
+
+    expect(response).toEqual({ stopReason: "end_turn" });
+    expect(authenticate).toHaveBeenCalledTimes(1);
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+    expect(prompt).toHaveBeenCalledTimes(2);
+  });
+
+  test("prompt auth_required with no satisfiable env_var method throws an actionable error", async () => {
+    const { proc, prompt, authenticate } = await setupAuthProcess({
+      env: {},
+      promptRejections: [authRequiredError],
+    });
+
+    const promise = proc.prompt("session-1", "hello");
+
+    await expect(promise).rejects.toThrow(
+      'ACP agent "codex" requires authentication',
+    );
+    await expect(promise).rejects.toThrow("Login with ChatGPT");
+    await expect(promise).rejects.toThrow(CODEX_VAR);
+    await expect(promise).rejects.toThrow(OPENAI_VAR);
+    expect(authenticate).not.toHaveBeenCalled();
+    expect(prompt).toHaveBeenCalledTimes(1);
+  });
+
+  test("a non-auth error propagates without authenticating", async () => {
+    const { proc, newSession, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      newSessionRejections: [{ code: -32603, message: "Internal error" }],
+    });
+
+    await expect(proc.createSession("/tmp/project")).rejects.toMatchObject({
+      code: -32603,
+    });
+    expect(authenticate).not.toHaveBeenCalled();
+    expect(newSession).toHaveBeenCalledTimes(1);
+  });
+
+  test("auth_required after the agent process exits throws the not-spawned error", async () => {
+    const { proc, authenticate, internals } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+    });
+
+    // Simulate handleProcessExit racing the auth retry: the agent dies as it
+    // rejects with auth_required, nulling the connection before the retry
+    // path runs. spawnedEnv stays satisfiable so only the connection guard
+    // can produce the failure.
+    (internals.connection as { newSession: unknown }).newSession = () => {
+      internals.connection = null;
+      return Promise.reject(authRequiredError);
+    };
+
+    await expect(proc.createSession("/tmp/project")).rejects.toThrow(
+      'ACP agent "codex" is not spawned',
+    );
+    expect(authenticate).not.toHaveBeenCalled();
+  });
+
+  test("satisfiability uses the env captured at spawn, not live process.env", async () => {
+    const { proc, authenticate } = await setupAuthProcess({
+      env: { [OPENAI_VAR]: "sk-test" },
+      newSessionRejections: [authRequiredError],
+    });
+
+    // Drift process.env after "spawn": the earlier-advertised codex var
+    // appears, but the spawned child never saw it, so the openai method
+    // must still be selected.
+    process.env[CODEX_VAR] = "sk-codex-late";
+    try {
+      await proc.createSession("/tmp/project");
+    } finally {
+      delete process.env[CODEX_VAR];
+    }
+
+    expect(authenticate).toHaveBeenCalledTimes(1);
+    expect(authenticate).toHaveBeenCalledWith({ methodId: "openai-api-key" });
+  });
+});

package/src/acp/__tests__/prepare-agent-env.test.ts

@@ -47,8 +47,7 @@ mock.module("../../tools/credentials/metadata-store.js", () => ({
     const existing = metadataStore.get(key);
     metadataStore.set(key, {
       allowedTools: policy?.allowedTools ?? existing?.allowedTools ?? [],
-      usageDescription:
-        policy?.usageDescription ?? existing?.usageDescription,
+      usageDescription: policy?.usageDescription ?? existing?.usageDescription,
     });
     return {
       credentialId: `cred-${key}`,
@@ -331,16 +330,91 @@ describe("prepareAgentEnv - gemini gating", () => {
   });
 });
 
-describe("prepareAgentEnv — non-claude commands", () => {
-  test("returns the config unchanged for codex-acp (no required env vars today)", async () => {
+describe("prepareAgentEnv - codex-acp gating", () => {
+  function seedVaultOpenaiKey(key: string): void {
+    vaultStore.set("acp/openai_api_key", key);
+  }
+
+  function seedVaultCodexKey(key: string): void {
+    vaultStore.set("acp/codex_api_key", key);
+  }
+
+  test("injects OPENAI_API_KEY from the vault via the broker when agent.env has no override", async () => {
+    seedVaultOpenaiKey("vault-fake-openai-AAA");
+
+    const prepared = await prepareAgentEnv({
+      command: "codex-acp",
+      args: [],
+    });
+
+    expect(prepared.env?.OPENAI_API_KEY).toBe("vault-fake-openai-AAA");
+  });
+
+  test("agent.env override wins over the vault entry and skips the broker (precedence pin)", async () => {
+    // Seed a vault value but no metadata: if the override path consulted the
+    // broker anyway, ensureAcpCredentialPolicy would create metadata here.
+    seedVaultOpenaiKey("vault-fake-openai-BBB");
+
+    const prepared = await prepareAgentEnv({
+      command: "codex-acp",
+      args: [],
+      env: { OPENAI_API_KEY: "config-fake-openai-CCC" },
+    });
+
+    expect(prepared.env?.OPENAI_API_KEY).toBe("config-fake-openai-CCC");
+    expect(metadataStore.has("acp/openai_api_key")).toBe(false);
+  });
+
+  test("a vault miss for both fields does NOT throw and spawns with env unchanged (keys are optional)", async () => {
+    const prepared = await prepareAgentEnv({
+      command: "codex-acp",
+      args: [],
+      env: { NO_COLOR: "1" },
+    });
+
+    expect(prepared.env).toEqual({ NO_COLOR: "1" });
+    expect(prepared.env).not.toHaveProperty("OPENAI_API_KEY");
+    expect(prepared.env).not.toHaveProperty("CODEX_API_KEY");
+  });
+
+  test("injects CODEX_API_KEY independently of OPENAI_API_KEY", async () => {
+    seedVaultCodexKey("vault-fake-codex-DDD");
+
     const prepared = await prepareAgentEnv({
       command: "codex-acp",
       args: [],
     });
 
-    expect(prepared.env).toEqual({});
+    expect(prepared.env?.CODEX_API_KEY).toBe("vault-fake-codex-DDD");
+    expect(prepared.env).not.toHaveProperty("OPENAI_API_KEY");
   });
 
+  test("injects both keys when both vault fields are present", async () => {
+    seedVaultOpenaiKey("vault-fake-openai-EEE");
+    seedVaultCodexKey("vault-fake-codex-FFF");
+
+    const prepared = await prepareAgentEnv({
+      command: "codex-acp",
+      args: [],
+    });
+
+    expect(prepared.env?.OPENAI_API_KEY).toBe("vault-fake-openai-EEE");
+    expect(prepared.env?.CODEX_API_KEY).toBe("vault-fake-codex-FFF");
+  });
+
+  test("gates on the resolved command BASENAME (custom agent id with full path still gets injection)", async () => {
+    seedVaultOpenaiKey("vault-fake-openai-GGG");
+
+    const prepared = await prepareAgentEnv({
+      command: "/data/.bun/bin/codex-acp",
+      args: [],
+    });
+
+    expect(prepared.env?.OPENAI_API_KEY).toBe("vault-fake-openai-GGG");
+  });
+});
+
+describe("prepareAgentEnv — non-claude commands", () => {
   test("returns the config unchanged for an unrecognized command basename", async () => {
     seedVaultToken("vault-HHH");