@vellumai/assistant 0.8.4 → 0.8.6

package/src/conversations/message-consolidation.ts

@@ -0,0 +1,404 @@
+/**
+ * Message consolidation — shared logic that maps the raw DB row sequence
+ * onto the displayed transcript.
+ *
+ * A single "display turn" in the UI can span multiple DB rows. During
+ * streaming the agent loop persists each API call as a separate
+ * assistant row (deferred consolidation, for prefix-cache stability),
+ * and intervening user rows often carry only tool_result blocks that
+ * are suppressed at display time. The two-pass collapse here is what
+ * turns the raw sequence into the rows the user actually sees.
+ *
+ *   Pass 1 — `mergeToolResultsIntoAssistantMessages`
+ *     For each user row, lift its `tool_result` (and `web_search_tool_result`)
+ *     blocks onto the preceding assistant row. If the user row had
+ *     nothing else (only tool_result + system_notice blocks), drop the
+ *     row entirely. Otherwise keep its non-tool-result content as a
+ *     real user message.
+ *
+ *   Pass 2 — `mergeConsecutiveAssistantMessages`
+ *     After pass 1 removes tool-result-only user separators, fold runs
+ *     of adjacent assistant rows onto the first row of the run
+ *     ("anchor"). Anchors keep their id, createdAt, and metadata;
+ *     subagent-notification metadata is promoted from later rows when
+ *     the anchor lacks it.
+ *
+ * `findDisplayTurnEndIndex` is the primitive any write-path needs when
+ * it has a single DB-row id and wants to know "which DB rows make up
+ * the display turn that this row anchors". It uses the same
+ * `isToolResultOnlyUserMessage` predicate that pass 1 uses to decide
+ * which user rows are suppressed, so the read-path and write-path
+ * agree on cluster boundaries without duplicating the merge code.
+ */
+
+import type { MessageRow } from "../memory/conversation-crud.js";
+import { getLogger } from "../util/logger.js";
+
+const log = getLogger("message-consolidation");
+
+// ── Block predicates ────────────────────────────────────────────────
+
+function isToolResultType(type: string): boolean {
+  return type === "tool_result" || type === "web_search_tool_result";
+}
+
+function isSystemNoticeText(block: Record<string, unknown>): boolean {
+  if (block.type !== "text") return false;
+  const text = typeof block.text === "string" ? block.text : "";
+  return (
+    text.startsWith("<system_notice>") && text.endsWith("</system_notice>")
+  );
+}
+
+/**
+ * True when a user row's content is exclusively tool_result blocks
+ * (and optional system_notice text). Such rows are suppressed by the
+ * read-path collapse — they exist in the DB to satisfy the LLM's
+ * tool_use ↔ tool_result pairing requirement but are never displayed
+ * to the user. Any write-path that walks DB rows in display order
+ * must treat them as part of the surrounding assistant turn.
+ */
+export function isToolResultOnlyUserMessage(msg: MessageRow): boolean {
+  if (msg.role !== "user") return false;
+  let blocks: unknown[];
+  try {
+    const parsed = JSON.parse(msg.content);
+    if (!Array.isArray(parsed)) return false;
+    blocks = parsed;
+  } catch {
+    return false;
+  }
+
+  let sawToolResult = false;
+  for (const block of blocks) {
+    if (
+      typeof block !== "object" ||
+      block === null ||
+      typeof (block as Record<string, unknown>).type !== "string"
+    ) {
+      return false;
+    }
+    const rec = block as Record<string, unknown>;
+    if (isToolResultType(rec.type as string)) {
+      sawToolResult = true;
+      continue;
+    }
+    if (isSystemNoticeText(rec)) {
+      continue;
+    }
+    return false;
+  }
+  return sawToolResult;
+}
+
+// ── Display-turn boundary lookup ────────────────────────────────────
+
+/**
+ * Returns the inclusive DB-row index where the display turn that
+ * contains `messages[startIdx]` ends.
+ *
+ * For non-assistant rows, returns `startIdx` unchanged — only assistant
+ * turns can span multiple rows.
+ *
+ * For assistant rows, advances past any consecutive rows that the
+ * read-path collapse would fold into the same display turn:
+ *   - another assistant row → part of the consecutive-assistant run, OR
+ *   - a tool-result-only user row → suppressed at display time, sits
+ *     between two halves of the same assistant turn.
+ *
+ * Stops at the first real user message (or any non-collapsible row),
+ * end-of-array, or invalid index.
+ *
+ * Mirror of the cluster boundary that `mergeConsecutiveAssistantMessages`
+ * (post pass-1) would produce — without paying the cost of parsing
+ * every content block.
+ */
+export function findDisplayTurnEndIndex(
+  messages: MessageRow[],
+  startIdx: number,
+): number {
+  if (startIdx < 0 || startIdx >= messages.length) return startIdx;
+  if (messages[startIdx]?.role !== "assistant") return startIdx;
+
+  let endIdx = startIdx;
+  while (endIdx + 1 < messages.length) {
+    const next = messages[endIdx + 1];
+    if (!next) break;
+    if (next.role === "assistant") {
+      endIdx += 1;
+      continue;
+    }
+    if (next.role === "user" && isToolResultOnlyUserMessage(next)) {
+      endIdx += 1;
+      continue;
+    }
+    break;
+  }
+  return endIdx;
+}
+
+// ── Pass 1: tool-result merging ─────────────────────────────────────
+
+/**
+ * Merge tool_result blocks from user messages into the preceding assistant
+ * message's content array. This lets renderHistoryContent's pendingToolUses
+ * map pair tool_use and tool_result blocks, preventing "unknown" tool names.
+ *
+ * User messages that consist entirely of tool_result blocks (and optional
+ * system_notice text) are removed from the output. Mixed messages (tool_result
+ * + real user text) keep only the non-tool-result blocks.
+ */
+export function mergeToolResultsIntoAssistantMessages(
+  messages: MessageRow[],
+): MessageRow[] {
+  // Index of the most recent assistant message in the output array.
+  let lastAssistantIdx = -1;
+  // Parsed content caches — lazily populated per assistant message.
+  const parsedAssistantContent = new Map<number, unknown[]>();
+
+  const result: MessageRow[] = [];
+
+  for (const msg of messages) {
+    if (msg.role === "assistant") {
+      lastAssistantIdx = result.length;
+      result.push(msg);
+      continue;
+    }
+
+    // Only process user messages — other roles pass through.
+    if (msg.role !== "user") {
+      result.push(msg);
+      continue;
+    }
+
+    let blocks: unknown[];
+    try {
+      const parsed = JSON.parse(msg.content);
+      if (!Array.isArray(parsed)) {
+        result.push(msg);
+        continue;
+      }
+      blocks = parsed;
+    } catch {
+      result.push(msg);
+      continue;
+    }
+
+    // Separate tool-result blocks from real user content.
+    const toolResultBlocks: unknown[] = [];
+    const otherBlocks: unknown[] = [];
+    for (const block of blocks) {
+      if (
+        typeof block === "object" &&
+        block !== null &&
+        typeof (block as Record<string, unknown>).type === "string"
+      ) {
+        const rec = block as Record<string, unknown>;
+        if (isToolResultType(rec.type as string)) {
+          toolResultBlocks.push(block);
+        } else if (isSystemNoticeText(rec)) {
+          // System notices don't count as user content — drop them when
+          // the message is otherwise tool-result-only.
+          otherBlocks.push(block);
+        } else {
+          otherBlocks.push(block);
+        }
+      } else {
+        otherBlocks.push(block);
+      }
+    }
+
+    // No tool results → pass through unchanged. System notices are only
+    // injected alongside tool results in the agent loop, so a pure user
+    // message (no tool_result blocks) should never be filtered — even if
+    // the user's text happens to look like a system_notice tag.
+    if (toolResultBlocks.length === 0) {
+      result.push(msg);
+      continue;
+    }
+
+    // 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];
+      let assistantContent = parsedAssistantContent.get(lastAssistantIdx);
+      if (!assistantContent) {
+        try {
+          const parsed = JSON.parse(assistant.content);
+          assistantContent = Array.isArray(parsed) ? parsed : [parsed];
+        } catch {
+          assistantContent = [];
+        }
+        parsedAssistantContent.set(lastAssistantIdx, assistantContent);
+      }
+      assistantContent.push(...toolResultBlocks);
+    }
+
+    // If the user message had only tool_result (+ system_notice) blocks,
+    // suppress it entirely. Otherwise keep the non-tool-result content.
+    const realUserContent = otherBlocks.filter(
+      (b) =>
+        !(
+          typeof b === "object" &&
+          b !== null &&
+          isSystemNoticeText(b as Record<string, unknown>)
+        ),
+    );
+    if (realUserContent.length > 0) {
+      result.push({ ...msg, content: JSON.stringify(otherBlocks) });
+    }
+    // else: tool-result-only → suppressed
+  }
+
+  // Write back any modified assistant message content.
+  for (const [idx, content] of parsedAssistantContent) {
+    result[idx] = { ...result[idx], content: JSON.stringify(content) };
+  }
+
+  return result;
+}
+
+// ── Pass 2: consecutive-assistant merging ───────────────────────────
+
+/** Parse a message's JSON content into an array of content blocks. */
+function parseContentBlocks(content: string): unknown[] {
+  try {
+    const parsed = JSON.parse(content);
+    return Array.isArray(parsed) ? parsed : [parsed];
+  } catch (err) {
+    log.warn(
+      { err },
+      "Failed to parse content blocks during assistant message merge",
+    );
+    return [];
+  }
+}
+
+/**
+ * Append content blocks from a donor message onto a target block array.
+ * Parses the donor's JSON content and pushes each block into `target`.
+ */
+function appendContentBlocks(target: unknown[], donorContent: string): void {
+  try {
+    const parsed = JSON.parse(donorContent);
+    if (Array.isArray(parsed)) {
+      target.push(...parsed);
+    } else {
+      target.push(parsed);
+    }
+  } catch (err) {
+    log.warn(
+      { err },
+      "Failed to parse donor content blocks during assistant message merge",
+    );
+  }
+}
+
+/**
+ * Promote metadata fields from a donor message to the surviving message
+ * when the survivor lacks them. Currently promotes `subagentNotification`.
+ * Returns a new MessageRow if promotion occurred, otherwise the original.
+ */
+function promoteMetadata(survivor: MessageRow, donor: MessageRow): MessageRow {
+  if (donor.metadata && survivor.metadata) {
+    try {
+      const survivorMeta = JSON.parse(survivor.metadata);
+      const donorMeta = JSON.parse(donor.metadata);
+      if (
+        !survivorMeta.subagentNotification &&
+        donorMeta.subagentNotification
+      ) {
+        survivorMeta.subagentNotification = donorMeta.subagentNotification;
+        return { ...survivor, metadata: JSON.stringify(survivorMeta) };
+      }
+    } catch (err) {
+      log.warn(
+        { err },
+        "Failed to parse metadata during assistant message merge",
+      );
+    }
+  } else if (donor.metadata && !survivor.metadata) {
+    return { ...survivor, metadata: donor.metadata };
+  }
+  return survivor;
+}
+
+/**
+ * Merge consecutive assistant messages into a single message at query time.
+ *
+ * During streaming, all assistant turns within one agent loop accumulate on
+ * a single client-side message. In the DB, each API turn is stored as a
+ * separate assistant row (consolidation is deferred to compaction for
+ * prefix-cache stability). This produces N separate assistant messages that
+ * the client would otherwise render as N individual rows — each showing
+ * "Completed 1 step" instead of one grouped "Completed N steps" accordion.
+ *
+ * This function concatenates the content block arrays of consecutive
+ * assistant messages (no intervening user messages after tool-result
+ * merging) into the first message of each run. The merged messages are
+ * removed from the output. This is query-time only — the DB is not
+ * modified.
+ *
+ * The first message in each run keeps its id, createdAt, and metadata so
+ * that attachment lookups, display timestamps, and subagent notifications
+ * continue to work. Metadata from later messages in the run (e.g.
+ * subagentNotification) is preserved by promoting it to the surviving
+ * message when the surviving message has no metadata of its own for that
+ * field.
+ */
+export function mergeConsecutiveAssistantMessages(messages: MessageRow[]): {
+  messages: MessageRow[];
+  /** Maps each surviving message ID → all original message IDs merged into it. */
+  mergedIdMap: Map<string, string[]>;
+} {
+  const result: MessageRow[] = [];
+  // Key = index in `result`, value = accumulated content blocks.
+  const pendingMerges = new Map<number, unknown[]>();
+  // Key = index in `result`, value = IDs of messages merged into the target.
+  const mergedIds = new Map<number, string[]>();
+
+  for (const msg of messages) {
+    const lastIdx = result.length - 1;
+    const isConsecutiveAssistant =
+      msg.role === "assistant" &&
+      lastIdx >= 0 &&
+      result[lastIdx].role === "assistant";
+
+    if (!isConsecutiveAssistant) {
+      result.push(msg);
+      continue;
+    }
+
+    // Track the donor message ID.
+    let ids = mergedIds.get(lastIdx);
+    if (!ids) {
+      ids = [];
+      mergedIds.set(lastIdx, ids);
+    }
+    ids.push(msg.id);
+
+    // Lazily parse the target's content on first merge.
+    let targetContent = pendingMerges.get(lastIdx);
+    if (!targetContent) {
+      targetContent = parseContentBlocks(result[lastIdx].content);
+      pendingMerges.set(lastIdx, targetContent);
+    }
+
+    appendContentBlocks(targetContent, msg.content);
+    result[lastIdx] = promoteMetadata(result[lastIdx], msg);
+  }
+
+  // Write back merged content for any messages that were targets.
+  for (const [idx, content] of pendingMerges) {
+    result[idx] = { ...result[idx], content: JSON.stringify(content) };
+  }
+
+  // Build the merged ID map keyed by surviving message ID.
+  const mergedIdMap = new Map<string, string[]>();
+  for (const [idx, ids] of mergedIds) {
+    mergedIdMap.set(result[idx].id, ids);
+  }
+
+  return { messages: result, mergedIdMap };
+}

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

@@ -212,3 +212,43 @@ export function discoverCes(): DiscoveryResult {
   }
   return discoverLocalCes();
 }
+
+/** How long to poll for the managed bootstrap socket before giving up. */
+const MANAGED_DISCOVERY_TIMEOUT_MS = 3_000;
+
+/** Delay between managed bootstrap-socket discovery attempts. */
+const MANAGED_DISCOVERY_INTERVAL_MS = 100;
+
+/**
+ * Discover CES, polling for the managed bootstrap socket with a short
+ * backoff before failing.
+ *
+ * The managed CES sidecar re-binds its bootstrap socket after each assistant
+ * session ends — it outlives any single assistant and accepts the next
+ * connection (see `credential-executor/src/managed-main.ts`). A reconnecting
+ * assistant (e.g. after a container restart) can therefore probe during the
+ * brief window before the socket is re-bound. A single existence check would
+ * race that window and incorrectly report the sidecar as unavailable; polling
+ * absorbs the gap.
+ *
+ * Local discovery is returned immediately — a missing binary will not appear
+ * by waiting, so there is nothing to poll for.
+ */
+export async function discoverCesWithRetry({
+  timeoutMs = MANAGED_DISCOVERY_TIMEOUT_MS,
+  intervalMs = MANAGED_DISCOVERY_INTERVAL_MS,
+}: { timeoutMs?: number; intervalMs?: number } = {}): Promise<DiscoveryResult> {
+  // Only the managed bootstrap socket is worth polling for; a missing local
+  // binary will not appear by waiting.
+  if (!getIsContainerized()) {
+    return discoverCes();
+  }
+
+  const deadline = Date.now() + timeoutMs;
+  let result = discoverCes();
+  while (result.mode === "unavailable" && Date.now() < deadline) {
+    await new Promise((resolve) => setTimeout(resolve, intervalMs));
+    result = discoverCes();
+  }
+  return result;
+}

package/src/credential-execution/process-manager.ts

@@ -29,7 +29,7 @@ import { ensureBun } from "../util/bun-runtime.js";
 import { getLogger } from "../util/logger.js";
 import type { CesTransport } from "./client.js";
 import {
-  discoverCes,
+  discoverCesWithRetry,
   discoverLocalCes,
   type DiscoveryResult,
   type LocalDiscoverySuccess,
@@ -118,7 +118,11 @@ export function createCesProcessManager(
         throw new Error("CES process manager is already running");
       }
 
-      discoveryResult = await discoverCes();
+      // Poll for the managed bootstrap socket with a short backoff. The CES
+      // sidecar re-binds its socket after each assistant session ends, so a
+      // reconnecting assistant can briefly race the re-bind; a single probe
+      // would otherwise fall back to local discovery and fail in managed mode.
+      discoveryResult = await discoverCesWithRetry();
       if (discoveryResult.mode === "unavailable") {
         // The managed sidecar bootstrap socket is not present — this happens
         // when the instance pre-dates the socket volume mount (e.g. existing