[eric] thinking pill shows full turn cost (input + output + subagent + tool work) with click-to-see input/output breakdown; reopening

completed chats no longer replays the typewriter (per-session lastSeq survives AgentChat remount so resume protocol stays at the high-water
  mark instead of last_seq=0)

backend/apps/agents/agent_manager.py

@@ -2428,6 +2428,49 @@ class AgentManager:
                     _turn_total_ms = int((time.time() - _turn_started_ts) * 1000)
                 if _turn_thinking_msg_id is None:
                     _turn_thinking_msg_id = uuid4().hex
+                # Combined token total for the pill — input + output for
+                # the parent turn PLUS any work delegated to subagents
+                # (browser agents, invoke-agent forks) and tool MCP
+                # servers that produced their own usage on this turn.
+                # The user-visible answer to "how big is this turn" is
+                # the all-in sum, not just the primary's output. We sum
+                # every reachable source:
+                #   - parent's input  (session.tokens["input"] —
+                #     ResultMessage.usage at line ~2886)
+                #   - parent's output (session.tokens["output"] — same
+                #     ResultMessage)
+                #   - every direct sub-session whose parent_session_id
+                #     points at this session (browser agents, sub-agent
+                #     forks, invoke-agent calls book their own usage at
+                #     subprocess return time — agent_manager.py:1365 +
+                #     browser_agent.py:1000-1001)
+                # This mirrors how billing accumulates per-turn — caches,
+                # tool MCP servers that talk to LLMs (e.g. summarizers),
+                # and subagent reasoning all show up under the parent's
+                # "session.tokens" once their result lands.
+                _parent_in = 0
+                _parent_out = 0
+                if isinstance(session.tokens, dict):
+                    _parent_in = int(session.tokens.get("input", 0) or 0)
+                    _parent_out = int(session.tokens.get("output", 0) or 0)
+                _children_in = 0
+                _children_out = 0
+                try:
+                    for _child in self.sessions.values():
+                        if getattr(_child, "parent_session_id", None) != session.id:
+                            continue
+                        _ct = getattr(_child, "tokens", None)
+                        if not isinstance(_ct, dict):
+                            continue
+                        _children_in += int(_ct.get("input", 0) or 0)
+                        _children_out += int(_ct.get("output", 0) or 0)
+                except Exception:
+                    pass
+                _turn_total_tokens: int | None = (
+                    _parent_in + _parent_out + _children_in + _children_out
+                )
+                if not _turn_total_tokens or _turn_total_tokens <= 0:
+                    _turn_total_tokens = None
                 consolidated = Message(
                     id=_turn_thinking_msg_id,
                     role="thinking",
@@ -2435,6 +2478,7 @@ class AgentManager:
                     branch_id=session.active_branch_id,
                     elapsed_ms=_turn_total_ms or None,
                     tokens=turn_tokens,
+                    input_tokens=_turn_total_tokens,
                     tool_count=_turn_tool_count or None,
                     # Persist Gemini thoughtSignature so we can re-attach
                     # it on the next request — this is what stops Google
@@ -2810,6 +2854,30 @@ class AgentManager:
                         except Exception:
                             pass
 
+                        # Pre-populate session.tokens BEFORE emitting the
+                        # final consolidated thinking pill. Order matters:
+                        # _emit_consolidated_thinking reads
+                        # session.tokens["input"]/["output"] for the
+                        # combined-total stamp on the pill. If we emit
+                        # first, the pill freezes with input=0 because
+                        # the ResultMessage hasn't been consumed yet
+                        # (the writes below at line ~2918 wouldn't
+                        # land until after the pill is already broadcast).
+                        try:
+                            _pre_usage = getattr(message, "usage", None) or {}
+                            if isinstance(_pre_usage, dict):
+                                _pre_in = int(_pre_usage.get("input_tokens", 0) or 0)
+                                _pre_create = int(_pre_usage.get("cache_creation_input_tokens", 0) or 0)
+                                _pre_read = int(_pre_usage.get("cache_read_input_tokens", 0) or 0)
+                                _pre_total_in = _pre_in + _pre_create + _pre_read
+                                _pre_out = int(_pre_usage.get("output_tokens", 0) or 0)
+                                if _pre_total_in > 0:
+                                    session.tokens["input"] = _pre_total_in
+                                if _pre_out > 0:
+                                    session.tokens["output"] = _pre_out
+                        except Exception:
+                            pass
+
                         # Final consolidated emission with the full
                         # duration + authoritative tokens. The frontend
                         # bubble freezes on this final value.

backend/apps/agents/models.py

@@ -64,6 +64,15 @@ class Message(BaseModel):
     # "Thought for 18s · 430 reasoning · 2.4K answer · 3 tools" label.
     answer_tokens: Optional[int] = None
     tool_count: Optional[int] = None
+    # Combined input+output token total for the turn that produced this
+    # thinking message — including all sub-work delegated to subagents
+    # (browser, invoke-agent) and tool MCP servers that report their own
+    # usage. Stored under `input_tokens` for back-compat with older
+    # session JSONs even though the value is now the full
+    # input+output+children sum. This is the "how big was this turn"
+    # number that drives the pill's "M tokens" segment. None when no
+    # usage data was captured.
+    input_tokens: Optional[int] = None
     # Gemini 2.5/3.x emit a `thoughtSignature` (an opaque encrypted
     # blob) on each thinking block, and Google rejects subsequent
     # multi-step requests with a 400 if the signature isn't echoed

frontend/src/app/pages/AgentChat/AgentChat.tsx

@@ -178,12 +178,33 @@ const AgentChat: React.FC<AgentChatProps> = ({ sessionId: sessionIdProp, onClose
 
   useEffect(() => {
     if (!id || isDraft) return;
-    const ws = createSessionWs(id);
-    ws.connect();
-    wsRef.current = ws;
-    dispatch(fetchSession(id));
+    let cancelled = false;
+    let ws: ReturnType<typeof createSessionWs> | null = null;
+    // Order matters: hydrate the persisted message list from REST FIRST,
+    // THEN connect the WS. The WS resume protocol replays buffered
+    // events starting at last_seq=0, which includes every stream_*
+    // event for messages that finished before the disconnect. The
+    // replay-skip guard in WebSocketManager._messageAlreadyComplete
+    // checks `session.messages` to decide whether to drop deltas — so
+    // if we connect first, the slice is empty when the replay arrives,
+    // the guard returns false, and the user sees the chat type itself
+    // out again. Awaiting fetchSession before connect makes the slice
+    // authoritative before any replay event lands.
+    (async () => {
+      try {
+        await dispatch(fetchSession(id));
+      } catch {
+        // Even if the REST hydrate fails, still connect — the WS resume
+        // protocol can hydrate from buffered events as a fallback.
+      }
+      if (cancelled) return;
+      ws = createSessionWs(id);
+      ws.connect();
+      wsRef.current = ws;
+    })();
     return () => {
-      ws.disconnect();
+      cancelled = true;
+      if (ws) ws.disconnect();
       wsRef.current = null;
     };
   }, [id, isDraft, dispatch]);

frontend/src/app/pages/AgentChat/MessageBubble.tsx

@@ -507,6 +507,11 @@ const ThinkingBubble: React.FC<{
   // disappear when the streaming bubble unmounts.
   persistedElapsedMs?: number;
   persistedTokens?: number;
+  // Server-stamped input-side total for the turn (fresh + cache-creation
+  // + cache-read). Used to render "M in" alongside the existing "K out"
+  // segment so the pill honestly reflects the full turn cost, not just
+  // output. Optional — turns with no SDK usage data (rare) skip it.
+  persistedInputTokens?: number;
   // Tool invocation count for this turn — drives the "3 tools used"
   // segment of the post-stream label.
   persistedToolCount?: number;
@@ -514,7 +519,7 @@ const ThinkingBubble: React.FC<{
   // pull request", "Drafting your email"). Replaces the static
   // "Thinking…" verb when present and the stream is still active.
   dynamicLabel?: string | null;
-}> = ({ content, isStreaming, persistedElapsedMs, persistedTokens, persistedToolCount, dynamicLabel }) => {
+}> = ({ content, isStreaming, persistedElapsedMs, persistedTokens, persistedInputTokens, persistedToolCount, dynamicLabel }) => {
   const c = useClaudeTokens();
 
   // Live timer is only used as a fallback when we don't yet have
@@ -620,23 +625,99 @@ const ThinkingBubble: React.FC<{
   // experiments showed it confused users (it counted both visible reply
   // text AND tool-call JSON arguments, making tool-heavy turns
   // misleadingly look like long answers).
-  const buildPostStreamLabel = () => {
-    const parts: string[] = [];
-    if (finalSeconds != null) {
-      parts.push(`Thought for ${fmtThoughtDuration(finalSeconds)}`);
-    } else {
-      parts.push('Thoughts');
+  // Backend stamps `input_tokens` as the all-in input+output+children
+  // total (parent's primary call PLUS every subagent and tool MCP that
+  // booked usage on this turn). Falls back to just-output (finalTokens)
+  // for legacy thinking messages that predate the combined-total field.
+  const combinedTotalTokens =
+    persistedInputTokens != null && persistedInputTokens > 0
+      ? persistedInputTokens
+      : finalTokens;
+  // Input/output split shown in the breakdown tooltip on click. We
+  // already have `finalTokens` (server-stamped output side) and
+  // `combinedTotalTokens` (input + output + children sum). The
+  // implied "input + children" portion is the difference. When the
+  // backend hasn't separated them yet (legacy data), we still show
+  // the total but skip the breakdown.
+  const tokenBreakdown = (() => {
+    if (combinedTotalTokens == null || combinedTotalTokens <= 0) return null;
+    if (finalTokens == null || finalTokens <= 0) {
+      // Total-only case (rare). No split available.
+      return { total: combinedTotalTokens, output: null as number | null, input: null as number | null };
     }
-    if (finalTokens != null) {
-      parts.push(`${fmtTokens(finalTokens)} tokens`);
+    const inputSide = Math.max(0, combinedTotalTokens - finalTokens);
+    return { total: combinedTotalTokens, output: finalTokens, input: inputSide };
+  })();
+
+  const renderPostStreamLabel = () => {
+    const segments: React.ReactNode[] = [];
+    segments.push(
+      <span key="duration">
+        {finalSeconds != null
+          ? `Thought for ${fmtThoughtDuration(finalSeconds)}`
+          : 'Thoughts'}
+      </span>
+    );
+    if (tokenBreakdown) {
+      const { total, input, output } = tokenBreakdown;
+      const tooltipBody = input != null && output != null ? (
+        <Box sx={{ p: 0.5, fontFamily: c.font.sans, fontSize: '0.78rem', lineHeight: 1.5 }}>
+          <Box sx={{ display: 'flex', justifyContent: 'space-between', gap: 2 }}>
+            <span>Input</span><span style={{ fontVariantNumeric: 'tabular-nums' }}>{input.toLocaleString()}</span>
+          </Box>
+          <Box sx={{ display: 'flex', justifyContent: 'space-between', gap: 2 }}>
+            <span>Output</span><span style={{ fontVariantNumeric: 'tabular-nums' }}>{output.toLocaleString()}</span>
+          </Box>
+          <Box sx={{ display: 'flex', justifyContent: 'space-between', gap: 2, mt: 0.25, pt: 0.25, borderTop: `1px solid ${c.border.subtle}`, fontWeight: 600 }}>
+            <span>Total</span><span style={{ fontVariantNumeric: 'tabular-nums' }}>{total.toLocaleString()}</span>
+          </Box>
+          <Box sx={{ mt: 0.5, color: c.text.ghost, fontSize: '0.7rem', fontStyle: 'italic' }}>
+            Input includes system prompt, history, tool defs, cache reads, and any subagent/tool work this turn.
+          </Box>
+        </Box>
+      ) : (
+        <Box sx={{ p: 0.5, fontFamily: c.font.sans, fontSize: '0.78rem' }}>
+          {total.toLocaleString()} tokens (input + output + children)
+        </Box>
+      );
+      segments.push(<span key="sep-1"> · </span>);
+      segments.push(
+        <Tooltip
+          key="tokens"
+          title={tooltipBody}
+          placement="top"
+          arrow
+          slotProps={{ tooltip: { sx: { bgcolor: c.bg.elevated, color: c.text.primary, border: `1px solid ${c.border.medium}`, maxWidth: 'none' } } }}
+        >
+          <Box
+            component="span"
+            onClick={(e) => { e.stopPropagation(); }}
+            sx={{
+              cursor: 'help',
+              borderBottom: `1px dotted ${c.border.medium}`,
+              '&:hover': { color: c.text.secondary },
+            }}
+          >
+            {fmtTokens(total)} tokens
+          </Box>
+        </Tooltip>
+      );
     }
     if (persistedToolCount != null && persistedToolCount > 0) {
-      parts.push(`${persistedToolCount} tool${persistedToolCount === 1 ? '' : 's'} used`);
+      segments.push(<span key="sep-2"> · </span>);
+      segments.push(
+        <span key="tools">{persistedToolCount} tool{persistedToolCount === 1 ? '' : 's'} used</span>
+      );
     }
-    return parts.join(' · ');
+    return segments;
   };
 
-  const label = isStreaming ? activeLabel : buildPostStreamLabel();
+  // Streaming gets a plain string label (the shimmer animation needs
+  // the text to flow through a single gradient mask, which only works
+  // on a flat string node). Post-stream uses the React-node renderer
+  // so the tokens segment can be wrapped in a Tooltip with the
+  // input/output breakdown.
+  const label: React.ReactNode = isStreaming ? activeLabel : renderPostStreamLabel();
 
   // Shimmer colors — use a bright mid-tone against the muted base to make
   // the sweep visible without being loud. The base color matches the
@@ -805,6 +886,7 @@ const MessageBubble: React.FC<Props> = React.memo(({ message, editing = false, o
         timestamp={message.timestamp}
         persistedElapsedMs={(message as any).elapsed_ms}
         persistedTokens={(message as any).tokens}
+        persistedInputTokens={(message as any).input_tokens}
         persistedToolCount={(message as any).tool_count}
         dynamicLabel={isStreaming ? dynamicTurnLabel : null}
       />

frontend/src/shared/state/agentsSlice.ts

@@ -30,6 +30,11 @@ export interface AgentMessage {
   // survives reload instead of decaying to "Thoughts".
   elapsed_ms?: number;
   tokens?: number;
+  // Server-stamped input-side token count for the turn (fresh
+  // input + cache_creation + cache_read). Populated on thinking
+  // messages so the pill can show "Thought for Ns · M in / K out"
+  // — which is the only honest answer to "how big was this turn".
+  input_tokens?: number;
   // Richer thinking-pill data: total post-thinking output tokens
   // (user-visible answer text + tool arguments) and tool invocation
   // count. Drives the "Thought for 18s · 430 reasoning · 2.4K answer

frontend/src/shared/ws/WebSocketManager.ts

@@ -36,6 +36,7 @@ const _getAuthTokenSafe = (): string => {
   try { return getAuthToken() || ''; } catch { return ''; }
 };
 
+
 const _genUuid = (): string => {
   // Avoid pulling in `crypto.randomUUID` for compat — this is a
   // disambiguator, not a security boundary, so a 96-bit hex string is
@@ -126,6 +127,15 @@ class WebSocketManager {
     this.skipStreamEvents = options?.skipStreamEvents ?? false;
     this.sessionId = options?.sessionId ?? null;
     this.connectionUuid = _genUuid();
+    // Seed lastSeq from the cross-mount persistent map so a fresh
+    // manager (created on every AgentChat remount via key={session.id})
+    // doesn't ask the server to replay events the previous manager
+    // already saw. This is the architectural fix for "completed chats
+    // re-type themselves on reopen": the server's resume protocol now
+    // sees a real high-water mark and has nothing to replay.
+    if (this.sessionId) {
+      this.lastSeq = _sessionLastSeq.get(this.sessionId) ?? 0;
+    }
   }
 
   private bufferDelta(sessionId: string, messageId: string, delta: string) {
@@ -362,6 +372,11 @@ class WebSocketManager {
     // session, so this is the high-water mark we send back on resume.
     if (typeof msg.seq === 'number' && msg.seq > this.lastSeq) {
       this.lastSeq = msg.seq;
+      // Mirror to the module-scope persistent map so the next fresh
+      // manager (next AgentChat remount) starts here, not at zero.
+      if (this.sessionId) {
+        _sessionLastSeq.set(this.sessionId, this.lastSeq);
+      }
     }
 
     // ----- Connection-scoped frames (no business-logic side effects) -----
@@ -395,8 +410,11 @@ class WebSocketManager {
         store.dispatch(fetchSession(session_id));
         // Reset lastSeq — the REST refetch is the new authoritative
         // baseline; subsequent server events with seq numbers will
-        // re-establish the high-water mark.
+        // re-establish the high-water mark. Also wipe the cross-mount
+        // persistent map so a remount during this gap window doesn't
+        // resurrect the stale value.
         this.lastSeq = 0;
+        _sessionLastSeq.delete(session_id);
       }
       return;
     }
@@ -484,29 +502,46 @@ class WebSocketManager {
         break;
 
       case 'agent:stream_start':
-        if (session_id && data.message_id) {
-          store.dispatch(streamStart({
-            sessionId: session_id,
-            messageId: data.message_id,
-            role: data.role,
-            toolName: data.tool_name,
-          }));
-        }
-        break;
-
       case 'agent:stream_delta':
-        if (session_id && data.message_id) {
-          this.bufferDelta(session_id, data.message_id, data.delta);
-        }
-        break;
-
       case 'agent:stream_end':
-        if (session_id && data.message_id) {
-          this.flushInterpolator(data.message_id);
-          store.dispatch(streamEnd({
-            sessionId: session_id,
-            messageId: data.message_id,
-          }));
+        // Replay-skip guard. The WS resume protocol replays buffered
+        // events from the ring buffer with seq > last_seq. When this
+        // manager is freshly constructed (every AgentChat mount,
+        // because of `key={session.id}`), last_seq is 0, so the server
+        // replays EVERY buffered stream_* event for the session.
+        // Without this guard, opening any chat with prior streaming
+        // turns animates the entire history through the typewriter
+        // interpolator on every reopen.
+        //
+        // The discriminator is `resumeAcked`: it flips to true when
+        // server:hello arrives, which the server sends AFTER the replay
+        // completes. Any stream_* event arriving while !resumeAcked is
+        // replay-from-buffer (historical) and can be dropped — the REST
+        // snapshot we awaited before connect is authoritative for any
+        // already-finalized message, and any genuinely live turn the
+        // server is pushing will continue emitting events after the ack.
+        if (!this.resumeAcked) break;
+        if (event === 'agent:stream_start') {
+          if (session_id && data.message_id) {
+            store.dispatch(streamStart({
+              sessionId: session_id,
+              messageId: data.message_id,
+              role: data.role,
+              toolName: data.tool_name,
+            }));
+          }
+        } else if (event === 'agent:stream_delta') {
+          if (session_id && data.message_id) {
+            this.bufferDelta(session_id, data.message_id, data.delta);
+          }
+        } else if (event === 'agent:stream_end') {
+          if (session_id && data.message_id) {
+            this.flushInterpolator(data.message_id);
+            store.dispatch(streamEnd({
+              sessionId: session_id,
+              messageId: data.message_id,
+            }));
+          }
         }
         break;
 
@@ -765,6 +800,40 @@ import { WS_BASE } from '@/shared/config';
 
 export const dashboardWs = new WebSocketManager(`${WS_BASE}/ws/dashboard`, { skipStreamEvents: true });
 
+// Per-session high-water mark for the resume protocol. Survives across
+// AgentChat mounts/unmounts so reopening a chat doesn't re-trigger a
+// full replay from the server's ring buffer.
+//
+// Why this exists: AgentChat uses `key={session.id}` on the embedded
+// instance inside AgentCard, so every expand/collapse remounts the
+// component, which constructs a fresh WebSocketManager. Without this
+// persistent map, each fresh manager starts at last_seq=0 and asks the
+// server for the entire buffered history. The server faithfully
+// replays it, the client renders the typewriter animation again, and
+// the user sees their completed chat "type itself out" on every reopen.
+//
+// Lifetime: tied to the JS module load, which means the page tab. Lost
+// on full app reload (intentional — that should re-hydrate from REST).
+// On backend restart the buffers are wiped anyway, so a stale
+// lastSeq pointing past the buffer top falls into the "fresh client"
+// path on the server (last_seq>0 but no buffer) which short-circuits
+// to a no-op replay. Safe.
+const _sessionLastSeq: Map<string, number> = new Map();
+
+export function getPersistedLastSeq(sessionId: string): number {
+  return _sessionLastSeq.get(sessionId) ?? 0;
+}
+
+export function setPersistedLastSeq(sessionId: string, seq: number): void {
+  if (seq > (_sessionLastSeq.get(sessionId) ?? 0)) {
+    _sessionLastSeq.set(sessionId, seq);
+  }
+}
+
+export function clearPersistedLastSeq(sessionId: string): void {
+  _sessionLastSeq.delete(sessionId);
+}
+
 export function createSessionWs(sessionId: string): WebSocketManager {
   return new WebSocketManager(`${WS_BASE}/ws/agents/${sessionId}`, { sessionId });
 }