diff --git a/app/src/providers/ChatRuntimeProvider.tsx b/app/src/providers/ChatRuntimeProvider.tsx index ae9c709df..8c9167ac4 100644 --- a/app/src/providers/ChatRuntimeProvider.tsx +++ b/app/src/providers/ChatRuntimeProvider.tsx @@ -4,6 +4,7 @@ import { useCallback, useEffect, useRef } from 'react'; import { requestUsageRefresh } from '../hooks/usageRefresh'; import { useRefetchSnapshotOnTurnEnd } from '../hooks/useRefetchSnapshotOnTurnEnd'; import { + type ChatDoneEvent, type ChatInferenceStartEvent, type ChatIterationStartEvent, type ChatSegmentEvent, @@ -42,6 +43,8 @@ import { formatTimelineEntry, promptFromArgsBuffer } from '../utils/toolTimeline const logChatRuntime = debug('openhuman:chat-runtime'); +type SegmentDelivery = { segments: Map }; + function rtLog(message: string, fields?: Record) { if (IS_PROD) return; if (fields && Object.keys(fields).length > 0) { @@ -54,6 +57,29 @@ function rtLog(message: string, fields?: Record | undefined { + return event.citations?.length ? { citations: event.citations } : undefined; +} + const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { const dispatch = useAppDispatch(); const { refetch: refetchSnapshot } = useRefetchSnapshotOnTurnEnd(); @@ -67,6 +93,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { ); const seenChatEventsRef = useRef>(new Map()); + const segmentDeliveriesRef = useRef>(new Map()); const proactiveThreadCreationPromiseRef = useRef | null>(null); const proactiveDispatchQueueRef = useRef>(Promise.resolve()); const toolTimelineRef = useRef(toolTimelineByThread); @@ -186,6 +213,24 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { return { ...entry, displayName: formatted.title, detail: formatted.detail }; }; + const finishChatDoneTurn = (event: ChatDoneEvent, path: string) => { + rtLog('refresh_usage_counter', { + thread: event.thread_id, + request: event.request_id, + reason: 'chat_done', + }); + requestUsageRefresh(); + rtLog('snapshot_refetch_queued', { + thread: event.thread_id, + request: event.request_id, + reason: 'chat_done', + path, + }); + refetchSnapshot(); + dispatch(endInferenceTurn({ threadId: event.thread_id })); + dispatch(setActiveThread(null)); + }; + const findPendingDelegationContext = ( entries: ToolTimelineEntry[], round: number @@ -484,9 +529,16 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { !markChatEventSeen(eventKey, { threadId: event.thread_id, requestId: event.request_id }) ) return; + const content = segmentText(event); + const deliveryKey = segmentDeliveryKey(event.thread_id, event.request_id); + const delivery = segmentDeliveriesRef.current.get(deliveryKey) ?? { + segments: new Map(), + }; + delivery.segments.set(event.segment_index, content); + segmentDeliveriesRef.current.set(deliveryKey, delivery); void dispatch( addInferenceResponse({ - content: segmentText(event), + content, threadId: event.thread_id, extraMetadata: event.citations?.length ? { citations: event.citations } : undefined, }) @@ -607,6 +659,11 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { output_tokens: event.total_output_tokens, }); + const deliveryKey = segmentDeliveryKey(event.thread_id, event.request_id); + const segmentDelivery = segmentDeliveriesRef.current.get(deliveryKey); + const completeSegmentDelivery = hasCompleteSegmentDelivery(event, segmentDelivery); + segmentDeliveriesRef.current.delete(deliveryKey); + dispatch( recordChatTurnUsage({ inputTokens: event.total_input_tokens, @@ -630,9 +687,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { addInferenceResponse({ content: event.full_response, threadId: event.thread_id, - extraMetadata: event.citations?.length - ? { citations: event.citations } - : undefined, + extraMetadata: chatDoneExtraMetadata(event), }) ).unwrap(); void dispatch( @@ -648,21 +703,42 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { error: error instanceof Error ? error.message : String(error), }); } - rtLog('refresh_usage_counter', { - thread: event.thread_id, - request: event.request_id, - reason: 'chat_done', - }); - requestUsageRefresh(); - rtLog('snapshot_refetch_queued', { - thread: event.thread_id, - request: event.request_id, - reason: 'chat_done', - path: 'proactive', - }); - refetchSnapshot(); - dispatch(endInferenceTurn({ threadId: event.thread_id })); - dispatch(setActiveThread(null)); + finishChatDoneTurn(event, 'proactive'); + })(); + return; + } + + if (!completeSegmentDelivery && event.full_response.length > 0) { + rtLog('chat_done_segment_reconcile', { + thread: event.thread_id, + request: event.request_id, + expected: event.segment_total, + received: segmentDelivery?.segments.size ?? 0, + full_len: event.full_response.length, + }); + void (async () => { + try { + await dispatch( + addInferenceResponse({ + content: event.full_response, + threadId: event.thread_id, + extraMetadata: chatDoneExtraMetadata(event), + }) + ).unwrap(); + void dispatch( + generateThreadTitleIfNeeded({ + threadId: event.thread_id, + assistantMessage: event.full_response, + }) + ); + } catch (error) { + rtLog('chat_done_reconcile_append_failed', { + thread: event.thread_id, + request: event.request_id, + error: error instanceof Error ? error.message : String(error), + }); + } + finishChatDoneTurn(event, 'segment_reconcile'); })(); return; } @@ -673,21 +749,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { assistantMessage: event.full_response, }) ); - rtLog('refresh_usage_counter', { - thread: event.thread_id, - request: event.request_id, - reason: 'chat_done', - }); - requestUsageRefresh(); - rtLog('snapshot_refetch_queued', { - thread: event.thread_id, - request: event.request_id, - reason: 'chat_done', - path: 'ordinary', - }); - refetchSnapshot(); - dispatch(endInferenceTurn({ threadId: event.thread_id })); - dispatch(setActiveThread(null)); + finishChatDoneTurn(event, 'ordinary'); }, onError: event => { const eventKey = `error:${event.thread_id}:${event.request_id ?? 'none'}:${event.error_type}:${event.message}`; @@ -702,6 +764,7 @@ const ChatRuntimeProvider = ({ children }: { children: React.ReactNode }) => { err: event.error_type, }); + segmentDeliveriesRef.current.delete(segmentDeliveryKey(event.thread_id, event.request_id)); dispatch(clearInferenceStatusForThread({ threadId: event.thread_id })); dispatch(clearStreamingAssistantForThread({ threadId: event.thread_id })); diff --git a/app/src/providers/__tests__/ChatRuntimeProvider.test.tsx b/app/src/providers/__tests__/ChatRuntimeProvider.test.tsx index 1acb4a2e0..90cba448d 100644 --- a/app/src/providers/__tests__/ChatRuntimeProvider.test.tsx +++ b/app/src/providers/__tests__/ChatRuntimeProvider.test.tsx @@ -127,7 +127,7 @@ describe('ChatRuntimeProvider — dedupe, proactive resolution, mid-turn invaria expect(usage.turns).toBe(1); // Snapshot refetch fired exactly once on the first chat_done — issue #924. - expect(mockRefetchSnapshot).toHaveBeenCalledTimes(1); + await waitFor(() => expect(mockRefetchSnapshot).toHaveBeenCalledTimes(1)); }); it('processes tool_call for different rounds as distinct events', () => { @@ -245,6 +245,85 @@ describe('ChatRuntimeProvider — dedupe, proactive resolution, mid-turn invaria }); describe('mid-turn streaming invariants', () => { + it('reconciles missing segment events from chat_done.full_response', async () => { + const listeners = renderProvider(); + + act(() => { + listeners.onSegment?.({ + thread_id: 't-segmented', + request_id: 'r-segmented', + full_response: 'Part one.', + segment_index: 0, + segment_total: 2, + }); + }); + + await waitFor(() => + expect(threadApi.appendMessage).toHaveBeenCalledWith( + 't-segmented', + expect.objectContaining({ content: 'Part one.', sender: 'agent' }) + ) + ); + + act(() => { + listeners.onDone?.({ + thread_id: 't-segmented', + request_id: 'r-segmented', + full_response: 'Part one.\n\nPart two.', + rounds_used: 1, + total_input_tokens: 10, + total_output_tokens: 20, + segment_total: 2, + }); + }); + + await waitFor(() => + expect(threadApi.appendMessage).toHaveBeenCalledWith( + 't-segmented', + expect.objectContaining({ content: 'Part one.\n\nPart two.', sender: 'agent' }) + ) + ); + expect(threadApi.appendMessage).toHaveBeenCalledTimes(2); + }); + + it('does not duplicate chat_done.full_response when all segments arrived', async () => { + const listeners = renderProvider(); + + act(() => { + listeners.onSegment?.({ + thread_id: 't-complete', + request_id: 'r-complete', + full_response: 'Part one.', + segment_index: 0, + segment_total: 2, + }); + listeners.onSegment?.({ + thread_id: 't-complete', + request_id: 'r-complete', + full_response: 'Part two.', + segment_index: 1, + segment_total: 2, + }); + }); + + await waitFor(() => expect(threadApi.appendMessage).toHaveBeenCalledTimes(2)); + + act(() => { + listeners.onDone?.({ + thread_id: 't-complete', + request_id: 'r-complete', + full_response: 'Part one.\n\nPart two.', + rounds_used: 1, + total_input_tokens: 10, + total_output_tokens: 20, + segment_total: 2, + }); + }); + + await waitFor(() => expect(mockRefetchSnapshot).toHaveBeenCalledTimes(1)); + expect(threadApi.appendMessage).toHaveBeenCalledTimes(2); + }); + it('accumulates text_delta chunks within the same request_id', () => { const listeners = renderProvider();