diff --git a/tui_gateway/server.py b/tui_gateway/server.py index e3fb585135..8db15b6f67 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1342,6 +1342,7 @@ def _(rid, params: dict) -> dict: stream_callback=_stream, ) + last_reasoning = None if isinstance(result, dict): if isinstance(result.get("messages"), list): with session["history_lock"]: @@ -1350,11 +1351,16 @@ def _(rid, params: dict) -> dict: session["history_version"] = history_version + 1 raw = result.get("final_response", "") status = "interrupted" if result.get("interrupted") else "error" if result.get("error") else "complete" + lr = result.get("last_reasoning") + if isinstance(lr, str) and lr.strip(): + last_reasoning = lr.strip() else: raw = str(result) status = "complete" payload = {"text": raw, "usage": _get_usage(agent), "status": status} + if last_reasoning: + payload["reasoning"] = last_reasoning rendered = render_message(raw, cols) if rendered: payload["rendered"] = rendered diff --git a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts index be27d5347f..c4f5628ca7 100644 --- a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts +++ b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts @@ -2,7 +2,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { createGatewayEventHandler } from '../app/createGatewayEventHandler.js' import { resetOverlayState } from '../app/overlayStore.js' -import { resetUiState } from '../app/uiStore.js' +import { getUiState, patchUiState, resetUiState } from '../app/uiStore.js' import { estimateTokensRough } from '../lib/text.js' import type { Msg } from '../types.js' @@ -348,4 +348,178 @@ describe('createGatewayEventHandler', () => { expect(appended[0]?.thinking).toBe(streamed) expect(appended[0]?.thinkingTokens).toBe(estimateTokensRough(streamed)) }) + + it('uses message.complete reasoning when no streamed reasoning ref', () => { + const appended: Msg[] = [] + const fromServer = 'recovered from last_reasoning' + + const refs = { + activeToolsRef: ref([] as { context?: string; id: string; name: string; startedAt?: number }[]), + bufRef: ref(''), + interruptedRef: ref(false), + lastStatusNoteRef: ref(''), + persistedToolLabelsRef: ref(new Set()), + protocolWarnedRef: ref(false), + reasoningRef: ref(''), + statusTimerRef: ref | null>(null), + toolTokenAccRef: ref(0), + toolCompleteRibbonRef: ref(null), + turnToolsRef: ref([] as string[]) + } + + const onEvent = createGatewayEventHandler({ + composer: { + dequeue: () => undefined, + queueEditRef: ref(null), + sendQueued: vi.fn() + }, + gateway: { + gw: { request: vi.fn() } as any, + rpc: vi.fn(async () => null) + }, + session: { + STARTUP_RESUME_ID: '', + colsRef: ref(80), + newSession: vi.fn(), + resetSession: vi.fn(), + setCatalog: vi.fn() + }, + system: { + bellOnComplete: false, + sys: vi.fn() + }, + transcript: { + appendMessage: (msg: Msg) => appended.push(msg), + setHistoryItems: vi.fn() + }, + turn: { + actions: { + clearReasoning: vi.fn(() => { + refs.reasoningRef.current = '' + refs.toolTokenAccRef.current = 0 + }), + endReasoningPhase: vi.fn(), + idle: vi.fn(() => { + refs.activeToolsRef.current = [] + }), + pruneTransient: vi.fn(), + pulseReasoningStreaming: vi.fn(), + pushActivity: vi.fn(), + pushTrail: vi.fn(), + scheduleReasoning: vi.fn(), + scheduleStreaming: vi.fn(), + setActivity: vi.fn(), + setReasoningTokens: vi.fn(), + setStreaming: vi.fn(), + setToolTokens: vi.fn(), + setTools: vi.fn(), + setTurnTrail: vi.fn() + }, + refs + } + } as any) + + onEvent({ + payload: { reasoning: fromServer, text: 'final answer' }, + type: 'message.complete' + } as any) + + expect(appended).toHaveLength(1) + expect(appended[0]?.thinking).toBe(fromServer) + expect(appended[0]?.thinkingTokens).toBe(estimateTokensRough(fromServer)) + }) + + it('merges message.complete usage into existing context fields', () => { + const appended: Msg[] = [] + + patchUiState({ + usage: { + calls: 1, + context_max: 100_000, + context_percent: 12, + context_used: 12_000, + input: 10, + output: 20, + total: 30 + } + }) + + const refs = { + activeToolsRef: ref([] as { context?: string; id: string; name: string; startedAt?: number }[]), + bufRef: ref(''), + interruptedRef: ref(false), + lastStatusNoteRef: ref(''), + persistedToolLabelsRef: ref(new Set()), + protocolWarnedRef: ref(false), + reasoningRef: ref(''), + statusTimerRef: ref | null>(null), + toolTokenAccRef: ref(0), + toolCompleteRibbonRef: ref(null), + turnToolsRef: ref([] as string[]) + } + + const onEvent = createGatewayEventHandler({ + composer: { + dequeue: () => undefined, + queueEditRef: ref(null), + sendQueued: vi.fn() + }, + gateway: { + gw: { request: vi.fn() } as any, + rpc: vi.fn(async () => null) + }, + session: { + STARTUP_RESUME_ID: '', + colsRef: ref(80), + newSession: vi.fn(), + resetSession: vi.fn(), + setCatalog: vi.fn() + }, + system: { + bellOnComplete: false, + sys: vi.fn() + }, + transcript: { + appendMessage: (msg: Msg) => appended.push(msg), + setHistoryItems: vi.fn() + }, + turn: { + actions: { + clearReasoning: vi.fn(() => { + refs.reasoningRef.current = '' + }), + endReasoningPhase: vi.fn(), + idle: vi.fn(), + pruneTransient: vi.fn(), + pulseReasoningStreaming: vi.fn(), + pushActivity: vi.fn(), + pushTrail: vi.fn(), + scheduleReasoning: vi.fn(), + scheduleStreaming: vi.fn(), + setActivity: vi.fn(), + setReasoningTokens: vi.fn(), + setStreaming: vi.fn(), + setToolTokens: vi.fn(), + setTools: vi.fn(), + setTurnTrail: vi.fn() + }, + refs + } + } as any) + + onEvent({ + payload: { + text: 'ok', + usage: { calls: 2, input: 50, output: 60, total: 110 } + }, + type: 'message.complete' + } as any) + + const u = getUiState().usage + expect(u.input).toBe(50) + expect(u.total).toBe(110) + expect(u.context_max).toBe(100_000) + expect(u.context_used).toBe(12_000) + expect(u.context_percent).toBe(12) + }) }) diff --git a/ui-tui/src/app/createGatewayEventHandler.ts b/ui-tui/src/app/createGatewayEventHandler.ts index 86bacdecb6..5541bf513f 100644 --- a/ui-tui/src/app/createGatewayEventHandler.ts +++ b/ui-tui/src/app/createGatewayEventHandler.ts @@ -631,7 +631,9 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: const p = ev.payload const finalText = (p?.rendered ?? p?.text ?? bufRef.current).trimStart() const persisted = persistedToolLabelsRef.current - const savedReasoning = reasoningRef.current.trim() + const streamedReasoning = reasoningRef.current.trim() + const payloadReasoning = String(p?.reasoning ?? '').trim() + const savedReasoning = streamedReasoning || payloadReasoning const savedReasoningTokens = savedReasoning ? estimateTokensRough(savedReasoning) : 0 const savedToolTokens = toolTokenAccRef.current @@ -666,7 +668,7 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: setStatus('ready') if (p?.usage) { - patchUiState({ usage: p.usage }) + patchUiState(state => ({ ...state, usage: { ...state.usage, ...p.usage } })) } if (queueEditRef.current !== null) { diff --git a/ui-tui/src/gatewayTypes.ts b/ui-tui/src/gatewayTypes.ts index 6eff78c58c..7fab065971 100644 --- a/ui-tui/src/gatewayTypes.ts +++ b/ui-tui/src/gatewayTypes.ts @@ -199,5 +199,9 @@ export type GatewayEvent = | { payload: SubagentEventPayload; session_id?: string; type: 'subagent.progress' } | { payload: SubagentEventPayload; session_id?: string; type: 'subagent.complete' } | { payload: { rendered?: string; text?: string }; session_id?: string; type: 'message.delta' } - | { payload?: { rendered?: string; text?: string; usage?: Usage }; session_id?: string; type: 'message.complete' } + | { + payload?: { reasoning?: string; rendered?: string; text?: string; usage?: Usage } + session_id?: string + type: 'message.complete' + } | { payload?: { message?: string }; session_id?: string; type: 'error' }