fix: ctx usage display
This commit is contained in:
Brooklyn Nicholson
parent
f81dba0da2
commit
39b1336d1f
@ -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
|
||||
|
||||
@ -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<string>()),
|
||||
protocolWarnedRef: ref(false),
|
||||
reasoningRef: ref(''),
|
||||
statusTimerRef: ref<ReturnType<typeof setTimeout> | null>(null),
|
||||
toolTokenAccRef: ref(0),
|
||||
toolCompleteRibbonRef: ref(null),
|
||||
turnToolsRef: ref([] as string[])
|
||||
}
|
||||
|
||||
const onEvent = createGatewayEventHandler({
|
||||
composer: {
|
||||
dequeue: () => undefined,
|
||||
queueEditRef: ref<number | null>(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<string>()),
|
||||
protocolWarnedRef: ref(false),
|
||||
reasoningRef: ref(''),
|
||||
statusTimerRef: ref<ReturnType<typeof setTimeout> | null>(null),
|
||||
toolTokenAccRef: ref(0),
|
||||
toolCompleteRibbonRef: ref(null),
|
||||
turnToolsRef: ref([] as string[])
|
||||
}
|
||||
|
||||
const onEvent = createGatewayEventHandler({
|
||||
composer: {
|
||||
dequeue: () => undefined,
|
||||
queueEditRef: ref<number | null>(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)
|
||||
})
|
||||
})
|
||||
|
||||
@ -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) {
|
||||
|
||||
@ -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' }
|
||||
|
||||
Loading…
Reference in New Issue
Block a user