diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f17a29fe..588a7078 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,7 +17,7 @@ concurrency: jobs: pytest: name: Full unit and contract suite - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # actions/checkout@v7 @@ -39,7 +39,7 @@ jobs: nim_benchmark_quality: name: NIM benchmark coverage, docstrings, and package smoke - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # actions/checkout@v7 diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 40c2f94f..0fd934d0 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -21,7 +21,7 @@ jobs: # Always-on, cross-platform property tests. Fast, deterministic, no native deps. property_tests: name: Hypothesis property tests - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # actions/checkout@v7 @@ -45,7 +45,7 @@ jobs: # so CI stays cheap; schedule/dispatch runs use a longer budget. coverage_guided: name: Atheris coverage-guided - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # actions/checkout@v7 diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml index 3daca9cf..cce6a976 100644 --- a/.github/workflows/security.yml +++ b/.github/workflows/security.yml @@ -27,7 +27,7 @@ jobs: codeql_analysis: name: CodeQL analysis if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: actions: read contents: read @@ -51,7 +51,7 @@ jobs: python_supply_chain: name: Python supply chain - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Checkout repository uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # actions/checkout@v7 diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 9d168b3b..34dd509f 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -366,6 +366,7 @@ class _FastMLSIJudgeAdapter: served_agent_id: str | None = None served_model: str | None = None served_usage: dict[str, Any] | None = None + served_usage_source: str | None = None served_output: str | None = None mode: str = "auto" allowed_agent_ids: set[str] | None = None @@ -381,10 +382,27 @@ def client(self) -> ModelClient: """Expose the existing gateway client capability to fast-mlsirm.""" return self.orchestrator.client + @staticmethod + def _usage_source_for_agent(agent: ModelAgent, usage: Any) -> str | None: + """Classify transport-boundary usage without guessing from token counts.""" + if not isinstance(usage, dict): + return None + return ( + "synthetic_mock" + if agent.base_url.startswith("mock://") + else "provider_reported" + ) + def complete(self, messages: list[ChatMessage], mode: str | None = None) -> dict[str, Any]: """Return one judge completion through the constrained adapter.""" if mode is not None and (type(mode) is not str or mode not in {"auto", "route", "conduct"}): raise ValueError("mode must be auto, route, or conduct") + served_agent_at_call: ModelAgent | None = None + + def _capture_serving_agent(agent: ModelAgent) -> None: + nonlocal served_agent_at_call + served_agent_at_call = agent + output, served_id, served_model, usage = self.orchestrator._invoke( self._agent(), messages, @@ -393,7 +411,36 @@ def complete(self, messages: list[ChatMessage], mode: str | None = None) -> dict allowed_agent_ids=self.allowed_agent_ids, eligibility_role="verifier", excluded_agent_ids=self.excluded_agent_ids, - ) + on_success=_capture_serving_agent, + ) + # Prefer the exact ModelAgent _invoke's on_success callback captured + # atomically at the moment it served this call. A concurrent admin + # request can replace TaskOrchestrator.candidates (add/patch/remove + # candidate) with a different base_url under the same served_id in the + # gap between that provider call completing and this method resolving + # provenance; since ModelAgent is frozen and pool mutation always + # reassigns self.candidates rather than mutating it in place, the + # captured reference is immune to that race, unlike a fresh + # self.orchestrator._agent(served_id) lookup against the live mutable + # pool (Devin review on PR #1002). + # + # _invoke may also fail over to a candidate outside this orchestrator's + # own pool (e.g. a test double standing in for the served agent), and a + # custom _invoke test double may not accept/call on_success at all; + # unresolvable provenance must fail closed to unknown (None, treated as + # unmeasured downstream) rather than raise and drop this + # otherwise-successful call's accounting entirely. + if served_agent_at_call is not None: + self.served_usage_source = self._usage_source_for_agent( + served_agent_at_call, usage + ) + else: + try: + served_agent = self.orchestrator._agent(served_id) + except KeyError: + self.served_usage_source = None + else: + self.served_usage_source = self._usage_source_for_agent(served_agent, usage) return self._completion_payload( output, served_id, served_model, usage, self.mode if mode is None else mode ) @@ -434,8 +481,12 @@ def complete_structured( # must survive validation failing on how to interpret its result. self.served_agent_id = agent.id self.served_model = agent.model + response_usage = response.get("usage") self.served_usage = ( - response.get("usage") if isinstance(response.get("usage"), dict) else None + dict(response_usage) if isinstance(response_usage, dict) else None + ) + self.served_usage_source = self._usage_source_for_agent( + agent, self.served_usage ) output = ModelClient._response_content(agent, response) return self._completion_payload( @@ -453,7 +504,8 @@ def _completion_payload( """Build the bounded adapter response shared by normal and structured calls.""" self.served_agent_id = served_id self.served_model = served_model - self.served_usage = usage + usage_snapshot = dict(usage) if isinstance(usage, dict) else None + self.served_usage = usage_snapshot self.served_output = output trace = [ { @@ -464,8 +516,8 @@ def _completion_payload( "output": output, } ] - if usage is not None: - trace[0]["usage"] = usage + if usage_snapshot is not None: + trace[0]["usage"] = dict(usage_snapshot) return { "answer": output, "mode": mode, @@ -7670,6 +7722,7 @@ def _invoke( allowed_agent_ids: set[str] | None = None, eligibility_role: str | None = None, excluded_agent_ids: set[str] | None = None, + on_success: Callable[[ModelAgent], None] | None = None, ) -> tuple[str, str, str, dict[str, Any] | None]: """Call an agent with bounded, safety-aware tool retry and failover. @@ -7680,6 +7733,21 @@ def _invoke( ``eligibility_role`` keeps operator exclusions tied to the role used to select the primary when the call's effort profile has a distinct name. + + ``on_success``, when given, receives the exact ``ModelAgent`` that served + the winning call before this method returns. ``ModelAgent`` is a frozen + dataclass and ``candidates``/``race_members`` here are call-local lists + snapshotted before any provider transport; a pool-mutation API + (add/patch/remove candidate) only ever reassigns ``self.candidates`` to a + new list object and never mutates an existing one in place (see + ``_agent``), so this reference stays valid even if a concurrent request + replaces the same agent id in the live pool while this call is still in + flight or between this method returning and a caller's own follow-up + lookup. Callers that need serving-transport provenance (e.g. classifying + usage as provider-reported vs. synthetic) must prefer this over a + post-hoc ``self._agent(served_id)`` lookup, which reads the live mutable + pool and can silently resolve to a different agent (Devin review on + PR #1002). """ required_tags = ("vision",) if self._source_image_parts(messages) else () prompt_context = self._prompt_interaction(messages) @@ -7762,6 +7830,11 @@ def call(agent: ModelAgent) -> tuple[str, str, str, dict[str, Any] | None]: outcome.completion_ms / 1000, output_tokens=output_tokens, ) + if on_success is not None: + for candidate in race_members: + if candidate.id == outcome.winner_endpoint_id: + on_success(candidate) + break return outcome.value retry_limit = min(self.tool_retry_attempts, MAX_TOOL_RETRY_ATTEMPTS) bounded_provider_response_failures = 0 @@ -7881,6 +7954,8 @@ def call(agent: ModelAgent) -> tuple[str, str, str, dict[str, Any] | None]: total_tokens=total_tokens, ) self._record_success(agent.id) + if on_success is not None: + on_success(agent) return output, agent.id, agent.model, usage if ( last_provider_response_error is not None @@ -8168,18 +8243,14 @@ def _model_judge_verification( "judge": "model", } verification.update(self._judge_adapter_accounting_fields(judge_adapter)) - # The adapter's provider-boundary capture is authoritative for - # usage. fast-mlsirm aggregates a missing trace usage into a - # non-empty zero-token mapping, which must not turn an unmeasured - # call into provider-reported zero spend here. - result_usage_has_positive_evidence = isinstance(result.usage, Mapping) and any( - type(result.usage.get(key)) is int and result.usage[key] > 0 - for key in ("prompt_tokens", "completion_tokens", "total_tokens") - ) - if result.usage and ( - judge_adapter.served_usage is not None or result_usage_has_positive_evidence + # Prefer the adapter's transport-boundary snapshot, including an + # authoritative all-zero provider report. Only fall back to a + # positive fast-mlsirm aggregate when no boundary usage survived. + if ( + "judge_usage" not in verification + and self._usage_has_positive_evidence(result.usage) ): - verification["judge_usage"] = result.usage + verification["judge_usage"] = dict(result.usage) verification["judge_orchestration_mode"] = result.orchestration_mode # The provider call has already completed by this point (result # is a real response, with judge_agent_id/judge_model/judge_usage @@ -8246,6 +8317,35 @@ def _model_judge_verification( **self._judge_adapter_accounting_fields(judge_adapter), } + @staticmethod + def _usage_has_positive_evidence(usage: Any) -> bool: + """Return whether Chat or Responses usage reports a positive token count. + + This remains the compatibility rule for aggregate usage whose origin + is unknown. A provider-boundary zero measurement is handled separately + by ``_usage_is_reported_token_mapping`` plus explicit provenance. + """ + if not isinstance(usage, Mapping): + return False + counts = ( + usage.get("prompt_tokens", usage.get("input_tokens")), + usage.get("completion_tokens", usage.get("output_tokens")), + usage.get("total_tokens"), + ) + return any(type(value) is int and value > 0 for value in counts) + + @staticmethod + def _usage_is_reported_token_mapping(usage: Any) -> bool: + """Validate complete non-negative Chat or Responses token counters.""" + if not isinstance(usage, Mapping): + return False + counts = ( + usage.get("prompt_tokens", usage.get("input_tokens")), + usage.get("completion_tokens", usage.get("output_tokens")), + usage.get("total_tokens"), + ) + return all(type(value) is int and value >= 0 for value in counts) + @staticmethod def _judge_adapter_accounting_fields( judge_adapter: "_FastMLSIJudgeAdapter | None", @@ -8270,15 +8370,18 @@ def _judge_adapter_accounting_fields( fields["judge_agent_id"] = judge_adapter.served_agent_id if judge_adapter.served_model is not None: fields["judge_model"] = judge_adapter.served_model - if judge_adapter.served_usage: - # A falsy served_usage (missing/invalid response usage) is left - # genuinely absent rather than fabricated as reported-zero - # (Devin review on #961, on this same fix): judge_agent_id/ - # judge_model above already keep a completed-but-unmeasured call - # attributable, and downstream budget/spend consumers derive an - # honest estimated fallback from the judge's own served_output - # text instead of trusting a fabricated "reported" usage dict. - fields["judge_usage"] = judge_adapter.served_usage + if TaskOrchestrator._usage_has_positive_evidence( + judge_adapter.served_usage + ) or ( + judge_adapter.served_usage_source == "provider_reported" + and TaskOrchestrator._usage_is_reported_token_mapping( + judge_adapter.served_usage + ) + ): + # Positive aggregate usage remains backward-compatible. Exact + # provider-boundary zero usage is accepted only with explicit + # provenance; synthetic mock/fast-mlsirm zero remains unmeasured. + fields["judge_usage"] = dict(judge_adapter.served_usage) if judge_adapter.served_output is not None: # The judge's own generated text, not the verifier_output text # it was judging (Devin review on #961, on this same fallback diff --git a/tests/test_batch_optimizer.py b/tests/test_batch_optimizer.py index 350384a6..5a9c0b86 100644 --- a/tests/test_batch_optimizer.py +++ b/tests/test_batch_optimizer.py @@ -73,9 +73,16 @@ def _orch(client: ModelClient | None = None) -> TaskOrchestrator: def test_batch_route_persists_runs_with_usage() -> None: + """Batch routing persists one run per task, each carrying the worker Batch API's own usage.""" client = _CountingClient() orchestrator = _orch(client) - records = orchestrator.batch_route([t["prompt"] for t in TASKS]) + # This regression measures the worker Batch API usage contract only. The + # full CI environment installs fast-mlsirm, whose optional model-judge call + # is a separate spend source; allowing it into this fixture would make the + # aggregate usage source correctly mixed/unavailable and stop testing the + # worker provenance this case is named for. + with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None): + records = orchestrator.batch_route([t["prompt"] for t in TASKS]) assert len(records) == 3 assert client.batch_calls == 1 and client.chat_calls == 0 # one batch, zero serial calls diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index cfa249ce..ac75940b 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -209,6 +209,65 @@ def count_messages(self, messages: list[dict], model: str = "") -> int: assert budget_contribution["model-x"] > 0 +def test_judge_adapter_zero_aggregate_served_usage_is_not_treated_as_reported() -> None: + """A non-empty-but-all-zero ``served_usage`` must stay genuinely absent. + + Devin review on `contextual-orchestrator#1002`: unlike `result.usage` + (fast-mlsirm's own aggregate, already guarded), `_FastMLSIJudgeAdapter`'s + `served_usage` comes straight from `ModelClient`'s provider-boundary + capture -- for the `mock://` transport that is exactly + `{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}` + (`ModelClient._mock_response`'s chat.completion branch), a non-empty + dict that is truthy in Python even though every count is 0. The plain + `if judge_adapter.served_usage:` check this function used to run + treated that as genuine reported-zero evidence, disagreeing with the + sibling `result.usage` guard a few lines above for the identical + fabricated shape. `_judge_adapter_accounting_fields` must reject it + exactly like `result.usage` already does. + """ + adapter = orchestrator_module._FastMLSIJudgeAdapter( + orchestrator=None, # type: ignore[arg-type] + text="task", + judge="model", + served_agent_id="general_agent", + served_model="model-x", + served_usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + served_output="judge rationale text", + ) + + fields = TaskOrchestrator._judge_adapter_accounting_fields(adapter) + + assert fields["judge_agent_id"] == "general_agent" + assert fields["judge_model"] == "model-x" + assert "judge_usage" not in fields + assert fields["judge_output_text"] == "judge rationale text" + + +def test_judge_adapter_positive_served_usage_is_still_reported() -> None: + """A genuine positive token count in ``served_usage`` must still count. + + The zero-aggregate guard above must not overcorrect into dropping real + provider-reported usage. + """ + adapter = orchestrator_module._FastMLSIJudgeAdapter( + orchestrator=None, # type: ignore[arg-type] + text="task", + judge="model", + served_agent_id="general_agent", + served_model="model-x", + served_usage={"prompt_tokens": 12, "completion_tokens": 34, "total_tokens": 46}, + served_output="judge rationale text", + ) + + fields = TaskOrchestrator._judge_adapter_accounting_fields(adapter) + + assert fields["judge_usage"] == { + "prompt_tokens": 12, + "completion_tokens": 34, + "total_tokens": 46, + } + + def test_free_conduct_keeps_model_judge_inside_zero_cost_pool() -> None: class _RecordingClient(_ScriptedClient): def __init__(self) -> None: @@ -450,13 +509,17 @@ def chat(self, agent: ModelAgent, messages: list, temperature: float | None = No def test_fast_mlsirm_path_is_used_when_available() -> None: + """Model judge verification drives the injected judge and keeps the adapter's own usage.""" + class _FakeJudge: def __init__(self, orchestrator, mode: str = "route", accept_threshold: float = 0.7) -> None: + """Store the injected adapter plus the requested mode and accept threshold.""" self.adapter = orchestrator self.mode = mode self.accept_threshold = accept_threshold def judge(self, **_) -> object: + """Drive one adapter call and return a fixed, accepted structured verdict.""" self.adapter.complete([{"role": "user", "content": "ping"}]) return type("Result", (), { "accepted": True, @@ -472,6 +535,7 @@ class _FormatError(Exception): class _Criterion: def __init__(self, criterion_id: str, description: str, weight: float) -> None: + """Store one rubric criterion's id, description, and weight.""" self.criterion_id = criterion_id self.description = description self.weight = weight @@ -504,7 +568,11 @@ def __init__(self, criterion_id: str, description: str, weight: float) -> None: # candidate pool -- "backup_judge" is not a real pool member here. assert result["judge_model"] == "backup-model" assert result["judge_orchestration_mode"] == "route" - assert result["judge_usage"] == {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7} + # The adapter's own transport-boundary capture (_invoke's returned usage) + # is now preferred over fast-mlsirm's result.usage aggregate when both + # carry evidence, so this reflects served_usage exactly as _invoke + # returned it -- not backfilled from result.usage's fuller breakdown. + assert result["judge_usage"] == {"total_tokens": 7} assert result["judge_criterion_scores"] == {"evidence_quality": 0.8, "risk_signal": 0.9} assert result["judge_irt_item_type"] == "dichotomous" assert result["judge_irt_row"] == [1, 1] diff --git a/tests/test_model_judge_usage_provenance_regression.py b/tests/test_model_judge_usage_provenance_regression.py new file mode 100644 index 00000000..c9645a52 --- /dev/null +++ b/tests/test_model_judge_usage_provenance_regression.py @@ -0,0 +1,322 @@ +"""Regression coverage for model-judge usage provenance. + +A provider may authoritatively report an all-zero usage mapping, while the +repository's mock transport and fast-mlsirm can also synthesize the same value +shape when no measured usage exists. Accounting must distinguish those origins +instead of guessing from token counts alone. The transport boundary is also a +snapshot boundary: changing getters and later mutation must not rewrite the +captured evidence. +""" + +from __future__ import annotations + +from unittest.mock import patch + +from contextual_orchestrator import ModelAgent, TaskOrchestrator +from contextual_orchestrator.orchestrator import ModelClient, _FastMLSIJudgeAdapter + + +ZERO_USAGE = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, +} +RESPONSES_ZERO_USAGE = { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, +} + + +class _ChangingUsageResponse(dict[str, object]): + """Expose a TOCTOU-sensitive ``get('usage')`` without altering other keys.""" + + def __init__(self) -> None: + """Seed a valid structured judge response and a zero usage-read counter.""" + super().__init__( + choices=[ + {"message": {"content": '{"decision":"ACCEPT","reason":"ok"}'}} + ] + ) + self.usage_reads = 0 + + def get(self, key: str, default: object = None) -> object: + """Return other keys normally; return zero usage once, then a changed value.""" + if key != "usage": + return super().get(key, default) + self.usage_reads += 1 + if self.usage_reads == 1: + return dict(ZERO_USAGE) + return { + "prompt_tokens": 91, + "completion_tokens": 37, + "total_tokens": 128, + } + + +def _adapter(*, usage_source: str | None) -> _FastMLSIJudgeAdapter: + """Build a judge adapter fixture with zero served_usage and the given provenance.""" + adapter = _FastMLSIJudgeAdapter( + orchestrator=None, # type: ignore[arg-type] + text="task", + judge="judge_agent", + served_agent_id="judge_agent", + served_model="judge-model", + served_usage=dict(ZERO_USAGE), + served_output="judge rationale", + ) + # Regression-first compatibility: the pre-repair adapter has no declared + # provenance field, so assigning it dynamically makes the RED exercise the + # accounting decision rather than fail during object construction. + adapter.served_usage_source = usage_source # type: ignore[attr-defined] + return adapter + + +def test_provider_reported_all_zero_usage_remains_reported() -> None: + """Exact provider evidence must not be converted into estimated spend.""" + fields = TaskOrchestrator._judge_adapter_accounting_fields( + _adapter(usage_source="provider_reported") + ) + + assert fields["judge_usage"] == ZERO_USAGE + + +def test_provider_reported_responses_zero_usage_remains_reported() -> None: + """Responses-style input/output counters are equally authoritative evidence.""" + adapter = _adapter(usage_source="provider_reported") + adapter.served_usage = dict(RESPONSES_ZERO_USAGE) + + fields = TaskOrchestrator._judge_adapter_accounting_fields(adapter) + + assert fields["judge_usage"] == RESPONSES_ZERO_USAGE + + +def test_synthetic_mock_all_zero_usage_remains_unmeasured() -> None: + """The mock transport's zero fill must not become provider evidence.""" + fields = TaskOrchestrator._judge_adapter_accounting_fields( + _adapter(usage_source="synthetic_mock") + ) + + assert "judge_usage" not in fields + + +def test_unknown_all_zero_usage_remains_unmeasured() -> None: + """Missing provenance fails closed instead of asserting reported zero.""" + fields = TaskOrchestrator._judge_adapter_accounting_fields( + _adapter(usage_source=None) + ) + + assert "judge_usage" not in fields + + +def test_structured_adapter_captures_provider_usage_source_at_transport_boundary() -> None: + """Structured judge capture records whether zero usage came from a provider.""" + provider_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.example/v1", + tags=("verification",), + ) + orchestrator = TaskOrchestrator([provider_agent]) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", provider_agent.id) + response = { + "choices": [{"message": {"content": '{"decision":"ACCEPT","reason":"ok"}'}}], + "usage": dict(ZERO_USAGE), + } + + with patch.object(orchestrator.client, "proxy_send", return_value=response): + adapter.complete_structured( + [{"role": "user", "content": "judge"}], + response_format={"type": "json_object"}, + ) + + assert adapter.served_usage_source == "provider_reported" # type: ignore[attr-defined] + + +def test_plain_adapter_captures_served_provider_usage_source() -> None: + """Fallback-aware plain completion classifies the agent that actually served.""" + provider_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.example/v1", + tags=("verification",), + ) + orchestrator = TaskOrchestrator([provider_agent]) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", provider_agent.id) + + with patch.object( + orchestrator, + "_invoke", + return_value=( + "judge rationale", + provider_agent.id, + provider_agent.model, + dict(ZERO_USAGE), + ), + ): + adapter.complete([{"role": "user", "content": "judge"}]) + + assert adapter.served_usage == ZERO_USAGE + assert adapter.served_usage_source == "provider_reported" # type: ignore[attr-defined] + + +def test_structured_adapter_reads_usage_once_at_transport_boundary() -> None: + """A changing getter cannot substitute different usage after validation.""" + provider_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.example/v1", + tags=("verification",), + ) + orchestrator = TaskOrchestrator([provider_agent]) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", provider_agent.id) + response = _ChangingUsageResponse() + + with patch.object(orchestrator.client, "proxy_send", return_value=response): + adapter.complete_structured( + [{"role": "user", "content": "judge"}], + response_format={"type": "json_object"}, + ) + + assert response.usage_reads == 1 + assert adapter.served_usage == ZERO_USAGE + + +def test_structured_adapter_snapshots_mutable_usage_alias() -> None: + """Later mutation of the provider response cannot rewrite captured evidence.""" + provider_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.example/v1", + tags=("verification",), + ) + orchestrator = TaskOrchestrator([provider_agent]) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", provider_agent.id) + reported_usage = dict(ZERO_USAGE) + response = { + "choices": [{"message": {"content": '{"decision":"ACCEPT","reason":"ok"}'}}], + "usage": reported_usage, + } + + with patch.object(orchestrator.client, "proxy_send", return_value=response): + adapter.complete_structured( + [{"role": "user", "content": "judge"}], + response_format={"type": "json_object"}, + ) + + reported_usage["prompt_tokens"] = 99 + reported_usage["total_tokens"] = 99 + assert adapter.served_usage == ZERO_USAGE + + +def test_judge_usage_source_survives_concurrent_pool_replacement_provider_to_mock() -> None: + """A provider's genuine all-zero usage must not be reclassified as synthetic. + + Simulates a concurrent admin pool update racing an in-flight judge call + (Devin review on PR #1002): while the provider's ``chat()`` request is + still executing, another thread replaces the served agent id's pool entry + with a ``mock://`` agent. ``_invoke``'s ``on_success`` callback captures + the exact (frozen) ``ModelAgent`` that served this call before that + replacement can be observed downstream, so the adapter's provenance + classification -- and therefore whether this call's zero usage is + honestly counted as measured spend -- must reflect what genuinely served + the call, never a fresh ``TaskOrchestrator.candidates`` lookup by id that + can race the replacement. + """ + provider_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.example/v1", + tags=("verification",), + ) + + class _ConcurrentReplacementClient(ModelClient): + def chat(self, agent: ModelAgent, messages: list, temperature: float | None = None) -> str: # type: ignore[override] + del messages, temperature + # The concurrent mutation: a same-id pool replacement landing + # while this provider call is still in flight, exactly like a + # different request thread patching the agent through the admin + # API. TaskOrchestrator never mutates self.candidates in place + # (only reassigns it), so this cannot retroactively change + # ``agent`` -- the point under test is whether the *adapter* + # still resolves provenance from that unaffected reference. + orchestrator.candidates = [ + ModelAgent( + "judge_agent", + "judge-model", + base_url="mock://replaced-concurrently", + tags=("verification",), + ) + ] + self._local.usage = dict(ZERO_USAGE) + return "judge rationale" + + orchestrator = TaskOrchestrator([provider_agent], client=_ConcurrentReplacementClient()) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", provider_agent.id) + + adapter.complete([{"role": "user", "content": "judge"}]) + + assert adapter.served_usage_source == "provider_reported" # type: ignore[attr-defined] + fields = TaskOrchestrator._judge_adapter_accounting_fields(adapter) + assert fields["judge_usage"] == ZERO_USAGE + + +def test_judge_usage_source_survives_concurrent_pool_replacement_mock_to_provider() -> None: + """The mock transport's synthetic zero fill must not become provider evidence. + + Mirror of the provider-to-mock race above, in the other direction: the + agent that actually served this call was ``mock://``, but a concurrent + pool update replaces the same id with a real provider agent while the + call is (conceptually) still in flight. Without the atomic capture, a + post-hoc ``self._agent(served_id)`` lookup would observe the replacement + and mislabel synthetic mock usage as authoritative provider evidence. + """ + mock_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="mock://catalog", + tags=("verification",), + ) + + class _ConcurrentReplacementClient(ModelClient): + def chat(self, agent: ModelAgent, messages: list, temperature: float | None = None) -> str: # type: ignore[override] + del messages, temperature + orchestrator.candidates = [ + ModelAgent( + "judge_agent", + "judge-model", + base_url="https://provider.replaced-concurrently/v1", + tags=("verification",), + ) + ] + self._local.usage = dict(ZERO_USAGE) + return "judge rationale" + + orchestrator = TaskOrchestrator([mock_agent], client=_ConcurrentReplacementClient()) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", mock_agent.id) + + adapter.complete([{"role": "user", "content": "judge"}]) + + assert adapter.served_usage_source == "synthetic_mock" # type: ignore[attr-defined] + fields = TaskOrchestrator._judge_adapter_accounting_fields(adapter) + assert "judge_usage" not in fields + + +def test_structured_adapter_marks_mock_zero_usage_as_synthetic() -> None: + """The identical mock value shape is explicitly non-authoritative.""" + mock_agent = ModelAgent( + "judge_agent", + "judge-model", + base_url="mock://catalog", + tags=("verification",), + ) + orchestrator = TaskOrchestrator([mock_agent]) + adapter = _FastMLSIJudgeAdapter(orchestrator, "task", mock_agent.id) + + adapter.complete_structured( + [{"role": "user", "content": "judge"}], + response_format={"type": "json_object"}, + ) + + assert adapter.served_usage == ZERO_USAGE + assert adapter.served_usage_source == "synthetic_mock" # type: ignore[attr-defined] diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 0eb661fb..ec449183 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -64,7 +64,7 @@ def test_unknown_tokenizer_uses_authoritative_provider_usage() -> None: embedding_token_counter=UnavailableEmbeddingTokenCounter(), ) - document = coordinator.complete_embeddings_batch(["synthetic input"]) + document = coordinator.complete_embeddings_batch(["synthetic input"], wait_timeout=1) assert document["status"] == "completed" assert document["total_tokens"] == len("synthetic input".encode("utf-8")) @@ -124,8 +124,17 @@ def test_unknown_tokenizer_byte_bound_never_becomes_recorded_usage(text) -> None embedding_token_counter=UnavailableEmbeddingTokenCounter(), ) - document = coordinator.complete_embeddings_batch([text]) + # The provider embedding backend completes asynchronously in a background + # worker thread (see ProviderEmbeddingBatchBackend._run_job); without an + # explicit wait_timeout, complete_embeddings_batch() can return before the + # job finishes, and _embeddings_batch_document_locked's not-is_complete + # early-return document omits "total_tokens" entirely. That is a race + # against thread scheduling, not a property of any particular input text. + # Every other provider-backed complete_embeddings_batch() call in this + # file passes wait_timeout for the same reason. + document = coordinator.complete_embeddings_batch([text], wait_timeout=1) + assert document["status"] == "completed" assert document["total_tokens"] == len(text.encode("utf-8")) diff --git a/tests/test_provider_usage_capture.py b/tests/test_provider_usage_capture.py index 6e5f89b8..abb14ef2 100644 --- a/tests/test_provider_usage_capture.py +++ b/tests/test_provider_usage_capture.py @@ -2,7 +2,9 @@ A gateway that already sees provider `usage` should bill on it, not a char heuristic. These assert reported completion_tokens flow into spend_analytics and are labeled, -while the mock path stays unavailable. +while the mock path stays unavailable. Provider-worker accounting is isolated from +the optional model-judge extra because that is a separate provider call with its own +accounting contract. """ from __future__ import annotations @@ -42,15 +44,19 @@ def test_take_usage_returns_and_clears() -> None: def test_reported_usage_preferred_and_labeled() -> None: + """Provider-reported completion usage wins over a char-count estimate and is labeled 'reported'.""" client = _ReportingClient(completion_tokens=50) orchestrator = TaskOrchestrator( [ModelAgent("general_agent", "priced-model", tags=("reasoning",))], client=client, price_per_million={"priced-model": 10.0}, ) - # Accounting contract, not dispatch: pin the single-step route path. + # Accounting contract, not dispatch: pin the single-step route path and + # disable the optional model judge so its independent call cannot consume + # or add usage to this worker-only assertion. orchestrator._triage_fn = lambda text: False - orchestrator.run([{"role": "user", "content": "do the work"}]) + with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None): + orchestrator.run([{"role": "user", "content": "do the work"}]) row = next(r for r in orchestrator.spend_analytics()["by_model"] if r["model"] == "priced-model") assert row["usage_source"] == "reported" @@ -59,10 +65,12 @@ def test_reported_usage_preferred_and_labeled() -> None: def test_reported_prompt_tokens_surface_in_totals() -> None: + """Provider-reported prompt tokens roll up into spend_analytics totals labeled 'reported'.""" client = _ReportingClient(completion_tokens=30) # also reports prompt_tokens=5 per call orchestrator = TaskOrchestrator([ModelAgent("general_agent", "priced-model", tags=("reasoning",))], client=client) orchestrator._triage_fn = lambda text: False # single-step route accounting - orchestrator.run([{"role": "user", "content": "route once"}]) + with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None): + orchestrator.run([{"role": "user", "content": "route once"}]) totals = orchestrator.spend_analytics()["totals"] assert totals["prompt_tokens_source"] == "reported" assert totals["prompt_tokens"] == 5 @@ -107,4 +115,4 @@ def test_conduct_all_steps_reported() -> None: if name.startswith("test_") and callable(fn): fn() print(f"ok {name}") - print("ok") + print("ok") \ No newline at end of file diff --git a/tests/test_spend_analytics.py b/tests/test_spend_analytics.py index 7ef56568..3d6f900e 100644 --- a/tests/test_spend_analytics.py +++ b/tests/test_spend_analytics.py @@ -4,6 +4,7 @@ import json import threading +from unittest.mock import patch import urllib.request from contextual_orchestrator import ModelAgent, TaskOrchestrator @@ -29,8 +30,17 @@ def _orchestrator(*, price: float | None = None) -> TaskOrchestrator: def test_exact_output_without_prompt_usage_is_explicitly_unavailable() -> None: - orchestrator = _orchestrator() - orchestrator.run([{"role": "user", "content": "account for this"}]) + """Output-only tokenizer usage leaves prompt/cost totals explicitly unavailable, not estimated.""" + # This test owns the raw-output tokenizer fallback contract, not the + # optional fast-mlsirm judge integration. Resolve that optional capability + # deterministically as unavailable so installing an extra package cannot + # change the evidence-source assertion from tokenizer to reported/mixed. + with patch( + "contextual_orchestrator.orchestrator._resolve_fast_mlsirm_components", + return_value=None, + ): + orchestrator = _orchestrator() + orchestrator.run([{"role": "user", "content": "account for this"}]) report = orchestrator.spend_analytics() row = report["by_model"][0] @@ -38,7 +48,7 @@ def test_exact_output_without_prompt_usage_is_explicitly_unavailable() -> None: assert report["totals"]["output_tokens"] > 0 assert report["totals"]["prompt_tokens"] is None assert report["totals"]["cost_usd"] is None - assert row["usage_source"] == "mixed" + assert row["usage_source"] == "tokenizer" assert row["cost_usd"] is None assert not any("estimated" in key for key in row | report["totals"])