@vellumai/assistant 0.10.3 → 0.10.4-dev.202607010028.1a4efcc

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

@@ -10,6 +10,7 @@ import type pino from "pino";
 import { v4 as uuid } from "uuid";
 
 import type { AgentEvent } from "../agent/loop.js";
+import { getThreadTs } from "../channels/slack-thread-store.js";
 import type {
   TurnChannelContext,
   TurnInterfaceContext,
@@ -18,49 +19,46 @@ import { getConfig } from "../config/loader.js";
 import { recordEstimate } from "../context/estimator-calibration.js";
 import { stripInjectionsForCompaction } from "../context/strip-injections.js";
 import { getCalibrationProviderKey } from "../context/token-estimator.js";
+import {
+  formatSlackTimezoneLabel,
+  type SlackMessageMetadata,
+  writeSlackMetadata,
+} from "../messaging/providers/slack/message-metadata.js";
 import {
   recordCompactionEndBestEffort,
   recordCompactionStartBestEffort,
-} from "../memory/compaction-log-store-clickhouse.js";
-import { projectAssistantMessage } from "../memory/conversation-attention-store.js";
+} from "../persistence/compaction-log-store-clickhouse.js";
+import { projectAssistantMessage } from "../persistence/conversation-attention-store.js";
 import {
   deleteMessageById,
   getConversation,
   getMessageById,
   messageMetadataSchema,
   provenanceFromTrustContext,
+  recordConversationPersistedSeq,
   reserveMessage,
   setConversationHistoryStrippedAt,
   setLastNotifiedInferenceProfile,
   updateMessageContent,
-} from "../memory/conversation-crud.js";
-import { syncMessageToDisk } from "../memory/conversation-disk-view.js";
-import { indexMessageNow } from "../memory/indexer.js";
+} from "../persistence/conversation-crud.js";
+import { syncMessageToDisk } from "../persistence/conversation-disk-view.js";
 import {
   backfillMessageIdOnLogs,
   buildProviderErrorResponsePayload,
   recordRequestLog,
   setAgentLoopExitReasonOnLatestLog,
-} from "../memory/llm-request-log-store.js";
-import { backfillMemoryRecallLogMessageId } from "../memory/memory-recall-log-store.js";
-import { backfillMemoryV2ActivationMessageId } from "../memory/memory-v2-activation-log-store.js";
-import { getThreadTs } from "../memory/slack-thread-store.js";
-import {
-  formatSlackTimezoneLabel,
-  type SlackMessageMetadata,
-  writeSlackMetadata,
-} from "../messaging/providers/slack/message-metadata.js";
+} from "../persistence/llm-request-log-store.js";
 import type { ContextWindowResult } from "../plugins/defaults/compaction/window-manager.js";
-import { backfillMemoryV3SelectionMessageId } from "../plugins/defaults/memory-v3-shadow/shadow-plugin.js";
+import { indexMessageNow } from "../plugins/defaults/memory/indexer.js";
+import { backfillMemoryRecallLogMessageId } from "../plugins/defaults/memory/memory-recall-log-store.js";
+import { backfillMemoryV2ActivationMessageId } from "../plugins/defaults/memory/memory-v2-activation-log-store.js";
+import { backfillMemoryV3SelectionMessageId } from "../plugins/defaults/memory/v3/shadow-plugin.js";
 import type {
   ContentBlock,
   ImageContent,
   Message,
 } from "../providers/types.js";
-import {
-  getCurrentSeq,
-  recordPersistedSeq,
-} from "../runtime/assistant-stream-state.js";
+import { getCurrentSeq } from "../runtime/assistant-stream-state.js";
 import { publishSyncInvalidation } from "../runtime/sync/sync-publisher.js";
 import { redactSecrets } from "../security/secret-scanner.js";
 import { extractDomain } from "../tools/network/domain-normalize.js";
@@ -76,6 +74,7 @@ import {
 import { ProviderError } from "../util/errors.js";
 import { faviconUrlForDomain } from "../util/favicon.js";
 import { getLogger } from "../util/logger.js";
+import { withSqliteRetry } from "../util/sqlite-retry.js";
 import type { DirectiveRequest } from "./assistant-attachments.js";
 import {
   cleanAssistantContent,
@@ -101,9 +100,16 @@ import type {
   WebSearchMetadata,
   WebSearchResultItem,
 } from "./message-types/web-activity.js";
+import type { TurnLatencyTracker } from "./turn-latency-tracker.js";
 
 const log = getLogger("agent-loop-handlers");
 
+function shouldPersistProviderErrorAsAssistantMessage(classified: {
+  code: string;
+}): boolean {
+  return classified.code !== "MANAGED_KEY_INVALID";
+}
+
 /**
  * Persist the history-stripped marker after the loop strips runtime injections
  * for compaction / overflow recovery. The marker is a durability hint, not
@@ -163,6 +169,7 @@ export interface EventHandlerState {
   readonly exchangeRawResponses: unknown[];
   model: string;
   providerErrorUserMessage: string | null;
+  persistProviderErrorAsAssistantMessage: boolean;
   lastAssistantMessageId: string | undefined;
   /**
    * True when `handleLlmCallStarted` has reserved an empty assistant row
@@ -209,6 +216,15 @@ export interface EventHandlerState {
     string,
     { startedAt: number; completedAt?: number }
   >;
+  /**
+   * Tracks tool_use_id → the `tool_use_preview_start` timestamp (Unix ms), the
+   * first byte of the tool call. Stamped before execution begins, so it lives
+   * in its own map rather than `toolCallTimestamps` (whose record requires an
+   * execution `startedAt`). Read when emitting `tool_use_start` and when
+   * annotating the persisted block so the user-perceived latency survives a
+   * refresh.
+   */
+  readonly toolPreviewStartedAt: Map<string, number>;
   /** The tool_use_id of the currently executing tool (set in handleToolUse, cleared in handleToolResult). */
   currentToolUseId: string | undefined;
   /** Maps confirmation requestId → tool_use_id for linking decisions to tools. */
@@ -296,6 +312,12 @@ export interface EventHandlerState {
    * consumed (deleted) when the end event dispatches.
    */
   readonly compactionStartMessages: Map<string, Message[]>;
+  /**
+   * Cursor into the turn's latency-mark list marking how far prior calls have
+   * already been serialized, so each `usage` event emits only its own call's
+   * latency segment. Advances on every `handleUsage`.
+   */
+  latencyCursor: number;
 }
 
 /** Immutable context shared across event handlers within a single agent loop run. */
@@ -321,6 +343,15 @@ export interface EventHandlerDeps {
     result: ContextWindowResult,
     messages: Message[],
   ) => Promise<void>;
+  /**
+   * Per-turn first-token latency instrumentation. The orchestrator stamps the
+   * turn-level marks; the agent loop stamps the per-call marks. `handleUsage`
+   * serializes the breakdown for each call and persists it on the request log.
+   * Optional: a pure observability hook that the production orchestrator always
+   * supplies, but test fixtures and any future caller may omit — `handleUsage`
+   * degrades gracefully when it's absent.
+   */
+  readonly latencyTracker?: TurnLatencyTracker;
 }
 
 // ── Factory ──────────────────────────────────────────────────────────
@@ -341,6 +372,7 @@ export function createEventHandlerState(): EventHandlerState {
     exchangeRawResponses: [],
     model: "",
     providerErrorUserMessage: null,
+    persistProviderErrorAsAssistantMessage: false,
     lastAssistantMessageId: undefined,
     assistantRowAwaitingFinalization: false,
     pendingToolResults: new Map(),
@@ -356,6 +388,7 @@ export function createEventHandlerState(): EventHandlerState {
     firstThinkingDeltaEmitted: false,
     lastCompletedToolName: undefined,
     toolCallTimestamps: new Map(),
+    toolPreviewStartedAt: new Map(),
     currentToolUseId: undefined,
     requestIdToToolUseId: new Map(),
     toolConfirmationOutcomes: new Map(),
@@ -372,6 +405,7 @@ export function createEventHandlerState(): EventHandlerState {
     currentThinkingTimestamps: [],
     lastPersistedContentSeq: undefined,
     compactionStartMessages: new Map(),
+    latencyCursor: 0,
   };
 }
 
@@ -505,6 +539,39 @@ function resetPartialPersistAccumulator(state: EventHandlerState): void {
   state.pendingPartialFlushPromise = undefined;
 }
 
+/**
+ * Persist an in-loop message-content write, retrying transient SQLite write
+ * contention (`SQLITE_BUSY`/`SQLITE_IOERR`) and swallowing a final failure so a
+ * lock held by another writer cannot abort the turn. Every in-loop write
+ * rewrites the full content snapshot of its assistant/tool-result row, so a
+ * dropped write is overwritten by a later write in the same turn (the
+ * end-of-turn finalize) or the next turn — missing one is a self-healing
+ * cosmetic gap, not corruption.
+ *
+ * Returns whether the write committed, so callers can gate dependent
+ * bookkeeping (e.g. advancing the persisted seq) on durable content.
+ */
+async function persistLoopMessageContent(
+  messageId: string,
+  contentJson: string,
+  op: string,
+  rlog: pino.Logger,
+): Promise<boolean> {
+  try {
+    await withSqliteRetry(() => updateMessageContent(messageId, contentJson), {
+      op,
+      context: { messageId },
+    });
+    return true;
+  } catch (err) {
+    rlog.error(
+      { err, messageId, op },
+      "in-loop message-content write failed after retries; continuing without interrupting the turn",
+    );
+    return false;
+  }
+}
+
 /** Flush `state.currentMessageContent` to the persisted assistant row. */
 async function flushAccumulatedContent(
   state: EventHandlerState,
@@ -525,18 +592,16 @@ async function flushAccumulatedContent(
   // again, but they are not part of this write.
   const flushedSeq = state.lastPersistedContentSeq;
 
-  try {
-    updateMessageContent(messageId, contentJson);
-    // Record only after the write commits, so the snapshot seq never
-    // claims content that is not yet durable.
-    if (flushedSeq != null) {
-      recordPersistedSeq(deps.ctx.conversationId, flushedSeq);
-    }
-  } catch (err) {
-    deps.rlog.warn(
-      { err, messageId },
-      "partial flush of accumulated assistant content failed; finalize at message_complete will recover",
-    );
+  const persisted = await persistLoopMessageContent(
+    messageId,
+    contentJson,
+    "partial_flush_assistant_content",
+    deps.rlog,
+  );
+  // Record only after the write commits, so the snapshot seq never
+  // claims content that is not yet durable.
+  if (persisted && flushedSeq != null) {
+    recordConversationPersistedSeq(deps.ctx.conversationId, flushedSeq);
   }
 }
 
@@ -897,6 +962,14 @@ export function handleToolUse(
   // fetched mid-tool (refresh / reconnect) carries it and clients can render a
   // running timer without having seen the live `tool_use_start` event.
   recordToolStartOnPersistedMessage(state, event.id, startedAt);
+  // Mirror the first-byte preview timestamp onto the same durable block so a
+  // mid-tool snapshot keeps the perceived-start anchor instead of falling back
+  // to execution start. The block exists now (message_complete wrote it before
+  // this tool event), unlike at `tool_use_preview_start` time.
+  const previewStartedAt = state.toolPreviewStartedAt.get(event.id);
+  if (previewStartedAt != null) {
+    recordToolPreviewStartOnPersistedMessage(state, event.id, previewStartedAt);
+  }
   const statusText = computeToolUseStatusText(event.name, event.input);
   deps.ctx.emitActivityState("tool_running", "tool_use_start", {
     requestId: deps.reqId,
@@ -910,6 +983,9 @@ export function handleToolUse(
     toolUseId: event.id,
     messageId: state.lastAssistantMessageId,
     startedAt,
+    // 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),
   });
   // `message_complete` always precedes tool events (see handleMessageComplete),
   // so this tool_use block is already durable in the assistant row. The
@@ -918,7 +994,7 @@ export function handleToolUse(
   // persisted seq to it. Without this the snapshot would advertise a seq below
   // an event it already incorporates, and a client applying `seq > snapshot.seq`
   // would replay this tool start.
-  recordPersistedSeq(deps.ctx.conversationId, getCurrentSeq());
+  recordConversationPersistedSeq(deps.ctx.conversationId, getCurrentSeq());
 }
 
 export function handleToolUsePreviewStart(
@@ -926,12 +1002,24 @@ export function handleToolUsePreviewStart(
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "tool_use_preview_start" }>,
 ): void {
+  // Stamp the first-byte timestamp on the server clock. The user-perceived
+  // latency timer anchors here, so clients can start rendering the tool card
+  // and ticking elapsed time the moment the call is recognized — well before
+  // its input finishes streaming (which can lag many seconds on a large input).
+  //
+  // We only record it in state here, not onto the persisted assistant row: the
+  // tool_use block does not exist yet (message_complete writes it after the
+  // stream ends, which is after this preview event). `handleToolUse` mirrors
+  // this timestamp onto the durable block once it exists.
+  const previewStartedAt = Date.now();
+  state.toolPreviewStartedAt.set(event.toolUseId, previewStartedAt);
   deps.onEvent({
     type: "tool_use_preview_start",
     toolUseId: event.toolUseId,
     toolName: event.toolName,
     conversationId: deps.ctx.conversationId,
     messageId: state.lastAssistantMessageId,
+    previewStartedAt,
   });
   const statusText = `Preparing ${friendlyToolName(event.toolName)}...`;
   deps.ctx.emitActivityState("tool_running", "preview_start", {
@@ -1132,11 +1220,15 @@ async function persistPendingToolResultRow(
   );
   // Serialize the content after the reservation resolves so the last of the
   // concurrent writers reflects the fullest batch.
-  updateMessageContent(
+  const persisted = await persistLoopMessageContent(
     rowId,
     JSON.stringify(buildToolResultBlocks(state.pendingToolResults)),
+    "persist_tool_result_row",
+    deps.rlog,
   );
-  recordPersistedSeq(deps.ctx.conversationId, seq);
+  if (persisted) {
+    recordConversationPersistedSeq(deps.ctx.conversationId, seq);
+  }
   const conv = getConversation(deps.ctx.conversationId);
   if (conv != null) {
     syncMessageToDisk(deps.ctx.conversationId, rowId, conv.createdAt);
@@ -1166,7 +1258,12 @@ export async function finalizePendingToolResultRow(
   const contentJson = JSON.stringify(
     buildToolResultBlocks(state.pendingToolResults),
   );
-  updateMessageContent(rowId, contentJson);
+  await persistLoopMessageContent(
+    rowId,
+    contentJson,
+    "finalize_tool_result_row",
+    rlog,
+  );
   // Sync the row to the JSONL disk view so it stays in lockstep with the DB.
   // `getConversation` returns `ConversationRow | null`, so `!= null` gates on a
   // real row (skipping the sync when the conversation was not found rather than
@@ -1464,7 +1561,66 @@ function recordToolStartOnPersistedMessage(
     if (rec.id !== toolUseId) continue;
     if (rec._startedAt === startedAt) return;
     rec._startedAt = startedAt;
-    updateMessageContent(messageId, JSON.stringify(content));
+    // Best-effort early stamp: `annotatePersistedAssistantMessage` re-stamps
+    // once every tool in the turn completes, so a transient `SQLITE_BUSY` here
+    // must not abort the turn — the end-of-turn write recovers it.
+    try {
+      updateMessageContent(messageId, JSON.stringify(content));
+    } catch (err) {
+      log.error(
+        { err, messageId },
+        "stamping tool start time failed; end-of-turn annotation will recover",
+      );
+    }
+    return;
+  }
+}
+
+/**
+ * Stamp `_previewStartedAt` (the first-byte timestamp) onto the durable
+ * tool_use block, mirroring `recordToolStartOnPersistedMessage`. Called from
+ * `handleToolUse` rather than `handleToolUsePreviewStart`: the block only exists
+ * once message_complete has written it, which happens after the preview event
+ * but before the tool event. Without this a `/messages` snapshot fetched
+ * mid-tool would lose the perceived-start anchor and clients would fall back to
+ * execution start — hiding the input-streaming gap the user actually waited
+ * through.
+ */
+function recordToolPreviewStartOnPersistedMessage(
+  state: EventHandlerState,
+  toolUseId: string,
+  previewStartedAt: number,
+): void {
+  const messageId = state.lastAssistantMessageId;
+  if (!messageId) return;
+
+  const row = getMessageById(messageId);
+  if (!row) return;
+
+  let content: ContentBlock[];
+  try {
+    content = JSON.parse(row.content) as ContentBlock[];
+  } catch {
+    return;
+  }
+
+  for (const block of content) {
+    if (block.type !== "tool_use") continue;
+    const rec = block as unknown as Record<string, unknown>;
+    if (rec.id !== toolUseId) continue;
+    if (rec._previewStartedAt === previewStartedAt) return;
+    rec._previewStartedAt = previewStartedAt;
+    // Best-effort early stamp, mirroring `recordToolStartOnPersistedMessage`:
+    // `annotatePersistedAssistantMessage` re-stamps at end of turn, so a
+    // transient `SQLITE_BUSY` here must not abort the turn.
+    try {
+      updateMessageContent(messageId, JSON.stringify(content));
+    } catch (err) {
+      log.error(
+        { err, messageId },
+        "stamping tool preview-start time failed; end-of-turn annotation will recover",
+      );
+    }
     return;
   }
 }
@@ -1507,6 +1663,11 @@ function annotatePersistedAssistantMessage(
         }
         modified = true;
       }
+      const previewStartedAt = state.toolPreviewStartedAt.get(id);
+      if (previewStartedAt != null) {
+        rec._previewStartedAt = previewStartedAt;
+        modified = true;
+      }
       const confirmation = state.toolConfirmationOutcomes.get(id);
       if (confirmation) {
         rec._confirmationDecision = confirmation.decision;
@@ -1580,6 +1741,10 @@ function annotatePersistedAssistantMessage(
   }
 
   if (modified) {
+    // This end-of-turn write is best-effort: the caller wraps it in a
+    // try/catch (and `dispatchAgentEvent` swallows `tool_result` handler
+    // errors), so a transient `SQLITE_BUSY` here is logged and the turn
+    // continues — it never reaches the turn-level catch.
     updateMessageContent(messageId, JSON.stringify(content));
   }
 
@@ -1618,6 +1783,8 @@ function handleError(
     buildConversationErrorMessage(deps.ctx.conversationId, classified),
   );
   state.providerErrorUserMessage = classified.userMessage;
+  state.persistProviderErrorAsAssistantMessage =
+    shouldPersistProviderErrorAsAssistantMessage(classified);
 }
 
 export function handleMaxTokensReached(
@@ -1688,15 +1855,22 @@ export async function handleMessageComplete(
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "message_complete" }>,
 ): Promise<void> {
-  // The model has now received the turn context, so persist any pending
-  // inference-profile-change notification. Guarded by the pending slot so it
-  // fires once per turn; a turn that fails before reaching delivery leaves the
-  // slot unconsumed and re-sends the notice next turn.
+  // A completed message means the model received the turn context, so record
+  // any pending inference-profile-change notification. Guarded by the pending
+  // slot so it fires once per turn; a turn that fails before reaching delivery
+  // leaves the slot unconsumed and re-sends the notice next turn.
   if (state.pendingNotifiedInferenceProfile != null) {
-    setLastNotifiedInferenceProfile(
-      deps.ctx.conversationId,
-      state.pendingNotifiedInferenceProfile,
-    );
+    try {
+      setLastNotifiedInferenceProfile(
+        deps.ctx.conversationId,
+        state.pendingNotifiedInferenceProfile,
+      );
+    } catch (err) {
+      deps.rlog.warn(
+        { conversationId: deps.ctx.conversationId, err },
+        "Failed to persist last notified inference profile (non-fatal)",
+      );
+    }
     deps.ctx.lastNotifiedInferenceProfile =
       state.pendingNotifiedInferenceProfile;
     state.pendingNotifiedInferenceProfile = null;
@@ -1808,7 +1982,12 @@ export async function handleMessageComplete(
     );
   }
   const contentJson = JSON.stringify(contentForPersistence);
-  updateMessageContent(assistantMessageId, contentJson);
+  const persisted = await persistLoopMessageContent(
+    assistantMessageId,
+    contentJson,
+    "finalize_assistant_message",
+    deps.rlog,
+  );
   state.assistantRowAwaitingFinalization = false;
   // The assistant row now holds the authoritative content (text + thinking +
   // tool_use blocks from `event.message`), and any drained tool-result rows
@@ -1816,10 +1995,15 @@ export async function handleMessageComplete(
   // delta's seq -- the highest stamped content event this row reflects -- so
   // recording it is honest. A drained tool result was stamped earlier in the
   // turn, so this seq already covers it; a call that streams no content (a
-  // pure tool call) advances instead via `tool_use_start`. `recordPersistedSeq`
-  // clamps monotonically, so a lower value here never regresses the seq.
-  if (state.lastPersistedContentSeq != null) {
-    recordPersistedSeq(deps.ctx.conversationId, state.lastPersistedContentSeq);
+  // pure tool call) advances instead via `tool_use_start`.
+  // `recordConversationPersistedSeq` clamps monotonically, so a lower value
+  // here never regresses the seq. Gate on `persisted` so a swallowed finalize
+  // write never advances the seq past content that is not durable.
+  if (persisted && state.lastPersistedContentSeq != null) {
+    recordConversationPersistedSeq(
+      deps.ctx.conversationId,
+      state.lastPersistedContentSeq,
+    );
   }
   // Reset the partial-persist mirror so subsequent calls in this turn
   // start with an empty running view.
@@ -2030,6 +2214,27 @@ function handleUsage(
     state.exchangeRawResponses.push(event.rawResponse);
   }
 
+  // Serialize this call's first-token latency segment and advance the cursor
+  // so the next call in the turn serializes only its own marks. Non-fatal: a
+  // tracking hiccup must never escalate into a turn-level throw.
+  let latencyBreakdownJson: string | undefined;
+  let ttftMs: number | undefined;
+  try {
+    const segment = deps.latencyTracker?.serializeSince(state.latencyCursor);
+    if (segment) {
+      state.latencyCursor = segment.cursor;
+      if (segment.breakdown) {
+        latencyBreakdownJson = JSON.stringify(segment.breakdown);
+        ttftMs = segment.breakdown.ttftMs ?? undefined;
+      }
+    }
+  } catch (err) {
+    deps.rlog.warn(
+      { err },
+      "Failed to serialize latency breakdown (non-fatal)",
+    );
+  }
+
   if (event.rawRequest && event.rawResponse) {
     try {
       recordRequestLog(
@@ -2039,6 +2244,7 @@ function handleUsage(
         undefined,
         providerName,
         "mainAgent",
+        latencyBreakdownJson,
       );
     } catch (err) {
       deps.rlog.warn({ err }, "Failed to persist LLM request log (non-fatal)");
@@ -2061,6 +2267,10 @@ function handleUsage(
         inputTokens: event.inputTokens,
         outputTokens: event.outputTokens,
         latencyMs: event.providerDurationMs,
+        // Time-to-first-token for this call (network + provider queue + first
+        // token), so the Logs-tab trace timeline reflects TTFT alongside the
+        // whole-call latency. Omitted when no token streamed (pure tool call).
+        ...(ttftMs != null ? { ttftMs } : {}),
       },
     },
   );
@@ -2127,6 +2337,13 @@ function handleProviderError(
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "provider_error" }>,
 ): void {
+  const classified = classifyConversationError(event.error, {
+    phase: "agent_loop",
+  });
+  if (!shouldPersistProviderErrorAsAssistantMessage(classified)) {
+    return;
+  }
+
   try {
     recordRequestLog(
       deps.ctx.conversationId,