@vellumai/assistant 0.12.3-staging.1 → 0.12.3-staging.2

package/src/credential-execution/ces-runtime.ts

@@ -9,10 +9,16 @@ import {
   setCesReconnect,
 } from "../security/secure-keys.js";
 import { getLogger } from "../util/logger.js";
-import { type CesClient, createCesClient } from "./client.js";
+import {
+  openCesRpcSession,
+  reconnectCesRpcSession,
+} from "./ces-connect.js";
+import {
+  type CesClient,
+  type CesClientHandshakeOptions,
+} from "./client.js";
 import {
   type CesProcessManager,
-  CesUnavailableError,
   createCesProcessManager,
 } from "./process-manager.js";
 import {
@@ -45,70 +51,69 @@ interface CesStartupResult {
 }
 
 /**
- * Start the CES process and perform the RPC handshake. Returns immediately with
- * handles to the in-flight initialization — callers don't need to await this
- * for startup to continue.
+ * Open the assistant's CES RPC client and perform the handshake. Returns
+ * immediately with handles to the in-flight initialization: callers don't
+ * need to await this for startup to continue.
  *
- * CES serves a multi-connection Unix socket, so this is called at the
- * process level and child processes may open their own connections.
+ * Claims reconnect ownership before any credential read so boot identity
+ * loading cannot open a second, identity-less session through the child
+ * entry point. CES serves a multi-connection Unix socket, so child
+ * processes in other address spaces still open their own connections.
  */
 function startCesProcess(config: AssistantConfig): CesStartupResult {
   const pm = createCesProcessManager({ assistantConfig: config });
   const abortController = new AbortController();
   let currentClient: CesClient | undefined;
+  let handshake: CesClientHandshakeOptions = {};
+
+  // Own this process's CES session before resolveManagedProxyContext()
+  // reads the API key. That read must not take the child open path.
+  setCesReconnect(async () => {
+    const client = await reconnectCesRpcSession(pm, handshake);
+    if (client) {
+      log.info("CES reconnection handshake accepted");
+    }
+    return client;
+  });
 
   const handshakePromise = (async (): Promise<CesClient | undefined> => {
     try {
-      const transport = await pm.start();
-      if (abortController.signal.aborted) {
-        throw new Error("CES initialization aborted during shutdown");
-      }
-      const client = createCesClient(transport);
-      currentClient = client;
       // Resolve the assistant API key so CES can use it for platform
       // credential materialisation. In managed mode the key is provisioned
-      // after hatch and stored in the credential store — CES can't read
-      // the env var, so we pass it via the handshake.
+      // after hatch and stored in the credential store. CES can't read
+      // the env var, so we pass it via the handshake. Reconnect ownership
+      // is already claimed, so this read uses HTTP or the encrypted store,
+      // not a second RPC session.
       const proxyCtx = await resolveManagedProxyContext();
+      if (abortController.signal.aborted) {
+        return undefined;
+      }
       const assistantId = getPlatformAssistantId();
-      const { accepted, reason } = await client.handshake({
+      handshake = {
         ...(proxyCtx.assistantApiKey
           ? { assistantApiKey: proxyCtx.assistantApiKey }
           : {}),
         ...(assistantId ? { assistantId } : {}),
+      };
+      const session = await openCesRpcSession({
+        processManager: pm,
+        signal: abortController.signal,
+        handshake,
       });
-      if (abortController.signal.aborted) {
-        client.close();
-        throw new Error("CES initialization aborted during shutdown");
-      }
-      if (accepted) {
+      currentClient = session?.client;
+      if (session) {
         log.info(
           "CES client initialized and handshake accepted (server-level)",
         );
-        return client;
       }
+      return session?.client;
+    } catch (err) {
       log.warn(
-        { reason },
-        "CES handshake rejected — CES tools will be unavailable",
+        { error: err instanceof Error ? err.message : String(err) },
+        "Failed to initialize CES client",
       );
-      client.close();
       currentClient = undefined;
-      await pm.stop();
-      return undefined;
-    } catch (err) {
-      if (err instanceof CesUnavailableError) {
-        log.info(
-          { reason: err.message },
-          "CES is not available — CES tools will be unavailable",
-        );
-      } else {
-        log.warn(
-          { error: err instanceof Error ? err.message : String(err) },
-          "Failed to initialize CES client — CES tools will be unavailable",
-        );
-      }
       await pm.stop().catch(() => {});
-      currentClient = undefined;
       return undefined;
     }
   })();
@@ -154,11 +159,11 @@ function updateClientRef(client: CesClient | undefined): void {
 }
 
 /**
- * Bring up the daemon's CES connection: start the process, run the handshake
- * (blocking up to a 20s timeout so credential reads can route through CES
- * before provider init), register the reconnection callback, and keep the live
- * client reference in sync. Non-fatal — on failure the daemon falls back to the
- * direct credential store.
+ * Open the assistant's CES RPC client: handshake (blocking up to a 20s
+ * timeout so credential reads can route through CES before provider init)
+ * and keep the live client reference in sync. Reconnect ownership is
+ * claimed before the identity read. Non-fatal: on failure the assistant
+ * falls back to the direct credential store.
  */
 export async function startCes(config: AssistantConfig): Promise<void> {
   const cesResult = startCesProcess(config);
@@ -171,7 +176,7 @@ export async function startCes(config: AssistantConfig): Promise<void> {
       timeoutMs: DEFAULT_CES_STARTUP_TIMEOUT_MS,
       onTimeout: () => {
         log.warn(
-          "CES handshake timed out after 20s — falling back to direct credential store",
+          "CES handshake timed out after 20s, falling back to direct credential store",
         );
       },
     });
@@ -193,56 +198,13 @@ export async function startCes(config: AssistantConfig): Promise<void> {
     }
   }
 
-  // Register CES reconnection callback so the credential layer can re-establish
-  // the connection when the transport dies, instead of falling back to the
-  // encrypted file store.
+  // Reconnect ownership is claimed inside startCesProcess before the
+  // identity read. Snapshotting the API key there (not here) avoids a
+  // second resolveManagedProxyContext() after setCesClient, which would
+  // read the key through CES while reconnecting.
   if (cesResult.processManager) {
     const pm = cesResult.processManager;
 
-    // Snapshot the managed-proxy context and assistant ID at CES startup so the
-    // reconnect closure below never calls back into `resolveManagedProxyContext()`.
-    // That function reads the assistant API key via `getSecureKeyAsync()`, which
-    // — once `setCesClient()` has resolved the backend to CES RPC — routes the
-    // read through CES itself. During a reconnect the old transport is dead and
-    // a new one is being set up by this very closure, so the nested credential
-    // read recursively awaits its own in-flight reconnection and deadlocks until
-    // `CREDENTIAL_OP_TIMEOUT_MS` (45s) fires. That 45-second stall delays every
-    // CES restart and causes dependent credential reads (e.g. Meet's STT
-    // provider resolution) to return `undefined` during the window. API key
-    // rotation uses the `updateAssistantApiKey` RPC on the live client, not a
-    // reconnect, so caching at startup is safe.
-    const startupProxyCtx = await resolveManagedProxyContext();
-    const startupAssistantId = getPlatformAssistantId();
-
-    setCesReconnect(async () => {
-      try {
-        await pm.stop();
-        const transport = await pm.start();
-        const newClient = createCesClient(transport);
-        const { accepted, reason } = await newClient.handshake({
-          ...(startupProxyCtx.assistantApiKey
-            ? { assistantApiKey: startupProxyCtx.assistantApiKey }
-            : {}),
-          ...(startupAssistantId ? { assistantId: startupAssistantId } : {}),
-        });
-        if (accepted) {
-          log.info("CES reconnection handshake accepted");
-          return newClient;
-        }
-        log.warn({ reason }, "CES reconnection handshake rejected");
-        newClient.close();
-        await pm.stop().catch(() => {});
-        return undefined;
-      } catch (err) {
-        log.warn(
-          { error: err instanceof Error ? err.message : String(err) },
-          "CES reconnection attempt failed",
-        );
-        await pm.stop().catch(() => {});
-        return undefined;
-      }
-    });
-
     // Proactive reconnect: when the transport dies (socket close, process
     // exit), start a retry-with-backoff loop immediately instead of waiting
     // for the next credential operation to trigger the lazy reconnect path.

package/src/credential-execution/executable-discovery.ts

@@ -53,9 +53,7 @@ export type DiscoveryResult = ManagedDiscoverySuccess | DiscoveryFailure;
  * processes can each connect. The actual connection is made later by
  * `CesProcessManager.start()`.
  */
-export function discoverManagedCes():
-  | ManagedDiscoverySuccess
-  | DiscoveryFailure {
+export function discoverCes(): DiscoveryResult {
   const socketPath = getCesSocketPath();
 
   if (!isNamedPipePath(socketPath) && !existsSync(socketPath)) {
@@ -68,14 +66,6 @@ export function discoverManagedCes():
   return { mode: "managed", socketPath };
 }
 
-/**
- * Discover CES for the current process. Local and managed topologies use
- * the same bootstrap socket path.
- */
-export function discoverCes(): DiscoveryResult {
-  return discoverManagedCes();
-}
-
 /** How long to poll for the CES socket before giving up. */
 const MANAGED_DISCOVERY_TIMEOUT_MS = 3_000;
 

package/src/daemon/__tests__/conversation-surfaces-launch.test.ts

@@ -440,7 +440,12 @@ describe("handleSurfaceAction — launch_conversation dispatch", () => {
     // claimed to eliminate that round-trip; this test enforces it.
     nextKeyStoreResult = { conversationId: "conv-pending-set" };
 
-    const ctx = makeContext();
+    const settleStructuralWait = mock(() => true);
+    const ctx = makeContext({
+      modeSessions: {
+        settleStructuralWait,
+      } as unknown as Conversation["modeSessions"],
+    });
     registerCardSurface(ctx, "surface-pending");
     // Simulate `ui_show` having stamped a pending entry for this surface
     // (which it does for any interactive card, including persistent ones).
@@ -473,6 +478,10 @@ describe("handleSurfaceAction — launch_conversation dispatch", () => {
     // Pending entry was deleted so subsequent sibling clicks on the same
     // persistent card aren't blocked behind a stale "owes-an-answer" flag.
     expect(ctx.pendingSurfaceActions.has("surface-pending")).toBe(false);
+    expect(settleStructuralWait).toHaveBeenCalledWith(
+      { kind: "surface", responseId: "surface-pending" },
+      { status: "completed", endReason: "surface_launch_settled" },
+    );
 
     await processStartedPromise;
     resolveProcess();

package/src/daemon/__tests__/inference-profile-notification.test.ts

@@ -47,6 +47,7 @@ mock.module("../../plugins/defaults/memory/v2/activation-log-store.js", () => ({
 }));
 
 // ── Imports (after mocks) ────────────────────────────────────────────────────
+import { mockUnownedModeSessions } from "../../__tests__/helpers/mock-conversation.js";
 import type { AgentEvent } from "../../agent/loop.js";
 import type { Message } from "../../providers/types.js";
 import type {
@@ -68,6 +69,7 @@ function makeDeps(): EventHandlerDeps {
       provider: { name: "mock-provider" },
       currentTurnSurfaces: [],
       trustContext: undefined,
+      modeSessions: mockUnownedModeSessions(),
     } as unknown as EventHandlerDeps["ctx"],
     onEvent: () => {},
     reqId: "req-profile-notify",

package/src/daemon/browser-mode-session.test.ts

@@ -0,0 +1,402 @@
+import { afterEach, describe, expect, test } from "bun:test";
+
+import { setOverridesForTesting } from "../__tests__/feature-flag-test-helpers.js";
+import type { ModeSession } from "../api/mode-session.js";
+import { isSessionGroupsEnabled } from "../config/session-groups-gate.js";
+import { BrowserModeSessionProducer } from "./browser-mode-session.js";
+import type {
+  ModeSessionSourceHandle,
+  ModeSessionTerminalDisposition,
+} from "./conversation-mode-session.js";
+
+function createCoordinator() {
+  let owner: ModeSession | undefined;
+  let activation = 0;
+  const activated: ModeSessionSourceHandle[] = [];
+  const retired: Array<{
+    handle: ModeSessionSourceHandle;
+    disposition?: ModeSessionTerminalDisposition;
+  }> = [];
+  const activity: Array<{ turnId: string; at: number }> = [];
+
+  const coordinator = {
+    activateSource(input: {
+      sourceId: string;
+      generation: number;
+      mode: "browser";
+      sourceStartedAt: number;
+    }): ModeSessionSourceHandle {
+      activation += 1;
+      const handle = {
+        sourceId: input.sourceId,
+        generation: input.generation,
+        activation,
+        id: `browser-${activation}`,
+        mode: input.mode,
+      } as const;
+      activated.push(handle);
+      return handle;
+    },
+    claimTurn(
+      _turnId: string,
+      handle: ModeSessionSourceHandle,
+      _at: number,
+    ): ModeSession {
+      owner ??= { id: handle.id, mode: handle.mode };
+      return owner;
+    },
+    getTurnOwner: () => owner,
+    recordActivity: (turnId: string, at: number) => {
+      activity.push({ turnId, at });
+      return true;
+    },
+    retireSource: (
+      handle: ModeSessionSourceHandle,
+      disposition?: ModeSessionTerminalDisposition,
+    ) => {
+      retired.push({ handle, disposition });
+      return true;
+    },
+  };
+
+  return {
+    coordinator,
+    activated,
+    retired,
+    activity,
+    clearOwner: () => {
+      owner = undefined;
+    },
+    setOwner: (next: ModeSession) => {
+      owner = next;
+    },
+  };
+}
+
+describe("BrowserModeSessionProducer", () => {
+  afterEach(() => {
+    setOverridesForTesting({});
+  });
+
+  test.each([false, true])(
+    "retires failed admission while preserving a partial turn owner (%s)",
+    (partial) => {
+      const state = createCoordinator();
+      const claim = state.coordinator.claimTurn;
+      state.coordinator.claimTurn = (...args) => {
+        if (partial) {
+          claim(...args);
+        }
+        throw new Error("tracking unavailable");
+      };
+      const producer = new BrowserModeSessionProducer(state.coordinator);
+      expect(() =>
+        producer.beginOperation({
+          turnId: "turn-123",
+          lifecycle: "action",
+          at: 100,
+        }),
+      ).toThrow("tracking unavailable");
+      expect(state.retired).toHaveLength(1);
+      expect(state.retired[0]?.disposition).toMatchObject({
+        status: "interrupted",
+        endReason: "tracking_failed",
+      });
+      if (partial) {
+        expect(state.coordinator.getTurnOwner()).toBeDefined();
+      } else {
+        expect(state.coordinator.getTurnOwner()).toBeUndefined();
+      }
+    },
+  );
+
+  test("admits new tracking only while the shared feature flag is enabled", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(
+      state.coordinator,
+      1,
+      isSessionGroupsEnabled,
+    );
+
+    setOverridesForTesting({});
+    expect(isSessionGroupsEnabled()).toBe(false);
+    expect(
+      producer.beginOperation({
+        turnId: "turn-disabled",
+        lifecycle: "action",
+        at: 100,
+      }),
+    ).toBeUndefined();
+    expect(state.activated).toEqual([]);
+
+    setOverridesForTesting({ "session-groups": true });
+    const admitted = producer.beginOperation({
+      turnId: "turn-enabled",
+      lifecycle: "action",
+      at: 110,
+    });
+    expect(admitted?.owner).toEqual({ id: "browser-1", mode: "browser" });
+
+    setOverridesForTesting({ "session-groups": false });
+    state.clearOwner();
+    const close = producer.beginOperation({
+      turnId: "turn-close",
+      lifecycle: "terminal",
+      at: 120,
+    });
+    expect(close?.owner).toEqual({ id: "browser-1", mode: "browser" });
+    expect(
+      producer.finishOperation(close, {
+        at: 130,
+        isError: false,
+        cancelled: false,
+        terminalReason: "browser_closed",
+      }),
+    ).toBe(true);
+    expect(state.retired).toHaveLength(1);
+  });
+
+  test("claims the first action and records repeated successful activity", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+
+    const first = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 100,
+    });
+    expect(first?.owner).toEqual({ id: "browser-1", mode: "browser" });
+    expect(state.activated).toHaveLength(1);
+    expect(
+      producer.finishOperation(first, {
+        at: 120,
+        isError: false,
+        cancelled: false,
+      }),
+    ).toBe(true);
+    expect(state.activity).toEqual([{ turnId: "turn-1", at: 120 }]);
+
+    const repeated = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 130,
+    });
+    expect(repeated?.owner.id).toBe("browser-1");
+    expect(state.activated).toHaveLength(1);
+  });
+
+  test("status and a terminal operation without an active run are inert", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+
+    expect(
+      producer.beginOperation({
+        turnId: "turn-1",
+        lifecycle: "status",
+        at: 100,
+      }),
+    ).toBeUndefined();
+    expect(
+      producer.beginOperation({
+        turnId: "turn-1",
+        lifecycle: "terminal",
+        at: 110,
+      }),
+    ).toBeUndefined();
+    expect(state.activated).toEqual([]);
+  });
+
+  test("keeps an existing non-browser owner without activating a nested run", () => {
+    const state = createCoordinator();
+    state.setOwner({ id: "live-1", mode: "live_vision" });
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+
+    const token = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 100,
+    });
+    expect(token).toMatchObject({ owner: { id: "live-1" } });
+    expect(token?.handle).toBeUndefined();
+    expect(state.activated).toEqual([]);
+    expect(
+      producer.finishOperation(token, {
+        at: 120,
+        isError: false,
+        cancelled: false,
+      }),
+    ).toBe(true);
+    expect(state.activity).toEqual([{ turnId: "turn-1", at: 120 }]);
+  });
+
+  test("successful close retires the active run with a completed disposition", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+    producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 100,
+    });
+    state.clearOwner();
+
+    const close = producer.beginOperation({
+      turnId: "turn-close",
+      lifecycle: "terminal",
+      at: 200,
+    });
+    expect(close?.owner.id).toBe("browser-1");
+    expect(
+      producer.finishOperation(close, {
+        at: 220,
+        isError: false,
+        cancelled: false,
+        terminalReason: "browser_closed",
+      }),
+    ).toBe(true);
+    expect(state.retired).toEqual([
+      {
+        handle: expect.objectContaining({ id: "browser-1", activation: 1 }),
+        disposition: {
+          status: "completed",
+          endReason: "browser_closed",
+        },
+      },
+    ]);
+  });
+
+  test("a foreign-owned terminal operation retires only the captured browser run", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+    producer.beginOperation({
+      turnId: "turn-browser",
+      lifecycle: "action",
+      at: 100,
+    });
+    state.setOwner({ id: "live-1", mode: "live_vision" });
+
+    const close = producer.beginOperation({
+      turnId: "turn-live",
+      lifecycle: "terminal",
+      at: 200,
+    });
+    expect(close).toMatchObject({
+      owner: { id: "live-1", mode: "live_vision" },
+      handle: { id: "browser-1", activation: 1 },
+    });
+    expect(
+      producer.finishOperation(close, {
+        at: 220,
+        isError: false,
+        cancelled: false,
+        terminalReason: "browser_detached",
+      }),
+    ).toBe(true);
+    expect(state.retired.at(-1)).toEqual({
+      handle: expect.objectContaining({ id: "browser-1", activation: 1 }),
+      disposition: {
+        status: "completed",
+        endReason: "browser_detached",
+      },
+    });
+
+    state.clearOwner();
+    producer.beginOperation({
+      turnId: "turn-browser-2",
+      lifecycle: "action",
+      at: 300,
+    });
+    expect(
+      producer.finishOperation(close, {
+        at: 320,
+        isError: true,
+        cancelled: false,
+      }),
+    ).toBe(false);
+    expect(state.retired).toHaveLength(1);
+  });
+
+  test("a failed close keeps the session active until a successful retry", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+    const action = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 100,
+    });
+    const close = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "terminal",
+      at: 110,
+    });
+    expect(
+      producer.finishOperation(close, {
+        at: 120,
+        isError: true,
+        cancelled: false,
+      }),
+    ).toBe(true);
+    expect(state.retired).toHaveLength(0);
+    expect(state.activity).toEqual([{ turnId: "turn-1", at: 120 }]);
+
+    const retry = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "terminal",
+      at: 130,
+    });
+    expect(retry?.handle).toEqual(action?.handle);
+    producer.finishOperation(retry, {
+      at: 140,
+      isError: false,
+      cancelled: false,
+    });
+    expect(state.retired[0]?.disposition).toEqual({
+      status: "completed",
+      endReason: "browser_closed",
+    });
+  });
+
+  test("cancellation retires the exact run while accepted output remains on its captured turn", () => {
+    const state = createCoordinator();
+    const producer = new BrowserModeSessionProducer(state.coordinator);
+    const stale = producer.beginOperation({
+      turnId: "turn-1",
+      lifecycle: "action",
+      at: 100,
+    });
+    expect(
+      producer.finishOperation(stale, {
+        at: 120,
+        isError: true,
+        cancelled: true,
+      }),
+    ).toBe(true);
+    expect(state.retired[0]?.disposition).toEqual({
+      status: "interrupted",
+      endReason: "browser_operation_cancelled",
+    });
+
+    state.clearOwner();
+    const current = producer.beginOperation({
+      turnId: "turn-2",
+      lifecycle: "action",
+      at: 200,
+    });
+    expect(current?.handle?.activation).toBe(2);
+    expect(
+      producer.finishOperation(stale, {
+        at: 210,
+        isError: true,
+        cancelled: true,
+      }),
+    ).toBe(false);
+    expect(
+      producer.finishOperation(stale, {
+        at: 220,
+        isError: false,
+        cancelled: false,
+      }),
+    ).toBe(true);
+    expect(state.activity).toEqual([{ turnId: "turn-1", at: 220 }]);
+    expect(state.retired).toHaveLength(1);
+  });
+});