"""What the proxy writes must outlive the translation that follows it. Claude Code receives an Anthropic-shaped response, which has nowhere to put OpenAI's cached_tokens, cache_write_tokens or reasoning_tokens. If those are not captured before the translation, the only remaining record of them is a bill. """ from __future__ import annotations import contextlib import json from pathlib import Path from types import SimpleNamespace import pytest from workflow_bench import litellm_usage_callback, provider_usage from workflow_bench.litellm_usage_callback import USAGE_LOG_ENV_VAR, ProviderUsageLogger from workflow_bench.model_gateway import ( OpenAIGateway, USAGE_CALLBACK_MODULE, openai_litellm_config, write_openai_litellm_config, ) from workflow_bench.provider_usage import ( ANTHROPIC, LITELLM_NORMALIZED, USAGE_ENV_VARS, normalize_usage, ) class _Usage: """Stands in for the provider usage model LiteLLM hands the callback.""" def __init__(self, payload: dict) -> None: self._payload = payload def model_dump(self) -> dict: return self._payload def _openai_response(usage: dict) -> SimpleNamespace: return SimpleNamespace( id="resp_68f2c1", # The model that actually answered, which is not the role the caller asked for. model="gpt-5.6-sol-2026-08-01", usage=_Usage(usage), ) # The shape a callback actually receives: LiteLLM normalises usage into its own # Chat-Completions-style object before any logger sees it, so an OpenAI reply # arrives as prompt_tokens / prompt_tokens_details. Confirmed against a real # proxy in tests/test_mock_provider.py; a fixture in the wire shape would test # an object this code path never gets. NATIVE = { "prompt_tokens": 48_000, "prompt_tokens_details": {"cached_tokens": 44_000, "cache_write_tokens": 1_000}, "completion_tokens": 900, "completion_tokens_details": {"reasoning_tokens": 640}, } @pytest.fixture def logged(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): log = tmp_path / "provider_usage.jsonl" monkeypatch.setenv(USAGE_LOG_ENV_VAR, str(log)) def emit(usage: dict) -> dict: ProviderUsageLogger()._append( "success", {"model": "claude-sonnet-4-5", "custom_llm_provider": "openai", "call_type": "responses"}, _openai_response(usage), 0.0, 1.0, ) return json.loads(log.read_text().splitlines()[-1]) return emit def test_native_openai_usage_survives_the_anthropic_translation(logged) -> None: event = logged(NATIVE) native = event["native_usage"] # Verbatim: the fields an Anthropic-shaped response cannot carry. assert native["prompt_tokens_details"]["cached_tokens"] == 44_000 assert native["prompt_tokens_details"]["cache_write_tokens"] == 1_000 assert native["completion_tokens_details"]["reasoning_tokens"] == 640 assert event["response_id"] == "resp_68f2c1" def test_the_actual_model_is_recorded_separately_from_the_requested_role(logged) -> None: """Pricing must follow what answered, not what the caller named.""" event = logged(NATIVE) assert event["requested_model"] == "claude-sonnet-4-5" assert event["actual_model"] == "gpt-5.6-sol-2026-08-01" assert "cell_id" not in event, "a proxy-wide variable cannot identify a cell" def test_the_captured_event_normalizes_with_openai_arithmetic(logged) -> None: """Capture and normalization must agree end to end, not just in isolation.""" event = logged(NATIVE) # The provider the LOG recorded, not one the test supplies - passing # OPENAI_RESPONSES by hand here is what hid the adapter-key mismatch. # LITELLM_NORMALIZED, not OPENAI_RESPONSES: a proxy callback never sees the # upstream body. Measured against a real gateway - the Responses adapter # found none of its keys there and reported every field unknown. assert event["provider"] == LITELLM_NORMALIZED assert event["provider_label"] == "openai" usage = normalize_usage(event["provider"], event["native_usage"]) assert usage.total_input_tokens == 48_000 assert usage.ordinary_input_tokens == 3_000 assert usage.cache_read_input_tokens == 44_000 assert usage.complete def test_usage_without_details_normalizes_to_unknown_rather_than_zero(logged) -> None: """The mutation the accounting must not survive: dropped details, silent zeros.""" stripped = {k: v for k, v in NATIVE.items() if k != "prompt_tokens_details"} event = logged(stripped) usage = normalize_usage(event["provider"], event["native_usage"]) assert usage.cache_read_input_tokens is None assert usage.ordinary_input_tokens is None assert not usage.complete def test_a_failed_request_is_still_accounted_for(logged, tmp_path: Path) -> None: """The money was spent whether or not the cell produced an artifact.""" import asyncio logger = ProviderUsageLogger() args = ({"model": "claude-sonnet-4-5"}, _openai_response(NATIVE), 0.0, 1.0) # Every hook LiteLLM can call, not the private helper underneath them: the # sync failure hook was missing entirely and _append could never show that. logger.log_failure_event(*args) asyncio.run(logger.async_log_failure_event(*args)) events = [json.loads(line) for line in (tmp_path / "provider_usage.jsonl").read_text().splitlines()] assert len(events) == 2, "both failure hooks must record" assert all(e["status"] == "failure" for e in events) def test_every_public_outcome_hook_records(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """Overriding a subset silently drops whichever path LiteLLM actually uses.""" import asyncio monkeypatch.setenv(USAGE_LOG_ENV_VAR, str(tmp_path / "usage.jsonl")) logger = ProviderUsageLogger() args = ({"model": "m"}, _openai_response(NATIVE), 0.0, 1.0) logger.log_success_event(*args) logger.log_failure_event(*args) asyncio.run(logger.async_log_success_event(*args)) asyncio.run(logger.async_log_failure_event(*args)) events = [json.loads(line) for line in (tmp_path / "usage.jsonl").read_text().splitlines()] assert [e["status"] for e in events] == ["success", "failure", "success", "failure"] def test_the_logger_never_raises_into_the_proxy(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """Accounting is evidence, not control flow.""" monkeypatch.setenv(USAGE_LOG_ENV_VAR, str(tmp_path / "missing-dir" / "usage.jsonl")) ProviderUsageLogger()._append("success", {}, object(), 0.0, 1.0) def test_no_log_is_written_when_the_destination_is_unset(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv(USAGE_LOG_ENV_VAR, raising=False) ProviderUsageLogger()._append("success", {}, _openai_response(NATIVE), 0.0, 1.0) assert not list(tmp_path.iterdir()) def test_the_generated_config_loads_the_callback_from_beside_itself(tmp_path: Path) -> None: """LiteLLM resolves the dotted path relative to the config directory.""" config = write_openai_litellm_config(tmp_path / "litellm.yaml", ["gpt-5.6-sol"]) assert openai_litellm_config(["gpt-5.6-sol"])["litellm_settings"]["callbacks"] == [ f"{USAGE_CALLBACK_MODULE}.handler" ] installed = config.parent / f"{USAGE_CALLBACK_MODULE}.py" assert installed.is_file(), "the proxy cannot import a callback that was never placed" # Importing it, not grepping it: a text search passes even when the module # cannot load, which is exactly how a package-relative import survived # review here. This is the deployment configuration, so load it the way the # proxy does - by path, as a top-level module. import importlib.util spec = importlib.util.spec_from_file_location(USAGE_CALLBACK_MODULE, installed) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) assert isinstance(module.handler, module.ProviderUsageLogger) def test_the_gateway_forwards_the_usage_environment_into_the_proxy( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: """The proxy is a separate process with a constructed environment. Popen(env=...) replaces the parent environment rather than extending it, so a variable the callback reads is simply absent unless the gateway forwards it by name. Without this the accounting looks configured and silently records nothing on every request - the in-process tests above cannot see that, because they never cross the subprocess boundary. """ for name in USAGE_ENV_VARS: monkeypatch.setenv(name, f"value-for-{name}") captured: dict[str, dict[str, str]] = {} class _Popen: def __init__(self, *_a, **kwargs): captured["env"] = kwargs["env"] raise RuntimeError("stop before launching a real proxy") # The console-script resolver runs before Popen and is absent in this # environment (the same reason two gateway tests fail here); the argv it # builds is not what this test is about. monkeypatch.setattr( "workflow_bench.model_gateway.litellm_proxy_argv", lambda **_k: ["/bin/true"], ) monkeypatch.setattr("workflow_bench.model_gateway.subprocess.Popen", _Popen) gateway = OpenAIGateway( openai_api_key="sk-test", model_names=["gpt-5.6-sol"], work_dir=tmp_path, ) with contextlib.suppress(Exception): gateway.__enter__() env = captured.get("env") assert env is not None, "the proxy was never constructed" for name in USAGE_ENV_VARS: assert env.get(name) == f"value-for-{name}", f"{name} never reached the proxy" # The credential allowlist is still an allowlist, not the parent environment. assert "PATH" in env and len(env) < 40 def test_an_unresolvable_provider_is_refused_rather_than_guessed() -> None: """LiteLLM says "openai" for Chat Completions too, and it counts differently.""" from workflow_bench.provider_usage import canonical_provider # Every openai call reaching this callback has already been normalised by # LiteLLM, whatever endpoint it used - the observed call_type for a Claude # Code request through the gateway is "anthropic_messages". The adapter has # to match the object in hand, not the protocol on the wire. assert canonical_provider("openai", "responses") == LITELLM_NORMALIZED assert canonical_provider("openai", "anthropic_messages") == LITELLM_NORMALIZED assert canonical_provider("anthropic", "completion") == ANTHROPIC # An unrecognised provider is still refused rather than guessed. assert canonical_provider("some-new-provider", "responses") is None assert canonical_provider(None, None) is None def test_request_identity_cannot_come_from_the_proxy_environment() -> None: """One proxy serves the whole sweep, so its environment identifies the sweep. attach_openai_gateway wraps all of _run_sweep, and cells run concurrently under --workers, interleaving requests through that single process. Any variable forwarded at launch is therefore constant for every event it ever records. Pinned so a future change does not reintroduce a per-cell environment variable that would silently stamp one value on every request. """ assert USAGE_ENV_VARS == ( "GITNEXUS_BENCH_PROVIDER_USAGE", "GITNEXUS_BENCH_SWEEP_ID", ), "a per-cell variable here would be constant across concurrent cells" def test_a_request_records_its_session_so_attribution_stays_possible(logged) -> None: """The per-request half of identity, recorded even when the provider omits it.""" event = logged(NATIVE) assert "session_id" in event, "absent attribution is still a fact about the run" def test_the_callback_imports_the_way_litellm_actually_loads_it(tmp_path: Path) -> None: """By path, as a top-level module, with no parent package and no sys.path entry. LiteLLM resolves a dotted callback through spec_from_file_location against the config directory, so the copied file is not part of workflow_bench when it runs. A relative or sibling import therefore raises ImportError and the proxy exits before becoming ready - which the in-package tests cannot see, because they import it as workflow_bench.litellm_usage_callback. """ import importlib.util import shutil source = Path(litellm_usage_callback.__file__) installed = tmp_path / f"{USAGE_CALLBACK_MODULE}.py" shutil.copy(source, installed) spec = importlib.util.spec_from_file_location(USAGE_CALLBACK_MODULE, installed) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) # ImportError here is the proxy refusing to start assert hasattr(module, "handler") def test_the_callbacks_copied_constants_match_the_canonical_ones() -> None: """The copies are deliberate; drifting apart silently is not. The callback cannot import from the package (see the test above), so it carries its own literals. These assertions are what keep the duplication honest. """ assert litellm_usage_callback.USAGE_LOG_ENV_VAR == provider_usage.USAGE_LOG_ENV_VAR assert litellm_usage_callback.SWEEP_ID_ENV_VAR == provider_usage.SWEEP_ID_ENV_VAR for label, call_type in ( ("openai", "responses"), ("openai", "completion"), ("openai", None), ("anthropic", "completion"), ("mystery", "responses"), ): assert litellm_usage_callback.canonical_provider(label, call_type) == provider_usage.canonical_provider( label, call_type ), f"resolver drifted for {label!r}/{call_type!r}"