@samitouri / QOSami_Agents_Zero / commits / 4fb3d255

Add provider accounting to context window

Show provider-reported price, cache hit, and input/output tokens in the context window popover while omitting unavailable price and cache rows. Drain OpenRouter terminal usage through plugin extensions without changing Responses behavior, and keep prompt-bucket reconciliation accurate.

Alessandro committed Aug 24, 2026 at 17:05 UTC 4fb3d2555f669be6413b469128ada95d6ab8768a
11 files changed +504 -34
plugins/_context_window/AGENTS.md
+16 -1
@@ -8,7 +8,8 @@
8 ## Ownership
9
10 - `helpers/usage.py` owns per-prompt bucket measurement and reconciliation.
11 -- `extensions/python/` records prompt parts at their source extension points.
11 +- `extensions/python/` records prompt parts at their source extension points,
12 + preserves terminal streamed usage, and captures optional provider usage.
13 - `api/context_window.py` exposes the active chat's token usage and effective
14 model limit without returning prompt content.
15 - `webui/` and `extensions/webui/` own the Alpine store, indicator, popover,
@@ -26,6 +27,19 @@
27 fragments use a bounded, content-addressed, runtime-only cache.
28 - Bucket totals reconcile to the already-stored prompt token total; the
29 unclaimed remainder belongs to System prompt.
30 +- If the history ledger would consume the whole prompt estimate, recompute only
31 + the rendered message portion before reconciliation; ordinary prompt builds
32 + keep the fast ledger path.
33 +- The prompt estimate never guesses provider-specific image token costs or
34 + counts embedded image bytes as text.
35 +- Provider price, cache hit, and input/output tokens form a flat summary without
36 + diagnostic detail rows.
37 +- Provider rows are exposed only when the provider or transport reports their
38 + values; unavailable price and cache data render no row.
39 +- Streamed OpenRouter main turns request LiteLLM's terminal usage event. The
40 + response callback still runs normally; only an actual Chat Completions result
41 + restores the accepted response after the accounting tail is drained.
42 +- Responses API turns keep their native result and callback behavior unchanged.
43 - Older chats without a stored breakdown show the explanatory empty state.
44 - `_model_config` supplies the effective model limit and the
45 `model-context-strip-end` WebUI slot; it does not own this feature's state.
@@ -35,6 +49,7 @@
49 ## Work Guidance
50
51 - Keep prompt accounting out of rendered-text heuristics.
52 +- Keep provider-reported usage separate from the six estimated context buckets.
53 - Keep the API response limited to counts needed by the UI.
54 - Preserve the upward, right-aligned popover geometry used beside the model and
55 profile selectors.
plugins/_context_window/README.md
+8
@@ -7,3 +7,11 @@ Extras, and Free space.
7
8 Older chats gain the detailed breakdown after their next model turn. Mobile and
9 desktop visibility can be changed under **Settings > Interface**.
10 +
11 +When a model provider reports usage, the popover shows price, cache-hit rate,
12 +and input/output tokens. Unreported price and cache data are omitted. The
13 +context breakdown does not guess model-specific image token costs.
14 +
15 +For streamed OpenRouter chat calls, the plugin requests and drains LiteLLM's
16 +terminal usage event after Agent Zero has accepted the response. Price remains
17 +hidden when LiteLLM does not report a cost or map the selected model.
plugins/_context_window/api/context_window.py
+5 -1
@@ -1,5 +1,8 @@
1 from helpers.api import ApiHandler, Input, Output, Request
2 -from plugins._context_window.helpers.usage import usage_snapshot
2 +from plugins._context_window.helpers.usage import (
3 + latest_provider_usage,
4 + usage_snapshot,
5 +)
6 from plugins._model_config.helpers.model_config import get_chat_model_config
7
8
@@ -15,4 +18,5 @@ class ContextWindow(ApiHandler):
18 "tokens": max(int(window.get("tokens") or 0), 0),
19 "context_window": max(int(config.get("ctx_length") or 0), 0),
20 "usage": usage_snapshot(window.get("usage")),
21 + "provider_usage": latest_provider_usage(agent),
22 }
plugins/_context_window/extensions/python/_functions/agent/Agent/call_chat_model_turn/end/_90_record_provider_usage.py new
+9
@@ -0,0 +1,9 @@
1 +from helpers.extension import Extension
2 +from plugins._context_window.helpers.usage import capture_provider_usage
3 +
4 +
5 +class RecordProviderUsage(Extension):
6 + def execute(self, data: dict | None = None, **kwargs):
7 + payload = data if isinstance(data, dict) else {}
8 + if not payload.get("exception"):
9 + capture_provider_usage(self.agent, payload.get("result"))
plugins/_context_window/extensions/python/_functions/models/LiteLLMChatWrapper/unified_turn/end/_90_restore_provider_response.py new
+24
@@ -0,0 +1,24 @@
1 +from helpers.extension import Extension
2 +from helpers.llm_result import LLMResult
3 +
4 +
5 +RESPONSE_KEY = "_context_window_accepted_response"
6 +
7 +
8 +class RestoreProviderResponse(Extension):
9 + def execute(self, data: dict, **kwargs):
10 + response = data.pop(RESPONSE_KEY, None)
11 + result = data.get("result")
12 + if (
13 + response is None
14 + or data.get("exception")
15 + or not isinstance(result, LLMResult)
16 + or result.mode != "chat_completions"
17 + ):
18 + return
19 +
20 + result.response = response
21 + if not result.output_items:
22 + result.output_items = LLMResult.from_chat(
23 + response=response, reasoning=result.reasoning
24 + ).output_items
plugins/_context_window/extensions/python/_functions/models/LiteLLMChatWrapper/unified_turn/start/_10_drain_provider_usage.py new
+30
@@ -0,0 +1,30 @@
1 +from helpers.extension import Extension
2 +
3 +
4 +RESPONSE_KEY = "_context_window_accepted_response"
5 +
6 +
7 +class DrainProviderUsage(Extension):
8 + def execute(self, data: dict, **kwargs):
9 + model = data["args"][0]
10 + callback = data["kwargs"].get("response_callback")
11 + if callback is None:
12 + return
13 +
14 + if (
15 + model.provider == "openrouter"
16 + and data["kwargs"].get("explicit_caching")
17 + ):
18 + data["kwargs"]["stream_options"] = {
19 + **model.kwargs.get("stream_options", {}),
20 + **data["kwargs"].get("stream_options", {}),
21 + "include_usage": True,
22 + }
23 +
24 + async def drain_callback(chunk: str, full: str):
25 + response = await callback(chunk, full)
26 + if response is not None and RESPONSE_KEY not in data:
27 + data[RESPONSE_KEY] = response
28 + return None
29 +
30 + data["kwargs"]["response_callback"] = drain_callback
plugins/_context_window/extensions/webui/model-context-strip-end/context-window.html
+67 -17
@@ -30,7 +30,10 @@
30 aria-label="Context window usage">
31 <div class="context-window-header">
32 <strong>Context window</strong>
33 - <span x-text="$store.contextWindow.usage.summary"></span>
33 + <div class="context-window-summary">
34 + <span class="context-window-summary-tokens" x-text="$store.contextWindow.usage.summaryTokens"></span>
35 + <span class="context-window-summary-percent" x-text="$store.contextWindow.usage.summaryPercent"></span>
36 + </div>
37 </div>
38 <div class="context-window-meter" aria-hidden="true">
39 <span :style="$store.contextWindow.usage.meterStyle"></span>
@@ -39,7 +42,6 @@
42 <div class="context-window-rows">
43 <template x-for="row in $store.contextWindow.usage.rows" :key="row.key">
44 <div class="context-window-row">
42 - <span class="context-window-dot" :style="row.dotStyle"></span>
45 <span x-text="row.label"></span>
46 <span class="context-window-value context-window-tokens" x-text="row.tokensLabel"></span>
47 <span class="context-window-value" x-text="row.percentLabel"></span>
@@ -50,6 +52,28 @@
52 <template x-if="$store.contextWindow.usage.missingBreakdown">
53 <div class="context-window-empty">Breakdown available after the next message.</div>
54 </template>
55 + <template x-if="$store.contextWindow.usage.provider.hasData">
56 + <div class="context-window-provider">
57 + <template x-if="$store.contextWindow.usage.provider.price.hasData">
58 + <div class="context-window-provider-row context-window-provider-price">
59 + <span>Price</span>
60 + <span class="context-window-value" x-text="$store.contextWindow.usage.provider.price.label"></span>
61 + </div>
62 + </template>
63 + <template x-if="$store.contextWindow.usage.provider.cache.hasData">
64 + <div class="context-window-provider-row">
65 + <span>Cache hit</span>
66 + <span class="context-window-value" x-text="$store.contextWindow.usage.provider.cache.label"></span>
67 + </div>
68 + </template>
69 + <template x-if="$store.contextWindow.usage.provider.tokens">
70 + <div class="context-window-provider-row">
71 + <span>Tokens In/Out</span>
72 + <span class="context-window-value" x-text="$store.contextWindow.usage.provider.tokens"></span>
73 + </div>
74 + </template>
75 + </div>
76 + </template>
77 </div>
78 </div>
79 </template>
@@ -99,7 +123,7 @@
123 .context-window-button > span {
124 position: relative;
125 z-index: 1;
102 - font-family: var(--font-family-code);
126 + font-family: var(--font-family-main);
127 font-size: .42rem;
128 font-weight: 600;
129 letter-spacing: -.03em;
@@ -117,27 +141,59 @@
141 box-shadow: 0 16px 38px rgba(0, 0, 0, .3);
142 }
143 .context-window-header,
120 - .context-window-row {
144 + .context-window-row,
145 + .context-window-provider-row {
146 display: grid;
147 align-items: center;
148 }
149 .context-window-header {
125 - grid-template-columns: minmax(0, 1fr) auto;
126 - gap: .75rem;
150 + gap: .4rem;
151 font-size: .86rem;
152 }
129 - .context-window-header > span,
153 + .context-window-summary,
154 .context-window-value {
131 - font-family: var(--font-family-code);
155 + font-family: var(--font-family-main);
156 font-variant-numeric: tabular-nums;
157 }
134 - .context-window-header > span {
158 + .context-window-summary {
159 + display: flex;
160 + align-items: baseline;
161 + justify-content: space-between;
162 + gap: .75rem;
163 + color: color-mix(in srgb, var(--color-text) 72%, transparent);
164 + font-size: .76rem;
165 + }
166 + .context-window-summary-percent {
167 + color: var(--color-text);
168 + }
169 + .context-window-summary-tokens,
170 + .context-window-tokens {
171 + font-family: var(--font-family-code);
172 + }
173 + .context-window-provider {
174 + margin-top: .7rem;
175 + padding-top: .7rem;
176 + border-top: 1px solid var(--color-border);
177 + }
178 + .context-window-provider-row {
179 + grid-template-columns: minmax(0, 1fr) auto;
180 + gap: .75rem;
181 + min-height: 1.35rem;
182 color: color-mix(in srgb, var(--color-text) 72%, transparent);
183 font-size: .76rem;
184 }
185 + .context-window-provider-price {
186 + margin-bottom: .25rem;
187 + color: var(--color-text);
188 + font-size: .82rem;
189 + }
190 + .context-window-provider-price .context-window-value {
191 + font-size: 1rem;
192 + font-weight: 600;
193 + }
194 .context-window-meter {
195 height: .45rem;
140 - margin: .85rem 0 .75rem;
196 + margin: .4rem 0 .75rem;
197 overflow: hidden;
198 border-radius: 999px;
199 background: color-mix(in srgb, var(--color-text) 9%, transparent);
@@ -159,7 +215,7 @@
215 font-size: .8rem;
216 }
217 .context-window-row {
162 - grid-template-columns: .55rem minmax(0, 1fr) auto 3rem;
218 + grid-template-columns: minmax(0, 1fr) auto 3rem;
219 gap: .5rem;
220 min-height: 1.5rem;
221 color: color-mix(in srgb, var(--color-text) 78%, transparent);
@@ -172,12 +228,6 @@
228 .context-window-tokens {
229 color: color-mix(in srgb, var(--color-text) 62%, transparent);
230 }
175 - .context-window-dot {
176 - width: .55rem;
177 - height: .55rem;
178 - border-radius: 50%;
179 - background: var(--color-highlight);
180 - }
231 @media (max-width: 25rem) {
232 .context-window-popover {
233 right: 1.25rem;
plugins/_context_window/helpers/usage.py
+110 -1
@@ -1,11 +1,14 @@
1 import hashlib
2 +import math
3 from typing import Any
4
5 from helpers import files, history, skills, tokens
6 +from helpers.llm_result import result_from_metadata
7
8
9 PARTS_KEY = "context_window_usage"
10 CACHE_KEY = "_context_window_usage_cache"
11 +PROVIDER_USAGE_KEY = "context_window_provider_usage"
12 USAGE_KEYS = (
13 "messages",
14 "system_tools",
@@ -43,6 +46,7 @@ def capture_context(agent: Any, loop_data: Any) -> None:
46 return
47
48 output = list(getattr(loop_data, "history_output", None) or [])
49 + parts["_history_output"] = output
50 skill_output = [message for message in output if skills.skill_instruction_name(message)]
51 skill_tokens = _output_tokens(agent, "history_skills", skill_output)
52 parts["messages"] = max(_history_tokens(agent, output) - skill_tokens, 0)
@@ -78,9 +82,20 @@ def finalize(agent: Any) -> None:
82 if not isinstance(parts, dict) or not isinstance(window, dict):
83 return
84
85 + history_output = parts.pop("_history_output", None)
86 total = _non_negative_int(window.get("tokens"))
87 usage = {key: _non_negative_int(parts.get(key)) for key in MEASURED_KEYS}
88 measured_total = sum(usage.values())
89 + if total and measured_total >= total and isinstance(history_output, list):
90 + message_output = [
91 + message
92 + for message in history_output
93 + if not skills.skill_instruction_name(message)
94 + ]
95 + usage["messages"] = _output_tokens(
96 + agent, "history_messages", message_output
97 + )
98 + measured_total = sum(usage.values())
99 if measured_total > total and measured_total:
100 usage = _scale_to_total(usage, total, measured_total)
101 measured_total = total
@@ -98,7 +113,73 @@ def usage_snapshot(value: Any) -> dict[str, int]:
113 return {key: _non_negative_int(value.get(key)) for key in USAGE_KEYS}
114
115
101 -def _parts(agent: Any) -> dict[str, int] | None:
116 +def capture_provider_usage(agent: Any, result: Any) -> None:
117 + if agent is None or result is None or not hasattr(result, "usage"):
118 + return
119 +
120 + snapshot = provider_usage_snapshot(getattr(result, "usage", None))
121 + agent.set_data(
122 + PROVIDER_USAGE_KEY,
123 + snapshot if snapshot else {"available": False},
124 + )
125 +
126 +
127 +def latest_provider_usage(agent: Any) -> dict[str, int | float]:
128 + data = getattr(agent, "data", None)
129 + if isinstance(data, dict) and PROVIDER_USAGE_KEY in data:
130 + stored = data.get(PROVIDER_USAGE_KEY)
131 + if isinstance(stored, dict) and stored.get("available") is False:
132 + return {}
133 + return provider_usage_snapshot(stored)
134 +
135 + all_messages = getattr(getattr(agent, "history", None), "all_messages", None)
136 + if not callable(all_messages):
137 + return {}
138 + for message in reversed(all_messages()):
139 + if not getattr(message, "ai", False):
140 + continue
141 + result = result_from_metadata(getattr(message, "metadata", None))
142 + if result:
143 + return provider_usage_snapshot(result.usage)
144 + return {}
145 +
146 +
147 +def provider_usage_snapshot(value: Any) -> dict[str, int | float]:
148 + if not isinstance(value, dict):
149 + return {}
150 +
151 + input_details = {
152 + **_mapping(value.get("prompt_tokens_details")),
153 + **_mapping(value.get("input_tokens_details")),
154 + }
155 + result: dict[str, int | float] = {}
156 + fields = {
157 + "input_tokens": (value.get("input_tokens"), value.get("prompt_tokens")),
158 + "cached_tokens": (
159 + input_details.get("cached_tokens"),
160 + input_details.get("cache_read_tokens"),
161 + value.get("cache_read_input_tokens"),
162 + value.get("cached_tokens"),
163 + ),
164 + "output_tokens": (
165 + value.get("output_tokens"),
166 + value.get("completion_tokens"),
167 + ),
168 + }
169 + for key, values in fields.items():
170 + number = _optional_non_negative_int(*values)
171 + if number is not None:
172 + result[key] = number
173 +
174 + cost = _optional_non_negative_float(
175 + value.get("cost"), value.get("response_cost")
176 + )
177 + if cost is not None:
178 + result["cost"] = cost
179 + return result
180 +
181 +
182 +def _parts(agent: Any) -> dict[str, Any] | None:
183 params = _temporary_params(agent)
184 value = params.get(PARTS_KEY) if params is not None else None
185 return value if isinstance(value, dict) else None
@@ -160,6 +241,34 @@ def _non_negative_int(value: Any) -> int:
241 return 0
242
243
244 +def _mapping(value: Any) -> dict[str, Any]:
245 + return value if isinstance(value, dict) else {}
246 +
247 +
248 +def _optional_non_negative_int(*values: Any) -> int | None:
249 + for value in values:
250 + if value is None:
251 + continue
252 + try:
253 + return max(int(value), 0)
254 + except (TypeError, ValueError):
255 + continue
256 + return None
257 +
258 +
259 +def _optional_non_negative_float(*values: Any) -> float | None:
260 + for value in values:
261 + if value is None:
262 + continue
263 + try:
264 + number = float(value)
265 + except (TypeError, ValueError):
266 + continue
267 + if math.isfinite(number):
268 + return max(number, 0)
269 + return None
270 +
271 +
272 def _scale_to_total(values: dict[str, int], total: int, current: int) -> dict[str, int]:
273 scaled = {key: value * total // current for key, value in values.items()}
274 remainder = total - sum(scaled.values())
plugins/_context_window/tests/test_context_window.py
+181 -4
@@ -1,17 +1,23 @@
1 +import sys
2 from pathlib import Path
3 from types import SimpleNamespace
4
5 import pytest
6 +from langchain_core.messages import HumanMessage, SystemMessage
7
8 +
9 +ROOT = Path(__file__).resolve().parents[3]
10 +if str(ROOT) not in sys.path:
11 + sys.path.insert(0, str(ROOT))
12 +
13 +import models
14 from agent import Agent, LoopData
7 -from helpers import extension, history
15 +from helpers import extension, extract_tools, history, litellm_transport
16 +from helpers.llm_result import LLMResult
17 from plugins._context_window.api.context_window import ContextWindow
18 from plugins._context_window.helpers import usage
19
20
12 -ROOT = Path(__file__).resolve().parents[3]
13 -
14 -
21 class _Log:
22 def set_progress(self, _message: str) -> None:
23 pass
@@ -80,6 +86,7 @@ async def test_usage_follows_prompt_sources_and_reconciles_to_total(monkeypatch)
86 assert sum(breakdown.values()) == window["tokens"]
87 assert all(breakdown[key] > 0 for key in usage.USAGE_KEYS)
88 assert usage.PARTS_KEY not in loop_data.params_temporary
89 + assert "history_messages" not in agent.data[usage.CACHE_KEY]
90
91
92 @pytest.mark.asyncio
@@ -115,6 +122,7 @@ async def test_api_returns_only_counts_and_effective_limit(monkeypatch):
122 "system_prompt": 0,
123 "extras": 0,
124 },
125 + "provider_usage": {},
126 }
127 assert "text" not in result
128
@@ -147,6 +155,34 @@ def test_webui_and_accounting_are_plugin_owned():
155 assert "right: 1.25rem" in component
156 assert "width: min(17rem, calc(100vw - 3rem))" in component
157 assert 'label: "Free space"' in context_store
158 + assert "Last model call" not in component
159 + assert ">Price<" in component
160 + assert ">Cache hit<" in component
161 + assert ">Tokens In/Out<" in component
162 + assert "context-window-cache-meter" not in component
163 + assert "price: {" in context_store
164 + assert "hasData: cost !== null" in context_store
165 + assert 'label: cost === null ? "" : formatCost(cost)' in context_store
166 + assert "usage.provider.price.hasData" in component
167 + assert "usage.provider.price.label" in component
168 + assert 'value < 0.001 ? "<$0.001"' in context_store
169 + assert "maximumSignificantDigits: 3" in context_store
170 + assert "border-top: 1px solid var(--color-border)" in component
171 + assert " → " in context_store
172 + assert "summaryTokens" in context_store
173 + assert 'summaryPercent: `${percentLabel} used`' in context_store
174 + assert "formatTokens(output)} tok" not in context_store
175 + assert "context-window-summary-tokens" in component
176 + assert "context-window-summary-percent" in component
177 + assert "font-family: var(--font-family-main)" in component
178 + assert "<details" not in component
179 + assert "Reasoning" not in component
180 + assert "Images sent" not in component
181 + assert "provider did not split out their token cost" not in context_store
182 + assert "cached / input" in context_store
183 + assert "Math.round(cachePercent)" in context_store
184 + assert "context-window-dot" not in component
185 + assert "dotStyle" not in context_store
186 assert "Breakdown available after the next message." in component
187 assert "startswith(" not in helper
188 assert "rpartition(" not in helper
@@ -156,6 +192,9 @@ def test_source_prompt_extensions_are_registered():
192 expected = {
193 "_functions/agent/Agent/prepare_prompt/start": "ResetContextUsage",
194 "_functions/agent/Agent/prepare_prompt/end": "StoreContextUsage",
195 + "_functions/agent/Agent/call_chat_model_turn/end": "RecordProviderUsage",
196 + "_functions/models/LiteLLMChatWrapper/unified_turn/start": "DrainProviderUsage",
197 + "_functions/models/LiteLLMChatWrapper/unified_turn/end": "RestoreProviderResponse",
198 "message_loop_prompts_after": "CaptureContextUsage",
199 }
200 for point, class_name in expected.items():
@@ -178,6 +217,75 @@ def test_source_prompt_extensions_are_registered():
217 assert any(cls.__name__ == recorder for cls in classes)
218
219
220 +@pytest.mark.asyncio
221 +async def test_chat_stream_drains_terminal_provider_usage(monkeypatch):
222 + response = '{"tool_name":"response","tool_args":{"text":"done"}}'
223 + chunks = [
224 + {"choices": [{"delta": {"content": response}, "message": {}}]},
225 + {"choices": [{"delta": {"content": " ignored"}, "message": {}}]},
226 + {
227 + "choices": [],
228 + "usage": {
229 + "prompt_tokens": 12_000,
230 + "prompt_tokens_details": {"cached_tokens": 9_000},
231 + "completion_tokens": 80,
232 + },
233 + "_hidden_params": {"response_cost": 0.0042},
234 + },
235 + ]
236 + consumed = []
237 +
238 + async def stream():
239 + for chunk in chunks:
240 + consumed.append(chunk)
241 + yield chunk
242 +
243 + async def fake_acompletion(*args, **kwargs):
244 + assert kwargs["stream_options"] == {"include_usage": True}
245 + return stream()
246 +
247 + async def fake_rate_limiter(*args, **kwargs):
248 + return None
249 +
250 + callback_calls = []
251 +
252 + async def response_callback(chunk: str, full: str):
253 + callback_calls.append((chunk, full))
254 + return full if extract_tools.extract_tool_request(full) else None
255 +
256 + monkeypatch.setattr(litellm_transport, "acompletion", fake_acompletion)
257 + monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter)
258 + wrapper = models.LiteLLMChatWrapper(
259 + model="test-model",
260 + provider="openrouter",
261 + model_config=None,
262 + api_base="https://openrouter.ai/api/v1",
263 + )
264 +
265 + result = await wrapper.unified_turn(
266 + messages=[
267 + SystemMessage(content="stable instructions"),
268 + HumanMessage(content="question"),
269 + ],
270 + response_callback=response_callback,
271 + explicit_caching=True,
272 + )
273 +
274 + assert consumed == chunks
275 + assert callback_calls == [
276 + (response, response),
277 + (" ignored", response + " ignored"),
278 + ]
279 + assert result.response == response
280 + assert result.output_items[0].type == "message"
281 + assert result.usage == {
282 + "prompt_tokens": 12_000,
283 + "prompt_tokens_details": {"cached_tokens": 9_000},
284 + "completion_tokens": 80,
285 + "cost": 0.0042,
286 + }
287 +
288 +
289 def test_prompt_fragment_cache_is_bounded_and_content_addressed(monkeypatch):
290 calls = []
291 agent = SimpleNamespace(data={}, loop_data=LoopData())
@@ -248,3 +356,72 @@ def test_history_ledger_changes_without_invalidating_fragment_cache(monkeypatch)
356 assert second["messages"] == 400 - second["skills"]
357 assert calls.count("stable tools") == 1
358 assert len(agent.data[usage.CACHE_KEY]) == 3
359 +
360 +
361 +def test_rendered_history_fallback_preserves_system_prompt_bucket(monkeypatch):
362 + data = {}
363 + output = [{"ai": False, "content": "short message"}]
364 + agent = SimpleNamespace(
365 + DATA_NAME_CTX_WINDOW="ctx_window",
366 + data=data,
367 + loop_data=LoopData(),
368 + history=SimpleNamespace(get_tokens=lambda: 10_000),
369 + _build_context_message=lambda *args, **kwargs: [],
370 + get_data=lambda key: data.get(key),
371 + set_data=lambda key, value: data.__setitem__(key, value),
372 + )
373 + loop_data = SimpleNamespace(
374 + history_output=output,
375 + protocol_persistent={},
376 + protocol_temporary={},
377 + extras_persistent={},
378 + extras_temporary={},
379 + )
380 + monkeypatch.setattr(
381 + usage.tokens,
382 + "approximate_prompt_tokens",
383 + lambda text: len(text),
384 + )
385 +
386 + usage.reset(agent)
387 + usage.capture_context(agent, loop_data)
388 + data[agent.DATA_NAME_CTX_WINDOW] = {"tokens": 100}
389 + usage.finalize(agent)
390 +
391 + breakdown = data[agent.DATA_NAME_CTX_WINDOW]["usage"]
392 + assert breakdown["messages"] == len("user: short message")
393 + assert breakdown["system_prompt"] > 0
394 + assert sum(breakdown.values()) == 100
395 +
396 +
397 +def test_provider_usage_is_optional():
398 + data = {}
399 + agent = SimpleNamespace(
400 + data=data,
401 + history=SimpleNamespace(all_messages=lambda: []),
402 + set_data=lambda key, value: data.__setitem__(key, value),
403 + )
404 + result = LLMResult.from_chat(
405 + response="done",
406 + usage={
407 + "prompt_tokens": 12_000,
408 + "prompt_tokens_details": {"cached_tokens": 9_000},
409 + "completion_tokens": 80,
410 + "cost": 0.0123,
411 + },
412 + )
413 +
414 + usage.capture_provider_usage(agent, result)
415 +
416 + assert usage.latest_provider_usage(agent) == {
417 + "input_tokens": 12_000,
418 + "cached_tokens": 9_000,
419 + "output_tokens": 80,
420 + "cost": 0.0123,
421 + }
422 +
423 + usage.capture_provider_usage(agent, LLMResult.from_chat(response="no usage"))
424 + assert usage.latest_provider_usage(agent) == {}
425 + assert usage.provider_usage_snapshot(
426 + {"input_tokens": 100, "cached_tokens": None, "cost": None}
427 + ) == {"input_tokens": 100}
plugins/_context_window/webui/context-window-store.js
+52 -9
@@ -5,13 +5,18 @@ import { store as preferencesStore } from "/components/sidebar/bottom/preference
5
6 const API_PATH = "/plugins/_context_window/context_window";
7 const ROWS = [
8 - { key: "messages", label: "Messages", opacity: 1 },
9 - { key: "system_tools", label: "System tools", opacity: 0.88 },
10 - { key: "skills", label: "Skills", opacity: 0.76 },
11 - { key: "mcp_tools", label: "MCP tools", opacity: 0.64 },
12 - { key: "system_prompt", label: "System prompt", opacity: 0.52 },
13 - { key: "extras", label: "Extras", opacity: 0.4 },
8 + { key: "messages", label: "Messages" },
9 + { key: "system_tools", label: "System tools" },
10 + { key: "skills", label: "Skills" },
11 + { key: "mcp_tools", label: "MCP tools" },
12 + { key: "system_prompt", label: "System prompt" },
13 + { key: "extras", label: "Extras" },
14 ];
15 +const COST_FORMATTER = new Intl.NumberFormat("en-US", {
16 + style: "currency",
17 + currency: "USD",
18 + maximumSignificantDigits: 3,
19 +});
20
21 preferencesStore.registerUiControlVisibility("contextWindowUsage", {
22 mobile: true,
@@ -31,6 +36,44 @@ function formatPercent(value) {
36 return `${Number.isInteger(rounded) ? rounded.toFixed(0) : rounded.toFixed(1)}%`;
37 }
38
39 +function optionalNumber(value, key) {
40 + if (!value || !Object.prototype.hasOwnProperty.call(value, key)) return null;
41 + const number = Number(value[key]);
42 + return Number.isFinite(number) && number >= 0 ? number : null;
43 +}
44 +
45 +function formatCost(value) {
46 + if (value === 0) return "$0";
47 + return value < 0.001 ? "<$0.001" : COST_FORMATTER.format(value);
48 +}
49 +
50 +function buildProviderUsage(value = {}) {
51 + const input = optionalNumber(value, "input_tokens");
52 + const cached = optionalNumber(value, "cached_tokens");
53 + const output = optionalNumber(value, "output_tokens");
54 + const cost = optionalNumber(value, "cost");
55 +
56 + const tokenSummary = input === null && output === null
57 + ? ""
58 + : `${input === null ? "–" : formatTokens(input)} → ${output === null ? "–" : formatTokens(output)}`;
59 +
60 + const cachePercent = input > 0 && cached !== null
61 + ? Math.min((cached / input) * 100, 100)
62 + : null;
63 + return {
64 + hasData: cost !== null || cachePercent !== null || Boolean(tokenSummary),
65 + price: {
66 + hasData: cost !== null,
67 + label: cost === null ? "" : formatCost(cost),
68 + },
69 + cache: {
70 + hasData: cachePercent !== null,
71 + label: cachePercent === null ? "" : `${Math.round(cachePercent)}%`,
72 + },
73 + tokens: tokenSummary,
74 + };
75 +}
76 +
77 function buildUsage(data = {}) {
78 const tokens = Math.max(Number(data.tokens) || 0, 0);
79 const contextWindow = Math.max(Number(data.context_window) || 0, 0);
@@ -43,7 +86,6 @@ function buildUsage(data = {}) {
86 ...row,
87 tokensLabel: formatTokens(rowTokens),
88 percentLabel: formatPercent(rowPercent),
46 - dotStyle: `opacity:${row.opacity}`,
89 };
90 });
91 const hasBreakdown = rows.some(row => Number(breakdown[row.key]) > 0);
@@ -55,7 +97,6 @@ function buildUsage(data = {}) {
97 label: "Free space",
98 tokensLabel: formatTokens(freeTokens),
99 percentLabel: formatPercent(freePercent),
58 - dotStyle: "opacity:0.24",
100 });
101 }
102 const percentLabel = formatPercent(percent);
@@ -66,8 +107,10 @@ function buildUsage(data = {}) {
107 ariaLabel: `Context window ${percentLabel} used`,
108 ringLabel: contextWindow ? `${Math.round(percent)}%` : "–",
109 ringDasharray: `${Math.min(percent, 100)} 100`,
69 - summary: `${formatTokens(tokens)}/${contextWindow ? formatTokens(contextWindow) : "–"} (${percentLabel})`,
110 + summaryTokens: `${formatTokens(tokens)}/${contextWindow ? formatTokens(contextWindow) : "–"} tokens`,
111 + summaryPercent: `${percentLabel} used`,
112 meterStyle: `width:${Math.min(percent, 100)}%`,
113 + provider: buildProviderUsage(data.provider_usage),
114 };
115 }
116
tests/test_stream_tool_early_stop.py
+2 -1
@@ -449,7 +449,8 @@ async def test_unified_turn_stops_chat_stream_after_text_tool_request(monkeypatc
449 async def response_callback(chunk: str, full: str):
450 return full if extract_tools.extract_tool_request(full) else None
451
452 - result = await wrapper.unified_turn(
452 + result = await wrapper.unified_turn.__wrapped__(
453 + wrapper,
454 messages=[],
455 response_callback=response_callback,
456 )