@townco/ui 0.1.109 → 0.1.111

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -21,7 +21,7 @@ export function useChatMessages(client, startSession) {
21
21
  const addToolCallToCurrentMessage = useChatStore((state) => state.addToolCallToCurrentMessage);
22
22
  const updateToolCallInCurrentMessage = useChatStore((state) => state.updateToolCallInCurrentMessage);
23
23
  const addHookNotificationToCurrentMessage = useChatStore((state) => state.addHookNotificationToCurrentMessage);
24
- const addSourcesToCurrentMessage = useChatStore((state) => state.addSourcesToCurrentMessage);
24
+ const addSourcesByToolCallId = useChatStore((state) => state.addSourcesByToolCallId);
25
25
  const truncateMessagesFrom = useChatStore((state) => state.truncateMessagesFrom);
26
26
  // Track the current assistant message ID for cancellation
27
27
  const currentAssistantMessageIdRef = useRef(null);
@@ -206,7 +206,7 @@ export function useChatMessages(client, startSession) {
206
206
  // Sources chunk - citation sources from tool calls
207
207
  logger.debug("Received sources chunk", { chunk });
208
208
  // Add sources to current assistant message for citation rendering
209
- addSourcesToCurrentMessage(chunk.sources);
209
+ addSourcesByToolCallId(chunk.sources);
210
210
  }
211
211
  }
212
212
  // Ensure streaming state is cleared even if no explicit isComplete was received
@@ -259,7 +259,7 @@ export function useChatMessages(client, startSession) {
259
259
  addToolCallToCurrentMessage,
260
260
  updateToolCallInCurrentMessage,
261
261
  addHookNotificationToCurrentMessage,
262
- addSourcesToCurrentMessage,
262
+ addSourcesByToolCallId,
263
263
  ]);
264
264
  /**
265
265
  * Cancel the current agent turn
@@ -446,7 +446,7 @@ export function useChatMessages(client, startSession) {
446
446
  }
447
447
  else if (chunk.type === "sources") {
448
448
  logger.debug("Received sources chunk", { chunk });
449
- addSourcesToCurrentMessage(chunk.sources);
449
+ addSourcesByToolCallId(chunk.sources);
450
450
  }
451
451
  }
452
452
  if (!streamCompleted) {
@@ -497,7 +497,7 @@ export function useChatMessages(client, startSession) {
497
497
  addToolCallToCurrentMessage,
498
498
  updateToolCallInCurrentMessage,
499
499
  addHookNotificationToCurrentMessage,
500
- addSourcesToCurrentMessage,
500
+ addSourcesByToolCallId,
501
501
  ]);
502
502
  return {
503
503
  messages,
@@ -0,0 +1,28 @@
1
+ import type { SubagentMessage } from "../schemas/tool-call.js";
2
+ export interface UseSubagentStreamOptions {
3
+ /** Sub-agent HTTP port */
4
+ port: number;
5
+ /** Sub-agent session ID */
6
+ sessionId: string;
7
+ /** Base host (defaults to localhost) */
8
+ host?: string;
9
+ }
10
+ export interface UseSubagentStreamReturn {
11
+ /** Accumulated messages from the sub-agent */
12
+ messages: SubagentMessage[];
13
+ /** Whether the stream is currently active */
14
+ isStreaming: boolean;
15
+ /** Error message if connection failed */
16
+ error: string | null;
17
+ }
18
+ /**
19
+ * Hook to connect directly to a sub-agent's SSE endpoint and stream messages.
20
+ *
21
+ * This hook:
22
+ * - Connects to the sub-agent's HTTP server at the given port
23
+ * - Subscribes to the /events SSE endpoint with the session ID
24
+ * - Parses incoming session/update notifications
25
+ * - Extracts text chunks and tool calls
26
+ * - Returns accumulated messages for display
27
+ */
28
+ export declare function useSubagentStream(options: UseSubagentStreamOptions | null): UseSubagentStreamReturn;
@@ -0,0 +1,256 @@
1
+ import { createLogger } from "@townco/core";
2
+ import { useCallback, useEffect, useRef, useState } from "react";
3
+ const logger = createLogger("subagent-stream");
4
+ /**
5
+ * Hook to connect directly to a sub-agent's SSE endpoint and stream messages.
6
+ *
7
+ * This hook:
8
+ * - Connects to the sub-agent's HTTP server at the given port
9
+ * - Subscribes to the /events SSE endpoint with the session ID
10
+ * - Parses incoming session/update notifications
11
+ * - Extracts text chunks and tool calls
12
+ * - Returns accumulated messages for display
13
+ */
14
+ export function useSubagentStream(options) {
15
+ const [messages, setMessages] = useState([]);
16
+ // Start as streaming=true if options provided, since we're about to connect
17
+ const [_isStreaming, setIsStreaming] = useState(!!options);
18
+ const [hasCompleted, setHasCompleted] = useState(false);
19
+ const [error, setError] = useState(null);
20
+ const abortControllerRef = useRef(null);
21
+ const currentMessageRef = useRef(null);
22
+ const updateTimeoutRef = useRef(null);
23
+ // Throttled update to prevent excessive re-renders
24
+ const scheduleUpdate = useCallback(() => {
25
+ // If there's already a pending timeout, let it handle the update
26
+ // (it will read the latest currentMessageRef value)
27
+ if (updateTimeoutRef.current)
28
+ return;
29
+ updateTimeoutRef.current = setTimeout(() => {
30
+ updateTimeoutRef.current = null;
31
+ if (currentMessageRef.current) {
32
+ setMessages([{ ...currentMessageRef.current }]);
33
+ }
34
+ }, 250); // Batch updates every 250ms
35
+ }, []);
36
+ // Process incoming SSE message from sub-agent
37
+ // Defined BEFORE connectToSubagent so it's available in the closure
38
+ const processSSEMessage = useCallback((data) => {
39
+ try {
40
+ const message = JSON.parse(data);
41
+ logger.debug("Processing SSE message", {
42
+ method: message.method,
43
+ hasParams: !!message.params,
44
+ });
45
+ // Check if this is a session/update notification
46
+ if (message.method === "session/update" && message.params?.update) {
47
+ const update = message.params.update;
48
+ logger.debug("Got session update", {
49
+ sessionUpdate: update.sessionUpdate,
50
+ });
51
+ if (update.sessionUpdate === "agent_message_chunk") {
52
+ // Handle text chunk
53
+ const content = update.content;
54
+ if (content?.type === "text" && typeof content.text === "string") {
55
+ if (currentMessageRef.current) {
56
+ currentMessageRef.current.content += content.text;
57
+ // Add to contentBlocks - append to last text block or create new one
58
+ const blocks = currentMessageRef.current.contentBlocks ?? [];
59
+ const lastBlock = blocks[blocks.length - 1];
60
+ if (lastBlock && lastBlock.type === "text") {
61
+ lastBlock.text += content.text;
62
+ }
63
+ else {
64
+ blocks.push({ type: "text", text: content.text });
65
+ }
66
+ currentMessageRef.current.contentBlocks = blocks;
67
+ scheduleUpdate();
68
+ }
69
+ }
70
+ }
71
+ else if (update.sessionUpdate === "tool_call") {
72
+ // Handle new tool call
73
+ const toolCall = {
74
+ id: update.toolCallId ?? `tc-${Date.now()}`,
75
+ title: update.title ?? "Tool call",
76
+ prettyName: update._meta?.prettyName,
77
+ icon: update._meta?.icon,
78
+ status: update.status ?? "pending",
79
+ content: [],
80
+ };
81
+ if (currentMessageRef.current) {
82
+ currentMessageRef.current.toolCalls = [
83
+ ...(currentMessageRef.current.toolCalls ?? []),
84
+ toolCall,
85
+ ];
86
+ // Add to contentBlocks for interleaved display
87
+ const blocks = currentMessageRef.current.contentBlocks ?? [];
88
+ blocks.push({ type: "tool_call", toolCall });
89
+ currentMessageRef.current.contentBlocks = blocks;
90
+ scheduleUpdate();
91
+ }
92
+ }
93
+ else if (update.sessionUpdate === "tool_call_update") {
94
+ // Handle tool call update (status change, completion)
95
+ if (currentMessageRef.current?.toolCalls) {
96
+ const toolCallId = update.toolCallId;
97
+ const updateToolCall = (tc) => tc.id === toolCallId
98
+ ? {
99
+ ...tc,
100
+ status: update.status ?? tc.status,
101
+ content: update.content ?? tc.content,
102
+ }
103
+ : tc;
104
+ currentMessageRef.current.toolCalls =
105
+ currentMessageRef.current.toolCalls.map(updateToolCall);
106
+ // Also update in contentBlocks
107
+ if (currentMessageRef.current.contentBlocks) {
108
+ currentMessageRef.current.contentBlocks =
109
+ currentMessageRef.current.contentBlocks.map((block) => block.type === "tool_call"
110
+ ? { ...block, toolCall: updateToolCall(block.toolCall) }
111
+ : block);
112
+ }
113
+ scheduleUpdate();
114
+ }
115
+ }
116
+ }
117
+ }
118
+ catch (err) {
119
+ logger.error("Failed to parse sub-agent SSE message", {
120
+ error: err instanceof Error ? err.message : String(err),
121
+ });
122
+ }
123
+ }, [scheduleUpdate]);
124
+ const connectToSubagent = useCallback(async (port, sessionId, host, protocol) => {
125
+ const baseUrl = `${protocol}//${host}:${port}`;
126
+ logger.info("Connecting to sub-agent SSE", { baseUrl, sessionId });
127
+ setIsStreaming(true);
128
+ setError(null);
129
+ // Create abort controller for cleanup
130
+ const abortController = new AbortController();
131
+ abortControllerRef.current = abortController;
132
+ try {
133
+ logger.info("Fetching SSE endpoint", {
134
+ url: `${baseUrl}/events`,
135
+ sessionId,
136
+ });
137
+ const response = await fetch(`${baseUrl}/events`, {
138
+ method: "GET",
139
+ headers: {
140
+ "X-Session-ID": sessionId,
141
+ },
142
+ signal: abortController.signal,
143
+ });
144
+ logger.info("SSE response received", {
145
+ status: response.status,
146
+ ok: response.ok,
147
+ });
148
+ if (!response.ok) {
149
+ throw new Error(`SSE connection failed: HTTP ${response.status}`);
150
+ }
151
+ if (!response.body) {
152
+ throw new Error("Response body is null");
153
+ }
154
+ logger.info("Sub-agent SSE connection opened, starting to read stream");
155
+ // Read the SSE stream
156
+ const reader = response.body.getReader();
157
+ const decoder = new TextDecoder();
158
+ let buffer = "";
159
+ // Initialize current message
160
+ currentMessageRef.current = {
161
+ id: `subagent-${Date.now()}`,
162
+ content: "",
163
+ toolCalls: [],
164
+ contentBlocks: [],
165
+ isStreaming: true,
166
+ };
167
+ setMessages([currentMessageRef.current]);
168
+ while (true) {
169
+ const { done, value } = await reader.read();
170
+ if (done) {
171
+ logger.debug("Sub-agent SSE stream closed");
172
+ break;
173
+ }
174
+ // Decode the chunk and add to buffer
175
+ buffer += decoder.decode(value, { stream: true });
176
+ // Process complete SSE messages
177
+ const lines = buffer.split("\n");
178
+ buffer = lines.pop() || ""; // Keep incomplete line in buffer
179
+ let currentEvent = { event: "message", data: "" };
180
+ for (const line of lines) {
181
+ if (line.startsWith("event:")) {
182
+ currentEvent.event = line.substring(6).trim();
183
+ }
184
+ else if (line.startsWith("data:")) {
185
+ currentEvent.data = line.substring(5).trim();
186
+ }
187
+ else if (line === "") {
188
+ // Empty line signals end of event
189
+ if (currentEvent.event === "message" && currentEvent.data) {
190
+ processSSEMessage(currentEvent.data);
191
+ }
192
+ // Reset for next event
193
+ currentEvent = { event: "message", data: "" };
194
+ }
195
+ }
196
+ }
197
+ }
198
+ catch (err) {
199
+ if (err instanceof Error && err.name === "AbortError") {
200
+ logger.debug("Sub-agent SSE stream aborted");
201
+ }
202
+ else {
203
+ const errorMessage = err instanceof Error
204
+ ? err.message
205
+ : "Failed to connect to sub-agent";
206
+ logger.error("Sub-agent SSE error", { error: errorMessage });
207
+ setError(errorMessage);
208
+ }
209
+ }
210
+ finally {
211
+ // Mark streaming as complete
212
+ if (currentMessageRef.current) {
213
+ currentMessageRef.current.isStreaming = false;
214
+ setMessages((prev) => prev.map((m) => m.id === currentMessageRef.current?.id
215
+ ? { ...m, isStreaming: false }
216
+ : m));
217
+ }
218
+ setHasCompleted(true);
219
+ setIsStreaming(false);
220
+ abortControllerRef.current = null;
221
+ logger.debug("Sub-agent stream completed");
222
+ }
223
+ }, [processSSEMessage]);
224
+ // Extract values from options (memoized to avoid dependency issues)
225
+ const port = options?.port;
226
+ const sessionId = options?.sessionId;
227
+ const host = options?.host ??
228
+ (typeof window !== "undefined" ? window.location.hostname : "localhost");
229
+ const protocol = typeof window !== "undefined" ? window.location.protocol : "http:";
230
+ // Connect when options change
231
+ useEffect(() => {
232
+ if (!port || !sessionId) {
233
+ return;
234
+ }
235
+ // Reset state for new connection
236
+ setMessages([]);
237
+ setError(null);
238
+ setHasCompleted(false);
239
+ setIsStreaming(true);
240
+ connectToSubagent(port, sessionId, host, protocol);
241
+ // Cleanup on unmount or options change
242
+ return () => {
243
+ if (abortControllerRef.current) {
244
+ abortControllerRef.current.abort();
245
+ abortControllerRef.current = null;
246
+ }
247
+ if (updateTimeoutRef.current) {
248
+ clearTimeout(updateTimeoutRef.current);
249
+ updateTimeoutRef.current = null;
250
+ }
251
+ };
252
+ }, [port, sessionId, host, protocol, connectToSubagent]);
253
+ // Derive streaming status: streaming if we haven't completed yet
254
+ const effectiveIsStreaming = !hasCompleted;
255
+ return { messages, isStreaming: effectiveIsStreaming, error };
256
+ }
@@ -1,7 +1,7 @@
1
1
  import { createLogger } from "@townco/core";
2
2
  import { useEffect } from "react";
3
3
  import { useChatStore } from "../store/chat-store.js";
4
- const _logger = createLogger("use-tool-calls", "debug");
4
+ const logger = createLogger("use-tool-calls", "debug");
5
5
  /**
6
6
  * Hook to track and manage tool calls from ACP sessions
7
7
  *
@@ -16,23 +16,46 @@ export function useToolCalls(client) {
16
16
  const updateToolCall = useChatStore((state) => state.updateToolCall);
17
17
  const addToolCallToCurrentMessage = useChatStore((state) => state.addToolCallToCurrentMessage);
18
18
  const updateToolCallInCurrentMessage = useChatStore((state) => state.updateToolCallInCurrentMessage);
19
+ const addSourcesByToolCallId = useChatStore((state) => state.addSourcesByToolCallId);
19
20
  useEffect(() => {
20
21
  if (!client)
21
22
  return;
23
+ logger.info("useToolCalls: subscribing to session updates");
22
24
  // Subscribe to session updates for tool calls
23
25
  const unsubscribe = client.onSessionUpdate((update) => {
26
+ logger.info("useToolCalls: received session update", {
27
+ type: update.type,
28
+ sessionId: update.sessionId,
29
+ });
24
30
  if (update.type === "tool_call") {
31
+ logger.info("useToolCalls: processing tool_call", {
32
+ toolCallId: update.toolCall.id,
33
+ title: update.toolCall.title,
34
+ });
25
35
  // Add to session-level tool calls (for sidebar)
26
36
  addToolCall(update.sessionId, update.toolCall);
27
37
  // Also add to current assistant message (for inline display)
28
38
  addToolCallToCurrentMessage(update.toolCall);
29
39
  }
30
40
  else if (update.type === "tool_call_update") {
41
+ logger.info("useToolCalls: processing tool_call_update", {
42
+ toolCallId: update.toolCallUpdate.id,
43
+ });
31
44
  // Update session-level tool calls (for sidebar)
32
45
  updateToolCall(update.sessionId, update.toolCallUpdate);
33
46
  // Also update in current assistant message (for inline display)
34
47
  updateToolCallInCurrentMessage(update.toolCallUpdate);
35
48
  }
49
+ else if (update.type === "sources") {
50
+ logger.info("useToolCalls: processing sources", {
51
+ sourcesCount: update.sources.length,
52
+ sourceIds: update.sources.map((s) => s.id),
53
+ toolCallIds: update.sources.map((s) => s.toolCallId),
54
+ });
55
+ // Distribute sources to the correct messages based on toolCallId
56
+ // This handles sources restored during session replay
57
+ addSourcesByToolCallId(update.sources);
58
+ }
36
59
  });
37
60
  return () => {
38
61
  unsubscribe();
@@ -43,6 +66,7 @@ export function useToolCalls(client) {
43
66
  updateToolCall,
44
67
  addToolCallToCurrentMessage,
45
68
  updateToolCallInCurrentMessage,
69
+ addSourcesByToolCallId,
46
70
  ]);
47
71
  return {
48
72
  toolCalls,
@@ -33,6 +33,7 @@ export interface ChatStore {
33
33
  messages: DisplayMessage[];
34
34
  isStreaming: boolean;
35
35
  streamingStartTime: number | null;
36
+ sources: Source[];
36
37
  toolCalls: Record<string, ToolCall[]>;
37
38
  totalBilled: {
38
39
  inputTokens: number;
@@ -56,6 +57,7 @@ export interface ChatStore {
56
57
  addMessage: (message: DisplayMessage) => void;
57
58
  updateMessage: (id: string, updates: Partial<DisplayMessage>) => void;
58
59
  clearMessages: () => void;
60
+ addSources: (sources: Source[]) => void;
59
61
  setIsStreaming: (streaming: boolean) => void;
60
62
  setStreamingStartTime: (time: number | null) => void;
61
63
  addToolCall: (sessionId: string, toolCall: ToolCall) => void;
@@ -64,6 +66,7 @@ export interface ChatStore {
64
66
  updateToolCallInCurrentMessage: (update: ToolCallUpdate) => void;
65
67
  addHookNotificationToCurrentMessage: (notification: HookNotification) => void;
66
68
  addSourcesToCurrentMessage: (sources: Source[]) => void;
69
+ addSourcesByToolCallId: (sources: Source[]) => void;
67
70
  setInputValue: (value: string) => void;
68
71
  setInputSubmitting: (submitting: boolean) => void;
69
72
  addFileAttachment: (file: InputState["attachedFiles"][number]) => void;
@@ -99,6 +99,7 @@ export const useChatStore = create((set) => ({
99
99
  messages: [],
100
100
  isStreaming: false,
101
101
  streamingStartTime: null,
102
+ sources: [],
102
103
  toolCalls: {},
103
104
  totalBilled: {
104
105
  inputTokens: 0,
@@ -229,7 +230,18 @@ export const useChatStore = create((set) => ({
229
230
  currentContext: newCurrentContext,
230
231
  };
231
232
  }),
232
- clearMessages: () => set({ messages: [] }),
233
+ clearMessages: () => set({ messages: [], sources: [] }),
234
+ addSources: (incomingSources) => set((state) => {
235
+ if (!incomingSources || incomingSources.length === 0)
236
+ return state;
237
+ // Merge/dedupe by source id (ids are globally unique within a session)
238
+ const byId = new Map();
239
+ for (const s of state.sources)
240
+ byId.set(s.id, s);
241
+ for (const s of incomingSources)
242
+ byId.set(s.id, s);
243
+ return { sources: Array.from(byId.values()) };
244
+ }),
233
245
  setIsStreaming: (streaming) => set({ isStreaming: streaming }),
234
246
  setStreamingStartTime: (time) => set({ streamingStartTime: time }),
235
247
  addToolCall: (sessionId, toolCall) => set((state) => ({
@@ -421,7 +433,16 @@ export const useChatStore = create((set) => ({
421
433
  isStreaming: false,
422
434
  sources,
423
435
  };
424
- return { messages: [...state.messages, newMessage] };
436
+ // Also store in session-level sources
437
+ const byId = new Map();
438
+ for (const s of state.sources)
439
+ byId.set(s.id, s);
440
+ for (const s of sources)
441
+ byId.set(s.id, s);
442
+ return {
443
+ messages: [...state.messages, newMessage],
444
+ sources: [...byId.values()],
445
+ };
425
446
  }
426
447
  const messages = [...state.messages];
427
448
  const lastAssistantMsg = messages[lastAssistantIndex];
@@ -436,7 +457,117 @@ export const useChatStore = create((set) => ({
436
457
  sourcesCount: sources.length,
437
458
  totalSources: (lastAssistantMsg.sources?.length || 0) + sources.length,
438
459
  });
439
- return { messages };
460
+ // Also store in session-level sources
461
+ const byId = new Map();
462
+ for (const s of state.sources)
463
+ byId.set(s.id, s);
464
+ for (const s of sources)
465
+ byId.set(s.id, s);
466
+ return { messages, sources: [...byId.values()] };
467
+ }),
468
+ addSourcesByToolCallId: (sources) => set((state) => {
469
+ logger.info("addSourcesByToolCallId called", {
470
+ sourcesCount: sources.length,
471
+ messagesCount: state.messages.length,
472
+ sourceDetails: sources.map((s) => ({
473
+ id: s.id,
474
+ toolCallId: s.toolCallId,
475
+ title: s.title?.slice(0, 30),
476
+ })),
477
+ });
478
+ // Log all messages and their tool calls for debugging
479
+ logger.info("Current messages state", {
480
+ messages: state.messages.map((m, i) => ({
481
+ index: i,
482
+ role: m.role,
483
+ contentLength: m.content?.length || 0,
484
+ toolCallIds: m.toolCalls?.map((tc) => tc.id) || [],
485
+ existingSources: m.sources?.length || 0,
486
+ })),
487
+ });
488
+ // Group sources by toolCallId
489
+ const sourcesByToolCallId = new Map();
490
+ for (const source of sources) {
491
+ if (!source.toolCallId) {
492
+ logger.warn("Source missing toolCallId", { sourceId: source.id });
493
+ continue;
494
+ }
495
+ const existing = sourcesByToolCallId.get(source.toolCallId) || [];
496
+ existing.push(source);
497
+ sourcesByToolCallId.set(source.toolCallId, existing);
498
+ }
499
+ if (sourcesByToolCallId.size === 0) {
500
+ logger.warn("No sources with toolCallId to distribute");
501
+ return state;
502
+ }
503
+ // Find messages containing tool calls with matching IDs and add sources
504
+ const messages = [...state.messages];
505
+ let sourcesAdded = 0;
506
+ for (const [toolCallId, toolCallSources] of sourcesByToolCallId) {
507
+ // Find the message containing this tool call
508
+ const messageIndex = messages.findIndex((msg) => msg.role === "assistant" &&
509
+ msg.toolCalls?.some((tc) => tc.id === toolCallId));
510
+ logger.info("Looking for message with toolCallId", {
511
+ toolCallId,
512
+ foundIndex: messageIndex,
513
+ sourcesCount: toolCallSources.length,
514
+ });
515
+ if (messageIndex !== -1) {
516
+ const msg = messages[messageIndex];
517
+ if (msg) {
518
+ messages[messageIndex] = {
519
+ ...msg,
520
+ sources: [...(msg.sources || []), ...toolCallSources],
521
+ };
522
+ sourcesAdded += toolCallSources.length;
523
+ logger.info("Added sources to message", {
524
+ messageIndex,
525
+ messageId: msg.id,
526
+ toolCallId,
527
+ sourcesAdded: toolCallSources.length,
528
+ totalSourcesNow: messages[messageIndex]?.sources?.length,
529
+ messageContentPreview: msg.content.slice(0, 100),
530
+ });
531
+ }
532
+ }
533
+ else {
534
+ // Tool call not found in any message - add to last assistant message as fallback
535
+ logger.warn("Tool call not found in any message, using fallback", {
536
+ toolCallId,
537
+ });
538
+ const lastAssistantIndex = messages.findLastIndex((msg) => msg.role === "assistant");
539
+ if (lastAssistantIndex !== -1) {
540
+ const msg = messages[lastAssistantIndex];
541
+ if (msg) {
542
+ messages[lastAssistantIndex] = {
543
+ ...msg,
544
+ sources: [...(msg.sources || []), ...toolCallSources],
545
+ };
546
+ sourcesAdded += toolCallSources.length;
547
+ logger.info("Added sources to last assistant message (fallback)", {
548
+ lastAssistantIndex,
549
+ sourcesAdded: toolCallSources.length,
550
+ });
551
+ }
552
+ }
553
+ else {
554
+ logger.error("No assistant message found to add sources to");
555
+ }
556
+ }
557
+ }
558
+ logger.info("Finished distributing sources", {
559
+ totalSources: sources.length,
560
+ sourcesAdded,
561
+ toolCallIds: Array.from(sourcesByToolCallId.keys()),
562
+ });
563
+ // Always keep a session-level registry as well, so inline citations can resolve
564
+ // even if we couldn't associate sources to a specific message during replay.
565
+ const byId = new Map();
566
+ for (const s of state.sources)
567
+ byId.set(s.id, s);
568
+ for (const s of sources)
569
+ byId.set(s.id, s);
570
+ return { messages, sources: [...byId.values()] };
440
571
  }),
441
572
  updateToolCall: (sessionId, update) => set((state) => {
442
573
  const sessionToolCalls = state.toolCalls[sessionId] || [];
@@ -5,6 +5,7 @@ import { cn } from "../lib/utils.js";
5
5
  import { SandboxFileSystemProvider } from "../providers/SandboxFileSystemProvider.js";
6
6
  import { FileSystemView } from "./FileSystemView.js";
7
7
  import { SourceListItem } from "./SourceListItem.js";
8
+ import { Tabs, TabsContent, TabsList, TabsTrigger } from "./Tabs.js";
8
9
  import { TodoList } from "./TodoList.js";
9
10
  export const TodoTabContent = React.forwardRef(({ todos = [], className, ...props }, ref) => {
10
11
  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
76
77
  });
77
78
  FilesTabContent.displayName = "FilesTabContent";
78
79
  export const SourcesTabContent = React.forwardRef(({ sources = [], highlightedSourceId, className, ...props }, ref) => {
80
+ const { usedSources, consultedSources } = React.useMemo(() => {
81
+ const used = [];
82
+ const consulted = [];
83
+ for (const s of sources) {
84
+ if (s.usedIn && s.usedIn.length > 0)
85
+ used.push(s);
86
+ else
87
+ consulted.push(s);
88
+ }
89
+ return { usedSources: used, consultedSources: consulted };
90
+ }, [sources]);
91
+ const defaultSourcesTab = usedSources.length > 0 ? "used" : "consulted";
92
+ const [activeSourcesTab, setActiveSourcesTab] = React.useState(defaultSourcesTab);
93
+ // Keep tab selection sensible as sources stream in.
94
+ React.useEffect(() => {
95
+ setActiveSourcesTab(defaultSourcesTab);
96
+ }, [defaultSourcesTab]);
79
97
  // Show empty state if no sources
80
98
  if (sources.length === 0) {
81
99
  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: [
82
100
  _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." })
83
101
  ] }));
84
102
  }
85
- 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))) }));
103
+ return (_jsxs(Tabs, { value: activeSourcesTab, onValueChange: (v) => setActiveSourcesTab(v), className: cn("flex flex-col gap-2", className), children: [
104
+ _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: [
105
+ _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, ")"] })
106
+ ] }) }), _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." })) })
107
+ ] }));
86
108
  });
87
109
  SourcesTabContent.displayName = "SourcesTabContent";
88
110
  export const DatabaseTabContent = React.forwardRef(({ data, className, ...props }, ref) => {