@vellumai/assistant 0.8.8 → 0.8.9-dev.202606091926.ebb2d62

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

@@ -17,7 +17,6 @@ import type {
 import { getConfig } from "../config/loader.js";
 import { recordEstimate } from "../context/estimator-calibration.js";
 import { getCalibrationProviderKey } from "../context/token-estimator.js";
-import type { ContextWindowResult } from "../context/window-manager.js";
 import { projectAssistantMessage } from "../memory/conversation-attention-store.js";
 import {
   deleteMessageById,
@@ -46,6 +45,7 @@ import {
   type SlackMessageMetadata,
   writeSlackMetadata,
 } from "../messaging/providers/slack/message-metadata.js";
+import type { ContextWindowResult } from "../plugins/defaults/compaction/window-manager.js";
 import type {
   ContentBlock,
   ImageContent,
@@ -76,10 +76,8 @@ import {
   cleanAssistantContent,
   drainDirectiveDisplayBuffer,
 } from "./assistant-attachments.js";
-import type {
-  AgentLoopConversationContext,
-  AssistantSurface,
-} from "./conversation-agent-loop.js";
+import type { Conversation } from "./conversation.js";
+import type { AssistantSurface } from "./conversation-agent-loop.js";
 import {
   buildConversationErrorMessage,
   classifyConversationError,
@@ -284,6 +282,15 @@ export interface EventHandlerState {
    * debounce so a mid-turn snapshot reflects what the user is watching live.
    */
   currentMessageContent: ContentBlock[];
+  /**
+   * Per-thinking-block timing for the in-flight LLM call, in stream order:
+   * entry `i` is the `i`-th thinking block mirrored into `currentMessageContent`.
+   * `startedAt` is stamped when a thinking block opens; `completedAt` is the
+   * last reasoning delta fused into it. Used to stamp `_startedAt`/`_completedAt`
+   * onto the authoritative thinking blocks at `message_complete`, mirroring how
+   * tool calls persist timing. Reset alongside `currentMessageContent`.
+   */
+  currentThinkingTimestamps: { startedAt: number; completedAt: number }[];
   /**
    * `seq` of the most recent streamed content delta mirrored into
    * `currentMessageContent`. Recorded as the conversation's persisted `seq`
@@ -295,29 +302,11 @@ export interface EventHandlerState {
    * never claims content a flush has not yet written.
    */
   lastPersistedContentSeq: number | undefined;
-  /**
-   * Whether the workspace top-level block should be (re)injected on this
-   * turn. Compaction's prepare phase strips the workspace / NOW.md / PKB
-   * blocks off the tail, so it is set after any successful compaction to
-   * force the workspace overview back in. On an ordinary turn the block is
-   * already present in history, so it defaults `false` to avoid burning
-   * tokens re-injecting it redundantly.
-   */
-  shouldInjectWorkspace: boolean;
-  /**
-   * Whether the reducer has compacted `ctx.messages`, gating the Slack
-   * chronological-transcript override on re-injection. The captured
-   * transcript is the full persisted history; blindly replaying it after
-   * compaction would overwrite the reduced messages and undo compaction, so
-   * once this is `true` the override falls back to the reduced
-   * `ctx.messages`.
-   */
-  reducerCompacted: boolean;
 }
 
 /** Immutable context shared across event handlers within a single agent loop run. */
 export interface EventHandlerDeps {
-  readonly ctx: AgentLoopConversationContext;
+  readonly ctx: Conversation;
   readonly onEvent: (msg: ServerMessage) => void;
   readonly reqId: string;
   readonly isFirstMessage: boolean;
@@ -391,9 +380,8 @@ export function createEventHandlerState(): EventHandlerState {
     pendingPartialFlushTimer: undefined,
     pendingPartialFlushPromise: undefined,
     currentMessageContent: [],
+    currentThinkingTimestamps: [],
     lastPersistedContentSeq: undefined,
-    shouldInjectWorkspace: false,
-    reducerCompacted: false,
   };
 }
 
@@ -442,6 +430,37 @@ export function buildPersistedAssistantContent(
   });
 }
 
+/**
+ * Stamp `_startedAt`/`_completedAt` onto the content's thinking blocks from the
+ * per-block timing captured while streaming, matched by position: the `i`-th
+ * thinking block in `content` takes `timings[i]`. The `_`-prefixed fields are
+ * vellum-internal metadata (never sent to providers), mirroring how tool calls
+ * persist timing; `renderHistoryContent` reads them back onto the wire schema's
+ * `startedAt`/`completedAt`.
+ *
+ * Timing is only captured when reasoning is streamed (`streamThinking`), so a
+ * turn with thinking disabled — or any block without a matching entry — is left
+ * unstamped and the UI hides its duration, exactly as a tool call with no
+ * timing does.
+ */
+export function stampThinkingTiming(
+  content: ContentBlock[],
+  timings: ReadonlyArray<{ startedAt: number; completedAt: number }>,
+): ContentBlock[] {
+  if (timings.length === 0) return content;
+  let thinkingIdx = 0;
+  return content.map((block) => {
+    if (block.type !== "thinking") return block;
+    const timing = timings[thinkingIdx++];
+    if (!timing) return block;
+    return {
+      ...block,
+      _startedAt: timing.startedAt,
+      _completedAt: timing.completedAt,
+    } as ContentBlock;
+  });
+}
+
 /** Append a streamed text chunk to `state.currentMessageContent`, fusing into tail text block. */
 function appendTextToCurrentMessage(
   state: EventHandlerState,
@@ -468,15 +487,19 @@ function appendThinkingToCurrentMessage(
   thinking: string,
 ): void {
   if (thinking.length === 0) return;
+  const now = Date.now();
   const tail = state.currentMessageContent.at(-1);
   if (tail && tail.type === "thinking") {
     tail.thinking = tail.thinking + thinking;
+    const timing = state.currentThinkingTimestamps.at(-1);
+    if (timing) timing.completedAt = now;
   } else {
     state.currentMessageContent.push({
       type: "thinking",
       thinking,
       signature: "",
     });
+    state.currentThinkingTimestamps.push({ startedAt: now, completedAt: now });
   }
 }
 
@@ -487,6 +510,7 @@ function resetPartialPersistAccumulator(state: EventHandlerState): void {
     state.pendingPartialFlushTimer = undefined;
   }
   state.currentMessageContent = [];
+  state.currentThinkingTimestamps = [];
   state.lastPersistedContentSeq = undefined;
   state.pendingPartialFlushPromise = undefined;
 }
@@ -655,9 +679,7 @@ function computeToolUseStatusText(
   return `Running ${friendlyToolName(name)}`;
 }
 
-function resolveAssistantReplyTimestampTimezone(
-  ctx: AgentLoopConversationContext,
-): string {
+function resolveAssistantReplyTimestampTimezone(ctx: Conversation): string {
   const config = getConfig();
   return resolveTurnTimezoneContext({
     configuredUserTimeZone: config.ui?.userTimezone ?? null,
@@ -876,9 +898,14 @@ export function handleToolUse(
   if (event.name === "app_create" || event.name === "app_refresh") {
     state.appBuildToolUsedThisRun = true;
   }
-  state.toolCallTimestamps.set(event.id, { startedAt: Date.now() });
+  const startedAt = Date.now();
+  state.toolCallTimestamps.set(event.id, { startedAt });
   state.currentToolUseId = event.id;
   state.currentTurnToolUseIds.push(event.id);
+  // Stamp the start time onto the already-durable tool_use block so a snapshot
+  // 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);
   const statusText = computeToolUseStatusText(event.name, event.input);
   deps.ctx.emitActivityState("tool_running", "tool_use_start", {
     requestId: deps.reqId,
@@ -891,6 +918,7 @@ export function handleToolUse(
     conversationId: deps.ctx.conversationId,
     toolUseId: event.id,
     messageId: state.lastAssistantMessageId,
+    startedAt,
   });
   // `message_complete` always precedes tool events (see handleMessageComplete),
   // so this tool_use block is already durable in the assistant row. The
@@ -1272,8 +1300,9 @@ export async function handleToolResult(
   });
 
   // Record tool completion timestamp
+  const completedAt = Date.now();
   const ts = state.toolCallTimestamps.get(event.toolUseId);
-  if (ts) ts.completedAt = Date.now();
+  if (ts) ts.completedAt = completedAt;
   state.currentToolUseId = undefined;
 
   // Capture risk metadata when present. autoApproved is true when the tool
@@ -1392,6 +1421,7 @@ export async function handleToolResult(
     approvalReason: event.approvalReason,
     riskThreshold: event.riskThreshold,
     activityMetadata: event.activityMetadata,
+    completedAt,
   });
 
   // Capture the seq synchronously (before the persist await) so it reflects the
@@ -1409,6 +1439,44 @@ export async function handleToolResult(
   }
 }
 
+/**
+ * Stamp `_startedAt` onto the in-flight tool_use block in the persisted
+ * assistant message the moment the tool begins. The block is already durable
+ * (message_complete precedes tool events), so without this a `/messages`
+ * snapshot fetched mid-tool would carry no start time and clients could not
+ * render a running elapsed-time counter until the whole turn finished. The
+ * full timing + risk annotation still happens in
+ * `annotatePersistedAssistantMessage` once every tool in the turn completes.
+ */
+function recordToolStartOnPersistedMessage(
+  state: EventHandlerState,
+  toolUseId: string,
+  startedAt: 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._startedAt === startedAt) return;
+    rec._startedAt = startedAt;
+    updateMessageContent(messageId, JSON.stringify(content));
+    return;
+  }
+}
+
 /**
  * After all tools for the current turn complete, fetch the persisted assistant
  * message, annotate its tool_use blocks with timing and confirmation metadata,
@@ -1506,6 +1574,13 @@ function annotatePersistedAssistantMessage(
         display: surface.display,
         ...(surface.persistent ? { persistent: true } : {}),
         ...(surface.toolCallId ? { toolCallId: surface.toolCallId } : {}),
+        // Daemon-only commit-timing activation tag, persisted so
+        // restoreSurfaceStateFromHistory can rehydrate it after a reload. This
+        // block lives only in server-side conversation history, never in the
+        // client `ui_surface_show` message.
+        ...(surface.activationMoment
+          ? { activationMoment: surface.activationMoment }
+          : {}),
       } as unknown as ContentBlock);
     }
     modified = true;
@@ -1659,6 +1734,8 @@ export async function handleMessageComplete(
       deps.ctx.conversationId,
       state.pendingNotifiedInferenceProfile,
     );
+    deps.ctx.lastNotifiedInferenceProfile =
+      state.pendingNotifiedInferenceProfile;
     state.pendingNotifiedInferenceProfile = null;
   }
 
@@ -1745,10 +1822,13 @@ export async function handleMessageComplete(
   // redacted) via the shared helper. The partial-persist flush uses
   // the same helper with `surfaces=[]` so a mid-turn snapshot lands in
   // the same shape as the finalize.
-  const contentForPersistence = buildPersistedAssistantContent(
-    event.message.content as ContentBlock[],
-    deps.ctx.currentTurnSurfaces,
-    state.toolActivityMetadata,
+  const contentForPersistence = stampThinkingTiming(
+    buildPersistedAssistantContent(
+      event.message.content as ContentBlock[],
+      deps.ctx.currentTurnSurfaces,
+      state.toolActivityMetadata,
+    ),
+    state.currentThinkingTimestamps,
   );
 
   // The row was reserved at `llm_call_started` (with channel metadata
@@ -1781,6 +1861,7 @@ export async function handleMessageComplete(
   // Reset the partial-persist mirror so subsequent calls in this turn
   // start with an empty running view.
   state.currentMessageContent = [];
+  state.currentThinkingTimestamps = [];
   state.lastPersistedContentSeq = undefined;
 
   // ── Indexing + attention projection ──
@@ -2302,17 +2383,14 @@ export async function dispatchAgentEvent(
         // so re-injection re-applies onto the stripped history even when the
         // pipeline ran but did not compact. When it did compact, commit the
         // durable result (DB-record fields, Slack provenance, SSE) — which
-        // overwrites `ctx.messages` with the compacted history — and flip the
-        // per-turn re-injection guards the orchestrator reads. This runs
+        // overwrites `ctx.messages` with the compacted history. This runs
         // before the loop's `reinject` hook (the loop awaits this dispatch),
-        // so the guards are set in time. A failed durable commit re-throws
-        // below to abort the turn rather than re-injecting against
-        // half-applied state.
+        // so the committed history is in place in time. A failed durable
+        // commit re-throws below to abort the turn rather than re-injecting
+        // against half-applied state.
         deps.ctx.messages = event.basis;
         if (event.result.compacted) {
           await deps.applyCompaction(event.result, event.basis);
-          state.reducerCompacted = true;
-          state.shouldInjectWorkspace = true;
         }
         break;
       case "history_stripped":