| 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 |
| 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 | |
| 21 | class _Log: |
| 22 | def set_progress(self, _message: str) -> None: |
| 23 | pass |
| 24 | |
| 25 | |
| 26 | @pytest.mark.asyncio |
| 27 | async def test_usage_follows_prompt_sources_and_reconciles_to_total(monkeypatch): |
| 28 | agent = object.__new__(Agent) |
| 29 | loop_data = LoopData() |
| 30 | agent.loop_data = loop_data |
| 31 | agent.context = SimpleNamespace(log=_Log()) |
| 32 | agent.history = history.History(agent) |
| 33 | agent.data = {} |
| 34 | agent.history.add_message(False, "User asks a question.") |
| 35 | agent.history.add_message(True, "Assistant answers.") |
| 36 | agent.history.add_message( |
| 37 | False, |
| 38 | { |
| 39 | "tool_name": "skills_tool", |
| 40 | "tool_result": "Skill instructions without a special heading.", |
| 41 | "skill_instructions": { |
| 42 | "name": "test-skill", |
| 43 | "content_included": True, |
| 44 | }, |
| 45 | }, |
| 46 | ) |
| 47 | |
| 48 | system_parts = { |
| 49 | "system_prompt": "Main instructions without a special heading.", |
| 50 | "system_tools": "Tool definitions without a special heading.", |
| 51 | "mcp_tools": "Remote definitions without a special heading.", |
| 52 | "skills": "Available skill names without a special heading.", |
| 53 | } |
| 54 | |
| 55 | async def get_system_prompt(_loop_data): |
| 56 | for key in ("system_tools", "mcp_tools", "skills"): |
| 57 | usage.record_prompt(agent, key, system_parts[key]) |
| 58 | return list(system_parts.values()) |
| 59 | |
| 60 | def read_prompt(prompt_file: str, **kwargs) -> str: |
| 61 | if prompt_file == "agent.context.protocol.md": |
| 62 | return "[PROTOCOL]\n" + kwargs["protocol"] |
| 63 | if prompt_file == "agent.context.extras.md": |
| 64 | return "[EXTRAS]\n" + kwargs["extras"] |
| 65 | raise AssertionError(f"Unexpected prompt: {prompt_file}") |
| 66 | |
| 67 | async def call_extensions(extension_point: str, agent=None, **kwargs): |
| 68 | if extension_point == "message_loop_prompts_after": |
| 69 | current = kwargs["loop_data"] |
| 70 | current.protocol_persistent["project"] = "Project instructions." |
| 71 | current.extras_temporary["time"] = "Current time." |
| 72 | usage.capture_context(agent, current) |
| 73 | |
| 74 | agent.get_system_prompt = get_system_prompt |
| 75 | agent.read_prompt = read_prompt |
| 76 | monkeypatch.setattr(extension, "call_extensions_async", call_extensions) |
| 77 | monkeypatch.setattr(history.History, "_get_max_embeds", lambda self: 0) |
| 78 | |
| 79 | usage.reset(agent) |
| 80 | await Agent.prepare_prompt.__wrapped__(agent, loop_data) |
| 81 | usage.finalize(agent) |
| 82 | |
| 83 | window = agent.get_data(Agent.DATA_NAME_CTX_WINDOW) |
| 84 | breakdown = window["usage"] |
| 85 | assert tuple(breakdown) == usage.USAGE_KEYS |
| 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 |
| 93 | async def test_api_returns_only_counts_and_effective_limit(monkeypatch): |
| 94 | agent = SimpleNamespace( |
| 95 | DATA_NAME_CTX_WINDOW="ctx_window", |
| 96 | get_data=lambda _key: { |
| 97 | "text": "private prompt", |
| 98 | "tokens": 120, |
| 99 | "usage": {"messages": 42}, |
| 100 | }, |
| 101 | ) |
| 102 | handler = object.__new__(ContextWindow) |
| 103 | handler.use_context = lambda _context_id: SimpleNamespace( |
| 104 | streaming_agent=None, |
| 105 | agent0=agent, |
| 106 | ) |
| 107 | monkeypatch.setattr( |
| 108 | "plugins._context_window.api.context_window.get_chat_model_config", |
| 109 | lambda _agent: {"ctx_length": 128_000}, |
| 110 | ) |
| 111 | |
| 112 | result = await handler.process({"context": "ctx-1"}, SimpleNamespace()) |
| 113 | |
| 114 | assert result == { |
| 115 | "tokens": 120, |
| 116 | "context_window": 128_000, |
| 117 | "usage": { |
| 118 | "messages": 42, |
| 119 | "system_tools": 0, |
| 120 | "skills": 0, |
| 121 | "mcp_tools": 0, |
| 122 | "system_prompt": 0, |
| 123 | "extras": 0, |
| 124 | }, |
| 125 | "provider_usage": {}, |
| 126 | } |
| 127 | assert "text" not in result |
| 128 | |
| 129 | |
| 130 | def test_webui_and_accounting_are_plugin_owned(): |
| 131 | model_switcher = ( |
| 132 | ROOT |
| 133 | / "plugins/_model_config/extensions/webui/chat-input-progress-start/model-switcher.html" |
| 134 | ).read_text(encoding="utf-8") |
| 135 | model_store = (ROOT / "plugins/_model_config/webui/switcher-mixin.js").read_text( |
| 136 | encoding="utf-8" |
| 137 | ) |
| 138 | component = ( |
| 139 | ROOT |
| 140 | / "plugins/_context_window/extensions/webui/model-context-strip-end/context-window.html" |
| 141 | ).read_text(encoding="utf-8") |
| 142 | context_store = ( |
| 143 | ROOT / "plugins/_context_window/webui/context-window-store.js" |
| 144 | ).read_text(encoding="utf-8") |
| 145 | helper = (ROOT / "plugins/_context_window/helpers/usage.py").read_text( |
| 146 | encoding="utf-8" |
| 147 | ) |
| 148 | refresh_hook = ( |
| 149 | ROOT |
| 150 | / "plugins/_context_window/extensions/webui/apply_snapshot_before/refresh-context-window.js" |
| 151 | ).read_text(encoding="utf-8") |
| 152 | |
| 153 | assert 'id="model-context-strip-end"' in model_switcher |
| 154 | assert "contextWindowUsage" not in model_switcher |
| 155 | assert "contextUsage" not in model_store |
| 156 | assert "Context window" in component |
| 157 | assert "position: static" in component |
| 158 | assert "width: min(19rem, calc(100vw - 2rem))" in component |
| 159 | assert "right: 1.25rem" in component |
| 160 | assert "width: min(17rem, calc(100vw - 3rem))" in component |
| 161 | assert 'label: "Free space"' in context_store |
| 162 | assert "Last model call" not in component |
| 163 | assert ">Price<" in component |
| 164 | assert ">Cache hit<" in component |
| 165 | assert ">Tokens In/Out<" in component |
| 166 | assert "context-window-cache-meter" not in component |
| 167 | assert "price: {" in context_store |
| 168 | assert "hasData: cost !== null" in context_store |
| 169 | assert 'label: cost === null ? "" : formatCost(cost)' in context_store |
| 170 | assert "usage.provider.price.hasData" in component |
| 171 | assert "usage.provider.price.label" in component |
| 172 | assert 'value < 0.001 ? "<$0.001"' in context_store |
| 173 | assert "maximumSignificantDigits: 3" in context_store |
| 174 | assert "border-top: 1px solid var(--color-border)" in component |
| 175 | assert " → " in context_store |
| 176 | assert "summaryTokens" in context_store |
| 177 | assert 'summaryPercent: `${percentLabel} used`' in context_store |
| 178 | assert "formatTokens(output)} tok" not in context_store |
| 179 | assert "context-window-summary-tokens" in component |
| 180 | assert "context-window-summary-percent" in component |
| 181 | assert "font-family: var(--font-family-main)" in component |
| 182 | assert "<details" not in component |
| 183 | assert "Reasoning" not in component |
| 184 | assert "Images sent" not in component |
| 185 | assert "provider did not split out their token cost" not in context_store |
| 186 | assert "cached / input" in context_store |
| 187 | assert "Math.round(cachePercent)" in context_store |
| 188 | assert "context-window-dot" not in component |
| 189 | assert "dotStyle" not in context_store |
| 190 | assert "Breakdown available after the next message." in component |
| 191 | assert "startswith(" not in helper |
| 192 | assert "rpartition(" not in helper |
| 193 | assert 'item?.type !== "agent"' in refresh_hook |
| 194 | assert "Number(item.agentno || 0) !== 0" in refresh_hook |
| 195 | assert "generationKey === lastGenerationKey" in refresh_hook |
| 196 | |
| 197 | |
| 198 | def test_source_prompt_extensions_are_registered(): |
| 199 | expected = { |
| 200 | "_functions/agent/Agent/prepare_prompt/start": "ResetContextUsage", |
| 201 | "_functions/agent/Agent/prepare_prompt/end": "StoreContextUsage", |
| 202 | "_functions/agent/Agent/call_chat_model_turn/end": "RecordProviderUsage", |
| 203 | "_functions/models/LiteLLMChatWrapper/unified_turn/start": "DrainProviderUsage", |
| 204 | "_functions/models/LiteLLMChatWrapper/unified_turn/end": "RestoreProviderResponse", |
| 205 | "message_loop_prompts_after": "CaptureContextUsage", |
| 206 | } |
| 207 | for point, class_name in expected.items(): |
| 208 | classes = extension._get_extension_classes(point) # type: ignore[attr-defined] |
| 209 | assert any(cls.__name__ == class_name for cls in classes) |
| 210 | |
| 211 | system_prompt_classes = { |
| 212 | cls.__name__: cls |
| 213 | for cls in extension._get_extension_classes("system_prompt") # type: ignore[attr-defined] |
| 214 | } |
| 215 | for owner, recorder in { |
| 216 | "ToolsPrompt": "RecordSystemToolsUsage", |
| 217 | "MCPToolsPrompt": "RecordMcpToolsUsage", |
| 218 | "SkillsPrompt": "RecordSkillsUsage", |
| 219 | }.items(): |
| 220 | builder = system_prompt_classes[owner].execute.__globals__["build_prompt"] |
| 221 | module = builder.__wrapped__.__module__.replace(".", "/") |
| 222 | point = f"_functions/{module}/build_prompt/end" |
| 223 | classes = extension._get_extension_classes(point) # type: ignore[attr-defined] |
| 224 | assert any(cls.__name__ == recorder for cls in classes) |
| 225 | |
| 226 | |
| 227 | @pytest.mark.asyncio |
| 228 | async def test_chat_stream_drains_terminal_provider_usage(monkeypatch): |
| 229 | response = '{"tool_name":"response","tool_args":{"text":"done"}}' |
| 230 | chunks = [ |
| 231 | {"choices": [{"delta": {"content": response}, "message": {}}]}, |
| 232 | {"choices": [{"delta": {"content": " ignored"}, "message": {}}]}, |
| 233 | { |
| 234 | "choices": [], |
| 235 | "usage": { |
| 236 | "prompt_tokens": 12_000, |
| 237 | "prompt_tokens_details": {"cached_tokens": 9_000}, |
| 238 | "completion_tokens": 80, |
| 239 | }, |
| 240 | "_hidden_params": {"response_cost": 0.0042}, |
| 241 | }, |
| 242 | ] |
| 243 | consumed = [] |
| 244 | |
| 245 | async def stream(): |
| 246 | for chunk in chunks: |
| 247 | consumed.append(chunk) |
| 248 | yield chunk |
| 249 | |
| 250 | async def fake_acompletion(*args, **kwargs): |
| 251 | assert kwargs["stream_options"] == {"include_usage": True} |
| 252 | return stream() |
| 253 | |
| 254 | async def fake_rate_limiter(*args, **kwargs): |
| 255 | return None |
| 256 | |
| 257 | callback_calls = [] |
| 258 | |
| 259 | async def response_callback(chunk: str, full: str): |
| 260 | callback_calls.append((chunk, full)) |
| 261 | return full if extract_tools.extract_tool_request(full) else None |
| 262 | |
| 263 | monkeypatch.setattr(litellm_transport, "acompletion", fake_acompletion) |
| 264 | monkeypatch.setattr(models, "apply_rate_limiter", fake_rate_limiter) |
| 265 | wrapper = models.LiteLLMChatWrapper( |
| 266 | model="test-model", |
| 267 | provider="openrouter", |
| 268 | model_config=None, |
| 269 | api_base="https://openrouter.ai/api/v1", |
| 270 | ) |
| 271 | |
| 272 | result = await wrapper.unified_turn( |
| 273 | messages=[ |
| 274 | SystemMessage(content="stable instructions"), |
| 275 | HumanMessage(content="question"), |
| 276 | ], |
| 277 | response_callback=response_callback, |
| 278 | explicit_caching=True, |
| 279 | ) |
| 280 | |
| 281 | assert consumed == chunks |
| 282 | assert callback_calls == [ |
| 283 | (response, response), |
| 284 | (" ignored", response + " ignored"), |
| 285 | ] |
| 286 | assert result.response == response |
| 287 | assert result.output_items[0].type == "message" |
| 288 | assert result.usage == { |
| 289 | "prompt_tokens": 12_000, |
| 290 | "prompt_tokens_details": {"cached_tokens": 9_000}, |
| 291 | "completion_tokens": 80, |
| 292 | "cost": 0.0042, |
| 293 | } |
| 294 | |
| 295 | |
| 296 | def test_prompt_fragment_cache_is_bounded_and_content_addressed(monkeypatch): |
| 297 | calls = [] |
| 298 | agent = SimpleNamespace(data={}, loop_data=LoopData()) |
| 299 | monkeypatch.setattr( |
| 300 | usage.tokens, |
| 301 | "approximate_prompt_tokens", |
| 302 | lambda text: calls.append(text) or len(text), |
| 303 | ) |
| 304 | |
| 305 | usage.reset(agent) |
| 306 | usage.record_prompt(agent, "system_tools", "same prompt") |
| 307 | usage.record_prompt(agent, "system_tools", "same prompt") |
| 308 | usage.record_prompt(agent, "system_tools", "changed prompt") |
| 309 | |
| 310 | assert calls == ["same prompt", "changed prompt"] |
| 311 | cache = agent.data[usage.CACHE_KEY] |
| 312 | assert len(cache) == 1 |
| 313 | assert cache["prompt:system_tools"][1] == len("changed prompt") |
| 314 | assert all(len(value[0]) == 64 for value in cache.values()) |
| 315 | |
| 316 | |
| 317 | def test_history_ledger_changes_without_invalidating_fragment_cache(monkeypatch): |
| 318 | calls = [] |
| 319 | history_tokens = 1_000 |
| 320 | agent = SimpleNamespace( |
| 321 | data={}, |
| 322 | loop_data=LoopData(), |
| 323 | history=SimpleNamespace(get_tokens=lambda: history_tokens), |
| 324 | _build_context_message=lambda *args, **kwargs: [], |
| 325 | ) |
| 326 | skill_message = { |
| 327 | "ai": False, |
| 328 | "content": { |
| 329 | "tool_name": "skills_tool", |
| 330 | "tool_result": "Loaded skill body.", |
| 331 | "skill_instructions": { |
| 332 | "name": "test-skill", |
| 333 | "content_included": True, |
| 334 | }, |
| 335 | }, |
| 336 | } |
| 337 | loop_data = SimpleNamespace( |
| 338 | history_output=[skill_message], |
| 339 | protocol_persistent={}, |
| 340 | protocol_temporary={}, |
| 341 | extras_persistent={}, |
| 342 | extras_temporary={}, |
| 343 | ) |
| 344 | monkeypatch.setattr( |
| 345 | usage.tokens, |
| 346 | "approximate_prompt_tokens", |
| 347 | lambda text: calls.append(text) or len(text), |
| 348 | ) |
| 349 | |
| 350 | usage.reset(agent) |
| 351 | usage.record_prompt(agent, "system_tools", "stable tools") |
| 352 | usage.capture_context(agent, loop_data) |
| 353 | first = dict(agent.loop_data.params_temporary[usage.PARTS_KEY]) |
| 354 | |
| 355 | history_tokens = 400 |
| 356 | agent.loop_data.params_temporary = {} |
| 357 | usage.reset(agent) |
| 358 | usage.record_prompt(agent, "system_tools", "stable tools") |
| 359 | usage.capture_context(agent, loop_data) |
| 360 | second = agent.loop_data.params_temporary[usage.PARTS_KEY] |
| 361 | |
| 362 | assert first["messages"] == 1_000 - first["skills"] |
| 363 | assert second["messages"] == 400 - second["skills"] |
| 364 | assert calls.count("stable tools") == 1 |
| 365 | assert len(agent.data[usage.CACHE_KEY]) == 3 |
| 366 | |
| 367 | |
| 368 | def test_rendered_history_fallback_preserves_system_prompt_bucket(monkeypatch): |
| 369 | data = {} |
| 370 | output = [{"ai": False, "content": "short message"}] |
| 371 | agent = SimpleNamespace( |
| 372 | DATA_NAME_CTX_WINDOW="ctx_window", |
| 373 | data=data, |
| 374 | loop_data=LoopData(), |
| 375 | history=SimpleNamespace(get_tokens=lambda: 10_000), |
| 376 | _build_context_message=lambda *args, **kwargs: [], |
| 377 | get_data=lambda key: data.get(key), |
| 378 | set_data=lambda key, value: data.__setitem__(key, value), |
| 379 | ) |
| 380 | loop_data = SimpleNamespace( |
| 381 | history_output=output, |
| 382 | protocol_persistent={}, |
| 383 | protocol_temporary={}, |
| 384 | extras_persistent={}, |
| 385 | extras_temporary={}, |
| 386 | ) |
| 387 | monkeypatch.setattr( |
| 388 | usage.tokens, |
| 389 | "approximate_prompt_tokens", |
| 390 | lambda text: len(text), |
| 391 | ) |
| 392 | |
| 393 | usage.reset(agent) |
| 394 | usage.capture_context(agent, loop_data) |
| 395 | data[agent.DATA_NAME_CTX_WINDOW] = {"tokens": 100} |
| 396 | usage.finalize(agent) |
| 397 | |
| 398 | breakdown = data[agent.DATA_NAME_CTX_WINDOW]["usage"] |
| 399 | assert breakdown["messages"] == len("user: short message") |
| 400 | assert breakdown["system_prompt"] > 0 |
| 401 | assert sum(breakdown.values()) == 100 |
| 402 | |
| 403 | |
| 404 | def test_provider_usage_is_optional(): |
| 405 | data = {} |
| 406 | agent = SimpleNamespace( |
| 407 | data=data, |
| 408 | history=SimpleNamespace(all_messages=lambda: []), |
| 409 | set_data=lambda key, value: data.__setitem__(key, value), |
| 410 | ) |
| 411 | result = LLMResult.from_chat( |
| 412 | response="done", |
| 413 | usage={ |
| 414 | "prompt_tokens": 12_000, |
| 415 | "prompt_tokens_details": {"cached_tokens": 9_000}, |
| 416 | "completion_tokens": 80, |
| 417 | "cost": 0.0123, |
| 418 | }, |
| 419 | ) |
| 420 | |
| 421 | usage.capture_provider_usage(agent, result) |
| 422 | |
| 423 | assert usage.latest_provider_usage(agent) == { |
| 424 | "input_tokens": 12_000, |
| 425 | "cached_tokens": 9_000, |
| 426 | "output_tokens": 80, |
| 427 | "cost": 0.0123, |
| 428 | } |
| 429 | |
| 430 | usage.capture_provider_usage(agent, LLMResult.from_chat(response="no usage")) |
| 431 | assert usage.latest_provider_usage(agent) == {} |
| 432 | assert usage.provider_usage_snapshot( |
| 433 | {"input_tokens": 100, "cached_tokens": None, "cost": None} |
| 434 | ) == {"input_tokens": 100} |