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 b705d904c174..dcb0ddaa2f67 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 @@ -112,6 +112,8 @@ const COMPACTION_RESUME_EVENT_TYPES = new Set([ 'reasoning.available', 'moa.reference', 'moa.aggregating', + 'moa.progress', + 'moa.phase', 'tool.start', 'tool.progress', 'tool.generating', @@ -551,6 +553,36 @@ export function useGatewayEventHandler(deps: GatewayEventDeps) { } else if (event.type === 'moa.aggregating') { // Status transition only; the aggregator's reply arrives via the normal // message stream. No reasoning/transcript mutation here. + if (isActiveEvent) { + setPetActivity({ reasoning: true }) + } + } else if (event.type === 'moa.progress') { + // Live reference fan-out progress ("refs k/n") — surfaced in the same + // reasoning disclosure the references land in. These lines arrive + // BEFORE any moa.reference event (references are only emitted once the + // whole fan-out completes), and the first moa.reference replaces the + // block, so the progress trail is self-cleaning. + if (sessionId && typeof payload?.refs_done === 'number' && typeof payload?.refs_total === 'number') { + const label = coerceGatewayText(payload?.label) + const line = label + ? `◇ MoA refs ${payload.refs_done}/${payload.refs_total} — ${label}\n` + : `◇ MoA refs ${payload.refs_done}/${payload.refs_total}\n` + appendReasoningDelta(sessionId, line, payload.refs_done <= 1) + flushQueuedDeltas(sessionId) + } + + if (isActiveEvent) { + setPetActivity({ reasoning: true }) + } + } else if (event.type === 'moa.phase') { + // Phase transition — currently only phase="aggregator" (fan-out done, + // aggregator acting). Append a one-line marker; the first + // moa.reference that follows replaces the whole block. + if (sessionId && payload?.phase === 'aggregator') { + appendReasoningDelta(sessionId, '◇ MoA aggregating…\n', false) + flushQueuedDeltas(sessionId) + } + if (isActiveEvent) { setPetActivity({ reasoning: true }) } diff --git a/apps/desktop/src/app/session/hooks/use-message-stream/moa-progress-event.test.tsx b/apps/desktop/src/app/session/hooks/use-message-stream/moa-progress-event.test.tsx new file mode 100644 index 000000000000..1e4c0d37867e --- /dev/null +++ b/apps/desktop/src/app/session/hooks/use-message-stream/moa-progress-event.test.tsx @@ -0,0 +1,123 @@ +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 : '' +} + +beforeEach(() => { + handleEvent = null + latestState = null +}) + +afterEach(() => { + cleanup() + vi.restoreAllMocks() +}) + +describe('useMessageStream moa.progress / moa.phase surfacing', () => { + it('shows refs k/n lines in the reasoning block as references complete', async () => { + await mountStream() + + emit('message.start') + emit('moa.progress', { label: 'model-a', refs_done: 1, refs_total: 3 }) + emit('moa.progress', { label: 'model-b', refs_done: 2, refs_total: 3 }) + + const text = reasoningText() + expect(text).toContain('MoA refs 1/3 — model-a') + expect(text).toContain('MoA refs 2/3 — model-b') + }) + + it('restarts the progress block on the first ref of a new fan-out', async () => { + await mountStream() + + emit('message.start') + emit('moa.progress', { label: 'stale', refs_done: 1, refs_total: 2 }) + emit('moa.progress', { label: 'stale-2', refs_done: 2, refs_total: 2 }) + // A later turn's fan-out starts over at 1/N — old lines must not linger. + emit('moa.progress', { label: 'fresh', refs_done: 1, refs_total: 2 }) + + const text = reasoningText() + expect(text).toContain('MoA refs 1/2 — fresh') + expect(text).not.toContain('stale') + }) + + it('appends the aggregating marker on moa.phase and ignores unknown phases', async () => { + await mountStream() + + emit('message.start') + emit('moa.progress', { label: 'model-a', refs_done: 1, refs_total: 1 }) + emit('moa.phase', { phase: 'reference', refs_done: 1, refs_total: 1 }) + expect(reasoningText()).not.toContain('aggregating') + + emit('moa.phase', { aggregator: 'agg-model', phase: 'aggregator', refs_done: 1, refs_total: 1 }) + expect(reasoningText()).toContain('MoA aggregating…') + }) + + it('a following moa.reference replaces the progress trail (self-cleaning)', async () => { + await mountStream() + + emit('message.start') + emit('moa.progress', { label: 'model-a', refs_done: 1, refs_total: 1 }) + emit('moa.phase', { phase: 'aggregator', refs_done: 1, refs_total: 1 }) + emit('moa.reference', { count: 1, index: 1, label: 'model-a', text: 'advice-a' }) + + const text = reasoningText() + expect(text).toContain('Reference 1/1 — model-a') + expect(text).not.toContain('MoA refs') + expect(text).not.toContain('aggregating') + }) +}) diff --git a/apps/desktop/src/lib/chat-messages.ts b/apps/desktop/src/lib/chat-messages.ts index 00d1c73a64d4..8541c0bb8098 100644 --- a/apps/desktop/src/lib/chat-messages.ts +++ b/apps/desktop/src/lib/chat-messages.ts @@ -94,6 +94,10 @@ export type GatewayEventPayload = { label?: string index?: number aggregator?: string + // moa.progress / moa.phase (Mixture of Agents fan-out progress relay) + refs_done?: number + refs_total?: number + phase?: string // message.complete — signals the final text was already previewed via // interim_assistant_callback, so the UI can settle instead of duplicating. response_previewed?: boolean diff --git a/ui-tui/src/__tests__/moaProgressActivity.test.ts b/ui-tui/src/__tests__/moaProgressActivity.test.ts new file mode 100644 index 000000000000..1d76556d2028 --- /dev/null +++ b/ui-tui/src/__tests__/moaProgressActivity.test.ts @@ -0,0 +1,100 @@ +import { beforeEach, describe, expect, it } from 'vitest' + +import { createGatewayEventHandler } from '../app/createGatewayEventHandler.js' +import { resetOverlayState } from '../app/overlayStore.js' +import { turnController } from '../app/turnController.js' +import { getTurnState, resetTurnState } from '../app/turnStore.js' +import { patchUiState, resetUiState } from '../app/uiStore.js' +import type { Msg } from '../types.js' + +const ref = (current: T) => ({ current }) + +const buildCtx = (appended: Msg[]) => + ({ + composer: { + dequeue: () => undefined, + queueEditRef: ref(null), + sendQueued: () => undefined, + setInput: () => undefined + }, + gateway: { + gw: { request: () => undefined }, + rpc: async () => null + }, + session: { + STARTUP_RESUME_ID: '', + colsRef: ref(80), + newSession: () => undefined, + resetSession: () => undefined, + resumeById: () => undefined, + setCatalog: () => undefined + }, + submission: { + submitRef: { current: () => undefined } + }, + system: { + bellOnComplete: false, + sys: () => undefined + }, + transcript: { + appendMessage: (msg: Msg) => appended.push(msg), + panel: () => undefined, + setHistoryItems: () => undefined + }, + voice: { + setProcessing: () => undefined, + setRecording: () => undefined, + setVoiceEnabled: () => undefined + } + }) as any + +const activityTexts = () => getTurnState().activity.map(item => item.text) + +describe('moa.progress / moa.phase activity surface', () => { + beforeEach(() => { + resetOverlayState() + resetUiState() + resetTurnState() + turnController.fullReset() + patchUiState({ showReasoning: true }) + }) + + it('shows "MoA: refs k/n" as each reference completes, replacing in place', () => { + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ payload: {}, type: 'message.start' } as any) + onEvent({ payload: { label: 'model-a', refs_done: 1, refs_total: 3 }, type: 'moa.progress' } as any) + + expect(activityTexts()).toContain('MoA: refs 1/3') + + onEvent({ payload: { label: 'model-b', refs_done: 2, refs_total: 3 }, type: 'moa.progress' } as any) + + const texts = activityTexts() + expect(texts).toContain('MoA: refs 2/3') + // Replaced in place — the stale 1/3 line must not linger alongside 2/3. + expect(texts).not.toContain('MoA: refs 1/3') + }) + + it('swaps the progress line for aggregator copy on moa.phase', () => { + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ payload: {}, type: 'message.start' } as any) + onEvent({ payload: { label: 'model-a', refs_done: 3, refs_total: 3 }, type: 'moa.progress' } as any) + onEvent({ payload: { phase: 'aggregator', refs_done: 3, refs_total: 3 }, type: 'moa.phase' } as any) + + const texts = activityTexts() + expect(texts).toContain('MoA: aggregating…') + expect(texts).not.toContain('MoA: refs 3/3') + }) + + it('ignores malformed payloads (missing counters / unknown phase)', () => { + const onEvent = createGatewayEventHandler(buildCtx([])) + + onEvent({ payload: {}, type: 'message.start' } as any) + const before = activityTexts().length + onEvent({ payload: { label: 'model-a' }, type: 'moa.progress' } as any) + onEvent({ payload: { phase: 'reference' }, type: 'moa.phase' } as any) + + expect(activityTexts().length).toBe(before) + }) +}) diff --git a/ui-tui/src/app/createGatewayEventHandler.ts b/ui-tui/src/app/createGatewayEventHandler.ts index 8527e083a580..d8630e216736 100644 --- a/ui-tui/src/app/createGatewayEventHandler.ts +++ b/ui-tui/src/app/createGatewayEventHandler.ts @@ -1005,6 +1005,25 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: // through the normal message stream. No committed transcript entry. return + case 'moa.progress': + // Live fan-out progress — one activity line, replaced in place as each + // reference completes ("MoA: refs 2/3"), so the user sees movement + // during the (potentially long) reference phase without transcript spam. + if (typeof ev.payload?.refs_done === 'number' && typeof ev.payload?.refs_total === 'number') { + turnController.pushActivity(`MoA: refs ${ev.payload.refs_done}/${ev.payload.refs_total}`, 'info', 'MoA') + } + + return + + case 'moa.phase': + // Phase transition — currently only phase="aggregator" (fan-out done, + // aggregator acting). Swap the progress line for aggregator copy. + if (ev.payload?.phase === 'aggregator') { + turnController.pushActivity('MoA: aggregating…', 'info', 'MoA') + } + + return + case 'tool.progress': if (ev.payload?.preview && ev.payload.name) { turnController.recordToolProgress(ev.payload.name, ev.payload.preview) diff --git a/ui-tui/src/gatewayTypes.ts b/ui-tui/src/gatewayTypes.ts index cf84abfddf7f..ff287440120e 100644 --- a/ui-tui/src/gatewayTypes.ts +++ b/ui-tui/src/gatewayTypes.ts @@ -608,6 +608,16 @@ export type GatewayEvent = type: 'moa.reference' } | { payload?: { aggregator?: string }; session_id?: string; type: 'moa.aggregating' } + | { + payload?: { label?: string; refs_done?: number; refs_total?: number } + session_id?: string + type: 'moa.progress' + } + | { + payload?: { aggregator?: string; phase?: string; refs_done?: number; refs_total?: number } + session_id?: string + type: 'moa.phase' + } | { payload: { name?: string; preview?: string }; session_id?: string; type: 'tool.progress' } | { payload: { name?: string }; session_id?: string; type: 'tool.generating' } | {