feat: improve chat ui messages (#132)

apps/web/hooks/use-lazy-load-messages.ts

@@ -1,8 +1,9 @@
-import { useCallback } from 'react';
+import { useCallback, useEffect, useRef } from 'react';
 import { listTaskSessionMessages } from '@/lib/api';
 import { useAppStore } from '@/components/state-provider';
 
 export function useLazyLoadMessages(sessionId: string | null) {
+  // Use refs for values that should not trigger callback recreation
   const hasMore = useAppStore((state) =>
     sessionId ? state.messages.metaBySession[sessionId]?.hasMore ?? false : false
   );
@@ -12,16 +13,24 @@ export function useLazyLoadMessages(sessionId: string | null) {
   const isLoading = useAppStore((state) =>
     sessionId ? state.messages.metaBySession[sessionId]?.isLoading ?? false : false
   );
+
+  // Store current values in refs to avoid recreating loadMore on every state change
+  const stateRef = useRef({ hasMore, oldestCursor, isLoading });
+  useEffect(() => {
+    stateRef.current = { hasMore, oldestCursor, isLoading };
+  }, [hasMore, oldestCursor, isLoading]);
+
   const prependMessages = useAppStore((state) => state.prependMessages);
   const setMessagesMetadata = useAppStore((state) => state.setMessagesMetadata);
 
+  // Stable loadMore - only depends on sessionId and store actions
   const loadMore = useCallback(async () => {
+    const { hasMore, isLoading, oldestCursor } = stateRef.current;
+
     if (!sessionId || !hasMore || isLoading || !oldestCursor) {
-      console.log('[useLazyLoadMessages] Cannot load more:', { sessionId, hasMore, isLoading, oldestCursor });
       return 0;
     }
 
-    console.log('[useLazyLoadMessages] Loading more messages before:', oldestCursor);
     setMessagesMetadata(sessionId, { isLoading: true });
     try {
       const response = await listTaskSessionMessages(sessionId, {
@@ -29,21 +38,9 @@ export function useLazyLoadMessages(sessionId: string | null) {
         before: oldestCursor,
         sort: 'desc',
       });
-      console.log('[useLazyLoadMessages] Received response:', {
-        count: response.messages?.length,
-        hasMore: response.has_more,
-        cursor: response.cursor,
-      });
       const orderedMessages = [...(response.messages ?? [])].reverse();
       // After reversing, orderedMessages[0] is the oldest message in this batch
       const newOldestCursor = orderedMessages[0]?.id ?? null;
-      console.log('[useLazyLoadMessages] Prepending messages:', {
-        count: orderedMessages.length,
-        newOldestCursor,
-        hasMore: response.has_more,
-        firstMsgId: orderedMessages[0]?.id,
-        lastMsgId: orderedMessages[orderedMessages.length - 1]?.id,
-      });
       prependMessages(sessionId, orderedMessages, {
         hasMore: response.has_more,
         oldestCursor: newOldestCursor,
@@ -54,7 +51,7 @@ export function useLazyLoadMessages(sessionId: string | null) {
       setMessagesMetadata(sessionId, { isLoading: false });
       return 0;
     }
-  }, [sessionId, hasMore, isLoading, oldestCursor, prependMessages, setMessagesMetadata]);
+  }, [sessionId, prependMessages, setMessagesMetadata]);
 
   return { loadMore, hasMore, isLoading };
 }

apps/web/hooks/use-processed-messages.ts

@@ -1,5 +1,27 @@
 import { useMemo } from 'react';
-import type { Message, ClarificationRequestMetadata } from '@/lib/types/http';
+import type { Message, ClarificationRequestMetadata, MessageType } from '@/lib/types/http';
+import type { ToolCallMetadata } from '@/components/task/chat/types';
+
+// Activity types that get grouped together
+const ACTIVITY_MESSAGE_TYPES: Set<MessageType> = new Set([
+  'thinking',
+  'tool_call',
+  'tool_edit',
+  'tool_read',
+  'tool_execute',
+  'tool_search',
+]);
+
+export type TurnGroup = {
+  type: 'turn_group';
+  id: string;
+  turnId: string | null;
+  messages: Message[];
+};
+
+export type RenderItem =
+  | { type: 'message'; message: Message }
+  | TurnGroup;
 
 export function useProcessedMessages(
   messages: Message[],
@@ -35,6 +57,33 @@ export function useProcessedMessages(
     return map;
   }, [messages]);
 
+  // Build a map of parent_tool_call_id → child messages for subagent nesting.
+  // When an agent sends parent_tool_use_id, child tools should be nested under the parent Task.
+  const childrenByParentToolCallId = useMemo(() => {
+    const map = new Map<string, Message[]>();
+    for (const message of messages) {
+      const metadata = message.metadata as ToolCallMetadata | undefined;
+      const parentId = metadata?.parent_tool_call_id;
+      if (parentId) {
+        const children = map.get(parentId) || [];
+        children.push(message);
+        map.set(parentId, children);
+      }
+    }
+    return map;
+  }, [messages]);
+
+  // Build a set of message IDs that are children of a subagent (should not render at top level)
+  const subagentChildIds = useMemo(() => {
+    const set = new Set<string>();
+    for (const children of childrenByParentToolCallId.values()) {
+      for (const child of children) {
+        set.add(child.id);
+      }
+    }
+    return set;
+  }, [childrenByParentToolCallId]);
+
   // Find pending clarification request (the most recent one that is still pending)
   // This will be rendered in the input area, not in the message list
   const pendingClarification = useMemo(() => {
@@ -55,8 +104,14 @@ export function useProcessedMessages(
   // Permission requests are only hidden if they have a matching tool_call in the messages
   // (so they can be displayed inline). Standalone permissions are shown as separate messages.
   // Pending clarification requests are also hidden (shown in input area instead).
+  // Child messages of subagents are hidden at top level (rendered nested under parent).
   const visibleMessages = useMemo(() => {
     return messages.filter((message) => {
+      // Hide child messages of subagents (they render nested under parent Task tool)
+      if (subagentChildIds.has(message.id)) {
+        return false;
+      }
+
       // Hide pending clarification request (it's shown in input area)
       if (message.type === 'clarification_request') {
         const metadata = message.metadata as ClarificationRequestMetadata | undefined;
@@ -71,6 +126,10 @@ export function useProcessedMessages(
         message.type === 'message' ||
         message.type === 'content' ||
         message.type === 'tool_call' ||
+        message.type === 'tool_read' ||
+        message.type === 'tool_edit' ||
+        message.type === 'tool_execute' ||
+        message.type === 'tool_search' ||
         message.type === 'progress' ||
         message.type === 'status' ||
         message.type === 'error' ||
@@ -90,7 +149,7 @@ export function useProcessedMessages(
 
       return false;
     });
-  }, [messages, toolCallIds]);
+  }, [messages, toolCallIds, subagentChildIds]);
 
   // Create a synthetic "user" message for the task description
   const taskDescriptionMessage: Message | null = useMemo(() => {
@@ -130,11 +189,69 @@ export function useProcessedMessages(
     return visibleMessages.filter((c) => c.author_type !== 'user').length;
   }, [visibleMessages]);
 
+  // Group consecutive activity messages (thinking + tool types) with same turn_id
+  // Only create groups if 2+ activity messages exist consecutively with the same turn_id
+  // Messages without turn_id are NOT grouped
+  const groupedItems = useMemo<RenderItem[]>(() => {
+    const items: RenderItem[] = [];
+    let currentGroup: Message[] = [];
+    let currentTurnId: string | null = null;
+
+    const flushGroup = () => {
+      if (currentGroup.length >= 2) {
+        // Create a turn group - use first message ID for unique ID
+        const groupTurnId = currentGroup[0].turn_id;
+        const firstMsgId = currentGroup[0].id;
+        items.push({
+          type: 'turn_group',
+          // Use first message ID to ensure unique group IDs even within same turn
+          id: `turn-group-${firstMsgId}`,
+          turnId: groupTurnId ?? null,
+          messages: currentGroup,
+        });
+      } else if (currentGroup.length === 1) {
+        // Single message, don't group
+        items.push({ type: 'message', message: currentGroup[0] });
+      }
+      currentGroup = [];
+      currentTurnId = null;
+    };
+
+    for (const message of allMessages) {
+      const isActivity = message.type && ACTIVITY_MESSAGE_TYPES.has(message.type);
+      const messageTurnId = message.turn_id ?? null;
+
+      // Only group activity messages that have an actual turn_id
+      if (isActivity && messageTurnId) {
+        // Check if this continues the current group (same turn_id)
+        if (currentGroup.length > 0 && currentTurnId === messageTurnId) {
+          currentGroup.push(message);
+        } else {
+          // Flush previous group and start new one
+          flushGroup();
+          currentGroup = [message];
+          currentTurnId = messageTurnId;
+        }
+      } else {
+        // Non-activity message or activity without turn_id breaks the group
+        flushGroup();
+        items.push({ type: 'message', message });
+      }
+    }
+
+    // Flush any remaining group
+    flushGroup();
+
+    return items;
+  }, [allMessages]);
+
   return {
     visibleMessages,
     allMessages,
+    groupedItems,
     toolCallIds,
     permissionsByToolCallId,
+    childrenByParentToolCallId,
     todoItems,
     agentMessageCount,
     pendingClarification,