From 5c82a2d653231f1adbb7ee27bd99acc7fe41fa73 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Thu, 27 Nov 2025 13:26:31 +0000
Subject: [PATCH] Align assistant metadata display with message content

---
 packages/ui/src/components/message-item.tsx   | 32 +++++++-----
 .../ui/src/components/message-stream-v2.tsx   | 49 +++++++++++++++++--
 .../ui/src/styles/messaging/message-base.css  |  2 +-
 3 files changed, 67 insertions(+), 16 deletions(-)

diff --git a/packages/ui/src/components/message-item.tsx b/packages/ui/src/components/message-item.tsx
index 39cf9b84..6ba7995e 100644
--- a/packages/ui/src/components/message-item.tsx
+++ b/packages/ui/src/components/message-item.tsx
@@ -14,6 +14,7 @@ interface MessageItemProps {
   orderedParts: ClientPart[]
   onRevert?: (messageId: string) => void
   onFork?: (messageId?: string) => void
+  showAgentMeta?: boolean
 }
 
 export default function MessageItem(props: MessageItemProps) {
@@ -166,17 +167,26 @@ export default function MessageItem(props: MessageItemProps) {
 
   return (
 
-    <div class={containerClass()}>
-      <div class={`flex justify-between items-center gap-2.5 ${isUser() ? "pb-0.5" : "pb-0"}`}>
-        <div class="flex flex-col">
-          <Show when={isUser()}>
-            <span class="font-semibold text-xs text-[var(--message-user-border)]">You</span>
-          </Show>
-          <Show when={!isUser()}>
-            <span class="text-xs text-[var(--message-assistant-border)]">Assistant</span>
-          </Show>
-        </div>
-        <div class="flex items-center gap-2">
+      <div class={containerClass()}>
+        <div class={`flex justify-between items-center gap-2.5 ${isUser() ? "pb-0.5" : "pb-0"}`}>
+          <div class="flex flex-col">
+            <Show when={isUser()}>
+              <span class="font-semibold text-xs text-[var(--message-user-border)]">You</span>
+            </Show>
+            <Show when={!isUser()}>
+              <div class="flex flex-wrap items-center gap-2 text-xs text-[var(--message-assistant-border)]">
+                <span class="font-semibold">Assistant</span>
+                <Show when={props.showAgentMeta && (agentIdentifier() || modelIdentifier())}>
+                  <span class="message-step-meta-inline">
+                    <Show when={agentIdentifier()}>{(value) => <span class="font-medium text-[var(--message-assistant-border)]">Agent: {value()}</span>}</Show>
+                    <Show when={modelIdentifier()}>{(value) => <span class="font-medium text-[var(--message-assistant-border)]">Model: {value()}</span>}</Show>
+                  </span>
+                </Show>
+              </div>
+            </Show>
+          </div>
+          <div class="flex items-center gap-2">
+
           <Show when={isUser() && props.onRevert}>
             <button
               class="bg-transparent border border-[var(--border-base)] text-[var(--text-muted)] cursor-pointer px-3 py-0.5 rounded text-xs font-semibold leading-none transition-all duration-200 flex items-center justify-center h-6 hover:bg-[var(--surface-hover)] hover:border-[var(--accent-primary)] hover:text-[var(--accent-primary)] active:scale-95"
diff --git a/packages/ui/src/components/message-stream-v2.tsx b/packages/ui/src/components/message-stream-v2.tsx
index 9a0dcf08..4d76edb8 100644
--- a/packages/ui/src/components/message-stream-v2.tsx
+++ b/packages/ui/src/components/message-stream-v2.tsx
@@ -3,6 +3,7 @@ import MessageItem from "./message-item"
 import ToolCall from "./tool-call"
 import Kbd from "./kbd"
 import type { MessageInfo, ClientPart } from "../types/message"
+import { partHasRenderableText } from "../types/message"
 import { getSessionInfo, sessions, setActiveParentSession, setActiveSession } from "../stores/sessions"
 import { showCommandPalette } from "../stores/command-palette"
 import { messageStoreBus } from "../stores/message-v2/bus"
@@ -120,6 +121,7 @@ interface ContentDisplayItem {
   parts: ClientPart[]
   messageInfo?: MessageInfo
   isQueued: boolean
+  showAgentMeta?: boolean
 }
 
 interface ToolDisplayItem {
@@ -144,6 +146,7 @@ type ReasoningDisplayItem = {
   key: string
   part: ClientPart
   messageInfo?: MessageInfo
+  showAgentMeta?: boolean
 }
 
 type MessageBlockItem = ContentDisplayItem | ToolDisplayItem | StepDisplayItem | ReasoningDisplayItem
@@ -280,11 +283,16 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
       const items: MessageBlockItem[] = []
       let segmentIndex = 0
       let pendingParts: ClientPart[] = []
+      let agentMetaAttached = record.role !== "assistant"
 
       const flushContent = () => {
         if (pendingParts.length === 0) return
         const segmentKey = makeInstanceCacheKey(instanceId, `${record.id}:segment:${segmentIndex}`)
         segmentIndex += 1
+        const shouldShowAgentMeta =
+          record.role === "assistant" &&
+          !agentMetaAttached &&
+          pendingParts.some((part) => partHasRenderableText(part))
         let cached = messageItemCache.get(segmentKey)
         if (!cached) {
           cached = {
@@ -294,6 +302,7 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
             parts: pendingParts.slice(),
             messageInfo,
             isQueued,
+            showAgentMeta: shouldShowAgentMeta,
           }
           messageItemCache.set(segmentKey, cached)
         } else {
@@ -301,6 +310,10 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
           cached.parts = pendingParts.slice()
           cached.messageInfo = messageInfo
           cached.isQueued = isQueued
+          cached.showAgentMeta = shouldShowAgentMeta
+        }
+        if (shouldShowAgentMeta) {
+          agentMetaAttached = true
         }
         items.push(cached)
         usedMessageKeys.add(segmentKey)
@@ -341,8 +354,6 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
 
         if (part.type === "step-start") {
           flushContent()
-          const key = makeInstanceCacheKey(instanceId, `${record.id}:${part.id ?? partIndex}:${part.type}`)
-          items.push({ type: part.type, key, part, messageInfo })
           return
         }
 
@@ -359,7 +370,11 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
           flushContent()
           if (showThinking && reasoningHasRenderableContent(part)) {
             const key = makeInstanceCacheKey(instanceId, `${record.id}:${part.id ?? partIndex}:reasoning`)
-            items.push({ type: "reasoning", key, part, messageInfo })
+            const showAgentMeta = record.role === "assistant" && !agentMetaAttached
+            if (showAgentMeta) {
+              agentMetaAttached = true
+            }
+            items.push({ type: "reasoning", key, part, messageInfo, showAgentMeta })
           }
           return
         }
@@ -622,6 +637,7 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
                         instanceId={props.instanceId}
                         sessionId={props.sessionId}
                         isQueued={(item as ContentDisplayItem).isQueued}
+                        showAgentMeta={(item as ContentDisplayItem).showAgentMeta}
                         onRevert={props.onRevert}
                         onFork={props.onFork}
                       />
@@ -699,6 +715,7 @@ export default function MessageStreamV2(props: MessageStreamV2Props) {
                         messageInfo={(item as ReasoningDisplayItem).messageInfo}
                         instanceId={props.instanceId}
                         sessionId={props.sessionId}
+                        showAgentMeta={(item as ReasoningDisplayItem).showAgentMeta}
                       />
                     </Match>
                   </Switch>
@@ -868,6 +885,7 @@ interface ReasoningCardProps {
   messageInfo?: MessageInfo
   instanceId: string
   sessionId: string
+  showAgentMeta?: boolean
 }
 
 function ReasoningCard(props: ReasoningCardProps) {
@@ -879,6 +897,21 @@ function ReasoningCard(props: ReasoningCardProps) {
     return date.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" })
   }
 
+  const agentIdentifier = () => {
+    const info = props.messageInfo
+    if (!info || info.role !== "assistant") return ""
+    return info.mode || ""
+  }
+
+  const modelIdentifier = () => {
+    const info = props.messageInfo
+    if (!info || info.role !== "assistant") return ""
+    const modelID = info.modelID || ""
+    const providerID = info.providerID || ""
+    if (modelID && providerID) return `${providerID}/${modelID}`
+    return modelID
+  }
+
   const reasoningText = () => {
     const part = props.part as any
     if (!part) return ""
@@ -925,7 +958,15 @@ function ReasoningCard(props: ReasoningCardProps) {
         aria-expanded={expanded()}
         aria-label={expanded() ? "Collapse thinking" : "Expand thinking"}
       >
-        <span class="message-reasoning-label">Thinking</span>
+        <span class="message-reasoning-label flex flex-wrap items-center gap-2">
+          <span>Thinking</span>
+          <Show when={props.showAgentMeta && (agentIdentifier() || modelIdentifier())}>
+            <span class="message-step-meta-inline">
+              <Show when={agentIdentifier()}>{(value) => <span class="font-medium text-[var(--message-assistant-border)]">Agent: {value()}</span>}</Show>
+              <Show when={modelIdentifier()}>{(value) => <span class="font-medium text-[var(--message-assistant-border)]">Model: {value()}</span>}</Show>
+            </span>
+          </Show>
+        </span>
         <span class="message-reasoning-meta">
           <span class="message-reasoning-indicator">{expanded() ? "Hide" : "View"}</span>
           <span class="message-reasoning-time">{timestamp()}</span>
diff --git a/packages/ui/src/styles/messaging/message-base.css b/packages/ui/src/styles/messaging/message-base.css
index 9c6b4459..7a1ec9e4 100644
--- a/packages/ui/src/styles/messaging/message-base.css
+++ b/packages/ui/src/styles/messaging/message-base.css
@@ -172,7 +172,7 @@
 }
 
 .message-step-meta-inline {
-  @apply flex flex-wrap gap-2 text-[11px];
+  @apply inline-flex flex-wrap items-center gap-2 text-[11px] font-medium;
   color: var(--message-assistant-border);
 }