From 288bd62e80ea62523113b2b245e713c3444f6af4 Mon Sep 17 00:00:00 2001 From: Abhyudaya Baiju Tharakan Date: Mon, 26 Oct 2026 21:28:50 +0530 Subject: [PATCH 1/5] feat(lean): add /lean routes that keep Haiku requests cached and under the long-prompt price /lean/v1/messages, /lean/v1/chat/completions and /lean/v1/responses serve the unified routes. For Anthropic API models the upstream body gains automatic prompt caching, server-side clearing of old tool results, deterministic shell-output condensing and head-and-tail clipping of oversized results. A prompt counted above 100,000 tokens is refused with each protocol's own context-overflow error, so the harness compacts instead of paying the long-prompt price. Every change depends only on the request, so later turns repeat earlier bytes and replayed thinking blocks stay valid. --- .env.example | 12 + README.md | 6 + app.py | 5 + docs/lean.md | 161 +++++++++++ middleware/lean_path.py | 43 +++ routes/lean.py | 55 ++++ routes/unified.py | 11 +- services/lean_mode.py | 497 ++++++++++++++++++++++++++++++++ services/native_token_count.py | 39 ++- tests/outbound_inventory.json | 20 +- tests/test_anthropic_worker.mjs | 10 + tests/test_lean_mode.py | 427 +++++++++++++++++++++++++++ worker/api-paths.mjs | 1 + worker/container-env.mjs | 10 + 14 files changed, 1288 insertions(+), 9 deletions(-) create mode 100644 docs/lean.md create mode 100644 middleware/lean_path.py create mode 100644 routes/lean.py create mode 100644 services/lean_mode.py create mode 100644 tests/test_lean_mode.py diff --git a/.env.example b/.env.example index 862a5e25..ba9694c6 100644 --- a/.env.example +++ b/.env.example @@ -85,6 +85,18 @@ NANOGPT_KEY_CHECK_EVERY_REQUESTS=50 NANOGPT_KEY_REJECTED_COOLDOWN_SECONDS=60 NAVYAI_MAX_REQUEST_BYTES=33554432 +# Lean endpoint (/lean/v1/...) for Claude API models; see docs/lean.md. +# LEAN_ENDPOINT_ENABLED=true +# LEAN_PROMPT_TOKEN_LIMIT=100000 +# LEAN_CLEAR_TRIGGER_TOKENS=80000 +# LEAN_CLEAR_KEEP_TOOL_USES=3 +# LEAN_CLEAR_AT_LEAST_TOKENS=20000 +# LEAN_CLEAR_EXCLUDE_TOOLS= +# LEAN_CACHE_TTL=5m +# LEAN_CONDENSE_SHELL_OUTPUT=true +# LEAN_TOOL_RESULT_MAX_CHARS=100000 +# LEAN_STYLE=none + # Process-local recovery for normalized transports. Fidelity-first raw routes # remain retry- and circuit-free. The default two half-open probes support two # concurrent Codex instances while keeping recovery pressure bounded. State diff --git a/README.md b/README.md index ba4f2efb..57e6b7fb 100644 --- a/README.md +++ b/README.md @@ -349,6 +349,12 @@ the proxy with `ANTHROPIC_BASE_URL`. See [protocol translation](docs/protocol-translation.md) for the coverage and limits. +For Claude Haiku 5.5 in a coding harness, use the `/lean` base path +(`/lean/v1/messages`, `/lean/v1/chat/completions`, `/lean/v1/responses`). It +turns on prompt caching, clears old tool results server-side, condenses shell +output and keeps every prompt under the 100,000-token long-prompt price. See +[lean endpoint](docs/lean.md). + ```bash # OpenAI-compatible endpoint curl -X POST "http://localhost:1400/openai/v1/chat/completions" \ diff --git a/app.py b/app.py index f27af2ef..2c72ec87 100644 --- a/app.py +++ b/app.py @@ -10,6 +10,7 @@ from env_loader import load_runtime_env from error_handlers import init_error_handlers from middleware.external_origin import TrustedExternalOriginMiddleware +from middleware.lean_path import LeanPathMiddleware from route_helpers import ( api_auth_required, apply_cors_headers, @@ -30,6 +31,7 @@ from routes.gateway_mcp import register_gateway_mcp_routes from routes.gateway_batches import register_gateway_batch_routes from routes.context_pages import register_context_page_routes +from routes.lean import register_lean_routes from routes.optimized import register_optimized_routes from routes.proxy import register_proxy_routes from routes.status_page import register_status_routes @@ -105,6 +107,8 @@ def create_app() -> Flask: # normalization and wins when Flask constructs public URLs. app.wsgi_app = TrustedExternalOriginMiddleware(app.wsgi_app) app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1) + if os.environ.get("LEAN_ENDPOINT_ENABLED", "").strip().lower() not in {"0", "false", "no", "off"}: + app.wsgi_app = LeanPathMiddleware(app.wsgi_app) app.secret_key = runtime_secrets["FLASK_SECRET_KEY"] app.config["JWT_SECRET"] = runtime_secrets["JWT_SECRET"] @@ -146,6 +150,7 @@ def create_app() -> Flask: register_intelligence_media_routes(app, csrf, AuthService, ProxyService) register_free_routes(app, csrf, AuthService, MetricsService, ProxyService) register_optimized_routes(app, csrf, AuthService, MetricsService, ProxyService) + register_lean_routes(app) register_status_routes(app, csrf, AuthService, ProxyService) register_core_routes(app) register_dashboard_security_routes(app) diff --git a/docs/lean.md b/docs/lean.md new file mode 100644 index 00000000..828bbb73 --- /dev/null +++ b/docs/lean.md @@ -0,0 +1,161 @@ +# Lean endpoint + +`/lean/v1/messages`, `/lean/v1/chat/completions` and `/lean/v1/responses` are the +unified routes with extra cost controls for models served by the Claude API +(`anthropic:` models). They are written for coding harnesses on Claude Haiku 5.5, +whose price depends on prompt length: + +| Haiku 5.5, per million tokens | Prompt of 100,000 tokens or fewer | Longer prompt | +|---|---|---| +| Input | $0.10 | $0.50 | +| Cache read | $0.01 | $0.05 | +| Cache write (5 minutes) | $0.125 | $0.625 | +| Output | $0.50 | $2.50 | + +The prompt length counts every input token, cache reads and writes included, and +a request one token over the limit pays the higher price on all of its tokens +([pricing](https://platform.claude.com/docs/en/about-claude/pricing#long-context-pricing)). +Lean mode keeps prompts cached and under that limit. + +Requests to other providers behave exactly as on `/v1` and carry +`X-MultiLLM-Lean: inactive`. `/lean/v1/messages/count_tokens` and `/lean/v1/models` +are the `/v1` routes unchanged, so a harness can use one base URL. + +## What lean mode does + +Every change depends only on the request itself, so each later request of a +conversation repeats the earlier one's bytes up to its new turn. That keeps prompt +caching working, and keeps Haiku 5.5's replayed thinking blocks valid: Haiku 5.5 +rejects a thinking block whose earlier conversation changed. Checked against the +Claude API: a second identical request read its whole 2,006-token prefix from the +cache, and a thinking block Haiku wrote through `/lean` was accepted when the next +turn replayed it through `/lean`, but rejected when the same turn went to `/v1` +with the raw shell output. + +| Stage | Default | What it does | +|---|---|---| +| Prompt caching | on | Adds automatic caching (`cache_control` at the top level) and a breakpoint after the system prompt, or after the last tool when there is none. Skipped when the caller placed its own `cache_control`, as Claude Code does. | +| Tool result clearing | on, from 80,000 tokens | Adds `context_management` with `clear_tool_uses_20250919`: once the prompt passes the trigger, the API replaces old tool results with placeholders and keeps the 3 most recent, clearing at least 20,000 tokens at a time. The API applies it before the model reads the prompt, the harness keeps its full history, and it never invalidates thinking blocks. Skipped when the caller sends its own `context_management` or a compaction block. | +| Shell output | on | Results of shell tools (`Bash`, `bash`, `shell`, `exec_command`, `run_shell_command` and similar) lose ANSI escapes and progress-bar redraws; runs of four or more identical lines become one line and a count, and runs of blank lines become one. Other tools' results, file reads included, keep their exact bytes, so an edit tool can still match them. | +| Oversized results | 100,000 characters | A longer tool result keeps its first and last 50,000 characters with a marker in between. | +| Prompt limit | 100,000 tokens | A prompt estimated near the limit is counted exactly with `count_tokens`, context editing included. Above the limit the request is refused with the error the harness treats as a full context window, so it compacts and retries instead of paying the higher price. | +| Concise answers | off | Adds a short instruction to the system prompt: lead with the answer, no filler, code and commands exact. Turn it on with `LEAN_STYLE=concise` or per request with `X-MultiLLM-Lean-Style: concise` (`caveman` is accepted too). | + +Lean mode never changes the model, effort, thinking settings or `max_tokens`. + +### The refusal + +| Route | Response | +|---|---| +| `/lean/v1/messages` | 400, Anthropic envelope, `invalid_request_error`, `prompt is too long: N tokens > 100000 maximum`. Claude Code compacts on this message. | +| `/lean/v1/chat/completions` | 400, `{"error": {"code": "context_length_exceeded", ...}}` with a message naming the maximum context length. | +| `/lean/v1/responses` with `stream: true` | A `response.failed` event with `context_length_exceeded`, the form Codex reads. Without streaming, the Chat Completions body above. | + +A count that fails (for example, rate-limited) lets the request through and +reports `X-MultiLLM-Lean-Prompt-Count: estimate`. + +### Response headers + +| Header | Meaning | +|---|---| +| `X-MultiLLM-Lean` | `applied` for a Claude API request, `inactive` for other providers | +| `X-MultiLLM-Lean-Rules` | Version of the rules above | +| `X-MultiLLM-Lean-Stages` | Stages that changed the request, for example `cache,clear-tool-uses,condense-shell` | +| `X-MultiLLM-Lean-Prompt-Tokens` | Prompt size: counted by the API near the limit, otherwise a high local estimate | +| `X-MultiLLM-Lean-Prompt-Count` | `provider` or `estimate` | +| `X-MultiLLM-Lean-Condensed-Chars` | Characters removed from tool results | + +The API reports what it cleared in the response's `context_management` field and +cache use in `usage.cache_read_input_tokens` and `usage.cache_creation_input_tokens`. + +## Harness setup + +Set each harness to compact before 100,000 tokens, so the refusal stays a safety +net. + +### Claude Code + +```bash +export ANTHROPIC_BASE_URL="$PROXY_BASE_URL/lean" +export ANTHROPIC_API_KEY="$MULTILLM_API_KEY" +export ANTHROPIC_MODEL="anthropic:claude-haiku-5-5" +export ANTHROPIC_DEFAULT_HAIKU_MODEL="anthropic:claude-haiku-5-5" +export CLAUDE_CODE_AUTO_COMPACT_WINDOW=100000 +claude +``` + +Claude Code places its own cache breakpoints, so lean mode leaves caching to it. + +### Codex + +```toml +model = "anthropic:claude-haiku-5-5" +model_provider = "multillm-lean" +model_context_window = 100000 +model_auto_compact_token_limit = 85000 + +[model_providers.multillm-lean] +name = "MultiLLM lean" +base_url = "https://multillm-proxy.cserules.workers.dev/lean/v1" +env_key = "MULTILLM_API_KEY" +wire_api = "responses" +``` + +### opencode + +```json +{ + "provider": { + "multillm-lean": { + "npm": "@ai-sdk/openai-compatible", + "options": { + "baseURL": "https://multillm-proxy.cserules.workers.dev/lean/v1", + "apiKey": "{env:MULTILLM_API_KEY}" + }, + "models": { + "anthropic:claude-haiku-5-5": { "limit": { "context": 100000, "output": 32000 } } + } + } + } +} +``` + +Chat Completions and Responses requests reach the Claude API translated, and the +translation drops caller `cache_control`, so for Codex and opencode lean mode's +caching is what makes cache reads happen at all. + +## Settings + +| Variable | Default | | +|---|---|---| +| `LEAN_ENDPOINT_ENABLED` | `true` | `false` stops serving `/lean` | +| `LEAN_PROMPT_TOKEN_LIMIT` | `100000` | `0` turns the refusal off | +| `LEAN_CLEAR_TRIGGER_TOKENS` | `80000` | `0` turns tool result clearing off | +| `LEAN_CLEAR_KEEP_TOOL_USES` | `3` | Recent tool uses that are never cleared | +| `LEAN_CLEAR_AT_LEAST_TOKENS` | `20000` | Minimum cleared per edit, so a cache rewrite is worth it | +| `LEAN_CLEAR_EXCLUDE_TOOLS` | empty | Comma-separated tool names whose results are never cleared | +| `LEAN_CACHE_TTL` | `5m` | `1h` for sessions with long pauses; 1-hour writes cost 2x input | +| `LEAN_CONDENSE_SHELL_OUTPUT` | `true` | | +| `LEAN_TOOL_RESULT_MAX_CHARS` | `100000` | `0` turns clipping off | +| `LEAN_STYLE` | `none` | `concise` adds the concise-answer instruction | + +Keep a conversation on one base URL and its settings fixed. Lean mode sends shell +output condensed, so moving a conversation between `/v1` and `/lean`, or changing +`LEAN_CONDENSE_SHELL_OUTPUT`, `LEAN_TOOL_RESULT_MAX_CHARS` or `LEAN_STYLE` partway +through, changes bytes the model already saw. On accounts created on or after +31 August 2026 the next request that replays a thinking block then fails with +``Invalid `signature` in `thinking` block. The block is bound to a different +conversation``; start a new conversation. The cache and clearing settings change +only cost: one cache rewrite. + +## Notes + +- Cleared tool results are gone for the model, not for the harness: if the model + needs one again it reruns the tool. +- The Claude API's prompt cache belongs to the account. The key pool keeps using + the first key until it is refused ([key failover](anthropic.md#spare-keys-and-failover)); a conversation + that moves to the spare key starts with a cold cache, and its earlier thinking + blocks are dropped rather than rejected. +- Research behind these choices, including what Claude Code, Codex, opencode, + Cline and Gemini CLI do, is summarised in the pull request that added this + page. diff --git a/middleware/lean_path.py b/middleware/lean_path.py new file mode 100644 index 00000000..0d6e5c35 --- /dev/null +++ b/middleware/lean_path.py @@ -0,0 +1,43 @@ +"""Serve `/lean/v1/...` from the unified `/v1/...` routes in lean mode.""" + +from collections.abc import Callable, Iterable +from typing import Any + +WSGIStartResponse = Callable[..., Any] +WSGIApplication = Callable[[dict[str, Any], WSGIStartResponse], Iterable[bytes]] + +LEAN_ENVIRON_KEY = "multillm.lean" +LEAN_PREFIX = "/lean" +LEAN_PATHS = frozenset( + { + "/v1/chat/completions", + "/v1/messages", + "/v1/messages/count_tokens", + "/v1/models", + "/v1/responses", + } +) + + +class LeanPathMiddleware: + """Rewrite a lean path to its unified route and mark the request as lean. + + Rewriting before Flask routes the request keeps every path-based rule (rate + limits, body limits, error envelopes, cache affinity) identical to `/v1`. + The mark lives in the WSGI environ, which a client cannot set. + """ + + def __init__(self, application: WSGIApplication) -> None: + self.application = application + + def __call__( + self, + environ: dict[str, Any], + start_response: WSGIStartResponse, + ) -> Iterable[bytes]: + environ.pop(LEAN_ENVIRON_KEY, None) + path = environ.get("PATH_INFO", "") + if path.startswith(f"{LEAN_PREFIX}/") and path[len(LEAN_PREFIX):] in LEAN_PATHS: + environ["PATH_INFO"] = path[len(LEAN_PREFIX):] + environ[LEAN_ENVIRON_KEY] = True + return self.application(environ, start_response) diff --git a/routes/lean.py b/routes/lean.py new file mode 100644 index 00000000..0099b526 --- /dev/null +++ b/routes/lean.py @@ -0,0 +1,55 @@ +"""Lean mode responses: report headers and the protocol's own "prompt is too long" error.""" + +from __future__ import annotations + +import json + +from flask import Response, jsonify, request + +from error_handlers import get_request_id +from services.lean_mode import LeanPromptTooLong, lean_response_headers +from services.protocol_translation import anthropic_error_body + + +def prompt_too_long_response(error: LeanPromptTooLong): + """Refuse in the shape each harness treats as a context overflow, so it compacts. + + Claude Code compacts on Anthropic's "prompt is too long" error. OpenAI-style + clients match `context_length_exceeded`; Codex reads it from a streamed + `response.failed` event. + """ + request_id = get_request_id() + if request.path.startswith("/v1/messages"): + return jsonify(anthropic_error_body(400, str(error), "invalid_request_error", request_id)), 400 + message = ( + f"This model's maximum context length is {error.limit} tokens. However, your " + f"messages resulted in {error.tokens} tokens ({error}). Compact the " + "conversation and try again." + ) + payload = request.get_json(silent=True) + if request.path == "/v1/responses" and isinstance(payload, dict) and payload.get("stream"): + event = { + "type": "response.failed", + "sequence_number": 0, + "response": { + "object": "response", + "status": "failed", + "error": {"code": "context_length_exceeded", "message": message}, + }, + } + return Response( + f"event: response.failed\ndata: {json.dumps(event)}\n\n", + mimetype="text/event-stream", + ) + error_body = { + "message": message, + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded", + } + return jsonify({"error": error_body, "request_id": request_id}), 400 + + +def register_lean_routes(app) -> None: + app.after_request(lean_response_headers) + app.register_error_handler(LeanPromptTooLong, prompt_too_long_response) diff --git a/routes/unified.py b/routes/unified.py index 529821af..fa953a89 100644 --- a/routes/unified.py +++ b/routes/unified.py @@ -68,6 +68,7 @@ from services.managed_turn import bounded_managed_timeout as bounded_timeout from routes.cascades import dispatch_unified_cascade, register_cascade_admin_routes, validate_cascade_target from services.context_optimizer import ContextOptimizationResult +from services.lean_mode import merge_beta_header, prepare_lean_request from services.model_registry import ModelRegistry from services.model_cooldown import ModelCooldownCapacity, ModelCooldownExhausted from routes.provider_credentials import ( @@ -307,7 +308,14 @@ def _send_native_request( ): """Send one body to a provider's native Responses or Messages endpoint.""" from services.managed_turn import prepare_dispatch_kwargs - raw_body = json.dumps(body).encode("utf-8") + lean = prepare_lean_request( + body, + provider=provider, + endpoint=endpoint, + request_headers=headers_source, + base_urls=app.config["API_BASE_URLS"], + ) + raw_body = json.dumps(lean.body).encode("utf-8") def send_request(token: str): headers = proxy_service_cls.prepare_headers( @@ -317,6 +325,7 @@ def send_request(token: str): upstream_path=endpoint, ) _merge_request_headers(headers, cache_headers) + merge_beta_header(headers, lean.betas) default_url = f"{app.config['API_BASE_URLS'][provider].rstrip('/')}/{endpoint}" request_kwargs = { "method": "POST", diff --git a/services/lean_mode.py b/services/lean_mode.py new file mode 100644 index 00000000..760aeb1a --- /dev/null +++ b/services/lean_mode.py @@ -0,0 +1,497 @@ +"""Lean mode: keep Claude API requests cached and under Haiku's long-prompt price. + +`/lean/v1/...` serves the unified Messages, Chat Completions and Responses routes +(middleware/lean_path.py). For a request that reaches the Claude API, the body is +changed only in ways that give every later request of the conversation the same +bytes up to its new turn, so prompt caching and preserved thinking keep working: + +* shell tool output is condensed by rules that read only that output; +* an oversized tool result keeps its head and tail; +* prompt caching is turned on when the caller placed no breakpoints; +* old tool results are cleared by server-side context editing, which never + invalidates thinking blocks; +* a prompt that would still pay the long-prompt price is refused with the + caller's own "prompt is too long" error, so the harness compacts and retries. + +Changing these rules changes the bytes of conversations already in flight and +invalidates their replayed thinking blocks; bump LEAN_RULES_VERSION when you do. +""" + +from __future__ import annotations + +import logging +import os +import re +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import Any + +from flask import g, has_request_context, request + +from middleware.lean_path import LEAN_ENVIRON_KEY +from providers.protocols import MESSAGES +from services.native_token_count import NativeTokenCountError, count_claude_prompt_tokens + +logger = logging.getLogger(__name__) + +LEAN_RULES_VERSION = "1" +CONTEXT_MANAGEMENT_BETA = "context-management-2025-06-27" +CLEAR_TOOL_USES = "clear_tool_uses_20250919" +STYLE_HEADER = "X-MultiLLM-Lean-Style" +# Haiku 5.5 bills a request whose prompt is over 100,000 tokens, cache reads and +# writes included, at five times the price of the whole request. +DEFAULT_PROMPT_TOKEN_LIMIT = 100_000 +DEFAULT_CLEAR_TRIGGER_TOKENS = 80_000 +DEFAULT_CLEAR_KEEP_TOOL_USES = 3 +DEFAULT_CLEAR_AT_LEAST_TOKENS = 20_000 +DEFAULT_TOOL_RESULT_MAX_CHARS = 100_000 +# The local estimate (about three characters per token) can read a third low: +# Haiku 5.5's tokenizer and tool definitions add tokens it does not see. Counting +# from 40% of the limit leaves that margin before a prompt can cross it uncounted. +COUNT_FROM_LIMIT_SHARE = 0.4 +MEDIA_BLOCK_TOKENS = 1_600 +REPEATED_LINE_MIN_RUN = 4 +SHELL_TOOL_NAMES = frozenset( + { + "bash", + "exec_command", + "execute_command", + "local_shell", + "run_command", + "run_shell_command", + "run_terminal_cmd", + "shell", + "shell_command", + "terminal", + } +) +CONCISE_STYLE = ( + "Response style: be brief. Lead with the answer or the next action. Skip " + "preamble, filler, apologies and restating the request. Keep code, commands, " + "file paths, error messages, numbers and negations exact and complete. Write " + "plain full sentences for risky or irreversible steps and whenever the user " + "seems unsure." +) +_ANSI_ESCAPE = re.compile( + r"\x1b\[[0-?]*[ -/]*[@-~]" # CSI: colours, cursor movement, erase line + r"|\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)" # OSC: titles, hyperlinks + r"|\x1b[@-Z\\-_]" # other two-byte escapes +) + + +class LeanPromptTooLong(Exception): + """The prompt, after every lean stage, would still pay the long-prompt price.""" + + def __init__(self, tokens: int, limit: int): + super().__init__(f"prompt is too long: {tokens} tokens > {limit} maximum") + self.tokens = tokens + self.limit = limit + + +@dataclass(frozen=True) +class LeanPolicy: + prompt_token_limit: int = DEFAULT_PROMPT_TOKEN_LIMIT + clear_trigger_tokens: int = DEFAULT_CLEAR_TRIGGER_TOKENS + clear_keep_tool_uses: int = DEFAULT_CLEAR_KEEP_TOOL_USES + clear_at_least_tokens: int = DEFAULT_CLEAR_AT_LEAST_TOKENS + clear_exclude_tools: tuple[str, ...] = () + cache_ttl: str = "5m" + condense_shell_output: bool = True + tool_result_max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS + style: str = "none" + + +@dataclass(frozen=True) +class LeanRequest: + body: dict + active: bool = False + betas: tuple[str, ...] = () + stages: tuple[str, ...] = () + prompt_tokens: int | None = None + prompt_count: str = "" + condensed_chars: int = 0 + + +def _env_int(environ: Mapping[str, str], name: str, default: int) -> int: + value = str(environ.get(name, "")).strip() + try: + return max(0, int(value)) if value else default + except ValueError: + return default + + +def _env_flag(environ: Mapping[str, str], name: str, default: bool) -> bool: + value = str(environ.get(name, "")).strip().lower() + if not value: + return default + return value in {"1", "true", "yes", "on"} + + +def _style(value: Any) -> str | None: + normalized = str(value or "").strip().lower() + if normalized in {"concise", "caveman"}: + return "concise" + if normalized in {"none", "off"}: + return "none" + return None + + +def lean_policy(environ: Mapping[str, str] | None = None) -> LeanPolicy: + environ = os.environ if environ is None else environ + exclude = str(environ.get("LEAN_CLEAR_EXCLUDE_TOOLS", "")) + return LeanPolicy( + prompt_token_limit=_env_int(environ, "LEAN_PROMPT_TOKEN_LIMIT", DEFAULT_PROMPT_TOKEN_LIMIT), + clear_trigger_tokens=_env_int(environ, "LEAN_CLEAR_TRIGGER_TOKENS", DEFAULT_CLEAR_TRIGGER_TOKENS), + clear_keep_tool_uses=_env_int(environ, "LEAN_CLEAR_KEEP_TOOL_USES", DEFAULT_CLEAR_KEEP_TOOL_USES), + clear_at_least_tokens=_env_int(environ, "LEAN_CLEAR_AT_LEAST_TOKENS", DEFAULT_CLEAR_AT_LEAST_TOKENS), + clear_exclude_tools=tuple(name.strip() for name in exclude.split(",") if name.strip()), + cache_ttl="1h" if str(environ.get("LEAN_CACHE_TTL", "")).strip().lower() == "1h" else "5m", + condense_shell_output=_env_flag(environ, "LEAN_CONDENSE_SHELL_OUTPUT", True), + tool_result_max_chars=_env_int(environ, "LEAN_TOOL_RESULT_MAX_CHARS", DEFAULT_TOOL_RESULT_MAX_CHARS), + style=_style(environ.get("LEAN_STYLE")) or "none", + ) + + +def lean_requested() -> bool: + return has_request_context() and request.environ.get(LEAN_ENVIRON_KEY) is True + + +def _header(headers: Mapping[str, Any] | None, name: str) -> str: + for key, value in (headers or {}).items(): + if str(key).lower() == name.lower() and isinstance(value, str): + return value.strip() + return "" + + +def condense_shell_output(text: str) -> str: + """Drop what a terminal shows but a reader does not need, deterministically. + + Removes ANSI escapes, keeps only the last redraw of a carriage-return + progress line, collapses a run of identical lines to one line and a count, + and collapses runs of blank lines. + """ + condensed = _ANSI_ESCAPE.sub("", text) + if "\r" in condensed: + condensed = condensed.replace("\r\n", "\n") + condensed = "\n".join( + next((part for part in reversed(line.split("\r")) if part), "") + for line in condensed.split("\n") + ) + lines = condensed.split("\n") + kept: list[str] = [] + index = 0 + while index < len(lines): + line = lines[index] + end = index + 1 + while end < len(lines) and lines[end] == line: + end += 1 + run = end - index + if not line.strip(): + kept.extend([line] if run < 3 else [""]) + elif run >= REPEATED_LINE_MIN_RUN: + kept.extend([line, f"[previous line repeated {run - 1} more times]"]) + else: + kept.extend(lines[index:end]) + index = end + return "\n".join(kept) + + +def clip_tool_text(text: str, max_chars: int) -> str: + """Keep the head and tail of an oversized tool result, marking the cut.""" + if max_chars <= 0 or len(text) <= max_chars: + return text + head = max_chars // 2 + tail = max_chars - head + omitted = len(text) - head - tail + return ( + f"{text[:head]}\n[MultiLLM lean: {omitted} characters omitted from the middle " + f"of this tool result]\n{text[-tail:]}" + ) + + +def _tool_names(messages: list) -> dict[str, str]: + names: dict[str, str] = {} + for message in messages: + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + content = message.get("content") + for block in content if isinstance(content, list) else (): + if isinstance(block, dict) and block.get("type") == "tool_use" and isinstance(block.get("id"), str): + names[block["id"]] = str(block.get("name") or "") + return names + + +def _condense_text(text: str, *, shell: bool, max_chars: int, stages: set[str]) -> str: + result = text + if shell: + condensed = condense_shell_output(result) + if condensed != result: + stages.add("condense-shell") + result = condensed + clipped = clip_tool_text(result, max_chars) + if clipped != result: + stages.add("clip-tool-result") + result = clipped + return result + + +def _condense_tool_result(block: dict, *, shell: bool, max_chars: int, stages: set[str]) -> dict: + content = block.get("content") + if isinstance(content, str): + condensed = _condense_text(content, shell=shell, max_chars=max_chars, stages=stages) + return block if condensed == content else {**block, "content": condensed} + if not isinstance(content, list): + return block + parts = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text" and isinstance(part.get("text"), str): + condensed = _condense_text(part["text"], shell=shell, max_chars=max_chars, stages=stages) + parts.append(part if condensed == part["text"] else {**part, "text": condensed}) + else: + parts.append(part) + return block if parts == content else {**block, "content": parts} + + +def _condense_tool_results(messages: Any, policy: LeanPolicy, stages: set[str]) -> tuple[Any, int]: + """Condense every tool result by rules that read only that result and its tool's name.""" + if not isinstance(messages, list): + return messages, 0 + names = _tool_names(messages) + before = _text_length(messages) + condensed_messages = [] + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(message, dict) or message.get("role") != "user" or not isinstance(content, list): + condensed_messages.append(message) + continue + blocks = [] + for block in content: + if isinstance(block, dict) and block.get("type") == "tool_result": + tool_name = names.get(block.get("tool_use_id"), "").lower() + block = _condense_tool_result( + block, + shell=policy.condense_shell_output and tool_name in SHELL_TOOL_NAMES, + max_chars=policy.tool_result_max_chars, + stages=stages, + ) + blocks.append(block) + condensed_messages.append(message if blocks == content else {**message, "content": blocks}) + if condensed_messages == messages: + return messages, 0 + return condensed_messages, before - _text_length(condensed_messages) + + +def _text_length(value: Any) -> int: + if isinstance(value, str): + return len(value) + if isinstance(value, dict): + return sum(_text_length(item) for item in value.values()) + if isinstance(value, list): + return sum(_text_length(item) for item in value) + return 0 + + +def _has_cache_control(value: Any) -> bool: + if isinstance(value, dict): + return "cache_control" in value or any(_has_cache_control(item) for item in value.values()) + if isinstance(value, list): + return any(_has_cache_control(item) for item in value) + return False + + +def _with_style(body: dict) -> dict: + style = {"type": "text", "text": CONCISE_STYLE} + system = body.get("system") + if isinstance(system, str) and system: + blocks = [{"type": "text", "text": system}, style] + elif isinstance(system, list): + blocks = [*system, style] + else: + blocks = [style] + return {**body, "system": blocks} + + +def _with_cache_breakpoints(body: dict, ttl: str) -> tuple[dict, bool]: + """Cache the conversation, plus the tools and system prompt on their own. + + The top-level marker moves with the conversation's last block. The second + breakpoint keeps the tools and system prompt cached when context editing + clears an old tool result and the rest of the prefix must be written again. + """ + if _has_cache_control(body): + return body, False + marker = {"type": "ephemeral"} if ttl == "5m" else {"type": "ephemeral", "ttl": ttl} + cached = {**body, "cache_control": dict(marker)} + system = body.get("system") + tools = body.get("tools") + if isinstance(system, str) and system: + cached["system"] = [{"type": "text", "text": system, "cache_control": dict(marker)}] + elif isinstance(system, list) and system and isinstance(system[-1], dict): + cached["system"] = [*system[:-1], {**system[-1], "cache_control": dict(marker)}] + elif isinstance(tools, list) and tools and isinstance(tools[-1], dict): + cached["tools"] = [*tools[:-1], {**tools[-1], "cache_control": dict(marker)}] + return cached, True + + +def _blocks(messages: Any) -> Iterable[dict]: + for message in messages if isinstance(messages, list) else (): + content = message.get("content") if isinstance(message, dict) else None + for block in content if isinstance(content, list) else (): + if isinstance(block, dict): + yield block + + +def _with_tool_clearing(body: dict, policy: LeanPolicy) -> tuple[dict, bool]: + """Clear old tool results server-side once the prompt passes the trigger. + + Skipped when the caller manages context itself, and for signed compaction + blocks, which cannot be combined with context editing. + """ + blocks = list(_blocks(body.get("messages"))) + uses_tools = bool(body.get("tools")) or any(block.get("type") == "tool_result" for block in blocks) + if ( + policy.clear_trigger_tokens <= 0 + or not uses_tools + or "context_management" in body + or "compaction" in body + or any(block.get("type") == "compaction" for block in blocks) + ): + return body, False + edit: dict[str, Any] = { + "type": CLEAR_TOOL_USES, + "trigger": {"type": "input_tokens", "value": policy.clear_trigger_tokens}, + "keep": {"type": "tool_uses", "value": policy.clear_keep_tool_uses}, + } + if policy.clear_at_least_tokens: + edit["clear_at_least"] = {"type": "input_tokens", "value": policy.clear_at_least_tokens} + if policy.clear_exclude_tools: + edit["exclude_tools"] = list(policy.clear_exclude_tools) + return {**body, "context_management": {"edits": [edit]}}, True + + +def estimate_prompt_tokens(body: Mapping[str, Any]) -> int: + """A rough estimate that decides when to count exactly: three characters per + token, images at a flat size.""" + characters = 0 + media = 0 + + def walk(value: Any) -> None: + nonlocal characters, media + if isinstance(value, dict): + if value.get("type") in {"image", "document"}: + media += 1 + return + for key, item in value.items(): + if key != "cache_control": + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + elif isinstance(value, str): + characters += len(value) + + for key in ("system", "tools", "messages"): + walk(body.get(key)) + return characters // 3 + media * MEDIA_BLOCK_TOKENS + + +def merged_beta_value(existing: str, betas: Iterable[str]) -> str: + values = [item.strip() for item in str(existing or "").split(",") if item.strip()] + for beta in betas: + if beta not in values: + values.append(beta) + return ",".join(values) + + +def merge_beta_header(headers: dict, betas: Iterable[str]) -> None: + """Add lean betas to an upstream header dict without dropping the caller's.""" + betas = tuple(betas) + if not betas: + return + name = next((key for key in headers if key.lower() == "anthropic-beta"), "Anthropic-Beta") + headers[name] = merged_beta_value(headers.get(name, ""), betas) + + +def _prompt_size( + body: dict, + betas: tuple[str, ...], + request_headers: Mapping[str, Any] | None, + base_urls: Mapping[str, str], + limit: int, +) -> tuple[int, str]: + estimate = estimate_prompt_tokens(body) + if limit <= 0 or estimate < limit * COUNT_FROM_LIMIT_SHARE: + return estimate, "estimate" + beta_header = merged_beta_value(_header(request_headers, "anthropic-beta"), betas) + try: + return count_claude_prompt_tokens(body, beta_header, base_urls), "provider" + except NativeTokenCountError as error: + # Counting is a guard, not a gate: a failed count lets the request through. + logger.warning("Lean prompt count failed status=%s; sending uncounted", error.status) + return estimate, "estimate" + + +def prepare_lean_request( + body: dict, + *, + provider: str, + endpoint: str, + request_headers: Mapping[str, Any] | None, + base_urls: Mapping[str, str], +) -> LeanRequest: + """Apply lean mode to a Claude API Messages body; other requests pass unchanged. + + Raises LeanPromptTooLong when the prepared prompt would pay the long-prompt price. + """ + if not lean_requested(): + return LeanRequest(body=body) + if provider != "anthropic" or endpoint != MESSAGES: + g.lean_request = LeanRequest(body=body) + return g.lean_request + policy = lean_policy() + stages: set[str] = set() + prepared = dict(body) + messages, condensed_chars = _condense_tool_results(body.get("messages"), policy, stages) + if messages is not body.get("messages"): + prepared["messages"] = messages + if (_style(_header(request_headers, STYLE_HEADER)) or policy.style) == "concise": + prepared = _with_style(prepared) + stages.add("concise-style") + prepared, cached = _with_cache_breakpoints(prepared, policy.cache_ttl) + if cached: + stages.add("cache") + prepared, clearing = _with_tool_clearing(prepared, policy) + betas = (CONTEXT_MANAGEMENT_BETA,) if clearing else () + if clearing: + stages.add("clear-tool-uses") + tokens, source = _prompt_size(prepared, betas, request_headers, base_urls, policy.prompt_token_limit) + decision = LeanRequest( + body=prepared, + active=True, + betas=betas, + stages=tuple(sorted(stages)), + prompt_tokens=tokens, + prompt_count=source, + condensed_chars=max(0, condensed_chars), + ) + g.lean_request = decision + if source == "provider" and tokens > policy.prompt_token_limit: + raise LeanPromptTooLong(tokens, policy.prompt_token_limit) + return decision + + +def lean_response_headers(response): + """Report what lean mode did, so callers can see why a request stayed cheap.""" + if not lean_requested(): + return response + decision = getattr(g, "lean_request", None) + active = decision is not None and decision.active + response.headers["X-MultiLLM-Lean"] = "applied" if active else "inactive" + response.headers["X-MultiLLM-Lean-Rules"] = LEAN_RULES_VERSION + if active: + response.headers["X-MultiLLM-Lean-Stages"] = ",".join(decision.stages) or "none" + response.headers["X-MultiLLM-Lean-Prompt-Tokens"] = str(decision.prompt_tokens) + response.headers["X-MultiLLM-Lean-Prompt-Count"] = decision.prompt_count + if decision.condensed_chars: + response.headers["X-MultiLLM-Lean-Condensed-Chars"] = str(decision.condensed_chars) + return response diff --git a/services/native_token_count.py b/services/native_token_count.py index fb0b38be..51d351b0 100644 --- a/services/native_token_count.py +++ b/services/native_token_count.py @@ -16,6 +16,8 @@ MAX_RESPONSE_BYTES = 64 * 1024 COUNT_FIELDS = frozenset({"model", "messages", "system", "tools", "tool_choice", "thinking"}) COUNT_HEADERS = frozenset({"anthropic-version", "anthropic-beta", "user-agent", "x-request-id"}) +CLAUDE_COUNT_PATH = "/v1/messages/count_tokens" +CLAUDE_COUNT_FIELDS = COUNT_FIELDS | {"context_management"} class NativeTokenCountError(Exception): @@ -100,7 +102,6 @@ def count_native_tokens( # share no credentials or model state through this adapter. from services.auth_service import AuthService from services.model_registry import ModelRegistry - from services.proxy_service import ProxyService model_id = payload.get("model") if not isinstance(model_id, str): @@ -122,11 +123,43 @@ def count_native_tokens( # Proxy credentials and caller-supplied provider credentials are not part of # this configured count adapter's header contract. incoming = {key: value for key, value in request_headers.items() if key.lower() in COUNT_HEADERS} + body = {key: value for key, value in payload.items() if key in COUNT_FIELDS} + body["model"] = model + return _post_count(provider, url, path, credential, incoming, body) + + +def count_claude_prompt_tokens( + body: Mapping, beta_header: str, base_urls: Mapping[str, str] +) -> int: + """Count a Claude API Messages body exactly as it will be sent. + + The count includes `context_management`, so it reports the prompt after + server-side context editing: the size the long-prompt price applies to. + """ + from services.auth_service import AuthService + + credential = AuthService.get_api_key("anthropic") + if not credential: + raise NativeTokenCountError("No Claude API key is configured", 503) + url = _count_url(base_urls["anthropic"], CLAUDE_COUNT_PATH) + incoming = {"anthropic-beta": beta_header} if beta_header else {} + count_body = {key: value for key, value in body.items() if key in CLAUDE_COUNT_FIELDS} + return _post_count("anthropic", url, CLAUDE_COUNT_PATH, credential, incoming, count_body) + + +def _post_count( + provider: str, + url: str, + path: str, + credential: str, + incoming: Mapping[str, str], + body: dict, +) -> int: + from services.proxy_service import ProxyService + headers = ProxyService.prepare_headers(incoming, provider, credential, upstream_path=path) headers["Content-Type"] = "application/json" headers["Accept"] = "application/json" - body = {key: value for key, value in payload.items() if key in COUNT_FIELDS} - body["model"] = model session = ProxyService._get_provider_session(provider, raw_passthrough=True) response = None try: diff --git a/tests/outbound_inventory.json b/tests/outbound_inventory.json index 72d3ac81..f4a6957f 100644 --- a/tests/outbound_inventory.json +++ b/tests/outbound_inventory.json @@ -2,7 +2,7 @@ "app.py": { "kind": "non-provider", "reason": "Runtime module pinned independently of HTTP client spelling; network attempts require runtime admission.", - "sha256": "38480cbcaf363f62177f1449b73b12c6c1b379dbb938e0096d6ff3eebb65921e" + "sha256": "ed56be46ef0772dec11bdc238e25e28ea977fc10979102c3a655a7c87fb3808c" }, "cloudflare-worker.mjs": { "kind": "internal", @@ -402,7 +402,7 @@ "routes/unified.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "bcc073d8fc97c65c2b00b1d7b069284ddfcd1834ab610c640e0d1644a439415a" + "sha256": "38b545840293971d167268ebecd8095e638f0717fe5fc4782925602e5eaa47d6" }, "routes/unified_bridge.py": { "kind": "non-provider", @@ -997,7 +997,7 @@ "services/native_token_count.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "46cd0e4dff7d3d123fe710c1cc00dcdf39a57a56ce1008446c94c8f4975282ac" + "sha256": "8f8f6f026a4094be7f0a18a7876483f02f47430a8c006c43a9416dbbe8d404ed" }, "services/observability_adapters.py": { "kind": "non-provider", @@ -1522,7 +1522,7 @@ "worker/api-paths.mjs": { "kind": "non-provider", "reason": "Runtime module pinned independently of HTTP client spelling; network attempts require runtime admission.", - "sha256": "378a091ba1b4f13ce2e75516eaa425a01c1655624970aac5bfd59627e5cf4a54" + "sha256": "ebf33474e9839780cb05b3fe652dc26b174f37eff5123a7a120593f80cde2c7b" }, "worker/auto-routes-d1.mjs": { "kind": "non-provider", @@ -1572,7 +1572,7 @@ "worker/container-env.mjs": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "f5dc040ed498889a81747ddde2bf7ea7750af3fffa3396e75e82c92438326a9a" + "sha256": "47bf9837f223cb05e6a42facb7d37e4784cf7656631227314fefb32da8e42cdb" }, "worker/context-canary.mjs": { "kind": "non-provider", @@ -2458,5 +2458,15 @@ "kind": "non-provider", "reason": "Runtime module pinned independently of HTTP client spelling; network attempts require runtime admission.", "sha256": "f681974644eb5340debf6d4c0b072bb8dfd5ffcad18a652611d9bdf888d595c3" + }, + "routes/lean.py": { + "kind": "non-provider", + "reason": "Runtime module pinned independently of HTTP client spelling; network attempts require runtime admission.", + "sha256": "425514a91851a5a4a209b3abe50bff97085630cebafa606c92b43f11929f0b23" + }, + "services/lean_mode.py": { + "kind": "guarded", + "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", + "sha256": "d54cd894d2c937c6ffb32c64da9f7d6648ae4ea004d6193220800ee55194231c" } } diff --git a/tests/test_anthropic_worker.mjs b/tests/test_anthropic_worker.mjs index 55714dee..144439c8 100644 --- a/tests/test_anthropic_worker.mjs +++ b/tests/test_anthropic_worker.mjs @@ -34,6 +34,16 @@ test("Anthropic native routes use edge API handling", () => { assert.equal(isApiRequestPath("/anthropic-other/v1/messages"), false); }); +test("lean routes use edge API handling and the container receives lean settings", () => { + assert.equal(isApiRequestPath("/lean/v1/messages"), true); + assert.equal(isApiRequestPath("/lean/v1/chat/completions"), true); + assert.equal(isApiRequestPath("/leaner/v1/messages"), false); + const forwarded = collectContainerEnv({ LEAN_STYLE: "concise", LEAN_PROMPT_TOKEN_LIMIT: "90000", LEAN_UNKNOWN: "x" }); + assert.equal(forwarded.LEAN_STYLE, "concise"); + assert.equal(forwarded.LEAN_PROMPT_TOKEN_LIMIT, "90000"); + assert.equal(forwarded.LEAN_UNKNOWN, undefined); +}); + test("Anthropic preflight allows Claude headers without waking the container", async () => { const { default: worker } = await loadWorkerModule(); const origin = "https://chat.example"; diff --git a/tests/test_lean_mode.py b/tests/test_lean_mode.py new file mode 100644 index 00000000..ce0da01f --- /dev/null +++ b/tests/test_lean_mode.py @@ -0,0 +1,427 @@ +"""Lean mode: cache, clear, condense and budget Claude API requests without editing history.""" + +import json +import os +from unittest.mock import patch + +import pytest + +from services import lean_mode +from services.lean_mode import ( + CONCISE_STYLE, + CONTEXT_MANAGEMENT_BETA, + LeanPolicy, + clip_tool_text, + condense_shell_output, + estimate_prompt_tokens, + lean_policy, + merge_beta_header, +) +from services.native_token_count import NativeTokenCountError +from tests.test_anthropic_provider import MESSAGE, MODEL, SELECTOR, upstream +from tests.test_protocol_routes import ADMIN, ANTHROPIC_STREAM, sse_upstream +from tests.unified_api_test_case import UnifiedApiTestCase + +SHELL_OUTPUT = ( + "\x1b[32mCompiling\x1b[0m crate\n" + "downloading 10%\rdownloading 55%\rdownloading 100%\n" + + "warning: unused import\n" * 6 + + "\n\n\n\n" + + "\x1b]0;title\x07error: failed\n" +) +TOOLS = [{"name": "Bash", "description": "Run a command", "input_schema": {"type": "object"}}] + + +def tool_turn(call_id, output, name="Bash"): + return [ + { + "role": "assistant", + "content": [{"type": "tool_use", "id": call_id, "name": name, "input": {"command": "cargo build"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": call_id, "content": output}]}, + ] + + +def conversation(turns): + messages = [{"role": "user", "content": "Build the project"}] + for index in range(turns): + messages.extend(tool_turn(f"call_{index}", SHELL_OUTPUT)) + return {"model": MODEL, "max_tokens": 4096, "system": "You are a coding agent.", "tools": TOOLS, "messages": messages} + + +class LeanContext: + """Run prepare_lean_request inside a lean request context.""" + + def __init__(self, app): + self.app = app + + def prepare(self, body, headers=None): + with self.app.test_request_context("/v1/messages", headers=headers or {}): + from flask import request + + request.environ[lean_mode.LEAN_ENVIRON_KEY] = True + return lean_mode.prepare_lean_request( + body, + provider="anthropic", + endpoint="v1/messages", + request_headers=headers or {}, + base_urls={"anthropic": "https://api.anthropic.com"}, + ) + + +@pytest.fixture +def lean(): + from flask import Flask + + return LeanContext(Flask(__name__)) + + +def test_shell_output_is_condensed_deterministically_and_idempotently(): + condensed = condense_shell_output(SHELL_OUTPUT) + assert condensed == ( + "Compiling crate\n" + "downloading 100%\n" + "warning: unused import\n" + "[previous line repeated 5 more times]\n" + "\n" + "error: failed\n" + ) + assert condense_shell_output(condensed) == condensed + assert condense_shell_output("ok\n\nok\n") == "ok\n\nok\n" + + +def test_oversized_tool_text_keeps_head_and_tail(): + text = "h" * 60 + "m" * 100 + "t" * 40 + clipped = clip_tool_text(text, 100) + assert clipped.startswith("h" * 50) + assert clipped.endswith("m" * 10 + "t" * 40) + assert "[MultiLLM lean: 100 characters omitted from the middle of this tool result]" in clipped + assert clip_tool_text(text, 0) == text + assert clip_tool_text(text, 500) == text + + +def test_policy_reads_environment_with_safe_defaults(): + assert lean_policy({}) == LeanPolicy() + policy = lean_policy( + { + "LEAN_PROMPT_TOKEN_LIMIT": "90000", + "LEAN_CLEAR_TRIGGER_TOKENS": "0", + "LEAN_CLEAR_EXCLUDE_TOOLS": "Read, Grep", + "LEAN_CACHE_TTL": "1h", + "LEAN_CONDENSE_SHELL_OUTPUT": "false", + "LEAN_TOOL_RESULT_MAX_CHARS": "bad", + "LEAN_STYLE": "caveman", + } + ) + assert policy.prompt_token_limit == 90_000 + assert policy.clear_trigger_tokens == 0 + assert policy.clear_exclude_tools == ("Read", "Grep") + assert policy.cache_ttl == "1h" + assert policy.condense_shell_output is False + assert policy.tool_result_max_chars == lean_mode.DEFAULT_TOOL_RESULT_MAX_CHARS + assert policy.style == "concise" + + +def test_prepared_body_caches_clears_and_condenses_only_shell_results(lean): + body = conversation(1) + body["messages"].extend(tool_turn("read_1", "\x1b[1mexact file bytes\x1b[0m \n", name="Read")) + decision = lean.prepare(body) + prepared = decision.body + assert decision.active + assert decision.betas == (CONTEXT_MANAGEMENT_BETA,) + assert decision.stages == ("cache", "clear-tool-uses", "condense-shell") + assert prepared["cache_control"] == {"type": "ephemeral"} + assert prepared["system"] == [ + {"type": "text", "text": "You are a coding agent.", "cache_control": {"type": "ephemeral"}} + ] + assert prepared["context_management"] == { + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 80_000}, + "keep": {"type": "tool_uses", "value": 3}, + "clear_at_least": {"type": "input_tokens", "value": 20_000}, + } + ] + } + shell_result = prepared["messages"][2]["content"][0]["content"] + assert shell_result == condense_shell_output(SHELL_OUTPUT) + # A file read keeps its exact bytes, escapes and trailing spaces included. + assert prepared["messages"][4]["content"][0]["content"] == "\x1b[1mexact file bytes\x1b[0m \n" + assert decision.condensed_chars > 0 + assert body["messages"][2]["content"][0]["content"] == SHELL_OUTPUT + + +def test_every_later_turn_repeats_the_earlier_prepared_bytes(lean): + """Cache hits and preserved thinking both need an unchanged prefix.""" + earlier = lean.prepare(conversation(2)).body + later = lean.prepare(conversation(3)).body + + def dumped(value): + return json.dumps(value, sort_keys=False, separators=(",", ":")) + + assert dumped(later["system"]) == dumped(earlier["system"]) + assert dumped(later["tools"]) == dumped(earlier["tools"]) + assert dumped(later["messages"][: len(earlier["messages"])]) == dumped(earlier["messages"]) + + +def test_caller_cache_and_context_management_are_left_alone(lean): + body = conversation(1) + body["system"] = [{"type": "text", "text": "Own cache", "cache_control": {"type": "ephemeral"}}] + body["context_management"] = {"edits": [{"type": "clear_thinking_20251015"}]} + decision = lean.prepare(body) + assert "cache_control" not in decision.body + assert decision.body["system"] == body["system"] + assert decision.body["context_management"] == body["context_management"] + assert decision.betas == () + + +def test_signed_compaction_blocks_are_not_combined_with_context_editing(lean): + body = conversation(1) + body["messages"].insert(0, {"role": "assistant", "content": [{"type": "compaction", "content": "summary"}]}) + assert lean.prepare(body).betas == () + + +def test_plain_chat_without_tools_only_gets_caching(lean): + decision = lean.prepare({"model": MODEL, "max_tokens": 64, "messages": [{"role": "user", "content": "Hi"}]}) + assert decision.stages == ("cache",) + assert "context_management" not in decision.body + assert decision.body["cache_control"] == {"type": "ephemeral"} + + +def test_concise_style_is_opt_in_and_cached_with_the_system_prompt(lean): + assert "concise-style" not in lean.prepare(conversation(1)).stages + decision = lean.prepare(conversation(1), headers={"X-MultiLLM-Lean-Style": "caveman"}) + assert "concise-style" in decision.stages + assert decision.body["system"] == [ + {"type": "text", "text": "You are a coding agent."}, + {"type": "text", "text": CONCISE_STYLE, "cache_control": {"type": "ephemeral"}}, + ] + + +def test_tools_only_request_caches_its_last_tool(lean): + body = {"model": MODEL, "max_tokens": 64, "tools": TOOLS, "messages": [{"role": "user", "content": "Hi"}]} + with patch.dict(os.environ, {"LEAN_CACHE_TTL": "1h"}): + prepared = lean.prepare(body).body + marker = {"type": "ephemeral", "ttl": "1h"} + assert prepared["cache_control"] == marker + assert prepared["tools"][-1]["cache_control"] == marker + + +def test_estimate_counts_media_at_a_flat_size(): + image = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "A" * 900_000}} + body = {"messages": [{"role": "user", "content": [image, {"type": "text", "text": "abc" * 10}]}]} + # Text, roles and block types are counted; the 900 KB of base64 is not. + assert estimate_prompt_tokens(body) == len("user" "text" + "abc" * 10) // 3 + lean_mode.MEDIA_BLOCK_TOKENS + + +def test_beta_header_merge_keeps_caller_betas_once(): + headers = {"anthropic-beta": "claude-code-20250219, context-management-2025-06-27"} + merge_beta_header(headers, (CONTEXT_MANAGEMENT_BETA,)) + assert headers == {"anthropic-beta": "claude-code-20250219,context-management-2025-06-27"} + empty = {} + merge_beta_header(empty, (CONTEXT_MANAGEMENT_BETA,)) + assert empty == {"Anthropic-Beta": CONTEXT_MANAGEMENT_BETA} + + +def test_near_limit_prompt_is_counted_and_refused_above_the_price_limit(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=100_001) as count: + with pytest.raises(lean_mode.LeanPromptTooLong) as raised: + lean.prepare(body, headers={"anthropic-beta": "claude-code-20250219"}) + assert str(raised.value) == "prompt is too long: 100001 tokens > 100000 maximum" + counted_body, beta_header, _ = count.call_args.args + assert counted_body["context_management"]["edits"][0]["type"] == "clear_tool_uses_20250919" + assert beta_header == "claude-code-20250219,context-management-2025-06-27" + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=99_000): + decision = lean.prepare(body) + assert (decision.prompt_tokens, decision.prompt_count) == (99_000, "provider") + + +def test_failed_count_lets_the_request_through(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + with patch.object(lean_mode, "count_claude_prompt_tokens", side_effect=NativeTokenCountError("down", 429)): + decision = lean.prepare(body) + assert decision.prompt_count == "estimate" + + +def test_claude_count_sends_context_management_with_the_lean_beta(): + from services import native_token_count + + captured = {} + + class CountResponse: + status_code = 200 + + def iter_content(self, chunk_size): + yield b'{"input_tokens": 1234}' + + def close(self): + pass + + def provider_http(provider, method, url, **kwargs): + captured.update(provider=provider, url=url, **kwargs) + return CountResponse() + + body = {**conversation(1), "context_management": {"edits": []}, "cache_control": {"type": "ephemeral"}, "stream": True} + with patch("services.auth_service.AuthService.get_api_key", return_value="claude-key"), patch( + "services.provider_dispatch.provider_http", provider_http + ): + count = native_token_count.count_claude_prompt_tokens( + body, CONTEXT_MANAGEMENT_BETA, {"anthropic": "https://api.anthropic.com"} + ) + assert count == 1234 + assert captured["url"] == "https://api.anthropic.com/v1/messages/count_tokens" + assert captured["headers"]["X-Api-Key"] == "claude-key" + assert captured["headers"]["Anthropic-Beta"] == CONTEXT_MANAGEMENT_BETA + assert set(captured["json"]) == {"model", "messages", "system", "tools", "context_management"} + + +class LeanRouteTest(UnifiedApiTestCase): + def setUp(self): + super().setUp() + os.environ["ANTHROPIC_API_KEY"] = "anthropic-provider-key" + + def post(self, path, body, response=None, headers=None): + with patch( + "app.ProxyService.make_request", + return_value=response if response is not None else upstream(MESSAGE), + ) as send: + result = self.client.post(path, json=body, headers=headers or ADMIN) + result.get_data() + return result, send + + def test_lean_messages_send_a_cached_cleared_body(self): + body = {**conversation(1), "model": SELECTOR} + result, send = self.post( + "/lean/v1/messages", + body, + headers={"X-Api-Key": "admin-test-key", "Anthropic-Beta": "test-beta"}, + ) + self.assertEqual(result.status_code, 200) + kwargs = send.call_args.kwargs + self.assertEqual(kwargs["url"], "https://api.anthropic.com/v1/messages") + self.assertEqual(kwargs["headers"]["Anthropic-Beta"], f"test-beta,{CONTEXT_MANAGEMENT_BETA}") + sent = json.loads(kwargs["data"]) + self.assertEqual(sent["cache_control"], {"type": "ephemeral"}) + self.assertIn("context_management", sent) + self.assertEqual(sent["messages"][2]["content"][0]["content"], condense_shell_output(SHELL_OUTPUT)) + self.assertEqual(result.headers["X-MultiLLM-Lean"], "applied") + self.assertEqual(result.headers["X-MultiLLM-Lean-Stages"], "cache,clear-tool-uses,condense-shell") + self.assertEqual(result.headers["X-MultiLLM-Lean-Prompt-Count"], "estimate") + self.assertEqual(result.get_json()["content"], MESSAGE["content"]) + + def test_plain_v1_messages_are_unchanged(self): + body = {**conversation(1), "model": SELECTOR} + result, send = self.post("/v1/messages", body, headers={"X-Api-Key": "admin-test-key"}) + self.assertEqual(result.status_code, 200) + sent = json.loads(send.call_args.kwargs["data"]) + self.assertNotIn("cache_control", sent) + self.assertNotIn("context_management", sent) + self.assertEqual(sent["messages"][2]["content"][0]["content"], SHELL_OUTPUT) + self.assertNotIn("X-MultiLLM-Lean", result.headers) + + def test_lean_chat_translates_then_condenses_shell_tool_output(self): + result, send = self.post( + "/lean/v1/chat/completions", + { + "model": SELECTOR, + "tools": [{"type": "function", "function": {"name": "bash", "parameters": {"type": "object"}}}], + "messages": [ + {"role": "system", "content": "Be useful"}, + {"role": "user", "content": "Build"}, + { + "role": "assistant", + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "bash", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": SHELL_OUTPUT}, + ], + }, + ) + self.assertEqual(result.status_code, 200) + sent = json.loads(send.call_args.kwargs["data"]) + self.assertEqual(sent["system"][0]["cache_control"], {"type": "ephemeral"}) + self.assertEqual(sent["messages"][-1]["content"][0]["content"], condense_shell_output(SHELL_OUTPUT)) + self.assertEqual(result.headers["X-MultiLLM-Lean"], "applied") + + def test_lean_responses_stream_is_cached(self): + result, send = self.post( + "/lean/v1/responses", + {"model": SELECTOR, "input": "Hello", "stream": True}, + response=sse_upstream(ANTHROPIC_STREAM), + ) + self.assertEqual(result.status_code, 200) + sent = json.loads(send.call_args.kwargs["data"]) + self.assertEqual(sent["cache_control"], {"type": "ephemeral"}) + self.assertEqual(result.headers["X-MultiLLM-Lean-Stages"], "cache") + + def oversized(self): + return patch.object(lean_mode, "count_claude_prompt_tokens", return_value=150_000) + + def large_input(self): + return "x" * 200_000 + + def test_over_limit_messages_get_anthropic_prompt_too_long(self): + with self.oversized(): + result, send = self.post( + "/lean/v1/messages", + {"model": SELECTOR, "max_tokens": 64, "messages": [{"role": "user", "content": self.large_input()}]}, + headers={"X-Api-Key": "admin-test-key"}, + ) + send.assert_not_called() + self.assertEqual(result.status_code, 400) + error = result.get_json() + self.assertEqual(error["type"], "error") + self.assertEqual(error["error"]["type"], "invalid_request_error") + self.assertEqual(error["error"]["message"], "prompt is too long: 150000 tokens > 100000 maximum") + self.assertEqual(result.headers["X-MultiLLM-Lean-Prompt-Tokens"], "150000") + + def test_over_limit_chat_gets_context_length_exceeded(self): + with self.oversized(): + result, send = self.post( + "/lean/v1/chat/completions", + {"model": SELECTOR, "messages": [{"role": "user", "content": self.large_input()}]}, + ) + send.assert_not_called() + self.assertEqual(result.status_code, 400) + error = result.get_json()["error"] + self.assertEqual(error["code"], "context_length_exceeded") + self.assertIn("maximum context length is 100000 tokens", error["message"]) + self.assertIn("prompt is too long", error["message"]) + + def test_over_limit_responses_stream_fails_like_openai(self): + with self.oversized(): + result, send = self.post( + "/lean/v1/responses", + {"model": SELECTOR, "input": self.large_input(), "stream": True}, + ) + send.assert_not_called() + self.assertEqual(result.status_code, 200) + self.assertEqual(result.mimetype, "text/event-stream") + frame = result.get_data(as_text=True) + self.assertTrue(frame.startswith("event: response.failed\ndata: ")) + event = json.loads(frame.split("data: ", 1)[1]) + self.assertEqual(event["response"]["error"]["code"], "context_length_exceeded") + + def test_unknown_lean_path_is_not_routed(self): + result = self.client.post("/lean/v1/embeddings", json={}, headers=ADMIN) + # The provider catch-all answers, as it does for any unknown prefix. + self.assertEqual(result.status_code, 400) + self.assertIn("Unsupported API provider: lean", result.get_data(as_text=True)) + + +class LeanDisabledRouteTest(UnifiedApiTestCase): + def setUp(self): + super().setUp() + os.environ["LEAN_ENDPOINT_ENABLED"] = "false" + self.app = self.app_module.create_app() + self.client = self.app.test_client() + + def test_disabled_lean_prefix_is_not_routed(self): + result = self.client.post("/lean/v1/messages", json={"model": SELECTOR}, headers=ADMIN) + self.assertEqual(result.status_code, 400) + self.assertIn("Unsupported API provider: lean", result.get_data(as_text=True)) diff --git a/worker/api-paths.mjs b/worker/api-paths.mjs index 6f796595..3c1a3309 100644 --- a/worker/api-paths.mjs +++ b/worker/api-paths.mjs @@ -15,6 +15,7 @@ const API_ROUTE_PREFIXES = new Set([ "hyperbolic", "intelligence", "kimi-code", + "lean", "linkapi", "mimo", "nanogpt", diff --git a/worker/container-env.mjs b/worker/container-env.mjs index e946a3aa..5a646990 100644 --- a/worker/container-env.mjs +++ b/worker/container-env.mjs @@ -129,6 +129,16 @@ const DIRECT_ENV_KEYS = [ "MAX_OUTPUT_TOKENS", "OPTIMIZER_MAX_REQUEST_BYTES", "OPTIMIZER_SUMMARY_TIMEOUT_SECONDS", + "LEAN_ENDPOINT_ENABLED", + "LEAN_PROMPT_TOKEN_LIMIT", + "LEAN_CLEAR_TRIGGER_TOKENS", + "LEAN_CLEAR_KEEP_TOOL_USES", + "LEAN_CLEAR_AT_LEAST_TOKENS", + "LEAN_CLEAR_EXCLUDE_TOOLS", + "LEAN_CACHE_TTL", + "LEAN_CONDENSE_SHELL_OUTPUT", + "LEAN_TOOL_RESULT_MAX_CHARS", + "LEAN_STYLE", "GLM_AUTO_OPTIMIZE", "GLM_AUTO_OPTIMIZE_TRIGGER_TOKENS", "GLM_AUTO_OPTIMIZE_KEEP_RECENT_TURNS", From 728a4e00994c56f374d593ff618dd653a68e1fd1 Mon Sep 17 00:00:00 2001 From: Abhyudaya Baiju Tharakan Date: Mon, 26 Oct 2026 21:34:36 +0530 Subject: [PATCH 2/5] fix(lean): apply the price limit to the prompt before clearing by default Anthropic's pricing page does not say whether a prompt shortened by server-side context editing is priced by its size before or after clearing. count_tokens reports both, so the guard now refuses on the larger size unless LEAN_PRICE_LIMIT_BASIS=cleared, and reports the size after clearing in its own header. The docs note that a harness stops after repeated failed compactions when one message alone is too large. --- .env.example | 1 + docs/lean.md | 8 ++++++-- services/lean_mode.py | 30 ++++++++++++++++++++++++------ services/native_token_count.py | 20 ++++++++++++-------- tests/outbound_inventory.json | 6 +++--- tests/test_lean_mode.py | 23 ++++++++++++++++++----- worker/container-env.mjs | 1 + 7 files changed, 65 insertions(+), 24 deletions(-) diff --git a/.env.example b/.env.example index ba9694c6..71dc6dbd 100644 --- a/.env.example +++ b/.env.example @@ -88,6 +88,7 @@ NAVYAI_MAX_REQUEST_BYTES=33554432 # Lean endpoint (/lean/v1/...) for Claude API models; see docs/lean.md. # LEAN_ENDPOINT_ENABLED=true # LEAN_PROMPT_TOKEN_LIMIT=100000 +# LEAN_PRICE_LIMIT_BASIS=original # LEAN_CLEAR_TRIGGER_TOKENS=80000 # LEAN_CLEAR_KEEP_TOOL_USES=3 # LEAN_CLEAR_AT_LEAST_TOKENS=20000 diff --git a/docs/lean.md b/docs/lean.md index 828bbb73..95e2fd89 100644 --- a/docs/lean.md +++ b/docs/lean.md @@ -38,7 +38,7 @@ with the raw shell output. | Tool result clearing | on, from 80,000 tokens | Adds `context_management` with `clear_tool_uses_20250919`: once the prompt passes the trigger, the API replaces old tool results with placeholders and keeps the 3 most recent, clearing at least 20,000 tokens at a time. The API applies it before the model reads the prompt, the harness keeps its full history, and it never invalidates thinking blocks. Skipped when the caller sends its own `context_management` or a compaction block. | | Shell output | on | Results of shell tools (`Bash`, `bash`, `shell`, `exec_command`, `run_shell_command` and similar) lose ANSI escapes and progress-bar redraws; runs of four or more identical lines become one line and a count, and runs of blank lines become one. Other tools' results, file reads included, keep their exact bytes, so an edit tool can still match them. | | Oversized results | 100,000 characters | A longer tool result keeps its first and last 50,000 characters with a marker in between. | -| Prompt limit | 100,000 tokens | A prompt estimated near the limit is counted exactly with `count_tokens`, context editing included. Above the limit the request is refused with the error the harness treats as a full context window, so it compacts and retries instead of paying the higher price. | +| Prompt limit | 100,000 tokens | A prompt estimated near the limit is counted exactly with `count_tokens`, which reports its size before and after context editing. Above the limit the request is refused with the error the harness treats as a full context window, so it compacts and retries instead of paying the higher price. The pricing page does not say whether the size before or after clearing picks the price, so by default the limit applies to the size before clearing; set `LEAN_PRICE_LIMIT_BASIS=cleared` once your bill shows that cleared prompts are priced by their smaller size. | | Concise answers | off | Adds a short instruction to the system prompt: lead with the answer, no filler, code and commands exact. Turn it on with `LEAN_STYLE=concise` or per request with `X-MultiLLM-Lean-Style: concise` (`caveman` is accepted too). | Lean mode never changes the model, effort, thinking settings or `max_tokens`. @@ -52,7 +52,9 @@ Lean mode never changes the model, effort, thinking settings or `max_tokens`. | `/lean/v1/responses` with `stream: true` | A `response.failed` event with `context_length_exceeded`, the form Codex reads. Without streaming, the Chat Completions body above. | A count that fails (for example, rate-limited) lets the request through and -reports `X-MultiLLM-Lean-Prompt-Count: estimate`. +reports `X-MultiLLM-Lean-Prompt-Count: estimate`. If compacting cannot bring a +prompt under the limit, for example because the newest message alone is larger, +the harness reports the error: Claude Code stops after three failed compactions. ### Response headers @@ -63,6 +65,7 @@ reports `X-MultiLLM-Lean-Prompt-Count: estimate`. | `X-MultiLLM-Lean-Stages` | Stages that changed the request, for example `cache,clear-tool-uses,condense-shell` | | `X-MultiLLM-Lean-Prompt-Tokens` | Prompt size: counted by the API near the limit, otherwise a high local estimate | | `X-MultiLLM-Lean-Prompt-Count` | `provider` or `estimate` | +| `X-MultiLLM-Lean-Prompt-Tokens-After-Clearing` | Counted size after context editing, when it differs | | `X-MultiLLM-Lean-Condensed-Chars` | Characters removed from tool results | The API reports what it cleared in the response's `context_management` field and @@ -130,6 +133,7 @@ caching is what makes cache reads happen at all. |---|---|---| | `LEAN_ENDPOINT_ENABLED` | `true` | `false` stops serving `/lean` | | `LEAN_PROMPT_TOKEN_LIMIT` | `100000` | `0` turns the refusal off | +| `LEAN_PRICE_LIMIT_BASIS` | `original` | `cleared` applies the limit to the prompt after context editing | | `LEAN_CLEAR_TRIGGER_TOKENS` | `80000` | `0` turns tool result clearing off | | `LEAN_CLEAR_KEEP_TOOL_USES` | `3` | Recent tool uses that are never cleared | | `LEAN_CLEAR_AT_LEAST_TOKENS` | `20000` | Minimum cleared per edit, so a cache rewrite is worth it | diff --git a/services/lean_mode.py b/services/lean_mode.py index 760aeb1a..199b7eb7 100644 --- a/services/lean_mode.py +++ b/services/lean_mode.py @@ -99,6 +99,10 @@ class LeanPolicy: condense_shell_output: bool = True tool_result_max_chars: int = DEFAULT_TOOL_RESULT_MAX_CHARS style: str = "none" + # Whether the price limit applies to the prompt before ("original") or after + # ("cleared") context editing. Anthropic's pricing page does not say which + # size picks the price, so the default assumes the larger one does. + price_limit_basis: str = "original" @dataclass(frozen=True) @@ -109,6 +113,7 @@ class LeanRequest: stages: tuple[str, ...] = () prompt_tokens: int | None = None prompt_count: str = "" + prompt_tokens_after_clearing: int | None = None condensed_chars: int = 0 @@ -149,6 +154,11 @@ def lean_policy(environ: Mapping[str, str] | None = None) -> LeanPolicy: condense_shell_output=_env_flag(environ, "LEAN_CONDENSE_SHELL_OUTPUT", True), tool_result_max_chars=_env_int(environ, "LEAN_TOOL_RESULT_MAX_CHARS", DEFAULT_TOOL_RESULT_MAX_CHARS), style=_style(environ.get("LEAN_STYLE")) or "none", + price_limit_basis=( + "cleared" + if str(environ.get("LEAN_PRICE_LIMIT_BASIS", "")).strip().lower() == "cleared" + else "original" + ), ) @@ -417,18 +427,22 @@ def _prompt_size( betas: tuple[str, ...], request_headers: Mapping[str, Any] | None, base_urls: Mapping[str, str], - limit: int, -) -> tuple[int, str]: + policy: LeanPolicy, +) -> tuple[int, str, int | None]: + """Return the size the price limit applies to, how it was measured, and the + counted size after context editing.""" estimate = estimate_prompt_tokens(body) + limit = policy.prompt_token_limit if limit <= 0 or estimate < limit * COUNT_FROM_LIMIT_SHARE: - return estimate, "estimate" + return estimate, "estimate", None beta_header = merged_beta_value(_header(request_headers, "anthropic-beta"), betas) try: - return count_claude_prompt_tokens(body, beta_header, base_urls), "provider" + after, before = count_claude_prompt_tokens(body, beta_header, base_urls) except NativeTokenCountError as error: # Counting is a guard, not a gate: a failed count lets the request through. logger.warning("Lean prompt count failed status=%s; sending uncounted", error.status) - return estimate, "estimate" + return estimate, "estimate", None + return (before if policy.price_limit_basis == "original" else after), "provider", after def prepare_lean_request( @@ -464,7 +478,7 @@ def prepare_lean_request( betas = (CONTEXT_MANAGEMENT_BETA,) if clearing else () if clearing: stages.add("clear-tool-uses") - tokens, source = _prompt_size(prepared, betas, request_headers, base_urls, policy.prompt_token_limit) + tokens, source, after_clearing = _prompt_size(prepared, betas, request_headers, base_urls, policy) decision = LeanRequest( body=prepared, active=True, @@ -472,6 +486,7 @@ def prepare_lean_request( stages=tuple(sorted(stages)), prompt_tokens=tokens, prompt_count=source, + prompt_tokens_after_clearing=after_clearing, condensed_chars=max(0, condensed_chars), ) g.lean_request = decision @@ -492,6 +507,9 @@ def lean_response_headers(response): response.headers["X-MultiLLM-Lean-Stages"] = ",".join(decision.stages) or "none" response.headers["X-MultiLLM-Lean-Prompt-Tokens"] = str(decision.prompt_tokens) response.headers["X-MultiLLM-Lean-Prompt-Count"] = decision.prompt_count + after = decision.prompt_tokens_after_clearing + if after is not None and after != decision.prompt_tokens: + response.headers["X-MultiLLM-Lean-Prompt-Tokens-After-Clearing"] = str(after) if decision.condensed_chars: response.headers["X-MultiLLM-Lean-Condensed-Chars"] = str(decision.condensed_chars) return response diff --git a/services/native_token_count.py b/services/native_token_count.py index 51d351b0..e10d7098 100644 --- a/services/native_token_count.py +++ b/services/native_token_count.py @@ -75,7 +75,7 @@ def _reject_nonfinite(value: str) -> None: raise ValueError("Nonfinite JSON value") -def _read_count(response: requests.Response) -> int: +def _read_count(response: requests.Response) -> dict: raw = bytearray() for chunk in response.iter_content(chunk_size=8192): if len(raw) + len(chunk) > MAX_RESPONSE_BYTES: @@ -88,7 +88,7 @@ def _read_count(response: requests.Response) -> int: count = body.get("input_tokens") if isinstance(body, dict) else None if type(count) is not int or count < 0: raise NativeTokenCountError("Malformed native token count response") - return count + return body def count_native_tokens( @@ -125,16 +125,16 @@ def count_native_tokens( incoming = {key: value for key, value in request_headers.items() if key.lower() in COUNT_HEADERS} body = {key: value for key, value in payload.items() if key in COUNT_FIELDS} body["model"] = model - return _post_count(provider, url, path, credential, incoming, body) + return _post_count(provider, url, path, credential, incoming, body)["input_tokens"] def count_claude_prompt_tokens( body: Mapping, beta_header: str, base_urls: Mapping[str, str] -) -> int: +) -> tuple[int, int]: """Count a Claude API Messages body exactly as it will be sent. - The count includes `context_management`, so it reports the prompt after - server-side context editing: the size the long-prompt price applies to. + Returns the prompt size after and before server-side context editing; they + are equal when the body clears nothing. """ from services.auth_service import AuthService @@ -144,7 +144,11 @@ def count_claude_prompt_tokens( url = _count_url(base_urls["anthropic"], CLAUDE_COUNT_PATH) incoming = {"anthropic-beta": beta_header} if beta_header else {} count_body = {key: value for key, value in body.items() if key in CLAUDE_COUNT_FIELDS} - return _post_count("anthropic", url, CLAUDE_COUNT_PATH, credential, incoming, count_body) + counted = _post_count("anthropic", url, CLAUDE_COUNT_PATH, credential, incoming, count_body) + management = counted.get("context_management") + original = management.get("original_input_tokens") if isinstance(management, dict) else None + after = counted["input_tokens"] + return after, original if type(original) is int and original >= after else after def _post_count( @@ -154,7 +158,7 @@ def _post_count( credential: str, incoming: Mapping[str, str], body: dict, -) -> int: +) -> dict: from services.proxy_service import ProxyService headers = ProxyService.prepare_headers(incoming, provider, credential, upstream_path=path) diff --git a/tests/outbound_inventory.json b/tests/outbound_inventory.json index f4a6957f..0f9e0571 100644 --- a/tests/outbound_inventory.json +++ b/tests/outbound_inventory.json @@ -997,7 +997,7 @@ "services/native_token_count.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "8f8f6f026a4094be7f0a18a7876483f02f47430a8c006c43a9416dbbe8d404ed" + "sha256": "73c3de349f04b1524c63b8d19849306213d2e2c6dd9b5c459ebc652ce7a4c402" }, "services/observability_adapters.py": { "kind": "non-provider", @@ -1572,7 +1572,7 @@ "worker/container-env.mjs": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "47bf9837f223cb05e6a42facb7d37e4784cf7656631227314fefb32da8e42cdb" + "sha256": "44b80ae15dc92596588b2cde34c956835b9d696df28fcea3eb48479121e91d03" }, "worker/context-canary.mjs": { "kind": "non-provider", @@ -2467,6 +2467,6 @@ "services/lean_mode.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "d54cd894d2c937c6ffb32c64da9f7d6648ae4ea004d6193220800ee55194231c" + "sha256": "4840552159fca30603f1da07cfa008abe95ad82c1243a6302e2b5f73ca95195b" } } diff --git a/tests/test_lean_mode.py b/tests/test_lean_mode.py index ce0da01f..8626e9f6 100644 --- a/tests/test_lean_mode.py +++ b/tests/test_lean_mode.py @@ -227,18 +227,31 @@ def test_beta_header_merge_keeps_caller_betas_once(): def test_near_limit_prompt_is_counted_and_refused_above_the_price_limit(lean): body = conversation(1) body["messages"][0]["content"] = "x" * 200_000 - with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=100_001) as count: + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(100_001, 100_001)) as count: with pytest.raises(lean_mode.LeanPromptTooLong) as raised: lean.prepare(body, headers={"anthropic-beta": "claude-code-20250219"}) assert str(raised.value) == "prompt is too long: 100001 tokens > 100000 maximum" counted_body, beta_header, _ = count.call_args.args assert counted_body["context_management"]["edits"][0]["type"] == "clear_tool_uses_20250919" assert beta_header == "claude-code-20250219,context-management-2025-06-27" - with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=99_000): + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(99_000, 99_000)): decision = lean.prepare(body) assert (decision.prompt_tokens, decision.prompt_count) == (99_000, "provider") +def test_price_limit_applies_before_clearing_unless_configured(lean): + """Which size picks Haiku's price is undocumented, so the larger one is the default.""" + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(60_000, 120_000)): + with pytest.raises(lean_mode.LeanPromptTooLong) as raised: + lean.prepare(body) + assert raised.value.tokens == 120_000 + with patch.dict(os.environ, {"LEAN_PRICE_LIMIT_BASIS": "cleared"}): + decision = lean.prepare(body) + assert (decision.prompt_tokens, decision.prompt_tokens_after_clearing) == (60_000, 60_000) + + def test_failed_count_lets_the_request_through(lean): body = conversation(1) body["messages"][0]["content"] = "x" * 200_000 @@ -256,7 +269,7 @@ class CountResponse: status_code = 200 def iter_content(self, chunk_size): - yield b'{"input_tokens": 1234}' + yield b'{"input_tokens": 1234, "context_management": {"original_input_tokens": 5678}}' def close(self): pass @@ -272,7 +285,7 @@ def provider_http(provider, method, url, **kwargs): count = native_token_count.count_claude_prompt_tokens( body, CONTEXT_MANAGEMENT_BETA, {"anthropic": "https://api.anthropic.com"} ) - assert count == 1234 + assert count == (1234, 5678) assert captured["url"] == "https://api.anthropic.com/v1/messages/count_tokens" assert captured["headers"]["X-Api-Key"] == "claude-key" assert captured["headers"]["Anthropic-Beta"] == CONTEXT_MANAGEMENT_BETA @@ -360,7 +373,7 @@ def test_lean_responses_stream_is_cached(self): self.assertEqual(result.headers["X-MultiLLM-Lean-Stages"], "cache") def oversized(self): - return patch.object(lean_mode, "count_claude_prompt_tokens", return_value=150_000) + return patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(150_000, 150_000)) def large_input(self): return "x" * 200_000 diff --git a/worker/container-env.mjs b/worker/container-env.mjs index 5a646990..bda36bca 100644 --- a/worker/container-env.mjs +++ b/worker/container-env.mjs @@ -131,6 +131,7 @@ const DIRECT_ENV_KEYS = [ "OPTIMIZER_SUMMARY_TIMEOUT_SECONDS", "LEAN_ENDPOINT_ENABLED", "LEAN_PROMPT_TOKEN_LIMIT", + "LEAN_PRICE_LIMIT_BASIS", "LEAN_CLEAR_TRIGGER_TOKENS", "LEAN_CLEAR_KEEP_TOOL_USES", "LEAN_CLEAR_AT_LEAST_TOKENS", From 584d54008c2def3180e586a83a734b3e978b6475 Mon Sep 17 00:00:00 2001 From: Abhyudaya Baiju Tharakan Date: Mon, 26 Oct 2026 21:48:49 +0530 Subject: [PATCH 3/5] fix(lean): close the price guard's gaps for documents, uncountable prompts and pooled keys - Text documents and base64 PDFs are estimated by their contents instead of a flat size, so a large document is always counted. - URL and file sources, which count_tokens rejects, are judged by the estimate, and so is any prompt whose count fails: an estimate above the limit is refused rather than sent. - Server tools are left out of the count and added back at a fixed size, so requests with web search can still be counted. - The count uses the key that will send the request. Thinking blocks from another account are dropped before counting, so any other key could undercount. - Only block-level cache_control counts as a caller breakpoint; a tool schema property of that name no longer turns caching off. - The default limit is 99,000 tokens, since count_tokens is an estimate. --- .env.example | 2 +- docs/lean.md | 17 ++-- routes/unified.py | 11 +-- services/lean_mode.py | 161 +++++++++++++++++++++++---------- services/native_token_count.py | 7 +- tests/outbound_inventory.json | 6 +- tests/test_lean_mode.py | 79 +++++++++++++--- 7 files changed, 205 insertions(+), 78 deletions(-) diff --git a/.env.example b/.env.example index 71dc6dbd..3b3ab971 100644 --- a/.env.example +++ b/.env.example @@ -87,7 +87,7 @@ NAVYAI_MAX_REQUEST_BYTES=33554432 # Lean endpoint (/lean/v1/...) for Claude API models; see docs/lean.md. # LEAN_ENDPOINT_ENABLED=true -# LEAN_PROMPT_TOKEN_LIMIT=100000 +# LEAN_PROMPT_TOKEN_LIMIT=99000 # LEAN_PRICE_LIMIT_BASIS=original # LEAN_CLEAR_TRIGGER_TOKENS=80000 # LEAN_CLEAR_KEEP_TOOL_USES=3 diff --git a/docs/lean.md b/docs/lean.md index 95e2fd89..08bd4f62 100644 --- a/docs/lean.md +++ b/docs/lean.md @@ -38,7 +38,7 @@ with the raw shell output. | Tool result clearing | on, from 80,000 tokens | Adds `context_management` with `clear_tool_uses_20250919`: once the prompt passes the trigger, the API replaces old tool results with placeholders and keeps the 3 most recent, clearing at least 20,000 tokens at a time. The API applies it before the model reads the prompt, the harness keeps its full history, and it never invalidates thinking blocks. Skipped when the caller sends its own `context_management` or a compaction block. | | Shell output | on | Results of shell tools (`Bash`, `bash`, `shell`, `exec_command`, `run_shell_command` and similar) lose ANSI escapes and progress-bar redraws; runs of four or more identical lines become one line and a count, and runs of blank lines become one. Other tools' results, file reads included, keep their exact bytes, so an edit tool can still match them. | | Oversized results | 100,000 characters | A longer tool result keeps its first and last 50,000 characters with a marker in between. | -| Prompt limit | 100,000 tokens | A prompt estimated near the limit is counted exactly with `count_tokens`, which reports its size before and after context editing. Above the limit the request is refused with the error the harness treats as a full context window, so it compacts and retries instead of paying the higher price. The pricing page does not say whether the size before or after clearing picks the price, so by default the limit applies to the size before clearing; set `LEAN_PRICE_LIMIT_BASIS=cleared` once your bill shows that cleared prompts are priced by their smaller size. | +| Prompt limit | 99,000 tokens | A prompt estimated near the limit is counted with `count_tokens`, using the key that will send it, which reports its size before and after context editing. Anthropic calls the count an estimate that can differ slightly, so the default sits 1,000 tokens under the 100,000-token price line. Above the limit the request is refused with the error the harness treats as a full context window, so it compacts and retries instead of paying the higher price. The pricing page does not say whether the size before or after clearing picks the price, so by default the limit applies to the size before clearing; set `LEAN_PRICE_LIMIT_BASIS=cleared` once your bill shows that cleared prompts are priced by their smaller size. | | Concise answers | off | Adds a short instruction to the system prompt: lead with the answer, no filler, code and commands exact. Turn it on with `LEAN_STYLE=concise` or per request with `X-MultiLLM-Lean-Style: concise` (`caveman` is accepted too). | Lean mode never changes the model, effort, thinking settings or `max_tokens`. @@ -47,12 +47,17 @@ Lean mode never changes the model, effort, thinking settings or `max_tokens`. | Route | Response | |---|---| -| `/lean/v1/messages` | 400, Anthropic envelope, `invalid_request_error`, `prompt is too long: N tokens > 100000 maximum`. Claude Code compacts on this message. | +| `/lean/v1/messages` | 400, Anthropic envelope, `invalid_request_error`, `prompt is too long: N tokens > 99000 maximum`. Claude Code compacts on this message. | | `/lean/v1/chat/completions` | 400, `{"error": {"code": "context_length_exceeded", ...}}` with a message naming the maximum context length. | | `/lean/v1/responses` with `stream: true` | A `response.failed` event with `context_length_exceeded`, the form Codex reads. Without streaming, the Chat Completions body above. | -A count that fails (for example, rate-limited) lets the request through and -reports `X-MultiLLM-Lean-Prompt-Count: estimate`. If compacting cannot bring a +When the prompt cannot be counted, the local estimate decides, and an estimate +above the limit is refused: `count_tokens` rejects URL and file sources (lean +estimates them at 1,600 tokens per image and 10,000 per document), and a count +can fail, for example when rate-limited. Server tools such as web search are +left out of the count and added back at 1,000 tokens each, because +`count_tokens` rejects them. Such requests report +`X-MultiLLM-Lean-Prompt-Count: estimate`. If compacting cannot bring a prompt under the limit, for example because the newest message alone is larger, the harness reports the error: Claude Code stops after three failed compactions. @@ -63,7 +68,7 @@ the harness reports the error: Claude Code stops after three failed compactions. | `X-MultiLLM-Lean` | `applied` for a Claude API request, `inactive` for other providers | | `X-MultiLLM-Lean-Rules` | Version of the rules above | | `X-MultiLLM-Lean-Stages` | Stages that changed the request, for example `cache,clear-tool-uses,condense-shell` | -| `X-MultiLLM-Lean-Prompt-Tokens` | Prompt size: counted by the API near the limit, otherwise a high local estimate | +| `X-MultiLLM-Lean-Prompt-Tokens` | Prompt size the limit was checked against: counted by the API near the limit, otherwise a rough local estimate | | `X-MultiLLM-Lean-Prompt-Count` | `provider` or `estimate` | | `X-MultiLLM-Lean-Prompt-Tokens-After-Clearing` | Counted size after context editing, when it differs | | `X-MultiLLM-Lean-Condensed-Chars` | Characters removed from tool results | @@ -132,7 +137,7 @@ caching is what makes cache reads happen at all. | Variable | Default | | |---|---|---| | `LEAN_ENDPOINT_ENABLED` | `true` | `false` stops serving `/lean` | -| `LEAN_PROMPT_TOKEN_LIMIT` | `100000` | `0` turns the refusal off | +| `LEAN_PROMPT_TOKEN_LIMIT` | `99000` | `0` turns the refusal off | | `LEAN_PRICE_LIMIT_BASIS` | `original` | `cleared` applies the limit to the prompt after context editing | | `LEAN_CLEAR_TRIGGER_TOKENS` | `80000` | `0` turns tool result clearing off | | `LEAN_CLEAR_KEEP_TOOL_USES` | `3` | Recent tool uses that are never cleared | diff --git a/routes/unified.py b/routes/unified.py index fa953a89..ede1f92c 100644 --- a/routes/unified.py +++ b/routes/unified.py @@ -68,7 +68,7 @@ from services.managed_turn import bounded_managed_timeout as bounded_timeout from routes.cascades import dispatch_unified_cascade, register_cascade_admin_routes, validate_cascade_target from services.context_optimizer import ContextOptimizationResult -from services.lean_mode import merge_beta_header, prepare_lean_request +from services.lean_mode import check_lean_prompt, merge_beta_header, prepare_lean_request from services.model_registry import ModelRegistry from services.model_cooldown import ModelCooldownCapacity, ModelCooldownExhausted from routes.provider_credentials import ( @@ -309,15 +309,14 @@ def _send_native_request( """Send one body to a provider's native Responses or Messages endpoint.""" from services.managed_turn import prepare_dispatch_kwargs lean = prepare_lean_request( - body, - provider=provider, - endpoint=endpoint, - request_headers=headers_source, - base_urls=app.config["API_BASE_URLS"], + body, provider=provider, endpoint=endpoint, request_headers=headers_source ) raw_body = json.dumps(lean.body).encode("utf-8") def send_request(token: str): + check_lean_prompt( + lean, token, request_headers=headers_source, base_urls=app.config["API_BASE_URLS"] + ) headers = proxy_service_cls.prepare_headers( headers_source, provider, diff --git a/services/lean_mode.py b/services/lean_mode.py index 199b7eb7..6742db22 100644 --- a/services/lean_mode.py +++ b/services/lean_mode.py @@ -23,7 +23,7 @@ import os import re from collections.abc import Iterable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any from flask import g, has_request_context, request @@ -39,8 +39,9 @@ CLEAR_TOOL_USES = "clear_tool_uses_20250919" STYLE_HEADER = "X-MultiLLM-Lean-Style" # Haiku 5.5 bills a request whose prompt is over 100,000 tokens, cache reads and -# writes included, at five times the price of the whole request. -DEFAULT_PROMPT_TOKEN_LIMIT = 100_000 +# writes included, at five times the price of the whole request. The default sits +# 1,000 tokens lower because count_tokens is itself an estimate. +DEFAULT_PROMPT_TOKEN_LIMIT = 99_000 DEFAULT_CLEAR_TRIGGER_TOKENS = 80_000 DEFAULT_CLEAR_KEEP_TOOL_USES = 3 DEFAULT_CLEAR_AT_LEAST_TOKENS = 20_000 @@ -50,6 +51,11 @@ # from 40% of the limit leaves that margin before a prompt can cross it uncounted. COUNT_FROM_LIMIT_SHARE = 0.4 MEDIA_BLOCK_TOKENS = 1_600 +# count_tokens rejects URL and file sources and every server tool but the +# advisor, so those are estimated at these sizes instead. +REMOTE_DOCUMENT_TOKENS = 10_000 +SERVER_TOOL_TOKENS = 1_000 +UNCOUNTABLE_SOURCES = frozenset({"url", "file"}) REPEATED_LINE_MIN_RUN = 4 SHELL_TOOL_NAMES = frozenset( { @@ -301,12 +307,23 @@ def _text_length(value: Any) -> int: return 0 -def _has_cache_control(value: Any) -> bool: - if isinstance(value, dict): - return "cache_control" in value or any(_has_cache_control(item) for item in value.values()) - if isinstance(value, list): - return any(_has_cache_control(item) for item in value) - return False +def _has_cache_breakpoint(body: Mapping[str, Any]) -> bool: + """Whether the caller placed a breakpoint where the API reads one. + + Only blocks count: a tool schema or a tool call's input can have an ordinary + property named `cache_control`. + """ + if "cache_control" in body: + return True + candidates: list = [] + for key in ("system", "tools"): + if isinstance(body.get(key), list): + candidates.extend(body[key]) + for block in _blocks(body.get("messages")): + candidates.append(block) + if block.get("type") == "tool_result" and isinstance(block.get("content"), list): + candidates.extend(block["content"]) + return any(isinstance(block, dict) and "cache_control" in block for block in candidates) def _with_style(body: dict) -> dict: @@ -328,7 +345,7 @@ def _with_cache_breakpoints(body: dict, ttl: str) -> tuple[dict, bool]: breakpoint keeps the tools and system prompt cached when context editing clears an old tool result and the rest of the prefix must be written again. """ - if _has_cache_control(body): + if _has_cache_breakpoint(body): return body, False marker = {"type": "ephemeral"} if ttl == "5m" else {"type": "ephemeral", "ttl": ttl} cached = {**body, "cache_control": dict(marker)} @@ -379,17 +396,38 @@ def _with_tool_clearing(body: dict, policy: LeanPolicy) -> tuple[dict, bool]: return {**body, "context_management": {"edits": [edit]}}, True -def estimate_prompt_tokens(body: Mapping[str, Any]) -> int: - """A rough estimate that decides when to count exactly: three characters per - token, images at a flat size.""" +def _is_server_tool(tool: Any) -> bool: + return isinstance(tool, dict) and tool.get("type") not in (None, "custom") + + +def _estimate(body: Mapping[str, Any]) -> tuple[int, bool]: + """Estimate the prompt and say whether count_tokens can count it. + + Text counts three characters per token, an image a flat size, a text + document its characters and a base64 PDF a token per three decoded bytes. + """ characters = 0 - media = 0 + tokens = 0 + countable = True def walk(value: Any) -> None: - nonlocal characters, media + nonlocal characters, tokens, countable if isinstance(value, dict): - if value.get("type") in {"image", "document"}: - media += 1 + kind = value.get("type") + source = value.get("source") if isinstance(value.get("source"), dict) else {} + if kind in {"image", "document"}: + data = source.get("data") + if source.get("type") in UNCOUNTABLE_SOURCES: + countable = False + tokens += MEDIA_BLOCK_TOKENS if kind == "image" else REMOTE_DOCUMENT_TOKENS + elif kind == "image": + tokens += MEDIA_BLOCK_TOKENS + elif source.get("type") == "content": + walk(source.get("content")) + elif isinstance(data, str): + tokens += len(data) // 4 if source.get("type") == "base64" else len(data) // 3 + else: + tokens += REMOTE_DOCUMENT_TOKENS return for key, item in value.items(): if key != "cache_control": @@ -402,7 +440,31 @@ def walk(value: Any) -> None: for key in ("system", "tools", "messages"): walk(body.get(key)) - return characters // 3 + media * MEDIA_BLOCK_TOKENS + tools = body.get("tools") if isinstance(body.get("tools"), list) else [] + tokens += SERVER_TOOL_TOKENS * sum(1 for tool in tools if _is_server_tool(tool)) + return characters // 3 + tokens, countable + + +def estimate_prompt_tokens(body: Mapping[str, Any]) -> int: + """A rough estimate that decides when to count exactly.""" + return _estimate(body)[0] + + +def _count_body(body: dict) -> tuple[dict, int]: + """Drop the server tools count_tokens rejects; return the body and their allowance.""" + tools = body.get("tools") + if not isinstance(tools, list) or not any(_is_server_tool(tool) for tool in tools): + return body, 0 + client_tools = [tool for tool in tools if not _is_server_tool(tool)] + counted = {key: value for key, value in body.items() if key != "tools"} + if client_tools: + counted["tools"] = client_tools + choice = body.get("tool_choice") + if isinstance(choice, dict) and choice.get("type") == "tool" and not any( + tool.get("name") == choice.get("name") for tool in client_tools if isinstance(tool, dict) + ): + counted.pop("tool_choice", None) + return counted, SERVER_TOOL_TOKENS * (len(tools) - len(client_tools)) def merged_beta_value(existing: str, betas: Iterable[str]) -> str: @@ -422,27 +484,40 @@ def merge_beta_header(headers: dict, betas: Iterable[str]) -> None: headers[name] = merged_beta_value(headers.get(name, ""), betas) -def _prompt_size( - body: dict, - betas: tuple[str, ...], +def check_lean_prompt( + decision: LeanRequest, + credential: str | None, + *, request_headers: Mapping[str, Any] | None, base_urls: Mapping[str, str], - policy: LeanPolicy, -) -> tuple[int, str, int | None]: - """Return the size the price limit applies to, how it was measured, and the - counted size after context editing.""" - estimate = estimate_prompt_tokens(body) +) -> None: + """Refuse a prepared prompt that would pay the long-prompt price. + + Counts with the key that will send the request: thinking blocks from another + account are dropped before counting, so any other key can undercount. When + the prompt cannot be counted, the estimate decides. + """ + if not decision.active: + return + policy = lean_policy() limit = policy.prompt_token_limit - if limit <= 0 or estimate < limit * COUNT_FROM_LIMIT_SHARE: - return estimate, "estimate", None - beta_header = merged_beta_value(_header(request_headers, "anthropic-beta"), betas) - try: - after, before = count_claude_prompt_tokens(body, beta_header, base_urls) - except NativeTokenCountError as error: - # Counting is a guard, not a gate: a failed count lets the request through. - logger.warning("Lean prompt count failed status=%s; sending uncounted", error.status) - return estimate, "estimate", None - return (before if policy.price_limit_basis == "original" else after), "provider", after + estimate, countable = _estimate(decision.body) + tokens, source, after = estimate, "estimate", None + if limit > 0 and countable and estimate >= limit * COUNT_FROM_LIMIT_SHARE: + count_body, allowance = _count_body(decision.body) + beta_header = merged_beta_value(_header(request_headers, "anthropic-beta"), decision.betas) + try: + after, before = count_claude_prompt_tokens(count_body, beta_header, base_urls, credential=credential) + except NativeTokenCountError as error: + logger.warning("Lean prompt count failed status=%s; using the estimate", error.status) + else: + after, before = after + allowance, before + allowance + tokens, source = (before if policy.price_limit_basis == "original" else after), "provider" + g.lean_request = replace( + decision, prompt_tokens=tokens, prompt_count=source, prompt_tokens_after_clearing=after + ) + if limit > 0 and tokens > limit: + raise LeanPromptTooLong(tokens, limit) def prepare_lean_request( @@ -451,11 +526,10 @@ def prepare_lean_request( provider: str, endpoint: str, request_headers: Mapping[str, Any] | None, - base_urls: Mapping[str, str], ) -> LeanRequest: """Apply lean mode to a Claude API Messages body; other requests pass unchanged. - Raises LeanPromptTooLong when the prepared prompt would pay the long-prompt price. + check_lean_prompt then sizes the result against the price limit. """ if not lean_requested(): return LeanRequest(body=body) @@ -478,20 +552,14 @@ def prepare_lean_request( betas = (CONTEXT_MANAGEMENT_BETA,) if clearing else () if clearing: stages.add("clear-tool-uses") - tokens, source, after_clearing = _prompt_size(prepared, betas, request_headers, base_urls, policy) decision = LeanRequest( body=prepared, active=True, betas=betas, stages=tuple(sorted(stages)), - prompt_tokens=tokens, - prompt_count=source, - prompt_tokens_after_clearing=after_clearing, condensed_chars=max(0, condensed_chars), ) g.lean_request = decision - if source == "provider" and tokens > policy.prompt_token_limit: - raise LeanPromptTooLong(tokens, policy.prompt_token_limit) return decision @@ -505,8 +573,9 @@ def lean_response_headers(response): response.headers["X-MultiLLM-Lean-Rules"] = LEAN_RULES_VERSION if active: response.headers["X-MultiLLM-Lean-Stages"] = ",".join(decision.stages) or "none" - response.headers["X-MultiLLM-Lean-Prompt-Tokens"] = str(decision.prompt_tokens) - response.headers["X-MultiLLM-Lean-Prompt-Count"] = decision.prompt_count + if decision.prompt_tokens is not None: + response.headers["X-MultiLLM-Lean-Prompt-Tokens"] = str(decision.prompt_tokens) + response.headers["X-MultiLLM-Lean-Prompt-Count"] = decision.prompt_count after = decision.prompt_tokens_after_clearing if after is not None and after != decision.prompt_tokens: response.headers["X-MultiLLM-Lean-Prompt-Tokens-After-Clearing"] = str(after) diff --git a/services/native_token_count.py b/services/native_token_count.py index e10d7098..6786bfe5 100644 --- a/services/native_token_count.py +++ b/services/native_token_count.py @@ -129,16 +129,17 @@ def count_native_tokens( def count_claude_prompt_tokens( - body: Mapping, beta_header: str, base_urls: Mapping[str, str] + body: Mapping, beta_header: str, base_urls: Mapping[str, str], *, credential: str | None = None ) -> tuple[int, int]: """Count a Claude API Messages body exactly as it will be sent. Returns the prompt size after and before server-side context editing; they - are equal when the body clears nothing. + are equal when the body clears nothing. Pass the key that will send the + request: the API drops thinking blocks from other accounts before counting. """ from services.auth_service import AuthService - credential = AuthService.get_api_key("anthropic") + credential = credential or AuthService.get_api_key("anthropic") if not credential: raise NativeTokenCountError("No Claude API key is configured", 503) url = _count_url(base_urls["anthropic"], CLAUDE_COUNT_PATH) diff --git a/tests/outbound_inventory.json b/tests/outbound_inventory.json index 0f9e0571..b2b81302 100644 --- a/tests/outbound_inventory.json +++ b/tests/outbound_inventory.json @@ -402,7 +402,7 @@ "routes/unified.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "38b545840293971d167268ebecd8095e638f0717fe5fc4782925602e5eaa47d6" + "sha256": "652ff5a01cee4be43058d3420f388c28ee4535a8bb237379de2da6cc91f8a7ed" }, "routes/unified_bridge.py": { "kind": "non-provider", @@ -997,7 +997,7 @@ "services/native_token_count.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "73c3de349f04b1524c63b8d19849306213d2e2c6dd9b5c459ebc652ce7a4c402" + "sha256": "36b3dd65a9d49ee74a98461f33aaad0fee9de8f7e142d3f7cff90bd0112c5395" }, "services/observability_adapters.py": { "kind": "non-provider", @@ -2467,6 +2467,6 @@ "services/lean_mode.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "4840552159fca30603f1da07cfa008abe95ad82c1243a6302e2b5f73ca95195b" + "sha256": "b300d409c616896ca34ee98aa0be2ab7e407f0db63f66781bdc7a1aab3033064" } } diff --git a/tests/test_lean_mode.py b/tests/test_lean_mode.py index 8626e9f6..194c143a 100644 --- a/tests/test_lean_mode.py +++ b/tests/test_lean_mode.py @@ -60,13 +60,18 @@ def prepare(self, body, headers=None): from flask import request request.environ[lean_mode.LEAN_ENVIRON_KEY] = True - return lean_mode.prepare_lean_request( - body, - provider="anthropic", - endpoint="v1/messages", + decision = lean_mode.prepare_lean_request( + body, provider="anthropic", endpoint="v1/messages", request_headers=headers or {} + ) + lean_mode.check_lean_prompt( + decision, + "attempt-key", request_headers=headers or {}, base_urls={"anthropic": "https://api.anthropic.com"}, ) + from flask import g + + return g.lean_request @pytest.fixture @@ -230,13 +235,14 @@ def test_near_limit_prompt_is_counted_and_refused_above_the_price_limit(lean): with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(100_001, 100_001)) as count: with pytest.raises(lean_mode.LeanPromptTooLong) as raised: lean.prepare(body, headers={"anthropic-beta": "claude-code-20250219"}) - assert str(raised.value) == "prompt is too long: 100001 tokens > 100000 maximum" + assert str(raised.value) == "prompt is too long: 100001 tokens > 99000 maximum" counted_body, beta_header, _ = count.call_args.args + assert count.call_args.kwargs == {"credential": "attempt-key"} assert counted_body["context_management"]["edits"][0]["type"] == "clear_tool_uses_20250919" assert beta_header == "claude-code-20250219,context-management-2025-06-27" - with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(99_000, 99_000)): + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(98_000, 98_000)): decision = lean.prepare(body) - assert (decision.prompt_tokens, decision.prompt_count) == (99_000, "provider") + assert (decision.prompt_tokens, decision.prompt_count) == (98_000, "provider") def test_price_limit_applies_before_clearing_unless_configured(lean): @@ -252,12 +258,57 @@ def test_price_limit_applies_before_clearing_unless_configured(lean): assert (decision.prompt_tokens, decision.prompt_tokens_after_clearing) == (60_000, 60_000) -def test_failed_count_lets_the_request_through(lean): +def test_failed_count_falls_back_to_the_estimate(lean): body = conversation(1) body["messages"][0]["content"] = "x" * 200_000 - with patch.object(lean_mode, "count_claude_prompt_tokens", side_effect=NativeTokenCountError("down", 429)): + failure = NativeTokenCountError("down", 429) + with patch.object(lean_mode, "count_claude_prompt_tokens", side_effect=failure): decision = lean.prepare(body) - assert decision.prompt_count == "estimate" + assert (decision.prompt_count, decision.prompt_tokens < 99_000) == ("estimate", True) + body["messages"][0]["content"] = "x" * 400_000 + with pytest.raises(lean_mode.LeanPromptTooLong): + lean.prepare(body) + + +def test_documents_are_estimated_by_their_contents(): + text = {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "word " * 120_000}} + pdf = {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "A" * 400_000}} + nested = {"type": "document", "source": {"type": "content", "content": [{"type": "text", "text": "abc" * 1_000}]}} + assert estimate_prompt_tokens({"messages": [{"role": "user", "content": [text]}]}) >= 200_000 + assert estimate_prompt_tokens({"messages": [{"role": "user", "content": [pdf]}]}) >= 100_000 + assert estimate_prompt_tokens({"messages": [{"role": "user", "content": [nested]}]}) >= 1_000 + + +def test_uncountable_sources_are_judged_by_the_estimate(lean): + remote = {"type": "document", "source": {"type": "url", "url": "https://example.com/spec.pdf"}} + body = {"model": MODEL, "max_tokens": 64, "messages": [{"role": "user", "content": [remote, {"type": "text", "text": "x" * 300_000}]}]} + with patch.object(lean_mode, "count_claude_prompt_tokens") as count: + with pytest.raises(lean_mode.LeanPromptTooLong): + lean.prepare(body) + count.assert_not_called() + + +def test_server_tools_are_left_out_of_the_count_with_an_allowance(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + body["tools"] = [*TOOLS, {"type": "web_search_20250305", "name": "web_search"}] + body["tool_choice"] = {"type": "tool", "name": "web_search"} + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(70_000, 70_000)) as count: + decision = lean.prepare(body) + counted_body = count.call_args.args[0] + assert [tool["name"] for tool in counted_body["tools"]] == ["Bash"] + assert "tool_choice" not in counted_body + assert decision.prompt_tokens == 70_000 + lean_mode.SERVER_TOOL_TOKENS + assert decision.body["tools"][-1]["name"] == "web_search" + + +def test_schema_properties_named_cache_control_are_not_breakpoints(lean): + body = conversation(1) + body["tools"] = [{**TOOLS[0], "input_schema": {"type": "object", "properties": {"cache_control": {"type": "string"}}}}] + body["messages"][1]["content"][0]["input"] = {"cache_control": "no-store"} + decision = lean.prepare(body) + assert "cache" in decision.stages + assert decision.body["cache_control"] == {"type": "ephemeral"} def test_claude_count_sends_context_management_with_the_lean_beta(): @@ -379,18 +430,20 @@ def large_input(self): return "x" * 200_000 def test_over_limit_messages_get_anthropic_prompt_too_long(self): - with self.oversized(): + with self.oversized() as count: result, send = self.post( "/lean/v1/messages", {"model": SELECTOR, "max_tokens": 64, "messages": [{"role": "user", "content": self.large_input()}]}, headers={"X-Api-Key": "admin-test-key"}, ) send.assert_not_called() + # Counted with the key that would have sent the request. + self.assertEqual(count.call_args.kwargs, {"credential": "anthropic-provider-key"}) self.assertEqual(result.status_code, 400) error = result.get_json() self.assertEqual(error["type"], "error") self.assertEqual(error["error"]["type"], "invalid_request_error") - self.assertEqual(error["error"]["message"], "prompt is too long: 150000 tokens > 100000 maximum") + self.assertEqual(error["error"]["message"], "prompt is too long: 150000 tokens > 99000 maximum") self.assertEqual(result.headers["X-MultiLLM-Lean-Prompt-Tokens"], "150000") def test_over_limit_chat_gets_context_length_exceeded(self): @@ -403,7 +456,7 @@ def test_over_limit_chat_gets_context_length_exceeded(self): self.assertEqual(result.status_code, 400) error = result.get_json()["error"] self.assertEqual(error["code"], "context_length_exceeded") - self.assertIn("maximum context length is 100000 tokens", error["message"]) + self.assertIn("maximum context length is 99000 tokens", error["message"]) self.assertIn("prompt is too long", error["message"]) def test_over_limit_responses_stream_fails_like_openai(self): From b0a3397f496fa2eb890851698648f95cbca538b3 Mon Sep 17 00:00:00 2001 From: Abhyudaya Baiju Tharakan Date: Mon, 26 Oct 2026 21:50:58 +0530 Subject: [PATCH 4/5] fix(lean): satisfy the CI security and type checks Bandit read the shell= keyword of the condensing helpers as a subprocess call; the parameter is now terminal_output. Narrow the estimate's source and tools lookups for mypy. --- services/lean_mode.py | 18 ++++++++++-------- tests/outbound_inventory.json | 2 +- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/services/lean_mode.py b/services/lean_mode.py index 6742db22..035245d6 100644 --- a/services/lean_mode.py +++ b/services/lean_mode.py @@ -237,9 +237,9 @@ def _tool_names(messages: list) -> dict[str, str]: return names -def _condense_text(text: str, *, shell: bool, max_chars: int, stages: set[str]) -> str: +def _condense_text(text: str, *, terminal_output: bool, max_chars: int, stages: set[str]) -> str: result = text - if shell: + if terminal_output: condensed = condense_shell_output(result) if condensed != result: stages.add("condense-shell") @@ -251,17 +251,17 @@ def _condense_text(text: str, *, shell: bool, max_chars: int, stages: set[str]) return result -def _condense_tool_result(block: dict, *, shell: bool, max_chars: int, stages: set[str]) -> dict: +def _condense_tool_result(block: dict, *, terminal_output: bool, max_chars: int, stages: set[str]) -> dict: content = block.get("content") if isinstance(content, str): - condensed = _condense_text(content, shell=shell, max_chars=max_chars, stages=stages) + condensed = _condense_text(content, terminal_output=terminal_output, max_chars=max_chars, stages=stages) return block if condensed == content else {**block, "content": condensed} if not isinstance(content, list): return block parts = [] for part in content: if isinstance(part, dict) and part.get("type") == "text" and isinstance(part.get("text"), str): - condensed = _condense_text(part["text"], shell=shell, max_chars=max_chars, stages=stages) + condensed = _condense_text(part["text"], terminal_output=terminal_output, max_chars=max_chars, stages=stages) parts.append(part if condensed == part["text"] else {**part, "text": condensed}) else: parts.append(part) @@ -286,7 +286,7 @@ def _condense_tool_results(messages: Any, policy: LeanPolicy, stages: set[str]) tool_name = names.get(block.get("tool_use_id"), "").lower() block = _condense_tool_result( block, - shell=policy.condense_shell_output and tool_name in SHELL_TOOL_NAMES, + terminal_output=policy.condense_shell_output and tool_name in SHELL_TOOL_NAMES, max_chars=policy.tool_result_max_chars, stages=stages, ) @@ -414,7 +414,8 @@ def walk(value: Any) -> None: nonlocal characters, tokens, countable if isinstance(value, dict): kind = value.get("type") - source = value.get("source") if isinstance(value.get("source"), dict) else {} + raw_source = value.get("source") + source: dict[str, Any] = raw_source if isinstance(raw_source, dict) else {} if kind in {"image", "document"}: data = source.get("data") if source.get("type") in UNCOUNTABLE_SOURCES: @@ -440,7 +441,8 @@ def walk(value: Any) -> None: for key in ("system", "tools", "messages"): walk(body.get(key)) - tools = body.get("tools") if isinstance(body.get("tools"), list) else [] + raw_tools = body.get("tools") + tools: list[Any] = raw_tools if isinstance(raw_tools, list) else [] tokens += SERVER_TOOL_TOKENS * sum(1 for tool in tools if _is_server_tool(tool)) return characters // 3 + tokens, countable diff --git a/tests/outbound_inventory.json b/tests/outbound_inventory.json index b2b81302..5cde48d4 100644 --- a/tests/outbound_inventory.json +++ b/tests/outbound_inventory.json @@ -2467,6 +2467,6 @@ "services/lean_mode.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "b300d409c616896ca34ee98aa0be2ab7e407f0db63f66781bdc7a1aab3033064" + "sha256": "d31473a2991bcef8ef7216819383d7cebe865ba13d653795e884a4e55ba82cc8" } } From 58c58d6ac02e184c88ba6e04a64024eafd18e6c7 Mon Sep 17 00:00:00 2001 From: Abhyudaya Baiju Tharakan Date: Mon, 26 Oct 2026 21:58:30 +0530 Subject: [PATCH 5/5] fix(lean): keep thinking-block counts valid and let key rotation handle refused keys count_tokens runs the same thinking-block prefix check as generation, so a body that replays thinking blocks is now counted with its tools unchanged instead of without its server tools. The advisor tool, which count_tokens accepts, is no longer stripped. A count that refuses the key (401, 402, 403) no longer decides the request: it goes out so rotation can rest that key and count again with the next. Prompts with a PDF are always counted, since file size does not bound the text inside. The docs say plainly when the limit is an estimate rather than a guarantee. --- docs/lean.md | 21 ++++++++++++------ services/lean_mode.py | 42 +++++++++++++++++++++++++++-------- tests/outbound_inventory.json | 2 +- tests/test_lean_mode.py | 39 ++++++++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 17 deletions(-) diff --git a/docs/lean.md b/docs/lean.md index 08bd4f62..bcbbd2f9 100644 --- a/docs/lean.md +++ b/docs/lean.md @@ -52,12 +52,19 @@ Lean mode never changes the model, effort, thinking settings or `max_tokens`. | `/lean/v1/responses` with `stream: true` | A `response.failed` event with `context_length_exceeded`, the form Codex reads. Without streaming, the Chat Completions body above. | When the prompt cannot be counted, the local estimate decides, and an estimate -above the limit is refused: `count_tokens` rejects URL and file sources (lean -estimates them at 1,600 tokens per image and 10,000 per document), and a count -can fail, for example when rate-limited. Server tools such as web search are -left out of the count and added back at 1,000 tokens each, because -`count_tokens` rejects them. Such requests report -`X-MultiLLM-Lean-Prompt-Count: estimate`. If compacting cannot bring a +above the limit is refused. That happens for URL and file sources, which +`count_tokens` rejects (lean estimates them at 1,600 tokens per image and 10,000 +per document); for a request that replays thinking blocks alongside a server tool +such as web search, because `count_tokens` rejects server tools and dropping them +would fail its thinking-block check; and when a count fails, for example when +rate-limited. Such requests report `X-MultiLLM-Lean-Prompt-Count: estimate`, and +the limit is then a safety net, not a guarantee: the estimate can be low for a +URL document or an unusual tokenization. Without replayed thinking, server tools +are left out of the count and added back at 1,000 tokens each. A prompt with a +PDF is always counted, since its file size does not bound the text inside. A count +that refuses the key itself (401, 402 or 403) does not decide: the request goes +out so key rotation can rest that key and count again with the next one +(`X-MultiLLM-Lean-Prompt-Count: skipped`). If compacting cannot bring a prompt under the limit, for example because the newest message alone is larger, the harness reports the error: Claude Code stops after three failed compactions. @@ -69,7 +76,7 @@ the harness reports the error: Claude Code stops after three failed compactions. | `X-MultiLLM-Lean-Rules` | Version of the rules above | | `X-MultiLLM-Lean-Stages` | Stages that changed the request, for example `cache,clear-tool-uses,condense-shell` | | `X-MultiLLM-Lean-Prompt-Tokens` | Prompt size the limit was checked against: counted by the API near the limit, otherwise a rough local estimate | -| `X-MultiLLM-Lean-Prompt-Count` | `provider` or `estimate` | +| `X-MultiLLM-Lean-Prompt-Count` | `provider`, `estimate`, or `skipped` when the count refused the key | | `X-MultiLLM-Lean-Prompt-Tokens-After-Clearing` | Counted size after context editing, when it differs | | `X-MultiLLM-Lean-Condensed-Chars` | Characters removed from tool results | diff --git a/services/lean_mode.py b/services/lean_mode.py index 035245d6..e0f56ea3 100644 --- a/services/lean_mode.py +++ b/services/lean_mode.py @@ -56,6 +56,8 @@ REMOTE_DOCUMENT_TOKENS = 10_000 SERVER_TOOL_TOKENS = 1_000 UNCOUNTABLE_SOURCES = frozenset({"url", "file"}) +# A count refused for these reasons says the key is unusable, not the prompt. +KEY_REFUSAL_STATUSES = frozenset({401, 402, 403}) REPEATED_LINE_MIN_RUN = 4 SHELL_TOOL_NAMES = frozenset( { @@ -397,21 +399,26 @@ def _with_tool_clearing(body: dict, policy: LeanPolicy) -> tuple[dict, bool]: def _is_server_tool(tool: Any) -> bool: - return isinstance(tool, dict) and tool.get("type") not in (None, "custom") + """A server tool count_tokens rejects; it accepts client tools and the advisor.""" + kind = tool.get("type") if isinstance(tool, dict) else None + return kind not in (None, "custom") and not str(kind).startswith("advisor_") -def _estimate(body: Mapping[str, Any]) -> tuple[int, bool]: - """Estimate the prompt and say whether count_tokens can count it. +def _estimate(body: Mapping[str, Any]) -> tuple[int, bool, bool]: + """Estimate the prompt; say whether count_tokens can count it and whether it must. Text counts three characters per token, an image a flat size, a text document its characters and a base64 PDF a token per three decoded bytes. + A PDF's file size does not bound the text inside it, so a prompt with one + is always counted. """ characters = 0 tokens = 0 countable = True + must_count = False def walk(value: Any) -> None: - nonlocal characters, tokens, countable + nonlocal characters, tokens, countable, must_count if isinstance(value, dict): kind = value.get("type") raw_source = value.get("source") @@ -426,6 +433,7 @@ def walk(value: Any) -> None: elif source.get("type") == "content": walk(source.get("content")) elif isinstance(data, str): + must_count = must_count or source.get("type") == "base64" tokens += len(data) // 4 if source.get("type") == "base64" else len(data) // 3 else: tokens += REMOTE_DOCUMENT_TOKENS @@ -444,7 +452,7 @@ def walk(value: Any) -> None: raw_tools = body.get("tools") tools: list[Any] = raw_tools if isinstance(raw_tools, list) else [] tokens += SERVER_TOOL_TOKENS * sum(1 for tool in tools if _is_server_tool(tool)) - return characters // 3 + tokens, countable + return characters // 3 + tokens, countable, must_count def estimate_prompt_tokens(body: Mapping[str, Any]) -> int: @@ -453,9 +461,19 @@ def estimate_prompt_tokens(body: Mapping[str, Any]) -> int: def _count_body(body: dict) -> tuple[dict, int]: - """Drop the server tools count_tokens rejects; return the body and their allowance.""" + """Drop the server tools count_tokens rejects; return the body and their allowance. + + count_tokens runs the same thinking-block prefix check as generation, so a + body that replays thinking blocks is counted unchanged: without its server + tools it would fail the check, and with them count_tokens rejects it, which + leaves the estimate to decide. + """ tools = body.get("tools") - if not isinstance(tools, list) or not any(_is_server_tool(tool) for tool in tools): + if ( + not isinstance(tools, list) + or not any(_is_server_tool(tool) for tool in tools) + or any(block.get("type") in {"thinking", "redacted_thinking"} for block in _blocks(body.get("messages"))) + ): return body, 0 client_tools = [tool for tool in tools if not _is_server_tool(tool)] counted = {key: value for key, value in body.items() if key != "tools"} @@ -503,14 +521,20 @@ def check_lean_prompt( return policy = lean_policy() limit = policy.prompt_token_limit - estimate, countable = _estimate(decision.body) + estimate, countable, must_count = _estimate(decision.body) tokens, source, after = estimate, "estimate", None - if limit > 0 and countable and estimate >= limit * COUNT_FROM_LIMIT_SHARE: + if limit > 0 and countable and (must_count or estimate >= limit * COUNT_FROM_LIMIT_SHARE): count_body, allowance = _count_body(decision.body) beta_header = merged_beta_value(_header(request_headers, "anthropic-beta"), decision.betas) try: after, before = count_claude_prompt_tokens(count_body, beta_header, base_urls, credential=credential) except NativeTokenCountError as error: + if error.status in KEY_REFUSAL_STATUSES: + # Send anyway: the generation request's refusal lets key rotation + # rest this key, and the next attempt counts with the next key. + logger.warning("Lean prompt count refused the key status=%s; sending to rotate", error.status) + g.lean_request = replace(decision, prompt_tokens=estimate, prompt_count="skipped") + return logger.warning("Lean prompt count failed status=%s; using the estimate", error.status) else: after, before = after + allowance, before + allowance diff --git a/tests/outbound_inventory.json b/tests/outbound_inventory.json index 5cde48d4..6f4ccec0 100644 --- a/tests/outbound_inventory.json +++ b/tests/outbound_inventory.json @@ -2467,6 +2467,6 @@ "services/lean_mode.py": { "kind": "guarded", "reason": "Provider adapter, protocol helper or orchestration; delegates network I/O to the guarded provider clients.", - "sha256": "d31473a2991bcef8ef7216819383d7cebe865ba13d653795e884a4e55ba82cc8" + "sha256": "154f0eae820ef2a040d16f49e87cee9d42fb623ac3c17b9dc9986d39f1a86937" } } diff --git a/tests/test_lean_mode.py b/tests/test_lean_mode.py index 194c143a..a0e7f008 100644 --- a/tests/test_lean_mode.py +++ b/tests/test_lean_mode.py @@ -302,6 +302,45 @@ def test_server_tools_are_left_out_of_the_count_with_an_allowance(lean): assert decision.body["tools"][-1]["name"] == "web_search" +def test_a_refused_key_on_the_count_lets_rotation_try_the_next_key(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 400_000 + with patch.object(lean_mode, "count_claude_prompt_tokens", side_effect=NativeTokenCountError("refused", 401)): + decision = lean.prepare(body) + assert decision.prompt_count == "skipped" + + +def test_replayed_thinking_is_counted_with_its_server_tools_unchanged(lean): + """count_tokens runs generation's prefix check, so dropping tools would fail it.""" + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + body["tools"] = [*TOOLS, {"type": "web_search_20250305", "name": "web_search"}] + body["messages"][1]["content"].insert(0, {"type": "thinking", "thinking": "", "signature": "sig"}) + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(70_000, 70_000)) as count: + decision = lean.prepare(body) + assert count.call_args.args[0]["tools"] == decision.body["tools"] + assert decision.prompt_tokens == 70_000 + + +def test_the_advisor_tool_is_counted_not_stripped(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + body["tools"] = [*TOOLS, {"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-5-5"}] + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(70_000, 70_000)) as count: + decision = lean.prepare(body) + assert [tool["name"] for tool in count.call_args.args[0]["tools"]] == ["Bash", "advisor"] + assert decision.prompt_tokens == 70_000 + + +def test_a_pdf_is_always_counted(lean): + pdf = {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "A" * 4_000}} + body = {"model": MODEL, "max_tokens": 64, "messages": [{"role": "user", "content": [pdf, {"type": "text", "text": "Summarize"}]}]} + with patch.object(lean_mode, "count_claude_prompt_tokens", return_value=(120_000, 120_000)) as count: + with pytest.raises(lean_mode.LeanPromptTooLong): + lean.prepare(body) + count.assert_called_once() + + def test_schema_properties_named_cache_control_are_not_breakpoints(lean): body = conversation(1) body["tools"] = [{**TOOLS[0], "input_schema": {"type": "object", "properties": {"cache_control": {"type": "string"}}}}]