fix(chat): refresh context after runtime transitions

This commit is contained in:
程序员阿江(Relakkes)
2026-08-03 04:06:59 +08:00
parent 2262973a48
commit b4807697f8
9 changed files with 299 additions and 40 deletions
+4 -1
View File
@@ -1395,7 +1395,10 @@ export function ChatInput({ variant = 'default', compact = false }: ChatInputPro
runtimeSelectionKey={runtimeSelectionKey}
fallbackModelLabel={runtimeModelLabel}
compact={useCompactControls}
refreshNonce={sessionState?.compactCount ?? 0}
refreshNonce={
(sessionState?.compactCount ?? 0) +
(sessionState?.runtimeConfigReadyCount ?? 0)
}
/>
)}
{!isMemberSession && activeTabId && (
@@ -188,6 +188,42 @@ describe('ContextUsageIndicator request behavior', () => {
expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(2)
})
it('loads context when a new session finishes its first turn', async () => {
sessionsApiMock.getInspection
.mockResolvedValueOnce({
active: true,
status: baseInspection.status,
errors: { context: 'Context is not ready' },
})
.mockResolvedValueOnce(baseInspection)
const { rerender } = render(
<ContextUsageIndicator
sessionId="session-1"
chatState="thinking"
messageCount={1}
/>,
)
await waitFor(() => {
expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(1)
})
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('--')
rerender(
<ContextUsageIndicator
sessionId="session-1"
chatState="idle"
messageCount={2}
/>,
)
await waitFor(() => {
expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(2)
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('21%')
})
})
it('keeps the last context visible while the switched runtime is still starting', async () => {
const nextInspection = deferred<typeof baseInspection>()
sessionsApiMock.getInspection
@@ -218,12 +254,12 @@ describe('ContextUsageIndicator request behavior', () => {
/>,
)
// A runtime switch restarts the CLI. Keep the last confirmed percentage
// in place instead of replacing it with a long-running spinner while the
// replacement control channel comes online.
// A runtime switch restarts the CLI. Keep the last confirmed model and
// percentage together instead of relabeling stale usage as the new model.
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('21%')
expect(screen.queryByLabelText('Context usage loading')).not.toBeInTheDocument()
expect(screen.getByText('deepseek-reasoner')).toBeInTheDocument()
expect(screen.getByText('kimi-k2.6')).toBeInTheDocument()
expect(screen.queryByText('deepseek-reasoner')).not.toBeInTheDocument()
await act(async () => {
nextInspection.resolve({
@@ -240,6 +276,75 @@ describe('ContextUsageIndicator request behavior', () => {
await waitFor(() => {
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('12%')
expect(screen.getByText('deepseek-reasoner')).toBeInTheDocument()
})
})
it('keeps the last context until the replacement runtime signals a refresh', async () => {
sessionsApiMock.getInspection
.mockResolvedValueOnce(baseInspection)
.mockResolvedValueOnce({
active: true,
status: { ...baseInspection.status, model: 'deepseek-reasoner' },
errors: { context: 'CLI session stopped' },
})
.mockResolvedValueOnce({
...baseInspection,
status: { ...baseInspection.status, model: 'deepseek-reasoner' },
context: {
...baseInspection.context,
model: 'deepseek-reasoner',
percentage: 12,
},
})
const { rerender } = render(
<ContextUsageIndicator
sessionId="session-1"
chatState="idle"
messageCount={1}
runtimeSelectionKey="deepseek:deepseek-chat"
fallbackModelLabel="deepseek-chat"
refreshNonce={0}
/>,
)
await waitFor(() => {
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('21%')
})
rerender(
<ContextUsageIndicator
sessionId="session-1"
chatState="idle"
messageCount={1}
runtimeSelectionKey="deepseek:deepseek-reasoner"
fallbackModelLabel="deepseek-reasoner"
refreshNonce={0}
/>,
)
await waitFor(() => {
expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(2)
})
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('21%')
expect(screen.getByText('kimi-k2.6')).toBeInTheDocument()
expect(screen.queryByText('Context usage is unavailable for this session.')).not.toBeInTheDocument()
rerender(
<ContextUsageIndicator
sessionId="session-1"
chatState="idle"
messageCount={1}
runtimeSelectionKey="deepseek:deepseek-reasoner"
fallbackModelLabel="deepseek-reasoner"
refreshNonce={1}
/>,
)
await waitFor(() => {
expect(sessionsApiMock.getInspection).toHaveBeenCalledTimes(3)
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('12%')
expect(screen.getByText('deepseek-reasoner')).toBeInTheDocument()
})
})
@@ -334,6 +439,45 @@ describe('ContextUsageIndicator request behavior', () => {
expect(screen.queryByText('90%')).not.toBeInTheDocument()
})
it('does not show context retained from a different session', async () => {
sessionsApiMock.getInspection
.mockResolvedValueOnce(baseInspection)
.mockResolvedValueOnce({
...baseInspection,
status: { ...baseInspection.status, sessionId: 'session-2' },
context: {
...baseInspection.context,
percentage: 7,
},
})
const { rerender } = render(
<ContextUsageIndicator
sessionId="session-1"
chatState="idle"
messageCount={1}
/>,
)
await waitFor(() => {
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('21%')
})
rerender(
<ContextUsageIndicator
sessionId="session-2"
chatState="idle"
messageCount={1}
fallbackModelLabel="session-2-model"
/>,
)
expect(screen.getByTestId('context-usage-indicator')).not.toHaveTextContent('21%')
expect(screen.queryByText('kimi-k2.6')).not.toBeInTheDocument()
await waitFor(() => {
expect(screen.getByTestId('context-usage-indicator')).toHaveTextContent('7%')
})
})
it('forces a fresh inspection when refreshNonce bumps after a compaction (#743)', async () => {
// First request hangs — simulates an auto refresh that started just
// before the compact boundary and would resolve with pre-compact data.
@@ -16,8 +16,8 @@ type Props = {
compact?: boolean
/**
* Bump to force an immediate refresh that bypasses the auto-refresh
* throttle and any in-flight (possibly pre-compact) request. Used after
* context compaction so the meter recovers right away (#743).
* throttle and any stale in-flight request. Used after context compaction
* and after a replacement runtime confirms it has started.
*/
refreshNonce?: number
}
@@ -28,9 +28,8 @@ const ACTIVE_REFRESH_MS = 30_000
// racing a client abort that can strand loopback sockets on Windows.
const CONTEXT_REQUEST_TIMEOUT_MS = 30_000
const AUTO_REFRESH_MIN_INTERVAL_MS = 10_000
// Right after a compaction the CLI may still be busy finishing the turn, so
// the forced refresh can time out — retry once instead of keeping the stale
// pre-compact percentage on screen.
// Right after a completed turn, compaction, or runtime restart the CLI can
// still be settling, so retry the event-driven refresh once.
const FORCED_REFRESH_RETRY_MS = 5_000
function formatNumber(value: number | undefined) {
@@ -98,7 +97,7 @@ export function ContextUsageIndicator({
const [mobileDetailsOpen, setMobileDetailsOpen] = useState(false)
const requestSeq = useRef(0)
const contextIdentityRef = useRef('')
const contextDataIdentityRef = useRef('')
const contextDataSessionIdRef = useRef<string | undefined>(undefined)
const inFlightRequestRef = useRef<Promise<boolean> | null>(null)
const inFlightIdentityRef = useRef<string | null>(null)
const lastAutoRefreshAtRef = useRef(0)
@@ -141,9 +140,6 @@ export function ContextUsageIndicator({
const nextContext = inspection.context ?? inspection.contextEstimate ?? null
const nextSource = inspection.context ? 'live' : inspection.contextEstimate ? 'estimate' : null
const usageModel = inspection.usage?.models.find((model) => firstNonEmpty(model.displayName, model.model)) ?? null
contextDataIdentityRef.current = nextContext ? activeContextIdentity : ''
setContext(nextContext)
setContextSource(nextSource)
setInspectionModel(firstNonEmpty(
inspection.context?.model,
inspection.contextEstimate?.model,
@@ -151,18 +147,19 @@ export function ContextUsageIndicator({
usageModel?.displayName,
usageModel?.model,
) ?? null)
setError(nextContext ? null : inspection.errors?.context ?? null)
setUpdatedAt(Date.now())
if (nextContext) {
contextDataSessionIdRef.current = activeSessionId
setContext(nextContext)
setContextSource(nextSource)
setError(null)
setUpdatedAt(Date.now())
} else {
setError(inspection.errors?.context ?? null)
}
return nextContext !== null
})
.catch((err) => {
if (seq !== requestSeq.current || activeContextIdentity !== contextIdentityRef.current) return false
if (contextDataIdentityRef.current !== activeContextIdentity) {
contextDataIdentityRef.current = ''
setContext(null)
setContextSource(null)
setUpdatedAt(null)
}
setError(err instanceof Error ? err.message : String(err))
return false
})
@@ -178,14 +175,7 @@ export function ContextUsageIndicator({
return request
}, [draft, runtimeSelectionKey, sessionId])
// After a compaction the context shrinks server-side but nothing else
// re-reads it promptly (auto refreshes are throttled and stop once the
// session goes idle), leaving the pre-compact percentage on screen (#743).
// Force a fresh request, and retry once if the CLI was still busy.
const lastRefreshNonceRef = useRef(refreshNonce)
useEffect(() => {
if (refreshNonce === lastRefreshNonceRef.current) return
lastRefreshNonceRef.current = refreshNonce
const forceRefreshWithRetry = useCallback(() => {
let cancelled = false
let retryTimer: ReturnType<typeof setTimeout> | null = null
void refresh('force').then((ok) => {
@@ -198,7 +188,29 @@ export function ContextUsageIndicator({
cancelled = true
if (retryTimer) clearTimeout(retryTimer)
}
}, [refresh, refreshNonce])
}, [refresh])
// Compaction and runtime replacement both change context outside the normal
// message flow. Their completion signals bump this nonce so the meter reads
// the authoritative process.
const lastRefreshNonceRef = useRef(refreshNonce)
useEffect(() => {
if (refreshNonce === lastRefreshNonceRef.current) return
lastRefreshNonceRef.current = refreshNonce
return forceRefreshWithRetry()
}, [forceRefreshWithRetry, refreshNonce])
// A new session usually mounts while its first turn is already running.
// The eager inspection then races the CLI, and message-count refreshes can
// be swallowed by the auto-refresh throttle. The terminal idle transition
// is the first reliable point to request that session's real context.
const lastChatStateRef = useRef(chatState)
useEffect(() => {
const previousChatState = lastChatStateRef.current
lastChatStateRef.current = chatState
if (chatState !== 'idle' || previousChatState === 'idle') return
return forceRefreshWithRetry()
}, [chatState, forceRefreshWithRetry])
useEffect(() => {
const contextIdentity = `${sessionId}:${runtimeSelectionKey}`
@@ -209,9 +221,22 @@ export function ContextUsageIndicator({
lastAutoRefreshAtRef.current = 0
setError(null)
setInspectionModel(null)
if (contextDataSessionIdRef.current !== sessionId) {
contextDataSessionIdRef.current = undefined
setContext(null)
setContextSource(null)
setUpdatedAt(null)
}
}
void refresh('auto')
}, [messageCount, refresh, runtimeSelectionKey, sessionId])
}, [refresh, runtimeSelectionKey, sessionId])
const lastMessageCountRef = useRef(messageCount)
useEffect(() => {
if (lastMessageCountRef.current === messageCount) return
lastMessageCountRef.current = messageCount
void refresh('auto')
}, [messageCount, refresh])
useEffect(() => {
if (typeof document === 'undefined') return
@@ -236,7 +261,7 @@ export function ContextUsageIndicator({
return pickUsedContextCategory(context)
}, [context])
const displayContext = context
const displayContext = contextDataSessionIdRef.current === sessionId ? context : null
const hasPlaceholderContext = !displayContext && (
draft || (!loading && messageCount === 0 && (!error || isCliNotRunningError(error)))
)
@@ -256,12 +281,10 @@ export function ContextUsageIndicator({
: 'var(--color-surface-container-high)',
}
const displayPercent = displayContext ? formatPercent(percentage) : '--'
const currentContextIdentity = `${sessionId}:${runtimeSelectionKey}`
const isContextFromPreviousRuntime = Boolean(displayContext) &&
contextDataIdentityRef.current !== currentContextIdentity
const displayModel = isContextFromPreviousRuntime
? firstNonEmpty(fallbackModelLabel, inspectionModel, context?.model)
: firstNonEmpty(context?.model, inspectionModel, fallbackModelLabel)
const displayInspectionModel = !context || contextDataSessionIdRef.current === sessionId
? inspectionModel
: null
const displayModel = firstNonEmpty(displayContext?.model, displayInspectionModel, fallbackModelLabel)
const ariaLabel = displayContext
? t('contextIndicator.ariaLabel', { percent: formatPercent(percentage) })
: isPendingContext
+29
View File
@@ -3180,6 +3180,35 @@ describe('chatStore history mapping', () => {
})
})
it('bumps context refresh only for the runtime config currently selected', () => {
useSessionRuntimeStore.getState().setSelection(TEST_SESSION_ID, {
providerId: 'provider-b',
modelId: 'model-b',
effortLevel: 'high',
})
useChatStore.setState({
sessions: {
[TEST_SESSION_ID]: makeSession({ runtimeConfigReadyCount: 0 }),
},
})
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
type: 'runtime_config_applied',
providerId: 'provider-a',
modelId: 'model-a',
effortLevel: 'high',
})
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.runtimeConfigReadyCount).toBe(0)
useChatStore.getState().handleServerMessage(TEST_SESSION_ID, {
type: 'runtime_config_applied',
providerId: 'provider-b',
modelId: 'model-b',
effortLevel: 'high',
})
expect(useChatStore.getState().sessions[TEST_SESSION_ID]?.runtimeConfigReadyCount).toBe(1)
})
it('keeps AskUserQuestion permission requests out of the message list while tracking the pending request', () => {
useChatStore.setState({
sessions: {
+17
View File
@@ -122,6 +122,8 @@ export type PerSessionState = {
* Optional: legacy persisted sessions predate the field.
*/
compactCount?: number
/** Bumped when the server confirms the selected runtime is applied. */
runtimeConfigReadyCount?: number
/**
* Characters streamed by the assistant during the current turn (text,
* thinking, tool input). ÷4 approximates output tokens for the streaming
@@ -176,6 +178,7 @@ const DEFAULT_SESSION_STATE: PerSessionState = {
pendingComputerUsePermissions: {},
tokenUsage: { input_tokens: 0, output_tokens: 0 },
compactCount: 0,
runtimeConfigReadyCount: 0,
streamingResponseChars: 0,
elapsedSeconds: 0,
statusVerb: '',
@@ -2303,6 +2306,20 @@ export const useChatStore = create<ChatStore>((set, get) => ({
)
break
case 'runtime_config_applied': {
const selected = useSessionRuntimeStore.getState().selections[sessionId]
const matchesCurrentSelection = Boolean(selected) &&
(selected?.providerId ?? null) === msg.providerId &&
selected?.modelId === msg.modelId &&
selected?.effortLevel === msg.effortLevel
if (matchesCurrentSelection) {
update((session) => ({
runtimeConfigReadyCount: (session.runtimeConfigReadyCount ?? 0) + 1,
}))
}
break
}
case 'permission_mode_changed': {
// CLI 是权限模式的真相来源。这里把它恢复/切换后的权威值校正到本地镜像。
// 注意:只更新本地状态,**不要**走 setSessionPermissionMode —— 那会把
+6
View File
@@ -117,6 +117,12 @@ export type ServerMessage =
| { type: 'message_complete'; usage: TokenUsage }
| { type: 'thinking'; text: string }
| { type: 'status'; state: ChatState; verb?: string; attemptStart?: boolean }
| {
type: 'runtime_config_applied'
providerId: string | null
modelId: string
effortLevel?: string
}
// CLI 回传的权限模式变化(如 ExitPlanMode 退出 plan 后恢复、Shift+Tab)。
// 桌面端据此把选择器校正回 CLI 的真实权限,避免本地影子值漂移。
| { type: 'permission_mode_changed'; mode: PermissionMode }
+14 -1
View File
@@ -3729,7 +3729,11 @@ describe('WebSocket Chat Integration', () => {
}))
await waitUntil(
async () => messages.slice(switchStartIndex).some((msg) => msg.type === 'status' && msg.state === 'idle'),
async () => {
const switchMessages = messages.slice(switchStartIndex)
return switchMessages.some((msg) => msg.type === 'runtime_config_applied') &&
switchMessages.some((msg) => msg.type === 'status' && msg.state === 'idle')
},
`idle runtime switch completion for ${sessionId}`,
)
@@ -3747,6 +3751,15 @@ describe('WebSocket Chat Integration', () => {
.filter((msg) => msg.type === 'status')
.map((msg) => msg.state),
).toEqual(['idle'])
expect(
messages
.slice(switchStartIndex)
.find((msg) => msg.type === 'runtime_config_applied'),
).toMatchObject({
type: 'runtime_config_applied',
providerId: provider.id,
modelId: 'idle-sonnet',
})
expect(messages.slice(switchStartIndex).some((msg) => msg.type === 'error')).toBe(false)
} finally {
ws.close()
+8
View File
@@ -53,6 +53,8 @@ export type AttachmentRef = {
// Server → Client
// ============================================================================
export const RUNTIME_CONFIG_APPLIED_EVENT = 'runtime_config_applied' as const
export type ServerMessage =
| { type: 'connected'; sessionId: string }
| { type: 'session_state'; turnState: 'running' | 'idle' }
@@ -89,6 +91,12 @@ export type ServerMessage =
| { type: 'message_complete'; usage: TokenUsage }
| { type: 'thinking'; text: string }
| { type: 'status'; state: ChatState; verb?: string; attemptStart?: boolean }
| {
type: typeof RUNTIME_CONFIG_APPLIED_EVENT
providerId: string | null
modelId: string
effortLevel?: string
}
// CLI 是权限模式的唯一真相来源。当 CLI 内部 mode 变化(如 ExitPlanMode 后
// 恢复到进入 plan 前的模式、Shift+Tab 切换)时,把新模式回传给前端,让桌面端
// 选择器与 CLI 保持同步,而不是停留在本地影子值上。
+16
View File
@@ -14,6 +14,7 @@ import type {
StreamingFallbackCause,
TokenUsage,
} from './events.js'
import { RUNTIME_CONFIG_APPLIED_EVENT } from './events.js'
import * as os from 'node:os'
import {
ConversationStartupError,
@@ -1582,6 +1583,8 @@ async function handleSetRuntimeConfig(
const currentRuntimeVersion = runtimeOverrideVersions.get(sessionId) ?? 0
if (startupRuntimeVersion >= currentRuntimeVersion) {
await persistSessionRuntimeConfig(sessionId, nextOverride)
await pendingStartup
broadcastAppliedRuntimeConfig(sessionId)
return
}
@@ -1603,6 +1606,7 @@ async function handleSetRuntimeConfig(
}
await persistSessionRuntimeConfig(sessionId, nextOverride)
broadcastAppliedRuntimeConfig(sessionId)
}
async function restartSessionWithPermissionMode(
@@ -1705,6 +1709,17 @@ async function persistSessionRuntimeConfig(
})
}
function broadcastAppliedRuntimeConfig(sessionId: string): void {
const runtime = runtimeOverrides.get(sessionId)
if (!runtime) return
sendToSession(sessionId, {
type: RUNTIME_CONFIG_APPLIED_EVENT,
providerId: runtime.providerId,
modelId: runtime.modelId,
...(runtime.effort ? { effortLevel: runtime.effort } : {}),
})
}
async function resolveRuntimeRestartWorkDir(sessionId: string): Promise<string> {
const activeWorkDir = conversationService.getSessionWorkDir(sessionId)
if (activeWorkDir) return activeWorkDir
@@ -1732,6 +1747,7 @@ async function restartSessionWithRuntimeConfig(
await conversationService.startSession(sessionId, workDir, sdkUrl, runtimeSettings)
runtimeExitStoppedSessions.delete(sessionId)
broadcastAppliedRuntimeConfig(sessionId)
sendMessage(ws, { type: 'status', state: 'idle' })
console.log(`[WS] Restarted CLI for ${sessionId} with runtime override`)
} catch (err) {