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

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

@@ -0,0 +1,186 @@
+import type { ModeSession } from "../api/mode-session.js";
+import type { BrowserOperationLifecycle } from "../browser/operations.js";
+import type {
+  ConversationModeSessionCoordinator,
+  ModeSessionSourceHandle,
+  ModeSessionTerminalDisposition,
+} from "./conversation-mode-session.js";
+import { claimModeSessionTurn } from "./mode-session-tracking.js";
+
+export type { BrowserOperationLifecycle } from "../browser/operations.js";
+
+type ModeSessionCoordinator = Pick<
+  ConversationModeSessionCoordinator,
+  | "activateSource"
+  | "claimTurn"
+  | "getTurnOwner"
+  | "recordActivity"
+  | "retireSource"
+>;
+
+export interface BrowserOperationToken {
+  readonly turnId: string;
+  readonly lifecycle: Exclude<BrowserOperationLifecycle, "status">;
+  readonly owner: ModeSession;
+  readonly handle?: ModeSessionSourceHandle;
+}
+
+const BROWSER_SOURCE_ID = "browser";
+
+/** Maps typed browser operations onto the shared session coordinator. */
+export class BrowserModeSessionProducer {
+  readonly #coordinator: ModeSessionCoordinator;
+  readonly #generation: number;
+  readonly #isTrackingEnabled: () => boolean;
+  #activeHandle?: ModeSessionSourceHandle;
+
+  constructor(
+    coordinator: ModeSessionCoordinator,
+    generation = 1,
+    isTrackingEnabled: () => boolean = () => true,
+  ) {
+    this.#coordinator = coordinator;
+    this.#generation = generation;
+    this.#isTrackingEnabled = isTrackingEnabled;
+  }
+
+  beginOperation(input: {
+    turnId: string;
+    lifecycle: BrowserOperationLifecycle;
+    at: number;
+  }): BrowserOperationToken | undefined {
+    if (input.lifecycle === "status") {
+      return undefined;
+    }
+
+    const existingOwner = this.#coordinator.getTurnOwner(input.turnId);
+    if (existingOwner) {
+      const handle = this.#activeHandle;
+      return {
+        turnId: input.turnId,
+        lifecycle: input.lifecycle,
+        owner: existingOwner,
+        ...(handle &&
+        (input.lifecycle === "terminal" || handle.id === existingOwner.id)
+          ? { handle }
+          : {}),
+      };
+    }
+
+    if (input.lifecycle === "terminal") {
+      const handle = this.#activeHandle;
+      if (!handle) {
+        return undefined;
+      }
+      const owner = claimModeSessionTurn(
+        this.#coordinator,
+        input.turnId,
+        handle,
+        input.at,
+      );
+      if (!owner || owner.id !== handle.id) {
+        return undefined;
+      }
+      return {
+        turnId: input.turnId,
+        lifecycle: input.lifecycle,
+        owner,
+        handle,
+      };
+    }
+
+    if (!this.#isTrackingEnabled()) {
+      return undefined;
+    }
+
+    const handle = this.#coordinator.activateSource({
+      sourceId: BROWSER_SOURCE_ID,
+      generation: this.#generation,
+      mode: "browser",
+      sourceStartedAt: input.at,
+    });
+    if (!handle) {
+      return undefined;
+    }
+    const owner = claimModeSessionTurn(
+      this.#coordinator,
+      input.turnId,
+      handle,
+      input.at,
+    );
+    if (!owner) {
+      return undefined;
+    }
+    if (owner.id === handle.id) {
+      this.#activeHandle = handle;
+    }
+    return {
+      turnId: input.turnId,
+      lifecycle: input.lifecycle,
+      owner,
+      ...(owner.id === handle.id ? { handle } : {}),
+    };
+  }
+
+  finishOperation(
+    token: BrowserOperationToken | undefined,
+    outcome: {
+      at: number;
+      isError: boolean;
+      cancelled: boolean;
+      terminalReason?: "browser_closed" | "browser_detached";
+    },
+  ): boolean {
+    if (!token) {
+      return false;
+    }
+
+    const handle = token.handle;
+    if (outcome.cancelled) {
+      if (!handle || !this.#isCurrent(handle)) {
+        return false;
+      }
+      return this.#retire(handle, {
+        status: "interrupted",
+        endReason: "browser_operation_cancelled",
+      });
+    }
+
+    if (token.lifecycle === "terminal") {
+      if (!handle || !this.#isCurrent(handle)) {
+        return false;
+      }
+      if (!outcome.isError) {
+        return this.#retire(handle, {
+          status: "completed",
+          endReason: outcome.terminalReason ?? "browser_closed",
+        });
+      }
+    }
+
+    return this.#coordinator.recordActivity(token.turnId, outcome.at);
+  }
+
+  #isCurrent(handle: ModeSessionSourceHandle): boolean {
+    const active = this.#activeHandle;
+    return (
+      active?.sourceId === handle.sourceId &&
+      active.generation === handle.generation &&
+      active.activation === handle.activation &&
+      active.id === handle.id
+    );
+  }
+
+  #retire(
+    handle: ModeSessionSourceHandle,
+    disposition: ModeSessionTerminalDisposition,
+  ): boolean {
+    try {
+      return this.#coordinator.retireSource(handle, disposition);
+    } finally {
+      if (this.#isCurrent(handle)) {
+        this.#activeHandle = undefined;
+      }
+    }
+  }
+}

package/src/daemon/computer-use-mode-session.test.ts

@@ -0,0 +1,210 @@
+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 { ComputerUseModeSessionProducer } from "./computer-use-mode-session.js";
+import type { ModeSessionSourceHandle } from "./conversation-mode-session.js";
+
+function createCoordinator() {
+  let owner: ModeSession | undefined;
+  let activation = 0;
+  const activated: Array<{
+    sourceId: string;
+    generation: number;
+    mode: "computer_use";
+    sourceStartedAt: number;
+  }> = [];
+  const retired: ModeSessionSourceHandle[] = [];
+  const activity: Array<{ turnId: string; at: number }> = [];
+  const draining: string[] = [];
+
+  const coordinator = {
+    activateSource(input: {
+      sourceId: string;
+      generation: number;
+      mode: "computer_use";
+      sourceStartedAt: number;
+    }): ModeSessionSourceHandle {
+      activated.push(input);
+      activation += 1;
+      return {
+        sourceId: input.sourceId,
+        generation: input.generation,
+        activation,
+        id: `session-${activation}`,
+        mode: input.mode,
+      };
+    },
+    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) => {
+      retired.push(handle);
+      return true;
+    },
+    beginDraining: (turnId: string) => {
+      draining.push(turnId);
+      return true;
+    },
+  };
+
+  return {
+    coordinator,
+    activated,
+    retired,
+    activity,
+    draining,
+    clearOwner: () => {
+      owner = undefined;
+    },
+    setOwner: (next: ModeSession) => {
+      owner = next;
+    },
+  };
+}
+
+describe("ComputerUseModeSessionProducer", () => {
+  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 ComputerUseModeSessionProducer(state.coordinator);
+      expect(() =>
+        producer.recordAction({
+          turnId: "turn-123",
+          source: { sourceId: "computer", generation: 1 },
+          at: 100,
+        }),
+      ).toThrow("tracking unavailable");
+      expect(state.retired).toHaveLength(1);
+      expect(state.retired[0]?.mode).toBe("computer_use");
+      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 ComputerUseModeSessionProducer(
+      state.coordinator,
+      isSessionGroupsEnabled,
+    );
+    const source = { sourceId: "proxy-123", generation: 2 };
+
+    setOverridesForTesting({ "session-groups": false });
+    expect(
+      producer.recordAction({ turnId: "turn-disabled", source, at: 100 }),
+    ).toBeUndefined();
+    expect(state.activated).toEqual([]);
+
+    setOverridesForTesting({ "session-groups": true });
+    expect(
+      producer.recordAction({ turnId: "turn-enabled", source, at: 110 }),
+    ).toEqual({ id: "session-1", mode: "computer_use" });
+
+    setOverridesForTesting({ "session-groups": false });
+    expect(producer.endTask({ turnId: "turn-enabled", source })).toBe(true);
+    expect(state.retired).toHaveLength(1);
+    expect(state.draining).toEqual(["turn-enabled"]);
+  });
+
+  test("activates on the first action and records later action activity", () => {
+    const state = createCoordinator();
+    const producer = new ComputerUseModeSessionProducer(state.coordinator);
+    const source = { sourceId: "proxy-123", generation: 2 };
+
+    expect(
+      producer.recordAction({ turnId: "turn-123", source, at: 100 }),
+    ).toEqual({ id: "session-1", mode: "computer_use" });
+    expect(state.activated).toEqual([
+      expect.objectContaining({ ...source, sourceStartedAt: 100 }),
+    ]);
+
+    producer.recordAction({ turnId: "turn-123", source, at: 120 });
+    expect(state.activated).toHaveLength(1);
+    expect(state.activity).toEqual([{ turnId: "turn-123", at: 120 }]);
+  });
+
+  test("keeps an existing mode owner without activating a secondary run", () => {
+    const state = createCoordinator();
+    state.setOwner({ id: "live-123", mode: "live_vision" });
+    const producer = new ComputerUseModeSessionProducer(state.coordinator);
+
+    expect(
+      producer.recordAction({
+        turnId: "turn-123",
+        source: { sourceId: "proxy-123", generation: 0 },
+        at: 100,
+      }),
+    ).toEqual({ id: "live-123", mode: "live_vision" });
+    expect(state.activated).toEqual([]);
+    expect(state.activity).toEqual([{ turnId: "turn-123", at: 100 }]);
+  });
+
+  test("terminal end retires its exact handle and begins final reply drain", () => {
+    const state = createCoordinator();
+    const producer = new ComputerUseModeSessionProducer(state.coordinator);
+    const source = { sourceId: "proxy-123", generation: 3 };
+    producer.recordAction({ turnId: "turn-123", source, at: 100 });
+
+    expect(producer.endTask({ turnId: "turn-123", source })).toBe(true);
+    expect(state.retired).toEqual([
+      expect.objectContaining({
+        sourceId: "proxy-123",
+        generation: 3,
+        activation: 1,
+        id: "session-1",
+      }),
+    ]);
+    expect(state.draining).toEqual(["turn-123"]);
+  });
+
+  test("stale terminal callbacks cannot retire a replacement activation", () => {
+    const state = createCoordinator();
+    const producer = new ComputerUseModeSessionProducer(state.coordinator);
+    const oldSource = { sourceId: "proxy-123", generation: 3 };
+    producer.recordAction({ turnId: "turn-old", source: oldSource, at: 100 });
+    state.clearOwner();
+
+    const nextSource = { sourceId: "proxy-123", generation: 4 };
+    producer.recordAction({ turnId: "turn-new", source: nextSource, at: 200 });
+
+    expect(producer.endTask({ turnId: "turn-old", source: oldSource })).toBe(
+      false,
+    );
+    expect(state.retired).toEqual([]);
+    expect(producer.endTask({ turnId: "turn-new", source: nextSource })).toBe(
+      true,
+    );
+    expect(state.retired[0]).toMatchObject({
+      generation: 4,
+      activation: 2,
+      id: "session-2",
+    });
+  });
+});

package/src/daemon/computer-use-mode-session.ts

@@ -0,0 +1,105 @@
+import type { ModeSession } from "../api/mode-session.js";
+import type {
+  ConversationModeSessionCoordinator,
+  ModeSessionSourceHandle,
+} from "./conversation-mode-session.js";
+import { claimModeSessionTurn } from "./mode-session-tracking.js";
+
+type ModeSessionCoordinator = Pick<
+  ConversationModeSessionCoordinator,
+  | "activateSource"
+  | "beginDraining"
+  | "claimTurn"
+  | "getTurnOwner"
+  | "recordActivity"
+  | "retireSource"
+>;
+
+export interface ComputerUseSourceIdentity {
+  sourceId: string;
+  generation: number;
+}
+
+/** Maps validated host computer actions onto the shared session coordinator. */
+export class ComputerUseModeSessionProducer {
+  readonly #coordinator: ModeSessionCoordinator;
+  readonly #isTrackingEnabled: () => boolean;
+  #activeHandle?: ModeSessionSourceHandle;
+
+  constructor(
+    coordinator: ModeSessionCoordinator,
+    isTrackingEnabled: () => boolean = () => true,
+  ) {
+    this.#coordinator = coordinator;
+    this.#isTrackingEnabled = isTrackingEnabled;
+  }
+
+  recordAction(input: {
+    turnId: string;
+    source: ComputerUseSourceIdentity;
+    at: number;
+  }): ModeSession | undefined {
+    const existingOwner = this.#coordinator.getTurnOwner(input.turnId);
+    if (existingOwner) {
+      this.#coordinator.recordActivity(input.turnId, input.at);
+      return existingOwner;
+    }
+
+    if (!this.#isTrackingEnabled()) {
+      return undefined;
+    }
+
+    const handle = this.#coordinator.activateSource({
+      ...input.source,
+      mode: "computer_use",
+      sourceStartedAt: input.at,
+    });
+    if (!handle) {
+      return undefined;
+    }
+    const owner = claimModeSessionTurn(
+      this.#coordinator,
+      input.turnId,
+      handle,
+      input.at,
+    );
+    if (owner?.id === handle.id) {
+      this.#activeHandle = handle;
+    }
+    return owner;
+  }
+
+  endTask(input: {
+    turnId?: string;
+    source: ComputerUseSourceIdentity;
+    disposition?: {
+      status: "completed" | "interrupted";
+      endReason: string;
+    };
+  }): boolean {
+    const handle = this.#activeHandle;
+    if (
+      !handle ||
+      handle.sourceId !== input.source.sourceId ||
+      handle.generation !== input.source.generation
+    ) {
+      return false;
+    }
+
+    this.#activeHandle = undefined;
+    const retired = this.#coordinator.retireSource(
+      handle,
+      input.disposition ?? {
+        status: "completed",
+        endReason: "computer_use_ended",
+      },
+    );
+    if (
+      input.turnId &&
+      this.#coordinator.getTurnOwner(input.turnId)?.id === handle.id
+    ) {
+      this.#coordinator.beginDraining(input.turnId);
+    }
+    return retired;
+  }
+}

package/src/daemon/config-watcher.ts

@@ -190,12 +190,22 @@ export class ConfigWatcher {
         if (this.suppressReload) {
           return;
         }
+        const prevMcpGlobalMax = this.lastConfig?.tools.mcpGlobalMaxTools;
         try {
           const changed = await this.refreshConfigFromSources();
           if (changed) {
             evictConversationsForReload();
             refreshAuthenticatedApiRateLimit();
             publishConfigChanged();
+            const nextMcpGlobalMax = getConfig().tools.mcpGlobalMaxTools;
+            if (prevMcpGlobalMax !== nextMcpGlobalMax) {
+              reloadMcpServers().catch((err: unknown) => {
+                log.error(
+                  { err },
+                  "MCP reload after tools.mcpGlobalMaxTools change failed",
+                );
+              });
+            }
           }
         } catch (err) {
           log.error(

package/src/daemon/conversation-agent-loop-handlers.ts

@@ -14,6 +14,11 @@ import { onActivationToolCall } from "../activation/turn-hooks.js";
 import type { AgentEvent } from "../agent/loop.js";
 import { resolveComputerUseToolName } from "../api/computer-use-tool.js";
 import type { AnsweredQuestion } from "../api/events/question-answered.js";
+import type {
+  ToolActivityMetadata,
+  WebSearchMetadata,
+  WebSearchResultItem,
+} from "../api/events/tool-result.js";
 import type { AssistantEvent } from "../api/index.js";
 import type {
   TurnChannelContext,
@@ -147,11 +152,7 @@ import type {
   SurfaceAction,
   UiSurfaceShow,
 } from "./message-protocol.js";
-import type {
-  ToolActivityMetadata,
-  WebSearchMetadata,
-  WebSearchResultItem,
-} from "./message-types/web-activity.js";
+import { bestEffortModeSessionTracking } from "./mode-session-tracking.js";
 import { referenceMediaBlocksForPersist } from "./persist-media-references.js";
 import { buildProviderRejectionLogFields } from "./provider-rejection-log-fields.js";
 import { turnOrRestingTrust } from "./trust-context-types.js";
@@ -1205,7 +1206,7 @@ async function reserveInflightMessageRow(
   conversationId: string,
   role: "assistant" | "user",
   metadata: Record<string, unknown> | undefined,
-): Promise<{ id: string }> {
+): Promise<{ id: string; createdAt: number }> {
   const writer = createInflightContentWriter(conversationId);
   const reserved = await reserveMessage(
     conversationId,
@@ -1644,7 +1645,11 @@ export async function handleLlmCallStarted(
     }
   }
 
-  const metadata = buildAssistantChannelMetadata(state, deps);
+  const modeSession = deps.ctx.modeSessions.getTurnOwner(deps.reqId);
+  const metadata = {
+    ...buildAssistantChannelMetadata(state, deps),
+    ...(modeSession ? { modeSession } : {}),
+  };
   const reservedRow = await reserveInflightMessageRow(
     state,
     deps.ctx.conversationId,
@@ -1652,6 +1657,13 @@ export async function handleLlmCallStarted(
     metadata,
   );
   state.lastAssistantMessageId = reservedRow.id;
+  bestEffortModeSessionTracking("assistant row reservation", () =>
+    deps.ctx.modeSessions.trackPersistedRow(
+      deps.reqId,
+      reservedRow.id,
+      reservedRow.createdAt,
+    ),
+  );
   state.assistantRowAwaitingFinalization = true;
   // Fresh row → fresh accumulator. If an earlier (failed) LLM call
   // within the same run left partial state behind, the
@@ -1663,6 +1675,7 @@ export async function handleLlmCallStarted(
     type: "assistant_turn_start",
     messageId: reservedRow.id,
     conversationId: deps.ctx.conversationId,
+    ...(modeSession ? { modeSession } : {}),
   });
 }
 
@@ -1881,6 +1894,7 @@ export function handleToolUse(
     // Carry the first-byte timestamp through so a client that connected after
     // the preview event still anchors the perceived-latency timer to it.
     previewStartedAt: state.toolPreviewStartedAt.get(event.id),
+    modeSession: deps.ctx.modeSessions.getTurnOwner(deps.reqId),
   });
   // `message_complete` always precedes tool events (see handleMessageComplete),
   // so this tool_use block is already durable in the assistant row. The
@@ -2152,6 +2166,7 @@ function buildToolResultBlocks(
 function buildToolResultMetadata(
   deps: EventHandlerDeps,
 ): Record<string, unknown> {
+  const modeSession = deps.ctx.modeSessions.getTurnOwner(deps.reqId);
   return {
     ...provenanceFromTrustContext(turnOrRestingTrust(deps.ctx)),
     userMessageChannel: deps.turnChannelContext.userMessageChannel,
@@ -2159,6 +2174,7 @@ function buildToolResultMetadata(
     userMessageInterface: deps.turnInterfaceContext.userMessageInterface,
     assistantMessageInterface:
       deps.turnInterfaceContext.assistantMessageInterface,
+    ...(modeSession ? { modeSession } : {}),
   };
 }
 
@@ -2215,6 +2231,19 @@ async function persistPendingToolResultRow(
     deps.ctx.conversationId,
     buildToolResultMetadata(deps),
   );
+  bestEffortModeSessionTracking("tool result persistence", () => {
+    const row = getMessageById(rowId, deps.ctx.conversationId);
+    if (row) {
+      deps.ctx.modeSessions.trackPersistedRow(
+        deps.reqId,
+        rowId,
+        row.createdAt,
+        {
+          startsDisplayBoundary: false,
+        },
+      );
+    }
+  });
   // Snapshot the batch after the reservation resolves so the last of the
   // concurrent writers reflects the fullest batch. On-arrival writes go to
   // the in-flight delta file; the finalize seam folds the row inline.
@@ -2253,7 +2282,7 @@ export async function finalizePendingToolResultRow(
   conversationId: string,
   metadata: Record<string, unknown>,
   rlog: pino.Logger,
-): Promise<void> {
+): Promise<string | undefined> {
   if (state.pendingToolResults.size === 0) {
     return;
   }
@@ -2380,6 +2409,7 @@ export async function finalizePendingToolResultRow(
   }
   state.pendingToolResults.clear();
   state.pendingToolResultRowReservation = undefined;
+  return rowId;
 }
 
 export async function handleToolResult(
@@ -2487,6 +2517,7 @@ export async function handleToolResult(
       conversationId: deps.ctx.conversationId,
       messageId: state.lastAssistantMessageId,
       toolUseId: event.toolUseId,
+      modeSession: deps.ctx.modeSessions.getTurnOwner(deps.reqId),
     });
     // Capture the seq synchronously (before the persist await) so it reflects
     // the just-stamped tool_result event, then persist on arrival. A failure
@@ -2680,6 +2711,7 @@ export async function handleToolResult(
     answeredQuestion: event.answeredQuestion,
     errorCode: event.errorCode,
     completedAt,
+    modeSession: deps.ctx.modeSessions.getTurnOwner(deps.reqId),
   });
 
   // Capture the seq synchronously (before the persist await) so it reflects the
@@ -3167,12 +3199,28 @@ export async function handleMessageComplete(
   // row as it arrived (`persistPendingToolResultRow`); this rewrites it to the
   // full batch (covering the case where a mid-arrival write failed), indexes it
   // for memory recall, and clears the batch state.
-  await finalizePendingToolResultRow(
+  const toolResultRowId = await finalizePendingToolResultRow(
     state,
     deps.ctx.conversationId,
     buildToolResultMetadata(deps),
     deps.rlog,
   );
+  if (toolResultRowId) {
+    bestEffortModeSessionTracking("tool result finalization", () => {
+      const toolResultRow = getMessageById(
+        toolResultRowId,
+        deps.ctx.conversationId,
+      );
+      if (toolResultRow) {
+        deps.ctx.modeSessions.trackPersistedRow(
+          deps.reqId,
+          toolResultRowId,
+          toolResultRow.createdAt,
+          { startsDisplayBoundary: false },
+        );
+      }
+    });
+  }
 
   // Accumulate directives + warnings from the assistant content for
   // downstream attachment processing. `cleanAssistantContent` is also
@@ -3526,7 +3574,7 @@ function handleProviderError(
   try {
     recordRequestLog(
       deps.ctx.conversationId,
-      JSON.stringify(event.rawRequest),
+      JSON.stringify(event.rawRequest ?? null),
       JSON.stringify(buildProviderErrorResponsePayload(event.error)),
       undefined,
       event.actualProvider,