diff --git a/CHANGELOG.md b/CHANGELOG.md index 83a3be438..3359e6185 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,23 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) official Run NIM Anywhere terms, restoring fail-closed live benchmark execution through 2026-10-05 without treating prototype access as production pricing or licensing evidence. +- Structured synthesis on the main `orchestrator/free` serving path now feeds + the realtime fast-mlsirm judge, so judged quality — not just transport + success — becomes routing evidence on the path that carries most traffic. + This wires an already-implemented, already-grounded measurement pipeline + into a path that was silently not calling it; it introduces no new + technique. The grounding is the existing record in + [`docs/doctoring/measured-routing-evidence.md`](docs/doctoring/measured-routing-evidence.md) + ("Real-time judging before returning answers" and "Multi-layer + simple-structure measurement (fast-mlsirm)": Ong et al., 2024; Chen et al., + 2023; Zheng et al., 2023; Jeon et al., 2021), plus Baker (2001) in + [`docs/papers/README.md`](docs/papers/README.md) for the IRT ability + fitting. The call is observation-only: it never branches the served answer, + stays inside the request's own agent eligibility (explicit model pin, free + pool, ZDR policy, file replicas) minus agents this request already proved + unavailable, is skipped once the request's own spend exhausts the operator + budget, and records its own token usage on the run so budget and spend + analytics see it. - Workflow workers now preserve the caller message array exactly once, while the added envelope carries only the subtask and Conductor-style prior-step access list instead of duplicating the task or source attachments. diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 9d168b3b3..30b028761 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections import Counter, deque, OrderedDict -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from contextlib import contextmanager from contextvars import ContextVar, copy_context from concurrent.futures import ThreadPoolExecutor @@ -95,6 +95,15 @@ _REQUEST_ENDPOINT_IDENTITY: ContextVar[str | None] = ContextVar( "contextual_orchestrator_request_endpoint_identity", default=None ) +#: Agent ids the active request is allowed to send its content to, or None when +#: unrestricted. Set by :func:`_request_eligibility_scope` from the very same +#: allow-list the request's own serving path is already constrained by, so +#: request-scoped *side* calls that carry request content but take no +#: ``allowed_agent_ids`` argument of their own -- today the routing-evidence +#: embedding in :meth:`TaskOrchestrator._embedding_agent_id` -- honor it too. +_REQUEST_ELIGIBLE_AGENT_IDS: ContextVar[frozenset[str] | None] = ContextVar( + "contextual_orchestrator_request_eligible_agent_ids", default=None +) _INVALID_REQUESTED_MODEL = object() @@ -149,6 +158,24 @@ def _agent_matches_request_endpoint(agent: ModelAgent) -> bool: ) +@contextmanager +def _request_eligibility_scope(allowed_agent_ids: Iterable[str] | None): + """Narrow request-scoped side calls to the agents this request may already reach. + + Mirrors :meth:`TaskOrchestrator.routing_endpoint_scope`: ``None`` means + "no restriction to add" and leaves any enclosing scope untouched, so this + can only ever narrow, never widen. + """ + if allowed_agent_ids is None: + yield + return + token = _REQUEST_ELIGIBLE_AGENT_IDS.set(frozenset(allowed_agent_ids)) + try: + yield + finally: + _REQUEST_ELIGIBLE_AGENT_IDS.reset(token) + + def _request_endpoint_partition() -> str: """Return a non-reversible cache partition for the configured endpoint.""" identity = _REQUEST_ENDPOINT_IDENTITY.get() @@ -157,6 +184,37 @@ def _request_endpoint_partition() -> str: return "endpoint:" + hashlib.sha256(identity.encode("utf-8")).hexdigest() +def _request_evidence_partition() -> str: + """Return a cache partition for every restriction that picks the embedder. + + A routing-evidence vector is produced by whichever provider + :meth:`TaskOrchestrator._embedding_agent_id` resolves under the active + request's own restrictions, so a cached vector may only be reused by a + later request whose restrictions resolve the same way: a restricted + request can never read a vector an unrestricted -- or differently + restricted -- one produced (Devin review on #1032). ZDR is part of it + because :meth:`TaskOrchestrator._zdr_agent_allowed` narrows the same + ranking, matching what ``_cache_key`` and ``_triage_workflow_required`` + already partition on. + + This is the *eligibility* half of an evidence cache key and is not + sufficient on its own: the same restriction shape can resolve to a + different embedding member (an agent-pool change, or a measured-member + reordering inside one group), whose vectors live in a different space + entirely. :meth:`TaskOrchestrator._embed_cached` and + :meth:`TaskOrchestrator._descriptor_vector_cached` therefore key on this + partition *and* the resolved embedding member's own identity. + """ + eligible = _REQUEST_ELIGIBLE_AGENT_IDS.get() + return "\x1f".join( + ( + _request_endpoint_partition(), + "zdr_only" if _REQUEST_ZDR_ONLY.get() else "zdr_any", + "eligible:*" if eligible is None else "eligible:" + ",".join(sorted(eligible)), + ) + ) + + # content is usually str; multimodal vision messages use OpenAI content-parts lists. ChatMessage = dict[str, Any] ProviderDestination = tuple[int, tuple[Any, ...]] @@ -4425,6 +4483,9 @@ def _orchestrated_provider_completion( virtual selector may advance to another eligible provider only after an HTTP 413 proves that the prior provider rejected the request before generation; other synthesis failures remain single-shot and fail closed. + Once a synthesis succeeds, the realtime fast-mlsirm judge observes the + served answer for quality-ledger and psychometric routing evidence + only -- it never changes the response already decided above. """ response_request = endpoint == "responses" api_surface = "responses" if response_request else "chat.completions" @@ -4558,11 +4619,36 @@ def _orchestrated_provider_completion( _excluded_agent_ids=request_exclusions, _allowed_agent_ids=None if virtual_model else {final_agent.id}, ) - in_flight_tokens, in_flight_cost = self._trace_budget_spend(workflow["trace"]) - self._raise_if_spend_budget_exceeded( - additional_output_tokens=in_flight_tokens, - additional_cost_usd=in_flight_cost, - ) + # conduct()'s own verifier-role judge is a completed provider call + # that never appears in workflow["trace"], so every checkpoint on + # this request's not-yet-persisted spend has to fold it in the same + # way the persisted meter does (Devin review on #1032). + in_flight_judges = self._run_judge_accounting_blocks(workflow) + + def budget_checkpoint(trace: list[dict[str, Any]]) -> None: + """Block the next provider call, metering what this one already spent. + + ``_replace_workflow_run`` is the only path that moves spend onto + the budget meter, and this request persists nothing until it + succeeds -- so raising here used to drop every provider call + ``conduct`` had already made, including its verifier-role judge. + The next request was then admitted against understated spend and + could repeat that forever (Devin review on #1032). Metering the + completed work as an unserved run first makes the rejection cost + what it actually cost. + """ + tokens, cost = self._trace_budget_spend( + trace, completed_judges=in_flight_judges + ) + try: + self._raise_if_spend_budget_exceeded( + additional_output_tokens=tokens, additional_cost_usd=cost + ) + except BudgetExceededError: + self._meter_unserved_spend(task, trace, workflow.get("verification")) + raise + + budget_checkpoint(workflow["trace"]) evidence = "\n\n".join( f"Workflow step {step['id']} ({step['role']}):\n{step['output']}" @@ -4633,29 +4719,38 @@ def _orchestrated_provider_completion( self.AUTO_MODEL, self.FREE_MODEL, } - allowed_agent_ids = ({final_agent.id} if isinstance(required_agent_id, str) else ( - { - candidate.id - for candidate in self.agents - if self._is_general_free_agent(candidate) and self._zdr_agent_allowed(candidate) - } - if free_only + # The one effective restriction this request's synthesis actually runs + # under, covering *every* way a request can be pinned rather than the + # `_required_agent_id` special case alone: a caller-named explicit + # model is exactly as pinned as a required file provider, because + # `synthesis_candidates` below is literally `[final_agent]` whenever + # `virtual_model` is false. Leaving that case unrestricted let an + # explicitly pinned request's prompt and served answer reach an + # unrelated -- possibly less trusted -- provider through the + # observation-only judge/embedding calls that reuse this set (Devin + # review on #1032). `free_only` implies `virtual_model` (FREE_MODEL is + # itself a virtual name), so the free pool and the ZDR-filtered + # virtual pool remain the only unpinned outcomes. + allowed_agent_ids = ( + {final_agent.id} + if isinstance(required_agent_id, str) or not virtual_model else ( { + candidate.id + for candidate in self.agents + if self._is_general_free_agent(candidate) + and self._zdr_agent_allowed(candidate) + } + if free_only + else { candidate.id for candidate in self.agents if self._zdr_agent_allowed(candidate) } - if virtual_model - else None ) - )) + ) if replica_agent_ids is not None: - allowed_agent_ids = ( - replica_agent_ids - if allowed_agent_ids is None - else allowed_agent_ids & replica_agent_ids - ) + allowed_agent_ids = allowed_agent_ids & replica_agent_ids synthesis_candidates = ( self._failover_candidates( final_agent, @@ -4811,6 +4906,17 @@ def send_synthesis( response_format = chat_body.get("response_format") synthesis_started = time.perf_counter() + # Provider calls this request already paid for and then walked away + # from: a model whose synthesis *and* repair both violated the + # caller's schema before a different member of the same virtual pool + # succeeded. `synthesis_step`/`repair_step` below are rebound on every + # pass of this loop, so without accumulating them the earlier model's + # two completed calls vanished from the final trace, from every later + # budget checkpoint, from the persisted run, and from spend analytics + # -- real provider spend, silently unmetered (Devin review on #1032). + # They are budget/analytics rows only: served-answer latency and usage + # stay bound to the successful attempt alone, below. + failed_attempts: list[dict[str, Any]] = [] while True: synthesis_failure_recorded = False try: @@ -4828,10 +4934,22 @@ def send_synthesis( and not isinstance(exc, EffortProfileError) ): self._group_router.observe_failure(final_agent.id) + # A transport failure here (as opposed to the schema-violation + # exits below, already fixed in round 6) still raises before + # any later budget_checkpoint runs -- so conduct()'s workflow + # spend and any earlier candidates' failed_attempts in this + # same loop would otherwise vanish from the meter exactly like + # budget_checkpoint's own except-clause guards against + # (Devin review on #1032, round 7). + self._meter_unserved_spend( + task, + [*workflow["trace"], *failed_attempts], + workflow.get("verification"), + ) raise synthesis_output = provider_output(final_agent, raw) synthesis_step = { - "id": len(workflow["trace"]), + "id": len(workflow["trace"]) + len(failed_attempts), "role": "synthesizer", "agent_id": final_agent.id, "subtask": "Provider-facing structured synthesis", @@ -4852,13 +4970,7 @@ def send_synthesis( if contract_error is None: break - in_flight_tokens, in_flight_cost = self._trace_budget_spend( - [*workflow["trace"], synthesis_step] - ) - self._raise_if_spend_budget_exceeded( - additional_output_tokens=in_flight_tokens, - additional_cost_usd=in_flight_cost, - ) + budget_checkpoint([*workflow["trace"], *failed_attempts, synthesis_step]) repair_upstream = copy.deepcopy(upstream) repair_instruction = ( "The prior synthesis violated the caller's strict JSON Schema " @@ -4883,28 +4995,47 @@ def send_synthesis( repair_started = time.perf_counter() try: repaired, final_agent = send_synthesis(repair_upstream) - except ProviderUpstreamError as exc: - if not _is_request_too_large_error(exc): + except Exception as exc: + if not _is_request_too_large_error(exc) and not isinstance(exc, EffortProfileError): self._record_failure(final_agent.id) - if final_agent.group_name and not _is_request_too_large_error(exc): + if ( + final_agent.group_name + and not _is_request_too_large_error(exc) + and not isinstance(exc, EffortProfileError) + ): self._group_router.observe_failure(final_agent.id) + # synthesis_step is a real, paid-for call that produced the + # schema-violating output prompting this repair -- it has not + # yet reached failed_attempts (that only happens further + # below, once a repair response exists to check). No + # attempted_repair_step exists to add: this call itself never + # returned a response (Devin review on #1032, round 7). + self._meter_unserved_spend( + task, + [*workflow["trace"], *failed_attempts, synthesis_step], + workflow.get("verification"), + ) raise repaired_output = provider_output(final_agent, repaired) repair_error = _structured_output_error(repaired_output, response_format) + # Built whether or not it satisfied the schema: this call reached + # the provider and consumed tokens either way, so it is a real + # accounting row even when its output is discarded below. + attempted_repair_step: dict[str, Any] = { + "id": synthesis_step["id"] + 1, + "role": "repair", + "agent_id": final_agent.id, + "subtask": "Strict JSON Schema repair", + "access": [synthesis_step["id"]], + "latency_ms": round((time.perf_counter() - repair_started) * 1000, 2), + "output": repaired_output, + } + if isinstance(repaired.get("usage"), dict): + attempted_repair_step["usage"] = _canonical_provider_usage( + repaired["usage"], responses=response_request + ) if repair_error is None: - repair_step = { - "id": synthesis_step["id"] + 1, - "role": "repair", - "agent_id": final_agent.id, - "subtask": "Strict JSON Schema repair", - "access": [synthesis_step["id"]], - "latency_ms": round((time.perf_counter() - repair_started) * 1000, 2), - "output": repaired_output, - } - if isinstance(repaired.get("usage"), dict): - repair_step["usage"] = _canonical_provider_usage( - repaired["usage"], responses=response_request - ) + repair_step = attempted_repair_step raw = repaired synthesis_output = repaired_output break @@ -4913,7 +5044,24 @@ def send_synthesis( self._record_failure(failed_agent.id) if failed_agent.group_name: self._group_router.observe_failure(failed_agent.id) + # Tagged once and reused by every exit below -- fail closed on a + # pinned model, fail closed with no same-endpoint candidate left, + # or continue the failover loop -- so these two completed calls + # (real provider spend) always reach either the persisted trace + # (``failed_attempts``, below) or an unserved-spend meter row + # (``_meter_unserved_spend``) before this iteration ends. Only the + # continue path used to do so; both raises dropped them silently + # (Devin review on #1032, round 6). + rejected_attempts = ( + {**synthesis_step, "structured_output_error": contract_error}, + {**attempted_repair_step, "structured_output_error": repair_error}, + ) if not virtual_model: + self._meter_unserved_spend( + task, + [*workflow["trace"], *failed_attempts, *rejected_attempts], + workflow.get("verification"), + ) raise ProviderResponseError( "structured synthesis and repair violated response_format" ) @@ -4929,15 +5077,127 @@ def send_synthesis( None, ) if next_agent is None: + self._meter_unserved_spend( + task, + [*workflow["trace"], *failed_attempts, *rejected_attempts], + workflow.get("verification"), + ) raise ProviderResponseError( "every eligible model on the selected endpoint violated response_format" ) + # Only reached when another model is about to be tried, so these + # two completed calls are spend with no served answer attached. + # ``structured_output_error`` marks them so an auditor reading the + # persisted trace can tell a discarded attempt from the row that + # actually produced ``answer``. + failed_attempts.extend(rejected_attempts) final_agent = next_agent synthesis_started = time.perf_counter() + # The wall clock of the call whose output is actually served, not of + # the whole loop: on the schema-repair path `synthesis_started` still + # precedes the *rejected* first synthesis, so publishing that span + # would pair an inflated duration with the repair call's own (correct, + # smaller) usage and understate the serving model's real throughput in + # both measured ledgers -- the exact honesty this observation exists + # to provide (Devin review on #1032). `repair_step["latency_ms"]` + # times exactly the call that produced the answer and reported the + # usage published beside it. + synthesis_latency_seconds = ( + repair_step["latency_ms"] / 1000 + if repair_step is not None + else time.perf_counter() - synthesis_started + ) self._record_success(final_agent.id) if final_agent.group_name: - self._group_router.observe_success( - final_agent.id, time.perf_counter() - synthesis_started + self._group_router.observe_success(final_agent.id, synthesis_latency_seconds) + # Feed the judged-quality ledger (and, when a prompt_context is + # available, the psychometric router) from this already-decided + # answer -- observation-only, matching stream_route/_finalize_batch_row: + # the verdict is recorded for future routing evidence, never branched + # on here, since send_synthesis's own retry/failover already ran to + # completion by this point. + # + # Research grounding for the mechanism itself is already recorded -- + # this wiring adds no new technique, it connects an existing + # measurement pipeline to a path that was not calling it. See + # docs/doctoring/measured-routing-evidence.md ("Real-time judging + # before returning answers"; "Multi-layer simple-structure + # measurement (fast-mlsirm)") and Baker (2001) in + # docs/papers/README.md for the IRT ability fitting. + # `usage` follows the same rule as the latency above: the served + # call's own numbers, never a discarded attempt's. `trace` is the + # spend side, and carries every completed attempt. + usage = (repair_step or synthesis_step).get("usage") + trace = [ + *workflow["trace"], + *failed_attempts, + synthesis_step, + *([repair_step] if repair_step is not None else []), + ] + # This call must stay pinned to the same eligibility constraint the + # request's own synthesis already honored (an explicit model pin, + # the free pool, a ZDR-only virtual request, or a file-replica + # subset), minus every agent this request already proved unavailable + # -- otherwise an observation-only judge call could reach a provider + # the caller's own request was never allowed to use, or a + # known-failing one whose failure would record a false-negative + # quality observation against the answer that actually succeeded + # (Devin review on #1032). + # + # The already-decided, already-served answer above must never be + # discarded just because this purely observation-only extra call + # would push spend over budget -- unlike batch_route's pre-call + # gate (which blocks a not-yet-incurred worker call before the + # caller has anything), this call happens after the response is + # fully decided, matching stream_route's own no-budget-check + # precedent for already-committed output. An exhausted budget skips + # the extra judge call outright instead of raising (Devin review on + # #1032). The gate reuses _raise_if_spend_budget_exceeded so this + # request's own not-yet-persisted spend (workflow trace + conduct's + # own verifier-role judge + synthesis + repair) counts, exactly as + # the pre-synthesis checkpoint above already counts it. + # + # Accepted limitation (Devin review on #1032, informational): this is + # a pre-call gate, not a reservation, so the one judge call it permits + # can still finish above the allowance by its own unknown output size. + # That is the gateway's budget contract everywhere -- every + # _raise_if_spend_budget_exceeded call site admits a call whose output + # length nobody knows yet -- and the judge is not special. The only + # available tightening, a hard max_output_tokens cap sized to the + # remaining allowance, would truncate the judge's structured verdict + # mid-JSON, which _model_judge_verification fails closed on: a + # false-negative quality/IRT observation would then be recorded + # against an answer that actually succeeded, corrupting the exact + # ledger this call exists to feed. Overshoot is bounded to one judge + # call admitted while spend was still inside the cap, and the next + # request's own pre-call gate stops there. + judge_budget_exceeded = False + if self.policy.realtime_judge and ( + self.budget_max_output_tokens is not None + or self.budget_max_cost_usd is not None + ): + in_flight_tokens, in_flight_cost = self._trace_budget_spend( + trace, completed_judges=in_flight_judges + ) + try: + self._raise_if_spend_budget_exceeded( + additional_output_tokens=in_flight_tokens, + additional_cost_usd=in_flight_cost, + ) + except BudgetExceededError: + judge_budget_exceeded = True + realtime_verification: dict[str, Any] | None = None + if not judge_budget_exceeded: + realtime_verification = self._realtime_route_judge( + text=task, + answer=synthesis_output, + served_id=final_agent.id, + latency_seconds=synthesis_latency_seconds, + usage=usage, + free_only=free_only, + prompt_context=prompt_context, + allowed_agent_ids=allowed_agent_ids, + excluded_agent_ids=request_exclusions or None, ) if response_request: raw.setdefault("output_text", synthesis_output) @@ -4952,11 +5212,6 @@ def send_synthesis( elif "messages" in echo: echo["messages"] = copy.deepcopy(messages) workflow_run_id = f"run_{uuid.uuid4().hex}" - trace = [ - *workflow["trace"], - synthesis_step, - *([repair_step] if repair_step is not None else []), - ] record = self._with_effort_snapshot( { "workflow_run_id": workflow_run_id, @@ -4969,6 +5224,12 @@ def send_synthesis( "trace": trace, "policy_snapshot": self.policy.as_dict(), "verification": workflow.get("verification"), + # The extra observation-only judge above is a second, + # independent provider call: conduct()'s own verifier-step + # judge already occupies "verification", so its spend needs + # its own record slot to reach the budget meter and buyer + # analytics (Devin review on #1032). + "realtime_verification": realtime_verification, } ) self._replace_workflow_run(record) @@ -5436,9 +5697,24 @@ def _raise_if_spend_budget_exceeded( raise BudgetExceededError("spend budget exceeded", detail=budget) def _trace_budget_spend( - self, trace: list[dict[str, Any]] + self, + trace: list[dict[str, Any]], + *, + completed_judges: Sequence[Mapping[str, Any]] = (), ) -> tuple[int | None, float | None]: - """Return completed provider-call spend for a workflow budget checkpoint.""" + """Return completed provider-call spend for a workflow budget checkpoint. + + ``completed_judges`` carries judge calls this request has already made + whose spend lives *outside* ``trace``: ``conduct``'s own verifier-role + judge records itself in ``workflow["verification"]``, never as a trace + row, so a checkpoint counting trace rows alone lets a request that has + already consumed its whole allowance admit yet another provider call + (Devin review on #1032). Callers pass + :meth:`_run_judge_accounting_blocks`, the same presence test the + persisted-run meter uses, and each block is priced through the same + :meth:`_judge_block_output_tokens` -- so both the token and the cost + budget fail closed when that first judge's usage is unavailable. + """ model_by_agent = {agent.id: agent.model for agent in self.agents} counts: list[tuple[int, str]] = [] for step in trace: @@ -5449,6 +5725,13 @@ def _trace_budget_spend( if count is None: return None, None counts.append((count, model)) + for verification in completed_judges: + judge_model, judge_count = self._judge_block_output_tokens( + verification, model_by_agent + ) + if judge_count is None: + return None, None + counts.append((judge_count, judge_model)) output_tokens = sum(count for count, _model in counts) if any(model not in self.price_per_million for _count, model in counts): return output_tokens, None @@ -5458,6 +5741,51 @@ def _trace_budget_spend( ) return output_tokens, round(output_cost, 6) + def _meter_unserved_spend( + self, + prompt_text: str, + trace: list[dict[str, Any]], + verification: Mapping[str, Any] | None, + ) -> None: + """Meter provider calls a rejected request already made but never served. + + Reuses ``batch_route``'s ``pending_verification`` shape (Devin review + on #961) for the same reason it exists there: ``_replace_workflow_run`` + is the sole path onto the in-memory budget meter, so completed spend + that never reaches it silently vanishes and later requests are + admitted against understated totals. The marker keeps the row out of + ``_completed_workflow_runs``, and skipping ``_run_order``/audit/ + analytics keeps it out of every user-visible listing -- a rejected + request must never surface as a finished workflow (Devin review on + #1032). + + Spend this run cannot measure flips the meter to + ``blocked_unavailable`` through the same + ``_budget_unavailable_run_ids`` path any *served* run with the same + unmeasurable steps already takes: real money left the account either + way, and fail-closed is the budget contract's deliberate answer to + not knowing how much. + """ + run_id = f"run_{uuid.uuid4().hex}" + record = self._with_effort_snapshot( + { + "workflow_run_id": run_id, + "created_at": int(time.time()), + "mode": "conduct", + "policy_mode": "conduct", + "prompt_text": prompt_text, + "answer": "", + "cache_status": "bypass", + "trace": copy.deepcopy(trace), + "policy_snapshot": self.policy.as_dict(), + "verification": dict(verification) if verification else {}, + "pending_verification": True, + } + ) + self._replace_workflow_run(record) + if self._store is not None: + self._store.save("workflow_run", run_id, record) + def batch_route(self, prompts: list[str]) -> list[dict[str, Any]]: """Route many prompts through the provider's Batch API and persist each run. @@ -6363,6 +6691,8 @@ def _realtime_route_judge( usage: dict[str, Any] | None, free_only: bool, prompt_context: str | None = None, + allowed_agent_ids: set[str] | None = None, + excluded_agent_ids: set[str] | None = None, ) -> dict[str, Any]: """Judge one direct-route answer now and feed the quality ledger. @@ -6373,25 +6703,51 @@ def _realtime_route_judge( ``None`` when the caller has no single-attempt wall-clock timing to honestly attribute to this one answer (see ``ModelGroupRouter.observe_success``); the success/failure signal is - still recorded, just not a misleading latency sample. + still recorded, just not a misleading latency sample. ``allowed_agent_ids``/ + ``excluded_agent_ids`` forward the caller's own synthesis eligibility + constraints (e.g. an explicit model pin) to the verifier selection so + this observation-only call cannot reach a provider the request itself + was never allowed to use. ``allowed_agent_ids`` additionally opens a + :func:`_request_eligibility_scope` around both the verdict call and + the ledger write, because the psychometric observation embeds the + prompt itself for routing evidence and that embedding takes no + allow-list argument of its own (Devin review on #1032). """ output_tokens = self._usage_completion_tokens(usage) def _record(accepted: bool, irt_row: tuple[int, ...] = ()) -> None: - if accepted: - self._quality_router.observe_success( - served_id, latency_seconds, output_tokens=output_tokens - ) - else: - self._quality_router.observe_failure(served_id) - if prompt_context is not None: - self._observe_contextual_quality( - prompt_context, + """Write the ledgers best-effort: evidence must never cost an answer. + + Every caller runs this *after* its answer is already produced and + paid for, and the psychometric observation writes through + ``_StateStore.save``/``prune_keyed`` (and may embed the prompt), + so a storage or embedding failure here could discard a perfectly + good completed answer and, with it, the spend accounting that + depends on the caller reaching its own persistence step (Devin + review on #1032). Losing one routing observation is the cheaper + failure by far, so it is the one that happens. + """ + try: + if accepted: + self._quality_router.observe_success( + served_id, latency_seconds, output_tokens=output_tokens + ) + else: + self._quality_router.observe_failure(served_id) + if prompt_context is not None: + self._observe_contextual_quality( + prompt_context, + served_id, + accepted=accepted, + latency_seconds=latency_seconds, + output_tokens=output_tokens, + irt_row=irt_row, + ) + except Exception: # noqa: BLE001 - observation-only ledger write + _LOGGER.warning( + "routing quality observation for %s failed; answer unaffected", served_id, - accepted=accepted, - latency_seconds=latency_seconds, - output_tokens=output_tokens, - irt_row=irt_row, + exc_info=True, ) if not self.policy.realtime_judge: @@ -6402,18 +6758,23 @@ def _record(accepted: bool, irt_row: tuple[int, ...] = ()) -> None: "judge": "model", } fallback_report = {"verifier_output": answer} - base = self._model_judge_verification( - text, fallback_report, free_only=free_only - ) - accepted = bool(base.get("accepted")) - raw_irt_row = base.get("judge_irt_row") - irt_row = ( - tuple(raw_irt_row) - if isinstance(raw_irt_row, list) - and all(type(value) is int and value in (0, 1) for value in raw_irt_row) - else () - ) - _record(accepted, irt_row) + with _request_eligibility_scope(allowed_agent_ids): + base = self._model_judge_verification( + text, + fallback_report, + free_only=free_only, + allowed_agent_ids=allowed_agent_ids, + excluded_agent_ids=excluded_agent_ids, + ) + accepted = bool(base.get("accepted")) + raw_irt_row = base.get("judge_irt_row") + irt_row = ( + tuple(raw_irt_row) + if isinstance(raw_irt_row, list) + and all(type(value) is int and value in (0, 1) for value in raw_irt_row) + else () + ) + _record(accepted, irt_row) return base @staticmethod @@ -7153,39 +7514,254 @@ def _cache_put(self, cache: OrderedDict[str, Any], key: str, value: Any) -> None cache.popitem(last=False) def _embedding_agent_id(self) -> str | None: - """First measured embedding-capable member id, or None when unconfigured.""" + """First embedding-capable member this request may reach, or None. + + Every embedding this gateway makes for routing evidence + (:meth:`_embed_cached`, :meth:`_descriptor_vector_cached`) sends + request-derived text -- the prompt itself, in the psychometric and + semantic-affinity paths -- to the returned provider. It therefore + honors the active request's own eligibility exactly like the serving + path does: an explicit model pin, the free pool, or a file-replica + subset all narrow this choice, and ZDR/endpoint pinning already + narrow it inside :meth:`_ranked_agents`. Without this, an + observation-only call could hand a restricted request's prompt to an + unrelated provider the caller was never allowed to use (Devin review + on #1032). + + An eligible agent's own configured endpoint counts as reachable: that + provider is already serving this request's content, so an embedding + deployment behind the same endpoint discloses nothing new. Anything + else yields None, and every caller already degrades to + declaration-only evidence when embedding is unavailable. + + That endpoint-sharing allowance is privacy-only, though, and says + nothing about cost: when the eligible set is entirely free (exactly + what ``free_only`` builds -- see the ``allowed_agent_ids`` branch in + ``_orchestrated_provider_completion``), a *paid* embedding deployment + co-located with a free agent must not ride along, or a free request + incurs real, unmetered spend outside every budget check (Devin + review on #1032, round 6). An explicit paid pin's own eligible set is + never all-free, so its endpoint fallback keeps admitting a co-located + deployment regardless of price, unchanged from before this gate. + """ try: - return self.select_capability_agent("embedding").id + candidates = self._capability_agents("embedding") except (RuntimeError, ValueError): return None + eligible = _REQUEST_ELIGIBLE_AGENT_IDS.get() + if eligible is None: + return candidates[0].id + eligible_agents = [agent for agent in self.agents if agent.id in eligible] + endpoints = { + agent.base_url.rstrip("/").casefold() for agent in eligible_agents + } + free_scope = bool(eligible_agents) and all( + self._is_free_agent(agent) for agent in eligible_agents + ) + return next( + ( + agent.id + for agent in candidates + if agent.id in eligible + or ( + agent.base_url.rstrip("/").casefold() in endpoints + and (not free_scope or self._is_free_agent(agent)) + ) + ), + None, + ) - def _embed_cached(self, text: str) -> list[float] | None: - """Embedding vector for text via the configured embedding member; None on failure.""" + def _meter_embedding_spend( + self, embedder: ModelAgent, prompt_tokens: int | None + ) -> None: + """Record one cache-miss embedding call's real provider spend. + + Called only from the miss branch of ``_embed_cached`` and + ``_descriptor_vector_cached`` -- a cache hit reuses a vector this + already metered and must never be counted twice -- and only when + ``embed_with_usage`` returned an authoritative ``prompt_tokens``; + the mock transport and any provider that omits usage both return + None, and this stays a graceful no-op rather than guess a count + (Devin review on #1032). + + These calls have no enclosing task, record, or workflow_run_id to + attach to -- the public ``select_model_group_members`` can trigger + one with none of those in scope at all -- so this mints its own run + and writes it straight onto the meter via + :meth:`_replace_workflow_run`, the same primitive every other + spend-recording path in this file uses; it works identically + whether or not a live request exists. + + Reuses the ``pending_verification: True`` shape ``batch_route`` and + :meth:`_meter_unserved_spend` already established for exactly this + split: :meth:`_replace_workflow_run`/``spend_analytics`` iterate + ``_workflow_runs`` directly, so the real spend is always counted, + while the marker keeps this synthetic row out of ``_run_order`` and + every consumer that goes through :meth:`_completed_workflow_runs` + instead (run counts, analytics KPIs, admin listings) -- it was + never a request, so it must never look like one. + ``completion_tokens`` is explicit ``0`` (embeddings have no + completion), which keeps this row "reported" rather than + "unavailable" so it can never flip a real run's -- or the whole + meter's -- availability. + + That per-step availability isn't the whole story: the budget + meter also requires every model appearing in *any* run to be + priced whenever ``budget_max_cost_usd`` is set, with no exception + for a model whose output is provably always zero tokens. An + operator who adds a cost budget without ever having priced their + embedder (embedding spend was invisible before this fix, so there + was never a reason to) would otherwise see the *entire* meter flip + to ``blocked_unavailable`` the first time this fires. Skip + metering in that one case -- true zero-cost either way -- instead + of writing a row that blinds unrelated requests' enforcement. + + :meth:`_replace_workflow_run` is pure in-memory dict/Decimal + arithmetic under ``_budget_spend_lock`` -- nothing that plausibly + raises in normal operation, and shared by every other of its call + sites, so it is left unguarded here: a real bug in it should stay + loud, not get mislabeled a durability issue. ``self._store.save`` + is a real sqlite3 write, so a state-store outage there is caught + and logged rather than propagated -- reusing the same + ``except Exception: _LOGGER.warning(..., exc_info=True)`` + best-effort-write convention :meth:`_observe_contextual_quality` + already established for its own ``self._store.save`` call. The + failure is deliberately *not* routed through + ``_budget_unavailable_run_ids``: that flag means "this record's + usage is not derivable," which is a measurement gap, not a + durability one -- an embedding row's ``completion_tokens`` is + always the explicit ``0`` above, so its usage is always fully + known regardless of whether the write behind it later succeeds, + and by the time the write is attempted :meth:`_replace_workflow_run` + has already updated the in-process spend meter, so enforcement is + never blind to it. ``_budget_unavailable_run_ids`` is a blunt, + global switch -- flipping it here would freeze real + chat-completion budget enforcement org-wide over one transient + disk hiccup on a best-effort, zero-completion-token routing-evidence + row. The only thing actually lost on a write failure is *restart* + survivability of that one row; logging it is the honest response. + """ + if prompt_tokens is None: + return + if ( + self.budget_max_cost_usd is not None + and embedder.model not in self.price_per_million + ): + return + run_id = f"run_{uuid.uuid4().hex}" + record = self._with_effort_snapshot( + { + "workflow_run_id": run_id, + "created_at": int(time.time()), + "mode": "embedding", + "policy_mode": "embedding", + "prompt_text": "", + "answer": "", + "cache_status": "bypass", + "trace": [ + { + "id": 0, + "role": "embedding", + "agent_id": embedder.id, + "subtask": "Routing-evidence embedding", + "access": [], + "output": "", + "model_name": embedder.model, + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": 0, + }, + } + ], + "policy_snapshot": self.policy.as_dict(), + "pending_verification": True, + } + ) + self._replace_workflow_run(record) + if self._store is not None: + try: + self._store.save("workflow_run", run_id, record) + except Exception: # noqa: BLE001 - durable write is best-effort here + _LOGGER.warning( + "embedding spend persistence for run %s failed; in-memory " + "meter already updated, spend not durable across a restart", + run_id, + exc_info=True, + ) + + def _embed_cached( + self, text: str, embedding_member: str | None = None + ) -> list[float] | None: + """Embedding vector for text via the configured embedding member; None on failure. + + Keyed on two independent things, because either alone leaks: + + * :func:`_request_evidence_partition` -- the active request's own + restrictions, so a cache hit can only ever return a vector some + provider *this* request may itself reach produced. Without it a + restricted request read the cache before :meth:`_embedding_agent_id` + was consulted at all, and an earlier unrestricted request's vector + -- the routing evidence baked into it included -- crossed the + isolation boundary (Devin review on #1032). + * the resolved embedding member itself, which is the identity of the + *vector space* the entry lives in. An identical restriction shape + can still resolve to a different embedding model/deployment later + (an agent-pool change, or :meth:`_measured_member_order` reordering + the members of one eligible group), and a cosine between two + different embedding spaces is meaningless -- it would silently + corrupt affinity-based routing rather than fail (a later Devin + review on the same PR; a refinement of the eligibility partition + above, not a duplicate of it). + + ``embedding_member`` lets a caller that makes several comparable + embeddings resolve the member once and pin every vector in one + comparison to the same space; ``None`` resolves it here. + """ + if embedding_member is None: + embedding_member = self._embedding_agent_id() + if embedding_member is None: + return None digest = hashlib.sha256( - f"{_request_endpoint_partition()}\x1f{text}".encode("utf-8") + "\x1f".join( + (_request_evidence_partition(), f"embedder:{embedding_member}", text) + ).encode("utf-8") ).hexdigest() with self._evidence_lock: cached = self._task_vector_cache.get(digest) if cached is not None: return cached - embedding_member = self._embedding_agent_id() - if embedding_member is None: - return None try: - vectors = self.client.embed(self._agent(embedding_member), [text]) + embedder = self._agent(embedding_member) + vectors, prompt_tokens = self.client.embed_with_usage(embedder, [text]) except Exception: # noqa: BLE001 - similarity is best-effort evidence return None vector = vectors[0] if vectors else None if vector is not None: self._cache_put(self._task_vector_cache, digest, vector) + self._meter_embedding_spend(embedder, prompt_tokens) return vector - def _descriptor_vector_cached(self, agent: ModelAgent) -> list[float] | None: - """Cached embedding of one agent's operator-declared metadata document.""" + def _descriptor_vector_cached( + self, agent: ModelAgent, embedding_member: str | None = None + ) -> list[float] | None: + """Cached embedding of one agent's operator-declared metadata document. + + Keyed exactly like :meth:`_embed_cached`, on both the request's + restrictions and the resolved embedding member: both halves of a + semantic affinity have to come from a provider the active request may + reach, *and* from the same vector space, or the restricted request + still routes on an ineligible provider's evidence -- or on a cosine + between two unrelated embedding spaces (Devin reviews on #1032). + """ + if embedding_member is None: + embedding_member = self._embedding_agent_id() + if embedding_member is None: + return None fingerprint = hashlib.sha256( "\x1f".join( [ - _request_endpoint_partition(), + _request_evidence_partition(), + f"embedder:{embedding_member}", agent.id, self._agent_descriptor_text(agent), ] @@ -7195,18 +7771,17 @@ def _descriptor_vector_cached(self, agent: ModelAgent) -> list[float] | None: cached = self._descriptor_vector_cache.get(fingerprint) if cached is not None: return cached - embedding_member = self._embedding_agent_id() - if embedding_member is None: - return None try: - vectors = self.client.embed( - self._agent(embedding_member), [self._agent_descriptor_text(agent)] + embedder = self._agent(embedding_member) + vectors, prompt_tokens = self.client.embed_with_usage( + embedder, [self._agent_descriptor_text(agent)] ) except Exception: # noqa: BLE001 - similarity is best-effort evidence return None vector = vectors[0] if vectors else None if vector is not None: self._cache_put(self._descriptor_vector_cache, fingerprint, vector) + self._meter_embedding_spend(embedder, prompt_tokens) return vector def _semantic_affinities( @@ -7217,16 +7792,26 @@ def _semantic_affinities( Returns ``{agent_id: float|None}``; all values are None whenever there is no task text, no embedding-capable member, or embedding transport fails -- callers then fall back to declaration-only ordering. + + The embedding member is resolved once, here, and pinned into every + vector this comparison uses: a cosine is only meaningful between two + vectors from the same embedding space, and resolving per call would + both re-rank the pool once per candidate and leave the task vector + free to come from a different space than a descriptor vector when the + pool changes mid-comparison (Devin review on #1032). """ stripped = text.strip() if isinstance(text, str) else "" if not stripped or not agents: return {agent.id: None for agent in agents} - task_vector = self._embed_cached(stripped) + embedding_member = self._embedding_agent_id() + if embedding_member is None: + return {agent.id: None for agent in agents} + task_vector = self._embed_cached(stripped, embedding_member) if task_vector is None: return {agent.id: None for agent in agents} affinities: dict[str, float | None] = {} for agent in agents: - descriptor_vector = self._descriptor_vector_cached(agent) + descriptor_vector = self._descriptor_vector_cached(agent, embedding_member) affinities[agent.id] = ( None if descriptor_vector is None @@ -8601,41 +9186,74 @@ def _run_budget_output_by_model( if output_tokens is None: return {}, False output_by_model[model] = output_by_model.get(model, 0) + output_tokens - verification = record.get("verification") - if isinstance(verification, Mapping): - judge_agent_id = verification.get("judge_agent_id") - if judge_agent_id is not None: - # A completed judge call (judge_agent_id is only ever set - # once one has) whose response carried no valid usage must - # still count toward the budget meter, or a run of - # unmeasured judge calls could exceed a spend cap this - # conservative-by-design check exists to enforce. Fall back - # to the same estimate-from-real-text _step_output_tokens - # already applies to worker steps with no reported usage - # (Devin review on #961: an earlier revision of this fix - # fabricated a "reported" zero-token dict instead). Estimate - # from judge_output_text (the judge's own generated - # rationale), not verifier_output (the worker answer it was - # judging) -- a second Devin review on this same fallback - # caught estimating from the wrong side of the call. - judge_model = verification.get("judge_model") or model_by_agent.get( - judge_agent_id, "unknown" - ) - completion_tokens, _judge_reported = _step_output_tokens( - { - "usage": verification.get("judge_usage"), - "output": verification.get("judge_output_text", ""), - }, - self.token_counter, - judge_model, - ) - if completion_tokens is None: - return {}, False - output_by_model[judge_model] = ( - output_by_model.get(judge_model, 0) + completion_tokens - ) + for verification in self._run_judge_accounting_blocks(record): + judge_model, completion_tokens = self._judge_block_output_tokens( + verification, model_by_agent + ) + if completion_tokens is None: + return {}, False + output_by_model[judge_model] = ( + output_by_model.get(judge_model, 0) + completion_tokens + ) return output_by_model, True + def _judge_block_output_tokens( + self, verification: Mapping[str, Any], model_by_agent: Mapping[str, str] + ) -> tuple[str, int | None]: + """Return one completed judge call's model and authoritative output tokens. + + A completed judge call (``judge_agent_id`` is only ever set once one + has) whose response carried no valid usage must still count toward + the budget meter, or a run of unmeasured judge calls could exceed a + spend cap this conservative-by-design check exists to enforce. Fall + back to the same estimate-from-real-text ``_step_output_tokens`` + already applies to worker steps with no reported usage (Devin review + on #961: an earlier revision of this fix fabricated a "reported" + zero-token dict instead). Estimate from ``judge_output_text`` (the + judge's own generated rationale), not ``verifier_output`` (the worker + answer it was judging) -- a second Devin review on this same fallback + caught estimating from the wrong side of the call. + + Shared by the persisted-run meter (:meth:`_run_budget_output_by_model`) + and the in-flight checkpoint (:meth:`_trace_budget_spend`), so a judge + call already made by *this* request is accounted exactly as the same + call is once persisted. ``None`` tokens means unmeasurable; both + callers fail closed on it. + """ + judge_model = verification.get("judge_model") or model_by_agent.get( + verification["judge_agent_id"], "unknown" + ) + output_tokens, _judge_reported = _step_output_tokens( + { + "usage": verification.get("judge_usage"), + "output": verification.get("judge_output_text", ""), + }, + self.token_counter, + judge_model, + ) + return judge_model, output_tokens + + @staticmethod + def _run_judge_accounting_blocks( + record: Mapping[str, Any], + ) -> list[Mapping[str, Any]]: + """Every completed judge call recorded on one run, in record order. + + ``verification`` holds the workflow's own verifier-step judge; + ``realtime_verification`` holds the separate observation-only + realtime judge ``_orchestrated_provider_completion`` fires after + synthesis. Both are real, already-incurred provider calls, so both + must reach the budget meter and buyer-facing spend analytics + (Devin review on #1032). ``judge_agent_id`` is set only once a call + actually completed, so it is the presence test for both. + """ + blocks: list[Mapping[str, Any]] = [] + for key in ("verification", "realtime_verification"): + block = record.get(key) + if isinstance(block, Mapping) and block.get("judge_agent_id") is not None: + blocks.append(block) + return blocks + def _replace_workflow_run(self, record: dict[str, Any]) -> None: """Store one run and update its constant-time budget meter atomically.""" model_by_agent = {agent.id: agent.model for agent in self.candidates} @@ -8742,13 +9360,8 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> bucket["output_tokens"] += effective total_output_tokens += effective - verification = run.get("verification") - judge_agent_id = ( - verification.get("judge_agent_id") - if isinstance(verification, Mapping) - else None - ) - if judge_agent_id is not None: + for verification in self._run_judge_accounting_blocks(run): + judge_agent_id = verification["judge_agent_id"] # A completed judge call (judge_agent_id is only ever set # once one has) must stay visible here even when its # response carried no valid usage, or a real, incurred diff --git a/tests/test_measured_routing_evidence.py b/tests/test_measured_routing_evidence.py index 5a0c35a30..28264d70a 100644 --- a/tests/test_measured_routing_evidence.py +++ b/tests/test_measured_routing_evidence.py @@ -320,7 +320,7 @@ def fake_invoke(primary, messages, **kwargs): monkeypatch.setattr(orchestrator, "_invoke", fake_invoke) - def judge(text, fallback, *, free_only=False): + def judge(text, fallback, *, free_only=False, **_ignored): accepted = "strong" in fallback["verifier_output"] return { "accepted": accepted, diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index cfa249cea..1f700a7ff 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -291,6 +291,16 @@ def chat(self, agent: ModelAgent, messages: list, **kwargs: object) -> str: # t def test_explicit_structured_group_model_pins_every_provider_call() -> None: + """Every provider call, including the post-synthesis realtime judge, stays pinned. + + The realtime judge that ``_orchestrated_provider_completion`` now calls + once synthesis succeeds (observation-only) picks its own verifier from + the same group; the transport ledger already favors the member that + just served (``_group_router.observe_success`` ran immediately before), + so it lands on ``selected_member`` too -- covered here as a sixth + ``evidence_or_judge`` call after synthesis. + """ + class _RecordingClient(_ScriptedClient): def __init__(self) -> None: super().__init__('{"decision":"ACCEPT","reason":"Exact judge passed."}') @@ -331,6 +341,7 @@ def proxy_send(self, agent: ModelAgent, endpoint: str, body: dict) -> dict: # t assert client.calls_by_kind == [ *(('evidence_or_judge', selected.id) for _ in range(5)), ("synthesis", selected.id), + ("evidence_or_judge", selected.id), ] diff --git a/tests/test_orchestrated_completion_judge_observation.py b/tests/test_orchestrated_completion_judge_observation.py new file mode 100644 index 000000000..ae1f7529a --- /dev/null +++ b/tests/test_orchestrated_completion_judge_observation.py @@ -0,0 +1,1911 @@ +"""Realtime fast-mlsirm judge observation wired into structured synthesis. + +Covers ``_orchestrated_provider_completion``'s one success point calling +``_realtime_route_judge`` for its recording side effect only (quality ledger +and psychometric routing evidence), never branching on the verdict: + +- the call receives the actually-served answer/agent and the already + canonicalized usage (not the raw provider dict), and records one quality + success; +- ``policy.realtime_judge = False`` skips the judge call and every ledger + write entirely, exactly as the disabled route-path contract already does; +- the observation genuinely reaches ``PsychometricRoutingEvidence`` and can + move an evidenced candidate ahead of a higher-static-priority one on a + later ranking call, proving the routing gap this wiring closes; +- a request pinned to one explicit model keeps that pin through the judge + call *and* through the prompt embedding the observation performs, so + neither can reach a provider the request itself was never allowed to use; +- a budget rejection meters the provider calls this request already made + instead of forgetting them, an observation-ledger write failure never + costs the caller an already-generated answer, and a schema-repaired + answer publishes the repair call's own latency rather than a span that + also covers the synthesis attempt that was thrown away; +- a model whose structured synthesis *and* repair both failed before a + different member of the same virtual pool succeeded still reaches every + accounting path (trace, budget checkpoint, persisted run, spend analytics) + without polluting the served answer's own latency/usage attribution; +- a routing-evidence cache entry is never reused across a change of + embedding model, whose vectors live in a different space entirely. +""" + +from __future__ import annotations + +import logging +import sqlite3 +import sys +import time +from dataclasses import replace +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402 +from contextual_orchestrator.orchestrator import ( # noqa: E402 + BudgetExceededError, + EffortProfileError, + ProviderResponseError, + ProviderUpstreamError, + ReasoningEffortProfile, + _request_eligibility_scope, + _request_evidence_partition, +) + +_STUB_CONDUCT = { + "mode": "conduct", + "answer": "evidence", + "trace": [], + "verification": {"accepted": True, "reason": "test", "verifier_output": ""}, +} + + +class _ResponsesUsageClient: + """One fixed Responses-shaped answer with Responses-API usage keys only.""" + + def __init__(self, *, content: str, usage: dict[str, int]) -> None: + self._content = content + self._usage = usage + self.calls: list[tuple[str, str]] = [] + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Record the attempt and return a fixed provider-shaped response.""" + del payload + self.calls.append((agent.id, endpoint)) + return { + "id": "resp_test", + "object": "response", + "model": agent.model, + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": self._content}], + } + ], + "usage": dict(self._usage), + } + + proxy_send = proxy_send_once + + +def test_orchestrated_completion_wires_realtime_judge_with_canonicalized_usage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The new call gets the served answer/agent and canonicalized usage, and records success.""" + agent = ModelAgent("solo_agent", "mock-model", tags=("reasoning",)) + client = _ResponsesUsageClient( + content="final answer", + usage={"input_tokens": 7, "output_tokens": 13, "total_tokens": 20}, + ) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + captured: dict[str, object] = {} + original_judge = TaskOrchestrator._realtime_route_judge + + def _spy(self: TaskOrchestrator, **kwargs: object) -> dict[str, Any]: + captured.update(kwargs) + return original_judge(self, **kwargs) + + monkeypatch.setattr(TaskOrchestrator, "_realtime_route_judge", _spy) + monkeypatch.setattr( + orchestrator, + "_model_judge_verification", + lambda task, fallback, *, free_only=False, **_ignored: { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + }, + ) + + messages = [{"role": "user", "content": "hello world"}] + expected_prompt_context = TaskOrchestrator._prompt_interaction(messages) + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + assert result["output_text"] == "final answer" + assert client.calls == [("solo_agent", "responses")] + + assert captured["text"] == "hello world" + assert captured["answer"] == "final answer" + assert captured["served_id"] == "solo_agent" + assert captured["free_only"] is False + assert captured["prompt_context"] == expected_prompt_context + assert isinstance(captured["latency_seconds"], float) + assert captured["latency_seconds"] >= 0 + # The raw provider dict only has Responses-API keys; the judge must + # receive the already-canonicalized usage (with the completion_tokens + # alias _usage_completion_tokens actually reads), not raw.get("usage"). + assert captured["usage"] == { + "input_tokens": 7, + "output_tokens": 13, + "total_tokens": 20, + "prompt_tokens": 7, + "completion_tokens": 13, + } + + quality = orchestrator._quality_router.member_report("solo_agent") + assert quality["success_count"] == 1 + + +def test_disabled_realtime_judge_skips_judge_call_and_ledger_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``policy.realtime_judge = False`` means no judge call and no ledger write.""" + from dataclasses import replace + + agent = ModelAgent("worker_agent", "mock", tags=("reasoning",)) + orchestrator = TaskOrchestrator([agent]) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + def _explode(*args: object, **kwargs: object) -> dict[str, Any]: + raise AssertionError("model judge must not be called when realtime_judge is disabled") + + monkeypatch.setattr(orchestrator, "_model_judge_verification", _explode) + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + assert result["output_text"] == "[worker_agent] chat-mock" + assert orchestrator._quality_router.member_observation_count("worker_agent") == 0 + assert orchestrator._psychometric_router.has_observations() is False + + +def test_orchestrated_completion_observation_flows_into_psychometric_reordering( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A served answer's real observation can later re-rank a synthesizer partition.""" + import fast_mlsirm + + alpha = ModelAgent("candidate_alpha", "mock", tags=("reasoning",), priority=50) + beta = ModelAgent("candidate_beta", "mock", tags=("reasoning",), priority=1) + orchestrator = TaskOrchestrator([alpha, beta]) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + messages = [{"role": "user", "content": "shared prompt text"}] + task = orchestrator._latest_user_text(messages) + prompt_context = orchestrator._prompt_interaction(messages) + + assert orchestrator._psychometric_router.has_observations() is False + baseline = orchestrator._ranked_agents(task, "synthesizer", prompt_context=prompt_context) + assert [candidate.id for candidate in baseline] == ["candidate_alpha", "candidate_beta"] + + monkeypatch.setattr( + orchestrator, + "_model_judge_verification", + lambda task, fallback, *, free_only=False, **_ignored: { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + }, + ) + + # Force the lower-static-priority agent to be the one that actually + # serves, so a real evidence-first reorder (not the static order it + # already had) is what proves the observation moved routing. + result = orchestrator.proxy_completion( + {"input": "shared prompt text", "_required_agent_id": "candidate_beta"}, + endpoint="responses", + single_agent=False, + ) + assert result["output_text"] == "[candidate_beta] chat-mock" + assert orchestrator._psychometric_router.has_observations() is True + + # The fast-mlsirm native fit legitimately refuses to converge on a + # single-item response matrix (see PsychometricRoutingEvidence._fit_locked); + # stub only that numeric boundary -- exactly as + # test_fast_mlsirm_fit_uses_judge_acceptance_item_for_context_score does -- + # so the real observation recorded above is what drives a real re-rank. + class _Result: + convergence_status = "converged" + params = object() + model = "MLSRM" + + def fake_fit_experiment(fit_callable: object, responses: object, item_type: str, **kwargs: object) -> _Result: + del fit_callable, responses, item_type, kwargs + return _Result() + + def fake_predict(_params: object, factor_id: object, *, model: str) -> np.ndarray: + del _params, model + return np.array([[0.99]] * len(factor_id)).reshape(1, len(factor_id)) + + monkeypatch.setattr(fast_mlsirm, "fit_irt_experiment", fake_fit_experiment) + monkeypatch.setattr(fast_mlsirm, "predict_proba", fake_predict) + + reordered = orchestrator._ranked_agents(task, "synthesizer", prompt_context=prompt_context) + assert [candidate.id for candidate in reordered] == ["candidate_beta", "candidate_alpha"] + + +def test_explicit_model_pin_constrains_realtime_judge_to_selected_agent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An explicit model pin stays pinned through the observation-only judge too. + + Devin review (PR #1032): ``_realtime_route_judge`` must forward the same + ``allowed_agent_ids`` the request's own synthesis was already constrained + to, so this extra call can never reach an unrelated, higher-ranked + verifier the caller never selected. + """ + pinned = ModelAgent("pinned_agent", "mock", tags=("reasoning",), priority=1) + unrelated = ModelAgent("unrelated_verifier", "mock", tags=("reasoning",), priority=100) + orchestrator = TaskOrchestrator([pinned, unrelated]) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + captured: dict[str, object] = {} + + def _spy( + task: str, + fallback: dict[str, Any], + *, + free_only: bool = False, + allowed_agent_ids: set[str] | None = None, + excluded_agent_ids: set[str] | None = None, + ) -> dict[str, Any]: + captured["allowed_agent_ids"] = allowed_agent_ids + return { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + } + + monkeypatch.setattr(orchestrator, "_model_judge_verification", _spy) + + result = orchestrator.proxy_completion( + {"input": "hello world", "_required_agent_id": "pinned_agent"}, + endpoint="responses", + single_agent=False, + ) + + assert result["output_text"] == "[pinned_agent] chat-mock" + assert captured["allowed_agent_ids"] == {"pinned_agent"} + + +def test_exhausted_budget_skips_extra_realtime_judge_call_but_keeps_answer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An exhausted budget skips the extra judge call, never the already-good answer. + + Devin review (PR #1032): this purely observation-only call's own spend + must never discard the already-decided response above. Unlike + ``batch_route``'s pre-call gate (which blocks a not-yet-incurred worker + call before the caller has anything), this call happens after the + response is fully decided -- an exhausted budget skips it outright + instead of raising and losing an already-good answer. + + The budget here is never spent by a previous run: the whole allowance is + consumed by *this* request's own not-yet-persisted workflow + synthesis + spend, which is exactly the case ``budget_status()`` alone cannot see + (Devin review on #1032). + """ + agent = ModelAgent("worker_agent", "mock", tags=("reasoning",)) + orchestrator = TaskOrchestrator([agent]) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator.budget_max_output_tokens = 1 + assert orchestrator.budget_status()["exceeded"] is False + + def _explode(*args: object, **kwargs: object) -> dict[str, Any]: + raise AssertionError("realtime judge must be skipped once budget is already exceeded") + + monkeypatch.setattr(orchestrator, "_model_judge_verification", _explode) + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + assert result["output_text"] == "[worker_agent] chat-mock" + assert orchestrator._quality_router.member_observation_count("worker_agent") == 0 + assert orchestrator._psychometric_router.has_observations() is False + + +def test_realtime_judge_excludes_agents_this_request_already_proved_unavailable() -> None: + """A failed-over-away agent must not be picked as this request's judge. + + Devin review (PR #1032): ``request_exclusions`` holds every agent this + request already proved unavailable. Feeding the judge from + ``allowed_agent_ids`` alone leaves such an agent eligible; if it fails + the judge call, the resulting failure records a false-negative quality + observation against the answer that actually succeeded -- corrupting the + exact measurement this wiring exists to produce. + """ + + class _FirstAgentAlwaysFails: + """Fail every call to ``broken_agent``; serve ``healthy_agent`` normally.""" + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Raise for the broken agent, otherwise return a fixed answer.""" + del endpoint, payload + if agent.id == "broken_agent": + raise ProviderUpstreamError( + agent_id=agent.id, + model=agent.model, + error_code="model_not_found", + message="provider rejected the request with HTTP 404", + client_status=404, + provider_status=404, + retryable=False, + transport="passthrough", + ) + return { + "id": "resp_test", + "object": "response", + "model": agent.model, + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "served answer"}], + } + ], + } + + proxy_send = proxy_send_once + + broken = ModelAgent("broken_agent", "mock", tags=("reasoning",), priority=100) + healthy = ModelAgent("healthy_agent", "mock", tags=("reasoning",), priority=1) + orchestrator = TaskOrchestrator([broken, healthy], client=_FirstAgentAlwaysFails()) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + captured: dict[str, object] = {} + + def _spy( + task: str, + fallback: dict[str, Any], + *, + free_only: bool = False, + allowed_agent_ids: set[str] | None = None, + excluded_agent_ids: set[str] | None = None, + ) -> dict[str, Any]: + captured["excluded_agent_ids"] = excluded_agent_ids + captured["judge"] = next( + agent.id + for agent in orchestrator._ranked_agents(task, "verifier") + if allowed_agent_ids is None or agent.id in allowed_agent_ids + if excluded_agent_ids is None or agent.id not in excluded_agent_ids + ) + return { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + } + + orchestrator._model_judge_verification = _spy # type: ignore[method-assign] + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + assert result["output_text"] == "served answer" + assert captured["excluded_agent_ids"] == {"broken_agent"} + # Without the exclusion the higher-priority broken agent wins verifier + # ranking and would have taken the judge call. + assert captured["judge"] == "healthy_agent" + + +class _EmbeddingSpyClient: + """Serve a fixed Responses answer and record every embedding provider call.""" + + def __init__(self) -> None: + self.embed_calls: list[str] = [] + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Return one fixed provider-shaped response for any agent.""" + del endpoint, payload + return { + "id": "resp_test", + "object": "response", + "model": agent.model, + "output": [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "served answer"}], + } + ], + } + + proxy_send = proxy_send_once + + def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]: + """Record which provider was asked to embed request-derived text.""" + self.embed_calls.append(agent.id) + return [[0.1, 0.2, 0.3] for _ in texts] + + def embed_with_usage( + self, agent: ModelAgent, texts: list[str] + ) -> tuple[list[list[float]], int | None]: + """Same fixed vectors; no authoritative usage (mirrors the mock transport).""" + return self.embed(agent, texts), None + + +def _pinned_pool() -> tuple[ModelAgent, ModelAgent]: + """One pinnable chat model plus an embedding deployment on another provider.""" + return ( + ModelAgent( + "pinned_agent", + "pinned-model", + base_url="https://pinned.example/v1", + tags=("reasoning",), + priority=1, + ), + ModelAgent( + "unrelated_agent", + "unrelated-model", + base_url="https://unrelated.example/v1", + tags=("reasoning", "embedding"), + priority=100, + ), + ) + + +def test_explicit_structured_model_pin_constrains_realtime_judge( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A caller-named explicit model pins the judge exactly like ``_required_agent_id``. + + Devin review (PR #1032): the judge's allow-list is derived from + ``_required_agent_id`` alone, so an explicit structured model request -- + which pins synthesis to that one agent just as hard -- left the judge + unrestricted and could send the prompt and served answer to an unrelated + provider. + """ + pinned, unrelated = _pinned_pool() + orchestrator = TaskOrchestrator([pinned, unrelated], client=_EmbeddingSpyClient()) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + captured: dict[str, object] = {} + + def _spy( + task: str, + fallback: dict[str, Any], + *, + free_only: bool = False, + allowed_agent_ids: set[str] | None = None, + excluded_agent_ids: set[str] | None = None, + ) -> dict[str, Any]: + captured["allowed_agent_ids"] = allowed_agent_ids + captured["judge"] = next( + ( + agent.id + for agent in orchestrator._ranked_agents(task, "verifier") + if allowed_agent_ids is None or agent.id in allowed_agent_ids + ), + None, + ) + return { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + } + + monkeypatch.setattr(orchestrator, "_model_judge_verification", _spy) + + result = orchestrator.proxy_completion( + {"input": "hello world", "model": "pinned-model"}, + endpoint="responses", + single_agent=False, + ) + + assert result["output_text"] == "served answer" + # No _required_agent_id anywhere -- the pin came from `model` alone. + assert captured["allowed_agent_ids"] == {"pinned_agent"} + # Without the pin the higher-priority unrelated provider wins verifier + # ranking and would have taken the judge call. + assert captured["judge"] == "pinned_agent" + + +def test_explicit_structured_model_pin_blocks_ineligible_prompt_embedding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The observation's prompt embedding honors the request's own eligibility. + + Devin review (PR #1032): the contextual observation embeds the prompt for + routing evidence, and that embedding took no allow-list of its own -- so + a request pinned to one provider could still hand its prompt to an + unrelated embedding provider. The pinned request must reach no embedding + provider at all here; the same pool with an unpinned virtual request + still does, proving the block is the pin and not a missing deployment. + """ + pinned, unrelated = _pinned_pool() + + def _accept( + task: str, fallback: dict[str, Any], **_ignored: object + ) -> dict[str, Any]: + """Stand in for the judge verdict so only the embedding path is measured.""" + del task + return { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + } + + pinned_client = _EmbeddingSpyClient() + pinned_orchestrator = TaskOrchestrator([pinned, unrelated], client=pinned_client) + pinned_orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + monkeypatch.setattr(pinned_orchestrator, "_model_judge_verification", _accept) + + pinned_orchestrator.proxy_completion( + {"input": "confidential prompt", "model": "pinned-model"}, + endpoint="responses", + single_agent=False, + ) + assert pinned_client.embed_calls == [] + assert pinned_orchestrator._psychometric_router.has_observations() is True + + open_client = _EmbeddingSpyClient() + open_orchestrator = TaskOrchestrator([pinned, unrelated], client=open_client) + open_orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + monkeypatch.setattr(open_orchestrator, "_model_judge_verification", _accept) + + open_orchestrator.proxy_completion( + {"input": "confidential prompt"}, endpoint="responses", single_agent=False + ) + assert "unrelated_agent" in open_client.embed_calls + + +def test_eligibility_scope_narrows_embedding_to_reachable_providers() -> None: + """Embedding eligibility follows the provider, not the exact model id. + + The narrowing must block an unrelated provider without silently killing + routing evidence for every restricted request: an embedding deployment + behind an endpoint the request already reaches stays usable, an empty + allow-list yields no embedding at all, and no scope keeps the + unrestricted pick. + """ + chat = ModelAgent( + "chat_agent", + "chat-model", + base_url="https://same.example/v1", + tags=("reasoning",), + ) + same_provider = ModelAgent( + "same_embedder", + "embed-model", + base_url="https://same.example/v1/", + tags=("embedding",), + ) + other_provider = ModelAgent( + "other_embedder", + "other-embed-model", + base_url="https://other.example/v1", + tags=("embedding",), + priority=100, + ) + orchestrator = TaskOrchestrator([chat, same_provider, other_provider]) + + unrestricted = orchestrator._embedding_agent_id() + assert unrestricted == "other_embedder" + with _request_eligibility_scope({"chat_agent"}): + assert orchestrator._embedding_agent_id() == "same_embedder" + with _request_eligibility_scope(set()): + assert orchestrator._embedding_agent_id() is None + with _request_eligibility_scope(None): + assert orchestrator._embedding_agent_id() == unrestricted + + +def test_free_scope_embedding_endpoint_fallback_requires_free_pricing() -> None: + """A free-scoped eligible set's endpoint fallback must stay free-priced. + + Devin review (PR #1032, round 6): the endpoint-sharing fallback in + ``_embedding_agent_id`` is privacy-only -- it says nothing about cost. + ``free_only`` builds an eligible set that is entirely free (see the + ``allowed_agent_ids`` branch in ``_orchestrated_provider_completion``), + so admitting a *paid* embedding deployment purely because it shares an + endpoint with a free eligible agent lets a free request incur real, + unmetered spend outside every budget check. + """ + free_chat = ModelAgent( + "chat_free", + "chat-free-model", + base_url="https://shared.example/v1", + tags=("reasoning", "cost:free"), + ) + paid_embedder = ModelAgent( + "paid_embedder", + "paid-embed-model", + base_url="https://shared.example/v1/", + tags=("embedding",), + ) + orchestrator = TaskOrchestrator([free_chat, paid_embedder]) + with _request_eligibility_scope({free_chat.id}): + assert orchestrator._embedding_agent_id() is None + + # A genuinely free co-located embedder must still ride along -- the gate + # excludes the paid sibling, it does not over-block the free scope. + free_embedder = ModelAgent( + "free_embedder", + "free-embed-model", + base_url="https://shared.example/v1", + tags=("embedding", "cost:free"), + ) + orchestrator_with_free = TaskOrchestrator([free_chat, paid_embedder, free_embedder]) + with _request_eligibility_scope({free_chat.id}): + assert orchestrator_with_free._embedding_agent_id() == free_embedder.id + + +def test_paid_pin_embedding_endpoint_fallback_stays_privacy_only() -> None: + """An explicit paid pin's endpoint fallback is unaffected by the cost gate. + + The free-scope gate fires only when the *entire* eligible set is free. + An explicit paid pin's eligible set is a single paid agent, so its + endpoint-sharing fallback keeps admitting a co-located deployment + regardless of price -- the pre-existing privacy-only rationale, unchanged + by the fix above. + """ + paid_chat = ModelAgent( + "chat_paid", + "chat-paid-model", + base_url="https://shared.example/v1", + tags=("reasoning",), + ) + paid_embedder = ModelAgent( + "paid_embedder", + "paid-embed-model", + base_url="https://shared.example/v1/", + tags=("embedding",), + ) + orchestrator = TaskOrchestrator([paid_chat, paid_embedder]) + with _request_eligibility_scope({paid_chat.id}): + assert orchestrator._embedding_agent_id() == paid_embedder.id + + +def test_realtime_judge_spend_reaches_budget_meter_and_spend_analytics() -> None: + """The extra judge call's own tokens are metered, not silently free. + + Devin review (PR #1032): ``_realtime_route_judge``'s return value was + discarded, so a real, already-incurred provider call was invisible to + both the budget meter that gates *subsequent* requests and buyer-facing + ``spend_analytics``. + """ + agent = ModelAgent("worker_agent", "mock-model", tags=("reasoning",)) + client = _ResponsesUsageClient( + content="final answer", + usage={"input_tokens": 7, "output_tokens": 13, "total_tokens": 20}, + ) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator._model_judge_verification = lambda task, fallback, **_ignored: { # type: ignore[method-assign] + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + "judge_agent_id": "worker_agent", + "judge_model": "mock-model", + "judge_usage": {"prompt_tokens": 3, "completion_tokens": 29, "total_tokens": 32}, + } + + orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + run = next(iter(orchestrator._workflow_runs.values())) + assert run["realtime_verification"]["judge_agent_id"] == "worker_agent" + # The synthesis step reported 13 output tokens; the judge call reported + # 29 more. Both must land on the meter -- 13 alone is the bug. + assert orchestrator.budget_status()["spent_output_tokens"] == 13 + 29 + assert orchestrator._run_budget_output_by_model(run) == ({"mock-model": 13 + 29}, True) + + rows = {row["model"]: row for row in orchestrator.spend_analytics()["by_model"]} + assert rows["mock-model"]["output_tokens"] == 13 + 29 + + +def test_conduct_verification_judge_spend_counts_toward_realtime_judge_gate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """conduct's own verifier judge is current-request spend the gate must see. + + Devin review (PR #1032): ``conduct`` runs its own verifier-role judge and + records it in ``workflow["verification"]``, never as a trace row. A gate + built from trace rows plus synthesis/repair therefore missed it, so a + request whose allowance was already consumed by that first judge fired + the optional second one anyway. + + The budget here (64 judge tokens + 1) is crossed only by summing *both* + completed calls: the 13-token synthesis alone stays well inside it, which + is exactly why counting trace rows alone let the judge through. + """ + agent = ModelAgent("worker_agent", "mock-model", tags=("reasoning",)) + client = _ResponsesUsageClient( + content="final answer", + usage={"input_tokens": 7, "output_tokens": 13, "total_tokens": 20}, + ) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.conduct = lambda *args, **kwargs: { # type: ignore[method-assign] + **_STUB_CONDUCT, + "verification": { + "accepted": True, + "reason": "test", + "verifier_output": "", + "judge": "model", + "judge_agent_id": "worker_agent", + "judge_model": "mock-model", + "judge_usage": {"prompt_tokens": 3, "completion_tokens": 64, "total_tokens": 67}, + }, + } + orchestrator.budget_max_output_tokens = 65 + # Nothing is persisted yet: the whole overrun is this request's own + # in-flight spend, which budget_status() alone cannot see. + assert orchestrator.budget_status()["exceeded"] is False + + def _explode(*args: object, **kwargs: object) -> dict[str, Any]: + raise AssertionError("realtime judge must be skipped once budget is already exceeded") + + monkeypatch.setattr(orchestrator, "_model_judge_verification", _explode) + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + # The already-decided answer is still served; only the optional second + # judge is skipped. + assert result["output_text"] == "final answer" + assert orchestrator._quality_router.member_observation_count("worker_agent") == 0 + run = next(iter(orchestrator._workflow_runs.values())) + assert run["realtime_verification"] is None + + +def test_restricted_request_cannot_read_an_incompatible_scopes_cached_evidence() -> None: + """A cached routing vector never crosses the request-eligibility boundary. + + Devin review (PR #1032): the evidence caches were keyed on text and + endpoint alone and were read *before* ``_embedding_agent_id`` validated + anything, so a restricted request hitting a key an earlier unrestricted + request had filled inherited that ineligible provider's vector -- and the + routing evidence baked into it -- straight past the isolation boundary + the scope exists to draw. + + The cache still has to work for the restricted path it exists to make + cheap, so a second request under the *same* scope must still hit. + """ + chat = ModelAgent( + "chat_agent", + "chat-model", + base_url="https://chat.example/v1", + tags=("reasoning",), + ) + unrelated_embedder = ModelAgent( + "unrelated_embedder", + "embed-model", + base_url="https://unrelated.example/v1", + tags=("embedding",), + ) + client = _EmbeddingSpyClient() + orchestrator = TaskOrchestrator([chat, unrelated_embedder], client=client) + + # An earlier unrestricted request fills the cache from a provider a + # chat_agent-scoped request may not reach. + assert orchestrator._embed_cached("confidential prompt") == [0.1, 0.2, 0.3] + assert orchestrator._descriptor_vector_cached(chat) == [0.1, 0.2, 0.3] + assert client.embed_calls == ["unrelated_embedder", "unrelated_embedder"] + + with _request_eligibility_scope({"chat_agent"}): + assert orchestrator._embed_cached("confidential prompt") is None + assert orchestrator._descriptor_vector_cached(chat) is None + # No eligible embedder, so no provider call either -- the restricted + # request degrades to declaration-only evidence, it does not borrow. + assert client.embed_calls == ["unrelated_embedder", "unrelated_embedder"] + + # Same scope twice still hits the cache: the fix partitions the cache, it + # does not disable it. + with _request_eligibility_scope({"chat_agent", "unrelated_embedder"}): + assert orchestrator._embed_cached("confidential prompt") == [0.1, 0.2, 0.3] + assert orchestrator._embed_cached("confidential prompt") == [0.1, 0.2, 0.3] + assert client.embed_calls == [ + "unrelated_embedder", + "unrelated_embedder", + "unrelated_embedder", + ] + + +_JUDGED_WORKFLOW = { + "mode": "conduct", + "answer": "evidence", + "trace": [ + { + "id": 0, + "role": "worker", + "agent_id": "worker_agent", + "subtask": "do the work", + "access": [], + "output": "worker output", + "usage": {"prompt_tokens": 5, "completion_tokens": 60, "total_tokens": 65}, + } + ], + "verification": { + "accepted": True, + "reason": "test", + "verifier_output": "", + "judge": "model", + "judge_agent_id": "worker_agent", + "judge_model": "mock-model", + "judge_usage": {"prompt_tokens": 3, "completion_tokens": 45, "total_tokens": 48}, + }, +} + + +def test_budget_rejection_meters_the_spend_conduct_already_incurred() -> None: + """A rejected request still pays for the provider calls it already made. + + Devin review (PR #1032): ``_replace_workflow_run`` is the only path onto + the budget meter and this request persists nothing until it succeeds, so + the post-conduct checkpoint's raise dropped every call ``conduct`` had + already completed -- its workflow trace *and* its verifier-role judge. + The next request was then admitted against understated spend, burned the + same allowance again, and forgot it again, with no bound on the repeat. + + The 100-token cap here is crossed only by the two completed calls + together (60 worker + 45 judge), which is exactly the spend that used to + vanish. + """ + agent = ModelAgent("worker_agent", "mock-model", tags=("reasoning",)) + client = _ResponsesUsageClient( + content="never served", + usage={"input_tokens": 7, "output_tokens": 13, "total_tokens": 20}, + ) + orchestrator = TaskOrchestrator([agent], client=client) + conduct_calls: list[object] = [] + + def _conduct(*_args: object, **_kwargs: object) -> dict[str, Any]: + """Return one already-completed workflow and count the attempt.""" + conduct_calls.append(object()) + return {**_JUDGED_WORKFLOW} + + orchestrator.conduct = _conduct # type: ignore[method-assign] + orchestrator.budget_max_output_tokens = 100 + + with pytest.raises(BudgetExceededError): + orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + # The rejection never reached synthesis, but 60 + 45 output tokens had + # already left the wallet: the meter has to say so. + assert client.calls == [] + assert orchestrator.budget_status()["spent_output_tokens"] == 60 + 45 + assert orchestrator.budget_status()["exceeded"] is True + + # ...without the failed request ever surfacing as a finished workflow. + assert orchestrator.count_workflow_runs() == 0 + assert orchestrator.list_recent_runs() == [] + + # And the next request is stopped before it can burn the same allowance + # again -- against a meter that forgot, conduct ran once per request. + with pytest.raises(BudgetExceededError): + orchestrator.proxy_completion( + {"input": "hello again"}, endpoint="responses", single_agent=False + ) + assert len(conduct_calls) == 1 + assert orchestrator.budget_status()["spent_output_tokens"] == 60 + 45 + + +def test_observation_write_failure_never_discards_a_completed_answer( + tmp_path: Path, +) -> None: + """A failed routing-evidence write costs the observation, never the answer. + + Devin review (PR #1032): the observation-only judge runs after synthesis + has already succeeded but before the response and its workflow run are + persisted, and ``_observe_contextual_quality`` writes through the state + store. A store failure there therefore threw away a perfectly good + answer *and* kept the synthesis and judge spend it had already incurred + from ever reaching the ledger. + """ + agent = ModelAgent("worker_agent", "mock-model", tags=("reasoning",)) + client = _ResponsesUsageClient( + content="final answer", + usage={"input_tokens": 7, "output_tokens": 13, "total_tokens": 20}, + ) + orchestrator = TaskOrchestrator( + [agent], client=client, state_db=str(tmp_path / "state.db") + ) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator._model_judge_verification = lambda task, fallback, **_ignored: { # type: ignore[method-assign] + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + "judge_agent_id": "worker_agent", + "judge_model": "mock-model", + "judge_usage": {"prompt_tokens": 3, "completion_tokens": 29, "total_tokens": 32}, + } + assert orchestrator._store is not None + healthy_save = orchestrator._store.save + + def _failing_save( + kind: str, key: str | None, payload: dict[str, Any], **options: Any + ) -> None: + """Fail exactly the psychometric write; leave run persistence working.""" + if kind == "psychometric_observation": + raise sqlite3.OperationalError("disk I/O error") + healthy_save(kind, key, payload, **options) + + orchestrator._store.save = _failing_save # type: ignore[method-assign] + + result = orchestrator.proxy_completion( + {"input": "hello world"}, endpoint="responses", single_agent=False + ) + + assert result["output_text"] == "final answer" + # Accounting still lands: the completed judge call and the synthesis both + # reach the run record and the budget meter. + run = next(iter(orchestrator._workflow_runs.values())) + assert run["realtime_verification"]["judge_agent_id"] == "worker_agent" + assert orchestrator.budget_status()["spent_output_tokens"] == 13 + 29 + assert orchestrator._quality_router.member_report("worker_agent")["success_count"] == 1 + + +class _SchemaRepairClient: + """Answer one slow schema violation, then a fast valid repair.""" + + def __init__(self, *, first_delay: float) -> None: + self.first_delay = first_delay + self.calls = 0 + + def proxy_send( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Return the invalid first synthesis, then the valid repair.""" + del agent, endpoint, payload + self.calls += 1 + if self.calls == 1: + time.sleep(self.first_delay) + return { + "choices": [{"message": {"content": '{"input_count":6}'}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}, + } + return { + "choices": [{"message": {"content": '{"input_count":10}'}}], + "usage": {"prompt_tokens": 2, "completion_tokens": 4, "total_tokens": 6}, + } + + proxy_send_once = proxy_send + + +def test_repaired_answer_publishes_only_the_repair_calls_latency( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A repaired answer's throughput sample times the call that produced it. + + Devin review (PR #1032): ``synthesis_started`` precedes the *rejected* + first synthesis, so the published latency spanned both calls while the + usage published beside it covered only the repair -- understating the + serving model's real throughput in the very ledger this observation + exists to keep honest. + """ + agent = ModelAgent("worker_agent", "mock-model", tags=("reasoning",)) + client = _SchemaRepairClient(first_delay=0.05) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + monkeypatch.setattr( + orchestrator, + "_model_judge_verification", + lambda task, fallback, **_ignored: { + "accepted": True, + "reason": "stub verdict", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + }, + ) + + captured: dict[str, Any] = {} + original_judge = TaskOrchestrator._realtime_route_judge + + def _spy(self: TaskOrchestrator, **kwargs: Any) -> dict[str, Any]: + """Record the observation's arguments and run the real call.""" + captured.update(kwargs) + return original_judge(self, **kwargs) + + monkeypatch.setattr(TaskOrchestrator, "_realtime_route_judge", _spy) + + result = orchestrator.proxy_completion( + { + "model": "mock-model", + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "exact_count", + "strict": True, + "schema": { + "type": "object", + "properties": {"input_count": {"const": 10}}, + "required": ["input_count"], + "additionalProperties": False, + }, + }, + }, + }, + single_agent=False, + ) + + assert client.calls == 2 + assert result["choices"][0]["message"]["content"] == '{"input_count":10}' + run = orchestrator.get_workflow_run(result["orchestration"]["workflow_run_id"]) + repair_step = run["trace"][-1] + assert repair_step["role"] == "repair" + # The usage published beside the latency is the repair call's own, so the + # latency has to be the repair call's own too. + assert captured["usage"]["completion_tokens"] == 4 + assert captured["latency_seconds"] == pytest.approx(repair_step["latency_ms"] / 1000) + # The thrown-away first synthesis alone took 50ms; it is not folded in. + assert captured["latency_seconds"] < client.first_delay + + +_EXACT_TEN = { + "type": "json_schema", + "json_schema": { + "name": "exact_count", + "strict": True, + "schema": { + "type": "object", + "properties": {"input_count": {"const": 10}}, + "required": ["input_count"], + "additionalProperties": False, + }, + }, +} + + +class _SchemaFailoverClient: + """Violate the caller's schema on named members, satisfy it on the rest. + + Every response reports usage, and a failing member's synthesis and repair + report *different* token counts, so each individual attempt is + identifiable in the budget meter's per-model totals. + """ + + FAILED_SYNTHESIS_TOKENS = 11 + FAILED_REPAIR_TOKENS = 13 + SERVED_TOKENS = 17 + + def __init__(self, *failing_agent_ids: str) -> None: + self._failing_agent_ids = frozenset(failing_agent_ids) + self.calls: list[str] = [] + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Return a schema-violating answer for the failing members only.""" + del endpoint, payload + self.calls.append(agent.id) + if agent.id not in self._failing_agent_ids: + content, completion = '{"input_count": 10}', self.SERVED_TOKENS + else: + content = '{"input_count": 6}' + completion = ( + self.FAILED_SYNTHESIS_TOKENS + if self.calls.count(agent.id) == 1 + else self.FAILED_REPAIR_TOKENS + ) + return { + "choices": [{"message": {"role": "assistant", "content": content}}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": completion, + "total_tokens": 5 + completion, + }, + } + + proxy_send = proxy_send_once + + +def test_failed_schema_attempts_reach_every_spend_accounting_path() -> None: + """A discarded model's synthesis and repair are metered, not silently free. + + Devin review (PR #1032): in the virtual same-endpoint retry loop + ``synthesis_step``/``repair_step`` are rebound on every pass, so when one + model failed structured synthesis *and* its repair before a different + member of the same pool succeeded, that first model's two completed + provider calls disappeared from the final trace -- and therefore from the + budget checkpoint, the persisted run, and buyer-facing spend analytics. + Genuine provider spend went unmetered. + + Distinct from the rejected-request case an earlier round fixed: these are + attempts that completed and were *followed* by a successful one, not a + wholesale request rejection. + """ + first = ModelAgent( + "first_agent", "first-model", base_url="mock://pool", tags=("reasoning",) + ) + second = ModelAgent( + "second_agent", "second-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _SchemaFailoverClient(first.id) + orchestrator = TaskOrchestrator([first, second], client=client) + # Isolate synthesis accounting: the optional observation-only judge has + # its own already-covered metering path (test above). + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator._select_agent = lambda *args, **kwargs: first # type: ignore[method-assign] + orchestrator._failover_candidates = lambda *args, **kwargs: [first, second] # type: ignore[method-assign] + + result = orchestrator.proxy_completion( + { + "model": TaskOrchestrator.AUTO_MODEL, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + assert result["choices"][0]["message"]["content"] == '{"input_count": 10}' + assert client.calls == [first.id, first.id, second.id] + + run = next(iter(orchestrator._workflow_runs.values())) + trace = run["trace"] + assert [(step["role"], step["agent_id"]) for step in trace] == [ + ("synthesizer", first.id), + ("repair", first.id), + ("synthesizer", second.id), + ] + # Unique, contiguous step ids: `access` indexes the trace positionally + # (see get_access_report), so a rebound id would misattribute evidence. + assert [step["id"] for step in trace] == [0, 1, 2] + assert trace[1]["access"] == [0] + # The discarded attempts are marked, so an auditor reading the persisted + # run can tell them from the row that actually produced `answer`. + assert trace[0]["structured_output_error"] + assert trace[1]["structured_output_error"] + assert "structured_output_error" not in trace[2] + + # Every completed call reaches the meter and buyer-facing analytics. + failed_tokens = ( + _SchemaFailoverClient.FAILED_SYNTHESIS_TOKENS + + _SchemaFailoverClient.FAILED_REPAIR_TOKENS + ) + assert orchestrator._run_budget_output_by_model(run) == ( + { + "first-model": failed_tokens, + "second-model": _SchemaFailoverClient.SERVED_TOKENS, + }, + True, + ) + assert orchestrator.budget_status()["spent_output_tokens"] == ( + failed_tokens + _SchemaFailoverClient.SERVED_TOKENS + ) + rows = {row["model"]: row for row in orchestrator.spend_analytics()["by_model"]} + assert rows["first-model"]["output_tokens"] == failed_tokens + + # Response-facing attribution stays on the served call alone. + assert run["answer"] == '{"input_count": 10}' + assert result["usage"]["completion_tokens"] == _SchemaFailoverClient.SERVED_TOKENS + + +def test_failed_schema_attempt_spend_gates_the_next_budget_checkpoint() -> None: + """The discarded attempts count against this request's own in-flight budget. + + The same rebinding hid them from ``_trace_budget_spend``, so a request + that had already burned its allowance on a failed model kept firing + further provider calls. Both members violate the schema here, so the + second one's pre-repair checkpoint is reached with the first one's two + completed calls behind it: 11 + 13 + 11 = 35 crosses the 34-token + allowance, while the second member's own synthesis (11) alone does not -- + which is precisely why counting the current attempt alone let it through. + """ + first = ModelAgent( + "first_agent", "first-model", base_url="mock://pool", tags=("reasoning",) + ) + second = ModelAgent( + "second_agent", "second-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _SchemaFailoverClient(first.id, second.id) + orchestrator = TaskOrchestrator([first, second], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator._select_agent = lambda *args, **kwargs: first # type: ignore[method-assign] + orchestrator._failover_candidates = lambda *args, **kwargs: [first, second] # type: ignore[method-assign] + # Nothing persisted yet, so the whole overrun is in-flight spend. + orchestrator.budget_max_output_tokens = ( + _SchemaFailoverClient.FAILED_SYNTHESIS_TOKENS * 2 + + _SchemaFailoverClient.FAILED_REPAIR_TOKENS + - 1 + ) + assert orchestrator.budget_status()["exceeded"] is False + + with pytest.raises(BudgetExceededError): + orchestrator.proxy_completion( + { + "model": TaskOrchestrator.AUTO_MODEL, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + # The second member's synthesis ran (the checkpoint is pre-repair), and + # its repair was refused. Without the fix that fourth call is made. + assert client.calls == [first.id, first.id, second.id] + + +def test_pinned_terminal_schema_failure_meters_the_spend_it_already_incurred() -> None: + """An explicitly pinned model's terminal schema failure still meters spend. + + Devin review (PR #1032, round 6): when synthesis and repair both violate + the schema on an explicitly pinned model (``virtual_model`` is False, + ``synthesis_candidates == [final_agent]``), the loop raises + ``ProviderResponseError`` before ``failed_attempts.extend(...)`` ever + runs. Those two completed provider calls are real spend, but they never + reached the trace, the budget meter, or spend_analytics -- and no + workflow run is ever persisted for a request that never served an + answer, so ``count_workflow_runs`` must stay unaffected too. + """ + agent = ModelAgent( + "pinned_agent", "pinned-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _SchemaFailoverClient(agent.id) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + with pytest.raises(ProviderResponseError): + orchestrator.proxy_completion( + { + "model": agent.model, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + assert client.calls == [agent.id, agent.id] + assert orchestrator.budget_status()["spent_output_tokens"] == ( + _SchemaFailoverClient.FAILED_SYNTHESIS_TOKENS + + _SchemaFailoverClient.FAILED_REPAIR_TOKENS + ) + assert orchestrator.count_workflow_runs() == 0 + + +def test_pool_exhausted_schema_failure_meters_the_spend_it_already_incurred() -> None: + """A virtual pool's same-endpoint exhaustion also meters its spend. + + Devin review (PR #1032, round 6): distinct from the pinned-model case + above -- this exercises the ``next_agent is None`` exit, reached on a + virtual/``AUTO_MODEL`` request whose failover candidates yield no other + agent on the same endpoint. Same silent loss: the raise happened before + ``failed_attempts.extend(...)``. + """ + agent = ModelAgent( + "pool_agent", "pool-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _SchemaFailoverClient(agent.id) + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + orchestrator._select_agent = lambda *args, **kwargs: agent # type: ignore[method-assign] + orchestrator._failover_candidates = lambda *args, **kwargs: [agent] # type: ignore[method-assign] + + with pytest.raises(ProviderResponseError): + orchestrator.proxy_completion( + { + "model": TaskOrchestrator.AUTO_MODEL, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + assert client.calls == [agent.id, agent.id] + assert orchestrator.budget_status()["spent_output_tokens"] == ( + _SchemaFailoverClient.FAILED_SYNTHESIS_TOKENS + + _SchemaFailoverClient.FAILED_REPAIR_TOKENS + ) + assert orchestrator.count_workflow_runs() == 0 + + +class _TransportFailureClient: + """Every synthesis attempt raises before any provider response exists.""" + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """Simulate a transport failure: no response, nothing to meter here.""" + del agent, endpoint, payload + raise RuntimeError("connection reset by peer") + + proxy_send = proxy_send_once + + +_TRACED_CONDUCT = { + "mode": "conduct", + "answer": "evidence", + "trace": [ + { + "id": 0, + "role": "worker", + "agent_id": "worker_agent", + "subtask": "do the work", + "access": [], + "output": "worker output", + "usage": {"prompt_tokens": 2, "completion_tokens": 9, "total_tokens": 11}, + } + ], + "verification": {"accepted": True, "reason": "test", "verifier_output": ""}, +} + + +def test_synthesis_transport_failure_meters_the_spend_it_already_incurred() -> None: + """A transport failure on synthesis still meters ``conduct``'s prior spend. + + Devin review (PR #1032, round 7): distinct from round 6's two + schema-violation exits below this same ``except Exception`` clause -- a + transport failure (a raised ``ProviderUpstreamError`` or similar, as + opposed to a schema-violating response) raises before any + ``budget_checkpoint`` in this loop ever runs, so ``conduct()``'s + already-completed workflow trace vanished from the meter exactly like the + post-conduct checkpoint's own except-clause guards against. + """ + agent = ModelAgent( + "pinned_agent", "pinned-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _TransportFailureClient() + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_TRACED_CONDUCT) # type: ignore[method-assign] + + with pytest.raises(ProviderUpstreamError): + orchestrator.proxy_completion( + { + "model": agent.model, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + assert orchestrator.budget_status()["spent_output_tokens"] == 9 + assert orchestrator.count_workflow_runs() == 0 + + +class _RepairTransportFailureClient: + """Schema-violating synthesis, then a transport failure on the repair.""" + + SYNTHESIS_TOKENS = 11 + + def __init__(self) -> None: + self.calls = 0 + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """First call: completed but schema-violating. Second: raises.""" + del endpoint, payload + self.calls += 1 + if self.calls == 1: + return { + "choices": [{"message": {"role": "assistant", "content": '{"input_count": 6}'}}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": self.SYNTHESIS_TOKENS, + "total_tokens": 5 + self.SYNTHESIS_TOKENS, + }, + } + del agent + raise RuntimeError("connection reset by peer") + + proxy_send = proxy_send_once + + +def test_repair_transport_failure_meters_the_completed_synthesis_spend() -> None: + """A transport failure on repair still meters the synthesis it repairs. + + Devin review (PR #1032, round 7): ``synthesis_step`` is a real, paid-for + call that produced the schema-violating output prompting this repair. It + has not yet reached ``failed_attempts`` (only added once a repair + response exists to check), so a transport failure on the repair call + itself dropped it from the meter along with everything else in + ``workflow["trace"]``/``failed_attempts``. + """ + agent = ModelAgent( + "pinned_agent", "pinned-model", base_url="mock://pool", tags=("reasoning",) + ) + client = _RepairTransportFailureClient() + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + + with pytest.raises(ProviderUpstreamError): + orchestrator.proxy_completion( + { + "model": agent.model, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + ) + + assert client.calls == 2 + assert ( + orchestrator.budget_status()["spent_output_tokens"] + == _RepairTransportFailureClient.SYNTHESIS_TOKENS + ) + assert orchestrator.count_workflow_runs() == 0 + + +class _RepairEffortProfileFailureClient: + """Schema-violating synthesis, then an unconverted ``EffortProfileError`` on repair. + + ``apply_effort_profile`` runs before ``send_synthesis``'s own inner + ``try:`` in the structured-synthesis loop, so an ``EffortProfileError`` it + raises is never passed through ``classify_provider_failure`` -- it leaves + ``send_synthesis`` as a raw, unconverted exception. ``EffortProfileError`` + does not inherit from ``ProviderUpstreamError``, so the repair call site's + own ``except ProviderUpstreamError`` could never catch it either + (CodeRabbit review, PR #1032, round 10). + """ + + SYNTHESIS_TOKENS = 11 + + def __init__(self) -> None: + self.calls = 0 + self.effort_calls = 0 + + def apply_effort_profile( + self, + agent: ModelAgent, + payload: dict[str, Any], + profile: ReasoningEffortProfile | None, + *, + api_surface: str = "chat.completions", + ) -> dict[str, Any]: + """First call (initial synthesis): pass through. Second (repair): raise.""" + del agent, profile, api_surface + self.effort_calls += 1 + if self.effort_calls == 1: + return payload + raise EffortProfileError("provider reasoning_effort support is unproven") + + def proxy_send_once( + self, agent: ModelAgent, endpoint: str, payload: dict[str, Any] + ) -> dict[str, Any]: + """The only provider call that ever completes: schema-violating synthesis.""" + del agent, endpoint, payload + self.calls += 1 + return { + "choices": [{"message": {"role": "assistant", "content": '{"input_count": 6}'}}], + "usage": { + "prompt_tokens": 5, + "completion_tokens": self.SYNTHESIS_TOKENS, + "total_tokens": 5 + self.SYNTHESIS_TOKENS, + }, + } + + proxy_send = proxy_send_once + + +def test_repair_effort_profile_error_still_meters_the_completed_synthesis_spend() -> None: + """An unconverted ``EffortProfileError`` on repair still meters synthesis spend. + + CodeRabbit review (PR #1032, round 10): the repair call site's + ``except ProviderUpstreamError`` is narrower than the initial-call site's + own ``except Exception``, so an ``EffortProfileError`` from the repair + attempt (raised by ``apply_effort_profile``, which ``send_synthesis`` calls + before its own ``except Exception`` could ever convert it) skipped both + the accounting exclusion *and* the ``_meter_unserved_spend`` call entirely + -- the already-incurred synthesis spend silently vanished, exactly the + "spend lost" class rounds 4/6/7 already fixed for other exception types. + A misconfigured effort profile is also not the agent's fault, so it must + stay excluded from routing penalties exactly like the initial-call site + already excludes it. + """ + agent = ModelAgent( + "pinned_agent", + "pinned-model", + base_url="mock://pool", + tags=("reasoning",), + group_name="repair_pool", + ) + client = _RepairEffortProfileFailureClient() + orchestrator = TaskOrchestrator([agent], client=client) + orchestrator.policy = replace(orchestrator.policy, realtime_judge=False) + orchestrator.conduct = lambda *args, **kwargs: dict(_STUB_CONDUCT) # type: ignore[method-assign] + prior_beta = orchestrator._group_router._members[agent.id]["beta"] + + with pytest.raises(EffortProfileError): + orchestrator.proxy_completion( + { + "model": agent.model, + "messages": [{"role": "user", "content": "classify ten items"}], + "response_format": _EXACT_TEN, + }, + single_agent=False, + effort_profile=ReasoningEffortProfile(), + ) + + # (a) the EffortProfileError itself propagates -- the fix does not + # swallow it, it only ensures accounting runs first. + assert client.calls == 1 + assert client.effort_calls == 2 + # (b) the synthesis call's spend, already incurred before repair failed, + # reached the budget meter. + assert ( + orchestrator.budget_status()["spent_output_tokens"] + == _RepairEffortProfileFailureClient.SYNTHESIS_TOKENS + ) + assert orchestrator.spend_analytics()["totals"]["output_tokens"] == ( + _RepairEffortProfileFailureClient.SYNTHESIS_TOKENS + ) + assert orchestrator.count_workflow_runs() == 0 + # (c) neither the circuit breaker nor the group router's stability + # posterior recorded a failure for this agent. + assert agent.id not in orchestrator._circuit + assert orchestrator._group_router._members[agent.id]["beta"] == prior_beta + + +class _EmbeddingSpaceSpyClient: + """One distinct unit vector per embedding deployment, every call recorded.""" + + VECTORS = { + "first_embedder": [1.0, 0.0, 0.0], + "second_embedder": [0.0, 1.0, 0.0], + } + + def __init__(self) -> None: + self.embed_calls: list[str] = [] + + def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]: + """Return this deployment's own vector space, recording the caller.""" + self.embed_calls.append(agent.id) + return [list(self.VECTORS[agent.id]) for _ in texts] + + def embed_with_usage( + self, agent: ModelAgent, texts: list[str] + ) -> tuple[list[list[float]], int | None]: + """Same vector space; no authoritative usage (mirrors the mock transport).""" + return self.embed(agent, texts), None + + +def _embedding_space_pool() -> tuple[ModelAgent, ModelAgent, ModelAgent]: + """One chat model plus two embedding deployments in one eligible set.""" + return ( + ModelAgent( + "chat_agent", "chat-model", base_url="mock://pool", tags=("reasoning",) + ), + ModelAgent( + "first_embedder", + "first-embed-model", + base_url="mock://pool", + tags=("embedding",), + priority=10, + ), + ModelAgent( + "second_embedder", + "second-embed-model", + base_url="mock://pool", + tags=("embedding",), + priority=1, + ), + ) + + +def test_evidence_cache_is_not_reused_across_a_change_of_embedding_model() -> None: + """A cached vector never survives the embedder that produced it changing. + + Devin review (PR #1032): the eligibility partition + (``_request_evidence_partition``) is keyed on endpoint identity, ZDR mode + and the sorted allow-list -- the restrictions that *select* an embedder. + ``_embedding_agent_id`` can still resolve to a different embedding + model/deployment under an identical restriction shape (an agent-pool + change, or measured-member reordering inside one eligible group), and the + two evidence caches persist across that. A hit then returned a vector from + a different embedding space, and a cosine between two embedding spaces is + meaningless -- affinity-based routing was silently corrupted rather than + degraded. + + This is a refinement of the eligibility partitioning, not a duplicate: + the restriction shape below is *identical* before and after, which is + exactly why that partition alone cannot catch it. + """ + chat, first_embedder, second_embedder = _embedding_space_pool() + client = _EmbeddingSpaceSpyClient() + orchestrator = TaskOrchestrator([chat, first_embedder, second_embedder], client=client) + + assert orchestrator._embedding_agent_id() == first_embedder.id + partition_before = _request_evidence_partition() + assert orchestrator._embed_cached("routing evidence text") == [1.0, 0.0, 0.0] + assert orchestrator._descriptor_vector_cached(chat) == [1.0, 0.0, 0.0] + + # An operator re-ranks the pool. Same agents, same endpoint, same ZDR + # mode, no eligibility scope -- so the restriction shape does not move. + orchestrator.patch_agent("default", second_embedder.id, {"priority": 100}) + assert orchestrator._embedding_agent_id() == second_embedder.id + assert _request_evidence_partition() == partition_before + + assert orchestrator._embed_cached("routing evidence text") == [0.0, 1.0, 0.0] + assert orchestrator._descriptor_vector_cached(chat) == [0.0, 1.0, 0.0] + assert client.embed_calls == [ + first_embedder.id, + first_embedder.id, + second_embedder.id, + second_embedder.id, + ] + + # The cache still works: a repeat under the current embedder hits. + assert orchestrator._embed_cached("routing evidence text") == [0.0, 1.0, 0.0] + assert client.embed_calls[-1] == second_embedder.id + assert len(client.embed_calls) == 4 + + +def test_one_cosine_never_mixes_two_embedding_spaces() -> None: + """Both halves of an affinity come from one embedding member, resolved once. + + The task vector and every descriptor vector in a single + ``_semantic_affinities`` call are pinned to the member resolved at its + start, so a pool change landing mid-comparison cannot make one cosine + span two embedding spaces. + """ + chat, first_embedder, second_embedder = _embedding_space_pool() + client = _EmbeddingSpaceSpyClient() + orchestrator = TaskOrchestrator([chat, first_embedder, second_embedder], client=client) + + affinities = orchestrator._semantic_affinities("classify ten items", [chat]) + assert affinities[chat.id] == pytest.approx(1.0) + assert set(client.embed_calls) == {first_embedder.id} + + orchestrator.patch_agent("default", second_embedder.id, {"priority": 100}) + client.embed_calls.clear() + affinities = orchestrator._semantic_affinities("classify ten items", [chat]) + assert affinities[chat.id] == pytest.approx(1.0) + # Not one leftover call to the old embedder: a mixed pair would have + # yielded a cosine of 0.0 between two orthogonal spaces. + assert set(client.embed_calls) == {second_embedder.id} + + +class _PricedEmbeddingClient: + """An embedder that reports authoritative ``prompt_tokens`` on every call. + + Unlike ``_EmbeddingSpaceSpyClient``/``_EmbeddingSpyClient`` (which mirror + the ``mock://`` transport's ``prompt_tokens=None``), this fixture stands + in for a real, usage-reporting provider -- the case + ``_meter_embedding_spend`` exists to meter. + """ + + def __init__(self, prompt_tokens: int) -> None: + self.prompt_tokens = prompt_tokens + self.embed_calls: list[str] = [] + + def embed_with_usage( + self, agent: ModelAgent, texts: list[str] + ) -> tuple[list[list[float]], int | None]: + self.embed_calls.append(agent.id) + return [[0.5, 0.5, 0.5] for _ in texts], self.prompt_tokens + + def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]: + vectors, _prompt_tokens = self.embed_with_usage(agent, texts) + return vectors + + +def _paid_embedder() -> ModelAgent: + return ModelAgent( + "paid_embedder", + "paid-embed-model", + base_url="https://embed.example/v1", + tags=("embedding",), + ) + + +def test_cache_miss_with_paid_embedder_is_metered() -> None: + """A cache-miss embedding call with authoritative usage reaches the budget meter. + + Before this fix, ``_embed_cached`` called ``client.embed()`` (discarding + usage) and never touched ``_workflow_runs`` at all: real, incurred + provider spend on routing-evidence embeddings was completely invisible + to ``spend_analytics``/``budget_status`` (Devin review on #1032). + """ + embedder = _paid_embedder() + client = _PricedEmbeddingClient(prompt_tokens=37) + orchestrator = TaskOrchestrator([embedder], client=client) + + vector = orchestrator._embed_cached("routing text") + + assert vector == [0.5, 0.5, 0.5] + analytics = orchestrator.spend_analytics() + assert analytics["totals"]["prompt_tokens"] == 37 + by_model = {row["model"]: row for row in analytics["by_model"]} + assert by_model["paid-embed-model"]["output_tokens"] == 0 + assert by_model["paid-embed-model"]["step_count"] == 1 + + +def test_cache_hit_is_never_metered_again() -> None: + """A cache hit reuses the already-metered vector; it must not meter twice. + + This is the fix's core safety property: metering sits strictly on the + miss branch, after the ``if cached is not None: return cached`` + short-circuit, so a hit is never reachable from it. + """ + embedder = _paid_embedder() + client = _PricedEmbeddingClient(prompt_tokens=37) + orchestrator = TaskOrchestrator([embedder], client=client) + + assert orchestrator._embed_cached("routing text") == [0.5, 0.5, 0.5] + assert orchestrator._embed_cached("routing text") == [0.5, 0.5, 0.5] + + assert client.embed_calls == [embedder.id] + assert orchestrator.spend_analytics()["totals"]["prompt_tokens"] == 37 + + assert orchestrator._descriptor_vector_cached(embedder) == [0.5, 0.5, 0.5] + assert orchestrator._descriptor_vector_cached(embedder) == [0.5, 0.5, 0.5] + + assert client.embed_calls == [embedder.id, embedder.id] + assert orchestrator.spend_analytics()["totals"]["prompt_tokens"] == 74 + + +def test_select_model_group_members_meters_embedding_with_no_live_request() -> None: + """A request-independent caller still gets its embedding spend metered. + + ``select_model_group_members`` is the confirmed real call site that + reaches ``_embed_cached``/``_descriptor_vector_cached`` with no + enclosing task, workflow record, or request-eligibility context at all + (no ``_request_eligibility_scope`` is entered anywhere in this test) -- + ``_meter_embedding_spend`` must not assume any of those exist. + """ + chat = ModelAgent( + "chat_agent", "chat-model", base_url="mock://pool", tags=("reasoning",) + ) + embedder = _paid_embedder() + client = _PricedEmbeddingClient(prompt_tokens=11) + orchestrator = TaskOrchestrator([chat, embedder], client=client) + + selected = orchestrator.select_model_group_members( + [chat], text="classify ten items" + ) + + assert [agent.id for agent in selected] == [chat.id] + assert client.embed_calls # the task text and/or a descriptor were embedded + expected_prompt_tokens = len(client.embed_calls) * 11 + assert ( + orchestrator.spend_analytics()["totals"]["prompt_tokens"] + == expected_prompt_tokens + ) + + +def test_unpriced_embedder_under_cost_budget_does_not_blind_meter() -> None: + """An unpriced embedder must never flip the whole meter to blocked_unavailable. + + Isolates the exact gap: every *general-chat* model is priced (so + ``budget_status()``'s own ``candidate_prices_available`` check, which + only looks at ``_is_general_chat_agent`` models, already passes) but the + embedder -- structurally unable to need a price before this fix, since + embedding calls never produced a workflow run at all -- is not. + ``completion_tokens`` is always ``0`` for an embedding step, which + satisfies ``_replace_workflow_run``'s *per-step* availability check -- + but its second, independent gate (every model appearing in the run must + be priced whenever ``budget_max_cost_usd`` is set) still fails on the + unpriced embedder. Before this guard, that added the synthetic run's id + to ``_budget_unavailable_run_ids`` and flipped + ``budget_status()["enforcement_status"]`` to ``"blocked_unavailable"`` + for every request org-wide, even though this run's real contribution to + spend is mathematically $0 (an embedding call has no completion + tokens). The guard skips metering in that one case -- true zero-cost + either way -- instead of writing a row that blinds unrelated requests' + enforcement. + """ + chat = ModelAgent( + "chat_agent", "chat-model", base_url="mock://pool", tags=("reasoning",) + ) + embedder = _paid_embedder() + + unpriced_client = _PricedEmbeddingClient(prompt_tokens=37) + unpriced = TaskOrchestrator( + [chat, embedder], + client=unpriced_client, + budget_max_cost_usd=1.0, + price_per_million={"chat-model": 1.0}, + ) + unpriced._embed_cached("routing text") + assert unpriced.budget_status()["enforcement_status"] == "within_budget" + assert unpriced.spend_analytics()["totals"]["prompt_tokens"] == 0 + + # Scoped to the unpriced case only: once the embedder is priced too + # (even at $0), metering proceeds normally and the meter still reads + # within_budget -- this is not a blanket "budget enabled => never meter". + priced_client = _PricedEmbeddingClient(prompt_tokens=37) + priced = TaskOrchestrator( + [chat, embedder], + client=priced_client, + budget_max_cost_usd=1.0, + price_per_million={"chat-model": 1.0, "paid-embed-model": 0.0}, + ) + priced._embed_cached("routing text") + assert priced.spend_analytics()["totals"]["prompt_tokens"] == 37 + assert priced.budget_status()["enforcement_status"] == "within_budget" + + +def test_embedding_spend_persistence_failure_never_aborts_model_selection( + tmp_path: Path, caplog: pytest.LogCaptureFixture +) -> None: + """A durable-write failure behind an embedding spend never propagates. + + Devin review (PR #1032, comment 3921362518): ``_meter_embedding_spend`` + ends with an unguarded ``self._store.save("workflow_run", ...)`` -- a + real sqlite3 write. Any failure there (a state-store outage) propagated + straight through ``_embed_cached``/``_descriptor_vector_cached``, + through ``_semantic_affinities``, through every one of + ``_ranked_agents``'s call sites, and into model selection: a disk + hiccup on a best-effort routing-evidence row could abort real request + handling. + """ + embedder = _paid_embedder() + client = _PricedEmbeddingClient(prompt_tokens=37) + orchestrator = TaskOrchestrator( + [embedder], client=client, state_db=str(tmp_path / "state.db") + ) + assert orchestrator._store is not None + healthy_save = orchestrator._store.save + + def _failing_save( + kind: str, key: str | None, payload: dict[str, Any], **options: Any + ) -> None: + """Fail exactly the workflow_run write; leave everything else working.""" + if kind == "workflow_run": + raise sqlite3.OperationalError("disk I/O error") + healthy_save(kind, key, payload, **options) + + orchestrator._store.save = _failing_save # type: ignore[method-assign] + + with caplog.at_level(logging.WARNING): + vector = orchestrator._embed_cached("routing text") + + # (a) the caller never sees the exception -- it gets its vector back. + assert vector == [0.5, 0.5, 0.5] + + # (b) the failure doesn't poison the cache: a second call is a cache + # hit, not a second provider call. + vector_again = orchestrator._embed_cached("routing text") + assert vector_again == [0.5, 0.5, 0.5] + assert client.embed_calls == [embedder.id] + + # (c) the failure isn't silently discarded -- it's logged. + assert any( + record.levelno == logging.WARNING and "embedding spend" in record.getMessage() + for record in caplog.records + ) + + # Caveat 3: the durability failure never blinds the in-memory meter -- + # _replace_workflow_run already ran before the guarded write, so the + # spend is real and enforcement stays exact, not "unavailable". + assert orchestrator.spend_analytics()["totals"]["prompt_tokens"] == 37 + status = orchestrator.budget_status() + assert status["enforcement_status"] == "within_budget" + assert status["measurement_status"] == "measured" + + +if __name__ == "__main__": # pragma: no cover + sys.exit(pytest.main([__file__])) diff --git a/tests/test_passthrough_provider_failover.py b/tests/test_passthrough_provider_failover.py index 87751f2ac..478df7c12 100644 --- a/tests/test_passthrough_provider_failover.py +++ b/tests/test_passthrough_provider_failover.py @@ -26,6 +26,40 @@ from contextual_orchestrator.provider_errors import ProviderUpstreamError +@pytest.fixture(autouse=True) +def _stub_realtime_judge(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep this file's exact-sequence assertions about provider failover only. + + ``_orchestrated_provider_completion`` now calls the realtime fast-mlsirm + judge once synthesis succeeds (observation-only: it never changes the + response). ``SequencedProxyClient`` is a minimal transport double with no + ``outcomes`` entry for whichever agent gets selected as verifier, so a + real judge attempt would append an unplanned call to ``client.calls`` and + break the ``[agent_id for agent_id, _ in client.calls] == [...]`` + assertions this file is actually about. Patch at the class level (not + just ``_build``) so every inline ``TaskOrchestrator(...)`` construction in + this file is covered. + """ + + def _accept( + self: TaskOrchestrator, + task: str, + fallback: dict[str, Any], + *, + free_only: bool = False, + **_ignored: Any, + ) -> dict[str, Any]: + del self, task, free_only, _ignored + return { + "accepted": True, + "reason": "stubbed for passthrough failover coverage", + "verifier_output": fallback.get("verifier_output", ""), + "judge": "model", + } + + monkeypatch.setattr(TaskOrchestrator, "_model_judge_verification", _accept) + + class SequencedProxyClient: """Return one configured outcome per provider while recording attempts."""