From f56a0d4b97ea1b9afcaa92154f8069acdedc2684 Mon Sep 17 00:00:00 2001
From: scotttong <scott.tong@gmail.com>
Date: Thu, 30 Jul 2026 16:04:32 -0700
Subject: [PATCH] fix(ui): task-chat agent header uses the agent's assigned
 icon avatar

Revert the textless gradient capsule in TaskChatBubble to the legacy
Avatar + AgentIcon treatment: the adapter now maps the author agent's
assigned icon (agentMap icon) onto the message item, with an initials
fallback when no icon is set.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../components/task-chat/TaskChatBubble.tsx   | 35 ++++++++++---------
 .../components/task-chat/task-chat-adapter.ts | 16 ++-------
 .../task-chat/task-chat-fixtures.ts           |  4 +--
 .../components/task-chat/task-chat-model.ts   |  4 +--
 4 files changed, 26 insertions(+), 33 deletions(-)

diff --git a/ui/src/components/task-chat/TaskChatBubble.tsx b/ui/src/components/task-chat/TaskChatBubble.tsx
index 152d14100af..db353504ab9 100644
--- a/ui/src/components/task-chat/TaskChatBubble.tsx
+++ b/ui/src/components/task-chat/TaskChatBubble.tsx
@@ -1,26 +1,26 @@
 import { cn } from "@/lib/utils";
 import { MarkdownBody } from "@/components/MarkdownBody";
+import { Avatar, AvatarFallback } from "@/components/ui/avatar";
+import { AgentIcon } from "@/components/AgentIconPicker";
 import type { TaskChatMessageItem } from "./task-chat-model";
 
 interface TaskChatBubbleProps {
   item: TaskChatMessageItem;
 }
 
-/** Brand agent-gradient count (mirrors AgentCapsule's token pairs in index.css). */
-const GRADIENT_COUNT = 10;
-
-function capsuleGradient(index: number | undefined): string {
-  const n = Math.trunc(index ?? 1);
-  const idx = ((((n - 1) % GRADIENT_COUNT) + GRADIENT_COUNT) % GRADIENT_COUNT) + 1;
-  return `linear-gradient(to bottom, var(--agent-${idx}a), var(--agent-${idx}b))`;
+function initialsForName(name: string) {
+  const parts = name.trim().split(/\s+/);
+  if (parts.length >= 2) {
+    return (parts[0][0] + parts[parts.length - 1][0]).toUpperCase();
+  }
+  return name.slice(0, 2).toUpperCase();
 }
 
 /**
  * Author-typed message row — the primary legibility signal. Human messages sit
  * right in a solid accent bubble; agent messages sit left in a neutral card
- * bubble with a capsule-avatar author header (textless vertical brand-gradient
- * capsule + name · mode chip, per the v3/v6 prototype decisions); system
- * notices are centered and recede.
+ * bubble with an avatar author header (the agent's assigned icon + name · mode
+ * chip); system notices are centered and recede.
  */
 export function TaskChatBubble({ item }: TaskChatBubbleProps) {
   if (item.author === "system") {
@@ -36,12 +36,15 @@ export function TaskChatBubble({ item }: TaskChatBubbleProps) {
     <div className={cn("tc-enter-bubble group flex w-full flex-col gap-1", isHuman ? "items-end" : "items-start")}>
       {!isHuman && item.authorName ? (
         <span className="flex items-center gap-2 px-1">
-          <span
-            aria-hidden="true"
-            data-testid="task-chat-agent-capsule"
-            className="inline-block h-6 w-2.5 shrink-0 rounded-full"
-            style={{ background: capsuleGradient(item.agentGradient) }}
-          />
+          <Avatar size="sm" className="shrink-0" data-testid="task-chat-agent-avatar">
+            {item.agentIcon ? (
+              <AvatarFallback>
+                <AgentIcon icon={item.agentIcon} className="h-3.5 w-3.5" />
+              </AvatarFallback>
+            ) : (
+              <AvatarFallback>{initialsForName(item.authorName)}</AvatarFallback>
+            )}
+          </Avatar>
           <span className="text-sm font-semibold text-foreground">{item.authorName}</span>
           {item.modeLabel ? (
             <span className="rounded-full border border-border px-2 py-px text-(length:--text-micro) font-medium text-muted-foreground">
diff --git a/ui/src/components/task-chat/task-chat-adapter.ts b/ui/src/components/task-chat/task-chat-adapter.ts
index 4f0fd09cad8..106cfc634fc 100644
--- a/ui/src/components/task-chat/task-chat-adapter.ts
+++ b/ui/src/components/task-chat/task-chat-adapter.ts
@@ -27,16 +27,6 @@ function effectiveAgentId(comment: IssueChatComment): string | null {
   return comment.authorAgentId ?? comment.derivedAuthorAgentId ?? null;
 }
 
-/**
- * Deterministic agent-id → brand gradient index (1–10) for the capsule avatar,
- * so a given agent always renders the same capsule colors.
- */
-export function agentGradientIndex(agentId: string): number {
-  let hash = 0;
-  for (let i = 0; i < agentId.length; i++) hash = (hash * 31 + agentId.charCodeAt(i)) | 0;
-  return (Math.abs(hash) % 10) + 1;
-}
-
 function authorKind(comment: IssueChatComment): TaskChatAuthorKind {
   if (effectiveAgentId(comment)) return "agent";
   if (comment.authorType === "user") return "human";
@@ -60,11 +50,11 @@ export function commentsToTaskChatItems(
     if (comment.deletedAt) continue;
     const kind = authorKind(comment);
     let authorName: string | undefined;
-    let agentGradient: number | undefined;
+    let agentIcon: string | null | undefined;
     if (kind === "agent") {
       const agentId = effectiveAgentId(comment);
       authorName = (agentId && ctx.agentMap?.get(agentId)?.name) || "Agent";
-      agentGradient = agentId ? agentGradientIndex(agentId) : undefined;
+      agentIcon = agentId ? ctx.agentMap?.get(agentId)?.icon : undefined;
     } else if (kind === "human") {
       authorName =
         (comment.authorUserId && ctx.userLabelMap?.get(comment.authorUserId)) || undefined;
@@ -83,7 +73,7 @@ export function commentsToTaskChatItems(
       text: comment.body,
       timestamp: formatTimestamp(comment.createdAt),
       optimistic,
-      agentGradient,
+      agentIcon,
       modeLabel: kind === "agent" ? ctx.agentModeLabel : undefined,
     });
   }
diff --git a/ui/src/components/task-chat/task-chat-fixtures.ts b/ui/src/components/task-chat/task-chat-fixtures.ts
index 311f7cc28d1..ba5e08f6ec2 100644
--- a/ui/src/components/task-chat/task-chat-fixtures.ts
+++ b/ui/src/components/task-chat/task-chat-fixtures.ts
@@ -50,7 +50,7 @@ export function buildScenario(id: TaskChatStateId): TaskChatScenario {
         surface: "thread",
         items: [
           ...exchangePrefix(),
-          { id: "m-agent-1", kind: "message", author: "agent", authorName: AGENT, agentGradient: 3, modeLabel: "Agent mode", text: "On it — I'll add a token-bucket limiter and wire it into the login route.", timestamp: "2:31 PM" },
+          { id: "m-agent-1", kind: "message", author: "agent", authorName: AGENT, agentIcon: "bot", modeLabel: "Agent mode", text: "On it — I'll add a token-bucket limiter and wire it into the login route.", timestamp: "2:31 PM" },
         ],
       };
     case "thinking":
@@ -125,7 +125,7 @@ export function buildScenario(id: TaskChatStateId): TaskChatScenario {
               { id: "tool-done-2", kind: "tool", name: "Edit", target: "server/src/routes/auth.ts", toolKind: "edit", status: "completed", diff: { path: "server/src/routes/auth.ts", added: 34, removed: 3 } },
             ],
           },
-          { id: "m-done", kind: "message", author: "agent", authorName: AGENT, agentGradient: 3, modeLabel: "Agent mode", text: "Done — added a per-account token-bucket limiter and wired it into the login route. Tests pass.", timestamp: "2:34 PM" },
+          { id: "m-done", kind: "message", author: "agent", authorName: AGENT, agentIcon: "bot", modeLabel: "Agent mode", text: "Done — added a per-account token-bucket limiter and wired it into the login route. Tests pass.", timestamp: "2:34 PM" },
         ],
       };
     case "awaiting-approval":
diff --git a/ui/src/components/task-chat/task-chat-model.ts b/ui/src/components/task-chat/task-chat-model.ts
index c5008515009..f8c7909fe55 100644
--- a/ui/src/components/task-chat/task-chat-model.ts
+++ b/ui/src/components/task-chat/task-chat-model.ts
@@ -74,8 +74,8 @@ export interface TaskChatMessageItem {
    * chip in the agent header and under a sent human bubble (v6 decision).
    */
   modeLabel?: string;
-  /** Brand agent-gradient index (1–10) for the capsule avatar header. */
-  agentGradient?: number;
+  /** Assigned agent icon name (AgentIconName) for the avatar header. */
+  agentIcon?: string | null;
 }
 
 /** Collapsed chain-of-thought (ACP agent_thought_chunk). */