diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts index 30ceb1510562..b705d904c174 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream/gateway-event.ts @@ -524,7 +524,25 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { const cnt = typeof payload?.count === 'number' ? payload.count : undefined const header = idx && cnt ? `◇ Reference ${idx}/${cnt} — ${label}` : `◇ Reference — ${label}` const body = coerceThinkingText(payload?.text) - appendReasoningDelta(sessionId, `${header}\n${body}\n\n`, true) + const text = `${header}\n${body}\n\n` + if (idx === undefined || idx <= 1) { + // First reference: clear any stale reasoning left over from + // before this turn's references start, same as before. + appendReasoningDelta(sessionId, text, true) + } else { + // Later references must accumulate, not replace — otherwise + // each new reference wipes out the ones already shown (#64658). + // Queue-then-flush (rather than the streamed/batched queue path) + // applies it immediately, since each reference arrives as one + // complete block rather than incremental tokens. reasoning.delta + // cannot be mid-flight here: MoAChatCompletions.reference_callback + // (agent/moa_loop.py) fires "moa.reference" once per reference's + // already-complete text, with no concurrent token stream for the + // reference-gathering phase, so there is no in-flight delta to + // collide with in the shared queue bucket. + appendReasoningDelta(sessionId, text, false) + flushQueuedDeltas(sessionId) + } } if (isActiveEvent) { diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/moa-reference-event.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/moa-reference-event.test.tsx new file mode 100644 index 000000000000..b8237e09d22d --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/moa-reference-event.test.tsx @@ -0,0 +1,113 @@ +import { QueryClient } from '@tanstack/react-query' +import { act, cleanup, render, waitFor } from '@testing-library/react' +import { useEffect, useRef } from 'react' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import type { ClientSessionState } from '@/app/types' +import { createClientSessionState } from '@/lib/chat-runtime' +import type { RpcEvent } from '@/types/hermes' + +import { useMessageStream } from './index' + +const SID = 'session-1' +let handleEvent: ((event: RpcEvent) => void) | null = null +let latestState: ClientSessionState | null = null + +function Harness() { + const activeSessionIdRef = useRef(SID) + const sessionStateByRuntimeIdRef = useRef(new Map()) + const queryClientRef = useRef(new QueryClient()) + + const stream = useMessageStream({ + activeSessionIdRef, + hydrateFromStoredSession: vi.fn(async () => undefined), + queryClient: queryClientRef.current, + refreshHermesConfig: vi.fn(async () => undefined), + refreshSessions: vi.fn(async () => undefined), + sessionStateByRuntimeIdRef, + updateSessionState: (sessionId, updater) => { + const current = sessionStateByRuntimeIdRef.current.get(sessionId) ?? createClientSessionState() + const next = updater(current) + sessionStateByRuntimeIdRef.current.set(sessionId, next) + latestState = next + + return next + } + }) + + useEffect(() => { + handleEvent = stream.handleGatewayEvent + }, [stream.handleGatewayEvent]) + + return null +} + +async function mountStream() { + render() + await waitFor(() => expect(handleEvent).not.toBeNull()) +} + +function emit(type: RpcEvent['type'], payload: RpcEvent['payload'] = {}) { + act(() => handleEvent!({ payload, session_id: SID, type })) +} + +function reasoningText(): string { + const message = latestState?.messages.at(-1) + const part = message?.parts.find(p => p.type === 'reasoning') + + return part?.type === 'reasoning' ? part.text : '' +} + +describe('useMessageStream moa.reference accumulation (#64658)', () => { + beforeEach(() => { + handleEvent = null + latestState = null + }) + + afterEach(() => { + cleanup() + vi.restoreAllMocks() + }) + + it('keeps every reference model labelled block instead of only the latest one', async () => { + await mountStream() + + emit('moa.reference', { count: 2, index: 1, label: 'model-a', text: 'advice-a' }) + emit('moa.reference', { count: 2, index: 2, label: 'model-b', text: 'advice-b' }) + + const text = reasoningText() + + expect(text).toContain('model-a') + expect(text).toContain('advice-a') + expect(text).toContain('model-b') + expect(text).toContain('advice-b') + }) + + it('handles a single-reference MoA turn (count=1) without regression', async () => { + await mountStream() + + emit('moa.reference', { count: 1, index: 1, label: 'model-a', text: 'only-advice' }) + + const text = reasoningText() + + expect(text).toContain('model-a') + expect(text).toContain('only-advice') + }) + + it('accumulates three or more references in order', async () => { + await mountStream() + + emit('moa.reference', { count: 3, index: 1, label: 'model-a', text: 'advice-a' }) + emit('moa.reference', { count: 3, index: 2, label: 'model-b', text: 'advice-b' }) + emit('moa.reference', { count: 3, index: 3, label: 'model-c', text: 'advice-c' }) + + const text = reasoningText() + const orderOk = + text.indexOf('advice-a') < text.indexOf('advice-b') && text.indexOf('advice-b') < text.indexOf('advice-c') + + expect(text).toContain('advice-a') + expect(text).toContain('advice-b') + expect(text).toContain('advice-c') + expect(orderOk).toBe(true) + }) +})