From 042e0ac43cbdc40bbc2d13d667f61474cadf497d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 25 May 2026 07:37:03 +0000 Subject: [PATCH] Scope contextUsage to active checkpoint Three follow-ups on #5736 so the relaxed external-provider render gate does not show stale token / cache stats from a different model: 1) setCheckpoint now clears contextUsage on a real checkpoint change. setActiveThreadId and clearCheckpoint already did this; the most-traveled transition path (the user switching models from the picker) leaked the prior turn's counts because they were never cleared. 2) The external-selection branch in chat-page.tsx now also clears contextUsage at the same time it nulls ggufContextLength / activeNativePathToken. Without this an in-session switch from a local model to an external provider would visibly carry the previous local turn's counters into the new provider's bar. 3) exitCompare's rehydration is now scoped: restore the saved usage only when the message's modelId matches the active checkpoint AND, for local turns where a context window is known, when the saved total fits inside that window. Without this the bar could render a stale local-model usage on top of an external provider, or an oversized usage object that exceeds the now- active window. Typecheck clean. --- .../frontend/src/features/chat/chat-page.tsx | 41 ++++++++++++++++++- .../chat/stores/chat-runtime-store.ts | 10 +++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 46faaefe07..a06e47a01b 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1037,6 +1037,11 @@ export function ChatPage(): ReactElement { ggufMaxContextLength: null, ggufNativeContextLength: null, activeNativePathToken: null, + // External selection arrives mid-session: also clear any + // pre-existing per-turn usage from the previous model so the + // relaxed external-provider render gate does not show stale + // counts until the next completion overwrites them. + contextUsage: null, supportsReasoning: reasoningCaps.supportsReasoning, reasoningAlwaysOn: reasoningCaps.reasoningAlwaysOn, reasoningStyle: reasoningCaps.reasoningStyle, @@ -1161,7 +1166,11 @@ export function ChatPage(): ReactElement { if (!saved) return; viewBeforeCompareRef.current = null; navigate({ to: "/chat", search: saved }); - // Restore context usage from the active thread's last assistant message. + // Restore context usage from the active thread's last assistant + // message, but only if it was produced by the SAME checkpoint the + // user is now sitting on. Without this guard the relaxed render + // gate would show stale token / cache stats from a different + // provider or a local turn that exceeds the active GGUF window. const threadId = saved.thread ?? useChatRuntimeStore.getState().activeThreadId; if (threadId) { @@ -1175,7 +1184,35 @@ export function ChatPage(): ReactElement { const usage = metadata?.contextUsage as ReturnType< typeof useChatRuntimeStore.getState >["contextUsage"]; - if (usage) useChatRuntimeStore.getState().setContextUsage(usage); + if (!usage) return; + const store = useChatRuntimeStore.getState(); + const activeCheckpoint = store.params.checkpoint; + const usageModelId = + (usage as { modelId?: unknown }).modelId; + // Scope by modelId when the saved usage carries one (the + // chat-adapter stamps it on every external + local turn). + if ( + typeof usageModelId === "string" && + usageModelId && + activeCheckpoint && + usageModelId !== activeCheckpoint + ) { + return; + } + // For local llama-server turns, also require that the + // restored prompt count fits inside the active context + // window. Skip the check when the window is unknown (e.g. + // external provider, no ggufContextLength) so the existing + // external-provider rendering path stays intact. + const limit = store.ggufContextLength; + if ( + typeof limit === "number" && + limit > 0 && + (usage.totalTokens ?? 0) > limit + ) { + return; + } + store.setContextUsage(usage); }) .catch((error) => { if (!isExpectedBackgroundChatStorageError(error)) { diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 58969f5c70..25247d19ec 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -706,12 +706,22 @@ export const useChatRuntimeStore = create((set, get) => ({ // mount, and a stale persisted local id would race against the // freshly-loaded model. See LAST_EXTERNAL_CHECKPOINT_KEY notes. saveLastExternalCheckpoint(isExternalModelId(modelId) ? modelId : null); + // Clear any stale per-turn usage when the active model changes. + // The context bar render gate was relaxed for external providers + // (no longer requires ggufContextLength), so leaving a previous + // turn's counters around would visibly show old token / cache + // stats from a different model until the next completion + // overwrites them. setActiveThreadId / clearCheckpoint already + // do this; keep the invariant on the most-traveled transition + // path too. + const checkpointChanged = state.params.checkpoint !== modelId; return { params: { ...state.params, checkpoint: modelId, }, activeGgufVariant: ggufVariant ?? null, + ...(checkpointChanged ? { contextUsage: null } : {}), }; }), setActiveThreadId: (activeThreadId) =>