@vellumai/assistant 0.8.4 → 0.8.6

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

@@ -7,6 +7,7 @@
  */
 
 import type pino from "pino";
+import { v4 as uuid } from "uuid";
 
 import type { AgentEvent } from "../agent/loop.js";
 import type {
@@ -16,12 +17,16 @@ import type {
 import { getConfig } from "../config/loader.js";
 import { recordEstimate } from "../context/estimator-calibration.js";
 import { getCalibrationProviderKey } from "../context/token-estimator.js";
+import { projectAssistantMessage } from "../memory/conversation-attention-store.js";
 import {
+  deleteMessageById,
   getConversation,
   getMessageById,
+  messageMetadataSchema,
   provenanceFromTrustContext,
   updateMessageContent,
 } from "../memory/conversation-crud.js";
+import { indexMessageNow } from "../memory/indexer.js";
 import {
   backfillMessageIdOnLogs,
   buildProviderErrorResponsePayload,
@@ -40,15 +45,20 @@ import { defaultPersistenceTerminal } from "../plugins/defaults/persistence.js";
 import { DEFAULT_TIMEOUTS, runPipeline } from "../plugins/pipeline.js";
 import { getMiddlewaresFor } from "../plugins/registry.js";
 import type {
-  PersistAddResult,
   PersistArgs,
+  PersistReserveResult,
   PersistResult,
   TurnContext,
 } from "../plugins/types.js";
 import type { ContentBlock, ImageContent } from "../providers/types.js";
 import { isContextOverflowError } from "../providers/types.js";
+import { publishSyncInvalidation } from "../runtime/sync/sync-publisher.js";
 import { redactSecrets } from "../security/secret-scanner.js";
 import { extractDomain } from "../tools/network/domain-normalize.js";
+import {
+  buildPricingUsage,
+  resolveStructuredPricing,
+} from "../usage/pricing.js";
 import { ProviderError } from "../util/errors.js";
 import { faviconUrlForDomain } from "../util/favicon.js";
 import { getLogger } from "../util/logger.js";
@@ -62,10 +72,17 @@ import {
   buildConversationErrorMessage,
   classifyConversationError,
   isContextTooLarge,
+  maxTokensReachedClassification,
 } from "./conversation-error.js";
 import { isProviderOrderingError } from "./conversation-slash.js";
 import { resolveTurnTimezoneContext } from "./date-context.js";
-import type { ServerMessage } from "./message-protocol.js";
+import type {
+  CardSurfaceData,
+  ServerMessage,
+  SurfaceAction,
+  UiSurfaceShow,
+} from "./message-protocol.js";
+import { conversationMetadataSyncTag } from "./message-types/sync.js";
 import type {
   WebSearchMetadata,
   WebSearchResultItem,
@@ -144,12 +161,22 @@ export interface EventHandlerState {
    */
   contextTooLargeError: unknown;
   providerErrorUserMessage: string | null;
+  lastAssistantMessageId: string | undefined;
   /**
-   * First persisted assistant row in this run; history keeps this id when it
-   * merges tool-turn rows into one display bubble.
+   * True when `handleLlmCallStarted` has reserved an empty assistant row
+   * that has NOT yet been finalized via `handleMessageComplete`
+   * (`op:"updateContent"` + indexing + projection). Used by error/retry
+   * paths to detect a stranded reservation that must be cleaned up
+   * before the next LLM call reserves a fresh row — without it, every
+   * retryable failure (overflow, ordering, image overflow) and every
+   * terminal provider rejection would leak an empty assistant bubble
+   * into the transcript and mispoint downstream sync/projection.
+   *
+   * Cleared by `handleMessageComplete` on successful finalize, and by
+   * the synthetic-error branch in `conversation-agent-loop.ts` after it
+   * absorbs the reserved row into the error message.
    */
-  firstAssistantMessageId: string | undefined;
-  lastAssistantMessageId: string | undefined;
+  assistantRowAwaitingFinalization: boolean;
   readonly pendingToolResults: Map<string, PendingToolResult>;
   readonly persistedToolUseIds: Set<string>;
   readonly accumulatedDirectives: DirectiveRequest[];
@@ -249,8 +276,8 @@ export function createEventHandlerState(): EventHandlerState {
     imageTooLargeDetected: false,
     contextTooLargeError: null,
     providerErrorUserMessage: null,
-    firstAssistantMessageId: undefined,
     lastAssistantMessageId: undefined,
+    assistantRowAwaitingFinalization: false,
     pendingToolResults: new Map(),
     persistedToolUseIds: new Set(),
     accumulatedDirectives: [],
@@ -275,12 +302,6 @@ export function createEventHandlerState(): EventHandlerState {
   };
 }
 
-export function getClientDisplayMessageId(
-  state: EventHandlerState,
-): string | undefined {
-  return state.firstAssistantMessageId ?? state.lastAssistantMessageId;
-}
-
 // ── Shared Helper ────────────────────────────────────────────────────
 
 // providerNameOverride should be supplied when the caller already knows the
@@ -319,6 +340,9 @@ function emitLlmCallStartedIfNeeded(
 // tools the client discards it (extractCodePreview only handles app tools),
 // so we skip forwarding entirely to avoid transport/decode overhead.
 const APP_TOOL_NAMES = new Set(["app_create"]);
+const MAX_TOKENS_CONTINUE_PROMPT =
+  "Continue from where you stopped. Do not repeat content you've already sent.";
+const MAX_TOKENS_SURFACE_COMPLETION_SUMMARY = "Continue";
 
 // ── Friendly Tool Names ──────────────────────────────────────────────
 
@@ -401,6 +425,133 @@ function resolveAssistantReplyTimestampTimezone(
   }).effectiveTimezone;
 }
 
+/**
+ * Assemble the metadata envelope written to the assistant message row.
+ *
+ * Stamped at reserve time (before `provider.sendMessage`) so the row carries
+ * channel provenance from the moment it lands in SQLite, mirroring the
+ * snapshot that handleMessageComplete used to compute at end-of-turn. All
+ * inputs (channel context, trust context, turnStartedAt) are stable across
+ * the LLM call, so building this once at reserve is equivalent to building
+ * it at complete. Slack reply rows further stamp a `slackMeta` sub-object —
+ * the `channelTs` field stays absent here and is back-filled by
+ * `deliverReplyViaCallback` after the gateway returns the ts.
+ */
+function buildAssistantChannelMetadata(
+  state: EventHandlerState,
+  deps: EventHandlerDeps,
+): Record<string, unknown> {
+  const metadata: Record<string, unknown> = {
+    ...provenanceFromTrustContext(deps.ctx.trustContext),
+    userMessageChannel: deps.turnChannelContext.userMessageChannel,
+    assistantMessageChannel: deps.turnChannelContext.assistantMessageChannel,
+    userMessageInterface: deps.turnInterfaceContext.userMessageInterface,
+    assistantMessageInterface:
+      deps.turnInterfaceContext.assistantMessageInterface,
+    sentAt: state.turnStartedAt,
+  };
+
+  if (deps.turnChannelContext.assistantMessageChannel === "slack") {
+    const channelId = deps.ctx.trustContext?.requesterChatId;
+    if (channelId) {
+      const threadTs = getThreadTs(deps.ctx.conversationId);
+      const timestampTimezone = resolveAssistantReplyTimestampTimezone(
+        deps.ctx,
+      );
+      const timestampTimezoneLabel = formatSlackTimezoneLabel(
+        timestampTimezone,
+        { nowMs: state.turnStartedAt },
+      );
+      const partialSlackMeta: Partial<SlackMessageMetadata> = {
+        source: "slack",
+        eventKind: "message",
+        channelId,
+        ...(threadTs ? { threadTs } : {}),
+        timestampTimezone,
+        ...(timestampTimezoneLabel ? { timestampTimezoneLabel } : {}),
+      };
+      // `channelTs` is filled in by the post-send reconciliation step in
+      // `deliverReplyViaCallback`; cast through the Partial to satisfy
+      // the writer's type at this pre-send boundary.
+      metadata.slackMeta = writeSlackMetadata(
+        partialSlackMeta as SlackMessageMetadata,
+      );
+    }
+  }
+
+  return metadata;
+}
+
+/**
+ * Reserve an empty assistant row for the LLM call about to begin, stash
+ * its id on `state.lastAssistantMessageId`, and announce the boundary on
+ * the wire via `assistant_turn_start`.
+ *
+ * Awaited so the row exists and the client has the anchor id BEFORE any
+ * streaming delta arrives — every subsequent `deps.onEvent` in this LLM
+ * call stamps `messageId: state.lastAssistantMessageId`, and
+ * `handleMessageComplete` flushes the final content to the same row via
+ * `op: "updateContent"` instead of inserting a fresh one.
+ *
+ * Multi-LLM-call agent turns (LLM call → tool execution → LLM call) emit
+ * one `llm_call_started` per call, so each LLM call reserves its own row.
+ * The read-path `findDisplayTurnEndIndex` collapses consecutive assistant
+ * rows for the merged history view, matching today's per-call DB layout.
+ */
+export async function handleLlmCallStarted(
+  state: EventHandlerState,
+  deps: EventHandlerDeps,
+): Promise<void> {
+  // Clean up an orphaned reservation from a previous LLM call in this run
+  // that errored before `message_complete` could finalize it. This covers
+  // the retryable paths (overflow, ordering, image overflow) where the
+  // agent loop re-enters with a fresh `run()` and reserves another row;
+  // without this delete the failed-attempt row stays in the transcript as
+  // an empty assistant bubble. The finalized-row case is filtered out via
+  // the `assistantRowAwaitingFinalization` flag — `handleMessageComplete`
+  // clears it after the successful `updateContent`, so the previous call's
+  // committed row is never touched here.
+  //
+  // Direct `deleteMessageById` (not via the `persistence` pipeline) is
+  // intentional: a never-finalized reservation has no segments, no
+  // attachments, and no observable history — undoing it isn't a real
+  // persistence event for plugins to react to, so routing through the
+  // pipeline would only widen the mock surface for no observability win.
+  if (state.assistantRowAwaitingFinalization && state.lastAssistantMessageId) {
+    try {
+      deleteMessageById(state.lastAssistantMessageId);
+    } catch (err) {
+      // Non-fatal: a leaked empty row is preferable to a turn-level throw.
+      deps.rlog.warn(
+        { err, messageId: state.lastAssistantMessageId },
+        "Failed to clean up stranded reserved assistant row before new reservation",
+      );
+    }
+  }
+
+  const metadata = buildAssistantChannelMetadata(state, deps);
+  const reserveResult = (await runPipeline<PersistArgs, PersistResult>(
+    "persistence",
+    getMiddlewaresFor("persistence"),
+    defaultPersistenceTerminal,
+    {
+      op: "reserve",
+      conversationId: deps.ctx.conversationId,
+      role: "assistant",
+      metadata,
+    },
+    buildHandlerTurnContext(deps),
+    DEFAULT_TIMEOUTS.persistence,
+  )) as PersistReserveResult;
+  state.lastAssistantMessageId = reserveResult.message.id;
+  state.assistantRowAwaitingFinalization = true;
+  deps.onEvent({
+    type: "assistant_turn_start",
+    messageId: reserveResult.message.id,
+    conversationId: deps.ctx.conversationId,
+  });
+}
+
 // ── Individual Handlers ──────────────────────────────────────────────
 
 function handleTextDelta(
@@ -429,6 +580,7 @@ function handleTextDelta(
       type: "assistant_text_delta",
       text: drained.emitText,
       conversationId: deps.ctx.conversationId,
+      messageId: state.lastAssistantMessageId,
     });
     if (deps.shouldGenerateTitle) state.firstAssistantText += drained.emitText;
   }
@@ -466,6 +618,7 @@ function handleThinkingDelta(
     type: "assistant_thinking_delta",
     thinking: event.thinking,
     conversationId: deps.ctx.conversationId,
+    messageId: state.lastAssistantMessageId,
   });
 }
 
@@ -496,11 +649,12 @@ export function handleToolUse(
     input: event.input,
     conversationId: deps.ctx.conversationId,
     toolUseId: event.id,
+    messageId: state.lastAssistantMessageId,
   });
 }
 
 export function handleToolUsePreviewStart(
-  _state: EventHandlerState,
+  state: EventHandlerState,
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "tool_use_preview_start" }>,
 ): void {
@@ -509,6 +663,7 @@ export function handleToolUsePreviewStart(
     toolUseId: event.toolUseId,
     toolName: event.toolName,
     conversationId: deps.ctx.conversationId,
+    messageId: state.lastAssistantMessageId,
   });
   const statusText = `Preparing ${friendlyToolName(event.toolName)}...`;
   deps.ctx.emitActivityState(
@@ -521,7 +676,7 @@ export function handleToolUsePreviewStart(
 }
 
 function handleToolOutputChunk(
-  _state: EventHandlerState,
+  state: EventHandlerState,
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "tool_output_chunk" }>,
 ): void {
@@ -579,6 +734,7 @@ function handleToolOutputChunk(
       chunk: event.chunk,
       conversationId: deps.ctx.conversationId,
       toolUseId: event.toolUseId,
+      messageId: state.lastAssistantMessageId,
       subType: structured.subType,
       subToolName: structured.subToolName,
       subToolInput: structured.subToolInput,
@@ -591,12 +747,13 @@ function handleToolOutputChunk(
       chunk: event.chunk,
       conversationId: deps.ctx.conversationId,
       toolUseId: event.toolUseId,
+      messageId: state.lastAssistantMessageId,
     });
   }
 }
 
 export function handleInputJsonDelta(
-  _state: EventHandlerState,
+  state: EventHandlerState,
   deps: EventHandlerDeps,
   event: Extract<AgentEvent, { type: "input_json_delta" }>,
 ): void {
@@ -610,6 +767,7 @@ export function handleInputJsonDelta(
     content: event.accumulatedJson,
     conversationId: deps.ctx.conversationId,
     toolUseId: event.toolUseId,
+    messageId: state.lastAssistantMessageId,
   });
 }
 
@@ -736,6 +894,7 @@ export function handleToolResult(
     diff: event.diff,
     status: event.status,
     conversationId: deps.ctx.conversationId,
+    messageId: state.lastAssistantMessageId,
     imageData: imageDataList?.[0],
     imageDataList,
     toolUseId: event.toolUseId,
@@ -915,6 +1074,69 @@ function handleError(
   }
 }
 
+export function handleMaxTokensReached(
+  _state: EventHandlerState,
+  deps: EventHandlerDeps,
+  event: Extract<AgentEvent, { type: "max_tokens_reached" }>,
+): void {
+  const classified = maxTokensReachedClassification();
+  const surfaceId = `max_tokens_${uuid()}`;
+  const data: CardSurfaceData = {
+    title: "Response limit reached",
+    subtitle: "The partial response above was saved.",
+    body: classified.userMessage,
+  };
+  const actions: SurfaceAction[] = [
+    {
+      id: "relay_prompt",
+      label: "Continue",
+      style: "primary",
+      data: {
+        prompt: MAX_TOKENS_CONTINUE_PROMPT,
+        _completeSurface: true,
+        _completionSummary: MAX_TOKENS_SURFACE_COMPLETION_SUMMARY,
+      },
+    },
+  ];
+
+  deps.ctx.surfaceState.set(surfaceId, {
+    surfaceType: "card",
+    title: data.title,
+    data,
+    actions,
+  });
+  deps.ctx.currentTurnSurfaces.push({
+    surfaceId,
+    surfaceType: "card",
+    title: data.title,
+    data,
+    actions,
+    display: "inline",
+    persistent: true,
+  });
+
+  deps.rlog.warn(
+    {
+      conversationId: deps.ctx.conversationId,
+      stopReason: event.stopReason,
+      surfaceId,
+    },
+    "Surfacing max-tokens continuation card",
+  );
+
+  deps.onEvent({
+    type: "ui_surface_show",
+    conversationId: deps.ctx.conversationId,
+    surfaceId,
+    surfaceType: "card",
+    title: data.title,
+    data,
+    actions,
+    display: "inline",
+    persistent: true,
+  } as UiSurfaceShow);
+}
+
 export async function handleMessageComplete(
   state: EventHandlerState,
   deps: EventHandlerDeps,
@@ -929,6 +1151,7 @@ export async function handleMessageComplete(
       type: "assistant_text_delta",
       text: state.pendingDirectiveDisplayBuffer,
       conversationId: deps.ctx.conversationId,
+      messageId: state.lastAssistantMessageId,
     });
     if (deps.shouldGenerateTitle)
       state.firstAssistantText += state.pendingDirectiveDisplayBuffer;
@@ -1035,52 +1258,6 @@ export async function handleMessageComplete(
     } as unknown as ContentBlock);
   }
 
-  const assistantChannelMetadata: Record<string, unknown> = {
-    ...provenanceFromTrustContext(deps.ctx.trustContext),
-    userMessageChannel: deps.turnChannelContext.userMessageChannel,
-    assistantMessageChannel: deps.turnChannelContext.assistantMessageChannel,
-    userMessageInterface: deps.turnInterfaceContext.userMessageInterface,
-    assistantMessageInterface:
-      deps.turnInterfaceContext.assistantMessageInterface,
-    sentAt: state.turnStartedAt,
-  };
-
-  // When the assistant is replying through Slack, stamp a `slackMeta`
-  // sub-object so the transcript-rendering / thread-aware-context lookup
-  // can identify this row's thread without joining tables.
-  // Persistence happens BEFORE the Slack adapter sends the message, so
-  // Slack's authoritative `ts` (-> `channelTs`) is not yet known and is
-  // intentionally omitted here. The post-send reconciliation step in
-  // `deliverReplyViaCallback` writes `channelTs` back into this row once
-  // the gateway returns the Slack-assigned ts, restoring a fully-formed
-  // metadata envelope before any subsequent turn reads the row.
-  if (deps.turnChannelContext.assistantMessageChannel === "slack") {
-    const channelId = deps.ctx.trustContext?.requesterChatId;
-    if (channelId) {
-      const threadTs = getThreadTs(deps.ctx.conversationId);
-      const timestampTimezone = resolveAssistantReplyTimestampTimezone(
-        deps.ctx,
-      );
-      const timestampTimezoneLabel = formatSlackTimezoneLabel(
-        timestampTimezone,
-        { nowMs: state.turnStartedAt },
-      );
-      const partialSlackMeta: Partial<SlackMessageMetadata> = {
-        source: "slack",
-        eventKind: "message",
-        channelId,
-        ...(threadTs ? { threadTs } : {}),
-        timestampTimezone,
-        ...(timestampTimezoneLabel ? { timestampTimezoneLabel } : {}),
-      };
-      assistantChannelMetadata.slackMeta = writeSlackMetadata(
-        // `channelTs` is filled in by the post-send reconciliation step in
-        // `deliverReplyViaCallback`; cast through the Partial to satisfy
-        // the writer's type at this pre-send boundary.
-        partialSlackMeta as SlackMessageMetadata,
-      );
-    }
-  }
   // Redact known-pattern secrets from assistant text blocks before they are
   // written to durable storage. Non-text blocks (images, UI surfaces) pass
   // through unchanged. The live model history retains the original values.
@@ -1092,33 +1269,123 @@ export async function handleMessageComplete(
     return block;
   });
 
-  // Route the assistant-message persistence through the `persistence`
-  // pipeline. No `syncToDisk` here — the orchestrator separately invokes
-  // `syncMessageToDisk` on `state.lastAssistantMessageId` after the loop
-  // completes (see `conversation-agent-loop.ts::syncLastAssistantMessageToDisk`).
-  const assistantPersistResult = (await runPipeline<PersistArgs, PersistResult>(
+  // The row was reserved at `llm_call_started` (with channel metadata
+  // stamped at that point) and `state.lastAssistantMessageId` carries its
+  // id. Flush the final content via `updateContent` instead of inserting a
+  // new row. No `syncToDisk` flag here — the orchestrator separately
+  // invokes `syncMessageToDisk` on `state.lastAssistantMessageId` after
+  // the loop completes (see
+  // `conversation-agent-loop.ts::syncLastAssistantMessageToDisk`).
+  const assistantMessageId = state.lastAssistantMessageId;
+  if (!assistantMessageId) {
+    throw new Error(
+      "handleMessageComplete fired without a prior llm_call_started reserving an assistant row",
+    );
+  }
+  const contentJson = JSON.stringify(contentForPersistence);
+  await runPipeline<PersistArgs, PersistResult>(
     "persistence",
     getMiddlewaresFor("persistence"),
     defaultPersistenceTerminal,
     {
-      op: "add",
-      conversationId: deps.ctx.conversationId,
-      role: "assistant",
-      content: JSON.stringify(contentForPersistence),
-      metadata: assistantChannelMetadata,
+      op: "updateContent",
+      messageId: assistantMessageId,
+      content: contentJson,
     },
     buildHandlerTurnContext(deps),
     DEFAULT_TIMEOUTS.persistence,
-  )) as PersistAddResult;
-  const assistantMsg = assistantPersistResult.message;
-  state.firstAssistantMessageId ??= assistantMsg.id;
-  state.lastAssistantMessageId = assistantMsg.id;
+  );
+  state.assistantRowAwaitingFinalization = false;
+
+  // ── Indexing + attention projection (restored from the pre-B3 `add` path) ──
+  // `reserveMessage` + `updateMessageContent` are CRUD-only: they don't run
+  // the memory indexer or the attention-cursor projector. The pre-B3 path
+  // wrote the row via `addMessage`, which ran both as side-effects of the
+  // insert. Calling them here keeps the assistant row's external state
+  // (Qdrant segments, conversation attention cursor) in lockstep with the
+  // finalized content. Both are non-fatal — a memory hiccup must not
+  // escalate a successful generation into a turn-level throw. Indexing
+  // intentionally fires AFTER `updateContent` succeeds so we never index
+  // the empty reserved placeholder.
+  const finalizedRow = getMessageById(
+    assistantMessageId,
+    deps.ctx.conversationId,
+  );
+  if (finalizedRow) {
+    let provenanceTrustClass:
+      | "guardian"
+      | "trusted_contact"
+      | "unknown"
+      | undefined;
+    let automated: boolean | undefined;
+    if (finalizedRow.metadata) {
+      try {
+        const parsedMeta = messageMetadataSchema.safeParse(
+          JSON.parse(finalizedRow.metadata),
+        );
+        if (parsedMeta.success) {
+          provenanceTrustClass = parsedMeta.data.provenanceTrustClass;
+          automated = parsedMeta.data.automated;
+        }
+      } catch {
+        // Malformed metadata JSON — fall through with undefined fields,
+        // matching the legacy behavior in `addMessage`.
+      }
+    }
+    try {
+      await indexMessageNow(
+        {
+          messageId: assistantMessageId,
+          conversationId: deps.ctx.conversationId,
+          role: "assistant",
+          content: contentJson,
+          createdAt: finalizedRow.createdAt,
+          scopeId: "default",
+          provenanceTrustClass,
+          automated,
+        },
+        getConfig().memory,
+      );
+    } catch (err) {
+      deps.rlog.warn(
+        {
+          err,
+          conversationId: deps.ctx.conversationId,
+          messageId: assistantMessageId,
+        },
+        "Failed to index assistant message for memory (non-fatal)",
+      );
+    }
+    try {
+      const attentionStateChanged = projectAssistantMessage({
+        conversationId: deps.ctx.conversationId,
+        messageId: assistantMessageId,
+        messageAt: finalizedRow.createdAt,
+      });
+      if (attentionStateChanged) {
+        void publishSyncInvalidation([
+          conversationMetadataSyncTag(deps.ctx.conversationId),
+        ]);
+      }
+    } catch (err) {
+      deps.rlog.warn(
+        {
+          err,
+          conversationId: deps.ctx.conversationId,
+          messageId: assistantMessageId,
+        },
+        "Failed to project assistant message for attention tracking (non-fatal)",
+      );
+    }
+  }
 
   // Backfill message_id on all LLM request logs from this turn.
   // The agent loop is single-threaded per conversation, so all rows with
-  // message_id IS NULL belong to the current turn.
+  // message_id IS NULL belong to the current turn. The reserved id was
+  // available before the LLM call ran but the logs are inserted DURING
+  // the call, so the sweep still runs here.
   try {
-    backfillMessageIdOnLogs(deps.ctx.conversationId, assistantMsg.id);
+    backfillMessageIdOnLogs(deps.ctx.conversationId, assistantMessageId);
   } catch (err) {
     deps.rlog.warn(
       { err },
@@ -1127,7 +1394,10 @@ export async function handleMessageComplete(
   }
 
   try {
-    backfillMemoryRecallLogMessageId(deps.ctx.conversationId, assistantMsg.id);
+    backfillMemoryRecallLogMessageId(
+      deps.ctx.conversationId,
+      assistantMessageId,
+    );
   } catch (err) {
     deps.rlog.warn(
       { err },
@@ -1138,7 +1408,7 @@ export async function handleMessageComplete(
   try {
     backfillMemoryV2ActivationMessageId(
       deps.ctx.conversationId,
-      assistantMsg.id,
+      assistantMessageId,
     );
   } catch (err) {
     deps.rlog.warn(
@@ -1222,6 +1492,7 @@ function handleUsage(
         JSON.stringify(event.rawResponse),
         undefined,
         providerName,
+        "mainAgent",
       );
     } catch (err) {
       deps.rlog.warn({ err }, "Failed to persist LLM request log (non-fatal)");
@@ -1248,6 +1519,36 @@ function handleUsage(
     },
   );
   state.llmCallStartedEmitted = false;
+
+  // Emit a lightweight per-call usage progress event so clients can show
+  // live-updating token/cost metrics. This is a UI hint only — no DB writes.
+  const pricingUsage = buildPricingUsage({
+    providerName,
+    model: event.model,
+    inputTokens: event.inputTokens,
+    outputTokens: event.outputTokens,
+    cacheCreationInputTokens: event.cacheCreationInputTokens,
+    cacheReadInputTokens: event.cacheReadInputTokens,
+    rawResponse: event.rawResponse,
+  });
+  const pricing = resolveStructuredPricing(
+    providerName,
+    event.model,
+    pricingUsage,
+  );
+  const estimatedCost =
+    pricing.pricingStatus === "priced" && pricing.estimatedCostUsd != null
+      ? pricing.estimatedCostUsd
+      : 0;
+
+  deps.onEvent({
+    type: "usage_progress",
+    conversationId: deps.ctx.conversationId,
+    inputTokens: event.inputTokens,
+    outputTokens: event.outputTokens,
+    estimatedCost,
+    model: event.model,
+  });
 }
 
 /**
@@ -1287,6 +1588,7 @@ function handleProviderError(
       JSON.stringify(buildProviderErrorResponsePayload(event.error)),
       undefined,
       event.actualProvider,
+      "mainAgent",
     );
   } catch (err) {
     deps.rlog.warn(
@@ -1306,6 +1608,9 @@ export async function dispatchAgentEvent(
 ): Promise<void> {
   try {
     switch (event.type) {
+      case "llm_call_started":
+        await handleLlmCallStarted(state, deps);
+        break;
       case "text_delta":
         handleTextDelta(state, deps, event);
         break;
@@ -1346,6 +1651,7 @@ export async function dispatchAgentEvent(
           input: event.input,
           conversationId: deps.ctx.conversationId,
           toolUseId: event.toolUseId,
+          messageId: state.lastAssistantMessageId,
         });
         break;
       }
@@ -1425,6 +1731,7 @@ export async function dispatchAgentEvent(
           isError: event.isError,
           conversationId: deps.ctx.conversationId,
           toolUseId: event.toolUseId,
+          messageId: state.lastAssistantMessageId,
           ...(metadata ? { activityMetadata: { webSearch: metadata } } : {}),
         });
         break;
@@ -1432,6 +1739,9 @@ export async function dispatchAgentEvent(
       case "error":
         handleError(state, deps, event);
         break;
+      case "max_tokens_reached":
+        handleMaxTokensReached(state, deps, event);
+        break;
       case "provider_error":
         handleProviderError(deps, event);
         break;