diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index a1a379f4a..1120b9877 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -77,12 +77,14 @@ def __init__( ) else: self.batch_backend = batch_backend + self._embedding_backend_is_default = embedding_batch_backend is None self.embedding_batch_backend: EmbeddingBatchBackend = embedding_batch_backend or self._default_embedding_backend() # job_id -> submitted BatchJob (so poll/retrieve can be driven by id) self._batch_jobs: Dict[str, BatchJob] = {} # embeddings batch state: job handle + submitted requests + cached doc, # keyed by batch id so poll/retrieve is idempotent (usage recorded once). self._embedding_jobs: Dict[str, BatchJob] = {} + self._embedding_job_backends: Dict[str, EmbeddingBatchBackend] = {} self._embedding_requests: Dict[str, List[EmbeddingBatchRequest]] = {} self._embedding_input_counts: Dict[str, int] = {} self._embedding_part_counts: Dict[str, List[int]] = {} @@ -487,8 +489,14 @@ def submit_embeddings_batch( provider_name=provider_name, attribution=shared_attribution, ) - job = self.embedding_batch_backend.submit(requests, metadata=metadata) + backend = ( + self._default_embedding_backend() + if self._embedding_backend_is_default + else self.embedding_batch_backend + ) + job = backend.submit(requests, metadata=metadata) self._embedding_jobs[job.job_id] = job + self._embedding_job_backends[job.job_id] = backend self._embedding_requests[job.job_id] = requests self._embedding_input_counts[job.job_id] = len(inputs) self._embedding_part_counts[job.job_id] = part_counts @@ -702,7 +710,8 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: return cached job = self._require_embedding_job(batch_id) - status = self.embedding_batch_backend.poll(job) + backend = self._embedding_job_backends.get(batch_id, self.embedding_batch_backend) + status = backend.poll(job) if not status.get("is_complete"): return { "batch_id": batch_id, @@ -711,7 +720,7 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: "embeddings": None, } - items: List[EmbeddingBatchResultItem] = self.embedding_batch_backend.retrieve(job) + items: List[EmbeddingBatchResultItem] = backend.retrieve(job) requests = self._embedding_requests.get(batch_id, []) request_by_custom_id = {request.custom_id: request for request in requests} input_count = self._embedding_input_counts.get(batch_id, len(requests)) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 38a2a48a7..85dbb07e9 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -68,6 +68,8 @@ def _temperature_capability_rejection(exc: Exception) -> bool: return False try: body = exc.read() + except http.client.IncompleteRead as read_error: + body = read_error.partial except (OSError, ValueError): body = b"" body_bytes = body if isinstance(body, bytes) else str(body).encode("utf-8") @@ -220,6 +222,7 @@ def complete_structured( response = self.orchestrator.client.proxy_send(agent, "chat/completions", payload) output = ModelClient._response_content(agent, response) usage = response.get("usage") if isinstance(response.get("usage"), dict) else None + self.orchestrator._record_in_flight_provider_usage(agent, usage, output) return self._completion_payload(output, agent.id, usage, self.mode if mode is None else mode) def _completion_payload( @@ -313,6 +316,8 @@ def __post_init__(self) -> None: raise TypeError("local_credential_key must be a string") if scheme == "local" and not self.local_credential_key: raise ValueError("local:// gateway URLs require local_credential_key") + if scheme == "local" and not _is_local_provider_url(self.base_url): + raise ValueError("local:// gateway URLs require a well-formed explicit loopback endpoint") if self.local_credential_key and scheme != "local": raise ValueError("local_credential_key requires a local:// gateway URL") if not self.auth_scheme or type(self.auth_scheme) is not str: @@ -456,7 +461,7 @@ def as_dict(self) -> dict[str, Any]: # is a caller or configuration error and must not be retried. TRANSIENT_HTTP_STATUS = frozenset({408, 409, 425, 429, 500, 502, 503, 504}) LOCAL_PROVIDER_SCHEMES = frozenset({"local"}) -LOCAL_PROVIDER_HOSTS = frozenset({"127.0.0.1", "::1", "localhost", "host.docker.internal"}) +LOCAL_PROVIDER_HOSTS = frozenset({"127.0.0.1", "::1", "localhost"}) def _is_local_provider_url(base_url: str) -> bool: @@ -2018,14 +2023,15 @@ def close(self) -> None: class _StateStore: """Minimal write-through sqlite persistence for orchestrator runtime state. - ponytail: one generic table, no ORM. Keyed kinds (workflow_run, evaluation_run) + ponytail: one generic table, no ORM. Keyed kinds (workflow_run, evaluation_run, + budget_meter) upsert by key; stream kinds (analytics, audit) append. Stream rows grow unbounded on disk while the in-memory deques stay capped — add pruning if db size matters. Runtime values (kind, key, payload, limit) are always bound through SQLite placeholders so persisted prompts and identifiers cannot become SQL syntax. """ - _KEYED = {"workflow_run", "evaluation_run"} + _KEYED = {"workflow_run", "evaluation_run", "budget_meter"} _CREATE_RECORDS_SQL = ( "CREATE TABLE IF NOT EXISTS records (" "seq INTEGER PRIMARY KEY AUTOINCREMENT, kind TEXT NOT NULL, key TEXT, payload TEXT NOT NULL)" @@ -2169,6 +2175,9 @@ def __init__( self._audit_events: deque[dict[str, Any]] = deque(maxlen=256) self._run_order: deque[str] = deque(maxlen=128) self._workflow_run_lock = threading.Lock() + self._budget_spend_lock = threading.Lock() + self._budget_spent_output_tokens = 0 + self._budget_spent_cost_usd = 0.0 self._archived_spend: dict[str, Any] = { "run_count": 0, "estimated_prompt_tokens": 0, @@ -2192,6 +2201,15 @@ def __init__( self._commercial_report_cache_local = threading.local() if self._store is not None: self._reload_state() + restored = self.spend_analytics()["totals"] + self._budget_spent_output_tokens = max( + self._budget_spent_output_tokens, + restored["estimated_output_tokens"], + ) + self._budget_spent_cost_usd = max( + self._budget_spent_cost_usd, + restored["estimated_cost_usd"] or 0.0, + ) def close(self) -> None: """Release optional durable resources owned by this orchestrator.""" @@ -2248,6 +2266,14 @@ def provider_readiness_report( } def _reload_state(self) -> None: + budget_meter = self._store.load("budget_meter") + if budget_meter: + self._budget_spent_output_tokens = int( + budget_meter[-1].get("spent_output_tokens", 0) + ) + self._budget_spent_cost_usd = float( + budget_meter[-1].get("spent_cost_usd", 0.0) + ) for record in self._store.load("workflow_run"): self._remember_workflow_run(record) for evaluation in self._store.load("evaluation_run"): @@ -2339,6 +2365,11 @@ def proxy_completion( usage, responses=endpoint.strip("/") == "responses", ) + self._record_in_flight_provider_usage( + agent, + passthrough_step.get("usage"), + passthrough_output, + ) self._persist_workflow_run( { "workflow_run_id": f"run_{uuid.uuid4().hex}", @@ -2386,29 +2417,47 @@ def _orchestrated_provider_completion( disabled_model = requested_model if requested_model is not None else final_agent.model raise RuntimeError(f"requested model {disabled_model!r} is disabled") + with self._budget_spend_lock: + budget_before = ( + self._budget_spent_output_tokens, + self._budget_spent_cost_usd, + ) workflow = self.conduct(messages, preserve_messages=True, judge=True) - in_flight_output_tokens = 0 - in_flight_cost_usd = 0.0 + workflow_output_tokens = 0 + workflow_cost_usd = 0.0 model_by_agent = {agent.id: agent.model for agent in self.agents} - for row in workflow["trace"]: + metered_steps = list(workflow["trace"]) + verification = workflow.get("verification") + if isinstance(verification, dict) and isinstance( + verification.get("judge_usage"), dict + ): + metered_steps.append( + { + "agent_id": verification.get("judge_agent_id", "unknown"), + "usage": verification["judge_usage"], + "output": "", + } + ) + for row in metered_steps: output_tokens, _reported = _step_output_token_count(row) - in_flight_output_tokens += output_tokens - model = model_by_agent.get(row.get("served_agent_id") or row.get("agent_id")) + workflow_output_tokens += output_tokens + model = model_by_agent.get( + row.get("served_agent_id") or row.get("agent_id") + ) price = self.price_per_million.get(model) if model is not None else None if price is not None: - in_flight_cost_usd += output_tokens / 1_000_000 * price - verification = workflow.get("verification") - verification_usage = verification.get("judge_usage") if isinstance(verification, dict) else None - if isinstance(verification_usage, dict): - judge_output_tokens, _reported = _step_output_token_count( - {"usage": verification_usage, "output": ""} + workflow_cost_usd += output_tokens / 1_000_000 * price + with self._budget_spend_lock: + recorded_output_tokens = self._budget_spent_output_tokens - budget_before[0] + recorded_cost_usd = self._budget_spent_cost_usd - budget_before[1] + self._budget_spent_output_tokens += max( + 0, + workflow_output_tokens - recorded_output_tokens, + ) + self._budget_spent_cost_usd += max( + 0.0, + workflow_cost_usd - recorded_cost_usd, ) - in_flight_output_tokens += judge_output_tokens - judge_agent_id = verification.get("judge_agent_id") - judge_model = model_by_agent.get(judge_agent_id) - judge_price = self.price_per_million.get(judge_model) if judge_model is not None else None - if judge_price is not None: - in_flight_cost_usd += judge_output_tokens / 1_000_000 * judge_price evidence = "\n\n".join( f"Workflow step {row['id']} ({row['role']}):\n{row['output']}" for row in workflow["trace"] @@ -2438,10 +2487,7 @@ def _orchestrated_provider_completion( ) upstream["model"] = final_agent.model upstream["stream"] = False - self._raise_if_spend_budget_exceeded( - additional_output_tokens=in_flight_output_tokens, - additional_cost_usd=round(in_flight_cost_usd, 6), - ) + self._raise_if_spend_budget_exceeded() synthesis_started = time.perf_counter() raw = self.client.proxy_send(final_agent, "responses", upstream) synthesis_output = raw.get("output_text") @@ -2487,10 +2533,7 @@ def _orchestrated_provider_completion( upstream["model"] = final_agent.model upstream["messages"] = synthesis_messages upstream["stream"] = False - self._raise_if_spend_budget_exceeded( - additional_output_tokens=in_flight_output_tokens, - additional_cost_usd=round(in_flight_cost_usd, 6), - ) + self._raise_if_spend_budget_exceeded() synthesis_started = time.perf_counter() raw = self.client.proxy_send(final_agent, "chat/completions", upstream) synthesis_output = "" @@ -2515,6 +2558,11 @@ def _orchestrated_provider_completion( usage, responses=response_request, ) + self._record_in_flight_provider_usage( + final_agent, + synthesis_step.get("usage"), + synthesis_output, + ) trace = [*workflow["trace"], synthesis_step] workflow_run_id = f"run_{uuid.uuid4().hex}" record = { @@ -2599,6 +2647,7 @@ def stream_route(self, messages: list[ChatMessage], workflow_run_id: str | None parts.append(delta) yield delta answer = "".join(parts) + self._record_in_flight_provider_usage(agent, None, answer) record = { "workflow_run_id": workflow_run_id or f"run_{uuid.uuid4().hex}", "created_at": int(time.time()), @@ -2749,6 +2798,11 @@ def batch_route(self, prompts: list[str]) -> list[dict[str, Any]]: } if result.get("usage") is not None: row["usage"] = result["usage"] + self._record_in_flight_provider_usage( + agent, + row.get("usage"), + result["content"], + ) record = { "workflow_run_id": f"run_{uuid.uuid4().hex}", "created_at": int(time.time()), @@ -3235,10 +3289,15 @@ def _plan_generated(self, task: str) -> list[WorkflowStep]: "verifier step when correctness matters.\n" f"Available agents:\n{pool}" ) - raw = self.client.chat(planner, [ - {"role": "system", "content": system}, - {"role": "user", "content": task}, - ]) + raw, _served_id, _usage = self._invoke( + planner, + [ + {"role": "system", "content": system}, + {"role": "user", "content": task}, + ], + text=task, + role="thinker", + ) return self._parse_workflow_plan(raw) def _parse_workflow_plan(self, raw: str) -> list[WorkflowStep]: @@ -3363,13 +3422,17 @@ def _invoke( last_error: Exception | None = None for agent in candidates: try: + self._raise_if_spend_budget_exceeded() output = self.client.chat(agent, messages) + except BudgetExceededError: + raise except Exception as exc: # noqa: BLE001 - one agent failing routes to the next last_error = exc self._record_failure(agent.id) continue self._record_success(agent.id) usage = self.client.take_usage() if hasattr(self.client, "take_usage") else None + self._record_in_flight_provider_usage(agent, usage, output) return output, agent.id, usage if isinstance(last_error, ProviderResponseError): raise last_error @@ -3772,6 +3835,10 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> "estimated_cost_usd": cost, }) + with self._budget_spend_lock: + budget_output_tokens = self._budget_spent_output_tokens + budget_cost_usd = self._budget_spent_cost_usd + return { "measurement_status": "local_runtime_estimate", "source_note": ( @@ -3792,7 +3859,10 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> }, "by_model": rows, "unpriced_models": unpriced, - "budget": self._budget_block(total_output_tokens, round(total_cost, 6) if prices else None), + "budget": self._budget_block( + budget_output_tokens, + round(budget_cost_usd, 6) if prices else None, + ), } def _budget_block(self, spent_tokens: int, spent_cost: float | None) -> dict[str, Any]: @@ -3820,6 +3890,31 @@ def budget_status(self) -> dict[str, Any]: """Current spend-budget state (limits, spent, remaining, exceeded).""" return self.spend_analytics()["budget"] + def _record_in_flight_provider_usage( + self, + agent: ModelAgent, + usage: dict[str, Any] | None, + output: str, + ) -> None: + """Add one completed provider call to the process budget ledger.""" + output_tokens, _reported = _step_output_token_count( + {"usage": usage, "output": output} + ) + price = self.price_per_million.get(agent.model) + with self._budget_spend_lock: + self._budget_spent_output_tokens += output_tokens + if price is not None: + self._budget_spent_cost_usd += output_tokens / 1_000_000 * price + if self._store is not None: + self._store.save( + "budget_meter", + "process_budget", + { + "spent_output_tokens": self._budget_spent_output_tokens, + "spent_cost_usd": self._budget_spent_cost_usd, + }, + ) + def _raise_if_spend_budget_exceeded( self, *, diff --git a/docs/planning/adrs/0004-pr-review-merge-loop.md b/docs/planning/adrs/0004-pr-review-merge-loop.md index c63152d94..2ec1a5be3 100644 --- a/docs/planning/adrs/0004-pr-review-merge-loop.md +++ b/docs/planning/adrs/0004-pr-review-merge-loop.md @@ -340,6 +340,8 @@ For each repository, record branch, commit, PR URL, review result, check result, | Contextual PR #109 exact head `60d9cfc9be2ce0426ed37746eb9a2768b8f3455d` produced Strix run `31831835133`/job `94869100782` with a successful zero-finding report and artifact `9231362799`, but no `evidence-binding.json`; `run.json` contained only a local temporary target and no repository, PR head, job, report path, or digest binding. The report digest was `33a47fa5855600b393d92c3a1a77e0ac15adb99a55406f99e5b2efba93c17c18`, and the same-head publish step was skipped. | Keep the success as provider/content evidence only, not a clean exact-head security or Merge result. Require the central trusted workflow to publish a structured binding for repository, full PR head, run/job, report path, and digest; reject unbound success and re-run the exact head after protected-main workflow integration. | Reproduced 2026-08-15; central provenance dependency and protected Merge remain open | | The contextual exact-head code-scanning checks `Trivy` and `Scorecard` both became `neutral` because the PR's `security.yml` configuration was absent from protected `main` (`trivy-filesystem` and `supply-chain/branch-protection`), while the separate required `trivy-fs`/`scorecard` jobs were still pending or successful. | Keep the organization's CodeQL-only code-scanning rule unchanged; distinguish code-scanning alert comparison from required Security job results, record the missing configuration as a warning, and never treat neutral, skipped, or absent results as a pass. Re-audit after trusted workflow/ruleset integration. | Observed 2026-08-15; no local source bypass, governance/central workflow follow-up required | | Fast PR #816 exact head `355f93b27ba4a0cb141e86b0fbc9127681edb750` produced Strix run `31833030282`/job `94872991843`; NVIDIA NIM failed with `agents.exceptions.ModelBehaviorError: Tool execute not found in agent strix`, no penetration-test report was produced, and the required check failed closed after 599 seconds. | Keep this as provider/model-tool-contract evidence, not a source vulnerability or clean security result. Preserve the failure denominator, do not publish a status-only success or retry to hide it, and require central trusted workflow/provider repair plus a structured same-head artifact before normal Merge. | Observed 2026-08-15; central Strix dependency and protected Merge remain open | +| PR #765 review found `host.docker.internal` classified as a `local://` provider even though the egress validator correctly rejects its usual non-loopback Docker gateway address. | Keep authenticated local transport loopback-only under ADR 0012, remove the unreachable host classification, and reject non-loopback `local://` configuration at `ModelAgent` construction while retaining the DNS-resolution check against rebinding. | Decision recorded 2026-08-21; implementation and exact-head review/check evidence follow | +| PR #765 review found that auxiliary temperature-capability inspection could let `http.client.IncompleteRead` replace the provider's original HTTP error while reading its diagnostic body. | Treat an incomplete diagnostic body as insufficient capability evidence, preserve the original HTTP error for the caller, and cover the exact truncated-body branch without broad retry or fallback changes. | Decision recorded 2026-08-21; implementation and exact-head review/check evidence follow | ## Risks and Mitigations diff --git a/docs/planning/adrs/0014-gateway-owned-model-selection.md b/docs/planning/adrs/0014-gateway-owned-model-selection.md index a606a5a3c..d399c8b1d 100644 --- a/docs/planning/adrs/0014-gateway-owned-model-selection.md +++ b/docs/planning/adrs/0014-gateway-owned-model-selection.md @@ -81,6 +81,11 @@ because the provider response shape is richer. enforcement and spend totals remain cumulative without retaining every prompt, answer, or trace in process memory. A configured durable state store remains the long-term run-evidence boundary. +- Every completed provider call adds reported output usage (or the bounded + estimate when unavailable) to a synchronized process budget ledger before the + next planner, worker, verifier, judge, or synthesizer call. Failed workflows + therefore retain consumed spend even when no completed run can be persisted; + the configured state store checkpoints this compact meter across restarts. ## Consequences diff --git a/docs/planning/adrs/0015-auto-embedding-model-selection.md b/docs/planning/adrs/0015-auto-embedding-model-selection.md index 4f29b1f3a..67280b866 100644 --- a/docs/planning/adrs/0015-auto-embedding-model-selection.md +++ b/docs/planning/adrs/0015-auto-embedding-model-selection.md @@ -61,6 +61,10 @@ guess, or consumer-side fallback. cannot override either server-resolved identity. The standalone in-process backend remains a local test/development path; a configured provider path uses its injected embeddings backend and the resolved model. +6. The default batch backend resolves the current agent pool at submission time. + Runtime additions, disablement, and priority changes are therefore visible to + new jobs; every submitted job retains the backend instance that created it for + deterministic polling and retrieval. ## Contract and acceptance evidence diff --git a/tests/test_budget_enforcement.py b/tests/test_budget_enforcement.py index 3f3003a8c..1f4bdb36b 100644 --- a/tests/test_budget_enforcement.py +++ b/tests/test_budget_enforcement.py @@ -10,10 +10,13 @@ import json from pathlib import Path import sys +import tempfile import threading import urllib.error import urllib.request +import pytest + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402 @@ -75,6 +78,54 @@ def test_cost_budget_blocks() -> None: assert raised +def test_unpersisted_provider_usage_remains_in_the_budget_ledger() -> None: + agent = ModelAgent("general_agent", "priced-model", tags=("reasoning",)) + orchestrator = TaskOrchestrator( + [agent], + price_per_million={"priced-model": 1_000_000.0}, + budget_max_cost_usd=1.0, + ) + + orchestrator._record_in_flight_provider_usage( + agent, + {"completion_tokens": 1}, + "", + ) + + with pytest.raises(BudgetExceededError, match="spend budget exceeded"): + orchestrator._raise_if_spend_budget_exceeded() + + assert orchestrator.budget_status()["spent_cost_usd"] == 1.0 + + +def test_provider_budget_meter_survives_restart() -> None: + with tempfile.TemporaryDirectory() as directory: + state_db = str(Path(directory) / "state.db") + first = TaskOrchestrator( + [_agent()], + state_db=state_db, + budget_max_output_tokens=2, + ) + first._record_in_flight_provider_usage( + _agent(), + {"completion_tokens": 2}, + "", + ) + first.close() + + second = TaskOrchestrator( + [_agent()], + state_db=state_db, + budget_max_output_tokens=2, + ) + try: + assert second.budget_status()["spent_output_tokens"] == 2 + with pytest.raises(BudgetExceededError, match="spend budget exceeded"): + second._raise_if_spend_budget_exceeded() + finally: + second.close() + + def test_http_over_budget_returns_429() -> None: token = "budget_token" orchestrator = TaskOrchestrator([_agent()], budget_max_output_tokens=1) diff --git a/tests/test_local_gateway.py b/tests/test_local_gateway.py index aed2018b4..0caee0534 100644 --- a/tests/test_local_gateway.py +++ b/tests/test_local_gateway.py @@ -651,9 +651,17 @@ def test_local_responses_response_maps_reasoning_and_tool_calls() -> None: def test_local_provider_scheme_validation_rejects_remote_and_malformed_ports() -> None: assert _is_local_provider_url("local://127.0.0.1:8080/v1") + assert not _is_local_provider_url("local://host.docker.internal:8080/v1") assert not _is_local_provider_url("mlx://example.com:8080/v1") assert not _is_local_provider_url("mlx://127.0.0.1:8080/v1") assert not _is_local_provider_url("local://127.0.0.1:not-a-port/v1") + with pytest.raises(ValueError, match="explicit loopback endpoint"): + ModelAgent( + "docker_host_agent", + "local-model", + base_url="local://host.docker.internal:8080/v1", + local_credential_key="LOCAL_GATEWAY_TOKEN", + ) @pytest.mark.parametrize( diff --git a/tests/test_openai_passthrough.py b/tests/test_openai_passthrough.py index 57fea7aa3..9668692c8 100644 --- a/tests/test_openai_passthrough.py +++ b/tests/test_openai_passthrough.py @@ -184,6 +184,26 @@ def test_proxy_completion_blocks_before_structured_workflow_when_budget_is_excee ) +def test_structured_budget_stops_before_the_next_workflow_provider_call() -> None: + orchestrator = _build(budget_max_output_tokens=1) + + with ( + patch.object(orchestrator.client, "chat", wraps=orchestrator.client.chat) as chat, + patch.object(orchestrator.client, "proxy_send", wraps=orchestrator.client.proxy_send) as send, + pytest.raises(BudgetExceededError, match="spend budget exceeded"), + ): + orchestrator.proxy_completion( + { + "model": "mock-planner", + "messages": [{"role": "user", "content": "extract JSON"}], + "response_format": {"type": "json_object"}, + } + ) + + assert chat.call_count == 1 + send.assert_not_called() + + def test_model_client_request_settings_are_thread_local() -> None: client = _build().client previous_temperature = client.default_temperature diff --git a/tests/test_provider_embeddings.py b/tests/test_provider_embeddings.py index 02b4c6d36..69c33d5d2 100644 --- a/tests/test_provider_embeddings.py +++ b/tests/test_provider_embeddings.py @@ -115,6 +115,34 @@ def embed_many(selected: ModelAgent, inputs: list[str]) -> list[list[float]]: } +def test_default_embedding_backend_observes_runtime_agent_addition() -> None: + calls: list[tuple[str, list[str]]] = [] + + def embed_many(selected: ModelAgent, inputs: list[str]) -> list[list[float]]: + calls.append((selected.model, inputs)) + return [[1.0] for _input in inputs] + + orchestrator = SimpleNamespace( + candidates=[], + client=SimpleNamespace(embed_many=embed_many), + ) + coordinator = CostRoutingCoordinator(orchestrator) + orchestrator.candidates.append( + ModelAgent( + "runtime_embedding_agent", + "runtime-embedding-model", + "https://gateway.example/v1", + tags=("embedding",), + ) + ) + + result = coordinator.complete_embeddings_batch(["added later"]) + + assert calls == [("runtime-embedding-model", ["added later"])] + assert result["embeddings"][0]["embedding"] == [1.0] + assert result["provider"] == "gateway.example" + + def test_embedding_client_fails_closed_before_or_after_transport(monkeypatch) -> None: client = ModelClient(max_retries=0) mock_agent = ModelAgent("mock_embedding", "embedding-model", "mock://embedding") diff --git a/tests/test_responses_attribution_routing_http_honesty.py b/tests/test_responses_attribution_routing_http_honesty.py index 2e167260a..531e23fb7 100644 --- a/tests/test_responses_attribution_routing_http_honesty.py +++ b/tests/test_responses_attribution_routing_http_honesty.py @@ -122,6 +122,24 @@ def test_http_responses_rejects_routing_latency_tolerant_true() -> None: thread.join(timeout=5) +def test_http_responses_rejects_routing_unknown_key() -> None: + server, thread, port = _server() + try: + status, body = _post( + port, + { + "model": "mock-planner", + "input": "routing junk", + "routing": {"channel": "sync", "region": "us-east"}, + }, + ) + assert status == 400, body + assert "invalid_routing" in json.dumps(body) + finally: + server.shutdown() + thread.join(timeout=5) + + def test_http_responses_rejects_attribution_unknown_dimension() -> None: server, thread, port = _server() try: diff --git a/tests/test_temperature_capability_negotiation_honesty.py b/tests/test_temperature_capability_negotiation_honesty.py index ddaf0ec50..9e9d1f83a 100644 --- a/tests/test_temperature_capability_negotiation_honesty.py +++ b/tests/test_temperature_capability_negotiation_honesty.py @@ -2,6 +2,7 @@ from __future__ import annotations +import http.client import io import json import socket @@ -62,6 +63,29 @@ def test_non_negotiated_error_body_remains_available_to_the_caller() -> None: assert error.read() == expected +def test_incomplete_error_body_preserves_partial_capability_evidence() -> None: + """A truncated diagnostic must not replace the provider's original HTTP error.""" + partial = b"Unsupported value: temperature does not support 0.2 with this model" + + class _IncompleteBody: + def read(self) -> bytes: + raise http.client.IncompleteRead(partial) + + def close(self) -> None: + pass + + error = urllib.error.HTTPError( + "https://provider.example/v1/chat/completions", + 400, + "provider error", + {}, + _IncompleteBody(), + ) + + assert _temperature_capability_rejection(error) + assert error.read() == partial + + class _Response: """Minimal context-managed JSON response for transport tests."""