@townco/ui 0.1.109 → 0.1.110

package/dist/core/hooks/use-chat-messages.js

@@ -21,7 +21,7 @@ export function useChatMessages(client, startSession) {
     const addToolCallToCurrentMessage = useChatStore((state) => state.addToolCallToCurrentMessage);
     const updateToolCallInCurrentMessage = useChatStore((state) => state.updateToolCallInCurrentMessage);
     const addHookNotificationToCurrentMessage = useChatStore((state) => state.addHookNotificationToCurrentMessage);
-    const addSourcesToCurrentMessage = useChatStore((state) => state.addSourcesToCurrentMessage);
+    const addSourcesByToolCallId = useChatStore((state) => state.addSourcesByToolCallId);
     const truncateMessagesFrom = useChatStore((state) => state.truncateMessagesFrom);
     // Track the current assistant message ID for cancellation
     const currentAssistantMessageIdRef = useRef(null);
@@ -206,7 +206,7 @@ export function useChatMessages(client, startSession) {
                     // Sources chunk - citation sources from tool calls
                     logger.debug("Received sources chunk", { chunk });
                     // Add sources to current assistant message for citation rendering
-                    addSourcesToCurrentMessage(chunk.sources);
+                    addSourcesByToolCallId(chunk.sources);
                 }
             }
             // Ensure streaming state is cleared even if no explicit isComplete was received
@@ -259,7 +259,7 @@ export function useChatMessages(client, startSession) {
         addToolCallToCurrentMessage,
         updateToolCallInCurrentMessage,
         addHookNotificationToCurrentMessage,
-        addSourcesToCurrentMessage,
+        addSourcesByToolCallId,
     ]);
     /**
      * Cancel the current agent turn
@@ -446,7 +446,7 @@ export function useChatMessages(client, startSession) {
                 }
                 else if (chunk.type === "sources") {
                     logger.debug("Received sources chunk", { chunk });
-                    addSourcesToCurrentMessage(chunk.sources);
+                    addSourcesByToolCallId(chunk.sources);
                 }
             }
             if (!streamCompleted) {
@@ -497,7 +497,7 @@ export function useChatMessages(client, startSession) {
         addToolCallToCurrentMessage,
         updateToolCallInCurrentMessage,
         addHookNotificationToCurrentMessage,
-        addSourcesToCurrentMessage,
+        addSourcesByToolCallId,
     ]);
     return {
         messages,

package/dist/core/hooks/use-tool-calls.js

@@ -1,7 +1,7 @@
 import { createLogger } from "@townco/core";
 import { useEffect } from "react";
 import { useChatStore } from "../store/chat-store.js";
-const _logger = createLogger("use-tool-calls", "debug");
+const logger = createLogger("use-tool-calls", "debug");
 /**
  * Hook to track and manage tool calls from ACP sessions
  *
@@ -16,23 +16,46 @@ export function useToolCalls(client) {
     const updateToolCall = useChatStore((state) => state.updateToolCall);
     const addToolCallToCurrentMessage = useChatStore((state) => state.addToolCallToCurrentMessage);
     const updateToolCallInCurrentMessage = useChatStore((state) => state.updateToolCallInCurrentMessage);
+    const addSourcesByToolCallId = useChatStore((state) => state.addSourcesByToolCallId);
     useEffect(() => {
         if (!client)
             return;
+        logger.info("useToolCalls: subscribing to session updates");
         // Subscribe to session updates for tool calls
         const unsubscribe = client.onSessionUpdate((update) => {
+            logger.info("useToolCalls: received session update", {
+                type: update.type,
+                sessionId: update.sessionId,
+            });
             if (update.type === "tool_call") {
+                logger.info("useToolCalls: processing tool_call", {
+                    toolCallId: update.toolCall.id,
+                    title: update.toolCall.title,
+                });
                 // Add to session-level tool calls (for sidebar)
                 addToolCall(update.sessionId, update.toolCall);
                 // Also add to current assistant message (for inline display)
                 addToolCallToCurrentMessage(update.toolCall);
             }
             else if (update.type === "tool_call_update") {
+                logger.info("useToolCalls: processing tool_call_update", {
+                    toolCallId: update.toolCallUpdate.id,
+                });
                 // Update session-level tool calls (for sidebar)
                 updateToolCall(update.sessionId, update.toolCallUpdate);
                 // Also update in current assistant message (for inline display)
                 updateToolCallInCurrentMessage(update.toolCallUpdate);
             }
+            else if (update.type === "sources") {
+                logger.info("useToolCalls: processing sources", {
+                    sourcesCount: update.sources.length,
+                    sourceIds: update.sources.map((s) => s.id),
+                    toolCallIds: update.sources.map((s) => s.toolCallId),
+                });
+                // Distribute sources to the correct messages based on toolCallId
+                // This handles sources restored during session replay
+                addSourcesByToolCallId(update.sources);
+            }
         });
         return () => {
             unsubscribe();
@@ -43,6 +66,7 @@ export function useToolCalls(client) {
         updateToolCall,
         addToolCallToCurrentMessage,
         updateToolCallInCurrentMessage,
+        addSourcesByToolCallId,
     ]);
     return {
         toolCalls,

package/dist/core/store/chat-store.d.ts

@@ -33,6 +33,7 @@ export interface ChatStore {
     messages: DisplayMessage[];
     isStreaming: boolean;
     streamingStartTime: number | null;
+    sources: Source[];
     toolCalls: Record<string, ToolCall[]>;
     totalBilled: {
         inputTokens: number;
@@ -56,6 +57,7 @@ export interface ChatStore {
     addMessage: (message: DisplayMessage) => void;
     updateMessage: (id: string, updates: Partial<DisplayMessage>) => void;
     clearMessages: () => void;
+    addSources: (sources: Source[]) => void;
     setIsStreaming: (streaming: boolean) => void;
     setStreamingStartTime: (time: number | null) => void;
     addToolCall: (sessionId: string, toolCall: ToolCall) => void;
@@ -64,6 +66,7 @@ export interface ChatStore {
     updateToolCallInCurrentMessage: (update: ToolCallUpdate) => void;
     addHookNotificationToCurrentMessage: (notification: HookNotification) => void;
     addSourcesToCurrentMessage: (sources: Source[]) => void;
+    addSourcesByToolCallId: (sources: Source[]) => void;
     setInputValue: (value: string) => void;
     setInputSubmitting: (submitting: boolean) => void;
     addFileAttachment: (file: InputState["attachedFiles"][number]) => void;

package/dist/core/store/chat-store.js

@@ -99,6 +99,7 @@ export const useChatStore = create((set) => ({
     messages: [],
     isStreaming: false,
     streamingStartTime: null,
+    sources: [],
     toolCalls: {},
     totalBilled: {
         inputTokens: 0,
@@ -229,7 +230,18 @@ export const useChatStore = create((set) => ({
             currentContext: newCurrentContext,
         };
     }),
-    clearMessages: () => set({ messages: [] }),
+    clearMessages: () => set({ messages: [], sources: [] }),
+    addSources: (incomingSources) => set((state) => {
+        if (!incomingSources || incomingSources.length === 0)
+            return state;
+        // Merge/dedupe by source id (ids are globally unique within a session)
+        const byId = new Map();
+        for (const s of state.sources)
+            byId.set(s.id, s);
+        for (const s of incomingSources)
+            byId.set(s.id, s);
+        return { sources: Array.from(byId.values()) };
+    }),
     setIsStreaming: (streaming) => set({ isStreaming: streaming }),
     setStreamingStartTime: (time) => set({ streamingStartTime: time }),
     addToolCall: (sessionId, toolCall) => set((state) => ({
@@ -421,7 +433,16 @@ export const useChatStore = create((set) => ({
                 isStreaming: false,
                 sources,
             };
-            return { messages: [...state.messages, newMessage] };
+            // Also store in session-level sources
+            const byId = new Map();
+            for (const s of state.sources)
+                byId.set(s.id, s);
+            for (const s of sources)
+                byId.set(s.id, s);
+            return {
+                messages: [...state.messages, newMessage],
+                sources: [...byId.values()],
+            };
         }
         const messages = [...state.messages];
         const lastAssistantMsg = messages[lastAssistantIndex];
@@ -436,7 +457,117 @@ export const useChatStore = create((set) => ({
             sourcesCount: sources.length,
             totalSources: (lastAssistantMsg.sources?.length || 0) + sources.length,
         });
-        return { messages };
+        // Also store in session-level sources
+        const byId = new Map();
+        for (const s of state.sources)
+            byId.set(s.id, s);
+        for (const s of sources)
+            byId.set(s.id, s);
+        return { messages, sources: [...byId.values()] };
+    }),
+    addSourcesByToolCallId: (sources) => set((state) => {
+        logger.info("addSourcesByToolCallId called", {
+            sourcesCount: sources.length,
+            messagesCount: state.messages.length,
+            sourceDetails: sources.map((s) => ({
+                id: s.id,
+                toolCallId: s.toolCallId,
+                title: s.title?.slice(0, 30),
+            })),
+        });
+        // Log all messages and their tool calls for debugging
+        logger.info("Current messages state", {
+            messages: state.messages.map((m, i) => ({
+                index: i,
+                role: m.role,
+                contentLength: m.content?.length || 0,
+                toolCallIds: m.toolCalls?.map((tc) => tc.id) || [],
+                existingSources: m.sources?.length || 0,
+            })),
+        });
+        // Group sources by toolCallId
+        const sourcesByToolCallId = new Map();
+        for (const source of sources) {
+            if (!source.toolCallId) {
+                logger.warn("Source missing toolCallId", { sourceId: source.id });
+                continue;
+            }
+            const existing = sourcesByToolCallId.get(source.toolCallId) || [];
+            existing.push(source);
+            sourcesByToolCallId.set(source.toolCallId, existing);
+        }
+        if (sourcesByToolCallId.size === 0) {
+            logger.warn("No sources with toolCallId to distribute");
+            return state;
+        }
+        // Find messages containing tool calls with matching IDs and add sources
+        const messages = [...state.messages];
+        let sourcesAdded = 0;
+        for (const [toolCallId, toolCallSources] of sourcesByToolCallId) {
+            // Find the message containing this tool call
+            const messageIndex = messages.findIndex((msg) => msg.role === "assistant" &&
+                msg.toolCalls?.some((tc) => tc.id === toolCallId));
+            logger.info("Looking for message with toolCallId", {
+                toolCallId,
+                foundIndex: messageIndex,
+                sourcesCount: toolCallSources.length,
+            });
+            if (messageIndex !== -1) {
+                const msg = messages[messageIndex];
+                if (msg) {
+                    messages[messageIndex] = {
+                        ...msg,
+                        sources: [...(msg.sources || []), ...toolCallSources],
+                    };
+                    sourcesAdded += toolCallSources.length;
+                    logger.info("Added sources to message", {
+                        messageIndex,
+                        messageId: msg.id,
+                        toolCallId,
+                        sourcesAdded: toolCallSources.length,
+                        totalSourcesNow: messages[messageIndex]?.sources?.length,
+                        messageContentPreview: msg.content.slice(0, 100),
+                    });
+                }
+            }
+            else {
+                // Tool call not found in any message - add to last assistant message as fallback
+                logger.warn("Tool call not found in any message, using fallback", {
+                    toolCallId,
+                });
+                const lastAssistantIndex = messages.findLastIndex((msg) => msg.role === "assistant");
+                if (lastAssistantIndex !== -1) {
+                    const msg = messages[lastAssistantIndex];
+                    if (msg) {
+                        messages[lastAssistantIndex] = {
+                            ...msg,
+                            sources: [...(msg.sources || []), ...toolCallSources],
+                        };
+                        sourcesAdded += toolCallSources.length;
+                        logger.info("Added sources to last assistant message (fallback)", {
+                            lastAssistantIndex,
+                            sourcesAdded: toolCallSources.length,
+                        });
+                    }
+                }
+                else {
+                    logger.error("No assistant message found to add sources to");
+                }
+            }
+        }
+        logger.info("Finished distributing sources", {
+            totalSources: sources.length,
+            sourcesAdded,
+            toolCallIds: Array.from(sourcesByToolCallId.keys()),
+        });
+        // Always keep a session-level registry as well, so inline citations can resolve
+        // even if we couldn't associate sources to a specific message during replay.
+        const byId = new Map();
+        for (const s of state.sources)
+            byId.set(s.id, s);
+        for (const s of sources)
+            byId.set(s.id, s);
+        return { messages, sources: [...byId.values()] };
     }),
     updateToolCall: (sessionId, update) => set((state) => {
         const sessionToolCalls = state.toolCalls[sessionId] || [];

package/dist/gui/components/ChatPanelTabContent.js

@@ -5,6 +5,7 @@ import { cn } from "../lib/utils.js";
 import { SandboxFileSystemProvider } from "../providers/SandboxFileSystemProvider.js";
 import { FileSystemView } from "./FileSystemView.js";
 import { SourceListItem } from "./SourceListItem.js";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "./Tabs.js";
 import { TodoList } from "./TodoList.js";
 export const TodoTabContent = React.forwardRef(({ todos = [], className, ...props }, ref) => {
     return (_jsx("div", { ref: ref, className: cn("h-full", className), ...props, children: _jsx(TodoList, { todos: todos, className: "h-full" }) }));
@@ -76,13 +77,34 @@ export const FilesTabContent = React.forwardRef(({ files = [], provider, session
 });
 FilesTabContent.displayName = "FilesTabContent";
 export const SourcesTabContent = React.forwardRef(({ sources = [], highlightedSourceId, className, ...props }, ref) => {
+    const { usedSources, consultedSources } = React.useMemo(() => {
+        const used = [];
+        const consulted = [];
+        for (const s of sources) {
+            if (s.usedIn && s.usedIn.length > 0)
+                used.push(s);
+            else
+                consulted.push(s);
+        }
+        return { usedSources: used, consultedSources: consulted };
+    }, [sources]);
+    const defaultSourcesTab = usedSources.length > 0 ? "used" : "consulted";
+    const [activeSourcesTab, setActiveSourcesTab] = React.useState(defaultSourcesTab);
+    // Keep tab selection sensible as sources stream in.
+    React.useEffect(() => {
+        setActiveSourcesTab(defaultSourcesTab);
+    }, [defaultSourcesTab]);
     // Show empty state if no sources
     if (sources.length === 0) {
         return (_jsxs("div", { ref: ref, className: cn("flex flex-col items-center justify-center h-full text-center py-8 max-w-sm mx-auto", className), ...props, children: [
                 _jsx(Globe, { className: "size-8 text-muted-foreground opacity-50 mb-3" }), _jsx("p", { className: "text-paragraph text-muted-foreground", children: "No sources yet" }), _jsx("p", { className: "text-paragraph-sm text-muted-foreground/70 mt-1", children: "Sources will appear when your agent searches the web or fetches data." })
             ] }));
     }
-    return (_jsx("div", { ref: ref, className: cn("space-y-2", className), ...props, children: sources.map((source) => (_jsx(SourceListItem, { source: source, isSelected: source.id === highlightedSourceId }, source.id))) }));
+    return (_jsxs(Tabs, { value: activeSourcesTab, onValueChange: (v) => setActiveSourcesTab(v), className: cn("flex flex-col gap-2", className), children: [
+            _jsx("div", { className: cn("sticky top-0 z-10 -mx-4 -mt-4 mb-4 px-4 bg-card border-b border-border", "before:content-[''] before:absolute before:left-0 before:right-0 before:-top-4 before:h-4 before:bg-card before:pointer-events-none"), children: _jsxs(TabsList, { className: "w-full h-auto p-0 bg-transparent rounded-none justify-start", children: [
+                        _jsxs(TabsTrigger, { value: "used", disabled: usedSources.length === 0, className: cn("flex-1 !text-xs !px-2 !py-2 rounded-none", "bg-transparent shadow-none", "border-b-2 border-transparent -mb-px", "text-muted-foreground hover:text-foreground", "data-[state=active]:bg-transparent data-[state=active]:shadow-none", "data-[state=active]:text-foreground data-[state=active]:border-foreground"), children: ["Referenced (", usedSources.length, ")"] }), _jsxs(TabsTrigger, { value: "consulted", disabled: consultedSources.length === 0, className: cn("flex-1 !text-xs !px-2 !py-2 rounded-none", "bg-transparent shadow-none", "border-b-2 border-transparent -mb-px", "text-muted-foreground hover:text-foreground", "data-[state=active]:bg-transparent data-[state=active]:shadow-none", "data-[state=active]:text-foreground data-[state=active]:border-foreground"), children: ["Reviewed (", consultedSources.length, ")"] })
+                    ] }) }), _jsx(TabsContent, { value: "used", className: "mt-0", children: usedSources.length > 0 ? (_jsx("div", { ref: ref, className: "space-y-2", ...props, children: usedSources.map((source) => (_jsx(SourceListItem, { source: source, isSelected: source.id === highlightedSourceId }, source.id))) })) : (_jsx("div", { ref: ref, className: "text-xs text-muted-foreground", ...props, children: "No referenced sources yet." })) }), _jsx(TabsContent, { value: "consulted", className: "mt-0", children: consultedSources.length > 0 ? (_jsx("div", { ref: ref, className: "space-y-2", ...props, children: consultedSources.map((source) => (_jsx(SourceListItem, { source: source, isSelected: source.id === highlightedSourceId }, source.id))) })) : (_jsx("div", { ref: ref, className: "text-xs text-muted-foreground", ...props, children: "No reviewed sources." })) })
+        ] }));
 });
 SourcesTabContent.displayName = "SourcesTabContent";
 export const DatabaseTabContent = React.forwardRef(({ data, className, ...props }, ref) => {

package/dist/gui/components/ChatView.js

@@ -99,13 +99,49 @@ export function ChatView({ client, initialSessionId, error: initError, debuggerU
     const todos = useChatStore(selectTodosForCurrentSession);
     const _latestContextSize = useChatStore((state) => state.latestContextSize);
     const { resolvedTheme, setTheme } = useTheme();
+    const citationUsageBySourceId = React.useMemo(() => {
+        const map = new Map();
+        // Match both [[N]] and [N] (while avoiding markdown links/refs), mirroring remark-citations.
+        const citationRegex = /\[\[(\d+)\]\]|\[(\d+)\](?![:(])/g;
+        const stripCode = (markdown) => markdown
+            // fenced code blocks
+            .replace(/```[\s\S]*?```/g, "")
+            // inline code
+            .replace(/`[^`]*`/g, "");
+        for (let messageIndex = 0; messageIndex < messages.length; messageIndex++) {
+            const message = messages[messageIndex];
+            if (!message?.content)
+                continue;
+            const content = stripCode(message.content);
+            citationRegex.lastIndex = 0;
+            let match;
+            let citationIndex = 0; // global citation index within THIS message
+            // biome-ignore lint/suspicious/noAssignInExpressions: Standard regex exec loop pattern
+            while ((match = citationRegex.exec(content)) !== null) {
+                const sourceId = match[1] ?? match[2];
+                if (!sourceId)
+                    continue;
+                const existing = map.get(sourceId) ?? [];
+                existing.push({
+                    messageId: message.id,
+                    messageIndex,
+                    citationIndex,
+                });
+                map.set(sourceId, existing);
+                citationIndex++;
+            }
+        }
+        return map;
+    }, [messages]);
     // Collect all sources from all messages for the sources panel
     const allSources = React.useMemo(() => {
-        const sources = [];
+        const byId = new Map();
         for (const message of messages) {
             if (message.sources) {
                 for (const source of message.sources) {
-                    sources.push({
+                    if (byId.has(source.id))
+                        continue;
+                    byId.set(source.id, {
                         id: source.id,
                         title: source.title,
                         url: source.url,
@@ -116,8 +152,15 @@ export function ChatView({ client, initialSessionId, error: initError, debuggerU
                 }
             }
         }
-        return sources;
-    }, [messages]);
+        // Attach citation usage info for the Sources sidebar.
+        for (const [sourceId, usedIn] of citationUsageBySourceId.entries()) {
+            const existing = byId.get(sourceId);
+            if (existing) {
+                existing.usedIn = usedIn;
+            }
+        }
+        return Array.from(byId.values());
+    }, [messages, citationUsageBySourceId]);
     // Update document title with agent name and animated dots when streaming
     useDocumentTitle(agentName);
     // Initialize browser logger to capture console output and send to server

package/dist/gui/components/Message.js

@@ -100,7 +100,7 @@ export const Message = React.forwardRef(({ message, role: roleProp, layout, clas
         }
         return undefined;
     }, [role, autoScroll]);
-    return (_jsx("article", { ref: messageRef, "aria-label": `${role} message`, "data-message-id": messageId, className: cn(messageVariants({ role, layout }), className), style: {
+    return (_jsx("article", { ref: messageRef, id: messageId ? `message-${messageId}` : undefined, "aria-label": `${role} message`, "data-message-id": messageId, className: cn(messageVariants({ role, layout }), className), style: {
             minHeight: minHeight !== undefined ? `${minHeight}px` : undefined,
         }, ...props, children: children }));
 });

package/dist/gui/components/MessageContent.js

@@ -31,10 +31,31 @@ const messageContentVariants = cva("w-full rounded-2xl text-[var(--font-size)] f
         variant: "default",
     },
 });
+// Match both [[N]] and [N] (while avoiding markdown links/refs), mirroring `remark-citations`.
+const CITATION_REGEX = /\[\[(\d+)\]\]|\[(\d+)\](?![:(])/g;
+function countCitations(text) {
+    CITATION_REGEX.lastIndex = 0;
+    let count = 0;
+    while (CITATION_REGEX.exec(text) !== null)
+        count++;
+    return count;
+}
 export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStreaming: isStreamingProp, message, thinkingDisplayStyle = "collapsible", className, children, ...props }, ref) => {
     // Get streaming start time and current model from store
     const streamingStartTime = useChatStore((state) => state.streamingStartTime);
     const _currentModel = useChatStore((state) => state.currentModel);
+    const sessionSources = useChatStore((state) => state.sources);
+    const sourcesForResponse = React.useMemo(() => {
+        // Merge/dedupe global + message sources by id so inline citations can resolve after reload.
+        const byId = new Map();
+        for (const s of sessionSources)
+            byId.set(s.id, s);
+        if (message?.sources) {
+            for (const s of message.sources)
+                byId.set(s.id, s);
+        }
+        return Array.from(byId.values());
+    }, [sessionSources, message?.sources]);
     // Use smart rendering if message is provided and no custom children
     const useSmartRendering = message && !children;
     // Derive props from message if using smart rendering
@@ -189,18 +210,20 @@ export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStr
                         const hasHookPositions = hookNotifications.some((n) => n.contentPosition !== undefined);
                         if (!hasHookPositions) {
                             // No positions - render hooks at top, then content, then tool calls
-                            return (_jsxs(_Fragment, { children: [hookNotifications.length > 0 && (_jsx("div", { className: "flex flex-col gap-2 mb-1", children: hookNotifications.map((notification) => (_jsx(HookNotification, { notification: notification }, notification.id))) })), _jsx("div", { children: _jsx(Response, { content: message.content, isStreaming: message.isStreaming, showEmpty: false, sources: message.sources ?? [] }) }), groupedToolCalls.length > 0 && (_jsx("div", { className: "flex flex-col gap-2 mb-1", children: groupedToolCalls.map((item, index) => renderToolCallOrGroup(item, index)) })), shouldShowThinkingIndicator && (_jsx("div", { className: "flex flex-col my-4 rounded-md px-1 -mx-1 w-fit", children: _jsx("div", { className: "flex items-center gap-1.5", children: _jsx("span", { className: "text-paragraph-sm text-text-secondary/70", children: "Thinking..." }) }) }))] }));
+                            return (_jsxs(_Fragment, { children: [hookNotifications.length > 0 && (_jsx("div", { className: "flex flex-col gap-2 mb-1", children: hookNotifications.map((notification) => (_jsx(HookNotification, { notification: notification }, notification.id))) })), _jsx("div", { children: _jsx(Response, { content: message.content, isStreaming: message.isStreaming, showEmpty: false, sources: sourcesForResponse, messageId: message.id, citationIndexOffset: 0 }) }), groupedToolCalls.length > 0 && (_jsx("div", { className: "flex flex-col gap-2 mb-1", children: groupedToolCalls.map((item, index) => renderToolCallOrGroup(item, index)) })), shouldShowThinkingIndicator && (_jsx("div", { className: "flex flex-col my-4 rounded-md px-1 -mx-1 w-fit", children: _jsx("div", { className: "flex items-center gap-1.5", children: _jsx("span", { className: "text-paragraph-sm text-text-secondary/70", children: "Thinking..." }) }) }))] }));
                         }
                         // Hooks have positions - render them inline with content
                         const elements = [];
                         let currentPosition = 0;
+                        let citationOffset = 0;
                         hookNotifications.forEach((notification) => {
                             const position = notification.contentPosition ?? message.content.length;
                             // Add text before this hook notification
                             if (position > currentPosition) {
                                 const textChunk = message.content.slice(currentPosition, position);
                                 if (textChunk) {
-                                    elements.push(_jsx("div", { children: _jsx(Response, { content: textChunk, isStreaming: false, showEmpty: false, sources: message.sources ?? [] }) }, `text-before-hook-${notification.id}`));
+                                    elements.push(_jsx("div", { children: _jsx(Response, { content: textChunk, isStreaming: false, showEmpty: false, sources: sourcesForResponse, messageId: message.id, citationIndexOffset: citationOffset }) }, `text-before-hook-${notification.id}`));
+                                    citationOffset += countCitations(textChunk);
                                 }
                             }
                             // Add hook notification
@@ -211,7 +234,8 @@ export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStr
                         if (currentPosition < message.content.length) {
                             const remainingText = message.content.slice(currentPosition);
                             if (remainingText) {
-                                elements.push(_jsx("div", { children: _jsx(Response, { content: remainingText, isStreaming: message.isStreaming, showEmpty: false, sources: message.sources ?? [] }) }, "text-end-hooks"));
+                                elements.push(_jsx("div", { children: _jsx(Response, { content: remainingText, isStreaming: message.isStreaming, showEmpty: false, sources: sourcesForResponse, messageId: message.id, citationIndexOffset: citationOffset }) }, "text-end-hooks"));
+                                citationOffset += countCitations(remainingText);
                             }
                         }
                         // Add tool calls at the end
@@ -247,6 +271,7 @@ export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStr
                     // Sort by position
                     positionedItems.sort((a, b) => a.position - b.position);
                     let currentPosition = 0;
+                    let citationOffset = 0;
                     let currentBatch = [];
                     let currentBatchId;
                     let currentBatchTitle;
@@ -278,7 +303,8 @@ export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStr
                                 const itemId = positionedItem.type === "toolCall"
                                     ? positionedItem.item.id
                                     : positionedItem.item.id;
-                                elements.push(_jsx("div", { children: _jsx(Response, { content: textChunk, isStreaming: false, showEmpty: false, sources: message.sources ?? [] }) }, `text-before-${itemId}`));
+                                elements.push(_jsx("div", { children: _jsx(Response, { content: textChunk, isStreaming: false, showEmpty: false, sources: sourcesForResponse, messageId: message.id, citationIndexOffset: citationOffset }) }, `text-before-${itemId}`));
+                                citationOffset += countCitations(textChunk);
                             }
                         }
                         if (positionedItem.type === "hookNotification") {
@@ -332,7 +358,8 @@ export const MessageContent = React.forwardRef(({ role: roleProp, variant, isStr
                     if (currentPosition < message.content.length) {
                         const remainingText = message.content.slice(currentPosition);
                         if (remainingText) {
-                            elements.push(_jsx("div", { children: _jsx(Response, { content: remainingText, isStreaming: message.isStreaming, showEmpty: false, sources: message.sources ?? [] }) }, "text-end"));
+                            elements.push(_jsx("div", { children: _jsx(Response, { content: remainingText, isStreaming: message.isStreaming, showEmpty: false, sources: sourcesForResponse, messageId: message.id, citationIndexOffset: citationOffset }) }, "text-end"));
+                            citationOffset += countCitations(remainingText);
                         }
                     }
                     // Render preliminary (selecting) tool calls at the end, grouped

package/dist/gui/components/Response.d.ts

@@ -15,5 +15,16 @@ export interface ResponseProps extends React.HTMLAttributes<HTMLDivElement> {
     emptyMessage?: string;
     /** Citation sources for rendering inline citations */
     sources?: Source[];
+    /**
+     * Optional message id. If provided, we add stable per-citation anchors so the
+     * Sources sidebar can scroll to the exact cited spot.
+     */
+    messageId?: string;
+    /**
+     * When a message is rendered in multiple `Response` chunks (due to tool calls/hooks),
+     * citations need a global index across the full message. This is the starting index
+     * for citations rendered by THIS chunk.
+     */
+    citationIndexOffset?: number;
 }
 export declare const Response: React.NamedExoticComponent<ResponseProps & React.RefAttributes<HTMLDivElement>>;

package/dist/gui/components/Response.js

@@ -1,35 +1,76 @@
 "use client";
 import { jsx as _jsx } from "react/jsx-runtime";
+import { createLogger } from "@townco/core";
 import * as React from "react";
 import remarkGfm from "remark-gfm";
 import { Streamdown } from "streamdown";
 import { remarkCitations } from "../lib/remark-citations.js";
 import { cn } from "../lib/utils.js";
 import { CitationChip } from "./CitationChip.js";
-export const Response = React.memo(React.forwardRef(({ content, isStreaming = false, showEmpty = true, emptyMessage = "", sources = [], className, ...props }, ref) => {
+const logger = createLogger("Response", "debug");
+export const Response = React.memo(React.forwardRef(({ content, isStreaming = false, showEmpty = true, emptyMessage = "", sources = [], messageId, citationIndexOffset = 0, className, ...props }, ref) => {
+    // State to force remount when sources change
+    const [remountKey, setRemountKey] = React.useState(0);
+    // Log every render to track sources changes
+    logger.info("Response render", {
+        contentLength: content.length,
+        contentPreview: content.slice(0, 100),
+        sourcesCount: sources.length,
+        sourceIds: sources.map((s) => s.id),
+        remountKey,
+    });
+    // Increment remount key when sources change to force Streamdown remount
+    React.useEffect(() => {
+        setRemountKey((prev) => {
+            logger.info("Sources changed, incrementing remountKey", {
+                sourcesCount: sources.length,
+                sourceIds: sources.map((s) => s.id),
+                oldRemountKey: prev,
+            });
+            return prev + 1;
+        });
+    }, [sources]);
     // Memoize the remark plugins array to prevent re-renders
     // remarkGfm adds support for tables, strikethrough, autolinks, task lists
     const remarkPlugins = React.useMemo(() => [remarkGfm, remarkCitations], []);
     // Memoize the custom components to prevent re-renders
     // We use 'as Record<string, unknown>' because Streamdown's Components type
     // doesn't include all the props we need
-    const customComponents = React.useMemo(() => ({
-        // Custom span component that intercepts citation markers
-        // The remark-citations plugin creates spans with class "citation-marker"
-        // and data-citation-id attribute
-        // Note: Streamdown may pass 'class' instead of 'className'
-        span: (spanProps) => {
-            const { className, class: classAttr, "data-citation-id": citationId, children, ...restProps } = spanProps;
-            // Check both className and class (Streamdown may use either)
-            const cssClass = className || classAttr;
-            // Check if this span is a citation marker
-            if (cssClass === "citation-marker" && citationId) {
-                return _jsx(CitationChip, { sourceId: citationId, sources: sources });
-            }
-            // Otherwise render as normal span
-            return (_jsx("span", { className: cssClass, ...restProps, children: children }));
-        },
-    }), [sources]);
+    const customComponents = React.useMemo(() => {
+        // Each `Response` instance renders citations in-order. We assign anchors using a
+        // running local index + the provided chunk offset.
+        const citationCounter = { current: 0 };
+        return {
+            // Custom span component that intercepts citation markers
+            // The remark-citations plugin creates spans with class "citation-marker"
+            // and data-citation-id attribute
+            // Note: Streamdown may pass 'class' instead of 'className'
+            span: (spanProps) => {
+                const { className, class: classAttr, "data-citation-id": citationId, children, ...restProps } = spanProps;
+                // Check both className and class (Streamdown may use either)
+                const cssClass = className || classAttr;
+                // Check if this span is a citation marker
+                if (cssClass === "citation-marker" && citationId) {
+                    const globalCitationIndex = citationIndexOffset + citationCounter.current++;
+                    logger.info("Rendering citation marker", {
+                        citationId,
+                        sourcesAvailable: sources.length,
+                        sourceIds: sources.map((s) => s.id),
+                        foundSource: sources.find((s) => s.id === citationId)
+                            ? "yes"
+                            : "no",
+                    });
+                    // Anchor wrapper (when messageId is provided)
+                    if (messageId) {
+                        return (_jsx("span", { id: `citation-${messageId}-${globalCitationIndex}`, "data-citation-source-id": citationId, "data-citation-anchor": "true", className: "inline-flex rounded-md", children: _jsx(CitationChip, { sourceId: citationId, sources: sources }) }));
+                    }
+                    return _jsx(CitationChip, { sourceId: citationId, sources: sources });
+                }
+                // Otherwise render as normal span
+                return (_jsx("span", { className: cssClass, ...restProps, children: children }));
+            },
+        };
+    }, [sources, messageId, citationIndexOffset]);
     // Show empty state during streaming if no content yet
     if (!content && isStreaming && showEmpty) {
         return (_jsx("div", { ref: ref, className: cn("opacity-70 italic text-paragraph-sm", className), ...props, children: emptyMessage }));
@@ -37,7 +78,7 @@ export const Response = React.memo(React.forwardRef(({ content, isStreaming = fa
     if (!content) {
         return null;
     }
-    return (_jsx("div", { ref: ref, ...props, children: _jsx(Streamdown, { className: cn("size-full", "[&>*:first-child]:mt-0 [&>*:last-child]:mb-0", "[&_code]:whitespace-pre-wrap [&_code]:break-words", "[&_pre]:max-w-full [&_pre]:overflow-x-auto", className), remarkPlugins: remarkPlugins, components: customComponents, children: content }) }));
+    return (_jsx("div", { ref: ref, ...props, children: _jsx(Streamdown, { className: cn("size-full", "[&>*:first-child]:mt-0 [&>*:last-child]:mb-0", "[&_code]:whitespace-pre-wrap [&_code]:break-words", "[&_pre]:max-w-full [&_pre]:overflow-x-auto", className), remarkPlugins: remarkPlugins, components: customComponents, children: content }, `streamdown-${remountKey}`) }));
 }), (prevProps, nextProps) => prevProps.content === nextProps.content &&
     prevProps.sources === nextProps.sources);
 Response.displayName = "Response";

package/dist/gui/components/SourceListItem.d.ts

@@ -7,6 +7,16 @@ export interface SourceItem {
     sourceName: string;
     favicon?: string;
     timestamp?: string;
+    /**
+     * Where this source is cited in the transcript.
+     * `messageIndex` is 0-based within the current rendered message list.
+     */
+    usedIn?: Array<{
+        messageId: string;
+        messageIndex: number;
+        /** 0-based citation occurrence index within the message */
+        citationIndex: number;
+    }>;
 }
 export interface SourceListItemProps extends React.ButtonHTMLAttributes<HTMLButtonElement> {
     source: SourceItem;

package/dist/gui/components/SourceListItem.js

@@ -2,6 +2,38 @@ import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
 import * as React from "react";
 import { cn } from "../lib/utils.js";
 export const SourceListItem = React.forwardRef(({ source, isSelected, className, ...props }, ref) => {
+    const usedIn = source.usedIn ?? [];
+    const highlightCitation = React.useCallback((el) => {
+        // Clear any previous highlight
+        document
+            .querySelectorAll("[data-citation-highlight='true']")
+            .forEach((n) => {
+            n.removeAttribute("data-citation-highlight");
+            n.classList.remove("ring-2", "ring-ring", "ring-offset-2", "ring-offset-background");
+        });
+        el.setAttribute("data-citation-highlight", "true");
+        el.classList.add("ring-2", "ring-ring", "ring-offset-2", "ring-offset-background");
+        window.setTimeout(() => {
+            // Only remove if it's still the highlighted element.
+            if (el.getAttribute("data-citation-highlight") === "true") {
+                el.removeAttribute("data-citation-highlight");
+                el.classList.remove("ring-2", "ring-ring", "ring-offset-2", "ring-offset-background");
+            }
+        }, 1400);
+    }, []);
+    const scrollToCitationOrdinal = React.useCallback((sourceId, ordinal) => {
+        const nodes = Array.from(document.querySelectorAll(`[data-citation-anchor="true"][data-citation-source-id="${CSS.escape(sourceId)}"]`));
+        const el = nodes[ordinal];
+        if (!el)
+            return false;
+        el.scrollIntoView({
+            behavior: "smooth",
+            block: "center",
+            inline: "nearest",
+        });
+        window.setTimeout(() => highlightCitation(el), 150);
+        return true;
+    }, [highlightCitation]);
     return (_jsxs("button", { ref: ref, type: "button", className: cn(
         // Base styles matching FileSystemItem
         "group flex w-full text-left gap-2 items-start p-2 rounded-md cursor-pointer transition-colors text-paragraph-sm", 
@@ -10,12 +42,28 @@ export const SourceListItem = React.forwardRef(({ source, isSelected, className,
         // Focus state - matching FileSystemItem
         "focus-visible:outline-none focus-visible:ring-[3px] focus-visible:ring-border-dark", 
         // Selected state - matching FileSystemItem
-        isSelected && "bg-accent", className), onClick: () => window.open(source.url, "_blank"), ...props, children: [
-            _jsx("div", { className: "shrink-0 flex items-center h-5", children: _jsx("div", { className: "relative rounded-[3px] size-4 overflow-hidden bg-muted", children: source.favicon ? (_jsx("img", { alt: source.sourceName, className: "size-full object-cover", src: source.favicon })) : (_jsx("div", { className: "size-full bg-muted" })) }) }), _jsxs("div", { className: "flex flex-1 flex-col gap-1 min-w-0", children: [
+        isSelected && "bg-accent", className), onClick: () => {
+            // Always open the source URL. (Scrolling is handled via the usage pills.)
+            window.open(source.url, "_blank", "noopener,noreferrer");
+        }, ...props, children: [
+            _jsx("div", { className: "shrink-0 flex items-center h-5", children: _jsx("div", { className: "relative rounded-[3px] size-4 overflow-hidden bg-muted", children: source.favicon ? (_jsx("img", { alt: source.sourceName, className: "size-full object-cover", src: source.favicon })) : (_jsx("div", { className: "size-full bg-muted" })) }) }), _jsxs("div", { className: "flex flex-1 flex-col gap-0.5 min-w-0", children: [
                     _jsxs("div", { className: "text-paragraph-sm text-foreground truncate", children: [
                             _jsx("span", { className: "font-medium", children: source.sourceName }), _jsxs("span", { className: "text-muted-foreground", children: [" \u00B7 ", source.title] })
-                        ] }), _jsx("p", { className: "text-paragraph-sm text-muted-foreground line-clamp-3 break-all", children: source.snippet })
-                ] })
+                        ] }), usedIn.length > 0 && (_jsxs("div", { className: "flex flex-wrap items-center gap-1.5 pt-1", children: [
+                            _jsx("span", { className: "text-xs text-muted-foreground", children: "Referenced" }), Array.from({ length: Math.min(usedIn.length, 12) }, (_, idx) => idx).map((idx) => (
+                            // Not a <button> (nested buttons are invalid); we emulate button semantics.
+                            // biome-ignore lint/a11y/useSemanticElements: spans are required to avoid nested buttons
+                            _jsx("span", { role: "button", tabIndex: 0, className: cn("px-1.5 py-0.5 rounded-md", "text-[10px] leading-none font-medium", "bg-muted text-muted-foreground", "hover:bg-accent hover:text-foreground transition-colors", "focus-visible:outline-none focus-visible:ring-[3px] focus-visible:ring-border-dark", "cursor-pointer select-none"), onClick: (e) => {
+                                    e.preventDefault();
+                                    e.stopPropagation();
+                                    scrollToCitationOrdinal(source.id, idx);
+                                }, onKeyDown: (e) => {
+                                    if (e.key === "Enter" || e.key === " ") {
+                                        e.preventDefault();
+                                        e.stopPropagation();
+                                        scrollToCitationOrdinal(source.id, idx);
+                                    }
+                                }, "aria-label": `Scroll to citation ${idx + 1}`, children: idx + 1 }, `${source.id}-use-${idx}`))), usedIn.length > 12 && (_jsxs("span", { className: "text-xs text-muted-foreground", children: ["+", usedIn.length - 12, " more"] }))] }))] })
         ] }));
 });
 SourceListItem.displayName = "SourceListItem";

package/dist/sdk/schemas/session.d.ts

@@ -637,6 +637,74 @@ export declare const SessionUpdate: z.ZodUnion<readonly [z.ZodObject<{
         rawOutput: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
     }, z.core.$strip>;
     messageId: z.ZodOptional<z.ZodString>;
+}, z.core.$strip>, z.ZodObject<{
+    sessionId: z.ZodString;
+    status: z.ZodOptional<z.ZodEnum<{
+        active: "active";
+        connected: "connected";
+        connecting: "connecting";
+        disconnected: "disconnected";
+        error: "error";
+        idle: "idle";
+        streaming: "streaming";
+    }>>;
+    message: z.ZodOptional<z.ZodObject<{
+        id: z.ZodString;
+        role: z.ZodEnum<{
+            assistant: "assistant";
+            system: "system";
+            tool: "tool";
+            user: "user";
+        }>;
+        content: z.ZodArray<z.ZodDiscriminatedUnion<[z.ZodObject<{
+            type: z.ZodLiteral<"text">;
+            text: z.ZodString;
+        }, z.core.$strip>, z.ZodObject<{
+            type: z.ZodLiteral<"image">;
+            url: z.ZodOptional<z.ZodString>;
+            source: z.ZodOptional<z.ZodObject<{
+                type: z.ZodLiteral<"base64">;
+                media_type: z.ZodEnum<{
+                    "image/gif": "image/gif";
+                    "image/jpeg": "image/jpeg";
+                    "image/png": "image/png";
+                    "image/webp": "image/webp";
+                }>;
+                data: z.ZodString;
+            }, z.core.$strip>>;
+        }, z.core.$strip>, z.ZodObject<{
+            type: z.ZodLiteral<"file">;
+            name: z.ZodString;
+            path: z.ZodOptional<z.ZodString>;
+            url: z.ZodOptional<z.ZodString>;
+            mimeType: z.ZodString;
+            size: z.ZodOptional<z.ZodNumber>;
+        }, z.core.$strip>, z.ZodObject<{
+            type: z.ZodLiteral<"tool_call">;
+            id: z.ZodString;
+            name: z.ZodString;
+            arguments: z.ZodRecord<z.ZodString, z.ZodUnknown>;
+        }, z.core.$strip>, z.ZodObject<{
+            type: z.ZodLiteral<"tool_result">;
+            callId: z.ZodString;
+            result: z.ZodUnknown;
+            error: z.ZodOptional<z.ZodString>;
+        }, z.core.$strip>], "type">>;
+        timestamp: z.ZodISODateTime;
+        metadata: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
+    }, z.core.$strip>>;
+    error: z.ZodOptional<z.ZodString>;
+    _meta: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
+    type: z.ZodLiteral<"sources">;
+    sources: z.ZodArray<z.ZodObject<{
+        id: z.ZodString;
+        url: z.ZodString;
+        title: z.ZodString;
+        snippet: z.ZodOptional<z.ZodString>;
+        favicon: z.ZodOptional<z.ZodString>;
+        toolCallId: z.ZodString;
+        sourceName: z.ZodOptional<z.ZodString>;
+    }, z.core.$strip>>;
 }, z.core.$strip>, z.ZodObject<{
     sessionId: z.ZodString;
     status: z.ZodOptional<z.ZodEnum<{

package/dist/sdk/schemas/session.js

@@ -1,6 +1,6 @@
 import { z } from "zod";
 import { ToolCallSchema, ToolCallUpdateSchema, } from "../../core/schemas/tool-call.js";
-import { Message } from "./message.js";
+import { CitationSource, Message } from "./message.js";
 /**
  * Session status
  */
@@ -83,6 +83,14 @@ const ToolOutputSessionUpdate = BaseSessionUpdate.extend({
     }),
     messageId: z.string().optional(),
 });
+/**
+ * Session update with sources (sessionUpdate: "sources")
+ * Sent during replay to restore persisted citation sources
+ */
+const SourcesSessionUpdate = BaseSessionUpdate.extend({
+    type: z.literal("sources"),
+    sources: z.array(CitationSource),
+});
 /**
  * Generic session update
  */
@@ -96,5 +104,6 @@ export const SessionUpdate = z.union([
     ToolCallSessionUpdate,
     ToolCallUpdateSessionUpdate,
     ToolOutputSessionUpdate,
+    SourcesSessionUpdate,
     GenericSessionUpdate,
 ]);

package/dist/sdk/transports/http.d.ts

@@ -22,6 +22,7 @@ export declare class HttpTransport implements Transport {
     private options;
     private isReceivingMessages;
     private isInReplayMode;
+    private pendingReplayUpdates;
     private agentInfo?;
     constructor(options: HttpTransportOptions);
     connect(): Promise<void>;

package/dist/sdk/transports/http.js

@@ -25,6 +25,7 @@ export class HttpTransport {
     options;
     isReceivingMessages = false;
     isInReplayMode = false; // True during session replay, ignores non-replay streaming
+    pendingReplayUpdates = []; // Queue session updates during replay for late-subscribing callbacks
     agentInfo;
     constructor(options) {
         // Validate options at the boundary using Zod
@@ -565,6 +566,61 @@ export class HttpTransport {
     }
     onSessionUpdate(callback) {
         this.sessionUpdateCallbacks.add(callback);
+        // Replay any queued session updates from replay that arrived before callbacks were registered
+        // IMPORTANT: Use notifySessionUpdate to send to ALL registered callbacks, not just this one.
+        // This ensures that when multiple hooks register (e.g., message handler, useToolCalls),
+        // they ALL receive the queued updates, not just the first one to register.
+        if (this.pendingReplayUpdates.length > 0) {
+            logger.info("Replaying queued session updates to ALL callbacks", {
+                count: this.pendingReplayUpdates.length,
+                callbackCount: this.sessionUpdateCallbacks.size,
+                updateTypes: this.pendingReplayUpdates.map((u) => u.type),
+                updateDetails: this.pendingReplayUpdates.map((u) => {
+                    if (u.type === "tool_call") {
+                        return { type: "tool_call", toolCallId: u.toolCall?.id };
+                    }
+                    if (u.type === "sources") {
+                        return {
+                            type: "sources",
+                            count: u.sources?.length,
+                            ids: u.sources?.map((s) => s.id),
+                        };
+                    }
+                    if (u.type === "generic" && u.message) {
+                        const firstContent = u.message.content?.[0];
+                        return {
+                            type: "message",
+                            role: u.message.role,
+                            contentLength: firstContent?.type === "text" ? firstContent.text?.length : 0,
+                        };
+                    }
+                    return { type: u.type };
+                }),
+            });
+            // Copy and clear the queue immediately to prevent re-triggering when other callbacks register
+            const updatesToReplay = [...this.pendingReplayUpdates];
+            this.pendingReplayUpdates = [];
+            // Use setTimeout to ensure all callbacks have a chance to register before replaying
+            setTimeout(() => {
+                logger.info("Starting replay of queued updates to all callbacks", {
+                    updateCount: updatesToReplay.length,
+                    callbackCount: this.sessionUpdateCallbacks.size,
+                });
+                for (const update of updatesToReplay) {
+                    try {
+                        logger.info("Replaying update to all callbacks", {
+                            type: update.type,
+                        });
+                        // Notify ALL registered callbacks, not just the one that triggered this replay
+                        this.notifySessionUpdate(update);
+                    }
+                    catch (error) {
+                        logger.error("Error replaying session update", { error });
+                    }
+                }
+                logger.info("Finished replay of queued updates");
+            }, 0);
+        }
         return () => {
             this.sessionUpdateCallbacks.delete(callback);
         };
@@ -810,6 +866,28 @@ export class HttpTransport {
         try {
             const message = JSON.parse(data);
             logger.debug("Received SSE message", { message });
+            // Check if this is a sources message (custom extension to ACP)
+            const isSourcesMessage = message &&
+                typeof message === "object" &&
+                message.method === "session/update" &&
+                message.params?.update?.sessionUpdate === "sources";
+            if (isSourcesMessage) {
+                // Use console.warn directly for gui-console capture
+                console.warn("🟢 RECEIVED SOURCES SSE MESSAGE", {
+                    sourcesCount: message.params?.update?.sources?.length,
+                    isInReplayMode: this.isInReplayMode,
+                    callbackCount: this.sessionUpdateCallbacks.size,
+                });
+                // Handle sources directly without ACP schema validation
+                try {
+                    this.handleSessionNotification(message.params);
+                    console.warn("🟢 AFTER handleSessionNotification for sources");
+                }
+                catch (error) {
+                    console.error("🔴 ERROR in handleSessionNotification for sources", error);
+                }
+                return;
+            }
             // Validate the message is an ACP agent outgoing message
             const parseResult = acp.agentOutgoingMessageSchema.safeParse(message);
             if (!parseResult.success) {
@@ -843,18 +921,29 @@ export class HttpTransport {
      * Handle a session notification from the agent
      */
     handleSessionNotification(params) {
+        // Extract update type early to check if this is a sources notification
+        const paramsAny = params;
+        const updateAny = paramsAny.update;
+        const isSourcesNotification = updateAny?.sessionUpdate === "sources";
         // Skip processing if stream has been cancelled/completed
-        if (this.streamComplete) {
+        // BUT always allow sources through - they arrive after stream completion during replay
+        if (this.streamComplete && !isSourcesNotification) {
             logger.debug("Skipping session notification - stream complete/cancelled");
             return;
         }
+        if (this.streamComplete && isSourcesNotification) {
+            console.warn("🟢 Processing sources notification after stream complete");
+        }
         logger.debug("handleSessionNotification called", { params });
         // Extract content from the update
         const paramsExtended = params;
         const update = paramsExtended.update;
         const sessionId = this.currentSessionId || params.sessionId;
-        logger.debug("Update session type", {
+        logger.warn("📥 SSE UPDATE RECEIVED", {
             sessionUpdate: update?.sessionUpdate,
+            hasUpdate: !!update,
+            isInReplayMode: this.isInReplayMode,
+            callbackCount: this.sessionUpdateCallbacks.size,
         });
         // Handle sandbox file changes
         // Type assertion needed because TypeScript doesn't recognize this as a valid session update type
@@ -1064,9 +1153,18 @@ export class HttpTransport {
                 isReplay,
                 isInReplayMode: this.isInReplayMode,
             });
-            // During replay, notify directly since there's no active receive() consumer
+            // During replay, handle specially
             if (isReplay || this.isInReplayMode) {
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_call for late-subscribing callbacks", {
+                        toolCallId: toolCall.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool call as a chunk for ordered processing during live streaming
@@ -1273,8 +1371,16 @@ export class HttpTransport {
             // Check if this is replay (tool_call_update doesn't have isReplay in _meta,
             // but we can check if we're in replay mode)
             if (this.isInReplayMode) {
-                // During replay, notify directly since there's no active receive() consumer
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_call_update for late-subscribing callbacks", {
+                        toolCallId: toolCallUpdate.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool call update as a chunk for ordered processing
@@ -1406,7 +1512,16 @@ export class HttpTransport {
             };
             // During replay, notify directly; otherwise queue for ordered processing
             if (this.isInReplayMode) {
-                this.notifySessionUpdate(sessionUpdate);
+                // If no callbacks registered yet, queue for late-subscribing hooks
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    logger.debug("Queueing tool_output for late-subscribing callbacks", {
+                        toolCallId: toolOutput.id,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                }
+                else {
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
                 // Queue tool output as a chunk for ordered processing
@@ -1433,21 +1548,56 @@ export class HttpTransport {
             update.sessionUpdate === "sources") {
             // Sources notification - citation sources from tool calls
             const sourcesUpdate = update;
-            logger.debug("Received sources notification", {
+            console.warn("🔵 SOURCES in handleSessionNotification", {
                 sourcesCount: sourcesUpdate.sources.length,
+                isInReplayMode: this.isInReplayMode,
+                callbackCount: this.sessionUpdateCallbacks.size,
+                firstSourceId: sourcesUpdate.sources[0]?.id,
+                firstSourceToolCallId: sourcesUpdate.sources[0]?.toolCallId,
             });
-            // Create a sources chunk for the message queue
-            const sourcesChunk = {
+            // Create a sources session update
+            const sessionUpdate = {
                 type: "sources",
+                sessionId,
+                status: "active",
                 sources: sourcesUpdate.sources,
             };
-            // Queue for ordered processing
-            const resolver = this.chunkResolvers.shift();
-            if (resolver) {
-                resolver(sourcesChunk);
+            // During replay, handle sources specially
+            if (this.isInReplayMode) {
+                // If no callbacks are registered yet (React hooks haven't subscribed),
+                // queue the sources to be replayed when they do subscribe
+                if (this.sessionUpdateCallbacks.size === 0) {
+                    console.warn("🔵 QUEUEING sources for late-subscribing callbacks", {
+                        sourcesCount: sourcesUpdate.sources.length,
+                        queueLengthBefore: this.pendingReplayUpdates.length,
+                    });
+                    this.pendingReplayUpdates.push(sessionUpdate);
+                    console.warn("🔵 Queue length after:", {
+                        queueLengthAfter: this.pendingReplayUpdates.length,
+                    });
+                }
+                else {
+                    console.warn("🔵 NOTIFYING sources immediately (callbacks registered)", {
+                        sourcesCount: sourcesUpdate.sources.length,
+                        callbackCount: this.sessionUpdateCallbacks.size,
+                    });
+                    this.notifySessionUpdate(sessionUpdate);
+                }
             }
             else {
-                this.messageQueue.push(sourcesChunk);
+                // Create a sources chunk for the message queue
+                const sourcesChunk = {
+                    type: "sources",
+                    sources: sourcesUpdate.sources,
+                };
+                // Queue for ordered processing
+                const resolver = this.chunkResolvers.shift();
+                if (resolver) {
+                    resolver(sourcesChunk);
+                }
+                else {
+                    this.messageQueue.push(sourcesChunk);
+                }
             }
         }
         else if (update?.sessionUpdate === "agent_message_chunk") {
@@ -1539,8 +1689,18 @@ export class HttpTransport {
                             timestamp: new Date().toISOString(),
                         },
                     };
-                    // Notify as a complete message (for session replay or initial message)
-                    this.notifySessionUpdate(messageSessionUpdate);
+                    // During replay, queue messages along with tool_calls and sources
+                    // so everything is processed in the correct order
+                    if (isReplay && this.sessionUpdateCallbacks.size === 0) {
+                        logger.debug("Queueing assistant message for late-subscribing callbacks", {
+                            textLength: contentObj.text.length,
+                        });
+                        this.pendingReplayUpdates.push(messageSessionUpdate);
+                    }
+                    else {
+                        // Notify as a complete message (for session replay or initial message)
+                        this.notifySessionUpdate(messageSessionUpdate);
+                    }
                 }
             }
             // Send session update for:
@@ -1554,6 +1714,11 @@ export class HttpTransport {
         else if (update?.sessionUpdate === "user_message_chunk") {
             // Handle user message chunks (could be from replay or new messages)
             logger.debug("Received user_message_chunk", { update });
+            // Check if this is a replay
+            const isReplay = update._meta &&
+                typeof update._meta === "object" &&
+                "isReplay" in update._meta &&
+                update._meta.isReplay === true;
             const content = update.content;
             if (content && typeof content === "object") {
                 const contentObj = content;
@@ -1575,13 +1740,28 @@ export class HttpTransport {
                             timestamp: new Date().toISOString(),
                         },
                     };
-                    logger.debug("Notifying session update for user message");
-                    this.notifySessionUpdate(sessionUpdate);
+                    // During replay, queue messages along with tool_calls and sources
+                    // so everything is processed in the correct order
+                    if (isReplay && this.sessionUpdateCallbacks.size === 0) {
+                        logger.debug("Queueing user message for late-subscribing callbacks", {
+                            textLength: contentObj.text.length,
+                        });
+                        this.pendingReplayUpdates.push(sessionUpdate);
+                    }
+                    else {
+                        logger.debug("Notifying session update for user message");
+                        this.notifySessionUpdate(sessionUpdate);
+                    }
                 }
             }
         }
         else {
             // Handle other session updates
+            logger.warn("⚠️ UNHANDLED SESSION UPDATE - falling through to generic", {
+                sessionUpdate: update?.sessionUpdate,
+                updateKeys: update ? Object.keys(update) : [],
+                hasSources: update && "sources" in update,
+            });
             const sessionUpdate = {
                 type: "generic",
                 sessionId,

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@townco/ui",
-  "version": "0.1.109",
+  "version": "0.1.110",
   "type": "module",
   "main": "./dist/index.js",
   "types": "./dist/index.d.ts",
@@ -49,7 +49,7 @@
     "@radix-ui/react-slot": "^1.2.4",
     "@radix-ui/react-tabs": "^1.1.13",
     "@radix-ui/react-tooltip": "^1.2.8",
-    "@townco/core": "0.0.87",
+    "@townco/core": "0.0.88",
     "@types/mdast": "^4.0.4",
     "@uiw/react-json-view": "^2.0.0-alpha.39",
     "class-variance-authority": "^0.7.1",
@@ -67,7 +67,7 @@
     "zustand": "^5.0.8"
   },
   "devDependencies": {
-    "@townco/tsconfig": "0.1.106",
+    "@townco/tsconfig": "0.1.107",
     "@types/node": "^24.10.0",
     "@types/react": "^19.2.2",
     "@types/unist": "^3.0.3",