feat: improve agent message tool rendering

Author: 0xHexE

backend/Cargo.lock

@@ -21,7 +21,8 @@ export function AnimatedCollapse({
 }: AnimatedCollapseProps) {
   const [mounted, setMounted] = useState(expanded);
   const [contentHeight, setContentHeight] = useState(0);
-  const height = useSharedValue(expanded ? 0 : 0);
+  const [initialRender, setInitialRender] = useState(expanded);
+  const height = useSharedValue(0);
   const opacity = useSharedValue(expanded ? 1 : 0);
 
   const targetHeight = useMemo(() => {
@@ -39,6 +40,13 @@ export function AnimatedCollapse({
   }, [expanded]);
 
   useEffect(() => {
+    if (initialRender && targetHeight > 0) {
+      height.value = targetHeight;
+      opacity.value = 1;
+      setInitialRender(false);
+      return;
+    }
+    if (initialRender) return;
     if (!mounted) return;
 
     if (expanded) {
@@ -73,7 +81,7 @@ export function AnimatedCollapse({
   if (!mounted) return null;
 
   return (
-    <Animated.View style={[styles.container, style]}>
+    <Animated.View style={[styles.container, initialRender ? undefined : style]}>
       <View onLayout={handleLayout} style={styles.content}>
         {children}
       </View>

features/agent/components/message-list/animated-collapse.tsx

@@ -22,11 +22,13 @@ import { ToolCallGroup } from "./tool-call";
 interface AssistantMessageProps {
   message: ChatMessage;
   isDark: boolean;
+  toolCallsOverride?: ChatMessage["toolCalls"];
 }
 
 export const AssistantMessage = memo(function AssistantMessage({
   message,
   isDark,
+  toolCallsOverride,
 }: AssistantMessageProps) {
   const colorScheme = useColorScheme() ?? "light";
   const colors = Colors[colorScheme];
@@ -37,9 +39,10 @@ export const AssistantMessage = memo(function AssistantMessage({
     message.isStreaming,
   );
 
+  const toolCalls = toolCallsOverride ?? message.toolCalls;
   const hasThinking = !!message.thinking;
   const hasText = !!message.text;
-  const hasToolCalls = !!(message.toolCalls && message.toolCalls.length > 0);
+  const hasToolCalls = !!(toolCalls && toolCalls.length > 0);
   const hasError = !!message.errorMessage;
   const isStreaming = !!message.isStreaming;
   const isThinkingOnly = hasThinking && !hasText && !hasToolCalls && isStreaming;
@@ -70,7 +73,7 @@ export const AssistantMessage = memo(function AssistantMessage({
       {hasText && <View style={styles.textBlock}>{elements}</View>}
 
       {hasToolCalls && (
-        <ToolCallGroup toolCalls={message.toolCalls!} isDark={isDark} />
+        <ToolCallGroup toolCalls={toolCalls!} isDark={isDark} />
       )}
 
       {hasError && (

features/agent/components/message-list/assistant-message.tsx

@@ -8,6 +8,8 @@ interface CodePreviewProps {
   maxHeight?: number;
   startLine?: number;
   language?: string;
+  diffLanguage?: string;
+  showLineNumbers?: boolean;
 }
 
 type TokenKind =
@@ -160,12 +162,25 @@ function tokenizeGeneric(line: string): Segment[] {
   return [{ text: line || " ", kind: "plain" }];
 }
 
-function tokenizeDiff(line: string): Segment[] {
+function tintSegments(segments: Segment[], kind: TokenKind): Segment[] {
+  return segments.map((segment) =>
+    segment.kind === "plain" || segment.kind === "punctuation" ? { ...segment, kind } : segment,
+  );
+}
+
+function tokenizeDiff(line: string, diffLanguage?: string): Segment[] {
   if (line.startsWith("@@") || line.startsWith("diff ") || line.startsWith("index ")) {
     return [{ text: line || " ", kind: "diffMeta" }];
   }
-  if (line.startsWith("+")) return [{ text: line || " ", kind: "diffAdd" }];
-  if (line.startsWith("-")) return [{ text: line || " ", kind: "diffRemove" }];
+  if (line.startsWith("+")) {
+    return [{ text: "+", kind: "diffAdd" }, ...tintSegments(tokenizeLine(line.slice(1), diffLanguage), "diffAdd")];
+  }
+  if (line.startsWith("-")) {
+    return [{ text: "-", kind: "diffRemove" }, ...tintSegments(tokenizeLine(line.slice(1), diffLanguage), "diffRemove")];
+  }
+  if (line.startsWith(" ")) {
+    return [{ text: " ", kind: "plain" }, ...tokenizeLine(line.slice(1), diffLanguage)];
+  }
   return tokenizeGeneric(line);
 }
 
@@ -252,6 +267,21 @@ function tokenizeBash(line: string): Segment[] {
   );
 }
 
+function tokenizeYaml(line: string): Segment[] {
+  return tokenizeWithPattern(
+    line,
+    /(#.*$|"(?:[^"\\]|\\.)*"|'(?:[^'\\]|\\.)*'|\b\d+(?:\.\d+)?\b|\btrue\b|\bfalse\b|\bnull\b|\b[A-Za-z0-9_.-]+(?=\s*:)|[{}\[\],:-])/g,
+    (value) => {
+      if (value.startsWith("#")) return "comment";
+      if (value.startsWith('"') || value.startsWith("'")) return "string";
+      if (/^\d/.test(value)) return "number";
+      if (/^(true|false|null)$/.test(value)) return "keyword";
+      if (/^[A-Za-z0-9_.-]+$/.test(value)) return "property";
+      return "punctuation";
+    },
+  );
+}
+
 function tokenizeMarkup(line: string): Segment[] {
   return tokenizeWithPattern(
     line,
@@ -268,7 +298,7 @@ function tokenizeMarkup(line: string): Segment[] {
   );
 }
 
-function tokenizeLine(line: string, language?: string): Segment[] {
+function tokenizeLine(line: string, language?: string, diffLanguage?: string): Segment[] {
   const lang = normalizeLanguage(language);
 
   if (["ts", "tsx", "js", "jsx", "typescript", "javascript"].includes(lang)) {
@@ -277,8 +307,9 @@ function tokenizeLine(line: string, language?: string): Segment[] {
   if (lang === "json") return tokenizeJson(line);
   if (["py", "python"].includes(lang)) return tokenizePython(line);
   if (["bash", "sh", "shell", "zsh"].includes(lang)) return tokenizeBash(line);
+  if (["yaml", "yml"].includes(lang)) return tokenizeYaml(line);
   if (["html", "htm", "xml", "xhtml", "svg"].includes(lang)) return tokenizeMarkup(line);
-  if (["diff", "patch"].includes(lang)) return tokenizeDiff(line);
+  if (["diff", "patch"].includes(lang)) return tokenizeDiff(line, diffLanguage);
   return tokenizeGeneric(line);
 }
 
@@ -288,13 +319,15 @@ export const CodePreview = memo(function CodePreview({
   maxHeight,
   startLine = 1,
   language,
+  diffLanguage,
+  showLineNumbers = true,
 }: CodePreviewProps) {
   const colors = isDark ? Colors.dark : Colors.light;
   const tokenColors = useMemo(() => createTokenColors(isDark), [isDark]);
   const lines = useMemo(() => code.split("\n"), [code]);
 
   return (
-    <View style={[styles.container, { backgroundColor: colors.surfaceRaised, borderColor: colors.border }]}> 
+    <View style={[styles.container, { backgroundColor: colors.surfaceRaised, borderColor: colors.border }]}>
       <ScrollView
         style={maxHeight ? { maxHeight } : undefined}
         nestedScrollEnabled
@@ -303,15 +336,17 @@ export const CodePreview = memo(function CodePreview({
         <ScrollView horizontal showsHorizontalScrollIndicator={false}>
           <View>
             {lines.map((line, i) => {
-              const segments = tokenizeLine(line, language);
+              const segments = tokenizeLine(line, language, diffLanguage);
               return (
                 <View key={i} style={styles.row}>
-                  <View style={[styles.lineNoCol, { borderRightColor: colors.border }]}> 
-                    <Text style={[styles.lineNo, { color: colors.textTertiary }]}>
-                      {startLine + i}
-                    </Text>
-                  </View>
-                  <Text style={[styles.lineText, { color: tokenColors.plain }]}>
+                  {showLineNumbers ? (
+                    <View style={[styles.lineNoCol, { borderRightColor: colors.border }]}> 
+                      <Text style={[styles.lineNo, { color: colors.textTertiary }]}> 
+                        {startLine + i}
+                      </Text>
+                    </View>
+                  ) : null}
+                  <Text style={[styles.lineText, !showLineNumbers && styles.lineTextNoGutter, { color: tokenColors.plain }]}>
                     {segments.length ? segments.map((segment, idx) => (
                       <Text key={`${i}-${idx}`} style={{ color: tokenColors[segment.kind] }}>
                         {segment.text || (idx === 0 ? " " : "")}
@@ -357,4 +392,7 @@ const styles = StyleSheet.create({
     paddingHorizontal: 8,
     paddingVertical: 1,
   },
+  lineTextNoGutter: {
+    paddingLeft: 10,
+  },
 });

features/agent/components/message-list/code-preview.tsx

@@ -10,12 +10,12 @@ import {
   type NativeScrollEvent,
   type NativeSyntheticEvent,
 } from "react-native";
-import Animated, { FadeIn, FadeOut } from "react-native-reanimated";
+import Animated, { FadeIn, FadeOut, LinearTransition } from "react-native-reanimated";
 import { ArrowDown } from "lucide-react-native";
 import { useAgentSession } from "@pi-ui/client";
 import { Colors, Fonts } from "@/constants/theme";
 import { useColorScheme } from "@/hooks/use-color-scheme";
-import type { ChatMessage } from "../../types";
+import type { ChatMessage, ToolCallInfo } from "../../types";
 import { UserMessage } from "./user-message";
 import { AssistantMessage } from "./assistant-message";
 import { SystemMessage } from "./system-message";
@@ -24,6 +24,58 @@ interface MessageListProps {
   sessionId: string;
 }
 
+interface VisibleMessageItem {
+  key: string;
+  message: ChatMessage;
+  toolCalls?: ToolCallInfo[];
+  turnDurationMs?: number;
+}
+
+function mergeConsecutiveToolCalls(
+  messages: ChatMessage[],
+  turnDurations: Map<string, number>,
+): VisibleMessageItem[] {
+  const visible: VisibleMessageItem[] = [];
+  let anchor: VisibleMessageItem | null = null;
+
+  for (let index = 0; index < messages.length; index++) {
+    const msg = messages[index]!;
+    const hasText =
+      msg.text.length > 0 ||
+      (!!msg.errorMessage && msg.errorMessage.length > 0) ||
+      (!!msg.thinking && msg.thinking.length > 0);
+    const toolCalls = msg.toolCalls?.length ? msg.toolCalls : undefined;
+    const turnDurationMs = turnDurations.get(msg.id);
+
+    if (msg.role === "user" || msg.role === "system") {
+      anchor = null;
+      visible.push({ key: msg.id, message: msg, turnDurationMs });
+      continue;
+    }
+
+    if (hasText || !anchor || msg.isStreaming) {
+      const item: VisibleMessageItem = {
+        key: msg.id,
+        message: msg,
+        toolCalls,
+        turnDurationMs,
+      };
+      anchor = msg.isStreaming ? null : item;
+      visible.push(item);
+      continue;
+    }
+
+    if (toolCalls?.length && anchor) {
+      anchor.toolCalls = anchor.toolCalls?.length
+        ? [...anchor.toolCalls, ...toolCalls]
+        : [...toolCalls];
+      anchor.turnDurationMs = anchor.turnDurationMs ?? turnDurationMs;
+    }
+  }
+
+  return visible;
+}
+
 const SCROLL_THRESHOLD = 200;
 const INITIAL_RENDER_COUNT = 8;
 const RENDER_BATCH_COUNT = 4;
@@ -43,7 +95,6 @@ export const MessageList = memo(function MessageList({
   const messages = session.messages as ChatMessage[];
   const isStreaming = session.isStreaming;
 
-  const reversed = useMemo(() => [...messages].reverse(), [messages]);
   const prevMessageCountRef = useRef(messages.length);
 
   const turnDurations = useMemo(() => {
@@ -61,6 +112,12 @@ export const MessageList = memo(function MessageList({
     return map;
   }, [messages]);
 
+  const visibleItems = useMemo(
+    () => mergeConsecutiveToolCalls(messages, turnDurations),
+    [messages, turnDurations],
+  );
+  const reversed = useMemo(() => [...visibleItems].reverse(), [visibleItems]);
+
   useEffect(() => {
     if (!autoFollow) return;
     const countChanged = messages.length !== prevMessageCountRef.current;
@@ -102,17 +159,18 @@ export const MessageList = memo(function MessageList({
   }, [session]);
 
   const renderItem = useCallback(
-    ({ item }: ListRenderItemInfo<ChatMessage>) => (
+    ({ item }: ListRenderItemInfo<VisibleMessageItem>) => (
       <MessageItem
-        message={item}
+        message={item.message}
+        toolCalls={item.toolCalls}
         isDark={isDark}
-        turnDurationMs={turnDurations.get(item.id)}
+        turnDurationMs={item.turnDurationMs}
       />
     ),
-    [isDark, turnDurations],
+    [isDark],
   );
 
-  const keyExtractor = useCallback((item: ChatMessage) => item.id, []);
+  const keyExtractor = useCallback((item: VisibleMessageItem) => item.key, []);
 
   const listFooter = (
     <View style={styles.historyLoaderWrap}>
@@ -138,7 +196,7 @@ export const MessageList = memo(function MessageList({
     <View style={styles.root}>
       <FlatList
         ref={listRef}
-        data={reversed}
+        data={reversed as VisibleMessageItem[]}
         renderItem={renderItem}
         keyExtractor={keyExtractor}
         inverted
@@ -212,10 +270,12 @@ const TurnDivider = memo(function TurnDivider({
 
 const MessageItem = memo(function MessageItem({
   message,
+  toolCalls,
   isDark,
   turnDurationMs,
 }: {
   message: ChatMessage;
+  toolCalls?: ToolCallInfo[];
   isDark: boolean;
   turnDurationMs?: number;
 }) {
@@ -224,7 +284,7 @@ const MessageItem = memo(function MessageItem({
       case "user":
         return <UserMessage message={message} isDark={isDark} />;
       case "assistant":
-        return <AssistantMessage message={message} isDark={isDark} />;
+        return <AssistantMessage message={message} toolCallsOverride={toolCalls} isDark={isDark} />;
       case "system":
         return <SystemMessage message={message} isDark={isDark} />;
       default:
@@ -233,12 +293,17 @@ const MessageItem = memo(function MessageItem({
   })();
 
   return (
-    <View style={styles.itemWrap}>
+    <Animated.View
+      layout={LinearTransition.duration(180)}
+      entering={FadeIn.duration(160)}
+      exiting={FadeOut.duration(140)}
+      style={styles.itemWrap}
+    >
       {content}
       {typeof turnDurationMs === "number" && turnDurationMs > 0 && (
         <TurnDivider durationMs={turnDurationMs} isDark={isDark} />
       )}
-    </View>
+    </Animated.View>
   );
 });