From 79da85ba7f379b560f987f9c05098607b13368de Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 23 May 2026 07:24:05 +0000 Subject: [PATCH] Studio: render cache stats for external providers too Reviewer round on the original PR caught three asymmetric-fix sites where the producer side surfaced external prompt-cache stats but the consumer side still gated on ggufContextLength (which is only ever set for the local llama-server runtime). Result: the entire cache-stats PR shipped invisible for Anthropic / OpenAI Responses / Gemini, which is exactly the set of providers it was added for. - chat-page.tsx: drop the ggufContextLength precondition on the ContextUsageBar mount. The bar already tracks usage; let it decide what to render based on what it knows. - context-usage-bar.tsx: make `total` optional. When absent, drop the "/ total" ratio + percentage progress bar + "approaching limit" helper, and just show per-turn counters + cache stats. Bootstrap guard tightened so an all-zero, all-undefined state still renders nothing. - runtime-provider.tsx: external-provider rehydration was rejected by the `store.ggufContextLength` check. Keep the "fits inside window" sanity check when a local context window IS known, drop it when it isn't. - message-timing.tsx: the per-message timing popover used a separate "Cache hits" code path that only read llama-server's timings.cache_n. Fall through to custom.contextUsage for external providers, and add a parallel "Cache writes" line for Anthropic cache_creation events. --- .../assistant-ui/message-timing.tsx | 53 +++++++++++-- .../frontend/src/features/chat/chat-page.tsx | 8 +- .../chat/components/context-usage-bar.tsx | 76 +++++++++++++------ .../src/features/chat/runtime-provider.tsx | 11 ++- 4 files changed, 117 insertions(+), 31 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/message-timing.tsx b/studio/frontend/src/components/assistant-ui/message-timing.tsx index 4fb68d4b90..be8d7dc7a1 100644 --- a/studio/frontend/src/components/assistant-ui/message-timing.tsx +++ b/studio/frontend/src/components/assistant-ui/message-timing.tsx @@ -33,10 +33,27 @@ export const MessageTiming: FC<{ if (timing?.totalStreamTime === undefined) return null; - const serverTimings = ( + const custom = ( message.metadata as Record | undefined - )?.custom as { serverTimings?: Record } | undefined; - const st = serverTimings?.serverTimings; + )?.custom as + | { + serverTimings?: Record; + contextUsage?: { + cachedTokens?: number; + cacheWriteTokens?: number; + }; + } + | undefined; + const st = custom?.serverTimings; + // Cache-hit / cache-write counts. llama-server reports hits on + // timings.cache_n; external providers (Anthropic / OpenAI Responses / + // Gemini) report them on the include_usage envelope that the adapter + // normalizes into custom.contextUsage. Prefer the local-runtime value + // when present (so llama.cpp keeps populating the badge mid-stream) + // and fall back to the external envelope otherwise. + const cacheHits = (st?.cache_n ?? 0) || (custom?.contextUsage?.cachedTokens ?? 0); + // Anthropic-only: tokens written into the prompt cache on this turn. + const cacheWrites = custom?.contextUsage?.cacheWriteTokens ?? 0; // Guard unphysical tok/s: llama.cpp emits predicted_ms=0 on no-op // turns, blowing the rate up to Infinity. Require >=1 token AND a @@ -122,11 +139,19 @@ export const MessageTiming: FC<{ )} - {(st?.cache_n ?? 0) > 0 && ( + {cacheHits > 0 && (
Cache hits - {formatNumber(st!.cache_n)} + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)}
)} @@ -146,7 +171,7 @@ export const MessageTiming: FC<{ ) : ( <> - {/* Client-side metrics (safetensors fallback) */} + {/* Client-side metrics (safetensors + external provider fallback) */} {timing.firstTokenTime !== undefined && (
First token @@ -155,6 +180,22 @@ export const MessageTiming: FC<{
)} + {cacheHits > 0 && ( +
+ Cache hits + + {formatNumber(cacheHits)} + +
+ )} + {cacheWrites > 0 && ( +
+ Cache writes + + {formatNumber(cacheWrites)} + +
+ )}
Total diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index d4aa53e4c4..eaa8c87cc0 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -1491,9 +1491,15 @@ export function ChatPage(): ReactElement { ) : null}
- {view.mode === "single" && ggufContextLength && contextUsage ? ( + {view.mode === "single" && contextUsage ? ( { - if (total <= 0) return null; + const hasKnownLimit = typeof total === "number" && total > 0; + const hasUsageDetails = + promptTokens !== undefined || + completionTokens !== undefined || + (cached !== undefined && cached > 0) || + (cacheWrites !== undefined && cacheWrites > 0); - const percent = Math.min((used / total) * 100, 100); - const severity = getSeverityColor(percent); + // Nothing to show: no limit and no per-turn counters. + if (!hasKnownLimit && used <= 0 && !hasUsageDetails) return null; + + const percent = hasKnownLimit + ? Math.min((used / (total as number)) * 100, 100) + : null; + const severity = getSeverityColor(percent ?? 0); return (
-
- Context usage - - {percent.toFixed(1)}% - -
+ {hasKnownLimit && percent !== null ? ( +
+ Context usage + + {percent.toFixed(1)}% + +
+ ) : null} {promptTokens !== undefined && (
Prompt tokens @@ -122,18 +150,22 @@ export const ContextUsageBar: FC<{ )}
- Total + + {hasKnownLimit ? "Total" : "Total tokens"} + - {formatTokenCountFull(used)} / {formatTokenCountFull(total)} + {hasKnownLimit + ? `${formatTokenCountFull(used)} / ${formatTokenCountFull(total as number)}` + : formatTokenCountFull(used)}
- {percent > 85 && ( + {hasKnownLimit && percent !== null && percent > 85 ? (
Close to the context limit. Generation will stop at 100%. Increase Context Length in the chat Settings panel to keep going.
- )} + ) : null}
diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 6d136c978e..4c758e01c4 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -831,10 +831,17 @@ function useStudioRuntimeAdapters(): StudioRuntimeAdapters { } | undefined; const store = useChatRuntimeStore.getState(); + // External-provider threads have ggufContextLength === null because + // external picker selections don't carry a context-window number. + // Restore the persisted usage as long as it belongs to the active + // checkpoint; when a local GGUF context window IS known, also keep + // the original sanity check that the saved total fits inside it. + const withinLocalLimit = + !store.ggufContextLength || + (savedUsage?.totalTokens ?? 0) <= store.ggufContextLength; if ( savedUsage && - store.ggufContextLength && - savedUsage.totalTokens <= store.ggufContextLength && + withinLocalLimit && (!savedUsage.modelId || savedUsage.modelId === store.params.checkpoint) ) {