From 18cbeb907c87c50d548b5c62102080750baac972 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gerg=C5=91=20Magyar?= Date: Tue, 8 Sep 2026 18:22:04 +0100 Subject: [PATCH 1/2] feat(eval): record provider-native usage at the gateway instead of inferring it after translation (#3220) * feat(eval): record provider-native usage at the gateway, not after translation The benchmark reads token counts out of Claude Code's session output, which is Anthropic-shaped whatever actually served the request. That holds until the upstream is OpenAI, because the two providers do not merely name their fields differently - they mean opposite things by them: Anthropic: total_input = input_tokens + cache_creation + cache_read (input_tokens is the UNCACHED remainder; cache fields ADD) OpenAI: total_input = input_tokens ordinary = input_tokens - cached - cache_write (input_tokens is the WHOLE; cache fields are SUBSETS) Adding OpenAI's three double-counts; subtracting Anthropic's under-counts. One shared struct cannot be right for both, so the seam goes at the gateway, on the far side of the translation: a LiteLLM callback appends each upstream request's usage verbatim, along with the model that actually answered, the response id and the cell it belongs to. Normalization is derived offline from that record, so the derivation can be revisited without re-running a paid sweep. Two rules the tests encode literally. The native object is authoritative. The callback stores it unflattened, unrenamed and unsummed. Reasoning tokens are kept as the decomposition of output tokens they are, not added to them a second time. A field nobody reported is unknown, never zero. A stored cache_read of 0 used to mean either "the provider said zero" or "our adapter never looked" - the first says caching is not working, the second says we cannot tell. NormalizedUsage therefore uses None, and refuses to compute the ordinary portion when a term is missing rather than subtracting an invented zero. Mutation-checked three ways. Giving OpenAI Anthropic's arithmetic fails four tests. Making unknown fall back to zero fails the unknown test. Dropping input_tokens_details in the callback fails the end-to-end accounting test with "assert None == 3000" - it goes unknown rather than passing with zeros, which was the point of the exercise. The actual model is recorded separately from the requested role because several Claude role names map onto one upstream model here; pricing must follow what answered. Cost is deliberately NOT stored: prices change, and tokens plus a versioned pricing table can answer both what a past run cost and what the same usage would cost today, without rewriting historical evidence. The callback never raises. A cell that fails still spent money upstream, and losing the accounting because a log write failed is the worse outcome. Failed requests are recorded too. No caching configuration, model, skill or promotion change: this installs the thermometer without altering the experiment. 538 eval tests pass plus 27 gateway tests; ruff clean. The two test_model_gateway.py failures are environmental - litellm[proxy]'s console script is absent in this venv - and predate this branch. * fix(eval): drop the accidentally committed .venv symlink I symlinked eval/.venv at a sibling worktree's virtualenv to avoid rebuilding it, and git add -A committed the symlink. .gitignore lists ".venv/" with a trailing slash, which matches a directory and not a symlink, so nothing stopped it. That broke eval / containment (windows), where uv then refused to create the environment: "failed to create directory eval\\.venv: Cannot create a file when that file already exists". A machine-specific absolute path had no business in the tree in the first place. Removed, and .gitignore now also lists the bare name so the same slip cannot repeat. * Address PR review feedback (#3220) Forward the usage environment into the proxy. This is the one that mattered: the callback returns immediately when GITNEXUS_BENCH_PROVIDER_USAGE is absent, the proxy runs as its own process, and Popen(env=...) REPLACES the parent environment rather than extending it. The gateway's allowlist carried the OpenAI and master keys and nothing else, so the callback loaded, found no destination, and silently recorded nothing on every request. The accounting looked configured and measured nothing at all. My tests could not see it. They set the variable in-process and called the logger directly, so none of them ever crossed the subprocess boundary the feature actually runs behind. The new test drives OpenAIGateway.__enter__ with Popen captured and asserts each variable reaches the child - and that the result is still an allowlist rather than the inherited parent environment, since forwarding by name is what keeps the credential boundary explicit. Resolve the provider label into an adapter key. The callback recorded LiteLLM's custom_llm_provider, which is "openai", while the adapter table is keyed "openai-responses" - so nothing the logger wrote could have been normalized. The end-to-end test hid this by passing OPENAI_RESPONSES by hand instead of using the provider the log recorded; it now uses the logged value, which is what makes the mismatch visible. The label alone cannot pick an adapter: LiteLLM reports "openai" for Chat Completions as well, and the two report usage differently. canonical_provider combines the label with the call type and returns None when it cannot resolve one, so normalize_usage refuses rather than guessing token semantics. Both are stored - provider_label is what LiteLLM said, provider is the adapter key. The shared env-var names moved into provider_usage.py so model_gateway can import them without importing litellm, which only the in-proxy callback needs. Mutation-checked. Removing the forwarding loop fails the gateway test; using the raw label as the adapter key fails two. 656 eval tests pass, ruff clean. The two test_model_gateway.py failures are the environmental ones - litellm[proxy]'s console script is absent here, which is also why the new test patches the argv builder to reach Popen at all. * fix(eval): stop recording a cell id the proxy cannot know Setting out to build the correlation this PR was missing - cell usage as the sum of its upstream requests - turned up that the field it would have been built on cannot hold what its name claims. attach_openai_gateway wraps the whole sweep (runner.py:2122), so ONE proxy serves every cell, and its environment is fixed for that process's lifetime. Cells run concurrently under --workers and interleave requests through it. A cell id forwarded at launch is therefore the same constant on every event the callback ever writes - not an attribution, just a label that looks like one. Worse than absent, because a reader would trust it. So GITNEXUS_BENCH_CELL_ID is gone rather than left to be wired up later. What remains is honest about its scope: sweep_id is genuinely sweep-wide, and session_id is the per-request half - the only thing that can attribute a request to a cell, since anything read from the environment is shared by all of them. It is recorded even when the provider supplies nothing, because knowing attribution is unavailable is itself a fact about the run. Pinned by a test asserting the forwarded set contains no per-cell variable, so a later change does not reintroduce one and quietly stamp a single value across concurrent cells. What this leaves open, stated plainly: per-cell attribution is NOT built, and cannot be until a per-request identifier is available. Whether Claude Code propagates a session identifier through the proxy is unverified - determining it needs a real session against the gateway, which is a paid run. Sweep-level totals and per-request cache ratios do not need it, and those are what the caching question actually turns on. 658 eval tests pass, ruff clean; the two test_model_gateway.py failures remain environmental. * fix(eval): keep the usage callback importable the way LiteLLM loads it CI caught a regression I introduced: "ImportError: Could not import handler from provider_usage_callback", and the proxy exited before becoming ready. Moving the shared constants into provider_usage.py, I imported them from the callback with "from .provider_usage import ...". But LiteLLM resolves a dotted callback through spec_from_file_location against the config directory, so the copied file runs as a top-level module with no parent package and no sys.path entry - the relative import raises and the gateway never starts. The module's own docstring says it is deliberately self-contained for exactly this reason, and I broke that invariant while tidying. The in-package tests could not see it. They import workflow_bench.litellm_usage_callback, where the relative import resolves fine; the failure only exists on the path where the file is copied and loaded standalone. The callback carries its own literals again. Two tests keep that honest: one loads the copied file the way LiteLLM does - by path, as a top-level module - so an import that only works in-package fails there, and one asserts the copied constants and the provider resolver still agree with the canonical copies in provider_usage.py, so the deliberate duplication cannot drift silently. Mutation-checked: restoring the relative import reproduces CI's exact error. 660 eval tests pass locally; the two remaining test_model_gateway.py failures are the environmental ones (litellm[proxy]'s console script is absent here, which is also why this never reproduced locally). * test(eval): import the installed callback instead of grepping it Two review findings on the same weakness, both correct. The install test asserted "class ProviderUsageLogger" appeared in the copied file's text. That passes whenever the string is present, including when the module cannot load at all - which is precisely how a package-relative import got through review here and took the proxy down. It now loads the copy the way LiteLLM does, by path as a top-level module, and checks the handler instance the config actually names. The gateway-forwarding test built its work directory with tempfile.mkdtemp(), which nothing removed, so every run left the generated config and the copied callback behind in the system temp directory. It uses the pytest-managed tmp_path fixture like its neighbours. 660 eval tests pass; the two test_model_gateway.py failures are the environmental ones. * fix(eval): record failures on the synchronous callback path too ProviderUsageLogger overrode both async hooks and the sync SUCCESS hook, but not the sync failure hook. On that path failures fell through to CustomLogger's base implementation and were never appended - so a sweep recorded its successes and quietly understated what it spent, since a failed request is billed all the same. That contradicts the module's own stated reason for handling failures at all. The failure test could not have caught it: it called _append directly, which exercises neither public hook. Both failure tests now drive the hooks LiteLLM actually calls, and a new one walks all four - sync and async, success and failure - asserting each records in order. Removing the sync failure hook fails both. 661 eval tests pass; the two test_model_gateway.py failures remain environmental. --------- Co-authored-by: Gergo Magyar --- eval/.gitignore | 1 + eval/tests/test_provider_usage.py | 132 ++++++++ eval/tests/test_provider_usage_capture.py | 313 ++++++++++++++++++ eval/workflow_bench/litellm_usage_callback.py | 136 ++++++++ eval/workflow_bench/model_gateway.py | 43 ++- eval/workflow_bench/provider_usage.py | 188 +++++++++++ 6 files changed, 812 insertions(+), 1 deletion(-) create mode 100644 eval/tests/test_provider_usage.py create mode 100644 eval/tests/test_provider_usage_capture.py create mode 100644 eval/workflow_bench/litellm_usage_callback.py create mode 100644 eval/workflow_bench/provider_usage.py diff --git a/eval/.gitignore b/eval/.gitignore index d1ac9f241..8f1814dfa 100644 --- a/eval/.gitignore +++ b/eval/.gitignore @@ -14,3 +14,4 @@ build/ # Environment .env .venv/ +.venv diff --git a/eval/tests/test_provider_usage.py b/eval/tests/test_provider_usage.py new file mode 100644 index 000000000..c8e74a966 --- /dev/null +++ b/eval/tests/test_provider_usage.py @@ -0,0 +1,132 @@ +"""The two providers' accounting equations, encoded literally. + +Adding OpenAI's cache fields to its input_tokens double-counts, because they are +subsets of it. Subtracting Anthropic's under-counts, because they are additional +categories. A single generic struct cannot be right for both, so these tests +pin each equation rather than the field names. +""" + +from __future__ import annotations + +import pytest + +from workflow_bench.provider_usage import ( + ANTHROPIC, + OPENAI_RESPONSES, + UsageSemanticsError, + normalize_usage, +) + + +def _openai(input_tokens: int, cached: int | None = None, cache_write: int | None = None) -> dict: + details: dict[str, int] = {} + if cached is not None: + details["cached_tokens"] = cached + if cache_write is not None: + details["cache_write_tokens"] = cache_write + return { + "input_tokens": input_tokens, + "input_tokens_details": details, + "output_tokens": 300, + "output_tokens_details": {"reasoning_tokens": 250}, + } + + +def test_openai_uncached_request_is_all_ordinary_input() -> None: + usage = normalize_usage(OPENAI_RESPONSES, _openai(1000, cached=0, cache_write=0)) + assert usage.ordinary_input_tokens == 1000 + assert usage.total_input_tokens == 1000 + assert (usage.cache_read_input_tokens, usage.cache_write_input_tokens) == (0, 0) + + +def test_openai_cache_creation_keeps_the_parts_summing_to_input_tokens() -> None: + """The subsets must reconstruct the whole, never exceed it.""" + + usage = normalize_usage(OPENAI_RESPONSES, _openai(1000, cached=0, cache_write=400)) + assert usage.ordinary_input_tokens == 600 + assert ( + usage.ordinary_input_tokens + + usage.cache_read_input_tokens + + usage.cache_write_input_tokens + == usage.total_input_tokens + ) + + +def test_openai_cache_hit_plus_new_write_uses_the_documented_subtraction() -> None: + usage = normalize_usage(OPENAI_RESPONSES, _openai(10_000, cached=7_000, cache_write=1_000)) + assert usage.ordinary_input_tokens == 2_000 + assert usage.total_input_tokens == 10_000, "input_tokens is the whole, not a component" + + +def test_openai_reasoning_tokens_decompose_output_rather_than_adding_to_it() -> None: + usage = normalize_usage(OPENAI_RESPONSES, _openai(100, cached=0, cache_write=0)) + assert usage.output_tokens == 300 + assert usage.reasoning_output_tokens == 250 + assert usage.reasoning_output_tokens <= usage.output_tokens + + +def test_anthropic_uncached_total_is_just_input_tokens() -> None: + usage = normalize_usage( + ANTHROPIC, + {"input_tokens": 1000, "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, "output_tokens": 200}, + ) + assert usage.total_input_tokens == 1000 + assert usage.ordinary_input_tokens == 1000 + + +def test_anthropic_cached_total_adds_the_cache_categories() -> None: + """The opposite equation to OpenAI's, on deliberately identical numbers.""" + + usage = normalize_usage( + ANTHROPIC, + {"input_tokens": 2_000, "cache_creation_input_tokens": 1_000, + "cache_read_input_tokens": 7_000, "output_tokens": 200}, + ) + assert usage.total_input_tokens == 10_000 + assert usage.ordinary_input_tokens == 2_000 + + +def test_the_same_numbers_mean_different_totals_on_the_two_providers() -> None: + """The whole reason a shared struct is unsafe, in one assertion.""" + + openai = normalize_usage(OPENAI_RESPONSES, _openai(10_000, cached=7_000, cache_write=1_000)) + anthropic = normalize_usage( + ANTHROPIC, + {"input_tokens": 10_000, "cache_creation_input_tokens": 1_000, + "cache_read_input_tokens": 7_000, "output_tokens": 300}, + ) + assert openai.total_input_tokens == 10_000 + assert anthropic.total_input_tokens == 18_000 + assert openai.ordinary_input_tokens == 2_000 + assert anthropic.ordinary_input_tokens == 10_000 + + +def test_missing_native_cache_fields_are_unknown_and_never_zero() -> None: + """A zero we invented is indistinguishable from a zero the provider reported.""" + + usage = normalize_usage(OPENAI_RESPONSES, {"input_tokens": 1000, "output_tokens": 10}) + assert usage.cache_read_input_tokens is None + assert usage.cache_write_input_tokens is None + assert usage.ordinary_input_tokens is None, "cannot subtract what was never reported" + assert usage.total_input_tokens == 1000 + assert not usage.complete + assert "cache_read_input_tokens" in usage.unknown_fields + + +def test_an_absent_usage_object_is_entirely_unknown() -> None: + usage = normalize_usage(ANTHROPIC, None) + assert not usage.complete + assert usage.total_input_tokens is None + + +def test_an_unknown_provider_is_refused_rather_than_guessed() -> None: + with pytest.raises(UsageSemanticsError, match="refusing to guess"): + normalize_usage("some-new-provider", {"input_tokens": 1}) + + +def test_cache_subsets_larger_than_the_whole_are_rejected() -> None: + """Nonsense arithmetic must surface, not silently produce a negative.""" + + with pytest.raises(UsageSemanticsError, match="exceed input_tokens"): + normalize_usage(OPENAI_RESPONSES, _openai(100, cached=90, cache_write=50)) diff --git a/eval/tests/test_provider_usage_capture.py b/eval/tests/test_provider_usage_capture.py new file mode 100644 index 000000000..99e542e4d --- /dev/null +++ b/eval/tests/test_provider_usage_capture.py @@ -0,0 +1,313 @@ +"""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, + OPENAI_RESPONSES, + 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), + ) + + +NATIVE = { + "input_tokens": 48_000, + "input_tokens_details": {"cached_tokens": 44_000, "cache_write_tokens": 1_000}, + "output_tokens": 900, + "output_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["input_tokens_details"]["cached_tokens"] == 44_000 + assert native["input_tokens_details"]["cache_write_tokens"] == 1_000 + assert native["output_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. + assert event["provider"] == OPENAI_RESPONSES + 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 != "input_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 + + assert canonical_provider("openai", "responses") == OPENAI_RESPONSES + assert canonical_provider("openai", "completion") is None + assert canonical_provider("openai", None) is None + assert canonical_provider("anthropic", "completion") == ANTHROPIC + + +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}" diff --git a/eval/workflow_bench/litellm_usage_callback.py b/eval/workflow_bench/litellm_usage_callback.py new file mode 100644 index 000000000..843413765 --- /dev/null +++ b/eval/workflow_bench/litellm_usage_callback.py @@ -0,0 +1,136 @@ +"""Append each upstream request's usage exactly as the provider reported it. + +This runs INSIDE the LiteLLM proxy, on the far side of the translation that +turns an OpenAI response into the Anthropic shape Claude Code expects. That is +the only point that still knows which provider served the request, what model +actually answered, and what the native usage object said before its fields were +renamed into someone else's semantics. + +Deliberately self-contained: the proxy loads this file by path from the config +directory, so it cannot assume ``workflow_bench`` is importable. Normalization +lives in workflow_bench.provider_usage and runs offline over what this writes - +the native object is the evidence, and deriving from it here would mean the +derivation could not be revisited without re-running a paid sweep. + +Never raises. A cell that fails still spent money upstream, and losing the +accounting because the log write failed would be the worse outcome. +""" + +from __future__ import annotations + +import json +import os +import threading +from typing import Any + +from litellm.integrations.custom_logger import CustomLogger + +# Literals, not imports. LiteLLM loads this file BY PATH from the config +# directory via spec_from_file_location, so it has no parent package and the +# directory is not on sys.path - a relative or sibling import raises +# ImportError and the proxy refuses to start. workflow_bench.provider_usage +# holds the canonical copies and a test asserts these agree with them, which +# catches drift without coupling at import time. +USAGE_LOG_ENV_VAR = "GITNEXUS_BENCH_PROVIDER_USAGE" +SWEEP_ID_ENV_VAR = "GITNEXUS_BENCH_SWEEP_ID" + + +def canonical_provider(label, call_type): # noqa: ANN001, ANN201 + """Adapter key for the usage shape, or None when it cannot be resolved. + + Mirrors workflow_bench.provider_usage.canonical_provider; see the note + above for why this is a copy rather than an import. + """ + + if label == "openai" and call_type and "responses" in call_type: + return "openai-responses" + if label == "anthropic": + return "anthropic" + return None + +SCHEMA_VERSION = 1 +_LOCK = threading.Lock() + + +def _plain(value: Any) -> Any: + """Provider usage arrives as pydantic models; keep the shape, drop the class.""" + + for attr in ("model_dump", "dict"): + method = getattr(value, attr, None) + if callable(method): + try: + return method() + except Exception: + pass + if isinstance(value, dict): + return value + return None + + +class ProviderUsageLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + self._append("success", kwargs, response_obj, start_time, end_time) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + # Failed requests are billed too, and a sweep that only accounts for + # successes understates what it spent. + self._append("failure", kwargs, response_obj, start_time, end_time) + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + self._append("success", kwargs, response_obj, start_time, end_time) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + # The synchronous counterpart. Overriding only the success hook here + # recorded successes and let failures fall through to the base class, + # which accounts for nothing - and a failed request is still billed, so + # a sweep missing them understates what it spent. + self._append("failure", kwargs, response_obj, start_time, end_time) + + def _append(self, status, kwargs, response_obj, start_time, end_time) -> None: # noqa: ANN001 + path = os.environ.get(USAGE_LOG_ENV_VAR) + if not path: + return + try: + params = kwargs.get("litellm_params") or {} + call_type = kwargs.get("call_type") + provider_label = kwargs.get("custom_llm_provider") or params.get("custom_llm_provider") + metadata = params.get("metadata") or {} + event = { + "schema_version": SCHEMA_VERSION, + "status": status, + # Identity. The REQUESTED model is the caller's role name and the + # ACTUAL model is what answered; pricing must follow the second, + # because several roles map onto one upstream model here. + "requested_model": kwargs.get("model"), + "actual_model": getattr(response_obj, "model", None), + # Two fields, because they answer different questions. The raw + # label is what LiteLLM said; "provider" is the adapter key, + # which needs the call type too - LiteLLM reports "openai" for + # both Chat Completions and Responses and those report usage + # differently. Unresolvable stays None so normalize_usage + # refuses rather than guessing token semantics. + "provider_label": provider_label, + "provider": canonical_provider(provider_label, call_type), + "response_id": getattr(response_obj, "id", None), + "call_type": call_type, + "sweep_id": os.environ.get(SWEEP_ID_ENV_VAR), + # The per-request half of identity, and the only thing that can + # attribute a request to a cell: one proxy serves the whole + # sweep, so anything read from the environment is the same for + # every event. Recorded even when absent, because knowing the + # attribution is unavailable is itself a fact about the run. + "session_id": metadata.get("litellm_session_id") or metadata.get("session_id"), + "started_at": str(start_time), + "completed_at": str(end_time), + # Verbatim. Not flattened, not renamed, not summed. + "native_usage": _plain(getattr(response_obj, "usage", None)), + } + line = json.dumps(event, default=str) + "\n" + with _LOCK, open(path, "a", encoding="utf-8") as handle: + handle.write(line) + except Exception: + # Accounting is evidence, not control flow: never take the sweep down. + return + + +handler = ProviderUsageLogger() diff --git a/eval/workflow_bench/model_gateway.py b/eval/workflow_bench/model_gateway.py index e44da55ec..e798eb14a 100644 --- a/eval/workflow_bench/model_gateway.py +++ b/eval/workflow_bench/model_gateway.py @@ -31,6 +31,8 @@ from typing import Any import yaml +from .provider_usage import USAGE_ENV_VARS + ANTHROPIC_API_KEY_ENV = "GITNEXUS_BENCH_ANTHROPIC_API_KEY" LEGACY_ANTHROPIC_API_KEY_ENV = "GITNEXUS_BENCH_AUTH_TOKEN" OPENAI_API_KEY_ENV = "GITNEXUS_BENCH_OPENAI_API_KEY" @@ -173,6 +175,9 @@ def resolve_model_access( return ModelAccess(start_proxy=False) +USAGE_CALLBACK_MODULE = "provider_usage_callback" + + def openai_litellm_config(model_names: Sequence[str]) -> dict[str, Any]: seen: list[str] = [] for name in model_names: @@ -196,7 +201,15 @@ def openai_litellm_config(model_names: Sequence[str]) -> dict[str, Any]: } for name in seen ], - "litellm_settings": {"request_timeout": GATEWAY_REQUEST_TIMEOUT_S}, + "litellm_settings": { + "request_timeout": GATEWAY_REQUEST_TIMEOUT_S, + # Captures each upstream request's usage as the provider reported + # it, before translation renames OpenAI's fields into Anthropic's + # shape and loses which arithmetic applies. Resolved by LiteLLM + # relative to the config directory, which is why the module is + # copied next to the config rather than imported from the package. + "callbacks": [f"{USAGE_CALLBACK_MODULE}.handler"], + }, "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}, } @@ -204,9 +217,26 @@ def openai_litellm_config(model_names: Sequence[str]) -> dict[str, Any]: def write_openai_litellm_config(path: Path, model_names: Sequence[str]) -> Path: path.write_text(yaml.safe_dump(openai_litellm_config(model_names), sort_keys=False)) path.chmod(0o600) + _install_usage_callback(path.parent) return path +def _install_usage_callback(config_dir: Path) -> Path: + """Place the usage logger where LiteLLM resolves callbacks from. + + LiteLLM loads a dotted callback path as a file relative to the config + directory before falling back to a package import, and the proxy runs as + its own process that need not have this package on sys.path. Copying the + one module is what makes the callback resolvable in both cases. + """ + + source = Path(__file__).with_name("litellm_usage_callback.py") + destination = config_dir / f"{USAGE_CALLBACK_MODULE}.py" + destination.write_text(source.read_text()) + destination.chmod(0o600) + return destination + + def _free_loopback_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -302,6 +332,17 @@ class OpenAIGateway(AbstractContextManager["OpenAIGateway"]): "OPENAI_API_KEY": self.openai_api_key, "LITELLM_MASTER_KEY": self.auth_token, } + # The proxy is a separate process and Popen(env=...) REPLACES the + # parent environment rather than extending it, so anything the usage + # callback reads has to be forwarded by name. Without this the callback + # loads, finds no destination, and returns silently on every request - + # the accounting looks configured and records nothing. Forwarded + # individually rather than by inheriting the environment, because the + # allowlist above is the gateway's credential boundary. + for name in USAGE_ENV_VARS: + value = os.environ.get(name) + if value: + env[name] = value if os.name == "nt": # Windows subprocess DLL/socket initialization needs SystemRoot. # Keep the rest of the gateway's credential boundary explicit. diff --git a/eval/workflow_bench/provider_usage.py b/eval/workflow_bench/provider_usage.py new file mode 100644 index 000000000..782eb18fa --- /dev/null +++ b/eval/workflow_bench/provider_usage.py @@ -0,0 +1,188 @@ +"""Per-request usage as the provider reported it, plus a derived cross-provider view. + +The benchmark has been reading token counts out of Claude Code's session output, +which is Anthropic-shaped whatever actually served the request. That works until +the upstream is OpenAI, because the two providers do not merely name their fields +differently - they mean opposite things by them: + + Anthropic: total_input = input_tokens + + cache_creation_input_tokens + + cache_read_input_tokens + (input_tokens is only the UNCACHED remainder; cache fields ADD) + + OpenAI: total_input = input_tokens + ordinary = input_tokens - cached_tokens - cache_write_tokens + (input_tokens is the WHOLE; cache fields are SUBSETS) + +Adding OpenAI's three together double-counts; subtracting Anthropic's +under-counts. So the native object is authoritative and is stored verbatim, and +the normalized view is derived from it per provider. + +The second rule is that a field nobody reported is UNKNOWN, not zero. A stored +``cache_read = 0`` previously could mean either "the provider said zero" or "our +adapter never looked", and those two must never be written identically again: +the first says caching is not working, the second says we cannot tell. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +SCHEMA_VERSION = 1 + +# Read by the in-proxy callback and forwarded by the gateway that launches it. +# Defined here because this module is pure stdlib: model_gateway can import the +# name without importing litellm, which only the callback needs. +# +# Both are SWEEP-scoped, and that is a constraint rather than an oversight. +# attach_openai_gateway wraps the whole sweep (runner.py), so one proxy serves +# every cell and its environment is fixed for that proxy's lifetime - while +# cells run concurrently under --workers and interleave requests through it. An +# environment variable therefore cannot carry a per-cell identity: it would +# record one constant against every event. Attributing a request to a cell +# needs an identifier that travels WITH the request; see the session fields the +# callback records for the intended hook. +USAGE_LOG_ENV_VAR = "GITNEXUS_BENCH_PROVIDER_USAGE" +SWEEP_ID_ENV_VAR = "GITNEXUS_BENCH_SWEEP_ID" +USAGE_ENV_VARS = (USAGE_LOG_ENV_VAR, SWEEP_ID_ENV_VAR) + +ANTHROPIC = "anthropic" +OPENAI_RESPONSES = "openai-responses" + + +class UsageSemanticsError(ValueError): + """The native usage object does not satisfy its own provider's arithmetic.""" + + +@dataclass(frozen=True) +class NormalizedUsage: + """Cross-provider view. ``None`` means the provider did not report it. + + Deliberately not defaulted to 0: see the module docstring. Every consumer + that sums these has to decide what to do about unknown, and making it None + forces that decision to be explicit instead of silently counting zero. + """ + + ordinary_input_tokens: int | None + cache_read_input_tokens: int | None + cache_write_input_tokens: int | None + total_input_tokens: int | None + output_tokens: int | None + reasoning_output_tokens: int | None + + @property + def complete(self) -> bool: + return all( + value is not None + for value in ( + self.ordinary_input_tokens, + self.cache_read_input_tokens, + self.cache_write_input_tokens, + self.total_input_tokens, + self.output_tokens, + ) + ) + + @property + def unknown_fields(self) -> tuple[str, ...]: + return tuple( + name for name, value in sorted(vars(self).items()) if value is None + ) + + +def _int_or_none(source: Mapping[str, Any] | None, key: str) -> int | None: + """Absent, null, or non-numeric all read as unknown rather than zero.""" + + if not isinstance(source, Mapping): + return None + value = source.get(key) + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + + +def _normalize_openai_responses(usage: Mapping[str, Any]) -> NormalizedUsage: + """input_tokens is the WHOLE; cached and cache-write are subsets of it.""" + + total = _int_or_none(usage, "input_tokens") + details = usage.get("input_tokens_details") + cache_read = _int_or_none(details, "cached_tokens") + cache_write = _int_or_none(details, "cache_write_tokens") + output_details = usage.get("output_tokens_details") + + ordinary: int | None = None + if total is not None and cache_read is not None and cache_write is not None: + ordinary = total - cache_read - cache_write + if ordinary < 0: + raise UsageSemanticsError( + f"OpenAI cached ({cache_read}) + cache_write ({cache_write}) " + f"exceed input_tokens ({total})" + ) + return NormalizedUsage( + ordinary_input_tokens=ordinary, + cache_read_input_tokens=cache_read, + cache_write_input_tokens=cache_write, + total_input_tokens=total, + output_tokens=_int_or_none(usage, "output_tokens"), + # A decomposition of output_tokens, not an addition to it. + reasoning_output_tokens=_int_or_none(output_details, "reasoning_tokens"), + ) + + +def _normalize_anthropic(usage: Mapping[str, Any]) -> NormalizedUsage: + """input_tokens is the uncached REMAINDER; the cache fields add to it.""" + + ordinary = _int_or_none(usage, "input_tokens") + cache_read = _int_or_none(usage, "cache_read_input_tokens") + cache_write = _int_or_none(usage, "cache_creation_input_tokens") + + total: int | None = None + if ordinary is not None and cache_read is not None and cache_write is not None: + total = ordinary + cache_read + cache_write + return NormalizedUsage( + ordinary_input_tokens=ordinary, + cache_read_input_tokens=cache_read, + cache_write_input_tokens=cache_write, + total_input_tokens=total, + output_tokens=_int_or_none(usage, "output_tokens"), + reasoning_output_tokens=None, + ) + + +def canonical_provider(label: str | None, call_type: str | None) -> str | None: + """Map LiteLLM's provider label onto an adapter key, or None if unsure. + + LiteLLM reports ``custom_llm_provider`` as "openai" for both Chat + Completions and Responses, and those two report usage differently, so the + label alone cannot pick an adapter. The call type is what distinguishes + them. Returning None when it does not is deliberate: normalize_usage + refuses an unknown provider rather than guessing token semantics, which is + the whole point of keeping the native object authoritative. + """ + + if label == "openai" and call_type and "responses" in call_type: + return OPENAI_RESPONSES + if label in _ADAPTERS: + return label + return None + + +_ADAPTERS = { + ANTHROPIC: _normalize_anthropic, + OPENAI_RESPONSES: _normalize_openai_responses, +} + + +def normalize_usage(provider: str, native_usage: Mapping[str, Any] | None) -> NormalizedUsage: + """Derive the cross-provider view. Never mutates or replaces the native object.""" + + adapter = _ADAPTERS.get(provider) + if adapter is None: + raise UsageSemanticsError( + f"no usage adapter for provider {provider!r}; refusing to guess its token semantics" + ) + if not isinstance(native_usage, Mapping): + return NormalizedUsage(None, None, None, None, None, None) + return adapter(native_usage) From 376ed3bb4abda80f7a57cf963822002b5304f2f9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gerg=C5=91=20Magyar?= Date: Tue, 8 Sep 2026 19:06:03 +0100 Subject: [PATCH 2/2] perf(lock): probe this process's own start time once (#3222) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(lock): probe this process's own start time once `acquireFileLock` stamps the owner file with the acquiring process's start time so a later reclaimer can tell a live owner from pid reuse. That value cannot change while we are running, but it was re-probed on every acquisition — and on Windows the probe is a `powershell.exe` spawn plus a `Get-CimInstance Win32_Process` WMI query, which is the single most expensive step in taking an uncontended lock. Add `readProcessStartTimeCached` and make it the default reader in `acquireFileLock` and `resolveWatchDeps`. Only this process's own pid is cached: - A foreign pid is always re-probed. That process can exit and its pid be reused, which is precisely what the stamp exists to detect. - A failed probe is not cached. `acquireFileLock` throws when the start time is empty, so caching one transient failure would leave the process unable to take a lock for the rest of its life. `readProcessStartTime` itself is unchanged and still probes every call, so the existing timezone-pinning regression test keeps exercising the real `ps` invocation instead of passing off a cached value. Behavior is otherwise identical: same probe, same string, same stamp. Co-Authored-By: Claude Opus 5 (1M context) * chore(autofix): apply prettier + eslint fixes via /autofix command --------- Co-authored-by: Gergo Magyar Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- gitnexus/src/core/auto-sync/starter.ts | 4 +- gitnexus/src/storage/file-lock.ts | 8 ++-- gitnexus/src/utils/process-identity.ts | 19 +++++++++ gitnexus/test/unit/process-identity.test.ts | 45 +++++++++++++++++++++ 4 files changed, 71 insertions(+), 5 deletions(-) diff --git a/gitnexus/src/core/auto-sync/starter.ts b/gitnexus/src/core/auto-sync/starter.ts index 624e092ec..7d67dfa64 100644 --- a/gitnexus/src/core/auto-sync/starter.ts +++ b/gitnexus/src/core/auto-sync/starter.ts @@ -4,7 +4,7 @@ import path from 'node:path'; import { execFileSync } from 'node:child_process'; import { acquireFileLock, FileLockBusyError } from '../../storage/file-lock.js'; import { getGlobalDir } from '../../storage/repo-manager.js'; -import { isProcessAlive, readProcessStartTime } from '../../utils/process-identity.js'; +import { isProcessAlive, readProcessStartTimeCached } from '../../utils/process-identity.js'; import { loadAutoSyncConfig } from './config.js'; import { runAutoSyncOnce } from './runner.js'; import { getAutoSyncMutexPath, getAutoSyncWatchDir } from './state.js'; @@ -632,7 +632,7 @@ function resolveWatchDeps(deps: Partial = {}): AutoSyn return undefined; } }), - readProcessStartTime: deps.readProcessStartTime ?? readProcessStartTime, + readProcessStartTime: deps.readProcessStartTime ?? readProcessStartTimeCached, sleep: deps.sleep ?? ((ms) => diff --git a/gitnexus/src/storage/file-lock.ts b/gitnexus/src/storage/file-lock.ts index a857d6a47..2b2494ef8 100644 --- a/gitnexus/src/storage/file-lock.ts +++ b/gitnexus/src/storage/file-lock.ts @@ -3,7 +3,7 @@ import fs from 'node:fs/promises'; import os from 'node:os'; import path from 'node:path'; import { setTimeout as sleep } from 'node:timers/promises'; -import { isProcessAlive, readProcessStartTime } from '../utils/process-identity.js'; +import { isProcessAlive, readProcessStartTimeCached } from '../utils/process-identity.js'; const HOSTNAME = os.hostname(); @@ -46,7 +46,9 @@ export async function acquireFileLock( pid, ownerId: crypto.randomUUID(), processStartTime: - options.processStartTime ?? (options.readProcessStartTime ?? readProcessStartTime)(pid) ?? '', + options.processStartTime ?? + (options.readProcessStartTime ?? readProcessStartTimeCached)(pid) ?? + '', hostname: options.hostname ?? HOSTNAME, }; if (!owner.processStartTime) { @@ -69,7 +71,7 @@ export async function acquireFileLock( resolvedPath, owner, options.isProcessAlive ?? isProcessAlive, - options.readProcessStartTime ?? readProcessStartTime, + options.readProcessStartTime ?? readProcessStartTimeCached, ) ) { continue; diff --git a/gitnexus/src/utils/process-identity.ts b/gitnexus/src/utils/process-identity.ts index e7c99439c..8c00eab08 100644 --- a/gitnexus/src/utils/process-identity.ts +++ b/gitnexus/src/utils/process-identity.ts @@ -38,3 +38,22 @@ export function readProcessStartTime(pid: number): string | undefined { return undefined; } } + +let ownStartTime: string | undefined; + +/** + * `readProcessStartTime`, except this process's own start time is probed once. + * It cannot change while we are running, and every `acquireFileLock` — plus + * each retry attempt and each stale-lock reclaim guard — stamps the owner file + * with it. On Windows that probe is a `powershell.exe` spawn and a WMI query, + * so a process taking several locks pays it several times for one constant. + * + * A foreign pid is never cached: that process can exit and its pid can be + * reused, which is the very thing the stamp exists to detect. A failed probe + * is not cached either — one transient failure would otherwise leave the + * process unable to take a lock for its whole lifetime. + */ +export function readProcessStartTimeCached(pid: number): string | undefined { + if (pid !== process.pid) return readProcessStartTime(pid); + return (ownStartTime ??= readProcessStartTime(pid)); +} diff --git a/gitnexus/test/unit/process-identity.test.ts b/gitnexus/test/unit/process-identity.test.ts index 7b31f1d32..de3436882 100644 --- a/gitnexus/test/unit/process-identity.test.ts +++ b/gitnexus/test/unit/process-identity.test.ts @@ -4,8 +4,24 @@ import { isProcessAlive, readProcessStartTime } from '../../src/utils/process-id afterEach(() => { vi.restoreAllMocks(); + vi.doUnmock('node:child_process'); + vi.resetModules(); }); +/** + * Loads a fresh copy of the module (fresh memo) over a counted `execFileSync`, + * so "how many times did we actually shell out" is observable. `doMock` is not + * hoisted, so the statically imported functions used by the other tests keep + * the real implementation. + */ +async function withCountedProbe(probe: () => string) { + const execFileSync = vi.fn(probe); + vi.doMock('node:child_process', () => ({ execFileSync })); + vi.resetModules(); + const identity = await import('../../src/utils/process-identity.js'); + return { execFileSync, readProcessStartTimeCached: identity.readProcessStartTimeCached }; +} + describe('process identity', () => { it('treats only ESRCH as a dead process', () => { const kill = vi.spyOn(process, 'kill'); @@ -39,4 +55,33 @@ describe('process identity', () => { } }, ); + + it('probes this process once and re-probes a foreign pid every time', async () => { + const { execFileSync, readProcessStartTimeCached } = await withCountedProbe(() => 'STAMP\n'); + + expect(readProcessStartTimeCached(process.pid)).toBe('STAMP'); + expect(readProcessStartTimeCached(process.pid)).toBe('STAMP'); + // On Windows each probe is a powershell.exe spawn plus a WMI query. + expect(execFileSync).toHaveBeenCalledTimes(1); + + // A foreign process can exit and its pid be reused — caching that stamp + // would blind the reuse check the stamp exists for. + expect(readProcessStartTimeCached(process.pid + 1)).toBe('STAMP'); + expect(readProcessStartTimeCached(process.pid + 1)).toBe('STAMP'); + expect(execFileSync).toHaveBeenCalledTimes(3); + }); + + it('retries after a failed self probe instead of caching the failure', async () => { + const { execFileSync, readProcessStartTimeCached } = await withCountedProbe(() => 'STAMP\n'); + execFileSync.mockImplementationOnce(() => { + throw new Error('probe unavailable'); + }); + + // A cached failure would leave acquireFileLock throwing "Unable to + // determine process start time" for the rest of the process's life. + expect(readProcessStartTimeCached(process.pid)).toBeUndefined(); + expect(readProcessStartTimeCached(process.pid)).toBe('STAMP'); + expect(readProcessStartTimeCached(process.pid)).toBe('STAMP'); + expect(execFileSync).toHaveBeenCalledTimes(2); + }); });