@vellumai/assistant 0.12.3-staging.1 → 0.12.3

package/src/conversations/__tests__/message-consolidation.test.ts

@@ -3,12 +3,25 @@ import { describe, expect, test } from "bun:test";
 import type { MessageRow } from "../../persistence/conversation-crud.js";
 import { resolveMessageContentBlocks } from "../../persistence/message-content-file.js";
 import {
+  consolidateMessageRows,
   findDisplayTurnEndIndex,
   isToolResultOnlyUserMessage,
   mergeConsecutiveAssistantMessages,
   mergeToolResultsIntoAssistantMessages,
 } from "../message-consolidation.js";
 
+const sessionOwner = { mode: "browser", id: "session-123" } as const;
+
+function ownedMetadata(
+  owner: { mode: string; id: string } = sessionOwner,
+  sentAt?: number,
+): string {
+  return JSON.stringify({
+    modeSession: owner,
+    ...(sentAt !== undefined ? { sentAt } : {}),
+  });
+}
+
 function makeMsg(
   role: string,
   content: string,
@@ -280,6 +293,52 @@ describe("mergeToolResultsIntoAssistantMessages", () => {
     expect(userBlocks.map((b) => b.type)).toEqual(["text"]);
   });
 
+  test.each([null, ownedMetadata()])(
+    "preserves matching results across adjacent mixed user rows (metadata=%s)",
+    (metadata) => {
+      const calls = [
+        { type: "tool_use", id: "tool-first", name: "test", input: {} },
+        { type: "tool_use", id: "tool-second", name: "test", input: {} },
+      ] as const;
+      const firstResult = {
+        type: "tool_result",
+        tool_use_id: "tool-first",
+        content: "first",
+      } as const;
+      const secondResult = {
+        type: "tool_result",
+        tool_use_id: "tool-second",
+        content: "second",
+      } as const;
+      const userText = { type: "text", text: "Additional context" } as const;
+      const rows = [
+        makeMsg("assistant", JSON.stringify(calls), {
+          id: "assistant-123",
+          metadata,
+        }),
+        makeMsg("user", JSON.stringify([firstResult, userText]), {
+          id: "user-mixed",
+          metadata,
+        }),
+        makeMsg("user", JSON.stringify([secondResult]), {
+          id: "user-result",
+          metadata,
+        }),
+      ];
+      const result = consolidateMessageRows(rows);
+      expect(
+        result.messages.map(({ id, content }) => ({ id, content })),
+      ).toEqual([
+        { id: "assistant-123", content: [...calls, firstResult, secondResult] },
+        { id: "user-mixed", content: [userText] },
+      ]);
+      expect(result.mergedIdMap.get("assistant-123")).toEqual([
+        "user-mixed",
+        "user-result",
+      ]);
+    },
+  );
+
   test("passes plain user text through unchanged", () => {
     const messages = [makeMsg("user", "hi"), makeMsg("assistant", "hello")];
     const merged = mergeToolResultsIntoAssistantMessages(messages);
@@ -574,3 +633,93 @@ describe("provider-error boundaries", () => {
     expect(findDisplayTurnEndIndex(rows, 0)).toBe(0);
   });
 });
+
+describe("mode-session consolidation boundaries", () => {
+  test("does not merge marked and unmarked rows or different owners", () => {
+    const rows = [
+      makeMsg("assistant", "first", {
+        id: "owned",
+        metadata: ownedMetadata(),
+      }),
+      makeMsg("assistant", "unstamped", { id: "unstamped" }),
+      makeMsg("assistant", "different", {
+        id: "different",
+        metadata: ownedMetadata({ mode: "browser", id: "session-456" }),
+      }),
+    ];
+    expect(
+      mergeConsecutiveAssistantMessages(rows).messages.map((row) => row.id),
+    ).toEqual(["owned", "unstamped", "different"]);
+    expect(findDisplayTurnEndIndex(rows, 0)).toBe(0);
+  });
+
+  test("suppresses differently owned tool results without merging across their boundary", () => {
+    const assistant = makeMsg(
+      "assistant",
+      JSON.stringify([{ type: "tool_use", id: "tool-123", name: "test" }]),
+      { id: "assistant-123", metadata: ownedMetadata() },
+    );
+    const toolResult = makeMsg(
+      "user",
+      JSON.stringify([
+        { type: "tool_result", tool_use_id: "tool-123", content: "done" },
+      ]),
+      {
+        id: "result-123",
+        metadata: ownedMetadata({ mode: "browser", id: "session-456" }),
+      },
+    );
+    const laterToolResult = makeMsg(
+      "user",
+      JSON.stringify([
+        { type: "tool_result", tool_use_id: "tool-123", content: "later" },
+      ]),
+      { id: "result-456", metadata: ownedMetadata() },
+    );
+    const merged = mergeToolResultsIntoAssistantMessages([
+      assistant,
+      toolResult,
+      laterToolResult,
+    ]);
+    expect(merged.map((row) => row.id)).toEqual(["assistant-123"]);
+    expect(merged[0]?.content).toEqual(assistant.content);
+    expect(findDisplayTurnEndIndex([assistant, toolResult], 0)).toBe(0);
+  });
+
+  test("preserves donor aliases and earlier and later activity bounds", () => {
+    const rows = [
+      makeMsg("assistant", "anchor", {
+        id: "assistant-anchor",
+        createdAt: 200,
+        metadata: ownedMetadata(sessionOwner, 200),
+      }),
+      makeMsg(
+        "user",
+        JSON.stringify([
+          { type: "tool_result", tool_use_id: "tool-123", content: "done" },
+        ]),
+        {
+          id: "tool-result-earlier",
+          createdAt: 150,
+          metadata: ownedMetadata(sessionOwner, 150),
+        },
+      ),
+      makeMsg("assistant", "tail", {
+        id: "assistant-tail",
+        createdAt: 250,
+        metadata: ownedMetadata(sessionOwner, 300),
+      }),
+    ];
+    const consolidated = consolidateMessageRows(rows);
+    expect(consolidated.messages.map((row) => row.id)).toEqual([
+      "assistant-anchor",
+    ]);
+    expect(consolidated.mergedIdMap.get("assistant-anchor")).toEqual([
+      "tool-result-earlier",
+      "assistant-tail",
+    ]);
+    expect(consolidated.modeSessionActivityMap.get("assistant-anchor")).toEqual(
+      { firstAt: 150, lastAt: 300 },
+    );
+  });
+});

package/src/conversations/message-consolidation.ts

@@ -31,8 +31,13 @@
  * agree on cluster boundaries without duplicating the merge code.
  */
 
+import type { ModeSessionActivity } from "../api/mode-session.js";
 import type { MessageRow } from "../persistence/conversation-crud.js";
 import { isStandaloneAssistantMessage } from "../persistence/conversation-crud.js";
+import {
+  readMessageSentAt,
+  readModeSessionMetadata,
+} from "../persistence/message-metadata.js";
 import { assistantTextVisibilityOf } from "../persistence/user-facing-content.js";
 import type { ContentBlock } from "../providers/types.js";
 import { getLogger } from "../util/logger.js";
@@ -69,6 +74,15 @@ function sameAssistantTextVisibility(a: MessageRow, b: MessageRow): boolean {
   );
 }
 
+function sameModeSessionOwnership(a: MessageRow, b: MessageRow): boolean {
+  const aOwner = readModeSessionMetadata(a.metadata);
+  const bOwner = readModeSessionMetadata(b.metadata);
+  if (!aOwner || !bOwner) {
+    return aOwner === undefined && bOwner === undefined;
+  }
+  return aOwner.id === bOwner.id && aOwner.mode === bOwner.mode;
+}
+
 // ── Block predicates ────────────────────────────────────────────────
 
 function isToolResultType(type: string): boolean {
@@ -161,12 +175,17 @@ export function findDisplayTurnEndIndex(
     if (
       next.role === "assistant" &&
       !isStandaloneAssistantRow(next) &&
-      sameAssistantTextVisibility(messages[startIdx]!, next)
+      sameAssistantTextVisibility(messages[startIdx]!, next) &&
+      sameModeSessionOwnership(messages[startIdx]!, next)
     ) {
       endIdx += 1;
       continue;
     }
-    if (next.role === "user" && isToolResultOnlyUserMessage(next)) {
+    if (
+      next.role === "user" &&
+      isToolResultOnlyUserMessage(next) &&
+      sameModeSessionOwnership(messages[startIdx]!, next)
+    ) {
       endIdx += 1;
       continue;
     }
@@ -188,6 +207,7 @@ export function findDisplayTurnEndIndex(
  */
 export function mergeToolResultsIntoAssistantMessages(
   messages: MessageRow[],
+  mergedIdsByAssistant?: Map<string, string[]>,
 ): MessageRow[] {
   // Index of the most recent assistant message in the output array.
   let lastAssistantIdx = -1;
@@ -232,14 +252,20 @@ export function mergeToolResultsIntoAssistantMessages(
     // Append tool_result blocks to the preceding assistant message's content.
     // No-op at pagination boundaries (lastAssistantIdx < 0); orphan tool_results
     // are silently dropped by renderHistoryContent downstream either way.
-    if (lastAssistantIdx >= 0) {
-      const assistant = result[lastAssistantIdx];
+    const assistant =
+      lastAssistantIdx >= 0 ? result[lastAssistantIdx] : undefined;
+    const canMerge =
+      assistant !== undefined && sameModeSessionOwnership(assistant, msg);
+    if (canMerge) {
       let assistantContent = parsedAssistantContent.get(lastAssistantIdx);
       if (!assistantContent) {
         assistantContent = [...assistant.content];
         parsedAssistantContent.set(lastAssistantIdx, assistantContent);
       }
       assistantContent.push(...toolResultBlocks);
+      const mergedIds = mergedIdsByAssistant?.get(assistant.id) ?? [];
+      mergedIds.push(msg.id);
+      mergedIdsByAssistant?.set(assistant.id, mergedIds);
     }
 
     // If the user message had only tool_result (+ system_notice) blocks,
@@ -250,6 +276,9 @@ export function mergeToolResultsIntoAssistantMessages(
     if (realUserContent.length > 0) {
       result.push({ ...msg, content: otherBlocks });
     }
+    if (!canMerge) {
+      lastAssistantIdx = -1;
+    }
     // else: tool-result-only → suppressed
   }
 
@@ -337,7 +366,8 @@ export function mergeConsecutiveAssistantMessages(messages: MessageRow[]): {
       lastIdx >= 0 &&
       result[lastIdx].role === "assistant" &&
       !isStandaloneAssistantRow(result[lastIdx]) &&
-      sameAssistantTextVisibility(result[lastIdx], msg);
+      sameAssistantTextVisibility(result[lastIdx], msg) &&
+      sameModeSessionOwnership(result[lastIdx], msg);
 
     if (!isConsecutiveAssistant) {
       result.push(msg);
@@ -376,3 +406,70 @@ export function mergeConsecutiveAssistantMessages(messages: MessageRow[]): {
 
   return { messages: result, mergedIdMap };
 }
+
+export function consolidateMessageRows(messages: MessageRow[]): {
+  messages: MessageRow[];
+  mergedIdMap: Map<string, string[]>;
+  modeSessionActivityMap: Map<string, ModeSessionActivity>;
+} {
+  const rowsById = new Map(messages.map((message) => [message.id, message]));
+  const toolMergedIds = new Map<string, string[]>();
+  const toolMerged = mergeToolResultsIntoAssistantMessages(
+    messages,
+    toolMergedIds,
+  );
+  const consolidated = mergeConsecutiveAssistantMessages(toolMerged);
+  const mergedIdMap = new Map<string, string[]>();
+  const modeSessionActivityMap = new Map<string, ModeSessionActivity>();
+
+  for (const message of consolidated.messages) {
+    const assistantDonors = consolidated.mergedIdMap.get(message.id) ?? [];
+    const aliases = [
+      ...(toolMergedIds.get(message.id) ?? []),
+      ...assistantDonors.flatMap((id) => [
+        id,
+        ...(toolMergedIds.get(id) ?? []),
+      ]),
+    ];
+    if (aliases.length > 0) {
+      mergedIdMap.set(message.id, aliases);
+    }
+
+    const owner = readModeSessionMetadata(message.metadata);
+    if (!owner) {
+      continue;
+    }
+    const representedRows = [message.id, ...aliases]
+      .map((id) => rowsById.get(id))
+      .filter((row): row is MessageRow => row !== undefined);
+    if (
+      representedRows.some((row) => {
+        const rowOwner = readModeSessionMetadata(row.metadata);
+        return rowOwner?.id !== owner.id || rowOwner.mode !== owner.mode;
+      })
+    ) {
+      continue;
+    }
+    const activityTimes = representedRows
+      .map((row) => readMessageSentAt(row.metadata) ?? row.createdAt)
+      .filter(
+        (at) =>
+          Number.isInteger(at) &&
+          at >= 0 &&
+          Number.isFinite(new Date(at).getTime()),
+      );
+    if (activityTimes.length !== representedRows.length) {
+      continue;
+    }
+    modeSessionActivityMap.set(message.id, {
+      firstAt: Math.min(...activityTimes),
+      lastAt: Math.max(...activityTimes),
+    });
+  }
+
+  return {
+    messages: consolidated.messages,
+    mergedIdMap,
+    modeSessionActivityMap,
+  };
+}

package/src/credential-execution/ces-connect.test.ts

@@ -0,0 +1,69 @@
+/**
+ * Tests for the shared CES RPC session helper used by assistant boot and
+ * child-process credential resolution.
+ */
+import { describe, expect, test } from "bun:test";
+
+import { openCesRpcSession, reconnectCesRpcSession } from "./ces-connect.js";
+import { createCesProcessManager } from "./process-manager.js";
+
+describe("openCesRpcSession", () => {
+  test("returns undefined when discovery fails without polling", async () => {
+    const start = Date.now();
+    const pm = createCesProcessManager({
+      discover: async () => ({
+        mode: "unavailable",
+        reason: "missing test socket",
+      }),
+    });
+
+    const session = await openCesRpcSession({ processManager: pm });
+
+    expect(session).toBeUndefined();
+    expect(pm.isRunning()).toBe(false);
+    expect(Date.now() - start).toBeLessThan(1_000);
+  });
+
+  test("returns undefined when the abort signal is already fired", async () => {
+    const abort = new AbortController();
+    abort.abort();
+    const pm = createCesProcessManager({
+      discover: async () => {
+        throw new Error("discover should not run after abort");
+      },
+    });
+
+    const session = await openCesRpcSession({
+      processManager: pm,
+      signal: abort.signal,
+    });
+
+    expect(session).toBeUndefined();
+    expect(pm.isRunning()).toBe(false);
+  });
+
+  test("creates a process manager when none is provided", async () => {
+    const session = await openCesRpcSession({
+      discover: async () => ({
+        mode: "unavailable",
+        reason: "missing test socket",
+      }),
+    });
+
+    expect(session).toBeUndefined();
+  });
+
+  test("reconnectCesRpcSession stops the manager and returns undefined when discovery fails", async () => {
+    const pm = createCesProcessManager({
+      discover: async () => ({
+        mode: "unavailable",
+        reason: "missing test socket",
+      }),
+    });
+
+    const client = await reconnectCesRpcSession(pm);
+
+    expect(client).toBeUndefined();
+    expect(pm.isRunning()).toBe(false);
+  });
+});

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

@@ -0,0 +1,115 @@
+/**
+ * Open a CES RPC client over the shared bootstrap socket.
+ *
+ * Assistant boot and child-process credential reads use this helper. Both
+ * are CES API clients: discover `ces.sock`, connect, handshake, reconnect.
+ *
+ * They stay two entry points because boot must load handshake identity
+ * without talking to CES, then hand that identity to CES. Children open a
+ * session only when this process has not already claimed one (no live
+ * client and no reconnect owner).
+ */
+
+import type { AssistantConfig } from "../config/schema.js";
+import { getLogger } from "../util/logger.js";
+import {
+  type CesClient,
+  type CesClientHandshakeOptions,
+  createCesClient,
+} from "./client.js";
+import {
+  type CesProcessManager,
+  type CesProcessManagerConfig,
+  CesUnavailableError,
+  createCesProcessManager,
+} from "./process-manager.js";
+
+const log = getLogger("ces-connect");
+
+export interface CesRpcSession {
+  client: CesClient;
+  processManager: CesProcessManager;
+}
+
+export interface OpenCesRpcSessionOptions {
+  /**
+   * Reuse an existing process manager (reconnect after `stop()`). When
+   * omitted, a new manager is created.
+   */
+  processManager?: CesProcessManager;
+  handshake?: CesClientHandshakeOptions;
+  signal?: AbortSignal;
+  assistantConfig?: AssistantConfig;
+  discover?: CesProcessManagerConfig["discover"];
+}
+
+/**
+ * Discover the CES socket, connect, and complete the RPC handshake.
+ *
+ * Returns undefined when CES is missing, the handshake is rejected, the
+ * abort signal fires, or the transport fails. The process manager is
+ * stopped on those paths so a later reconnect can call `start()` again.
+ */
+export async function openCesRpcSession(
+  options: OpenCesRpcSessionOptions = {},
+): Promise<CesRpcSession | undefined> {
+  const pm =
+    options.processManager ??
+    createCesProcessManager({
+      assistantConfig: options.assistantConfig,
+      discover: options.discover,
+    });
+
+  const fail = async (): Promise<undefined> => {
+    await pm.stop().catch(() => {});
+    return undefined;
+  };
+
+  if (options.signal?.aborted) {
+    return fail();
+  }
+
+  try {
+    const transport = await pm.start();
+    if (options.signal?.aborted) {
+      return fail();
+    }
+
+    const client = createCesClient(transport);
+    const { accepted, reason } = await client.handshake(options.handshake);
+    if (options.signal?.aborted) {
+      client.close();
+      return fail();
+    }
+    if (!accepted) {
+      log.warn({ reason }, "CES handshake rejected");
+      client.close();
+      return fail();
+    }
+
+    return { client, processManager: pm };
+  } catch (err) {
+    if (err instanceof CesUnavailableError) {
+      log.info({ reason: err.message }, "CES is not available");
+    } else {
+      log.warn(
+        { error: err instanceof Error ? err.message : String(err) },
+        "Failed to open CES RPC session",
+      );
+    }
+    return fail();
+  }
+}
+
+/**
+ * Stop the current transport and open a new session on the same process
+ * manager. Used by assistant boot and child-process reconnect callbacks.
+ */
+export async function reconnectCesRpcSession(
+  processManager: CesProcessManager,
+  handshake?: CesClientHandshakeOptions,
+): Promise<CesClient | undefined> {
+  await processManager.stop();
+  const session = await openCesRpcSession({ processManager, handshake });
+  return session?.client;
+}