diff --git a/.env.example b/.env.example index 862a5e25..3b3ab971 100644 --- a/.env.example +++ b/.env.example @@ -85,6 +85,19 @@ 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=99000 +# LEAN_PRICE_LIMIT_BASIS=original +# 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..bcbbd2f9 --- /dev/null +++ b/docs/lean.md @@ -0,0 +1,177 @@ +# 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 | 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`. + +### The refusal + +| Route | Response | +|---|---| +| `/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. | + +When the prompt cannot be counted, the local estimate decides, and an estimate +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. + +### 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 the limit was checked against: counted by the API near the limit, otherwise a rough local 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 | + +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` | `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 | +| `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..ede1f92c 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 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 ( @@ -307,9 +308,15 @@ 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 + ) + 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, @@ -317,6 +324,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..e0f56ea3 --- /dev/null +++ b/services/lean_mode.py @@ -0,0 +1,610 @@ +"""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, replace +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. 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 +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 +# 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"}) +# 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( + { + "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" + # 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) +class LeanRequest: + body: dict + active: bool = False + betas: tuple[str, ...] = () + stages: tuple[str, ...] = () + prompt_tokens: int | None = None + prompt_count: str = "" + prompt_tokens_after_clearing: int | None = None + 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", + price_limit_basis=( + "cleared" + if str(environ.get("LEAN_PRICE_LIMIT_BASIS", "")).strip().lower() == "cleared" + else "original" + ), + ) + + +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, *, terminal_output: bool, max_chars: int, stages: set[str]) -> str: + result = text + if terminal_output: + 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, *, terminal_output: bool, max_chars: int, stages: set[str]) -> dict: + content = block.get("content") + if isinstance(content, str): + 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"], 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) + 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, + terminal_output=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_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: + 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_breakpoint(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 _is_server_tool(tool: Any) -> bool: + """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, 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, must_count + if isinstance(value, dict): + kind = value.get("type") + 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: + 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): + 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 + 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)) + 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, must_count + + +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. + + 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) + 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"} + 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: + 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 check_lean_prompt( + decision: LeanRequest, + credential: str | None, + *, + request_headers: Mapping[str, Any] | None, + base_urls: Mapping[str, str], +) -> 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 + estimate, countable, must_count = _estimate(decision.body) + tokens, source, after = estimate, "estimate", None + 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 + 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( + body: dict, + *, + provider: str, + endpoint: str, + request_headers: Mapping[str, Any] | None, +) -> LeanRequest: + """Apply lean mode to a Claude API Messages body; other requests pass unchanged. + + check_lean_prompt then sizes the result against the price limit. + """ + 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") + decision = LeanRequest( + body=prepared, + active=True, + betas=betas, + stages=tuple(sorted(stages)), + condensed_chars=max(0, condensed_chars), + ) + g.lean_request = decision + 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" + 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) + 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..6786bfe5 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): @@ -73,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: @@ -86,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( @@ -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,48 @@ 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)["input_tokens"] + + +def count_claude_prompt_tokens( + 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. 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 = 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) + 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} + 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( + provider: str, + url: str, + path: str, + credential: str, + incoming: Mapping[str, str], + body: dict, +) -> dict: + 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..6f4ccec0 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": "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": "46cd0e4dff7d3d123fe710c1cc00dcdf39a57a56ce1008446c94c8f4975282ac" + "sha256": "36b3dd65a9d49ee74a98461f33aaad0fee9de8f7e142d3f7cff90bd0112c5395" }, "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": "44b80ae15dc92596588b2cde34c956835b9d696df28fcea3eb48479121e91d03" }, "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": "154f0eae820ef2a040d16f49e87cee9d42fb623ac3c17b9dc9986d39f1a86937" } } 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..a0e7f008 --- /dev/null +++ b/tests/test_lean_mode.py @@ -0,0 +1,532 @@ +"""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 + 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 +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, 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 > 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=(98_000, 98_000)): + decision = lean.prepare(body) + assert (decision.prompt_tokens, decision.prompt_count) == (98_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_falls_back_to_the_estimate(lean): + body = conversation(1) + body["messages"][0]["content"] = "x" * 200_000 + failure = NativeTokenCountError("down", 429) + with patch.object(lean_mode, "count_claude_prompt_tokens", side_effect=failure): + decision = lean.prepare(body) + 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_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"}}}}] + 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(): + from services import native_token_count + + captured = {} + + class CountResponse: + status_code = 200 + + def iter_content(self, chunk_size): + yield b'{"input_tokens": 1234, "context_management": {"original_input_tokens": 5678}}' + + 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, 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 + 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, 150_000)) + + def large_input(self): + return "x" * 200_000 + + def test_over_limit_messages_get_anthropic_prompt_too_long(self): + 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 > 99000 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 99000 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..bda36bca 100644 --- a/worker/container-env.mjs +++ b/worker/container-env.mjs @@ -129,6 +129,17 @@ const DIRECT_ENV_KEYS = [ "MAX_OUTPUT_TOKENS", "OPTIMIZER_MAX_REQUEST_BYTES", "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", + "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",