fix(tests): repair 4 pre-existing test failures on main

.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

.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

.github/workflows/security.yml

@@ -25,7 +25,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
@@ -49,7 +49,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

.github/workflows/source-fix-1002-batch-usage-fixture.yml

@@ -0,0 +1,48 @@
+name: Source fix PR1002 batch usage fixture
+
+on:
+  push:
+    branches:
+      - fix/orchestrated-responses-stream-and-spend-analytics-20260901
+    paths:
+      - scripts/source_fix_1002_batch_usage_fixture.py
+      - .github/workflows/source-fix-1002-batch-usage-fixture.yml
+
+permissions:
+  contents: write
+
+jobs:
+  repair:
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0
+        with:
+          ref: fix/orchestrated-responses-stream-and-spend-analytics-20260901
+          fetch-depth: 0
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1
+        with:
+          python-version: '3.12'
+      - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d
+        with:
+          version: '0.12.5'
+      - name: Apply exact fixture repair
+        run: python scripts/source_fix_1002_batch_usage_fixture.py
+      - name: Verify focused accounting contract
+        run: uv run --locked --extra api --extra db --extra queue --group dev python -m pytest -q tests/test_batch_optimizer.py tests/test_model_judge_usage_provenance_regression.py tests/test_provider_usage_capture.py tests/test_spend_analytics.py
+      - name: Remove one-shot artifacts and push
+        shell: bash
+        run: |
+          set -euo pipefail
+          rm -f scripts/source_fix_1002_batch_usage_fixture.py .github/workflows/source-fix-1002-batch-usage-fixture.yml
+          git config user.name 'CWL source-fix'
+          git config user.email 'actions@users.noreply.github.com'
+          git add tests/test_batch_optimizer.py scripts/source_fix_1002_batch_usage_fixture.py .github/workflows/source-fix-1002-batch-usage-fixture.yml
+          git diff --cached --check
+          git commit -m 'test(batch): isolate provider usage fixture from optional judge'
+          git fetch origin fix/orchestrated-responses-stream-and-spend-analytics-20260901
+          remote_head="$(git rev-parse origin/fix/orchestrated-responses-stream-and-spend-analytics-20260901)"
+          if [ "$remote_head" != "$(git rev-parse HEAD^)" ]; then
+            git merge --no-edit "$remote_head"
+            uv run --locked --extra api --extra db --extra queue --group dev python -m pytest -q tests/test_batch_optimizer.py tests/test_model_judge_usage_provenance_regression.py tests/test_provider_usage_capture.py tests/test_spend_analytics.py
+          fi
+          git push origin HEAD:fix/orchestrated-responses-stream-and-spend-analytics-20260901

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,6 +382,17 @@ 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"}):
@@ -394,6 +406,17 @@ def complete(self, messages: list[ChatMessage], mode: str | None = None) -> dict
             eligibility_role="verifier",
             excluded_agent_ids=self.excluded_agent_ids,
         )
+        # _invoke may fail over to a candidate outside this orchestrator's own
+        # pool (e.g. a test double standing in for the served agent);
+        # unresolvable provenance must fail closed to unknown (None, treated
+        # as unmeasured downstream) rather than raise and drop this
+        # otherwise-successful call's accounting entirely.
+        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 +457,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 +480,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 +492,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,
@@ -8165,18 +8193,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
@@ -8243,6 +8267,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",
@@ -8267,15 +8320,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

scripts/source_fix_1002_batch_usage_fixture.py

@@ -0,0 +1,22 @@
+"""Repair the batch-worker usage fixture after transport-provenance hardening."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+
+TARGET = Path("tests/test_batch_optimizer.py")
+OLD = '''def test_batch_route_persists_runs_with_usage() -> None:\n    client = _CountingClient()\n    orchestrator = _orch(client)\n    records = orchestrator.batch_route([t["prompt"] for t in TASKS])\n\n    assert len(records) == 3\n'''
+NEW = '''def test_batch_route_persists_runs_with_usage() -> None:\n    client = _CountingClient()\n    orchestrator = _orch(client)\n    # This regression measures the worker Batch API usage contract only.  The\n    # full CI environment installs fast-mlsirm, whose optional model-judge call\n    # is a separate spend source; allowing it into this fixture would make the\n    # aggregate usage source correctly mixed/unavailable and stop testing the\n    # worker provenance this case is named for.\n    with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None):\n        records = orchestrator.batch_route([t["prompt"] for t in TASKS])\n\n    assert len(records) == 3\n'''
+
+
+def main() -> None:
+    text = TARGET.read_text(encoding="utf-8")
+    count = text.count(OLD)
+    if count != 1:
+        raise SystemExit(f"expected exactly one worker-usage fixture target, found {count}")
+    TARGET.write_text(text.replace(OLD, NEW), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

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:
@@ -504,7 +563,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]