From fcba8b86d5e42a3959c216f64575830431e53e54 Mon Sep 17 00:00:00 2001
From: Felix <24791380+vcfgv@users.noreply.github.com>
Date: Wed, 11 Feb 2026 17:10:53 +0800
Subject: [PATCH] fix(chat): improve message handling, fix type errors and
 migrate changes to enhance branch (#50)

---
 src/pages/Chat/ChatMessage.tsx  |  90 ++++++++--
 src/pages/Chat/ChatToolbar.tsx  |   4 +-
 src/pages/Chat/index.tsx        |  49 ++++--
 src/pages/Chat/message-utils.ts |  27 +--
 src/stores/chat.ts              | 291 +++++++++++++++++++++++++++++---
 5 files changed, 395 insertions(+), 66 deletions(-)

diff --git a/src/pages/Chat/ChatMessage.tsx b/src/pages/Chat/ChatMessage.tsx
index d84c09cf4..2fb2b43d9 100644
--- a/src/pages/Chat/ChatMessage.tsx
+++ b/src/pages/Chat/ChatMessage.tsx
@@ -16,25 +16,38 @@ interface ChatMessageProps {
   message: RawMessage;
   showThinking: boolean;
   isStreaming?: boolean;
+  streamingTools?: Array<{
+    id?: string;
+    toolCallId?: string;
+    name: string;
+    status: 'running' | 'completed' | 'error';
+    durationMs?: number;
+    summary?: string;
+  }>;
 }
 
 export const ChatMessage = memo(function ChatMessage({
   message,
   showThinking,
   isStreaming = false,
+  streamingTools = [],
 }: ChatMessageProps) {
   const isUser = message.role === 'user';
-  const isToolResult = message.role === 'toolresult';
+  const role = typeof message.role === 'string' ? message.role.toLowerCase() : '';
+  const isToolResult = role === 'toolresult' || role === 'tool_result';
   const text = extractText(message);
+  const hasText = text.trim().length > 0;
   const thinking = extractThinking(message);
   const images = extractImages(message);
   const tools = extractToolUse(message);
+  const visibleThinking = showThinking ? thinking : null;
+  const visibleTools = showThinking ? tools : [];
 
-  // Don't render empty tool results when thinking is hidden
-  if (isToolResult && !showThinking) return null;
+  // Never render tool result messages in chat UI
+  if (isToolResult) return null;
 
   // Don't render empty messages
-  if (!text && !thinking && images.length === 0 && tools.length === 0) return null;
+  if (!hasText && !visibleThinking && images.length === 0 && visibleTools.length === 0) return null;
 
   return (
     <div
@@ -56,23 +69,32 @@ export const ChatMessage = memo(function ChatMessage({
       </div>
 
       {/* Content */}
-      <div className={cn('max-w-[80%] space-y-2', isUser && 'items-end')}>
+      <div
+        className={cn(
+          'flex flex-col w-full max-w-[80%] space-y-2',
+          isUser ? 'items-end' : 'items-start',
+        )}
+      >
+        {showThinking && isStreaming && !isUser && streamingTools.length > 0 && (
+          <ToolStatusBar tools={streamingTools} />
+        )}
+
         {/* Thinking section */}
-        {showThinking && thinking && (
-          <ThinkingBlock content={thinking} />
+        {visibleThinking && (
+          <ThinkingBlock content={visibleThinking} />
         )}
 
         {/* Tool use cards */}
-        {showThinking && tools.length > 0 && (
+        {visibleTools.length > 0 && (
           <div className="space-y-1">
-            {tools.map((tool, i) => (
+            {visibleTools.map((tool, i) => (
               <ToolCard key={tool.id || i} name={tool.name} input={tool.input} />
             ))}
           </div>
         )}
 
         {/* Main text bubble */}
-        {text && (
+        {hasText && (
           <MessageBubble
             text={text}
             isUser={isUser}
@@ -99,6 +121,51 @@ export const ChatMessage = memo(function ChatMessage({
   );
 });
 
+function formatDuration(durationMs?: number): string | null {
+  if (!durationMs || !Number.isFinite(durationMs)) return null;
+  if (durationMs < 1000) return `${Math.round(durationMs)}ms`;
+  return `${(durationMs / 1000).toFixed(1)}s`;
+}
+
+function ToolStatusBar({
+  tools,
+}: {
+  tools: Array<{
+    id?: string;
+    toolCallId?: string;
+    name: string;
+    status: 'running' | 'completed' | 'error';
+    durationMs?: number;
+    summary?: string;
+  }>;
+}) {
+  return (
+    <div className="w-full rounded-lg border border-border/50 bg-muted/20 px-3 py-2 text-xs text-muted-foreground">
+      <div className="space-y-1">
+        {tools.map((tool) => {
+          const duration = formatDuration(tool.durationMs);
+          const statusLabel = tool.status === 'running' ? 'running' : (tool.status === 'error' ? 'error' : 'done');
+          return (
+            <div key={tool.toolCallId || tool.id || tool.name} className="flex flex-wrap items-center gap-2">
+              <span className={cn(
+                'inline-flex items-center gap-1 rounded-full px-2 py-0.5 text-[11px]',
+                tool.status === 'error' ? 'bg-destructive/10 text-destructive' : 'bg-foreground/5 text-muted-foreground',
+              )}>
+                <span className="font-mono">{tool.name}</span>
+                <span className="opacity-70">{statusLabel}</span>
+              </span>
+              {duration && <span className="text-[11px] opacity-70">{duration}</span>}
+              {tool.summary && (
+                <span className="truncate text-[11px]">{tool.summary}</span>
+              )}
+            </div>
+          );
+        })}
+      </div>
+    </div>
+  );
+}
+
 // ── Message Bubble ──────────────────────────────────────────────
 
 function MessageBubble({
@@ -124,6 +191,7 @@ function MessageBubble({
     <div
       className={cn(
         'relative rounded-2xl px-4 py-3',
+        !isUser && 'w-full',
         isUser
           ? 'bg-primary text-primary-foreground'
           : 'bg-muted',
@@ -205,7 +273,7 @@ function ThinkingBlock({ content }: { content: string }) {
   const [expanded, setExpanded] = useState(false);
 
   return (
-    <div className="rounded-lg border border-border/50 bg-muted/30 text-sm">
+    <div className="w-full rounded-lg border border-border/50 bg-muted/30 text-sm">
       <button
         className="flex items-center gap-2 w-full px-3 py-2 text-muted-foreground hover:text-foreground transition-colors"
         onClick={() => setExpanded(!expanded)}
diff --git a/src/pages/Chat/ChatToolbar.tsx b/src/pages/Chat/ChatToolbar.tsx
index 7c777d5cb..3f34b1ba1 100644
--- a/src/pages/Chat/ChatToolbar.tsx
+++ b/src/pages/Chat/ChatToolbar.tsx
@@ -38,12 +38,12 @@ export function ChatToolbar() {
           {/* Render all sessions; if currentSessionKey is not in the list, add it */}
           {!sessions.some((s) => s.key === currentSessionKey) && (
             <option value={currentSessionKey}>
-              {currentSessionKey === 'main' ? 'main' : currentSessionKey}
+              {currentSessionKey}
             </option>
           )}
           {sessions.map((s) => (
             <option key={s.key} value={s.key}>
-              {s.displayName || s.label || s.key}
+              {s.key}
             </option>
           ))}
         </select>
diff --git a/src/pages/Chat/index.tsx b/src/pages/Chat/index.tsx
index 7a4add702..32a357978 100644
--- a/src/pages/Chat/index.tsx
+++ b/src/pages/Chat/index.tsx
@@ -7,13 +7,13 @@
 import { useEffect, useRef, useState } from 'react';
 import { AlertCircle, Bot, MessageSquare, Sparkles } from 'lucide-react';
 import { Card, CardContent } from '@/components/ui/card';
-import { useChatStore } from '@/stores/chat';
+import { useChatStore, type RawMessage } from '@/stores/chat';
 import { useGatewayStore } from '@/stores/gateway';
 import { LoadingSpinner } from '@/components/common/LoadingSpinner';
 import { ChatMessage } from './ChatMessage';
 import { ChatInput } from './ChatInput';
 import { ChatToolbar } from './ChatToolbar';
-import { extractText } from './message-utils';
+import { extractImages, extractText, extractThinking, extractToolUse } from './message-utils';
 import { useTranslation } from 'react-i18next';
 
 export function Chat() {
@@ -27,6 +27,7 @@ export function Chat() {
   const error = useChatStore((s) => s.error);
   const showThinking = useChatStore((s) => s.showThinking);
   const streamingMessage = useChatStore((s) => s.streamingMessage);
+  const streamingTools = useChatStore((s) => s.streamingTools);
   const loadHistory = useChatStore((s) => s.loadHistory);
   const loadSessions = useChatStore((s) => s.loadSessions);
   const sendMessage = useChatStore((s) => s.sendMessage);
@@ -38,10 +39,16 @@ export function Chat() {
 
   // Load data when gateway is running
   useEffect(() => {
-    if (isGatewayRunning) {
-      loadHistory();
-      loadSessions();
-    }
+    if (!isGatewayRunning) return;
+    let cancelled = false;
+    (async () => {
+      await loadSessions();
+      if (cancelled) return;
+      await loadHistory();
+    })();
+    return () => {
+      cancelled = true;
+    };
   }, [isGatewayRunning, loadHistory, loadSessions]);
 
   // Auto-scroll on new messages or streaming
@@ -78,6 +85,14 @@ export function Chat() {
     : null;
   const streamText = streamMsg ? extractText(streamMsg) : (typeof streamingMessage === 'string' ? streamingMessage : '');
   const hasStreamText = streamText.trim().length > 0;
+  const streamThinking = streamMsg ? extractThinking(streamMsg) : null;
+  const hasStreamThinking = showThinking && !!streamThinking && streamThinking.trim().length > 0;
+  const streamTools = streamMsg ? extractToolUse(streamMsg) : [];
+  const hasStreamTools = showThinking && streamTools.length > 0;
+  const streamImages = streamMsg ? extractImages(streamMsg) : [];
+  const hasStreamImages = streamImages.length > 0;
+  const hasStreamToolStatus = showThinking && streamingTools.length > 0;
+  const shouldRenderStreaming = sending && (hasStreamText || hasStreamThinking || hasStreamTools || hasStreamImages || hasStreamToolStatus);
 
   return (
     <div className="flex flex-col -m-6" style={{ height: 'calc(100vh - 2.5rem)' }}>
@@ -106,20 +121,28 @@ export function Chat() {
               ))}
 
               {/* Streaming message */}
-              {sending && hasStreamText && (
+              {shouldRenderStreaming && (
                 <ChatMessage
-                  message={{
-                    role: 'assistant',
-                    content: streamMsg?.content ?? streamText,
-                    timestamp: streamMsg?.timestamp ?? streamingTimestamp,
-                  }}
+                  message={(streamMsg
+                    ? {
+                        ...(streamMsg as Record<string, unknown>),
+                        role: (typeof streamMsg.role === 'string' ? streamMsg.role : 'assistant') as RawMessage['role'],
+                        content: streamMsg.content ?? streamText,
+                        timestamp: streamMsg.timestamp ?? streamingTimestamp,
+                      }
+                    : {
+                        role: 'assistant',
+                        content: streamText,
+                        timestamp: streamingTimestamp,
+                      }) as RawMessage}
                   showThinking={showThinking}
                   isStreaming
+                  streamingTools={streamingTools}
                 />
               )}
 
               {/* Typing indicator when sending but no stream yet */}
-              {sending && !hasStreamText && (
+              {sending && !hasStreamText && !hasStreamThinking && !hasStreamTools && !hasStreamImages && !hasStreamToolStatus && (
                 <TypingIndicator />
               )}
             </>
diff --git a/src/pages/Chat/message-utils.ts b/src/pages/Chat/message-utils.ts
index e75804694..8419cd8ce 100644
--- a/src/pages/Chat/message-utils.ts
+++ b/src/pages/Chat/message-utils.ts
@@ -15,26 +15,25 @@ export function extractText(message: RawMessage | unknown): string {
   const content = msg.content;
 
   if (typeof content === 'string') {
-    return content;
+    return content.trim().length > 0 ? content : '';
   }
 
   if (Array.isArray(content)) {
     const parts: string[] = [];
     for (const block of content as ContentBlock[]) {
       if (block.type === 'text' && block.text) {
-        parts.push(block.text);
-      }
-      // tool_result blocks may have nested text
-      if (block.type === 'tool_result' && typeof block.content === 'string') {
-        parts.push(block.content);
+        if (block.text.trim().length > 0) {
+          parts.push(block.text);
+        }
       }
     }
-    return parts.join('\n\n');
+    const combined = parts.join('\n\n');
+    return combined.trim().length > 0 ? combined : '';
   }
 
   // Fallback: try .text field
   if (typeof msg.text === 'string') {
-    return msg.text;
+    return msg.text.trim().length > 0 ? msg.text : '';
   }
 
   return '';
@@ -54,11 +53,15 @@ export function extractThinking(message: RawMessage | unknown): string | null {
   const parts: string[] = [];
   for (const block of content as ContentBlock[]) {
     if (block.type === 'thinking' && block.thinking) {
-      parts.push(block.thinking);
+      const cleaned = block.thinking.trim();
+      if (cleaned) {
+        parts.push(cleaned);
+      }
     }
   }
 
-  return parts.length > 0 ? parts.join('\n\n') : null;
+  const combined = parts.join('\n\n').trim();
+  return combined.length > 0 ? combined : null;
 }
 
 /**
@@ -97,11 +100,11 @@ export function extractToolUse(message: RawMessage | unknown): Array<{ id: strin
 
   const tools: Array<{ id: string; name: string; input: unknown }> = [];
   for (const block of content as ContentBlock[]) {
-    if (block.type === 'tool_use' && block.name) {
+    if ((block.type === 'tool_use' || block.type === 'toolCall') && block.name) {
       tools.push({
         id: block.id || '',
         name: block.name,
-        input: block.input,
+        input: block.input ?? block.arguments,
       });
     }
   }
diff --git a/src/stores/chat.ts b/src/stores/chat.ts
index 648762f5d..e5f1fe7f2 100644
--- a/src/stores/chat.ts
+++ b/src/stores/chat.ts
@@ -14,17 +14,21 @@ export interface RawMessage {
   timestamp?: number;
   id?: string;
   toolCallId?: string;
+  toolName?: string;
+  details?: unknown;
+  isError?: boolean;
 }
 
 /** Content block inside a message */
 export interface ContentBlock {
-  type: 'text' | 'image' | 'thinking' | 'tool_use' | 'tool_result';
+  type: 'text' | 'image' | 'thinking' | 'tool_use' | 'tool_result' | 'toolCall' | 'toolResult';
   text?: string;
   thinking?: string;
   source?: { type: string; media_type: string; data: string };
   id?: string;
   name?: string;
   input?: unknown;
+  arguments?: unknown;
   content?: unknown;
 }
 
@@ -37,6 +41,16 @@ export interface ChatSession {
   model?: string;
 }
 
+export interface ToolStatus {
+  id?: string;
+  toolCallId?: string;
+  name: string;
+  status: 'running' | 'completed' | 'error';
+  durationMs?: number;
+  summary?: string;
+  updatedAt: number;
+}
+
 interface ChatState {
   // Messages
   messages: RawMessage[];
@@ -48,6 +62,7 @@ interface ChatState {
   activeRunId: string | null;
   streamingText: string;
   streamingMessage: unknown | null;
+  streamingTools: ToolStatus[];
   pendingFinal: boolean;
   lastUserMessageAt: number | null;
 
@@ -72,9 +87,20 @@ interface ChatState {
   clearError: () => void;
 }
 
+const DEFAULT_CANONICAL_PREFIX = 'agent:main';
+const DEFAULT_SESSION_KEY = `${DEFAULT_CANONICAL_PREFIX}:main`;
+
+function getCanonicalPrefixFromSessions(sessions: ChatSession[]): string | null {
+  const canonical = sessions.find((s) => s.key.startsWith('agent:'))?.key;
+  if (!canonical) return null;
+  const parts = canonical.split(':');
+  if (parts.length < 2) return null;
+  return `${parts[0]}:${parts[1]}`;
+}
+
 function isToolOnlyMessage(message: RawMessage | undefined): boolean {
   if (!message) return false;
-  if (message.role === 'toolresult') return true;
+  if (isToolResultRole(message.role)) return true;
 
   const content = message.content;
   if (!Array.isArray(content)) return false;
@@ -84,7 +110,7 @@ function isToolOnlyMessage(message: RawMessage | undefined): boolean {
   let hasNonToolContent = false;
 
   for (const block of content as ContentBlock[]) {
-    if (block.type === 'tool_use' || block.type === 'tool_result') {
+    if (block.type === 'tool_use' || block.type === 'tool_result' || block.type === 'toolCall' || block.type === 'toolResult') {
       hasTool = true;
       continue;
     }
@@ -100,6 +126,167 @@ function isToolOnlyMessage(message: RawMessage | undefined): boolean {
   return hasTool && !hasText && !hasNonToolContent;
 }
 
+function isToolResultRole(role: unknown): boolean {
+  if (!role) return false;
+  const normalized = String(role).toLowerCase();
+  return normalized === 'toolresult' || normalized === 'tool_result';
+}
+
+function extractTextFromContent(content: unknown): string {
+  if (typeof content === 'string') return content;
+  if (!Array.isArray(content)) return '';
+  const parts: string[] = [];
+  for (const block of content as ContentBlock[]) {
+    if (block.type === 'text' && block.text) {
+      parts.push(block.text);
+    }
+  }
+  return parts.join('\n');
+}
+
+function summarizeToolOutput(text: string): string | undefined {
+  const trimmed = text.trim();
+  if (!trimmed) return undefined;
+  const lines = trimmed.split(/\r?\n/).map((line) => line.trim()).filter(Boolean);
+  if (lines.length === 0) return undefined;
+  const summaryLines = lines.slice(0, 2);
+  let summary = summaryLines.join(' / ');
+  if (summary.length > 160) {
+    summary = `${summary.slice(0, 157)}...`;
+  }
+  return summary;
+}
+
+function normalizeToolStatus(rawStatus: unknown, fallback: 'running' | 'completed'): ToolStatus['status'] {
+  const status = typeof rawStatus === 'string' ? rawStatus.toLowerCase() : '';
+  if (status === 'error' || status === 'failed') return 'error';
+  if (status === 'completed' || status === 'success' || status === 'done') return 'completed';
+  return fallback;
+}
+
+function parseDurationMs(value: unknown): number | undefined {
+  if (typeof value === 'number' && Number.isFinite(value)) return value;
+  const parsed = typeof value === 'string' ? Number(value) : NaN;
+  return Number.isFinite(parsed) ? parsed : undefined;
+}
+
+function extractToolUseUpdates(message: unknown): ToolStatus[] {
+  if (!message || typeof message !== 'object') return [];
+  const msg = message as Record<string, unknown>;
+  const content = msg.content;
+  if (!Array.isArray(content)) return [];
+
+  const updates: ToolStatus[] = [];
+  for (const block of content as ContentBlock[]) {
+    if ((block.type !== 'tool_use' && block.type !== 'toolCall') || !block.name) continue;
+    updates.push({
+      id: block.id || block.name,
+      toolCallId: block.id,
+      name: block.name,
+      status: 'running',
+      updatedAt: Date.now(),
+    });
+  }
+
+  return updates;
+}
+
+function extractToolResultBlocks(message: unknown, eventState: string): ToolStatus[] {
+  if (!message || typeof message !== 'object') return [];
+  const msg = message as Record<string, unknown>;
+  const content = msg.content;
+  if (!Array.isArray(content)) return [];
+
+  const updates: ToolStatus[] = [];
+  for (const block of content as ContentBlock[]) {
+    if (block.type !== 'tool_result' && block.type !== 'toolResult') continue;
+    const outputText = extractTextFromContent(block.content ?? block.text ?? '');
+    const summary = summarizeToolOutput(outputText);
+    updates.push({
+      id: block.id || block.name || 'tool',
+      toolCallId: block.id,
+      name: block.name || block.id || 'tool',
+      status: normalizeToolStatus(undefined, eventState === 'delta' ? 'running' : 'completed'),
+      summary,
+      updatedAt: Date.now(),
+    });
+  }
+
+  return updates;
+}
+
+function extractToolResultUpdate(message: unknown, eventState: string): ToolStatus | null {
+  if (!message || typeof message !== 'object') return null;
+  const msg = message as Record<string, unknown>;
+  const role = typeof msg.role === 'string' ? msg.role.toLowerCase() : '';
+  if (!isToolResultRole(role)) return null;
+
+  const toolName = typeof msg.toolName === 'string' ? msg.toolName : (typeof msg.name === 'string' ? msg.name : '');
+  const toolCallId = typeof msg.toolCallId === 'string' ? msg.toolCallId : undefined;
+  const details = (msg.details && typeof msg.details === 'object') ? msg.details as Record<string, unknown> : undefined;
+  const rawStatus = (msg.status ?? details?.status);
+  const fallback = eventState === 'delta' ? 'running' : 'completed';
+  const status = normalizeToolStatus(rawStatus, fallback);
+  const durationMs = parseDurationMs(details?.durationMs ?? details?.duration ?? (msg as Record<string, unknown>).durationMs);
+
+  const outputText = (details && typeof details.aggregated === 'string')
+    ? details.aggregated
+    : extractTextFromContent(msg.content);
+  const summary = summarizeToolOutput(outputText) ?? summarizeToolOutput(String(details?.error ?? msg.error ?? ''));
+
+  const name = toolName || toolCallId || 'tool';
+  const id = toolCallId || name;
+
+  return {
+    id,
+    toolCallId,
+    name,
+    status,
+    durationMs,
+    summary,
+    updatedAt: Date.now(),
+  };
+}
+
+function mergeToolStatus(existing: ToolStatus['status'], incoming: ToolStatus['status']): ToolStatus['status'] {
+  const order: Record<ToolStatus['status'], number> = { running: 0, completed: 1, error: 2 };
+  return order[incoming] >= order[existing] ? incoming : existing;
+}
+
+function upsertToolStatuses(current: ToolStatus[], updates: ToolStatus[]): ToolStatus[] {
+  if (updates.length === 0) return current;
+  const next = [...current];
+  for (const update of updates) {
+    const key = update.toolCallId || update.id || update.name;
+    if (!key) continue;
+    const index = next.findIndex((tool) => (tool.toolCallId || tool.id || tool.name) === key);
+    if (index === -1) {
+      next.push(update);
+      continue;
+    }
+    const existing = next[index];
+    next[index] = {
+      ...existing,
+      ...update,
+      name: update.name || existing.name,
+      status: mergeToolStatus(existing.status, update.status),
+      durationMs: update.durationMs ?? existing.durationMs,
+      summary: update.summary ?? existing.summary,
+      updatedAt: update.updatedAt || existing.updatedAt,
+    };
+  }
+  return next;
+}
+
+function collectToolUpdates(message: unknown, eventState: string): ToolStatus[] {
+  const updates: ToolStatus[] = [];
+  const toolResultUpdate = extractToolResultUpdate(message, eventState);
+  if (toolResultUpdate) updates.push(toolResultUpdate);
+  updates.push(...extractToolResultBlocks(message, eventState));
+  updates.push(...extractToolUseUpdates(message));
+  return updates;
+}
+
 function hasNonToolAssistantContent(message: RawMessage | undefined): boolean {
   if (!message) return false;
   if (typeof message.content === 'string' && message.content.trim()) return true;
@@ -130,11 +317,12 @@ export const useChatStore = create<ChatState>((set, get) => ({
   activeRunId: null,
   streamingText: '',
   streamingMessage: null,
+  streamingTools: [],
   pendingFinal: false,
   lastUserMessageAt: null,
 
   sessions: [],
-  currentSessionKey: 'main',
+  currentSessionKey: DEFAULT_SESSION_KEY,
 
   showThinking: true,
   thinkingLevel: null,
@@ -160,33 +348,49 @@ export const useChatStore = create<ChatState>((set, get) => ({
           model: s.model ? String(s.model) : undefined,
         })).filter((s: ChatSession) => s.key);
 
-        // Normalize: the Gateway returns the main session with canonical key
-        // like "agent:main:main", but the frontend uses "main" for all RPC calls.
-        // Map the canonical main session key to "main" so the selector stays consistent.
-        const mainCanonicalPattern = /^agent:[^:]+:main$/;
-        const normalizedSessions = sessions.map((s) => {
-          if (mainCanonicalPattern.test(s.key)) {
-            return { ...s, key: 'main', displayName: s.displayName || 'main' };
+        const canonicalBySuffix = new Map<string, string>();
+        for (const session of sessions) {
+          if (!session.key.startsWith('agent:')) continue;
+          const parts = session.key.split(':');
+          if (parts.length < 3) continue;
+          const suffix = parts.slice(2).join(':');
+          if (suffix && !canonicalBySuffix.has(suffix)) {
+            canonicalBySuffix.set(suffix, session.key);
           }
-          return s;
-        });
+        }
 
-        // Deduplicate: if both "main" and "agent:X:main" existed, keep only one
+        // Deduplicate: if both short and canonical existed, keep canonical only
         const seen = new Set<string>();
-        const dedupedSessions = normalizedSessions.filter((s) => {
+        const dedupedSessions = sessions.filter((s) => {
+          if (!s.key.startsWith('agent:') && canonicalBySuffix.has(s.key)) return false;
           if (seen.has(s.key)) return false;
           seen.add(s.key);
           return true;
         });
 
-        set({ sessions: dedupedSessions });
-
-        // If currentSessionKey is 'main' and we now have sessions,
-        // ensure we stay on 'main' (no-op, but load history if needed)
         const { currentSessionKey } = get();
-        if (currentSessionKey === 'main' && !dedupedSessions.find((s) => s.key === 'main') && dedupedSessions.length > 0) {
-          // Main session not found at all — switch to the first available session
-          set({ currentSessionKey: dedupedSessions[0].key });
+        let nextSessionKey = currentSessionKey || DEFAULT_SESSION_KEY;
+        if (!nextSessionKey.startsWith('agent:')) {
+          const canonicalMatch = canonicalBySuffix.get(nextSessionKey);
+          if (canonicalMatch) {
+            nextSessionKey = canonicalMatch;
+          }
+        }
+        if (!dedupedSessions.find((s) => s.key === nextSessionKey) && dedupedSessions.length > 0) {
+          // Current session not found at all — switch to the first available session
+          nextSessionKey = dedupedSessions[0].key;
+        }
+
+        const sessionsWithCurrent = !dedupedSessions.find((s) => s.key === nextSessionKey) && nextSessionKey
+          ? [
+              ...dedupedSessions,
+              { key: nextSessionKey, displayName: nextSessionKey },
+            ]
+          : dedupedSessions;
+
+        set({ sessions: sessionsWithCurrent, currentSessionKey: nextSessionKey });
+
+        if (currentSessionKey !== nextSessionKey) {
           get().loadHistory();
         }
       }
@@ -203,6 +407,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
       messages: [],
       streamingText: '',
       streamingMessage: null,
+      streamingTools: [],
       activeRunId: null,
       error: null,
       pendingFinal: false,
@@ -216,12 +421,14 @@ export const useChatStore = create<ChatState>((set, get) => ({
 
   newSession: () => {
     // Generate a new unique session key and switch to it
-    const newKey = `session-${Date.now()}`;
+    const prefix = getCanonicalPrefixFromSessions(get().sessions) ?? DEFAULT_CANONICAL_PREFIX;
+    const newKey = `${prefix}:session-${Date.now()}`;
     set({
       currentSessionKey: newKey,
       messages: [],
       streamingText: '',
       streamingMessage: null,
+      streamingTools: [],
       activeRunId: null,
       error: null,
       pendingFinal: false,
@@ -247,11 +454,12 @@ export const useChatStore = create<ChatState>((set, get) => ({
       if (result.success && result.result) {
         const data = result.result;
         const rawMessages = Array.isArray(data.messages) ? data.messages as RawMessage[] : [];
+        const filteredMessages = rawMessages.filter((msg) => !isToolResultRole(msg.role));
         const thinkingLevel = data.thinkingLevel ? String(data.thinkingLevel) : null;
-        set({ messages: rawMessages, thinkingLevel, loading: false });
+        set({ messages: filteredMessages, thinkingLevel, loading: false });
         const { pendingFinal, lastUserMessageAt } = get();
         if (pendingFinal) {
-          const recentAssistant = [...rawMessages].reverse().find((msg) => {
+          const recentAssistant = [...filteredMessages].reverse().find((msg) => {
             if (msg.role !== 'assistant') return false;
             if (!hasNonToolAssistantContent(msg)) return false;
             if (lastUserMessageAt && msg.timestamp && msg.timestamp < lastUserMessageAt) return false;
@@ -291,6 +499,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
       error: null,
       streamingText: '',
       streamingMessage: null,
+      streamingTools: [],
       pendingFinal: false,
       lastUserMessageAt: userMsg.timestamp ?? null,
     }));
@@ -337,6 +546,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
   abortRun: async () => {
     const { currentSessionKey } = get();
     set({ sending: false, streamingText: '', streamingMessage: null, pendingFinal: false, lastUserMessageAt: null });
+    set({ streamingTools: [] });
 
     try {
       await window.electron.ipcRenderer.invoke(
@@ -362,19 +572,38 @@ export const useChatStore = create<ChatState>((set, get) => ({
     switch (eventState) {
       case 'delta': {
         // Streaming update - store the cumulative message
-        set({
-          streamingMessage: event.message ?? get().streamingMessage,
-        });
+        const updates = collectToolUpdates(event.message, eventState);
+        set((s) => ({
+          streamingMessage: (() => {
+            if (event.message && typeof event.message === 'object') {
+              const msgRole = (event.message as RawMessage).role;
+              if (isToolResultRole(msgRole)) return s.streamingMessage;
+            }
+            return event.message ?? s.streamingMessage;
+          })(),
+          streamingTools: updates.length > 0 ? upsertToolStatuses(s.streamingTools, updates) : s.streamingTools,
+        }));
         break;
       }
       case 'final': {
         // Message complete - add to history and clear streaming
         const finalMsg = event.message as RawMessage | undefined;
         if (finalMsg) {
+          const updates = collectToolUpdates(finalMsg, eventState);
+          if (isToolResultRole(finalMsg.role)) {
+            set((s) => ({
+              streamingText: '',
+              pendingFinal: true,
+              streamingTools: updates.length > 0 ? upsertToolStatuses(s.streamingTools, updates) : s.streamingTools,
+            }));
+            break;
+          }
           const toolOnly = isToolOnlyMessage(finalMsg);
           const hasOutput = hasNonToolAssistantContent(finalMsg);
           const msgId = finalMsg.id || (toolOnly ? `run-${runId}-tool-${Date.now()}` : `run-${runId}`);
           set((s) => {
+            const nextTools = updates.length > 0 ? upsertToolStatuses(s.streamingTools, updates) : s.streamingTools;
+            const streamingTools = hasOutput ? [] : nextTools;
             // Check if message already exists (prevent duplicates)
             const alreadyExists = s.messages.some(m => m.id === msgId);
             if (alreadyExists) {
@@ -383,12 +612,14 @@ export const useChatStore = create<ChatState>((set, get) => ({
                 streamingText: '',
                 streamingMessage: null,
                 pendingFinal: true,
+                streamingTools,
               } : {
                 streamingText: '',
                 streamingMessage: null,
                 sending: hasOutput ? false : s.sending,
                 activeRunId: hasOutput ? null : s.activeRunId,
                 pendingFinal: hasOutput ? false : true,
+                streamingTools,
               };
             }
             return toolOnly ? {
@@ -400,6 +631,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
               streamingText: '',
               streamingMessage: null,
               pendingFinal: true,
+              streamingTools,
             } : {
               messages: [...s.messages, {
                 ...finalMsg,
@@ -411,6 +643,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
               sending: hasOutput ? false : s.sending,
               activeRunId: hasOutput ? null : s.activeRunId,
               pendingFinal: hasOutput ? false : true,
+              streamingTools,
             };
           });
         } else {
@@ -428,6 +661,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
           activeRunId: null,
           streamingText: '',
           streamingMessage: null,
+          streamingTools: [],
           pendingFinal: false,
           lastUserMessageAt: null,
         });
@@ -439,6 +673,7 @@ export const useChatStore = create<ChatState>((set, get) => ({
           activeRunId: null,
           streamingText: '',
           streamingMessage: null,
+          streamingTools: [],
           pendingFinal: false,
           lastUserMessageAt: null,
         });