From f556862b867896fa184dd94c1b2100c7b7f67201 Mon Sep 17 00:00:00 2001 From: Brent Rager Date: Wed, 8 Jul 2026 15:47:24 -0400 Subject: [PATCH] th-1cc9fa: model-output ceiling clamp + raise starvation defaults (Python server parity) Python-server parity for the Rust server hardening (th-1cc9fa / th-562b6d): - Raise the anti-starvation turn sizing the runner applies: max_tokens 512 -> 8192 (DEFAULT_MAX_TOKENS), max_iterations 6 -> 20 (DEFAULT_MAX_ITERATIONS). The old chat-widget sizing starves reasoning models (empty replies / truncated tool loops). - New model_info.py: best-effort per-model output ceiling from the gateway's /model/info (map_model_info + model_output_ceiling), cached once per process, stdlib fetch, no extra deps. No key / any error / unknown model -> None -> unclamped. - Thread the resolved ceiling into AgentOptions.model_max_output so the engine clamps max_tokens to what the model can emit. Feature-detected against the pinned published core (which predates the field) so this stays green now and activates automatically once a core >= 1.3.2 is pinned. Tests: test_model_info.py (map + ceiling lookup: clamp/unknown/no-key/error/cache/url) and test_starvation_defaults.py (raised constants, runner sends 8192, iteration cap 20, + a forward-compat clamp test that skips on the pre-clamp core). 47 passed, 1 skipped. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/smooth_operator_server/model_info.py | 122 +++++++++++++++ .../src/smooth_operator_server/turn_runner.py | 41 ++++- python/server/tests/test_model_info.py | 141 ++++++++++++++++++ .../server/tests/test_starvation_defaults.py | 71 +++++++++ 4 files changed, 374 insertions(+), 1 deletion(-) create mode 100644 python/server/src/smooth_operator_server/model_info.py create mode 100644 python/server/tests/test_model_info.py create mode 100644 python/server/tests/test_starvation_defaults.py diff --git a/python/server/src/smooth_operator_server/model_info.py b/python/server/src/smooth_operator_server/model_info.py new file mode 100644 index 0000000..358a434 --- /dev/null +++ b/python/server/src/smooth_operator_server/model_info.py @@ -0,0 +1,122 @@ +"""Best-effort per-model **output-token ceiling** from the gateway's ``/model/info``. + +A budget/policy ``max_tokens`` can exceed what a model can physically emit — a +reasoning model then burns the whole budget on reasoning and returns EMPTY, or the +upstream 400s (e.g. ``groq-compound`` caps output at 8192). The actual clamp lives +in the engine (``smooth_operator_core.effective_max_tokens``); this module *sources* +the ceiling the turn runner threads into ``AgentOptions.model_max_output``. + +Mirrors the Rust server's ``admin.rs`` ``map_model_info`` / ``model_output_ceiling`` +(EPIC th-1cc9fa). Fetched at most once per process (module cache), reusing the same +gateway creds the turns use. **Best-effort**: no gateway key, any transport error, +an unknown model, or a model whose gateway entry has no positive ceiling ⇒ ``None`` +⇒ the engine leaves ``max_tokens`` unclamped (graceful passthrough, no behaviour +change). + +Zero extra runtime deps: the fetch uses ``urllib`` from the stdlib in a worker +thread. ``ponytail:`` stdlib GET is enough — the only consumer is one integer per +model; no need to pull in httpx just for this. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import urllib.request +from typing import Any, Awaitable, Callable + +#: Default OpenAI-compatible gateway (matches the Rust server's ``DEFAULT_GATEWAY_URL``). +DEFAULT_GATEWAY_URL = "https://llm.smoo.ai/v1" + +#: A seam for tests: given ``(url, key)`` return the parsed ``/model/info`` JSON, or +#: raise. Production uses :func:`_default_fetch`; tests inject a stub so the ceiling +#: path is exercised with no network. +Fetcher = Callable[[str, str | None], Awaitable[dict[str, Any]]] + +#: Process-wide cache of ``{model_name: max_output_tokens}``. ``None`` until the first +#: successful fetch — a failed/keyless attempt is NOT cached, so the next turn retries +#: (mirrors the Rust ``model_output_ceiling`` "cache success only" behaviour). +_cache: dict[str, int] | None = None + + +def map_model_info(payload: Any) -> dict[str, int]: + """Map a gateway ``/model/info`` payload + (``{ data: [{ model_name, model_info: { max_output_tokens, ... } }] }``) to + ``{model_name: max_output_tokens}``, keeping ONLY models that report a positive + integer ceiling. Entries without a name or a usable ceiling are skipped. Pure + + network-free so it's unit-testable on a sample payload. + + Unlike the Rust ``map_model_info`` (which also surfaces cost/tier/useCases for the + ``/admin/model-costs`` badge), the Python server has no model-costs route — the + only consumer is the ceiling clamp, so this maps just the ceiling. ``ponytail:`` + map what's read; add cost/tier here if a ``/admin/model-costs`` route ever lands.""" + out: dict[str, int] = {} + data = payload.get("data") if isinstance(payload, dict) else None + if not isinstance(data, list): + return out + for entry in data: + if not isinstance(entry, dict): + continue + name = entry.get("model_name") + if not isinstance(name, str) or not name: + continue + info = entry.get("model_info") + raw = info.get("max_output_tokens") if isinstance(info, dict) else None + # bool is an int subclass — reject True/False so a stray boolean ceiling + # never sneaks in as 1/0. + if isinstance(raw, bool) or not isinstance(raw, int): + continue + if raw > 0: + out[name] = raw + return out + + +def _default_fetch(url: str, key: str | None) -> dict[str, Any]: + """Blocking ``GET {url}`` with an optional bearer, parsed as JSON. Run via + :func:`asyncio.to_thread` so it never blocks the event loop.""" + req = urllib.request.Request(url) # noqa: S310 — fixed https gateway URL, not user input + if key: + req.add_header("Authorization", f"Bearer {key}") + with urllib.request.urlopen(req, timeout=10) as resp: # noqa: S310 + return json.loads(resp.read().decode("utf-8")) + + +async def model_output_ceiling( + model: str, + *, + gateway_url: str | None = None, + gateway_key: str | None = None, + fetch: Fetcher | None = None, +) -> int | None: + """The ``model``'s hard output ceiling (``max_output_tokens``) from the gateway, + or ``None`` when unknown. Threaded into ``AgentOptions.model_max_output`` so the + engine clamps ``max_tokens`` to what the model can emit (EPIC th-1cc9fa). + + Gateway URL/key default to ``SMOOAI_GATEWAY_URL`` / ``SMOOAI_GATEWAY_KEY`` (the + same creds :func:`smooth_operator_server.server._build_gateway_client` uses). + **No key ⇒ ``None`` with no network call** — a keyless server runs on a mock + client (tests) and never needs a live ceiling. Any fetch error ⇒ ``None`` and the + failure is not cached (next turn retries).""" + global _cache + if _cache is None: + key = gateway_key if gateway_key is not None else os.environ.get("SMOOAI_GATEWAY_KEY") + if not key: + return None # keyless → nothing to clamp against; skip the fetch entirely + base = gateway_url or os.environ.get("SMOOAI_GATEWAY_URL") or DEFAULT_GATEWAY_URL + url = f"{base.rstrip('/')}/model/info" + fetcher = fetch or (lambda u, k: asyncio.to_thread(_default_fetch, u, k)) + try: + payload = await fetcher(url, key) + except Exception: + return None # best-effort: gateway/transport/decode error ⇒ unclamped + _cache = map_model_info(payload) + ceiling = _cache.get(model) + return ceiling if isinstance(ceiling, int) and ceiling > 0 else None + + +def reset_cache() -> None: + """Drop the process-wide ceiling cache. For tests that exercise multiple fetch + outcomes; production fetches once and keeps it.""" + global _cache + _cache = None diff --git a/python/server/src/smooth_operator_server/turn_runner.py b/python/server/src/smooth_operator_server/turn_runner.py index feb24bb..afca254 100644 --- a/python/server/src/smooth_operator_server/turn_runner.py +++ b/python/server/src/smooth_operator_server/turn_runner.py @@ -13,7 +13,7 @@ from __future__ import annotations import json -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import Any, Callable from smooth_operator_core import ( @@ -31,11 +31,37 @@ from . import protocol from .confirmation import ConfirmationRegistry +from .model_info import model_output_ceiling from .session_store import MessageDirection, SessionStore #: Max prior turns replayed into the thread for memory (bounds context growth). MAX_PRIOR_MESSAGES = 50 +#: The engine's default model when the server pins none (matches +#: ``AgentOptions.model`` in smooth-operator-core). Used to look up the output +#: ceiling for the model the turn will actually send. +DEFAULT_MODEL = "claude-haiku-4-5" + +#: Per-call ``max_tokens`` sent to the gateway. Raised from the old chat-widget +#: default of 512 — that STARVES reasoning models (they exhaust the budget on +#: reasoning and return empty). Mirrors the Rust server's ``DEFAULT_MAX_TOKENS`` +#: 512→8192 (EPIC th-1cc9fa). The engine still clamps this DOWN to the model's real +#: output ceiling (:func:`model_output_ceiling`), so raising it is safe. +DEFAULT_MAX_TOKENS = 8192 + +#: Agent-loop iteration cap per turn. Raised from 6 — a tool-using reasoning turn +#: routinely needs more than six model round-trips. Mirrors the Rust server's +#: ``DEFAULT_MAX_ITERATIONS`` 6→20 (EPIC th-1cc9fa). +DEFAULT_MAX_ITERATIONS = 20 + +#: Whether the installed smooth-operator-core supports the ceiling clamp field. +#: The server pins the PUBLISHED core (see ``pyproject.toml``); a core predating the +#: clamp has no ``model_max_output`` on ``AgentOptions``, so passing it would raise +#: ``TypeError``. Feature-detect: thread the ceiling only when the field exists, so +#: this stays green on the pinned core and activates automatically once a core with +#: the clamp is released. (Release-ordering: core PR ships + publishes first.) +_CORE_SUPPORTS_CEILING = any(f.name == "model_max_output" for f in fields(AgentOptions)) + #: Top-K knowledge hits surfaced as auto-context citations (what grounded the #: answer). Matches the engine's auto-context injection and the TS/C#/Rust servers. AUTO_CONTEXT_LIMIT = 3 @@ -120,12 +146,25 @@ async def run( options_kwargs: dict[str, Any] = { "instructions": self._system_prompt, "knowledge": self._knowledge, + # Raised, anti-starvation sizing (see the module constants). The engine + # clamps max_tokens DOWN to the model's real output ceiling below. + "max_tokens": DEFAULT_MAX_TOKENS, + "max_iterations": DEFAULT_MAX_ITERATIONS, } if self._tools: options_kwargs["tools"] = self._tools if self._model is not None: options_kwargs["model"] = self._model + # Clamp max_tokens to what the resolved model can physically emit: look up + # its output ceiling from the gateway (best-effort; None ⇒ unclamped) and + # thread it into the engine's clamp. Feature-detected against the pinned + # core (see `_CORE_SUPPORTS_CEILING`). + if _CORE_SUPPORTS_CEILING: + ceiling = await model_output_ceiling(self._model or DEFAULT_MODEL) + if ceiling is not None: + options_kwargs["model_max_output"] = ceiling + # Write-confirmation HITL: when configured with tool patterns AND a registry # is present, install a HumanGate that parks the turn before a gated tool # runs (emit `write_confirmation_required`, await the client's verdict via diff --git a/python/server/tests/test_model_info.py b/python/server/tests/test_model_info.py new file mode 100644 index 0000000..382346c --- /dev/null +++ b/python/server/tests/test_model_info.py @@ -0,0 +1,141 @@ +"""Best-effort model output-token ceiling from the gateway's ``/model/info``. + +Mirrors the Rust server's ``map_model_info`` / ``model_output_ceiling`` tests in +``admin.rs``. The ceiling is what the turn runner threads into the engine's clamp +(EPIC th-1cc9fa). +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from smooth_operator_server import model_info + + +@pytest.fixture(autouse=True) +def _clear_cache(): + model_info.reset_cache() + yield + model_info.reset_cache() + + +# ── map_model_info (pure) ───────────────────────────────────────────────────── + +SAMPLE_PAYLOAD: dict[str, Any] = { + "data": [ + {"model_name": "claude-haiku-4-5", "model_info": {"max_output_tokens": 8192, "input_cost_per_token": 1e-6}}, + {"model_name": "big-context", "model_info": {"max_output_tokens": 65536}}, + {"model_name": "no-ceiling", "model_info": {"input_cost_per_token": 2e-6}}, # no max_output_tokens + ] +} + + +def test_map_extracts_positive_ceilings(): + out = model_info.map_model_info(SAMPLE_PAYLOAD) + assert out == {"claude-haiku-4-5": 8192, "big-context": 65536} + + +def test_map_skips_missing_name_and_nonpositive_and_nonint(): + payload = { + "data": [ + {"model_info": {"max_output_tokens": 100}}, # no model_name → skip + {"model_name": "zero", "model_info": {"max_output_tokens": 0}}, # 0 → skip + {"model_name": "neg", "model_info": {"max_output_tokens": -5}}, # negative → skip + {"model_name": "boolish", "model_info": {"max_output_tokens": True}}, # bool → skip + {"model_name": "floaty", "model_info": {"max_output_tokens": 1.5}}, # float → skip + {"model_name": "bare", "model_info": {}}, # no field → skip + {"model_name": "noinfo"}, # no model_info → skip + ] + } + assert model_info.map_model_info(payload) == {} + + +def test_map_tolerates_garbage_payloads(): + assert model_info.map_model_info({}) == {} + assert model_info.map_model_info({"data": "nope"}) == {} + assert model_info.map_model_info({"data": [None, 42, "x"]}) == {} + assert model_info.map_model_info(None) == {} + assert model_info.map_model_info([]) == {} + + +# ── model_output_ceiling (async, cached, best-effort) ───────────────────────── + + +async def _fetch_sample(url: str, key: str | None) -> dict[str, Any]: + return SAMPLE_PAYLOAD + + +@pytest.mark.asyncio +async def test_ceiling_looked_up_for_known_model(): + ceiling = await model_info.model_output_ceiling("big-context", gateway_key="sk-x", fetch=_fetch_sample) + assert ceiling == 65536 + + +@pytest.mark.asyncio +async def test_ceiling_none_for_unknown_model(): + ceiling = await model_info.model_output_ceiling("who-dis", gateway_key="sk-x", fetch=_fetch_sample) + assert ceiling is None + + +@pytest.mark.asyncio +async def test_no_key_skips_fetch_and_returns_none(): + called = False + + async def _boom(url: str, key: str | None) -> dict[str, Any]: + nonlocal called + called = True + raise AssertionError("must not fetch without a key") + + ceiling = await model_info.model_output_ceiling("big-context", gateway_key="", fetch=_boom) + assert ceiling is None + assert called is False + + +@pytest.mark.asyncio +async def test_fetch_error_returns_none_and_does_not_cache(): + attempts = 0 + + async def _flaky(url: str, key: str | None) -> dict[str, Any]: + nonlocal attempts + attempts += 1 + raise RuntimeError("gateway down") + + assert await model_info.model_output_ceiling("big-context", gateway_key="sk-x", fetch=_flaky) is None + # A failed fetch is NOT cached → the next call retries (mirrors the Rust behaviour). + assert await model_info.model_output_ceiling("big-context", gateway_key="sk-x", fetch=_flaky) is None + assert attempts == 2 + + +@pytest.mark.asyncio +async def test_success_is_cached_fetch_runs_once(): + calls = 0 + + async def _count(url: str, key: str | None) -> dict[str, Any]: + nonlocal calls + calls += 1 + return SAMPLE_PAYLOAD + + a = await model_info.model_output_ceiling("claude-haiku-4-5", gateway_key="sk-x", fetch=_count) + b = await model_info.model_output_ceiling("big-context", gateway_key="sk-x", fetch=_count) + assert a == 8192 + assert b == 65536 + assert calls == 1 # cached across models after the first success + + +@pytest.mark.asyncio +async def test_ceiling_url_uses_model_info_endpoint(): + seen: dict[str, Any] = {} + + async def _capture(url: str, key: str | None) -> dict[str, Any]: + seen["url"] = url + seen["key"] = key + return SAMPLE_PAYLOAD + + await model_info.model_output_ceiling( + "big-context", gateway_url="https://llm.smoo.ai/v1/", gateway_key="sk-x", fetch=_capture + ) + # Trailing slash trimmed; `/model/info` appended. + assert seen["url"] == "https://llm.smoo.ai/v1/model/info" + assert seen["key"] == "sk-x" diff --git a/python/server/tests/test_starvation_defaults.py b/python/server/tests/test_starvation_defaults.py new file mode 100644 index 0000000..2510e30 --- /dev/null +++ b/python/server/tests/test_starvation_defaults.py @@ -0,0 +1,71 @@ +"""The server's raised anti-starvation turn sizing (EPIC th-1cc9fa). + +The old chat-widget defaults (max_tokens=512, max_iterations=6) STARVE reasoning +models — they exhaust the token budget on reasoning and return empty, or run out of +loop iterations mid-tool-use. Raised to 8192 / 20, mirroring the Rust server's +``DEFAULT_MAX_TOKENS`` / ``DEFAULT_MAX_ITERATIONS``. The engine still clamps +``max_tokens`` DOWN to each model's real output ceiling, so raising it is safe. +""" + +from __future__ import annotations + +import pytest +from smooth_operator_core import MockLlmProvider + +from smooth_operator_server import turn_runner as tr +from smooth_operator_server.session_store import InMemorySessionStore +from smooth_operator_server.turn_runner import ( + DEFAULT_MAX_ITERATIONS, + DEFAULT_MAX_TOKENS, + TurnRunner, +) + + +def test_defaults_are_raised(): + assert DEFAULT_MAX_TOKENS == 8192 + assert DEFAULT_MAX_ITERATIONS == 20 + + +@pytest.mark.asyncio +async def test_turn_sends_raised_max_tokens(): + mock = MockLlmProvider() + mock.push_text("hello") + runner = TurnRunner(chat_client=mock, store=InMemorySessionStore()) + await runner.run(conversation_id="c1", request_id="r1", user_message="hi", sink=lambda _e: None) + assert mock.last_call is not None + # No gateway key in the test env ⇒ no ceiling ⇒ the raised default is sent as-is. + assert mock.last_call.kwargs["max_tokens"] == DEFAULT_MAX_TOKENS + + +@pytest.mark.asyncio +async def test_iteration_cap_is_raised(): + # Script MORE tool-calls than the cap; the loop must stop at DEFAULT_MAX_ITERATIONS + # (proving the cap is 20, not the engine's old 6/8). Each references an unknown + # tool, so every iteration makes exactly one model call then loops. + mock = MockLlmProvider() + for i in range(DEFAULT_MAX_ITERATIONS + 5): + mock.push_tool_call(f"call-{i}", "nonexistent_tool", "{}") + runner = TurnRunner(chat_client=mock, store=InMemorySessionStore()) + await runner.run(conversation_id="c2", request_id="r2", user_message="loop", sink=lambda _e: None) + assert mock.call_count == DEFAULT_MAX_ITERATIONS + + +@pytest.mark.skipif( + not tr._CORE_SUPPORTS_CEILING, + reason="installed smooth-operator-core predates the model_max_output clamp (release-ordering: core ships first)", +) +@pytest.mark.asyncio +async def test_ceiling_clamps_raised_max_tokens_end_to_end(monkeypatch): + # Forward-compat: once a core WITH the clamp is installed, a resolved ceiling + # below the raised 8192 default is threaded into AgentOptions and the engine + # sends the clamped value. Skipped on the pinned (pre-clamp) core. + async def _fake_ceiling(model: str) -> int: + return 4096 + + monkeypatch.setattr(tr, "model_output_ceiling", _fake_ceiling) + mock = MockLlmProvider() + mock.push_text("clamped") + runner = TurnRunner(chat_client=mock, store=InMemorySessionStore()) + await runner.run(conversation_id="c3", request_id="r3", user_message="hi", sink=lambda _e: None) + assert mock.last_call is not None + assert mock.last_call.kwargs["max_tokens"] == 4096