diff --git a/.gitignore b/.gitignore
index dd30cf213..50050edd5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,4 @@ tempcred.txt
# background-agent isolated worktree scratch state (nested git checkouts)
.claude/worktrees/
+.worktrees/
diff --git a/CHANGELOG.md b/CHANGELOG.md
index eedf1c4d2..d1a4d9687 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,38 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html)
### Fixed
+- 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.
+- Configured-gateway discovery now removes its blank bootstrap row after a
+ concrete catalog's chat candidates fail bounded readiness, so virtual
+ requests cannot bypass an authentication failure through an unprobed seed;
+ this retirement is process-local so a later startup can probe recovered
+ credentials, while explicit model pins still return their own typed error.
+- Queued embedding admissions now carry the durable registry's result
+ retention and the selected backend's polling cadence, so clients can poll
+ within the actual job lifecycle instead of guessing or failing closed on
+ missing lifecycle metadata.
+- Virtual structured workflows now exclude a same-endpoint candidate only
+ after both its synthesis and bounded repair violate the caller's schema,
+ then continue with the next eligible model on that endpoint. Explicit model
+ pins remain single-model and exhausted virtual pools return a typed error.
+- Configured-gateway runtime discovery now retains chat rows only after a
+ bounded structured-output probe, and virtual structured workflows share one
+ request-scoped missing-model exclusion set across evidence and synthesis.
+ Probe telemetry is separate from caller attempts, and explicit structured
+ requests keep their model pin throughout evidence, judgment, and synthesis.
+- Chat token accounting now uses valid provider usage or exact Rust raw-output
+ counts for ADR-declared tokenizer mappings. Unreconstructible prompts, tools,
+ multimodal input, unknown models, and missing stream usage are explicitly
+ unavailable; token-threshold routing stays synchronous, enabled budgets fail
+ closed, and API usage/cost fields no longer publish heuristic estimates
+ (ADR 0006).
+- Provider-embedding workers now propagate durable-claim renewal loss and
+ refresh ownership before terminal publication. Embedding token accounting
+ uses configured `pg_tiktoken` or the packaged Rust cl100k counter for exact
+ declared models, and otherwise fails closed without publishing estimated
+ usage or cost (ADR 0005). Chat accounting is governed separately by ADR 0006.
- The HTTP embedding endpoints (`/v1/embeddings`, `/v1/batch/embeddings`)
now correctly wire the coordinator's cheapest-price selection into an
*omitted* `model` (the common case), not only an explicitly-named model or
diff --git a/Dockerfile b/Dockerfile
index 3a153f9d7..47a017fc5 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -9,12 +9,29 @@
# see docs/kv-credentials.md for the bootstrap flow.
# Agents: defaults to the bundled mock pool; mount your own and set AGENTS_FILE:
# -v ./agents.json:/app/agents.json -e AGENTS_FILE=/app/agents.json
+ARG MATURIN_BUILDER_IMAGE=ghcr.io/pyo3/maturin@sha256:b6c8b59a0170b77eb31a35b56034abd39972483ad0ebfff344deaa42a85f3bd3
+FROM ${MATURIN_BUILDER_IMAGE} AS maturin-tools
+FROM rust:1.97.1-slim-bookworm@sha256:2775a09d208ff0d7c1f50490c45b62db929e87ba1dcbc3f2132ac71a704bcdd3 AS dependency-builder
+RUN apt-get update \
+ && apt-get install --no-install-recommends --yes build-essential \
+ && rm -rf /var/lib/apt/lists/*
+COPY --from=maturin-tools /usr/local/bin/uv /usr/local/bin/uv
+COPY --from=maturin-tools /usr/bin/maturin /usr/local/bin/maturin
+COPY requirements.lock /build/requirements.lock
+COPY rust/Cargo.toml rust/Cargo.lock /build/rust/
+COPY rust/token_counter/ /build/rust/token_counter/
+COPY contextual_orchestrator/ /build/contextual_orchestrator/
+RUN uv python install 3.12 \
+ && uv pip install --python 3.12 --require-hashes -r /build/requirements.lock --target /build/deps \
+ && maturin build --locked --release --manifest-path /build/rust/token_counter/Cargo.toml --out /build/wheels \
+ && uv pip install --python 3.12 /build/wheels/*.whl --target /build/deps
+
# python:3.12-slim
FROM python:3.12-slim@sha256:423ed6ab25b1921a477529254bfeeabf5855151dc2c3141699a1bfc852199fbf
WORKDIR /app
COPY pyproject.toml requirements.lock README.md LICENSE ./
-RUN pip install --no-cache-dir --require-hashes -r requirements.lock
+COPY --from=dependency-builder /build/deps/ /usr/local/lib/python3.12/site-packages/
COPY contextual_orchestrator/ /usr/local/lib/python3.12/site-packages/contextual_orchestrator/
COPY examples/ examples/
diff --git a/README.md b/README.md
index 60ffff831..0832ef456 100644
--- a/README.md
+++ b/README.md
@@ -154,7 +154,7 @@ Non-mock providers must use `https://` URLs and a **resolvable KV credential**
One public interface:
- `contextual-orchestrator` is the model-like control-plane candidate exposed to callers. `/v1/models` lists it first, followed by every configured worker candidate, including disabled candidates with their status.
-- `/v1/chat/completions` accepts normal chat messages, and `"stream": true` returns an OpenAI-compatible `text/event-stream` of `chat.completion.chunk` deltas terminated by `data: [DONE]`. `stream_options.include_usage=true` is accepted for ordinary chat streams and emits a usage-only chunk after the terminal stop chunk, labeled `usage_source: reported` when the provider returned usage or `usage_source: estimated` (never mislabeled `reported`) otherwise; single-agent `tools` passthrough accepts it the same way from the one non-streaming upstream call — the provider's own usage field when present, an honest estimate when the provider omits it; `response_format`-only structured passthrough (conduct mode) still rejects the combination before provider execution, since its usage comes from a multi-step workflow's cost ledger and may be unmeasured. In **route** mode the worker's tokens are streamed live as they arrive from the provider (real token streaming); in **conduct** mode the multi-step answer is produced then framed as deltas (a workflow can't honestly token-stream a synthesizer that hasn't run yet).
+- `/v1/chat/completions` accepts normal chat messages, and `"stream": true` returns an OpenAI-compatible `text/event-stream` of `chat.completion.chunk` deltas terminated by `data: [DONE]`. `stream_options.include_usage=true` is accepted for ordinary chat streams and emits a usage-only chunk after the terminal stop chunk. Valid provider counts carry `usage_source: reported` and `usage_measurement_status: measured`; missing or malformed counts carry `usage: null` and `usage_measurement_status: unavailable`. Single-agent `tools` passthrough follows the same rule and never reconstructs tool or multimodal framing. `response_format`-only structured passthrough (conduct mode) still rejects the combination before provider execution when workflow-level usage is unavailable. In **route** mode the worker's tokens are streamed live as they arrive from the provider (real token streaming); in **conduct** mode the multi-step answer is produced then framed as deltas (a workflow can't honestly token-stream a synthesizer that hasn't run yet).
- `TaskOrchestrator.complete()` decides whether to route to one worker or run a short workflow.
- `TaskOrchestrator.compare_to_baseline(prompts, mode)` (CLI `--eval PROMPT...`) measures the orchestration engine against a single-worker baseline — per-prompt and aggregate latency plus a structural coverage delta (contributing steps + verifier-pass presence). It is a measured tradeoff report, not a human-quality claim.
- Responses include orchestration mode metadata, and trusted callers can request the full trace for audit.
@@ -162,7 +162,7 @@ One public interface:
- The admin console can use [Clearfolio](https://github.com/ContextualWisdomLab/clearfolio) as its document viewer: pass `--clearfolio-url URL` (or `CONTEXTUAL_ORCHESTRATOR_CLEARFOLIO_URL`) and the Integrations view gains a Document Viewer card (open viewer / deep-link `{url}/viewer/{docId}`). Default: disabled, console unchanged.
- `/api/v1/provider_readiness/latest` reports provider liveness separately from an explicit chat readiness probe; `?refresh=true` re-probes instead of returning the cached result.
- `/api/v1/analytics_snapshots/latest` returns source-backed local KPI definitions (trace completeness, policy-safe run rate, successful chat requests, and related event-derived counts) from in-memory runtime state, localized via the same locale bundles as the admin console.
-- `/api/v1/spend_analytics/latest` exposes per-model token and cost spend aggregated from workflow runs. Output tokens use provider-reported `usage` when available and fall back to a ~4 chars/token estimate otherwise (each model row is labeled `usage_source: reported | mixed | estimated`); cost is computed only for models with an operator-supplied price (`TaskOrchestrator(price_per_million=...)`), otherwise reported as null with the model listed under `unpriced_models`. See [Observability & spend](#observability--spend).
+- `/api/v1/spend_analytics/latest` exposes per-model token and cost spend aggregated from workflow runs. Valid provider usage is authoritative; declared model IDs may use the packaged Rust tokenizer for exact raw textual output. Prompt framing, tools, multimodal input, unknown tokenizers, and missing native code remain unavailable. Cost is computed only when every required count and operator-supplied price is available. See [Observability & spend](#observability--spend).
- `/api/v1/sales_readiness/latest` exposes a local enterprise-pilot readiness gate for API compatibility, operator evidence, workflow traces, evaluation replay, security posture, analytics truthfulness, locale parity, and provider egress safety. It is process-local evidence, not a production compliance certificate.
- `/api/v1/commercial_readiness/latest` exposes a KRW 2,000,000,000 commercial due-diligence readiness gate. It is a buyer-review evidence snapshot, not a valuation guarantee or purchase commitment.
- `/api/v1/commercial_evidence_manifests/latest` shows the evidence gaps to resolve before commercial due diligence. The former `/api/v1/buyer_evidence_manifests/latest` route remains a deprecated compatibility alias.
@@ -201,15 +201,15 @@ See [docs/architecture.md](docs/architecture.md) for the source-backed analysis.
## Observability & spend
-Local spend observability, aggregated from in-memory workflow runs. It is honest by construction — estimates are labeled, and cost is only reported when a price is configured.
+Local spend observability, aggregated from in-memory workflow runs. It is honest by construction: counts are authoritative or explicitly unavailable, and cost is reported only when its required counts and prices are available.
```bash
curl -s http://127.0.0.1:8000/api/v1/spend_analytics/latest \
-H "authorization: Bearer $local_token" | jq '.totals, .by_model, .budget'
```
-- **Tokens.** `by_model[].output_tokens` uses the provider-reported `usage.completion_tokens` when a real worker returns it, and falls back to a `~4 chars/token` estimate otherwise. Each row carries `usage_source`: `reported` (all steps reported), `mixed`, or `estimated`. `estimated_output_tokens` is always the estimate, kept alongside for comparison. `measurement_status` is `local_runtime_estimate`, not production telemetry.
-- **Cost.** Supply a price table to turn tokens into money — `TaskOrchestrator(price_per_million={"gpt-5.5": 10.0})` (USD per 1M output tokens). Models without a price appear under `unpriced_models` with `estimated_cost_usd: null`. No prices are assumed or fabricated.
+- **Tokens.** `by_model[].output_tokens` uses provider-reported completion/output tokens first. For exact full model IDs declared by ADR 0006, a missing output count may use the packaged Rust tokenizer over raw textual output only. Rows carry `usage_source: reported | tokenizer | mixed | unavailable`; unavailable rows return `output_tokens: null`. Prompt tokens are provider-reported or null because chat framing is not reconstructed.
+- **Cost.** Supply a price table to turn authoritative output tokens into money — `TaskOrchestrator(price_per_million={"gpt-5.5": 10.0})` (USD per 1M output tokens). Models without a price appear under `unpriced_models`; `cost_usd` remains null when a price or required token count is unavailable. No prices or token counts are assumed.
- **Budget cap.** Set an operator cap to refuse runaway spend (default: no cap):
```bash
@@ -217,7 +217,7 @@ curl -s http://127.0.0.1:8000/api/v1/spend_analytics/latest \
--budget-max-output-tokens 2000000 --budget-max-cost-usd 50
```
- Or in code: `TaskOrchestrator(budget_max_output_tokens=..., budget_max_cost_usd=...)`. Once spend reaches a cap, the next run is refused — `run()` raises `BudgetExceededError` and `/v1/chat/completions` returns HTTP `429 budget_exceeded`. Current state is in `spend_analytics()["budget"]` (`enabled`, limits, `spent_*`, `remaining_*`, `exceeded`). Cost caps require a price table; token caps do not.
+ Or in code: `TaskOrchestrator(budget_max_output_tokens=..., budget_max_cost_usd=...)`. Once spend reaches a cap, the next run is refused — `run()` raises `BudgetExceededError` and `/v1/chat/completions` returns HTTP `429 budget_exceeded`. An enabled budget also fails closed when a required count or price is unavailable. Current state is in `spend_analytics()["budget"]` (`enabled`, limits, nullable `spent_*`/`remaining_*`, `measurement_status`, `enforcement_status`, `exceeded`). Cost caps require a complete price table; token caps require authoritative output counts.
- **Admin.** The `/admin` **Observability** view renders the totals and the per-model table (unpriced models show an `unpriced` chip).
These are process-local measured signals for a stdlib lab, not a billing system or production compliance data.
@@ -240,7 +240,9 @@ is read from a **KV config store**, never `os.getenv`.
first-class dimensions catalogued in `cost_attribution_dimensions`: **account,
service, upstream API/provider, model name, team, group, company**. Token
counts reuse `pg-llm-batch`'s `pg_tiktoken` counter when a Postgres DSN is
- configured, and fall back to a deterministic heuristic otherwise.
+ configured. Valid provider usage is authoritative; missing chat framing,
+ tool, multimodal, or unknown-tokenizer counts remain explicitly unavailable
+ instead of falling back to a deterministic heuristic.
- **Canonical Billing export.** Install the published `metering_billing`
producer SDK, create its durable outbox, and pass
`CanonicalUsageRecordSink(event_builder=build_contextual_usage_event,
diff --git a/contextual_orchestrator/__init__.py b/contextual_orchestrator/__init__.py
index b0b84839f..3a254768a 100644
--- a/contextual_orchestrator/__init__.py
+++ b/contextual_orchestrator/__init__.py
@@ -11,6 +11,7 @@
LocalEmbeddingBatchBackend,
PgLlmBatchBackend,
PgLlmBatchEmbeddingBackend,
+ ProviderEmbeddingBatchBackend,
RoutingDecision,
RoutingHints,
RoutingPolicy,
@@ -61,7 +62,14 @@
parse_reasoning_effort_profile,
snapshot_role_effort_catalog,
)
-from .token_counting import HeuristicTokenCounter, build_token_counter
+from .token_counting import (
+ NativeCl100kTokenCounter,
+ NativeExactTokenCounter,
+ TokenCountUnavailable,
+ UnavailableTokenCounter,
+ build_embedding_token_counter,
+ build_token_counter,
+)
from .response_cache import (
RedisResponseCacheProvider,
ResponseCacheProvider,
@@ -119,7 +127,11 @@
# config / tokens
"InMemoryConfigStore",
"get_config_store",
- "HeuristicTokenCounter",
+ "NativeCl100kTokenCounter",
+ "NativeExactTokenCounter",
+ "TokenCountUnavailable",
+ "UnavailableTokenCounter",
+ "build_embedding_token_counter",
"build_token_counter",
"ResponseCacheProvider",
"RedisResponseCacheProvider",
@@ -139,6 +151,7 @@
"EmbeddingBatchResultItem",
"LocalEmbeddingBatchBackend",
"PgLlmBatchEmbeddingBackend",
+ "ProviderEmbeddingBatchBackend",
"heuristic_embedding",
"build_embeddings_jsonl_body",
"cheapest_upstream",
diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py
index 091c38b8d..35398ce53 100644
--- a/contextual_orchestrator/__main__.py
+++ b/contextual_orchestrator/__main__.py
@@ -8,6 +8,7 @@
import sys
from dataclasses import replace
+from .chat_capability import is_chat_compatible_model_id
from .cost_ledger import PriceBook
from .cost_router import CostRoutingCoordinator
from .credentials import get_credential, register_credential
@@ -15,6 +16,7 @@
from .kv_config import InMemoryConfigStore
from .model_discovery import (
CONFIGURED_GATEWAY_CREDENTIAL_NAME,
+ DiscoveredModel,
PROVIDER_MODEL_SOURCES,
ProviderModelSource,
agent_from_discovered,
@@ -609,6 +611,33 @@ def _discover_models_command(argv: list[str]) -> None:
raise SystemExit(1)
+def _probe_configured_gateway_structured_chat(
+ orchestrator: TaskOrchestrator,
+ model: DiscoveredModel,
+) -> bool:
+ """Return whether one configured-gateway row proves bounded structured chat."""
+ agent = replace(agent_from_discovered(model), disabled=False)
+ payload = {
+ "model": agent.model,
+ "messages": [
+ {
+ "role": "user",
+ "content": 'Return only this JSON object: {"status":"ok"}.',
+ }
+ ],
+ "response_format": {"type": "json_object"},
+ "max_tokens": 8,
+ "stream": False,
+ }
+ try:
+ response = orchestrator.client.probe_structured_chat(agent, payload)
+ content = response["choices"][0]["message"]["content"]
+ parsed = json.loads(content)
+ except Exception: # noqa: BLE001 - startup capability probe is fail-closed
+ return False
+ return isinstance(parsed, dict) and parsed.get("status") == "ok"
+
+
def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, list[str]]:
"""Discover and activate routable chat models without runtime env transport.
@@ -625,11 +654,45 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l
for model in discovered
if not model.evidence_only and is_discovered_chat_candidate(model)
]
+ configured_gateway_probe_required = any(
+ model.provider_name == "configured_gateway"
+ and get_credential(model.credential_name) is not None
+ for model in chat_models
+ )
+ failed_configured_gateway_probe_ids: set[str] = set()
+ if configured_gateway_probe_required:
+ probed_chat_models = []
+ for model in chat_models:
+ if (
+ model.provider_name == "configured_gateway"
+ and not _probe_configured_gateway_structured_chat(orchestrator, model)
+ ):
+ failed_configured_gateway_probe_ids.add(agent_id_for(model))
+ else:
+ probed_chat_models.append(model)
+ chat_models = probed_chat_models
existing_by_id = {agent.id: agent for agent in orchestrator.candidates}
+ runtime_models = [
+ model
+ for model in discovered
+ if not model.evidence_only
+ and (
+ model in chat_models
+ or "embedding" in model.capabilities
+ or (
+ agent_id_for(model) in failed_configured_gateway_probe_ids
+ and agent_id_for(model) in existing_by_id
+ )
+ )
+ ]
+ discovered_chat_agent_ids = {agent_id_for(model) for model in chat_models}
agents = []
- for model in chat_models:
+ for model in runtime_models:
existing = existing_by_id.get(agent_id_for(model))
- routable = is_routable_discovered_model(model)
+ embedding_routable = "embedding" in model.capabilities and model.spend_admitted
+ spend_routable = is_routable_discovered_model(model) or embedding_routable
+ structured_routable = agent_id_for(model) not in failed_configured_gateway_probe_ids
+ routable = embedding_routable or (spend_routable and structured_routable)
if existing is None:
agents.append(replace(agent_from_discovered(model), disabled=not routable))
elif "discovered" not in existing.tags:
@@ -654,26 +717,56 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l
max_output_tokens=max_output_tokens,
context_window=context_window,
)
- if existing is not None and not routable:
- tags = (*existing.tags, "spend:blocked")
- if existing.disabled and "spend:blocked" not in existing.tags:
- tags = (*tags, "spend:blocked:preserve-disabled")
+ if existing is not None and (not routable or not structured_routable):
+ block_markers = {
+ "spend:blocked",
+ "spend:blocked:preserve-disabled",
+ "structured:blocked",
+ "structured:blocked:preserve-disabled",
+ }
+ preserve_disabled = existing.disabled and (
+ not block_markers.intersection(existing.tags)
+ or any(tag.endswith(":preserve-disabled") for tag in existing.tags)
+ )
+ blocked_tags = []
+ if not spend_routable:
+ blocked_tags.append("spend:blocked")
+ if not structured_routable:
+ blocked_tags.append("structured:blocked")
+ tags = (
+ *(tag for tag in existing.tags if tag not in block_markers),
+ *blocked_tags,
+ )
+ if preserve_disabled:
+ tags = (*tags, *(f"{tag}:preserve-disabled" for tag in blocked_tags))
agents.append(
replace(
existing,
- disabled=True,
+ disabled=not routable or preserve_disabled,
tags=tuple(dict.fromkeys(tags)),
)
)
- elif existing is not None and "spend:blocked" in existing.tags:
+ elif existing is not None and any(tag in existing.tags for tag in ("spend:blocked", "structured:blocked")):
agents.append(
replace(
existing,
- disabled="spend:blocked:preserve-disabled" in existing.tags,
+ disabled=any(
+ tag in existing.tags
+ for tag in (
+ "spend:blocked:preserve-disabled",
+ "structured:blocked:preserve-disabled",
+ )
+ ),
tags=tuple(
tag
for tag in existing.tags
- if tag not in {"spend:blocked", "spend:blocked:preserve-disabled"}
+ if tag
+ not in {
+ "spend:blocked",
+ "spend:blocked:preserve-disabled",
+ "structured:blocked",
+ "structured:blocked:preserve-disabled",
+ }
),
)
)
@@ -684,21 +777,39 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l
if agents
else {"added": [], "updated": []}
)
- if any(model.provider_name == "configured_gateway" for model in chat_models):
+ # Once the configured endpoint returned a concrete catalog, its blank
+ # bootstrap row must never remain callable. In particular, if every
+ # catalog row fails the structured/auth readiness probe, retaining the
+ # unprobed blank seed would route virtual traffic around that fail-closed
+ # decision and repeatedly surface the same authentication failure.
+ if any(model.provider_name == "configured_gateway" for model in discovered):
for agent in tuple(orchestrator.candidates):
+ has_ready_discovered_chat = any(
+ candidate.id != agent.id
+ and not candidate.disabled
+ and candidate.id in discovered_chat_agent_ids
+ for candidate in orchestrator.candidates
+ )
if (
agent.provider_name == "configured_gateway"
and not agent.model.strip()
- and any(
- candidate.id != agent.id and not candidate.disabled
- for candidate in orchestrator.candidates
+ and (
+ has_ready_discovered_chat
+ or bool(failed_configured_gateway_probe_ids)
)
):
- orchestrator.remove_agent("default", agent.id)
+ if any(
+ candidate.id != agent.id and not candidate.disabled
+ for candidate in orchestrator.candidates
+ ):
+ orchestrator.remove_agent("default", agent.id)
+ else:
+ orchestrator._retire_runtime_agent(agent.id)
has_real_runtime_agent = any(
not candidate.disabled
and not candidate.base_url.startswith("mock://")
and "bootstrap_seed" not in candidate.tags
+ and is_chat_compatible_model_id(candidate.model)
for candidate in orchestrator.agents
)
for candidate in tuple(orchestrator.agents):
@@ -870,6 +981,7 @@ def main(argv: list[str] | None = None) -> None:
budget_max_output_tokens=args.budget_max_output_tokens,
budget_max_cost_usd=args.budget_max_cost_usd,
cache_ttl=args.cache_ttl,
+ allow_empty_agents=args.auto_discover_model_agents,
role_effort_catalog=(
default_role_effort_catalog() if args.role_effort_catalog == "default" else None
),
diff --git a/contextual_orchestrator/admin.py b/contextual_orchestrator/admin.py
index 9717eeeed..b388f0aec 100644
--- a/contextual_orchestrator/admin.py
+++ b/contextual_orchestrator/admin.py
@@ -55,10 +55,10 @@
"observability_title": "Observability",
"spend_title": "Spend",
"spend_model": "Model",
- "spend_output_tokens": "Est. output tokens",
- "spend_prompt_tokens": "Est. prompt tokens",
+ "spend_output_tokens": "Output tokens",
+ "spend_prompt_tokens": "Prompt tokens",
"spend_steps": "Steps",
- "spend_cost": "Est. cost",
+ "spend_cost": "Cost (USD)",
"spend_runs": "Runs",
"settings_title": "Settings",
"session_title": "Operator session",
@@ -81,13 +81,15 @@
"readiness_summary_text": "Sales and commercial criteria passed: {pass}. Need attention: {warn}. Failed: {fail}. See the rows below to fix what failed.",
"readiness_source": "Readiness source",
"readiness_measurement_status": "Measurement status",
- "measurement_local_runtime_estimate": "Estimated on this server",
+ "measurement_measured": "Provider measured",
+ "measurement_exact_tokenizer": "Exact tokenizer",
+ "measurement_unavailable": "Unavailable",
"measurement_local_runtime_snapshot": "Measured on this server",
"measurement_local_runtime": "Generated locally on this server",
"measurement_estimate": "Estimated",
"measurement_unknown": "Unknown",
"spend_no_price": "No price set",
- "spend_no_price_action": "Add provider pricing to estimate cost.",
+ "spend_no_price_action": "Add provider pricing to calculate cost.",
"readiness_remediation_label": "Remediation",
"sales_readiness": "Sales readiness",
"sales_readiness_title": "Sales Readiness",
@@ -314,10 +316,10 @@
"observability_title": "관측",
"spend_title": "비용",
"spend_model": "모델",
- "spend_output_tokens": "추정 출력 토큰",
- "spend_prompt_tokens": "추정 입력 토큰",
+ "spend_output_tokens": "출력 토큰",
+ "spend_prompt_tokens": "입력 토큰",
"spend_steps": "단계",
- "spend_cost": "추정 비용",
+ "spend_cost": "비용 (USD)",
"spend_runs": "실행",
"settings_title": "설정",
"session_title": "운영자 세션",
@@ -340,13 +342,15 @@
"readiness_summary_text": "판매 및 상용 기준 통과 {pass}개, 주의 {warn}개, 실패 {fail}개. 아래 행에서 실패 항목을 해결하세요.",
"readiness_source": "준비 근거",
"readiness_measurement_status": "측정 상태",
- "measurement_local_runtime_estimate": "이 서버에서 추정됨",
+ "measurement_measured": "공급자 측정값",
+ "measurement_exact_tokenizer": "정확 토크나이저",
+ "measurement_unavailable": "사용 불가",
"measurement_local_runtime_snapshot": "이 서버에서 측정됨",
"measurement_local_runtime": "이 서버에서 생성됨",
"measurement_estimate": "추정값",
"measurement_unknown": "알 수 없음",
"spend_no_price": "가격 미설정",
- "spend_no_price_action": "비용을 추정하려면 공급자 가격을 추가하세요.",
+ "spend_no_price_action": "비용을 계산하려면 공급자 가격을 추가하세요.",
"readiness_remediation_label": "보완 조치",
"sales_readiness": "판매 준비도",
"sales_readiness_title": "판매 준비도",
@@ -1052,7 +1056,7 @@
- | Model | Est. output tokens | Steps | $/1M | Est. cost |
|---|
+ | Model | Output tokens | Steps | $/1M | Cost (USD) |
|---|
@@ -1289,7 +1293,9 @@
renderReadiness();
}
const MEASUREMENT_STATUS_KEYS = {
- local_runtime_estimate: "measurement_local_runtime_estimate",
+ measured: "measurement_measured",
+ exact_tokenizer: "measurement_exact_tokenizer",
+ unavailable: "measurement_unavailable",
local_runtime_snapshot: "measurement_local_runtime_snapshot",
estimate: "measurement_estimate",
unknown: "measurement_unknown"
@@ -1307,20 +1313,20 @@
if (statusEl) statusEl.textContent = statusLabel(spend.measurement_status);
const totalsEl = document.getElementById("spendTotals");
if (totalsEl) {
- const cost = totals.estimated_cost_usd == null ? "—" : ("$" + totals.estimated_cost_usd);
+ const cost = totals.cost_usd == null ? "—" : ("$" + totals.cost_usd);
totalsEl.innerHTML = [
[t("spend_runs") || "Runs", totals.run_count ?? 0],
- [t("spend_output_tokens") || "Est. output tokens", totals.estimated_output_tokens ?? 0],
- [t("spend_prompt_tokens") || "Est. prompt tokens", totals.estimated_prompt_tokens ?? 0],
- [t("spend_cost") || "Est. cost (USD)", cost]
+ [t("spend_output_tokens") || "Output tokens", totals.output_tokens ?? "—"],
+ [t("spend_prompt_tokens") || "Prompt tokens", totals.prompt_tokens ?? "—"],
+ [t("spend_cost") || "Cost (USD)", cost]
].map(([l, v]) => `${escapeHtml(l)}${escapeHtml(v)}
`).join("");
}
const rowsEl = document.getElementById("spendRows");
if (rowsEl) {
rowsEl.innerHTML = (spend.by_model || []).map(row => {
const price = row.price_per_million_usd == null ? "—" : escapeHtml(row.price_per_million_usd);
- const cost = row.estimated_cost_usd == null ? `${escapeHtml(t("spend_no_price"))}` : ("$" + escapeHtml(row.estimated_cost_usd));
- return `| ${escapeHtml(row.model)} | ${escapeHtml(row.estimated_output_tokens)} | ${escapeHtml(row.step_count)} | ${price} | ${cost} |
`;
+ const cost = row.cost_usd == null ? `${escapeHtml(t("spend_no_price"))}` : ("$" + escapeHtml(row.cost_usd));
+ return `| ${escapeHtml(row.model)} | ${escapeHtml(row.output_tokens ?? "—")} | ${escapeHtml(row.step_count)} | ${price} | ${cost} |
`;
}).join("") || `| ${t("no_trace")} |
`;
}
const noteEl = document.getElementById("spendNote");
diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py
index ac043c7aa..03f89ba3a 100644
--- a/contextual_orchestrator/api_contract.py
+++ b/contextual_orchestrator/api_contract.py
@@ -22,6 +22,61 @@
},
},
"schemas": {
+ "AuthoritativeUsage": {
+ "type": ["object", "null"],
+ "required": ["prompt_tokens", "completion_tokens"],
+ "properties": {
+ "prompt_tokens": {"type": "integer", "minimum": 0},
+ "completion_tokens": {"type": "integer", "minimum": 0},
+ "total_tokens": {"type": "integer", "minimum": 0},
+ },
+ },
+ "UsageCost": {
+ "type": "object",
+ "required": ["cost_amount", "currency_code", "measurement_status"],
+ "properties": {
+ "cost_amount": {"type": ["number", "null"]},
+ "currency_code": {"type": "string"},
+ "measurement_status": {
+ "type": "string",
+ "enum": ["measured", "unavailable"],
+ },
+ },
+ },
+ "ChatCompletionResponse": {
+ "type": "object",
+ "required": ["id", "object", "created", "model", "choices", "usage", "usage_measurement_status"],
+ "oneOf": [
+ {
+ "properties": {
+ "usage_measurement_status": {"const": "measured"},
+ "usage": {
+ "type": "object",
+ "required": ["prompt_tokens", "completion_tokens"],
+ },
+ }
+ },
+ {
+ "properties": {
+ "usage_measurement_status": {"const": "unavailable"},
+ "usage": {"type": "null"},
+ }
+ },
+ ],
+ "properties": {
+ "id": {"type": "string"},
+ "object": {"type": "string"},
+ "created": {"type": "integer"},
+ "model": {"type": "string"},
+ "choices": {"type": "array", "items": {"type": "object"}},
+ "usage": {"$ref": "#/components/schemas/AuthoritativeUsage"},
+ "usage_measurement_status": {
+ "type": "string",
+ "enum": ["measured", "unavailable"],
+ },
+ "orchestration": {"type": "object"},
+ },
+ },
"ModelGroupWrite": {
"type": "object",
"required": ["group_name", "member_agent_ids"],
@@ -169,7 +224,15 @@
},
},
"responses": {
- "200": {"description": "Chat completion or SSE response"},
+ "200": {
+ "description": "Chat completion or SSE response; missing provider usage is explicitly unavailable",
+ "content": {
+ "application/json": {
+ "schema": {"$ref": "#/components/schemas/ChatCompletionResponse"}
+ },
+ "text/event-stream": {"schema": {"type": "string"}},
+ },
+ },
"400": {"description": "Invalid request"},
},
}
@@ -985,7 +1048,13 @@
"input_part_counts, map_reduce}"
)
},
- "202": {"description": "Batch accepted; poll GET /v1/batch/embeddings/{batch_id}"},
+ "202": {
+ "description": (
+ "Batch accepted with registry-owned job_retention_ms and "
+ "backend-owned poll_after_ms; poll GET "
+ "/v1/batch/embeddings/{batch_id}"
+ )
+ },
"503": {"description": "No enabled embedding-capable agent is available"},
},
}
diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py
index a7bae84fb..21c7fc8a7 100644
--- a/contextual_orchestrator/batch_job_registry.py
+++ b/contextual_orchestrator/batch_job_registry.py
@@ -35,7 +35,11 @@
import dataclasses
import json
+import threading
+import time
+import weakref
from collections.abc import MutableMapping
+from contextlib import contextmanager
from typing import Any, Callable, Iterator, Optional
# Registry entries expire after this many seconds so abandoned jobs do
@@ -44,6 +48,66 @@
DEFAULT_RETENTION_SECONDS = 7 * 24 * 3600
+def _claim_renewal_interval_seconds(lease_seconds: float) -> float:
+ """Return the cadence the durable registry uses to observe claim ownership."""
+ return max(0.05, min(lease_seconds / 3, 1.0))
+
+
+class ClaimNotAcquired(RuntimeError):
+ """Another worker owns a non-blocking durable job claim."""
+
+
+class _ClaimLease:
+ def __init__(
+ self,
+ claim: Any = None,
+ *,
+ lease_seconds: float | None = None,
+ lost_ownership: threading.Event | None = None,
+ ) -> None:
+ self._claim = claim
+ self._lease_seconds = lease_seconds
+ self._lost_ownership = lost_ownership or threading.Event()
+
+ def mark_lost(self) -> None:
+ """Record that this worker can no longer prove claim ownership."""
+ self._lost_ownership.set()
+
+ def ensure_owned(self, *, refresh: bool = False) -> None:
+ """Fail closed unless this worker still owns the durable claim."""
+ if self._claim is None:
+ return
+ if self._lost_ownership.is_set():
+ raise ClaimNotAcquired("durable job claim ownership was lost")
+ try:
+ if refresh:
+ if self._lease_seconds is None:
+ raise ClaimNotAcquired("durable job claim lease is unavailable")
+ retained = self._claim.extend(
+ self._lease_seconds,
+ replace_ttl=True,
+ )
+ else:
+ retained = self._claim.owned()
+ except Exception as exc: # noqa: BLE001 - redis is optional.
+ self.mark_lost()
+ raise ClaimNotAcquired("durable job claim ownership is unavailable") from exc
+ if not retained:
+ self.mark_lost()
+ raise ClaimNotAcquired("durable job claim ownership was lost")
+
+ def atomic_identity(self) -> tuple[str, Any]:
+ """Return the Valkey lock key and token for one atomic fenced write."""
+ if self._claim is None:
+ raise ClaimNotAcquired("durable job claim identity is unavailable")
+ token = getattr(getattr(self._claim, "local", None), "token", None)
+ name = getattr(self._claim, "name", None)
+ if not name or token is None:
+ self.mark_lost()
+ raise ClaimNotAcquired("durable job claim identity is unavailable")
+ return str(name), token
+
+
def _encode(value: Any) -> str:
"""Serialize one registry value (dataclasses included) to JSON."""
if dataclasses.is_dataclass(value) and not isinstance(value, type):
@@ -130,12 +194,217 @@ class JobRegistryFactory:
def __init__(self, client: Any = None, *, retention_seconds: int = DEFAULT_RETENTION_SECONDS) -> None:
self._client = client
self._retention_seconds = retention_seconds
+ self._local_locks: weakref.WeakValueDictionary[str, threading.Lock] = (
+ weakref.WeakValueDictionary()
+ )
+ self._local_locks_guard = threading.Lock()
+
+ def lock(
+ self,
+ name: str,
+ key: str,
+ *,
+ lease_seconds: float | None = None,
+ renew_until_epoch: float | None = None,
+ ):
+ """Return an atomic shard claim with bounded lease and acquisition wait."""
+ lock_name = f"batch_job_registry:{name}:claim:{key}"
+ if self._client is not None:
+ if lease_seconds is None or lease_seconds <= 0:
+ raise ValueError("durable claim lease_seconds must be positive")
+ claim = self._client.lock(
+ lock_name,
+ timeout=lease_seconds,
+ blocking=True,
+ blocking_timeout=lease_seconds,
+ thread_local=False,
+ )
+
+ @contextmanager
+ def acquired_claim():
+ if not claim.acquire():
+ raise ClaimNotAcquired(lock_name)
+ stop_renewal = threading.Event()
+ lost_ownership = threading.Event()
+ lease = _ClaimLease(
+ claim,
+ lease_seconds=lease_seconds,
+ lost_ownership=lost_ownership,
+ )
+ renewal_thread = None
+ if renew_until_epoch is not None:
+
+ def renew_claim() -> None:
+ interval = _claim_renewal_interval_seconds(lease_seconds)
+ while not stop_renewal.wait(interval):
+ remaining = renew_until_epoch - time.time()
+ if remaining <= 0:
+ lease.mark_lost()
+ return
+ try:
+ # redis-py's Lock.extend script checks the
+ # claim token before replacing the TTL (CAS).
+ renewed = claim.extend(
+ max(0.05, min(lease_seconds, remaining)),
+ replace_ttl=True,
+ )
+ if not renewed:
+ lease.mark_lost()
+ return
+ except Exception: # noqa: BLE001 - redis is optional.
+ lease.mark_lost()
+ return
+
+ renewal_thread = threading.Thread(
+ target=renew_claim,
+ name="job-claim-renewal",
+ daemon=True,
+ )
+ renewal_thread.start()
+ try:
+ yield lease
+ finally:
+ stop_renewal.set()
+ if renewal_thread is not None:
+ renewal_thread.join()
+ try:
+ claim.release()
+ except Exception as exc: # noqa: BLE001 - redis is optional.
+ if renew_until_epoch is None or type(exc).__name__ != "LockNotOwnedError":
+ raise
+
+ return acquired_claim()
+ with self._local_locks_guard:
+ lock = self._local_locks.get(lock_name)
+ if lock is None:
+ lock = threading.Lock()
+ self._local_locks[lock_name] = lock
+
+ @contextmanager
+ def acquired_local_claim():
+ with lock:
+ yield _ClaimLease()
+
+ return acquired_local_claim()
+
+ def publish_provider_embedding_terminal(
+ self,
+ claim: _ClaimLease,
+ job_id: str,
+ *,
+ status: str,
+ results: Any = None,
+ usage: Any = None,
+ error: Any = None,
+ ) -> None:
+ """Atomically publish one durable terminal state while its claim is owned."""
+ if self._client is None:
+ raise RuntimeError("atomic terminal publication requires a durable registry")
+ if status not in {"completed", "failed"}:
+ raise ValueError("terminal status must be completed or failed")
+ lock_name, token = claim.atomic_identity()
+ script = """
+ if redis.call('get', KEYS[1]) ~= ARGV[1] then return 0 end
+ local current = redis.call('hget', KEYS[2], ARGV[2])
+ if current ~= ARGV[3] and current ~= ARGV[4] then return 0 end
+ if ARGV[6] ~= '' then redis.call('hset', KEYS[3], ARGV[2], ARGV[6]) end
+ if ARGV[7] ~= '' then redis.call('hset', KEYS[4], ARGV[2], ARGV[7]) end
+ if ARGV[8] ~= '' then redis.call('hset', KEYS[5], ARGV[2], ARGV[8]) end
+ redis.call('hset', KEYS[2], ARGV[2], ARGV[5])
+ for index = 2, 5 do redis.call('expire', KEYS[index], ARGV[9]) end
+ return 1
+ """
+ published = self._client.eval(
+ script,
+ 5,
+ lock_name,
+ "batch_job_registry:provider_embedding_states",
+ "batch_job_registry:provider_embedding_results",
+ "batch_job_registry:provider_embedding_usage",
+ "batch_job_registry:provider_embedding_errors",
+ token,
+ job_id,
+ _encode("running"),
+ _encode("queued"),
+ _encode(status),
+ "" if results is None else _encode(results),
+ "" if usage is None else _encode(usage),
+ "" if error is None else _encode(error),
+ self._retention_seconds,
+ )
+ if not published:
+ claim.mark_lost()
+ raise ClaimNotAcquired("durable job claim ownership was lost before publication")
+
+ def mark_provider_embedding_running(self, claim: _ClaimLease, job_id: str) -> None:
+ """Atomically retain a pending job and mark it running under its claim."""
+ if self._client is None:
+ raise RuntimeError("atomic state transition requires a durable registry")
+ lock_name, token = claim.atomic_identity()
+ script = """
+ if redis.call('get', KEYS[1]) ~= ARGV[1] then return 0 end
+ local current = redis.call('hget', KEYS[2], ARGV[2])
+ if current ~= ARGV[3] and current ~= ARGV[4] then return 0 end
+ redis.call('hset', KEYS[2], ARGV[2], ARGV[4])
+ redis.call('expire', KEYS[2], ARGV[5])
+ return 1
+ """
+ transitioned = self._client.eval(
+ script,
+ 2,
+ lock_name,
+ "batch_job_registry:provider_embedding_states",
+ token,
+ job_id,
+ _encode("queued"),
+ _encode("running"),
+ self._retention_seconds,
+ )
+ if not transitioned:
+ claim.mark_lost()
+ raise ClaimNotAcquired("durable job was cancelled or claim ownership was lost")
+
+ def cancel_provider_embedding(self, job_id: str, *, reason: str) -> bool:
+ """Atomically cancel a durable job unless terminal publication won."""
+ if self._client is None:
+ raise RuntimeError("atomic cancellation requires a durable registry")
+ script = """
+ local current = redis.call('hget', KEYS[1], ARGV[1])
+ if current ~= ARGV[2] and current ~= ARGV[3] and current ~= ARGV[4] then
+ return 0
+ end
+ redis.call('hset', KEYS[2], ARGV[1], ARGV[5])
+ redis.call('hset', KEYS[1], ARGV[1], ARGV[6])
+ redis.call('expire', KEYS[1], ARGV[7])
+ redis.call('expire', KEYS[2], ARGV[7])
+ return 1
+ """
+ return bool(
+ self._client.eval(
+ script,
+ 2,
+ "batch_job_registry:provider_embedding_states",
+ "batch_job_registry:provider_embedding_cancellations",
+ job_id,
+ _encode("reserved"),
+ _encode("queued"),
+ _encode("running"),
+ _encode({"reason": reason}),
+ _encode("cancelled"),
+ self._retention_seconds,
+ )
+ )
@property
def durable(self) -> bool:
"""True when registries survive a process restart."""
return self._client is not None
+ @property
+ def retention_seconds(self) -> int:
+ """Return the configured terminal-result retention contract."""
+ return self._retention_seconds
+
def mapping(self, name: str, *, decode: Optional[Callable[[Any], Any]] = None) -> MutableMapping:
"""Return the registry called ``name`` — a dict unless Valkey is configured."""
if self._client is None:
diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py
index ae0e3e27c..5650e441d 100644
--- a/contextual_orchestrator/batch_routing.py
+++ b/contextual_orchestrator/batch_routing.py
@@ -25,6 +25,7 @@
import hashlib
import json
import logging
+import threading
import time
import uuid
from contextlib import nullcontext
@@ -33,6 +34,12 @@
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Protocol
+from .batch_job_registry import (
+ ClaimNotAcquired,
+ JobRegistryFactory,
+ _claim_renewal_interval_seconds,
+)
+
_LOGGER = logging.getLogger(__name__)
_ROUTING_CATEGORY = "routing"
@@ -94,7 +101,7 @@ def __init__(self, config_store: Any) -> None:
def _batch_enabled(self) -> bool:
return bool(self._config.get(_ROUTING_CATEGORY, "batch_enabled", True))
- def decide(self, hints: RoutingHints, prompt_tokens: int = 0) -> RoutingDecision:
+ def decide(self, hints: RoutingHints, prompt_tokens: int | None = None) -> RoutingDecision:
"""Return the routing decision for one request."""
if not self._batch_enabled():
return RoutingDecision("sync", "batch routing disabled by config")
@@ -114,6 +121,10 @@ def decide(self, hints: RoutingHints, prompt_tokens: int = 0) -> RoutingDecision
return RoutingDecision("batch", "latency-tolerant request routed to batch")
batch_min_tokens = int(self._config.get(_ROUTING_CATEGORY, "batch_min_tokens", 0))
+ if batch_min_tokens and prompt_tokens is None:
+ return RoutingDecision(
+ "sync", "prompt token count unavailable; conservative sync path"
+ )
if batch_min_tokens and prompt_tokens >= batch_min_tokens:
return RoutingDecision(
"batch", f"prompt tokens {prompt_tokens} >= batch_min_tokens {batch_min_tokens}"
@@ -215,8 +226,8 @@ class BatchResultItem:
custom_id: str
answer: str
- prompt_tokens: int = 0
- completion_tokens: int = 0
+ prompt_tokens: Optional[int] = None
+ completion_tokens: Optional[int] = None
attribution: Dict[str, Any] = field(default_factory=dict)
model: str = "contextual-orchestrator"
mode: str = "auto"
@@ -541,7 +552,8 @@ async def _download() -> Dict[str, Any]:
items: List[BatchResultItem] = []
for entry in responses:
custom_id = entry.get("custom_id", "")
- body = (entry.get("response") or {}).get("body", {})
+ response = entry.get("response", {}) or {}
+ body = response.get("body", {}) or {}
answer = _extract_answer(body)
usage = body.get("usage", {}) or {}
prompt_tokens = usage.get("prompt_tokens")
@@ -600,10 +612,15 @@ class EmbeddingBatchRequest:
model: str = "contextual-orchestrator"
custom_id: str = field(default_factory=lambda: f"emb_{uuid.uuid4().hex}")
attribution: Dict[str, Any] = field(default_factory=dict)
+ metadata: Dict[str, Any] = field(default_factory=dict)
source_index: int = 0
part_index: int = 0
part_count: int = 1
token_count: int = 0
+ token_start: int = 0
+ token_end: int = 0
+ shard_index: int = 0
+ routing_agent_id: str | None = None
zdr_only: bool = False
agent_id: Optional[str] = None
@@ -649,6 +666,7 @@ class EmbeddingBatchBackend(Protocol):
"""Submit/poll/retrieve contract shared by every embeddings batch backend."""
name: str
+ poll_after_ms: int
def submit(
self, requests: List[EmbeddingBatchRequest], metadata: Optional[Dict[str, Any]] = None
@@ -716,10 +734,9 @@ def __init__(
)
def _count_tokens(self, text: str, model: str) -> int:
- if self._token_counter is not None:
- return int(self._token_counter.count_text(text, model))
- # Dependency-free fallback: count word-ish units.
- return len(text.split())
+ if self._token_counter is None:
+ raise RuntimeError("an authoritative embedding tokenizer is required")
+ return int(self._token_counter.count_text(text, model))
def submit(
self, requests: List[EmbeddingBatchRequest], metadata: Optional[Dict[str, Any]] = None
@@ -750,6 +767,411 @@ def retrieve(self, job: BatchJob) -> List[EmbeddingBatchResultItem]:
return self._results.get(job.job_id, [])
+class ProviderEmbeddingBatchBackend:
+ """Queue provider embedding work and expose a durable polling lifecycle."""
+
+ name = "provider"
+
+ def __init__(
+ self,
+ runner: Callable[[List[EmbeddingBatchRequest]], tuple[List[List[float]], int]],
+ *,
+ job_registry: Any = None,
+ max_concurrency: int = 1,
+ claim_lease_seconds: float | None = None,
+ execution_timeout_seconds: float | None = None,
+ ) -> None:
+ if type(max_concurrency) is not int or max_concurrency < 1:
+ raise ValueError("max_concurrency must be a positive integer")
+ self._runner = runner
+ self._max_concurrency = max_concurrency
+ self._executor: ThreadPoolExecutor | None = None
+ self._executor_lock = threading.Lock()
+ self._closed = threading.Event()
+ self._registry = job_registry or JobRegistryFactory()
+ if self._registry.durable and (
+ claim_lease_seconds is None or claim_lease_seconds <= 0
+ ):
+ raise ValueError("durable provider backend claim lease must be positive")
+ self._claim_lease_seconds = claim_lease_seconds
+ # This is the backend's actual durable-claim observation cadence (the
+ # same formula used by ``JobRegistryFactory.lock``), not an HTTP-layer
+ # polling guess. Admission responses expose it so callers do not have
+ # to invent their own retry interval.
+ self.poll_after_ms = int(
+ 1000
+ * _claim_renewal_interval_seconds(claim_lease_seconds or 0)
+ )
+ if execution_timeout_seconds is not None and execution_timeout_seconds <= 0:
+ raise ValueError("provider embedding execution timeout must be positive")
+ self._execution_timeout_seconds = (
+ execution_timeout_seconds
+ if execution_timeout_seconds is not None
+ else self._registry.retention_seconds
+ )
+ self._terminal_events: Dict[str, threading.Event] = {}
+ self._results: Dict[str, List[EmbeddingBatchResultItem]] = (
+ job_registry.mapping(
+ "provider_embedding_results", decode=lambda raw: EmbeddingBatchResultItem(**raw)
+ )
+ if job_registry is not None
+ else {}
+ )
+ self._requests = (
+ job_registry.mapping(
+ "provider_embedding_requests",
+ decode=lambda raw: EmbeddingBatchRequest(**raw),
+ )
+ if job_registry is not None
+ else {}
+ )
+ self._states = (
+ job_registry.mapping("provider_embedding_states")
+ if job_registry is not None
+ else {}
+ )
+ self._usage = (
+ job_registry.mapping("provider_embedding_usage")
+ if job_registry is not None
+ else {}
+ )
+ self._errors = (
+ job_registry.mapping("provider_embedding_errors")
+ if job_registry is not None
+ else {}
+ )
+ self._deadlines = (
+ job_registry.mapping("provider_embedding_deadlines")
+ if job_registry is not None
+ else {}
+ )
+ self._cancellations = (
+ job_registry.mapping("provider_embedding_cancellations")
+ if job_registry is not None
+ else {}
+ )
+ pending_job_ids = [
+ job_id
+ for job_id in list(self._states)
+ if self._states.get(job_id) in {"queued", "running"}
+ ]
+ if pending_job_ids:
+ self._executor = ThreadPoolExecutor(max_workers=self._max_concurrency)
+ for job_id in pending_job_ids:
+ self._terminal_events[job_id] = threading.Event()
+ self._executor.submit(copy_context().run, self._run_job, job_id)
+
+ def close(self) -> None:
+ """Release the bounded worker pool owned by this backend."""
+ self._closed.set()
+ with self._executor_lock:
+ executor, self._executor = self._executor, None
+ if executor is not None:
+ executor.shutdown(wait=False, cancel_futures=True)
+
+ def __enter__(self) -> "ProviderEmbeddingBatchBackend":
+ return self
+
+ def __exit__(self, *_exc: Any) -> None:
+ self.close()
+
+ def __del__(self) -> None: # pragma: no cover - interpreter timing varies
+ try:
+ self.close()
+ except Exception: # noqa: BLE001, S110 - interpreter teardown must not raise
+ pass
+
+ def submit(
+ self, requests: List[EmbeddingBatchRequest], metadata: Optional[Dict[str, Any]] = None
+ ) -> BatchJob:
+ """Persist a queued job and return immediately with a pollable handle."""
+ job = self.reserve(requests, metadata=metadata)
+ self.start(job)
+ return job
+
+ def reserve(
+ self, requests: List[EmbeddingBatchRequest], metadata: Optional[Dict[str, Any]] = None
+ ) -> BatchJob:
+ """Persist provider work without making it executable yet."""
+ if self._closed.is_set():
+ raise RuntimeError("provider embedding backend is closed")
+ job_id = f"providerembed_{uuid.uuid4().hex}"
+ self._requests[job_id] = list(requests)
+ self._states[job_id] = "reserved"
+ return BatchJob(
+ job_id=job_id,
+ backend=self.name,
+ status="reserved",
+ request_count=len(requests),
+ )
+
+ def start(self, job: BatchJob) -> None:
+ """Make a fully registered reservation executable."""
+ with self._executor_lock:
+ if self._closed.is_set():
+ raise RuntimeError("provider embedding backend is closed")
+ if self._states.get(job.job_id) != "reserved":
+ return
+ self._states[job.job_id] = "queued"
+ self._terminal_events[job.job_id] = threading.Event()
+ if self._executor is None:
+ self._executor = ThreadPoolExecutor(max_workers=self._max_concurrency)
+ self._executor.submit(copy_context().run, self._run_job, job.job_id)
+
+ def _run_job(self, job_id: str) -> None:
+ """Execute or reclaim one persisted job until it becomes terminal."""
+ while not self._closed.is_set() and self._states.get(job_id) in {"queued", "running"}:
+ try:
+ deadline_epoch = self._execution_deadline(job_id)
+ break
+ except ClaimNotAcquired:
+ threading.Event().wait(min(0.05, self._claim_lease_seconds))
+ else:
+ return
+ while (
+ not self._closed.is_set()
+ and self._states.get(job_id) in {"queued", "running"}
+ and time.time() < deadline_epoch
+ ):
+ try:
+ with self._registry.lock(
+ "provider_embedding_job_execution",
+ job_id,
+ lease_seconds=self._claim_lease_seconds,
+ renew_until_epoch=deadline_epoch,
+ ) as execution_claim:
+ self._run_claimed_job(job_id, execution_claim)
+ except ClaimNotAcquired:
+ remaining = max(0.0, deadline_epoch - time.time())
+ threading.Event().wait(min(0.05, remaining))
+ if (
+ not self._closed.is_set()
+ and time.time() >= deadline_epoch
+ and self._states.get(job_id) in {"queued", "running"}
+ ):
+ self._fail_expired_job(job_id)
+ if self._states.get(job_id) in {"completed", "failed", "cancelled"}:
+ event = self._terminal_events.pop(job_id, None)
+ if event is not None:
+ event.set()
+
+ def _execution_deadline(self, job_id: str) -> float:
+ """Persist a bounded lifetime beginning with the first execution claim."""
+ existing = self._deadlines.get(job_id)
+ if existing is not None:
+ return float(existing)
+ with self._registry.lock(
+ "provider_embedding_job_execution",
+ job_id,
+ lease_seconds=self._claim_lease_seconds,
+ ):
+ existing = self._deadlines.get(job_id)
+ if existing is None:
+ request_count = len(self._requests[job_id])
+ deadline = time.time() + self._execution_timeout_seconds * max(1, request_count)
+ set_if_absent = getattr(self._deadlines, "set_if_absent", None)
+ if callable(set_if_absent):
+ set_if_absent(job_id, deadline)
+ else:
+ self._deadlines[job_id] = deadline
+ existing = self._deadlines[job_id]
+ return float(existing)
+
+ def _fail_expired_job(self, job_id: str) -> None:
+ """Claim and atomically terminate provider work past its deadline."""
+ while not self._closed.is_set() and self._states.get(job_id) in {"queued", "running"}:
+ try:
+ with self._registry.lock(
+ "provider_embedding_job_execution",
+ job_id,
+ lease_seconds=self._claim_lease_seconds,
+ ) as execution_claim:
+ self._publish_terminal(
+ job_id,
+ execution_claim,
+ status="failed",
+ error={
+ "error_type": "TimeoutError",
+ "http_status": None,
+ "provider_code": "provider_embedding_deadline_exceeded",
+ "retryable": True,
+ "failed_shard_index": None,
+ },
+ )
+ return
+ except ClaimNotAcquired:
+ threading.Event().wait(0.05)
+
+ def _run_claimed_job(self, job_id: str, execution_claim: Any) -> None:
+ """Run one claim attempt and atomically publish its terminal outcome."""
+ execution_claim.ensure_owned()
+ if self._registry.durable:
+ self._registry.mark_provider_embedding_running(execution_claim, job_id)
+ else:
+ with self._registry.lock(
+ "provider_embedding_job_states",
+ job_id,
+ lease_seconds=self._claim_lease_seconds,
+ ):
+ if self._states.get(job_id) not in {"queued", "running"}:
+ return
+ self._states[job_id] = "running"
+ requests = list(self._requests[job_id])
+ try:
+ vectors, prompt_tokens = self._runner(requests)
+ execution_claim.ensure_owned()
+ if time.time() >= self._execution_deadline(job_id):
+ self._publish_terminal(
+ job_id,
+ execution_claim,
+ status="failed",
+ error={
+ "error_type": "TimeoutError",
+ "http_status": None,
+ "provider_code": "provider_embedding_deadline_exceeded",
+ "retryable": True,
+ "failed_shard_index": None,
+ },
+ )
+ return
+ if len(vectors) != len(requests):
+ raise ValueError("provider embedding batch result count did not match inputs")
+ dimensions = {len(vector) for vector in vectors}
+ if vectors and (dimensions == {0} or len(dimensions) != 1):
+ raise ValueError("provider embedding batch dimensions were inconsistent")
+ items = [
+ EmbeddingBatchResultItem(
+ custom_id=request.custom_id,
+ index=index,
+ embedding=vector,
+ prompt_tokens=0,
+ model=request.model,
+ )
+ for index, (request, vector) in enumerate(zip(requests, vectors, strict=True))
+ ]
+ self._publish_terminal(
+ job_id,
+ execution_claim,
+ status="completed",
+ results=items,
+ usage={"prompt_tokens": int(prompt_tokens)},
+ )
+ except ClaimNotAcquired:
+ raise
+ except Exception as exc: # noqa: BLE001 - polling exposes bounded failure metadata
+ error = {
+ "error_type": type(exc).__name__,
+ "http_status": getattr(
+ exc, "client_status", getattr(exc, "status_code", None)
+ ),
+ "provider_code": getattr(
+ exc, "error_code", getattr(exc, "provider_code", None)
+ ),
+ "retryable": bool(getattr(exc, "retryable", False)),
+ "failed_shard_index": getattr(exc, "failed_shard_index", None),
+ }
+ self._publish_terminal(
+ job_id, execution_claim, status="failed", error=error
+ )
+
+ def _publish_terminal(
+ self,
+ job_id: str,
+ execution_claim: Any,
+ *,
+ status: str,
+ results: Any = None,
+ usage: Any = None,
+ error: Any = None,
+ ) -> None:
+ """Publish terminal state atomically for durable registries."""
+ if self._registry.durable:
+ self._registry.publish_provider_embedding_terminal(
+ execution_claim,
+ job_id,
+ status=status,
+ results=results,
+ usage=usage,
+ error=error,
+ )
+ return
+ with self._registry.lock(
+ "provider_embedding_job_states",
+ job_id,
+ lease_seconds=self._claim_lease_seconds,
+ ):
+ execution_claim.ensure_owned()
+ if self._states.get(job_id) not in {"queued", "running"}:
+ raise ClaimNotAcquired("provider embedding job is already terminal")
+ if results is not None:
+ self._results[job_id] = results
+ if usage is not None:
+ self._usage[job_id] = usage
+ if error is not None:
+ self._errors[job_id] = error
+ self._states[job_id] = status
+
+ def wait(self, job: BatchJob, *, timeout: float) -> Dict[str, Any]:
+ """Wait within the caller's explicit deadline for a terminal state."""
+ event = self._terminal_events.get(job.job_id)
+ if event is not None:
+ event.wait(timeout=timeout)
+ return self.poll(job)
+
+ def poll(self, job: BatchJob) -> Dict[str, Any]:
+ """Return queued, running, completed, or failed without blocking."""
+ status = str(self._states.get(job.job_id, "failed"))
+ document = {
+ "job_id": job.job_id,
+ "status": status,
+ "is_complete": status in {"completed", "failed", "cancelled"},
+ }
+ if status == "failed":
+ document["failure"] = dict(self._errors.get(job.job_id, {}))
+ elif status == "cancelled":
+ document["cancellation"] = dict(self._cancellations.get(job.job_id, {}))
+ elif status == "completed":
+ document["usage"] = dict(self._usage.get(job.job_id, {}))
+ return document
+
+ def cancel(self, job: BatchJob, *, reason: str) -> Dict[str, Any]:
+ """Mark queued/running work cancelled and discard any late provider result."""
+ if self._registry.durable:
+ cancelled = self._registry.cancel_provider_embedding(job.job_id, reason=reason)
+ status = "cancelled" if cancelled else str(self._states.get(job.job_id, "failed"))
+ else:
+ with self._registry.lock(
+ "provider_embedding_job_states", job.job_id,
+ lease_seconds=self._claim_lease_seconds,
+ ):
+ status = str(self._states.get(job.job_id, "failed"))
+ if status not in {"completed", "failed", "cancelled"}:
+ self._cancellations[job.job_id] = {"reason": reason}
+ self._states[job.job_id] = "cancelled"
+ status = "cancelled"
+ if status == "cancelled":
+ event = self._terminal_events.pop(job.job_id, None)
+ if event is not None:
+ event.set()
+ return {
+ "job_id": job.job_id,
+ "status": status,
+ "is_complete": status in {"completed", "failed", "cancelled"},
+ **(
+ {"cancellation": dict(self._cancellations.get(job.job_id, {}))}
+ if status == "cancelled"
+ else {}
+ ),
+ }
+
+ def retrieve(self, job: BatchJob) -> List[EmbeddingBatchResultItem]:
+ """Return provider embeddings computed during submission."""
+ return self._results.get(job.job_id, [])
+
+ def usage(self, job: BatchJob) -> Dict[str, int]:
+ """Return provider-reported batch usage without per-input allocation."""
+ return dict(self._usage.get(job.job_id, {}))
+
class PgLlmBatchEmbeddingBackend:
"""Embeddings batch backend that submits to **pg-llm-batch** and retrieves.
@@ -771,6 +1193,14 @@ def __init__(
job_registry: Any = None,
) -> None:
self._client = client
+ poll_interval = getattr(client, "poll_interval_seconds", None)
+ self.poll_after_ms = (
+ int(poll_interval * 1000)
+ if isinstance(poll_interval, (int, float))
+ and not isinstance(poll_interval, bool)
+ and poll_interval > 0
+ else 0
+ )
self._endpoint_alias = endpoint_alias
self._endpoint = endpoint
self._assembler = payload_assembler
diff --git a/contextual_orchestrator/cost_ledger.py b/contextual_orchestrator/cost_ledger.py
index 67cc6e1e1..998d0665b 100644
--- a/contextual_orchestrator/cost_ledger.py
+++ b/contextual_orchestrator/cost_ledger.py
@@ -1411,6 +1411,8 @@ def record_usage(
cost_amount, currency, price_known = self.price_book.compute_cost(
provider, model, prompt_tokens, completion_tokens
)
+ if measurement_status == "unavailable" and self.price_book.get_price(provider, model) is None:
+ price_known = False
if measurement_status not in MEASUREMENT_STATUSES:
raise ValueError("measurement_status must be measured, estimated, or unavailable")
record = UsageRecord(
@@ -1580,11 +1582,11 @@ def rollup(
)
status = _measurement_status_of(row)
price_status = "known" if row.get("price_known") else "unknown"
- row_cost = Decimal(str(row.get("cost_amount", 0)))
+ row_cost = Decimal(str(row.get("cost_amount") or 0))
bucket["record_count"] += 1
- bucket["prompt_tokens"] += int(row.get("prompt_tokens", 0))
- bucket["completion_tokens"] += int(row.get("completion_tokens", 0))
- bucket["total_tokens"] += int(row.get("total_tokens", 0))
+ bucket["prompt_tokens"] += int(row.get("prompt_tokens") or 0)
+ bucket["completion_tokens"] += int(row.get("completion_tokens") or 0)
+ bucket["total_tokens"] += int(row.get("total_tokens") or 0)
bucket["cost_amount"] += row_cost
bucket["cost_amount_by_status"][status] += row_cost
bucket["record_count_by_status"][status] += 1
@@ -1667,7 +1669,7 @@ def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[
for row in rows:
status = _measurement_status_of(row)
price_status = "known" if row.get("price_known") else "unknown"
- row_cost = Decimal(str(row.get("cost_amount", 0)))
+ row_cost = Decimal(str(row.get("cost_amount") or 0))
cost_amount_by_status[status] += row_cost
record_count_by_status[status] += 1
cost_amount_by_price_status[price_status] += row_cost
@@ -1681,9 +1683,9 @@ def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[
)
return {
"record_count": len(rows),
- "prompt_tokens": sum(int(row.get("prompt_tokens", 0)) for row in rows),
- "completion_tokens": sum(int(row.get("completion_tokens", 0)) for row in rows),
- "total_tokens": sum(int(row.get("total_tokens", 0)) for row in rows),
+ "prompt_tokens": sum(int(row.get("prompt_tokens") or 0) for row in rows),
+ "completion_tokens": sum(int(row.get("completion_tokens") or 0) for row in rows),
+ "total_tokens": sum(int(row.get("total_tokens") or 0) for row in rows),
"cost_amount": (
None
if measurement_status == "unavailable"
@@ -1704,7 +1706,17 @@ def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[
def records(self, start: Optional[int] = None, end: Optional[int] = None) -> List[Dict[str, Any]]:
"""Return raw usage record rows in the optional window."""
- return self.store.query(start, end)
+ rows = self.store.query(start, end)
+ for row in rows:
+ if row.get("measurement_status") == "unavailable":
+ for field in (
+ "prompt_tokens",
+ "completion_tokens",
+ "total_tokens",
+ "cost_amount",
+ ):
+ row[field] = None
+ return rows
def _mark_inline_success(self) -> None:
with self._inline_health_lock:
diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py
index ffbe32a97..751c9e366 100644
--- a/contextual_orchestrator/cost_router.py
+++ b/contextual_orchestrator/cost_router.py
@@ -22,6 +22,7 @@
import re
from contextvars import ContextVar
from dataclasses import replace
+from threading import Lock
from typing import Any, Dict, List, Optional
from .batch_routing import (
@@ -35,14 +36,19 @@
EmbeddingBatchResultItem,
LocalBatchBackend,
LocalEmbeddingBatchBackend,
+ ProviderEmbeddingBatchBackend,
RoutingHints,
RoutingPolicy,
)
-from .batch_job_registry import JobRegistryFactory, build_job_registry
+from .batch_job_registry import ClaimNotAcquired, JobRegistryFactory, build_job_registry
from .cost_ledger import CostLedger, PriceBook, PriceEntry
from .kv_config import InMemoryConfigStore
from .model_discovery import _currency_is_comparable
-from .token_counting import HeuristicTokenCounter, build_token_counter
+from .token_counting import (
+ TokenCountUnavailable,
+ build_embedding_token_counter,
+ build_token_counter,
+)
_RACE_USAGE_CONTEXT: ContextVar[dict[str, Any] | None] = ContextVar(
@@ -52,6 +58,7 @@
_EMBEDDING_CONFIG_CATEGORY = "routing"
_DEFAULT_EMBEDDING_MAX_TOKENS_PER_REQUEST = 280_000
_DEFAULT_EMBEDDING_MAX_CHARS_PER_PART = 240_000
+_DEFAULT_EMBEDDING_MAX_INPUTS_PER_REQUEST = 1
_BATCH_LEDGER_SETTLEMENT_TIMEOUT_SECONDS = 1.0
_EMBEDDING_UNIT_RE = re.compile(r"\S+\s*|\s+", re.UNICODE)
@@ -75,6 +82,7 @@ def __init__(
price_book: Optional[PriceBook] = None,
ledger: Optional[CostLedger] = None,
token_counter: Any = None,
+ embedding_token_counter: Any = None,
routing_policy: Optional[RoutingPolicy] = None,
batch_backend: Optional[BatchBackend] = None,
embedding_batch_backend: Optional[EmbeddingBatchBackend] = None,
@@ -88,10 +96,17 @@ def __init__(
self._race_usage_context = _RACE_USAGE_CONTEXT
if hasattr(orchestrator, "_race_usage_sink"):
orchestrator._race_usage_sink = self._record_race_endpoint_usage
- self.token_counter = token_counter or (
- build_token_counter(postgres_dsn) if postgres_dsn else HeuristicTokenCounter()
- )
+ self.token_counter = token_counter or build_token_counter(postgres_dsn)
+ if embedding_token_counter is not None:
+ self.embedding_token_counter = embedding_token_counter
+ elif token_counter is not None:
+ self.embedding_token_counter = token_counter
+ else:
+ self.embedding_token_counter = build_embedding_token_counter(postgres_dsn)
self.policy = routing_policy or RoutingPolicy(self.config)
+ self._resolve_virtual_embedding_target = False
+ self._uses_default_embedding_backend = embedding_batch_backend is None
+ self._embedding_backend_lock = Lock()
# Job registries live in Valkey when the credential registry carries
# batch_job_registry_valkey_url, so submitted jobs survive a process
# restart; otherwise they are the historical in-process dicts. Built
@@ -110,12 +125,32 @@ def __init__(
)
else:
self.batch_backend = batch_backend
- self.embedding_batch_backend: EmbeddingBatchBackend = (
- embedding_batch_backend
- or LocalEmbeddingBatchBackend(
- token_counter=self.token_counter, job_registry=registry
+ if embedding_batch_backend is not None:
+ self.embedding_batch_backend = embedding_batch_backend
+ else:
+ try:
+ embedding_agents = orchestrator._capability_agents("embedding")
+ except (AttributeError, RuntimeError):
+ embedding_agents = []
+ remote_embedding_agents = [
+ agent for agent in embedding_agents if not agent.base_url.startswith("mock://")
+ ]
+ if remote_embedding_agents:
+ self._resolve_virtual_embedding_target = True
+ self.embedding_batch_backend = self._provider_embedding_backend()
+ else:
+ self.embedding_batch_backend = LocalEmbeddingBatchBackend(
+ token_counter=self.embedding_token_counter, job_registry=registry
+ )
+ self._embedding_backends = {
+ self.embedding_batch_backend.name: self.embedding_batch_backend
+ }
+ if self._uses_default_embedding_backend and "local" not in self._embedding_backends:
+ self._embedding_backends["local"] = LocalEmbeddingBatchBackend(
+ token_counter=self.embedding_token_counter, job_registry=registry
)
- )
+ if self._uses_default_embedding_backend and "provider" not in self._embedding_backends:
+ self._embedding_backends["provider"] = self._provider_embedding_backend()
self._job_registry = registry
# job_id -> submitted BatchJob (so poll/retrieve can be driven by id)
self._batch_jobs = registry.mapping("batch_jobs", decode=lambda raw: BatchJob(**raw))
@@ -123,6 +158,7 @@ def __init__(
# keyed by batch id so poll/retrieve is idempotent (usage recorded once).
self._embedding_jobs = registry.mapping("embedding_jobs", decode=lambda raw: BatchJob(**raw))
self._embedding_models = registry.mapping("embedding_models")
+ self._embedding_owners = registry.mapping("embedding_owners")
self._embedding_requests = registry.mapping(
"embedding_requests", decode=lambda raw: EmbeddingBatchRequest(**raw)
)
@@ -130,7 +166,14 @@ def __init__(
self._embedding_part_counts = registry.mapping("embedding_part_counts")
self._embedding_part_limits = registry.mapping("embedding_part_limits")
self._embedding_documents = registry.mapping("embedding_documents")
- self._batch_documents = registry.mapping("batch_documents")
+ for recovered_job in list(self._embedding_jobs.values()):
+ try:
+ recovered_backend = self._embedding_backend_for(recovered_job)
+ except RuntimeError:
+ continue
+ start_embedding_job = getattr(recovered_backend, "start", None)
+ if callable(start_embedding_job):
+ start_embedding_job(recovered_job)
def _run_local_batch(
self, messages: List[Dict[str, str]], mode: str, model: str
@@ -172,6 +215,132 @@ def _agent_provider_model(agent: Any, fallback_model: str) -> tuple[str, str]:
provider = agent.provider_name or _provider_from_base_url(agent.base_url)
return provider or "unknown", agent.model or fallback_model
+ def _run_embedding_shard(
+ self, agent: Any, requests: List[EmbeddingBatchRequest]
+ ) -> tuple[List[List[float]], int]:
+ texts = [request.input_text for request in requests]
+ if all(request.token_count > 0 or not request.input_text for request in requests):
+ return self.orchestrator.client.embed(agent, texts), sum(
+ request.token_count for request in requests
+ )
+ vectors, provider_tokens = self.orchestrator.client.embed_with_usage(agent, texts)
+ if provider_tokens is None:
+ raise TokenCountUnavailable(
+ "provider embedding response omitted authoritative usage"
+ )
+ return vectors, provider_tokens
+
+ def _provider_embedding_backend(self) -> ProviderEmbeddingBatchBackend:
+ client = getattr(self.orchestrator, "client", None)
+ client_timeout = float(getattr(client, "timeout", 0))
+ return ProviderEmbeddingBatchBackend(
+ self._run_provider_embeddings,
+ job_registry=self.job_registry,
+ max_concurrency=getattr(client, "local_concurrency", 1),
+ claim_lease_seconds=(
+ client_timeout
+ if self.job_registry.durable and client_timeout > 0
+ else None
+ ),
+ execution_timeout_seconds=client_timeout if client_timeout > 0 else None,
+ )
+
+ def _run_provider_embeddings(
+ self, requests: List[EmbeddingBatchRequest]
+ ) -> tuple[List[List[float]], int]:
+ if not requests:
+ return [], 0
+ first = requests[0]
+ agent = (
+ self.orchestrator._agent(first.agent_id)
+ if first.agent_id is not None
+ else self.orchestrator.select_capability_agent("embedding", first.model)
+ )
+ if any(
+ request.model != first.model or request.agent_id != first.agent_id
+ for request in requests
+ ):
+ raise RuntimeError("provider embedding batch must retain one selected route")
+ max_tokens, _max_chars, max_inputs = self._embedding_request_limits()
+ vectors: List[List[float]] = []
+ prompt_tokens = 0
+ shard: List[EmbeddingBatchRequest] = []
+ shard_tokens = 0
+ for request in requests:
+ request_tokens = request.token_count or len(request.input_text.encode("utf-8"))
+ if shard and (
+ len(shard) >= max_inputs or shard_tokens + request_tokens > max_tokens
+ ):
+ shard_vectors, shard_usage = self._run_embedding_shard(agent, shard)
+ vectors.extend(shard_vectors)
+ prompt_tokens += shard_usage
+ shard = []
+ shard_tokens = 0
+ shard.append(request)
+ shard_tokens += request_tokens
+ if shard:
+ shard_vectors, shard_usage = self._run_embedding_shard(agent, shard)
+ vectors.extend(shard_vectors)
+ prompt_tokens += shard_usage
+ return vectors, prompt_tokens
+
+ def _refresh_embedding_backend(self) -> None:
+ if not self._uses_default_embedding_backend or isinstance(
+ self.embedding_batch_backend, ProviderEmbeddingBatchBackend
+ ):
+ return
+ try:
+ embedding_agents = self.orchestrator._capability_agents("embedding")
+ except (AttributeError, RuntimeError):
+ return
+ remote_agents = [
+ agent
+ for agent in embedding_agents
+ if not agent.base_url.startswith("mock://")
+ ]
+ if remote_agents:
+ with self._embedding_backend_lock:
+ if isinstance(
+ self.embedding_batch_backend, ProviderEmbeddingBatchBackend
+ ):
+ return
+ self._resolve_virtual_embedding_target = True
+ self.embedding_batch_backend = self._embedding_backends["provider"]
+
+ def _embedding_backend_for(self, job: BatchJob) -> EmbeddingBatchBackend:
+ """Keep already-submitted jobs bound to the backend that owns them."""
+ backend = self._embedding_backends.get(job.backend)
+ if backend is None:
+ raise RuntimeError(f"embedding backend {job.backend!r} is unavailable")
+ return backend
+
+ def close_embedding_backends(self) -> None:
+ """Release every worker backend created during this coordinator's lifetime."""
+ closed: set[int] = set()
+ for backend in self._embedding_backends.values():
+ if id(backend) in closed:
+ continue
+ closed.add(id(backend))
+ close = getattr(backend, "close", None)
+ if callable(close):
+ close()
+
+ def _embedding_backend_for_route(
+ self, model: str, agent_id: Optional[str]
+ ) -> EmbeddingBatchBackend:
+ if not self._uses_default_embedding_backend:
+ return self.embedding_batch_backend
+ if agent_id is not None:
+ agents = [self.orchestrator._agent(agent_id)]
+ else:
+ try:
+ agents = self.orchestrator._capability_agents("embedding", model)
+ except (AttributeError, RuntimeError):
+ agents = []
+ if any(not agent.base_url.startswith("mock://") for agent in agents):
+ return self._embedding_backends["provider"]
+ return self._embedding_backends["local"]
+
def _cheapest_capability_candidate(self, candidates: List[Any]) -> Any:
"""Pick the lowest-priced member of a capability candidate list.
@@ -344,8 +513,8 @@ def _record_race_endpoint_usage(self, endpoint_id: str, value: Any) -> None:
model_name=context["model_name"],
provider_model=provider_model,
workflow_run_id=context["workflow_run_id"],
- prompt_tokens=counts[0],
- completion_tokens=counts[1],
+ prompt_tokens=counts[0] if counts else None,
+ completion_tokens=counts[1] if counts else None,
)
context["records"].append(record)
@@ -382,24 +551,20 @@ def complete(
``usage_record_id``. Batch requests are dispatched to the batch backend
and return a job envelope; their cost is recorded on retrieval.
- For multi-step workflows, each trace step records one ledger
- row. Rows backed by provider-reported token counts are labeled
- ``measurement_status="measured"``; rows that fall back to heuristic
- estimates are labeled ``"estimated"``. The original request prompt is
- attributed at most once across unreported rows: it lands in full on the
- first unreported step, and later unreported steps estimate only their
- own output tokens. Provider-reported prompt counts remain attached to
- their respective rows because each trace step is a separate billable
- provider call; they are neither deduplicated against nor replaced by
- the fallback estimate for a different call.
+ Each trace step backed by valid provider token counts is ``measured``.
+ A missing count is recorded with an ``unavailable`` status and numeric
+ storage sentinels; API usage and cost remain null.
"""
if not isinstance(cache_bypass, bool):
raise TypeError("cache_bypass must be a boolean")
if type(zdr_only) is not bool:
raise TypeError("zdr_only must be a boolean")
routing_hints = hints if isinstance(hints, RoutingHints) else RoutingHints.from_mapping(hints)
- prompt_tokens_estimate = self.token_counter.count_messages(messages, model_name)
- decision = self.policy.decide(routing_hints, prompt_tokens_estimate)
+ try:
+ prompt_tokens = self.token_counter.count_messages(messages, model_name)
+ except TokenCountUnavailable:
+ prompt_tokens = None
+ decision = self.policy.decide(routing_hints, prompt_tokens)
if decision.channel == "batch" and provider_request is None:
request = BatchRequest(
@@ -643,11 +808,21 @@ def complete(
client_usage_records = [
item for item in records if item.usage_record_id not in race_record_ids
]
- result["usage"] = {
- "prompt_tokens": sum(item.prompt_tokens for item in client_usage_records),
- "completion_tokens": sum(item.completion_tokens for item in client_usage_records),
- "total_tokens": sum(item.total_tokens for item in client_usage_records),
- }
+ client_measurement_available = all(
+ item.measurement_status == "measured" for item in client_usage_records
+ )
+ result["usage"] = (
+ {
+ "prompt_tokens": sum(item.prompt_tokens for item in client_usage_records),
+ "completion_tokens": sum(item.completion_tokens for item in client_usage_records),
+ "total_tokens": sum(item.total_tokens for item in client_usage_records),
+ }
+ if client_measurement_available
+ else None
+ )
+ result["usage_measurement_status"] = (
+ "measured" if client_measurement_available else "unavailable"
+ )
currencies = {item.currency_code for item in records}
statuses = {item.measurement_status for item in records}
aggregate_measurement_status = (
@@ -699,30 +874,21 @@ def _record_completion(
prompt_tokens: Optional[int] = None,
completion_tokens: Optional[int] = None,
usage_record_id: Optional[str] = None,
+ measurement_status: Optional[str] = None,
):
"""Record one completion's usage + cost and return its ledger record.
- ``prompt_tokens``/``completion_tokens`` carry provider-reported counts;
- when either is missing the ledger falls back to heuristic estimates and
- the row is labeled ``measurement_status="estimated"`` instead of
- ``"measured"``. The status is exposed on the completion's ``cost``
- payload, batch retrieval results, and analytics usage-record rows so a
- buyer can always tell provider-measured spend from estimated spend.
- For multi-step structured workflows the caller passes the original
- request ``messages`` only for the first unreported step, keeping the
- request prompt attributed at most once per completion; later unreported
- steps pass empty messages and estimate their own output alone.
+ Missing provider counts are unavailable. Zeroes are persisted only as
+ schema sentinels beside that status and are never exposed as measured
+ free usage or cost.
"""
provider, model = provider_model
- measurement_status = (
- "measured"
- if prompt_tokens is not None and completion_tokens is not None
- else "estimated"
- )
- if prompt_tokens is None:
- prompt_tokens = self.token_counter.count_messages(messages, model)
- if completion_tokens is None:
- completion_tokens = self.token_counter.count_text(answer, model)
+ if measurement_status is None:
+ measurement_status = "measured" if (
+ prompt_tokens is not None and completion_tokens is not None
+ ) else "unavailable"
+ prompt_tokens = prompt_tokens if measurement_status != "unavailable" and prompt_tokens is not None else 0
+ completion_tokens = completion_tokens if measurement_status != "unavailable" and completion_tokens is not None else 0
return self.ledger.record_usage(
provider=provider,
model=model,
@@ -771,9 +937,7 @@ def record_stream_usage(
)
statuses = {record.measurement_status for record in records}
measurement_status = (
- "unavailable" if "unavailable" in statuses
- else "estimated" if "estimated" in statuses
- else "measured"
+ "unavailable" if "unavailable" in statuses else "measured"
)
currencies = {record.currency_code for record in records}
price_known = all(record.price_known for record in records)
@@ -818,12 +982,16 @@ def submit_batch(
raise BatchModelSelectionError(
"no eligible model-group member is available for this batch request"
) from exc
- prompt_token_estimates = {
- request.custom_id: self.token_counter.count_messages(
- request.messages, request.model
- )
- for request in prepared_requests
- }
+ prompt_token_estimates: dict[str, int] = {}
+ for request in prepared_requests:
+ try:
+ count = self.token_counter.count_messages(
+ request.messages, request.model
+ )
+ if isinstance(count, int) and count >= 0:
+ prompt_token_estimates[request.custom_id] = count
+ except Exception:
+ pass
job = self.batch_backend.submit(prepared_requests, metadata=metadata)
job.owner_id = owner_id
job.prompt_token_estimates = prompt_token_estimates
@@ -942,9 +1110,27 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict
if not usage_valid and item.custom_id not in prompt_token_estimates:
original_request = request_by_custom_id.get(item.custom_id)
if original_request is not None:
- prompt_token_estimates[item.custom_id] = self.token_counter.count_messages(
- original_request.messages, item.model
- )
+ try:
+ prompt_token_estimates[item.custom_id] = self.token_counter.count_messages(
+ original_request.messages, item.model
+ )
+ except Exception:
+ pass
+ if usage_valid:
+ prompt_toks = item.prompt_tokens
+ comp_toks = item.completion_tokens
+ status = "measured"
+ else:
+ prompt_toks = prompt_token_estimates.get(item.custom_id)
+ if prompt_toks is not None:
+ try:
+ comp_toks = self.token_counter.count_text(item.answer, item.model)
+ except Exception:
+ comp_toks = 0
+ status = "estimated"
+ else:
+ comp_toks = None
+ status = "unavailable"
records.append(
self._record_completion(
messages=fallback_messages,
@@ -955,12 +1141,9 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict
model_name=item.model,
provider_model=self._resolve_batch_provider_model(item),
workflow_run_id=job.job_id,
- prompt_tokens=(
- item.prompt_tokens
- if usage_valid
- else prompt_token_estimates.get(item.custom_id)
- ),
- completion_tokens=item.completion_tokens if usage_valid else None,
+ prompt_tokens=prompt_toks,
+ completion_tokens=comp_toks,
+ measurement_status=status,
usage_record_id=self._batch_usage_record_id(
job_id, item.custom_id, "result", 0
),
@@ -989,7 +1172,15 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict
for record in records
]
currencies = {row["currency_code"] for row in record_rows}
+ statuses = {row["measurement_status"] for row in record_rows}
+ aggregate_measurement_status = (
+ "unavailable" if "unavailable" in statuses
+ else "estimated" if "estimated" in statuses
+ else "measured"
+ )
price_known = all(row.get("price_known", True) for row in record_rows)
+ cost_known = price_known and aggregate_measurement_status != "unavailable"
+ tokens_known = aggregate_measurement_status != "unavailable"
recorded.append(
{
"custom_id": item.custom_id,
@@ -998,25 +1189,24 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict
"usage_record_ids": [record.usage_record_id for record in records],
"cost_amount": (
round(sum(row["cost_amount"] for row in record_rows), 6)
- if price_known and len(currencies) == 1
+ if cost_known and len(currencies) == 1
else None
),
"currency_code": (
next(iter(currencies)) if len(currencies) == 1 else "MIXED"
),
"price_known": price_known,
- "prompt_tokens": sum(row["prompt_tokens"] for row in record_rows),
- "completion_tokens": sum(
- row["completion_tokens"] for row in record_rows
+ "prompt_tokens": (
+ sum(row["prompt_tokens"] for row in record_rows)
+ if tokens_known
+ else None
),
- "measurement_status": (
- "estimated"
- if any(
- row["measurement_status"] == "estimated"
- for row in record_rows
- )
- else "measured"
+ "completion_tokens": (
+ sum(row["completion_tokens"] for row in record_rows)
+ if tokens_known
+ else None
),
+ "measurement_status": aggregate_measurement_status,
**({"currency_components": [
{
"currency_code": currency,
@@ -1026,7 +1216,7 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict
), 6),
}
for currency in sorted(currencies)
- ]} if len(currencies) > 1 and price_known else {}),
+ ]} if len(currencies) > 1 and cost_known else {}),
}
)
if prompt_token_estimates != job.prompt_token_estimates:
@@ -1113,6 +1303,7 @@ def submit_embeddings_batch(
metadata: Optional[Dict[str, Any]] = None,
zdr_only: bool = False,
agent_id: Optional[str] = None,
+ owner_id: Optional[str] = None,
) -> BatchJob:
"""Submit a bulk embeddings batch to the configured embeddings backend.
@@ -1125,7 +1316,9 @@ def submit_embeddings_batch(
raise TypeError("zdr_only must be a boolean")
if agent_id is not None and (not isinstance(agent_id, str) or not agent_id):
raise TypeError("agent_id must be a non-empty string when provided")
+ self._refresh_embedding_backend()
resolved_model, resolved_agent_id = self._resolve_embedding_target(model, zdr_only, agent_id)
+ backend = self._embedding_backend_for_route(resolved_model, resolved_agent_id)
shared_attribution = dict(attribution or {})
requests, part_counts, part_limits = self._build_embedding_requests(
inputs,
@@ -1134,13 +1327,21 @@ def submit_embeddings_batch(
zdr_only=zdr_only,
agent_id=resolved_agent_id,
)
- job = self.embedding_batch_backend.submit(requests, metadata=metadata)
- self._embedding_jobs[job.job_id] = job
+ reserve = getattr(backend, "reserve", None)
+ start = getattr(backend, "start", None)
+ if callable(reserve) and callable(start):
+ job = reserve(requests, metadata=metadata)
+ else:
+ job = backend.submit(requests, metadata=metadata)
self._embedding_models[job.job_id] = resolved_model
+ self._embedding_owners[job.job_id] = owner_id
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
self._embedding_part_limits[job.job_id] = part_limits
+ self._embedding_jobs[job.job_id] = job
+ if callable(reserve) and callable(start):
+ start(job)
return job
def _resolve_embedding_target(
@@ -1162,10 +1363,16 @@ def _resolve_embedding_target(
pool-membership + privacy-tag validation), so that behavior is
unchanged here.
"""
- unspecified_model = model in {
- "contextual-orchestrator", getattr(self.orchestrator, "AUTO_MODEL", "")
+ virtual_models = {
+ "contextual-orchestrator",
+ getattr(self.orchestrator, "AUTO_MODEL", ""),
}
- if agent_id is None and not zdr_only and not unspecified_model:
+ unspecified_model = model in virtual_models
+ if (
+ agent_id is None
+ and not zdr_only
+ and not unspecified_model
+ ):
return model, None
selection_model = None if unspecified_model else model
with self.orchestrator.request_policy(zdr_only):
@@ -1188,7 +1395,7 @@ def _build_embedding_requests(
agent_id: Optional[str],
) -> tuple[List[EmbeddingBatchRequest], List[int], Dict[str, int]]:
"""Map original embedding inputs into token-budgeted provider parts."""
- max_tokens, max_chars = self._embedding_request_limits()
+ max_tokens, max_chars, max_inputs = self._embedding_request_limits()
requests: List[EmbeddingBatchRequest] = []
part_counts: List[int] = []
for source_index, text in enumerate(inputs):
@@ -1215,15 +1422,15 @@ def _build_embedding_requests(
return requests, part_counts, {
"max_tokens_per_part": max_tokens,
"max_chars_per_part": max_chars,
+ "max_inputs_per_request": max_inputs,
}
- def _embedding_request_limits(self) -> tuple[int, int]:
+ def _embedding_request_limits(self) -> tuple[int, int, int]:
"""Return configured per-provider-call embedding ceilings.
Azure's current embeddings limit is surfaced by LiteLLM as a 300,000
token request cap. The default stays below that ceiling and also applies
- a character guard so heuristic token counters cannot accidentally send a
- very long no-whitespace string as one provider request.
+ a character guard independent of tokenizer availability.
"""
max_tokens = _positive_int(
self.config.get(
@@ -1241,7 +1448,15 @@ def _embedding_request_limits(self) -> tuple[int, int]:
),
_DEFAULT_EMBEDDING_MAX_CHARS_PER_PART,
)
- return max_tokens, max_chars
+ max_inputs = _positive_int(
+ self.config.get(
+ _EMBEDDING_CONFIG_CATEGORY,
+ "embedding_max_inputs_per_request",
+ _DEFAULT_EMBEDDING_MAX_INPUTS_PER_REQUEST,
+ ),
+ _DEFAULT_EMBEDDING_MAX_INPUTS_PER_REQUEST,
+ )
+ return max_tokens, max_chars, max_inputs
def _split_embedding_input(
self,
@@ -1254,9 +1469,21 @@ def _split_embedding_input(
"""Split one original embedding input into provider-safe map parts."""
if text == "":
return [("", 0)]
- parts = self._force_token_safe_chunks(
- text, model=model, max_tokens=max_tokens, max_chars=max_chars
- )
+ try:
+ native_pack = getattr(self.embedding_token_counter, "pack_text", None)
+ if callable(native_pack) and len(text) <= max_chars:
+ return native_pack(text, model, max_tokens)
+ parts = self._force_token_safe_chunks(
+ text, model=model, max_tokens=max_tokens, max_chars=max_chars
+ )
+ except TokenCountUnavailable:
+ if (
+ self._resolve_virtual_embedding_target
+ and len(text) <= max_chars
+ and len(text.encode("utf-8")) <= max_tokens
+ ):
+ return [(text, 0)]
+ raise
return parts or [("", 0)]
def _force_token_safe_chunks(
@@ -1338,16 +1565,57 @@ def _force_token_safe_chunks(
)
def _count_embedding_tokens(self, text: str, model: str) -> int:
- """Count tokens for embedding split decisions, tolerating adapters."""
- try:
- value = int(self.token_counter.count_text(text, model))
- except Exception:
- value = len(text.split())
+ """Count embedding tokens authoritatively or propagate unavailability."""
+ value = int(self.embedding_token_counter.count_text(text, model))
if text and value <= 0:
- return 1
+ raise RuntimeError("an authoritative tokenizer returned a non-positive count")
return max(0, value)
- def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
+ def embeddings_batch_document(
+ self, batch_id: str, *, owner_id: Optional[str] = None
+ ) -> Dict[str, Any]:
+ """Materialize one owner-bound document under the shared job lock."""
+ job = self._require_embedding_job(batch_id, owner_id=owner_id)
+ raw_lease = getattr(self.orchestrator.client, "timeout", 30) or 30
+ lease_seconds = max(1.0, float(raw_lease))
+ try:
+ with self.job_registry.lock(
+ "embedding_document", batch_id, lease_seconds=lease_seconds
+ ):
+ document = self._embeddings_batch_document_locked(
+ batch_id, owner_id=owner_id
+ )
+ except ClaimNotAcquired:
+ cached = self._embedding_documents.get(batch_id)
+ if cached is not None:
+ document = cached
+ else:
+ document = {
+ "batch_id": batch_id,
+ "status": "in_progress",
+ "backend": job.backend,
+ "model": self._embedding_models.get(
+ batch_id, "contextual-orchestrator"
+ ),
+ "embeddings": None,
+ }
+ result = dict(document)
+ result["job_retention_ms"] = self.job_registry.retention_seconds * 1000
+ if result.get("embeddings") is None and result.get("status") not in {
+ "failed", "cancelled", "rejected"
+ }:
+ backend = self._embedding_backend_for(job)
+ poll_after_ms = getattr(backend, "poll_after_ms", None)
+ if type(poll_after_ms) is not int or poll_after_ms < 1:
+ raise RuntimeError(
+ "queued embedding backend omitted its polling cadence"
+ )
+ result["poll_after_ms"] = poll_after_ms
+ return result
+
+ def _embeddings_batch_document_locked(
+ self, batch_id: str, *, owner_id: Optional[str] = None
+ ) -> Dict[str, Any]:
"""Return the naruon-shaped batch document for ``batch_id``.
Polls the backend; once complete, retrieves the vectors, records one
@@ -1360,21 +1628,37 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
if cached is not None:
return cached
- job = self._require_embedding_job(batch_id)
+ job = self._require_embedding_job(batch_id, owner_id=owner_id)
requests = self._embedding_requests.get(batch_id, [])
model_name = self._embedding_models.get(batch_id, "contextual-orchestrator")
- status = self.embedding_batch_backend.poll(job)
+ backend = self._embedding_backend_for(job)
+ status = backend.poll(job)
if not status.get("is_complete"):
- return {
+ document = {
"batch_id": batch_id,
"status": status.get("status") or job.status,
"backend": job.backend,
"model": model_name,
"embeddings": None,
}
+ return document
+ terminal_status = str(status.get("status") or "failed")
+ if terminal_status != "completed":
+ document = {
+ "batch_id": batch_id,
+ "status": terminal_status,
+ "backend": job.backend,
+ "model": model_name,
+ "embeddings": None,
+ }
+ for detail in ("failure", "cancellation"):
+ if status.get(detail) is not None:
+ document[detail] = status[detail]
+ self._embedding_documents[batch_id] = document
+ return document
try:
- items: List[EmbeddingBatchResultItem] = self.embedding_batch_backend.retrieve(job)
+ items: List[EmbeddingBatchResultItem] = backend.retrieve(job)
except BatchDownloadError as exc:
# Deliberately NOT cached: an explicit download failure must stay
# retryable. Caching this under "completed" (as a bare `return []`
@@ -1395,30 +1679,45 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
part_counts = self._embedding_part_counts.get(batch_id, [1] * input_count)
part_limits = self._embedding_part_limits.get(batch_id, {})
ordered = sorted(items, key=lambda item: item.index)
+ usage = status.get("usage") if isinstance(status, dict) else None
+ provider_total_tokens = (
+ usage.get("prompt_tokens") if isinstance(usage, dict) else None
+ )
parts_by_source: Dict[int, List[Dict[str, Any]]] = {index: [] for index in range(input_count)}
for item in ordered:
request = request_by_custom_id.get(item.custom_id)
source_index = request.source_index if request else item.index
prompt_tokens = int(item.prompt_tokens)
if prompt_tokens <= 0 and request is not None:
- prompt_tokens = request.token_count or int(
- self.token_counter.count_text(request.input_text, item.model)
- )
+ prompt_tokens = request.token_count
+ if request.input_text and prompt_tokens <= 0:
+ try:
+ prompt_tokens = int(
+ self.embedding_token_counter.count_text(
+ request.input_text, item.model
+ )
+ )
+ except TokenCountUnavailable:
+ prompt_tokens = None
parts_by_source.setdefault(source_index, []).append(
{
"part_index": request.part_index if request else 0,
"embedding": item.embedding,
- "prompt_tokens": max(0, prompt_tokens),
+ "prompt_tokens": (
+ max(0, prompt_tokens) if prompt_tokens is not None else None
+ ),
"model": item.model,
"attribution": dict(request.attribution) if request else {},
+ "agent_id": request.agent_id if request else None,
}
)
embeddings: List[Dict[str, Any]] = []
- token_counts: List[int] = []
+ token_counts: List[int | None] = []
total_cost_amount = 0.0
price_known = True
currency_code = "USD"
+ aggregate_usage_recorded = False
for source_index in range(input_count):
parts = sorted(parts_by_source.get(source_index, []), key=lambda item: item["part_index"])
if not parts:
@@ -1426,11 +1725,51 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
token_counts.append(0)
continue
attribution = dict(parts[0]["attribution"])
+ authoritative_parts = all(part["prompt_tokens"] is not None for part in parts)
+ if not authoritative_parts:
+ if len(parts) != 1 or type(provider_total_tokens) is not int:
+ raise TokenCountUnavailable(
+ "provider embedding usage cannot be assigned to split inputs"
+ )
+ if not aggregate_usage_recorded:
+ agent_id = parts[0]["agent_id"]
+ if not agent_id:
+ raise TokenCountUnavailable(
+ "provider embedding usage omitted execution identity"
+ )
+ provider, model_name = self._agent_provider_model(
+ self.orchestrator._agent(agent_id), str(parts[0]["model"])
+ )
+ record = self.ledger.record_usage(
+ provider=provider,
+ model=model_name,
+ prompt_tokens=provider_total_tokens,
+ completion_tokens=0,
+ request_channel="batch",
+ route_mode="embedding",
+ workflow_run_id=batch_id,
+ attribution=dict(parts[0]["attribution"]),
+ usage_record_id=f"usage_embedding_{batch_id}_aggregate",
+ )
+ total_cost_amount += float(record.cost_amount)
+ currency_code = record.currency_code
+ aggregate_usage_recorded = True
+ token_counts.append(None)
+ embeddings.append(
+ {"index": source_index, "embedding": parts[0]["embedding"]}
+ )
+ continue
prompt_tokens = sum(int(part["prompt_tokens"]) for part in parts)
model_name = str(parts[0]["model"])
- provider = str(
- attribution.get("provider") or attribution.get("upstream_api") or "unknown"
- )
+ agent_id = parts[0]["agent_id"]
+ if agent_id:
+ provider, model_name = self._agent_provider_model(
+ self.orchestrator._agent(agent_id), model_name
+ )
+ else:
+ provider = str(
+ attribution.get("provider") or attribution.get("upstream_api") or "unknown"
+ )
record = self.ledger.record_usage(
provider=provider,
model=model_name,
@@ -1440,6 +1779,7 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
route_mode="embedding",
workflow_run_id=batch_id,
attribution=attribution,
+ usage_record_id=f"usage_embedding_{batch_id}_{source_index}",
)
total_cost_amount += float(record.cost_amount)
price_known = price_known and record.price_known
@@ -1461,7 +1801,11 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]:
"model": model_name,
"embeddings": embeddings,
"token_counts": token_counts,
- "total_tokens": sum(token_counts),
+ "total_tokens": (
+ provider_total_tokens
+ if any(value is None for value in token_counts)
+ else sum(value for value in token_counts if value is not None)
+ ),
"part_count": len(requests),
"input_part_counts": part_counts,
"map_reduce": {
@@ -1485,13 +1829,14 @@ def complete_embeddings_batch(
metadata: Optional[Dict[str, Any]] = None,
zdr_only: bool = False,
agent_id: Optional[str] = None,
+ wait_timeout: Optional[float] = None,
+ owner_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Submit an embeddings batch and return its document (one round-trip).
- For the local/in-process backend the batch completes synchronously, so
- this returns the finished ``completed`` document with vectors and cost.
- For an async backend (pg-llm-batch) it returns a ``{batch_id, status}``
- envelope the caller then polls via :meth:`embeddings_batch_document`.
+ Local backends complete immediately. Callers that require a synchronous
+ provider result pass ``wait_timeout``; a timed-out queued job is
+ cancelled so the synchronous surface does not leave orphaned work.
"""
job = self.submit_embeddings_batch(
inputs,
@@ -1500,12 +1845,20 @@ def complete_embeddings_batch(
metadata=metadata,
zdr_only=zdr_only,
agent_id=agent_id,
+ owner_id=owner_id,
)
- return self.embeddings_batch_document(job.job_id)
-
- def _require_embedding_job(self, batch_id: str) -> BatchJob:
+ backend = self._embedding_backend_for(job)
+ if wait_timeout is not None and hasattr(backend, "wait"):
+ status = backend.wait(job, timeout=wait_timeout)
+ if not status.get("is_complete") and hasattr(backend, "cancel"):
+ backend.cancel(job, reason="synchronous request deadline elapsed")
+ return self.embeddings_batch_document(job.job_id, owner_id=owner_id)
+
+ def _require_embedding_job(
+ self, batch_id: str, *, owner_id: Optional[str] = None
+ ) -> BatchJob:
job = self._embedding_jobs.get(batch_id)
- if job is None:
+ if job is None or self._embedding_owners.get(batch_id) != owner_id:
raise KeyError(f"embeddings batch job {batch_id!r} not found")
return job
diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py
index 5e12554b2..9284098a8 100644
--- a/contextual_orchestrator/orchestrator.py
+++ b/contextual_orchestrator/orchestrator.py
@@ -48,6 +48,7 @@
from .openrouter_uptime import OpenRouterUptimeCollector
from .benchmark_priors import resolve_quality_prior
from .endpoint_race import EndpointAttempt, EndpointEquivalenceContract, race_first_valid
+from .reasoning_effort_profile import EffortProfileError
from .provider_errors import (
ProviderUpstreamError,
classify_provider_failure,
@@ -85,6 +86,75 @@
apply_request_profile,
snapshot_role_effort_catalog,
)
+from .token_counting import TokenCountUnavailable, build_token_counter
+
+
+_REQUEST_ENDPOINT_AGENT_IDS: ContextVar[frozenset[str] | None] = ContextVar(
+ "contextual_orchestrator_request_endpoint_agent_ids", default=None
+)
+_REQUEST_ENDPOINT_IDENTITY: ContextVar[str | None] = ContextVar(
+ "contextual_orchestrator_request_endpoint_identity", default=None
+)
+_INVALID_REQUESTED_MODEL = object()
+
+
+class EndpointUnavailableError(ValueError):
+ """The requested configured endpoint cannot serve this request."""
+
+
+def normalize_endpoint_selector(value: str) -> str:
+ """Normalize an endpoint selector without ever using it as transport input."""
+ try:
+ parsed = urlparse(value)
+ scheme = parsed.scheme.casefold()
+ hostname = parsed.hostname
+ port = parsed.port
+ except ValueError as exc:
+ raise EndpointUnavailableError("endpoint_unavailable") from exc
+ if (
+ scheme not in {"http", "https"}
+ or not hostname
+ or parsed.username
+ or parsed.password
+ or parsed.params
+ or parsed.query
+ or parsed.fragment
+ ):
+ raise EndpointUnavailableError("endpoint_unavailable")
+ host = hostname.casefold()
+ if ":" in host:
+ host = f"[{host}]"
+ if port is not None and port != (443 if scheme == "https" else 80):
+ host = f"{host}:{port}"
+ path = parsed.path.rstrip("/").removesuffix("/v1")
+ return urlunsplit((scheme, host, path, "", ""))
+
+
+def _configured_endpoint_matches(value: str, normalized: str) -> bool:
+ """Return exact normalized equality for one already-configured transport."""
+ try:
+ return normalize_endpoint_selector(value) == normalized
+ except EndpointUnavailableError:
+ return False
+
+
+def _agent_matches_request_endpoint(agent: ModelAgent) -> bool:
+ """Revalidate both agent id and configured endpoint for the active request."""
+ endpoint_ids = _REQUEST_ENDPOINT_AGENT_IDS.get()
+ endpoint_identity = _REQUEST_ENDPOINT_IDENTITY.get()
+ if endpoint_ids is None or endpoint_identity is None:
+ return endpoint_ids is None and endpoint_identity is None
+ return agent.id in endpoint_ids and _configured_endpoint_matches(
+ agent.base_url, endpoint_identity
+ )
+
+
+def _request_endpoint_partition() -> str:
+ """Return a non-reversible cache partition for the configured endpoint."""
+ identity = _REQUEST_ENDPOINT_IDENTITY.get()
+ if identity is None:
+ return "endpoint:auto"
+ return "endpoint:" + hashlib.sha256(identity.encode("utf-8")).hexdigest()
# content is usually str; multimodal vision messages use OpenAI content-parts lists.
@@ -173,17 +243,22 @@ def __init__(
def _structured_output_error(
content: str, response_format: object
) -> str | None:
- """Return a bounded contract error for strict JSON Schema output."""
- if not isinstance(response_format, Mapping) or response_format.get("type") != "json_schema":
+ """Return a bounded contract error for JSON object or JSON Schema output."""
+ if not isinstance(response_format, Mapping):
+ return None
+ response_type = response_format.get("type")
+ if response_type not in {"json_object", "json_schema"}:
return None
specification = response_format.get("json_schema")
schema = specification.get("schema") if isinstance(specification, Mapping) else None
- if not isinstance(schema, Mapping):
+ if response_type == "json_schema" and not isinstance(schema, Mapping):
return "schema_missing"
try:
instance = json.loads(content)
except (TypeError, json.JSONDecodeError):
return "invalid_json"
+ if response_type == "json_object":
+ return None if isinstance(instance, Mapping) else "invalid_json_object"
validator_type = validator_for(schema)
try:
validator_type.check_schema(schema)
@@ -213,29 +288,30 @@ def _bind_provider_file_ids(
return result
-def estimate_tokens(text: str) -> int:
- """Rough token estimate (~4 chars/token). ponytail: heuristic, not a real tokenizer.
-
- Honest floor for spend analytics on mock/runtime text; replace with provider-reported
- usage when real workers return it.
- """
- return (len(text) + 3) // 4 if text else 0
-
-
-def _step_output_tokens(step: Mapping[str, Any]) -> tuple[int, bool]:
- """Return provider-reported output tokens or the existing text estimate."""
+def _step_output_tokens(
+ step: Mapping[str, Any], token_counter: Any, model: str
+) -> tuple[int | None, str]:
+ """Return reported or exact raw-output tokens with their evidence source."""
usage = step.get("usage")
if isinstance(usage, dict):
for key in ("completion_tokens", "output_tokens"):
reported = usage.get(key)
if type(reported) is int and reported >= 0:
- return reported, True
- return estimate_tokens(step.get("output", "")), False
+ return reported, "reported"
+ output = step.get("output")
+ if not isinstance(output, str):
+ return None, "unavailable"
+ try:
+ return token_counter.count_text(output, model), "tokenizer"
+ except TokenCountUnavailable:
+ return None, "unavailable"
-def _step_output_token_count(step: Mapping[str, Any]) -> int:
- """Return the output count used by in-flight structured budget checks."""
- return _step_output_tokens(step)[0]
+def _step_output_token_count(
+ step: Mapping[str, Any], token_counter: Any, model: str
+) -> int | None:
+ """Return the authoritative output count for an in-flight budget check."""
+ return _step_output_tokens(step, token_counter, model)[0]
def _cost_usd_decimal(output_tokens: int, price_per_million: float) -> Decimal:
@@ -292,6 +368,7 @@ class _FastMLSIJudgeAdapter:
served_output: str | None = None
mode: str = "auto"
allowed_agent_ids: set[str] | None = None
+ excluded_agent_ids: set[str] | None = None
@property
def contextual_orchestrator_contract(self) -> str:
@@ -314,6 +391,7 @@ def complete(self, messages: list[ChatMessage], mode: str | None = None) -> dict
role="judge",
allowed_agent_ids=self.allowed_agent_ids,
eligibility_role="verifier",
+ excluded_agent_ids=self.excluded_agent_ids,
)
return self._completion_payload(
output, served_id, served_model, usage, self.mode if mode is None else mode
@@ -658,7 +736,7 @@ def _eligible_role_effort_candidates(
def _is_general_chat_agent(agent: ModelAgent) -> bool:
"""Apply persisted provider capability tags before model-name fallback."""
- return is_general_chat_candidate(
+ return "structured:blocked" not in agent.tags and is_general_chat_candidate(
agent.model,
capabilities=(
tag.split(":", 1)[1]
@@ -1678,6 +1756,13 @@ def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]:
vector so unit tests can exercise cosine ordering without network
access; this fixture is never used against production traffic.
"""
+ vectors, _prompt_tokens = self.embed_with_usage(agent, texts)
+ return vectors
+
+ def embed_with_usage(
+ self, agent: ModelAgent, texts: list[str]
+ ) -> tuple[list[list[float]], int | None]:
+ """Return embeddings plus an authoritative provider input-token count when supplied."""
if not isinstance(texts, list) or not texts:
raise ValueError("texts must be a non-empty list of strings")
for item in texts:
@@ -1690,7 +1775,7 @@ def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]:
raw = [byte for byte in digest[: self.MOCK_EMBEDDING_DIMENSION]]
centered = [(value / 255.0) * 2.0 - 1.0 for value in raw]
vectors.append(centered)
- return vectors
+ return vectors, None
destination = self._validate_provider(agent) # pragma: no cover
payload = {"model": agent.model, "input": texts} # pragma: no cover
response = self._send_raw(agent, "embeddings", payload, destination) # pragma: no cover
@@ -1709,7 +1794,13 @@ def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]:
f"provider {agent.id} returned a non-numeric embedding vector"
)
vectors.append([float(value) for value in vector]) # pragma: no cover
- return vectors # pragma: no cover
+ usage = response.get("usage") if isinstance(response, dict) else None # pragma: no cover
+ prompt_tokens = None # pragma: no cover
+ if isinstance(usage, dict): # pragma: no cover
+ raw_tokens = usage.get("prompt_tokens", usage.get("input_tokens"))
+ if type(raw_tokens) is int and raw_tokens >= 0:
+ prompt_tokens = raw_tokens
+ return vectors, prompt_tokens # pragma: no cover
def chat(
self,
@@ -2259,6 +2350,18 @@ def proxy_send_once(
"""Send one passthrough attempt so cross-provider failover cannot amplify load."""
return self._proxy_send(agent, endpoint, payload, allow_transient_retries=False)
+ def probe_structured_chat(
+ self, agent: ModelAgent, payload: dict[str, Any]
+ ) -> dict[str, Any]:
+ """Run one bounded chat-readiness probe under separate telemetry."""
+ return self._proxy_send(
+ agent,
+ "chat/completions",
+ payload,
+ allow_transient_retries=False,
+ operation_kind="capability_probe",
+ )
+
def _proxy_send(
self,
agent: ModelAgent,
@@ -2266,6 +2369,7 @@ def _proxy_send(
payload: dict[str, Any],
*,
allow_transient_retries: bool,
+ operation_kind: str = "request",
) -> dict[str, Any]:
"""Apply the shared passthrough contract with a selectable retry policy."""
normalized_endpoint = endpoint.strip("/")
@@ -2287,16 +2391,25 @@ def _proxy_send(
"completions": "text_completion",
"responses": "generate_content",
}.get(normalized_endpoint, "generate_content")
+ trace_name = f"{operation_name} {agent.model}"
+ trace_attributes = {
+ "gen_ai.operation.name": operation_name,
+ "gen_ai.provider.name": agent.provider_name or parsed_provider.hostname or agent.id,
+ "gen_ai.request.model": agent.model,
+ "contextual_orchestrator.agent_id": agent.id,
+ "contextual_orchestrator.model_group": agent.group_name or agent.model,
+ "contextual_orchestrator.fallback_outcome": (
+ "not_attempted" if operation_kind == "capability_probe" else "not_observed"
+ ),
+ "server.address": parsed_provider.hostname or "",
+ "server.port": parsed_provider.port or (443 if parsed_provider.scheme == "https" else 80),
+ }
+ if operation_kind != "request":
+ trace_name = f"{operation_kind}.{trace_name}"
+ trace_attributes["contextual_orchestrator.operation_kind"] = operation_kind
with traced(
- f"{operation_name} {agent.model}",
- {
- "gen_ai.operation.name": operation_name,
- "gen_ai.provider.name": agent.provider_name or parsed_provider.hostname or agent.id,
- "gen_ai.request.model": agent.model,
- "contextual_orchestrator.agent_id": agent.id,
- "server.address": parsed_provider.hostname or "",
- "server.port": parsed_provider.port or (443 if parsed_provider.scheme == "https" else 80),
- },
+ trace_name,
+ trace_attributes,
):
if (
normalized_endpoint == "chat/completions"
@@ -2538,7 +2651,8 @@ def _mock_raw(
mock_content = (
"{}"
if isinstance(response_format, dict)
- and str(response_format.get("type", "")).strip().lower() == "json_schema"
+ and str(response_format.get("type", "")).strip().lower()
+ in {"json_object", "json_schema"}
else f"[{agent.id}] chat-mock"
)
echoed = {
@@ -2575,6 +2689,7 @@ def _mock_raw(
return {
"id": f"chatcmpl_mock_{agent.id}",
"object": "chat.completion",
+ "created": int(time.time()),
"model": agent.model,
"choices": [
{
@@ -3676,6 +3791,7 @@ def __init__(
role_effort_catalog: dict[str, ReasoningEffortProfile] | None = None,
pii_key_name: str = DEFAULT_PII_KEY_NAME,
allow_empty_agents: bool = False,
+ token_counter: Any = None,
) -> None:
# Optional durable model-group management: stored operator changes overlay the
# seed agents file at startup (stored rows win by id; stored-new rows append).
@@ -3720,6 +3836,7 @@ def __init__(
# production always uses the exact-schema implementation below.
self._triage_fn = self._triage_workflow_required
self.client = client or ModelClient()
+ self.token_counter = token_counter or build_token_counter()
# The cost coordinator installs this optional sink. Direct orchestrator
# callers still retain audit evidence without inventing price or usage.
self._race_usage_sink: Callable[[str, Any], None] | None = None
@@ -3763,6 +3880,7 @@ def __init__(
self._budget_spent_output_tokens = 0
self._budget_spent_cost_usd = Decimal(0)
self._budget_model_output_tokens: dict[str, int] = {}
+ self._budget_unavailable_run_ids: set[str] = set()
self._evaluation_runs: dict[str, dict[str, Any]] = {}
self._analytics_events: deque[dict[str, Any]] = deque(maxlen=256)
self._audit_events: deque[dict[str, Any]] = deque(maxlen=256)
@@ -3948,6 +4066,7 @@ def _reload_state(self) -> None:
"stream_options",
"_required_agent_id",
"_file_replicas",
+ "session_id",
}
)
@@ -4265,6 +4384,12 @@ def _orchestrated_provider_completion(
required_tags = ("vision",) if self._source_image_parts(messages) else ()
response_format_requested = bool(chat_body.get("response_format"))
requested_model = body.get("model")
+ virtual_model = requested_model in {
+ None,
+ "contextual-orchestrator",
+ self.AUTO_MODEL,
+ self.FREE_MODEL,
+ }
free_only = requested_model == self.FREE_MODEL
required_agent_id = body.get("_required_agent_id")
file_replicas = body.get("_file_replicas")
@@ -4355,27 +4480,23 @@ def _orchestrated_provider_completion(
raise RuntimeError(f"requested model {requested_model!r} is disabled")
self._raise_if_spend_budget_exceeded()
+ request_exclusions: set[str] = set()
workflow = self.conduct(
messages,
- model_name=self.FREE_MODEL if free_only else self.GATEWAY_DEFAULT_MODEL,
- )
- in_flight_tokens = sum(_step_output_token_count(step) for step in workflow["trace"])
- model_by_agent = {agent.id: agent.model for agent in self.agents}
- in_flight_cost = sum(
- _step_output_token_count(step)
- / 1_000_000
- * self.price_per_million[model]
- for step in workflow["trace"]
- if (
- model := model_by_agent.get(
- step.get("served_agent_id") or step.get("agent_id")
- )
- )
- in self.price_per_million
+ model_name=(
+ self.FREE_MODEL
+ if free_only
+ else self.GATEWAY_DEFAULT_MODEL
+ if virtual_model
+ else str(requested_model)
+ ),
+ _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=round(in_flight_cost, 6),
+ additional_cost_usd=in_flight_cost,
)
evidence = "\n\n".join(
@@ -4440,6 +4561,7 @@ def _orchestrated_provider_completion(
"stream": False,
}
)
+ active_profile = effort_profile or self._role_effort_profile("synthesizer")
virtual_model = requested_model in {
None,
"contextual-orchestrator",
@@ -4482,29 +4604,80 @@ def _orchestrated_provider_completion(
if virtual_model
else [final_agent]
)
+ synthesis_failure_recorded = False
+ synthesis_candidates = [
+ candidate
+ for candidate in synthesis_candidates
+ if candidate.id not in request_exclusions
+ ]
+ if final_agent.id in request_exclusions:
+ same_endpoint_candidates = [
+ candidate
+ for candidate in synthesis_candidates
+ if candidate.base_url.rstrip("/").casefold()
+ == final_agent.base_url.rstrip("/").casefold()
+ ]
+ if not same_endpoint_candidates:
+ raise ProviderUpstreamError(
+ agent_id=final_agent.id,
+ model=final_agent.model,
+ error_code="model_not_found",
+ message="every eligible model on the selected endpoint is unavailable",
+ client_status=404,
+ provider_status=404,
+ retryable=False,
+ transport="structured_synthesis",
+ )
+ final_agent = same_endpoint_candidates[0]
+ synthesis_candidates = same_endpoint_candidates
+
+ def provider_output(agent: ModelAgent, response: Mapping[str, Any]) -> str:
+ """Extract non-empty structured output from the attempted provider."""
+ if not response_request:
+ return ModelClient._response_content(agent, dict(response))
+ output = response.get("output_text")
+ if isinstance(output, str) and output:
+ return output
+ combined = "".join(
+ _responses_text(item.get("content"))
+ for item in response.get("output", [])
+ if isinstance(item, dict) and item.get("type") == "message"
+ )
+ if combined:
+ return combined
+ raise ProviderResponseError(
+ f"provider {agent.id} returned no structured response content"
+ )
def send_synthesis(
payload: dict[str, Any],
) -> tuple[dict[str, Any], ModelAgent]:
- """Send once per eligible provider, advancing only after a proven 413 rejection."""
- nonlocal final_agent
+ """Retry 413 broadly and stale virtual models only within one endpoint."""
+ nonlocal final_agent, synthesis_failure_recorded
seen_providers: set[str] = set()
preferred = final_agent
+ preferred_endpoint = preferred.base_url.rstrip("/").casefold()
+ last_model_not_found: ProviderUpstreamError | None = None
+ saw_request_too_large = False
ordered_candidates = [
- preferred,
+ *([preferred] if preferred.id not in request_exclusions else []),
*(
candidate
for candidate in synthesis_candidates
if candidate.id != preferred.id
+ and candidate.id not in request_exclusions
),
]
for candidate in ordered_candidates:
+ candidate_endpoint = candidate.base_url.rstrip("/").casefold()
+ if last_model_not_found is not None and candidate_endpoint != preferred_endpoint:
+ continue
provider_key = (
f"provider:{candidate.provider_name.casefold()}"
if candidate.provider_name.strip()
else f"endpoint:{candidate.base_url.rstrip('/').casefold()}"
)
- if provider_key in seen_providers:
+ if provider_key in seen_providers and candidate_endpoint != preferred_endpoint:
continue
seen_providers.add(provider_key)
# Keep the outer failure accounting attached to the provider
@@ -4528,70 +4701,92 @@ def send_synthesis(
send_once = getattr(self.client, "proxy_send_once", None)
if callable(send_once):
send = send_once
- return send(candidate, endpoint, candidate_payload), candidate
+ response = send(candidate, endpoint, candidate_payload)
+ provider_output(candidate, response)
+ return response, candidate
except Exception as exc: # noqa: BLE001 - provider trust boundary
request_too_large = _is_request_too_large_error(exc)
+ saw_request_too_large = saw_request_too_large or request_too_large
if request_too_large and not virtual_model:
raise ProviderRequestTooLargeError(
"request body exceeds provider limit"
) from exc
if not request_too_large:
- raise classify_provider_failure(
+ if (
+ virtual_model
+ and isinstance(exc, ProviderResponseError)
+ and candidate_endpoint == preferred_endpoint
+ ):
+ request_exclusions.add(candidate.id)
+ self._record_failure(candidate.id)
+ synthesis_failure_recorded = True
+ continue
+ classified = classify_provider_failure(
exc,
agent_id=candidate.id,
model=candidate.model,
- transport=endpoint,
- ) from None
+ transport="structured_synthesis",
+ )
+ if (
+ virtual_model
+ and classified.error_code == "model_not_found"
+ and candidate_endpoint == preferred_endpoint
+ ):
+ last_model_not_found = classified
+ request_exclusions.add(candidate.id)
+ self._record_failure(candidate.id)
+ synthesis_failure_recorded = True
+ continue
+ raise classified from None
+ if last_model_not_found is not None and not saw_request_too_large:
+ raise last_model_not_found
raise ProviderRequestTooLargeError(
"request body exceeds every eligible provider limit"
)
+ response_format = chat_body.get("response_format")
synthesis_started = time.perf_counter()
- try:
- raw, final_agent = send_synthesis(upstream)
- except ProviderUpstreamError as exc:
- if not _is_request_too_large_error(exc):
- self._record_failure(final_agent.id)
- if final_agent.group_name and not _is_request_too_large_error(exc):
- self._group_router.observe_failure(final_agent.id)
- raise
- def provider_output(response: Mapping[str, Any]) -> str:
- if not response_request:
- try:
- return ModelClient._response_content(final_agent, response)
- except RuntimeError:
- return ""
- output = response.get("output_text")
- if isinstance(output, str):
- return output
- return "".join(
- _responses_text(item.get("content"))
- for item in response.get("output", [])
- if isinstance(item, dict) and item.get("type") == "message"
+ while True:
+ synthesis_failure_recorded = False
+ try:
+ raw, final_agent = send_synthesis(upstream)
+ except Exception as exc:
+ if (
+ not _is_request_too_large_error(exc)
+ and not isinstance(exc, EffortProfileError)
+ and not synthesis_failure_recorded
+ ):
+ self._record_failure(final_agent.id)
+ 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)
+ raise
+ synthesis_output = provider_output(final_agent, raw)
+ synthesis_step = {
+ "id": len(workflow["trace"]),
+ "role": "synthesizer",
+ "agent_id": final_agent.id,
+ "subtask": "Provider-facing structured synthesis",
+ "access": [step["id"] for step in workflow["trace"]],
+ "latency_ms": round((time.perf_counter() - synthesis_started) * 1000, 2),
+ "output": synthesis_output,
+ }
+ if isinstance(raw.get("usage"), dict):
+ synthesis_step["usage"] = _canonical_provider_usage(
+ raw["usage"], responses=response_request
+ )
+ repair_step: dict[str, Any] | None = None
+ contract_error = _structured_output_error(synthesis_output, response_format)
+ if contract_error == "schema_missing":
+ raise ProviderResponseError(
+ "response_format.json_schema is missing a schema"
)
+ if contract_error is None:
+ break
- synthesis_output = provider_output(raw)
- synthesis_step: dict[str, Any] = {
- "id": len(workflow["trace"]),
- "role": "synthesizer",
- "agent_id": final_agent.id,
- "subtask": "Provider-facing structured synthesis",
- "access": [step["id"] for step in workflow["trace"]],
- "latency_ms": round((time.perf_counter() - synthesis_started) * 1000, 2),
- "output": synthesis_output,
- }
- if isinstance(raw.get("usage"), dict):
- synthesis_step["usage"] = _canonical_provider_usage(
- raw["usage"], responses=response_request
- )
- repair_step: dict[str, Any] | None = None
- response_format = chat_body.get("response_format")
- contract_error = _structured_output_error(synthesis_output, response_format)
- if contract_error == "schema_missing":
- raise ProviderResponseError(
- "response_format.json_schema is missing a schema"
- )
- if contract_error is not None:
in_flight_tokens, in_flight_cost = self._trace_budget_spend(
[*workflow["trace"], synthesis_step]
)
@@ -4629,29 +4824,51 @@ def provider_output(response: Mapping[str, Any]) -> str:
if final_agent.group_name and not _is_request_too_large_error(exc):
self._group_router.observe_failure(final_agent.id)
raise
- repaired_output = provider_output(repaired)
- if _structured_output_error(repaired_output, response_format) is not None:
- self._record_failure(final_agent.id)
- if final_agent.group_name:
- self._group_router.observe_failure(final_agent.id)
+ repaired_output = provider_output(final_agent, repaired)
+ repair_error = _structured_output_error(repaired_output, response_format)
+ 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
+ )
+ raw = repaired
+ synthesis_output = repaired_output
+ break
+
+ failed_agent = final_agent
+ self._record_failure(failed_agent.id)
+ if failed_agent.group_name:
+ self._group_router.observe_failure(failed_agent.id)
+ if not virtual_model:
raise ProviderResponseError(
"structured synthesis and repair violated response_format"
)
- 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
+ request_exclusions.add(failed_agent.id)
+ failed_endpoint = failed_agent.base_url.rstrip("/").casefold()
+ next_agent = next(
+ (
+ candidate
+ for candidate in synthesis_candidates
+ if candidate.id not in request_exclusions
+ and candidate.base_url.rstrip("/").casefold() == failed_endpoint
+ ),
+ None,
+ )
+ if next_agent is None:
+ raise ProviderResponseError(
+ "every eligible model on the selected endpoint violated response_format"
)
- raw = repaired
- synthesis_output = repaired_output
+ final_agent = next_agent
+ synthesis_started = time.perf_counter()
self._record_success(final_agent.id)
if final_agent.group_name:
self._group_router.observe_success(
@@ -4715,6 +4932,80 @@ def provider_output(response: Mapping[str, Any]) -> str:
}
return raw
+ @contextmanager
+ def routing_endpoint_scope(
+ self,
+ endpoint: str | None,
+ requested_model: Any,
+ *,
+ model_was_provided: bool = True,
+ ):
+ """Constrain this request to agents whose configured endpoint matches exactly."""
+ if endpoint is None:
+ yield
+ return
+ if not isinstance(endpoint, str) or not endpoint.strip():
+ raise EndpointUnavailableError("endpoint_unavailable")
+ normalized = normalize_endpoint_selector(endpoint.strip())
+ matching = frozenset(
+ agent.id
+ for agent in self.agents
+ if _configured_endpoint_matches(agent.base_url, normalized)
+ )
+ if not matching:
+ raise EndpointUnavailableError("endpoint_unavailable")
+ ids_token = _REQUEST_ENDPOINT_AGENT_IDS.set(matching)
+ identity_token = _REQUEST_ENDPOINT_IDENTITY.set(normalized)
+ try:
+ if not self._request_endpoint_supports_model(
+ self._normalize_endpoint_requested_model(
+ requested_model, model_was_provided=model_was_provided
+ )
+ ):
+ raise EndpointUnavailableError("endpoint_unavailable")
+ yield
+ finally:
+ _REQUEST_ENDPOINT_IDENTITY.reset(identity_token)
+ _REQUEST_ENDPOINT_AGENT_IDS.reset(ids_token)
+
+ def _normalize_endpoint_requested_model(
+ self, requested_model: Any, *, model_was_provided: bool
+ ) -> Any:
+ """Reuse request-model normalization without preempting later HTTP errors."""
+ if requested_model is None:
+ return _INVALID_REQUESTED_MODEL if model_was_provided else None
+ if type(requested_model) is not str:
+ return _INVALID_REQUESTED_MODEL
+ normalized = requested_model.strip()
+ if not normalized or len(normalized) > 256:
+ return _INVALID_REQUESTED_MODEL
+ return normalized
+
+ def _request_endpoint_supports_model(self, requested_model: Any) -> bool:
+ """Check endpoint-local eligibility without ranking or provider I/O."""
+ if requested_model is _INVALID_REQUESTED_MODEL:
+ return True
+ if requested_model in {
+ None,
+ self.GATEWAY_DEFAULT_MODEL,
+ self.AUTO_MODEL,
+ self.FREE_MODEL,
+ }:
+ free_only = requested_model == self.FREE_MODEL
+ return any(
+ not agent.disabled
+ and _agent_matches_request_endpoint(agent)
+ and self._zdr_agent_allowed(agent)
+ and _is_general_chat_agent(agent)
+ and (not free_only or self._is_general_free_agent(agent))
+ for agent in self.agents
+ )
+ try:
+ self._requested_agent(requested_model)
+ except ValueError:
+ return False
+ return True
+
def _requested_agent(self, requested_model: Any) -> ModelAgent | None:
"""Resolve an explicit model without silently serving a different model."""
if requested_model is None or requested_model in {
@@ -4728,6 +5019,7 @@ def _requested_agent(self, requested_model: Any) -> ModelAgent | None:
for candidate in self.candidates
if (
candidate.model == requested_model
+ and _agent_matches_request_endpoint(candidate)
and self._zdr_agent_allowed(candidate)
and (not _REQUEST_ZDR_ONLY.get() or not candidate.disabled)
)
@@ -4949,6 +5241,12 @@ def _cache_key(
"max_output_tokens": getattr(self.client, "max_output_tokens", None),
}
parameters = {**parameters, "zdr_only": _REQUEST_ZDR_ONLY.get()}
+ endpoint_partition = _request_endpoint_partition()
+ cache_partition = (
+ endpoint_partition
+ if cache_partition is None
+ else f"{cache_partition}|{endpoint_partition}"
+ )
return build_response_cache_key(
messages,
mode,
@@ -4971,6 +5269,8 @@ def run(
"""Execute completion and persist a workflow run with trace and policy evidence."""
if self.budget_max_output_tokens is not None or self.budget_max_cost_usd is not None:
budget = self.budget_status()
+ if budget.get("enforcement_status") == "blocked_unavailable":
+ raise BudgetExceededError("spend budget measurement unavailable", detail=budget)
if budget["exceeded"]:
raise BudgetExceededError("spend budget exceeded", detail=budget)
result = self.complete(
@@ -5036,8 +5336,8 @@ def run(
def _raise_if_spend_budget_exceeded(
self,
*,
- additional_output_tokens: int = 0,
- additional_cost_usd: float = 0.0,
+ additional_output_tokens: int | None = 0,
+ additional_cost_usd: float | None = 0.0,
) -> None:
"""Fail before another provider call would cross an operator budget."""
with self._budget_spend_lock:
@@ -5046,14 +5346,29 @@ def _raise_if_spend_budget_exceeded(
budget = self._budget_block(
spent_output_tokens,
float(spent_cost_decimal) if self.price_per_million else None,
+ measurement_available=not self._budget_unavailable_run_ids,
+ )
+ measurement_unavailable = (
+ (budget["max_output_tokens"] is not None and additional_output_tokens is None)
+ or (budget["max_cost_usd"] is not None and additional_cost_usd is None)
+ )
+ spent_tokens = (
+ spent_output_tokens + additional_output_tokens
+ if additional_output_tokens is not None
+ else None
)
- spent_tokens = spent_output_tokens + additional_output_tokens
spent_cost = budget["spent_cost_usd"]
effective_cost = (
- spent_cost + additional_cost_usd if spent_cost is not None else None
+ spent_cost + additional_cost_usd
+ if spent_cost is not None and additional_cost_usd is not None
+ else None
)
+ if measurement_unavailable or budget["enforcement_status"] == "blocked_unavailable":
+ detail = {**budget, "measurement_status": "unavailable"}
+ raise BudgetExceededError("spend budget measurement unavailable", detail=detail)
if budget["exceeded"] or (
budget["max_output_tokens"] is not None
+ and spent_tokens is not None
and spent_tokens >= budget["max_output_tokens"]
) or (
budget["max_cost_usd"] is not None
@@ -5062,19 +5377,26 @@ def _raise_if_spend_budget_exceeded(
):
raise BudgetExceededError("spend budget exceeded", detail=budget)
- def _trace_budget_spend(self, trace: list[dict[str, Any]]) -> tuple[int, float]:
+ def _trace_budget_spend(
+ self, trace: list[dict[str, Any]]
+ ) -> tuple[int | None, float | None]:
"""Return completed provider-call spend for a workflow budget checkpoint."""
model_by_agent = {agent.id: agent.model for agent in self.agents}
- output_tokens = sum(_step_output_token_count(step) for step in trace)
- output_cost = sum(
- _step_output_token_count(step) / 1_000_000 * self.price_per_million[model]
- for step in trace
- if (
- model := model_by_agent.get(
- step.get("served_agent_id") or step.get("agent_id")
- )
+ counts: list[tuple[int, str]] = []
+ for step in trace:
+ model = step.get("model_name") or model_by_agent.get(
+ step.get("served_agent_id") or step.get("agent_id"), "unknown"
)
- in self.price_per_million
+ count = _step_output_token_count(step, self.token_counter, model)
+ if count is None:
+ return None, None
+ counts.append((count, 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
+ output_cost = sum(
+ count / 1_000_000 * self.price_per_million[model]
+ for count, model in counts
)
return output_tokens, round(output_cost, 6)
@@ -5091,6 +5413,8 @@ def batch_route(self, prompts: list[str]) -> list[dict[str, Any]]:
"""
if self.budget_max_output_tokens is not None or self.budget_max_cost_usd is not None:
budget = self.budget_status()
+ if budget.get("enforcement_status") == "blocked_unavailable":
+ raise BudgetExceededError("spend budget measurement unavailable", detail=budget)
if budget["exceeded"]:
raise BudgetExceededError("spend budget exceeded", detail=budget)
selected = [(prompt, self._select_agent(prompt, "worker")) for prompt in prompts]
@@ -5847,6 +6171,20 @@ def remove_agent(self, agent_pool_id: str, worker_agent_id: str) -> dict[str, An
)
return {"removed": worker_agent_id}
+ def _retire_runtime_agent(self, worker_agent_id: str) -> None:
+ """Remove a bootstrap row for this process without persisting a tombstone."""
+ self._agent_in_pool("default", worker_agent_id)
+ self.candidates = [
+ agent for agent in self.candidates if agent.id != worker_agent_id
+ ]
+ self.agents = [agent for agent in self.candidates if not agent.disabled]
+ self._rebuild_budget_meter()
+ self._routers_forget_members({agent.id for agent in self.candidates})
+ self._append_audit_event(
+ "runtime_agent_retired",
+ {"agent_pool_id": "default", "worker_agent_id": worker_agent_id},
+ )
+
def route_once(
self,
messages: list[ChatMessage],
@@ -6047,6 +6385,8 @@ def conduct(
model_name: str = GATEWAY_DEFAULT_MODEL,
progress: Any = None,
workflow_run_id: str | None = None,
+ _excluded_agent_ids: set[str] | None = None,
+ _allowed_agent_ids: set[str] | None = None,
) -> dict[str, Any]:
"""Run a workflow, optionally persisting it under a supplied run id."""
self._raise_if_spend_budget_exceeded()
@@ -6054,9 +6394,10 @@ def conduct(
source_images = self._source_image_parts(messages)
required_tags = ("vision",) if source_images else ()
caller_instructions = "\n\n".join(
- message["content"]
+ instruction
for message in messages
- if message.get("role") == "system" and isinstance(message.get("content"), str)
+ if message.get("role") == "system"
+ if (instruction := _coerce_message_content_text(message.get("content")))
)
plan_source = "template"
if model_name not in {self.GATEWAY_DEFAULT_MODEL, self.AUTO_MODEL}:
@@ -6081,7 +6422,9 @@ def conduct(
}
requested_agent = self._requested_agent(model_name)
judge_agent_ids = (
- {
+ _allowed_agent_ids
+ if _allowed_agent_ids is not None
+ else {
candidate.id
for candidate in self.agents
if candidate.group_name == requested_agent.group_name
@@ -6118,16 +6461,7 @@ def conduct(
if progress is not None:
progress(step.role, "started")
prior = "\n\n".join(f"Step {i}: {outputs[i]}" for i in step.access)
- instruction = (
- f"Original task:\n{task}\n\nAccessed prior work:\n{prior}\n\n"
- f"Subtask:\n{step.subtask}"
- )
- user_content: str | list[dict[str, Any]] = instruction
- if source_images:
- user_content = [
- {"type": "text", "text": instruction},
- *copy.deepcopy(source_images),
- ]
+ instruction = f"Accessed prior work:\n{prior}\n\nSubtask:\n{step.subtask}"
step_messages = [
{
"role": "system",
@@ -6141,7 +6475,7 @@ def conduct(
*copy.deepcopy(messages),
{
"role": "user",
- "content": user_content,
+ "content": instruction,
},
]
start = time.perf_counter()
@@ -6150,7 +6484,10 @@ def conduct(
step_messages,
text=task,
role=step.role,
- allowed_agent_ids=free_ids if model_name == self.FREE_MODEL else None,
+ allowed_agent_ids=(
+ free_ids if model_name == self.FREE_MODEL else _allowed_agent_ids
+ ),
+ excluded_agent_ids=_excluded_agent_ids,
)
elapsed = (time.perf_counter() - start) * 1000
outputs[step.id] = output
@@ -6184,6 +6521,7 @@ def last_output(role: str) -> str:
verification,
free_only=model_name == self.FREE_MODEL,
allowed_agent_ids=judge_agent_ids,
+ excluded_agent_ids=_excluded_agent_ids,
)
answer = outputs[steps[-1].id]
if not verification["accepted"] and self.policy.verifier_required and last_output("worker"):
@@ -6196,6 +6534,7 @@ def last_output(role: str) -> str:
verification,
free_only=model_name == self.FREE_MODEL,
allowed_agent_ids=judge_agent_ids,
+ excluded_agent_ids=_excluded_agent_ids,
)
answer = outputs[steps[2].id] if not self.policy.verifier_required else outputs[steps[-1].id]
if not verification["accepted"] and self.policy.verifier_required:
@@ -6320,7 +6659,10 @@ def _plan_generated(self, task: str) -> list[WorkflowStep]:
f"roles={','.join(role for role in roles if agent.id in eligible_by_role[role])}, "
f"tags={', '.join(agent.tags) or 'none'}"
for agent in self.agents
- if any(agent.id in eligible_by_role[role] for role in roles)
+ if _is_general_chat_agent(agent)
+ and self._zdr_agent_allowed(agent)
+ and _agent_matches_request_endpoint(agent)
+ and any(agent.id in eligible_by_role[role] for role in roles)
)
system = (
"You are the workflow conductor. Decompose the user's task into a short workflow.\n"
@@ -6353,7 +6695,11 @@ def _parse_workflow_plan(self, raw: str) -> list[WorkflowStep]:
raw_steps = data.get("steps")
if not isinstance(raw_steps, list) or not (2 <= len(raw_steps) <= self.policy.max_workflow_steps):
raise ValueError(f"plan must have 2..{self.policy.max_workflow_steps} steps")
- known_agents = {agent.id: agent for agent in self.agents}
+ known_agents = {
+ agent.id: agent
+ for agent in self.agents
+ if _agent_matches_request_endpoint(agent)
+ }
steps: list[WorkflowStep] = []
for index, item in enumerate(raw_steps):
if int(item.get("id", -1)) != index:
@@ -6496,6 +6842,7 @@ def _ranked_agents(
agent
for agent in source
if not agent.disabled
+ and _agent_matches_request_endpoint(agent)
and self._zdr_agent_allowed(agent)
if (
not free_only
@@ -6756,7 +7103,9 @@ def _embedding_agent_id(self) -> str | None:
def _embed_cached(self, text: str) -> list[float] | None:
"""Embedding vector for text via the configured embedding member; None on failure."""
- digest = hashlib.sha256(text.encode("utf-8")).hexdigest()
+ digest = hashlib.sha256(
+ f"{_request_endpoint_partition()}\x1f{text}".encode("utf-8")
+ ).hexdigest()
with self._evidence_lock:
cached = self._task_vector_cache.get(digest)
if cached is not None:
@@ -6776,7 +7125,13 @@ def _embed_cached(self, text: str) -> list[float] | None:
def _descriptor_vector_cached(self, agent: ModelAgent) -> list[float] | None:
"""Cached embedding of one agent's operator-declared metadata document."""
fingerprint = hashlib.sha256(
- "\x1f".join([agent.id, self._agent_descriptor_text(agent)]).encode("utf-8")
+ "\x1f".join(
+ [
+ _request_endpoint_partition(),
+ agent.id,
+ self._agent_descriptor_text(agent),
+ ]
+ ).encode("utf-8")
).hexdigest()
with self._evidence_lock:
cached = self._descriptor_vector_cache.get(fingerprint)
@@ -6841,7 +7196,12 @@ def _triage_workflow_required(self, text: str) -> bool:
no evidence source exists at all. Verdicts are cached by content hash.
"""
digest = hashlib.sha256(
- (text + ("\x00zdr_only" if _REQUEST_ZDR_ONLY.get() else "")).encode("utf-8")
+ (
+ _request_endpoint_partition()
+ + "\x1f"
+ + text
+ + ("\x00zdr_only" if _REQUEST_ZDR_ONLY.get() else "")
+ ).encode("utf-8")
).hexdigest()
with self._evidence_lock:
cached = self._triage_cache.get(digest)
@@ -6858,7 +7218,9 @@ def _compute_triage_verdict(self, text: str) -> bool:
except RuntimeError:
candidates = []
if not candidates and not _REQUEST_ZDR_ONLY.get():
- candidates = list(self.agents)
+ candidates = [
+ agent for agent in self.agents if _agent_matches_request_endpoint(agent)
+ ]
if not candidates:
return False
triage_agent = candidates[0]
@@ -7249,6 +7611,7 @@ def _invoke(
role: str,
allowed_agent_ids: set[str] | None = None,
eligibility_role: str | None = None,
+ excluded_agent_ids: set[str] | None = None,
) -> tuple[str, str, str, dict[str, Any] | None]:
"""Call an agent with bounded, safety-aware tool retry and failover.
@@ -7278,6 +7641,12 @@ def _invoke(
allowed_agent_ids=allowed_agent_ids,
prompt_context=prompt_context,
)
+ if excluded_agent_ids:
+ candidates = [
+ candidate
+ for candidate in candidates
+ if candidate.id not in excluded_agent_ids
+ ]
if not candidates:
raise RuntimeError(f"no chat-compatible agent available for role={role}")
race_members = self._equivalent_race_members(candidates, capability="text")
@@ -7383,6 +7752,13 @@ def call(agent: ModelAgent) -> tuple[str, str, str, dict[str, Any] | None]:
raise
if isinstance(exc, ProviderUpstreamError):
last_upstream_error = exc
+ if (
+ excluded_agent_ids is not None
+ and exc.error_code == "model_not_found"
+ ):
+ excluded_agent_ids.add(agent.id)
+ self._record_failure(agent.id)
+ break
# The primary chat call is a bounded, side-effect-free
# model request, not a tool invocation: classify from
# the provider's own already-computed retryability
@@ -7588,7 +7964,7 @@ def _record_success(self, agent_id: str) -> None:
def _agent(self, agent_id: str) -> ModelAgent:
for agent in self.candidates:
- if agent.id == agent_id:
+ if agent.id == agent_id and _agent_matches_request_endpoint(agent):
return agent
raise KeyError(agent_id) # pragma: no cover
@@ -7660,6 +8036,7 @@ def _model_judge_verification(
*,
free_only: bool = False,
allowed_agent_ids: set[str] | None = None,
+ excluded_agent_ids: set[str] | None = None,
) -> dict[str, Any]:
"""Ask a model for a strict structured verdict and fail closed on uncertainty."""
verifier_output = fallback.get("verifier_output", "")
@@ -7692,6 +8069,7 @@ def _model_judge_verification(
agent
for agent in self._ranked_agents(task, "verifier", free_only=free_only)
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
)
# The judge is one bounded provider call. Do not pass the
# planning strategy ("template"/"generated") as an
@@ -7702,6 +8080,7 @@ def _model_judge_verification(
judge.id,
mode="route",
allowed_agent_ids=allowed_agent_ids,
+ excluded_agent_ids=excluded_agent_ids,
)
fast_judge = components.judge_cls(
judge_adapter,
@@ -8148,15 +8527,21 @@ def record_analytics_event(
if self._store is not None:
self._store.save("analytics", None, event)
- def _run_budget_output_by_model(self, record: Mapping[str, Any]) -> dict[str, int]:
- """Return the exact per-model output-token contribution of one run."""
+ def _run_budget_output_by_model(
+ self, record: Mapping[str, Any]
+ ) -> tuple[dict[str, int], bool]:
+ """Return authoritative per-model output tokens and availability."""
model_by_agent = {agent.id: agent.model for agent in self.candidates}
output_by_model: dict[str, int] = {}
for step in record.get("trace", []):
model = step.get("model_name") or model_by_agent.get(
step.get("served_agent_id") or step.get("agent_id"), "unknown"
)
- output_tokens, _reported = _step_output_tokens(step)
+ output_tokens, _source = _step_output_tokens(
+ step, self.token_counter, 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):
@@ -8175,19 +8560,23 @@ def _run_budget_output_by_model(self, record: Mapping[str, Any]) -> dict[str, in
# 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", ""),
- }
- )
- judge_model = verification.get("judge_model") or model_by_agent.get(
- judge_agent_id, "unknown"
+ },
+ 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
)
- return output_by_model
+ return output_by_model, True
def _replace_workflow_run(self, record: dict[str, Any]) -> None:
"""Store one run and update its constant-time budget meter atomically."""
@@ -8202,7 +8591,17 @@ def _replace_workflow_run(self, record: dict[str, Any]) -> None:
for sign, run in ((-1, previous), (1, record)):
if run is None:
continue
- for model, output_tokens in self._run_budget_output_by_model(run).items():
+ run_id_for_meter = run["workflow_run_id"]
+ output_by_model, available = self._run_budget_output_by_model(run)
+ available_for_budget = available and (
+ self.budget_max_cost_usd is None
+ or all(model in self.price_per_million for model in output_by_model)
+ )
+ if sign < 0:
+ self._budget_unavailable_run_ids.discard(run_id_for_meter)
+ elif not available_for_budget:
+ self._budget_unavailable_run_ids.add(run_id_for_meter)
+ for model, output_tokens in output_by_model.items():
before = self._budget_model_output_tokens.get(model, 0)
after = before + sign * output_tokens
price = self.price_per_million.get(model)
@@ -8221,10 +8620,19 @@ def _rebuild_budget_meter(self) -> None:
"""Reconcile the meter after a rare agent-pool identity change."""
with self._budget_spend_lock:
output_by_model: dict[str, int] = {}
+ unavailable_run_ids: set[str] = set()
for run in self._workflow_runs.values():
- for model, output_tokens in self._run_budget_output_by_model(run).items():
+ run_output, available = self._run_budget_output_by_model(run)
+ available_for_budget = available and (
+ self.budget_max_cost_usd is None
+ or all(model in self.price_per_million for model in run_output)
+ )
+ if not available_for_budget:
+ unavailable_run_ids.add(run["workflow_run_id"])
+ for model, output_tokens in run_output.items():
output_by_model[model] = output_by_model.get(model, 0) + output_tokens
self._budget_model_output_tokens = output_by_model
+ self._budget_unavailable_run_ids = unavailable_run_ids
self._budget_spent_output_tokens = sum(output_by_model.values())
self._budget_spent_cost_usd = sum(
(
@@ -8236,42 +8644,45 @@ def _rebuild_budget_meter(self) -> None:
)
def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> dict[str, Any]:
- """Estimated token and cost spend per model, aggregated from workflow runs.
-
- Tokens are ESTIMATED from runtime output text (~4 chars/token), not provider-reported
- usage. Cost is computed only for models with an operator-supplied price; models without
- one are reported under ``unpriced_models`` with a null cost. This is the honest local
- floor for spend observability, not a billing system.
- """
+ """Return provider-reported or exact-tokenizer output usage and cost."""
prices = {**self.price_per_million, **(price_per_million or {})}
model_by_agent = {agent.id: agent.model for agent in self.candidates}
by_model: dict[str, dict[str, Any]] = {}
total_output_tokens = 0
- total_prompt_tokens = 0
reported_prompt_tokens = 0
- any_reported_prompt = False
+ prompt_available = True
for run in self._workflow_runs.values():
- total_prompt_tokens += estimate_tokens(run.get("prompt_text", ""))
for step in run["trace"]:
model = step.get("model_name") or model_by_agent.get(
step.get("served_agent_id") or step.get("agent_id"), "unknown"
)
- estimated = estimate_tokens(step.get("output", ""))
usage = step.get("usage")
- reported_prompt = usage.get("prompt_tokens") if isinstance(usage, dict) else None
- if isinstance(reported_prompt, int):
+ reported_prompt = (
+ usage.get("prompt_tokens", usage.get("input_tokens"))
+ if isinstance(usage, dict)
+ else None
+ )
+ if type(reported_prompt) is int and reported_prompt >= 0:
reported_prompt_tokens += reported_prompt
- any_reported_prompt = True
- effective, is_reported = _step_output_tokens(step)
+ else:
+ prompt_available = False
+ effective, source = _step_output_tokens(step, self.token_counter, model)
bucket = by_model.setdefault(
- model, {"estimated_output_tokens": 0, "output_tokens": 0, "step_count": 0, "reported_steps": 0}
+ model,
+ {
+ "output_tokens": 0,
+ "step_count": 0,
+ "reported_steps": 0,
+ "tokenizer_steps": 0,
+ "unavailable_steps": 0,
+ },
)
- bucket["estimated_output_tokens"] += estimated
- bucket["output_tokens"] += effective
bucket["step_count"] += 1
- bucket["reported_steps"] += 1 if is_reported else 0
- total_output_tokens += effective
+ bucket[f"{source}_steps"] += 1
+ if effective is not None:
+ bucket["output_tokens"] += effective
+ total_output_tokens += effective
verification = run.get("verification")
judge_agent_id = (
@@ -8287,117 +8698,166 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) ->
# analytics entirely. Mirror the worker-step loop above
# exactly: estimate from the real judged text, and let
# _step_output_tokens report the honest reported/estimated
- # split instead of a fabricated "reported" dict (Devin
- # review on #961: an earlier revision of this fix fed a
- # zero-token usage that _step_output_tokens then counted
- # as genuinely provider-reported). Estimate from
+ # split instead of a fabricated "reported" dict. 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.
+ # not verifier_output (the worker answer it was judging).
judge_usage = verification.get("judge_usage")
judge_text = verification.get("judge_output_text", "")
judge_model = verification.get("judge_model") or model_by_agent.get(
judge_agent_id, "unknown"
)
- judge_estimated = estimate_tokens(judge_text)
- judge_output, judge_reported = _step_output_tokens(
- {"usage": judge_usage, "output": judge_text}
+ effective, source = _step_output_tokens(
+ {"usage": judge_usage, "output": judge_text},
+ self.token_counter,
+ judge_model,
+ )
+ reported_prompt = (
+ judge_usage.get("prompt_tokens", judge_usage.get("input_tokens"))
+ if isinstance(judge_usage, dict)
+ else None
)
- if judge_reported and isinstance(judge_usage, Mapping):
- judge_prompt = judge_usage.get("prompt_tokens")
- if type(judge_prompt) is int and judge_prompt >= 0:
- reported_prompt_tokens += judge_prompt
- any_reported_prompt = True
+ if type(reported_prompt) is int and reported_prompt >= 0:
+ reported_prompt_tokens += reported_prompt
+ else:
+ prompt_available = False
bucket = by_model.setdefault(
judge_model,
{
- "estimated_output_tokens": 0,
"output_tokens": 0,
"step_count": 0,
"reported_steps": 0,
+ "tokenizer_steps": 0,
+ "unavailable_steps": 0,
},
)
- bucket["estimated_output_tokens"] += judge_estimated
- bucket["output_tokens"] += judge_output
bucket["step_count"] += 1
- bucket["reported_steps"] += 1 if judge_reported else 0
- total_output_tokens += judge_output
+ bucket[f"{source}_steps"] += 1
+ if effective is not None:
+ bucket["output_tokens"] += effective
+ total_output_tokens += effective
rows: list[dict[str, Any]] = []
unpriced: list[str] = []
total_cost_usd = Decimal(0)
+ output_available = True
+ cost_available = True
for model, bucket in sorted(by_model.items()):
+ model_available = bucket["unavailable_steps"] == 0
+ output_available = output_available and model_available
price = prices.get(model)
cost_decimal = (
_cost_usd_decimal(bucket["output_tokens"], price)
- if price is not None
+ if price is not None and model_available
else None
)
cost = float(cost_decimal) if cost_decimal is not None else None
if price is None:
unpriced.append(model)
+ cost_available = False
+ elif not model_available:
+ cost_available = False
else:
total_cost_usd += cost_decimal
- if bucket["reported_steps"] == 0:
- usage_source = "estimated"
+ if not model_available:
+ usage_source = "unavailable"
elif bucket["reported_steps"] == bucket["step_count"]:
usage_source = "reported"
+ elif bucket["tokenizer_steps"] == bucket["step_count"]:
+ usage_source = "tokenizer"
else:
usage_source = "mixed"
rows.append({
"model": model,
- "estimated_output_tokens": bucket["estimated_output_tokens"],
- "output_tokens": bucket["output_tokens"],
+ "output_tokens": bucket["output_tokens"] if model_available else None,
"usage_source": usage_source,
"step_count": bucket["step_count"],
"price_per_million_usd": price,
- "estimated_cost_usd": cost,
+ "cost_usd": cost,
})
+ measurement_status = (
+ "unavailable"
+ if not output_available or not prompt_available
+ else "measured"
+ if all(row["usage_source"] == "reported" for row in rows)
+ else "exact_tokenizer"
+ )
+ candidate_prices_available = (
+ self.budget_max_cost_usd is None
+ or all(
+ agent.model in prices
+ for agent in self.agents
+ if _is_general_chat_agent(agent)
+ )
+ )
return {
- "measurement_status": "local_runtime_estimate",
+ "measurement_status": measurement_status,
"source_note": (
- "output_tokens use provider-reported usage when available (usage_source=reported/mixed) and "
- "fall back to a ~4 chars/token estimate otherwise; cost = output_tokens x operator-supplied price only."
+ "Provider usage is authoritative. Exact tokenizer counts apply only to "
+ "declared raw textual outputs; unreconstructible usage is unavailable."
),
"pricing_configured": bool(prices),
"totals": {
"run_count": len(self._workflow_runs),
- "estimated_output_tokens": total_output_tokens,
- "estimated_prompt_tokens": total_prompt_tokens,
- "reported_prompt_tokens": reported_prompt_tokens,
- "prompt_tokens_source": "reported" if any_reported_prompt else "estimated",
- "estimated_cost_usd": float(total_cost_usd) if prices else None,
+ "output_tokens": total_output_tokens if output_available else None,
+ "prompt_tokens": reported_prompt_tokens if prompt_available else None,
+ "prompt_tokens_source": "reported" if prompt_available else "unavailable",
+ "cost_usd": (
+ float(total_cost_usd) if prices and cost_available else None
+ ),
"currency": "USD",
},
"by_model": rows,
"unpriced_models": unpriced,
"budget": self._budget_block(
total_output_tokens,
- float(total_cost_usd) if prices else None,
+ float(total_cost_usd) if prices and cost_available else None,
+ measurement_available=output_available and candidate_prices_available,
),
}
- def _budget_block(self, spent_tokens: int, spent_cost: float | None) -> dict[str, Any]:
+ def _budget_block(
+ self,
+ spent_tokens: int,
+ spent_cost: float | None,
+ *,
+ measurement_available: bool = True,
+ ) -> dict[str, Any]:
token_limit = self.budget_max_output_tokens
cost_limit = self.budget_max_cost_usd
+ required_available = measurement_available and (
+ cost_limit is None or spent_cost is not None
+ )
exceeded = bool(
- (token_limit is not None and spent_tokens >= token_limit)
+ required_available
+ and ((token_limit is not None and spent_tokens >= token_limit)
or (cost_limit is not None and spent_cost is not None and spent_cost >= cost_limit)
+ )
)
+ enabled = token_limit is not None or cost_limit is not None
return {
- "enabled": token_limit is not None or cost_limit is not None,
+ "enabled": enabled,
"max_output_tokens": token_limit,
"max_cost_usd": cost_limit,
- "spent_output_tokens": spent_tokens,
- "spent_cost_usd": spent_cost,
- "remaining_output_tokens": max(0, token_limit - spent_tokens) if token_limit is not None else None,
+ "spent_output_tokens": spent_tokens if measurement_available else None,
+ "spent_cost_usd": spent_cost if required_available else None,
+ "remaining_output_tokens": (
+ max(0, token_limit - spent_tokens)
+ if token_limit is not None and measurement_available
+ else None
+ ),
"remaining_cost_usd": (
round(max(0.0, cost_limit - spent_cost), 6)
- if cost_limit is not None and spent_cost is not None else None
+ if cost_limit is not None and spent_cost is not None and required_available
+ else None
),
"exceeded": exceeded,
+ "measurement_status": "measured" if required_available else "unavailable",
+ "enforcement_status": (
+ "blocked_unavailable" if enabled and not required_available
+ else "exceeded" if exceeded
+ else "within_budget"
+ ),
}
def budget_status(self) -> dict[str, Any]:
@@ -8405,9 +8865,21 @@ def budget_status(self) -> dict[str, Any]:
with self._budget_spend_lock:
spent_tokens = self._budget_spent_output_tokens
spent_cost = float(self._budget_spent_cost_usd)
+ candidate_prices_available = (
+ self.budget_max_cost_usd is None
+ or all(
+ agent.model in self.price_per_million
+ for agent in self.agents
+ if _is_general_chat_agent(agent)
+ )
+ )
+ measurement_available = (
+ not self._budget_unavailable_run_ids and candidate_prices_available
+ )
return self._budget_block(
spent_tokens,
spent_cost if self.price_per_million else None,
+ measurement_available=measurement_available,
)
def analytics_snapshot(self, locale_bundles: dict[str, dict[str, str]] | None = None) -> dict[str, Any]:
@@ -14941,14 +15413,15 @@ def wrapper(self: TaskOrchestrator, *args: Any, **kwargs: Any) -> dict[str, Any]
def _pareto_front(results: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Configs not dominated on (quality up, cost down)."""
+ measured = [row for row in results if row.get("cost_usd") is not None]
front: list[dict[str, Any]] = []
- for a in results:
+ for a in measured:
dominated = any(
b is not a
and b["quality"] >= a["quality"]
and b["cost_usd"] <= a["cost_usd"]
and (b["quality"] > a["quality"] or b["cost_usd"] < a["cost_usd"])
- for b in results
+ for b in measured
)
if not dominated:
front.append(a)
@@ -14956,20 +15429,21 @@ def _pareto_front(results: list[dict[str, Any]]) -> list[dict[str, Any]]:
def _recommend_config(results: list[dict[str, Any]], cost_budget_usd: float | None) -> dict[str, Any] | None:
- if not results:
+ measured = [row for row in results if row.get("cost_usd") is not None]
+ if not measured:
return None
if cost_budget_usd is not None:
- affordable = [r for r in results if r["cost_usd"] <= cost_budget_usd]
+ affordable = [r for r in measured if r["cost_usd"] <= cost_budget_usd]
if affordable:
best = max(affordable, key=lambda r: (r["quality"], -r["cost_usd"]))
reason = "highest quality within cost budget"
else:
- best = min(results, key=lambda r: r["cost_usd"])
+ best = min(measured, key=lambda r: r["cost_usd"])
reason = "no config within budget; cheapest instead"
else:
# Maximize performance first, minimize cost as the tie-break (cheapest among the
# best-quality configs) — the honest reading of "max quality while min cost".
- best = max(results, key=lambda r: (r["quality"], -r["cost_usd"]))
+ best = max(measured, key=lambda r: (r["quality"], -r["cost_usd"]))
reason = "highest quality; cheapest among equal-quality configs"
return {"name": best["name"], "quality": best["quality"], "cost_usd": best["cost_usd"], "reason": reason}
@@ -15012,20 +15486,29 @@ def optimize_orchestration(
orchestrator = candidate["orchestrator"]
mode = candidate.get("mode", "auto")
quality = _score_config(orchestrator, tasks, quality_fn, mode, use_batch)
- cost = orchestrator.spend_analytics()["totals"]["estimated_cost_usd"] or 0.0
+ cost = orchestrator.spend_analytics()["totals"]["cost_usd"]
results.append({
"name": candidate["name"],
"mode": mode,
"quality": round(quality, 4),
- "cost_usd": round(cost, 6),
- "quality_per_usd": round(quality / cost, 2) if cost > 0 else None,
+ "cost_usd": round(cost, 6) if cost is not None else None,
+ "quality_per_usd": (
+ round(quality / cost, 2) if cost is not None and cost > 0 else None
+ ),
"task_count": len(tasks),
})
return {
"objective": "maximize quality, minimize cost",
"cost_budget_usd": cost_budget_usd,
- "results": sorted(results, key=lambda r: (-r["quality"], r["cost_usd"])),
+ "results": sorted(
+ results,
+ key=lambda r: (
+ -r["quality"],
+ r["cost_usd"] is None,
+ r["cost_usd"] if r["cost_usd"] is not None else float("inf"),
+ ),
+ ),
"pareto_front": [r["name"] for r in _pareto_front(results)],
"recommended": _recommend_config(results, cost_budget_usd),
}
@@ -15078,19 +15561,23 @@ def evaluate(config: dict[str, Any]) -> dict[str, Any]:
orchestrator = build_orchestrator(config)
mode = config.get("mode", "auto")
quality = _score_config(orchestrator, tasks, quality_fn, mode, use_batch)
- cost = orchestrator.spend_analytics()["totals"]["estimated_cost_usd"] or 0.0
+ cost = orchestrator.spend_analytics()["totals"]["cost_usd"]
result = {
"name": config_key,
"config": dict(config),
"quality": round(quality, 4),
- "cost_usd": round(cost, 6),
- "quality_per_usd": round(quality / cost, 2) if cost > 0 else None,
+ "cost_usd": round(cost, 6) if cost is not None else None,
+ "quality_per_usd": (
+ round(quality / cost, 2) if cost is not None and cost > 0 else None
+ ),
"task_count": len(tasks),
}
evaluated[config_key] = result
return result
def fitness(row: dict[str, Any]) -> tuple[int, float, float]:
+ if row["cost_usd"] is None:
+ return (0, row["quality"], float("-inf"))
affordable = 1 if cost_budget_usd is None or row["cost_usd"] <= cost_budget_usd else 0
return (affordable, row["quality"], -row["cost_usd"])
@@ -15146,8 +15633,7 @@ def chat_completion_response(
) -> dict[str, Any]: # pragma: no cover
"""Wrap orchestration output in an OpenAI-compatible chat completion response.
- ``usage`` carries the token counts recorded by the cost ledger; when absent
- the response reports zeros (the count could not be computed).
+ ``usage`` carries measured token counts. Absence remains explicit and null.
"""
orchestration = {
"workflow_run_id": result.get("workflow_run_id"),
@@ -15172,7 +15658,8 @@ def chat_completion_response(
"finish_reason": "stop",
}
],
- "usage": usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+ "usage": usage,
+ "usage_measurement_status": "measured" if usage is not None else "unavailable",
"orchestration": {key: value for key, value in orchestration.items() if value is not None},
}
@@ -15196,7 +15683,8 @@ def text_completion_response(
"finish_reason": "stop",
}
],
- "usage": usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
+ "usage": usage,
+ "usage_measurement_status": "measured" if usage is not None else "unavailable",
}
@@ -15264,13 +15752,20 @@ def chat_completion_chunks(
usage = result.get("usage")
cost = result.get("cost")
- if (
- include_usage
- and isinstance(cost, dict)
- and cost.get("measurement_status") == "measured"
- and isinstance(usage, dict)
- ):
- chunks.append({**base, "choices": [], "usage": usage})
+ if include_usage:
+ measured = (
+ isinstance(cost, dict)
+ and cost.get("measurement_status") == "measured"
+ and isinstance(usage, dict)
+ )
+ chunks.append(
+ {
+ **base,
+ "choices": [],
+ "usage": usage if measured else None,
+ "usage_measurement_status": "measured" if measured else "unavailable",
+ }
+ )
return chunks
diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py
index ecc763d3e..3b862792d 100644
--- a/contextual_orchestrator/provider_errors.py
+++ b/contextual_orchestrator/provider_errors.py
@@ -49,9 +49,17 @@
r"(?ix)(?:"
r"https?://|"
r"(?:^|[^0-9])(?:[0-9]{1,3}\.){3}[0-9]{1,3}(?:[^0-9]|$)|"
- r"\b(?:api[_ -]?key|authorization|bearer|password|secret|token|prompt|input|messages?)\b"
+ r"\b(?:api[_ -]?key|authorization|bearer|password|secret|token|prompt|input|messages?|content)\b"
r")"
)
+_SAFE_SCHEMA_DIAGNOSTIC = _re.compile(
+ r"['\"]?messages['\"]? must contain the word ['\"]?json['\"]?"
+ r"(?: in some form,)? to use "
+ r"(?:['\"]?response_format['\"]? of type ['\"]?json_object['\"]?|json_object)"
+ r"(?:\.|$)",
+ _re.IGNORECASE,
+)
+_SAFE_SCHEMA_ERROR_SUMMARY = "messages must mention json when response_format is json_object"
#: Upstream HTTP status -> ``(client_status, error_code, retryable)`` surface.
#: The client status is what this gateway returns; ``error_code`` follows the
@@ -148,6 +156,8 @@ def safe_provider_message(exc: BaseException) -> str | None:
for char in str(raw)
).strip()
collapsed = collapsed[:MAX_SAFE_MESSAGE_CHARS]
+ if _SAFE_SCHEMA_DIAGNOSTIC.search(collapsed):
+ return _SAFE_SCHEMA_ERROR_SUMMARY
if not collapsed or _SENSITIVE_PROVIDER_MESSAGE.search(collapsed):
return None
return collapsed
@@ -209,7 +219,18 @@ def classify_provider_failure(
secrets) can never reach callers or logs.
"""
if isinstance(exc, ProviderUpstreamError):
- return exc
+ if exc.transport == transport:
+ return exc
+ return ProviderUpstreamError(
+ agent_id=exc.agent_id,
+ model=exc.model,
+ error_code=exc.error_code,
+ message=str(exc),
+ client_status=exc.client_status,
+ provider_status=exc.provider_status,
+ retryable=exc.retryable,
+ transport=transport,
+ )
if isinstance(exc, urllib.error.HTTPError):
status = exc.code
client_status, error_code, retryable = PROVIDER_STATUS_SURFACES.get(
diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py
index 471a26493..eb6a77519 100644
--- a/contextual_orchestrator/server.py
+++ b/contextual_orchestrator/server.py
@@ -41,12 +41,13 @@
)
from .orchestrator import (
BudgetExceededError,
+ EndpointUnavailableError,
MAX_LOCAL_CONCURRENCY,
ProviderRequestTooLargeError,
ProviderResponseError,
ModelAgent,
TaskOrchestrator,
- estimate_tokens,
+ normalize_endpoint_selector,
_new_chat_completion_id,
_responses_to_chat_payload,
chat_completion_chunks,
@@ -67,6 +68,7 @@
detach_trace_context,
reset_session_id,
session_id_from_headers,
+ session_id_from_metadata,
session_id_from_request,
session_id_hash,
set_session_id,
@@ -165,6 +167,30 @@ class ResponsiveThreadingHTTPServer(ThreadingHTTPServer):
"""
request_queue_size = socket.SOMAXCONN
+ embedding_batch_backend: Any = None
+ embedding_backend_closer: Any = None
+
+ def _close_embedding_backend(self) -> None:
+ if callable(self.embedding_backend_closer):
+ self.embedding_backend_closer()
+ return
+ close = getattr(self.embedding_batch_backend, "close", None)
+ if callable(close):
+ close()
+
+ def shutdown(self) -> None:
+ """Stop accepting requests and release embedding worker threads."""
+ try:
+ super().shutdown()
+ finally:
+ self._close_embedding_backend()
+
+ def server_close(self) -> None:
+ """Close the listener and workers on normal or abnormal serve exit."""
+ try:
+ self._close_embedding_backend()
+ finally:
+ super().server_close()
# OpenAI request params forwarded verbatim to the provider on passthrough.
OPENAI_PASSTHROUGH_PARAM_KEYS = {
@@ -190,6 +216,7 @@ class ResponsiveThreadingHTTPServer(ThreadingHTTPServer):
ALLOWED_CHAT_KEYS = {
"model", "messages", "orchestration", "orchestration_mode", "mode",
"include_orchestration_trace", "stream", "attribution", "routing", "zdr_only",
+ "session_id",
# Tool-loop budget — accepted only for named unsupported error (no multi-step tool loop).
"max_tool_calls",
} | OPENAI_PASSTHROUGH_PARAM_KEYS
@@ -3144,7 +3171,9 @@ def _validate_attribution(attribution: Any) -> dict[str, Any] | None:
return cleaned or None
-def _validate_routing(routing: Any) -> dict[str, Any] | None:
+def _validate_routing(
+ routing: Any, *, allow_endpoint: bool = False
+) -> dict[str, Any] | None:
"""OpenAI-adjacent routing hints for sync vs batch channel selection.
Fail closed on shape so callers cannot smuggle non-boolean latency flags or
@@ -3155,7 +3184,10 @@ def _validate_routing(routing: Any) -> dict[str, Any] | None:
return None
if not isinstance(routing, dict):
raise RequestError(400, "invalid_routing", "routing must be an object")
- unknown = sorted(set(routing) - {"channel", "latency_tolerant", "priority"})
+ allowed = {"channel", "latency_tolerant", "priority"}
+ if allow_endpoint:
+ allowed.add("endpoint")
+ unknown = sorted(set(routing) - allowed)
if unknown:
raise RequestError(400, "invalid_routing", "routing contains unsupported keys", {"fields": unknown})
channel = routing.get("channel")
@@ -3199,6 +3231,30 @@ def _validate_routing(routing: Any) -> dict[str, Any] | None:
priority = routing.get("priority")
if isinstance(priority, str) and priority.strip():
cleaned["priority"] = priority.strip().lower()
+ if "endpoint" in routing:
+ endpoint = routing.get("endpoint")
+ if not isinstance(endpoint, str) or not endpoint.strip():
+ raise RequestError(
+ 400, "endpoint_unavailable", "routing.endpoint is unavailable"
+ )
+ try:
+ normalize_endpoint_selector(endpoint.strip())
+ except EndpointUnavailableError:
+ raise RequestError(
+ 400, "endpoint_unavailable", "routing.endpoint is unavailable"
+ ) from None
+ cleaned["endpoint"] = endpoint.strip()
+ if "endpoint" in cleaned and (
+ cleaned.get("channel") == "batch"
+ or cleaned.get("latency_tolerant") is True
+ ):
+ raise RequestError(
+ 400,
+ "invalid_routing",
+ "routing.endpoint cannot be combined with deferred batch routing",
+ )
+ if "endpoint" in cleaned:
+ cleaned["channel"] = "sync"
return cleaned if cleaned else {}
@@ -3763,13 +3819,18 @@ def _validate_chat_reasoning_effort(body: dict[str, Any]) -> None:
stripped = effort.strip().lower()
if not stripped:
return
+ if stripped == "auto":
+ # ``auto`` is the orchestrator-owned default used by consumers such
+ # as LineageWeave; it is not an OpenAI provider wire value.
+ body.pop("reasoning_effort", None)
+ return
if stripped in _OPENAI_REASONING_EFFORT_LEVELS:
body["reasoning_effort"] = stripped
return
raise RequestError(
400,
"invalid_reasoning_effort",
- "reasoning_effort must be one of none, minimal, low, medium, high "
+ "reasoning_effort must be one of auto, none, minimal, low, medium, high "
"on /v1/chat/completions",
)
@@ -5079,12 +5140,33 @@ def _response_payload(payload: dict[str, Any], include_trace: bool) -> dict[str,
return _strip_trace(public_payload)
+def _chat_usage_measurement_payload(payload: dict[str, Any]) -> dict[str, Any]:
+ """Add gateway usage provenance without rewriting provider chat content."""
+ normalized = dict(payload)
+ usage = normalized.get("usage")
+ prompt_tokens = usage.get("prompt_tokens") if isinstance(usage, dict) else None
+ completion_tokens = (
+ usage.get("completion_tokens") if isinstance(usage, dict) else None
+ )
+ measured = (
+ type(prompt_tokens) is int
+ and prompt_tokens >= 0
+ and type(completion_tokens) is int
+ and completion_tokens >= 0
+ )
+ if not measured:
+ normalized["usage"] = None
+ normalized["usage_measurement_status"] = (
+ "measured" if measured else "unavailable"
+ )
+ return normalized
+
+
def _chat_response_sse_chunks(
payload: dict[str, Any],
*,
model: str,
include_usage: bool,
- prompt_text: str,
) -> list[dict[str, Any]]:
"""Frame a completed provider-shaped chat response as OpenAI SSE chunks."""
completion_id = payload.get("id")
@@ -5166,21 +5248,35 @@ def _chat_response_sse_chunks(
for normal_chunk in chunks:
normal_chunk["usage"] = None
reported_usage = payload.get("usage")
- if isinstance(reported_usage, dict):
+ prompt_tokens = (
+ reported_usage.get("prompt_tokens", reported_usage.get("input_tokens"))
+ if isinstance(reported_usage, dict)
+ else None
+ )
+ completion_tokens = (
+ reported_usage.get("completion_tokens", reported_usage.get("output_tokens"))
+ if isinstance(reported_usage, dict)
+ else None
+ )
+ if (
+ type(prompt_tokens) is int
+ and prompt_tokens >= 0
+ and type(completion_tokens) is int
+ and completion_tokens >= 0
+ ):
usage = {**reported_usage, "usage_source": "reported"}
+ measurement_status = "measured"
else:
- completion_text = content if isinstance(content, str) else ""
- if tool_calls:
- completion_text += json.dumps(tool_calls, ensure_ascii=False)
- prompt_tokens = estimate_tokens(prompt_text)
- completion_tokens = estimate_tokens(completion_text)
- usage = {
- "prompt_tokens": prompt_tokens,
- "completion_tokens": completion_tokens,
- "total_tokens": prompt_tokens + completion_tokens,
- "usage_source": "estimated",
+ usage = None
+ measurement_status = "unavailable"
+ chunks.append(
+ {
+ **base,
+ "choices": [],
+ "usage": usage,
+ "usage_measurement_status": measurement_status,
}
- chunks.append({**base, "choices": [], "usage": usage})
+ )
return chunks
@@ -5706,7 +5802,11 @@ def do_GET(self) -> None: # noqa: N802
self._authorize("inference")
batch_id = path[len("/v1/batch/embeddings/"):]
try:
- self._send(coordinator.embeddings_batch_document(batch_id))
+ self._send(
+ coordinator.embeddings_batch_document(
+ batch_id, owner_id=security.principal_id(self.headers)
+ )
+ )
except KeyError:
self._send_error(404, "embeddings_batch_not_found", f"embeddings batch {batch_id} not found")
return
@@ -6329,6 +6429,7 @@ def do_DELETE(self) -> None: # noqa: N802
def do_POST(self) -> None: # noqa: N802
"""Dispatch authenticated completion, agent, and simulation writes."""
request_policy = None
+ endpoint_policy = None
try:
path = urllib.parse.urlparse(self.path).path
if path == "/admin/session":
@@ -6468,11 +6569,40 @@ def do_POST(self) -> None: # noqa: N802
zdr_only = _validate_zdr_only(body)
request_policy = orchestrator.request_policy(zdr_only)
request_policy.__enter__()
+ if path in {"/v1/chat/completions", "/v1/responses"}:
+ endpoint_routing = _validate_routing(
+ body.get("routing"), allow_endpoint=True
+ )
+ endpoint_policy = orchestrator.routing_endpoint_scope(
+ endpoint_routing.get("endpoint") if endpoint_routing else None,
+ body.get("model"),
+ model_was_provided="model" in body,
+ )
+ try:
+ endpoint_policy.__enter__()
+ except EndpointUnavailableError as exc:
+ endpoint_policy = None
+ raise RequestError(
+ 400,
+ "endpoint_unavailable",
+ "routing.endpoint is unavailable",
+ ) from exc
metadata_values = [
value
for key in ("metadata", "client_metadata")
if isinstance((value := body.get(key)), dict)
]
+ if "session_id" in body:
+ normalized_session_id = session_id_from_metadata(
+ {"session_id": body["session_id"]}
+ )
+ if normalized_session_id is None:
+ raise RequestError(
+ 400,
+ "invalid_session_id",
+ "session_id must be a non-empty string of at most 128 characters",
+ )
+ metadata_values.append({"session_id": normalized_session_id})
request_session_id = session_id_from_request(self.headers, *metadata_values)
if request_session_id != current_session_id():
self._bind_session(request_session_id)
@@ -6823,16 +6953,12 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
# "usage" key, so a provider may still omit it.
# _chat_response_sse_chunks (below) already frames that payload
# into a correctly-shaped terminal SSE chunk alongside tool_call
- # deltas, honestly labeling usage "reported" when the provider
- # sent it and "estimated" (never fabricated as "reported")
- # otherwise — the same fallback already exercised for the
- # non-tools streaming path — so there is nothing to fail closed
- # on here.
+ # deltas. Provider usage is measured when valid and explicitly
+ # unavailable otherwise; chat framing/tools are not reconstructed.
# response_format-only structured passthrough (conduct mode)
# is different: its usage comes from a multi-step workflow's
# cost ledger, which may be unmeasured, so it keeps failing
- # closed rather than let _chat_response_sse_chunks synthesize
- # an estimated figure for a workflow-level answer.
+ # closed when workflow-level usage is unavailable.
if stream and include_usage and not tool_loop:
raise RequestError(
400,
@@ -6884,7 +7010,9 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
)
else:
structured_messages = _validate_messages(body.get("messages"))
- structured_routing = _validate_routing(body.get("routing"))
+ structured_routing = _validate_routing(
+ body.get("routing"), allow_endpoint=True
+ )
if structured_routing and (
structured_routing.get("channel") == "batch"
or structured_routing.get("latency_tolerant") is True
@@ -6951,7 +7079,7 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
if include_trace and not trace_audited:
self._audit_trace_disclosure("/v1/chat/completions")
response_payload = (
- proxied
+ _chat_usage_measurement_payload(proxied)
if tool_loop
else _response_payload(proxied, include_trace)
)
@@ -6962,13 +7090,6 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
response_payload,
model=model_name,
include_usage=include_usage,
- prompt_text=json.dumps(
- {
- "messages": body.get("messages", []),
- "tools": body.get("tools"),
- },
- ensure_ascii=False,
- ),
)
)
)
@@ -6992,7 +7113,9 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
self._authorize_trace_access()
# stream + stream_options already coerced/validated before passthrough.
attribution = _validate_attribution(body.get("attribution"))
- routing = _validate_routing(body.get("routing"))
+ routing = _validate_routing(
+ body.get("routing"), allow_endpoint=True
+ )
# Require model — silent default to contextual-orchestrator hid
# which deployment the caller selected on the chat Completions path.
# The pool was validated before the structured/passthrough
@@ -7163,9 +7286,15 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
if not attribution.get("service"):
attribution["service"] = "embeddings_api"
started_at = time.perf_counter()
+ embedding_deadline = time.monotonic() + float(
+ orchestrator.client.timeout
+ )
document = None
last_embedding_error: Exception | None = None
for embedding_agent in embedding_agents:
+ remaining_timeout = embedding_deadline - time.monotonic()
+ if remaining_timeout <= 0:
+ break
attempt_started_at = time.perf_counter()
try:
document = self._run(lambda agent=embedding_agent: coordinator.complete_embeddings_batch(
@@ -7175,6 +7304,8 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
metadata={"actor_scope": "inference", "endpoint_alias": "embeddings"},
zdr_only=zdr_only,
agent_id=agent.id,
+ wait_timeout=remaining_timeout,
+ owner_id=security.principal_id(self.headers),
))
except Exception as exc: # noqa: BLE001 - measured member failover
last_embedding_error = exc
@@ -7185,7 +7316,12 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
embedding_agent.id,
time.perf_counter() - attempt_started_at,
)
- break
+ break
+ last_embedding_error = RuntimeError(
+ f"embedding member ended with {document.get('status', 'unknown')}"
+ )
+ orchestrator._group_router.observe_failure(embedding_agent.id)
+ document = None
if document is None:
raise RequestError(
503,
@@ -7272,6 +7408,7 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
metadata=submit_metadata,
zdr_only=zdr_only,
agent_id=agent.id,
+ owner_id=security.principal_id(self.headers),
))
except Exception as exc: # noqa: BLE001 - measured member failover
last_embedding_error = exc
@@ -7474,7 +7611,9 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
if "metadata" in body:
_validate_openai_metadata(body)
if "routing" in body:
- routing = _validate_routing(body.get("routing"))
+ routing = _validate_routing(
+ body.get("routing"), allow_endpoint=True
+ )
# Responses passthrough has no batch channel plane yet.
if routing and routing.get("channel") == "batch":
raise RequestError(
@@ -7642,7 +7781,9 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
responses_attribution.setdefault("model_name", body["model"])
responses_attribution.setdefault("service", "responses_api")
responses_routing = dict(
- _validate_routing(body.get("routing")) or {}
+ _validate_routing(
+ body.get("routing"), allow_endpoint=True
+ ) or {}
)
# Responses has no batch job envelope on this path;
# force the coordinator's synchronous contract even
@@ -7738,7 +7879,9 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
responses_messages,
mode="conduct",
attribution=responses_attribution,
- hints=_validate_routing(body.get("routing")),
+ hints=_validate_routing(
+ body.get("routing"), allow_endpoint=True
+ ),
model_name=body["model"],
provider_request=body,
provider_endpoint="responses",
@@ -7859,6 +8002,8 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di
traceback.print_exc()
self._send_error(500, "internal_error", "internal server error")
finally:
+ if endpoint_policy is not None:
+ endpoint_policy.__exit__(None, None, None)
if request_policy is not None:
request_policy.__exit__(None, None, None)
@@ -8567,7 +8712,10 @@ def _send_security_headers(self) -> None:
self.send_header("cache-control", "no-store")
self.send_header("x-frame-options", "DENY")
- return ResponsiveThreadingHTTPServer((host, port), Handler)
+ server = ResponsiveThreadingHTTPServer((host, port), Handler)
+ server.embedding_batch_backend = coordinator.embedding_batch_backend
+ server.embedding_backend_closer = coordinator.close_embedding_backends
+ return server
def serve(
@@ -8590,4 +8738,7 @@ def serve(
release_authority=release_authority,
)
print(f"listening on http://{host}:{port}")
- server.serve_forever()
+ try:
+ server.serve_forever()
+ finally:
+ server.server_close()
diff --git a/contextual_orchestrator/telemetry.py b/contextual_orchestrator/telemetry.py
index 737441a39..be1295947 100644
--- a/contextual_orchestrator/telemetry.py
+++ b/contextual_orchestrator/telemetry.py
@@ -3,7 +3,9 @@
from __future__ import annotations
import hashlib
+import ipaddress
import logging
+import re
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar, Token
@@ -39,6 +41,15 @@
# Match the OpenTelemetry SDK's default span-attribute budget so a single
# sequence-valued attribute cannot exceed the span's default evidence budget.
_MAX_ATTRIBUTE_SEQUENCE_ITEMS = 128
+_SAFE_SCHEMA_DIAGNOSTIC = re.compile(
+ r"messages (?:must contain the word ['\"]?json['\"]?(?: in some form,)? to use "
+ r"(?:['\"]?response_format['\"]? of type ['\"]?json_object['\"]?|json_object)"
+ r"(?:\.|$)|must mention json when response_format is json_object)",
+ re.IGNORECASE,
+)
+_SAFE_SCHEMA_ERROR_SUMMARY = (
+ "messages must mention json when response_format is json_object"
+)
_ALLOWED_ATTRIBUTE_KEYS = frozenset(
{
"gen_ai.operation.name",
@@ -51,7 +62,11 @@
"gen_ai.usage.total_tokens",
"contextual_orchestrator.agent_id",
"contextual_orchestrator.error_code",
+ "contextual_orchestrator.error_summary",
+ "contextual_orchestrator.fallback_outcome",
"contextual_orchestrator.latency_ms",
+ "contextual_orchestrator.model_group",
+ "contextual_orchestrator.operation_kind",
"contextual_orchestrator.provider_status_code",
"contextual_orchestrator.session_id_hash",
"server.address",
@@ -194,6 +209,13 @@ def _safe_attributes(
if isinstance(value, (list, tuple)):
continue
if isinstance(value, str):
+ if key == "server.address":
+ try:
+ ipaddress.ip_address(value)
+ except ValueError:
+ pass
+ else:
+ continue
result[key] = value[:256]
elif isinstance(value, (bool, int, float)):
result[key] = value
@@ -319,21 +341,44 @@ def traced(
try:
yield span
except Exception as exc:
- from .provider_errors import classify_provider_failure
+ from .provider_errors import classify_provider_failure, safe_provider_message
classified = classify_provider_failure(exc, agent_id="", model="")
failure_code = classified.error_code
provider_status = classified.provider_status
+ # Arbitrary provider prose can echo caller content even when it has
+ # no assignment-shaped marker. Recognize one exact schema contract,
+ # export our fixed wording, and reduce everything else to its stable
+ # package-owned code.
+ provider_summary = safe_provider_message(exc)
+ error_summary = (
+ _SAFE_SCHEMA_ERROR_SUMMARY
+ if provider_summary is not None
+ and _SAFE_SCHEMA_DIAGNOSTIC.search(provider_summary)
+ else failure_code
+ )
+ model_group = safe.get("contextual_orchestrator.model_group", "ungrouped")
+ fallback_outcome = safe.get(
+ "contextual_orchestrator.fallback_outcome", "not_observed"
+ )
if Status is not None and StatusCode is not None:
span.set_attribute("error.type", failure_code)
+ span.set_attribute(
+ "contextual_orchestrator.error_summary", error_summary
+ )
if provider_status is not None:
span.set_attribute(
"contextual_orchestrator.provider_status_code", provider_status
)
span.set_status(Status(StatusCode.ERROR))
_LOGGER.warning(
- "telemetry.operation_failed operation=%s error_type=%s",
+ "telemetry.operation_failed operation=%s error_type=%s "
+ "provider_status=%s error_summary=%r model_group=%s fallback_outcome=%s",
name,
failure_code,
+ provider_status,
+ error_summary,
+ model_group,
+ fallback_outcome,
)
raise
diff --git a/contextual_orchestrator/token_counting.py b/contextual_orchestrator/token_counting.py
index 0779cb1d1..48c5de47f 100644
--- a/contextual_orchestrator/token_counting.py
+++ b/contextual_orchestrator/token_counting.py
@@ -1,107 +1,178 @@
-"""Token counting seam for usage/cost accounting.
+"""Authoritative raw-text token counting for accounting boundaries.
-The cost ledger needs prompt/completion token counts on every completion. Two
-strategies are provided behind one :class:`TokenCounter`-compatible surface:
-
-* :class:`HeuristicTokenCounter` — a dependency-free estimator (the default).
- It approximates BPE token counts from whitespace/word structure so standalone
- runs and tests get stable, deterministic numbers without Postgres.
-* :class:`PgTiktokenAdapter` — delegates to ``pg_llm_batch.TokenCounter``
- (``pg_tiktoken`` running *inside* Postgres) when a DSN + the package are
- available, so counts match exactly what the batch engine bills against.
-
-Selection is centralised in :func:`build_token_counter`, which never reads the
-environment: the DSN is passed in by the caller.
+Provider-reported chat usage is authoritative. Local counters handle raw text
+only; they never reconstruct provider chat framing, tool schemas, or
+multimodal serialization. Exact full model identifiers select the packaged
+Rust tokenizer. Unknown identifiers and missing native code are explicitly
+unavailable rather than estimated.
"""
from __future__ import annotations
-import math
-import re
-from typing import Any, List, Optional, Protocol
-
-_WORD_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
-
-# Rough BPE expansion: sub-word models emit slightly more tokens than words.
-_TOKENS_PER_WORD = 1.3
+import importlib
+import operator
+from typing import Any, Optional, Protocol
+
+_CL100K_EMBEDDING_MODELS = frozenset(
+ {
+ "text-embedding-ada-002",
+ "text-embedding-3-small",
+ "text-embedding-3-large",
+ }
+)
+_CL100K_MODELS = _CL100K_EMBEDDING_MODELS | frozenset(
+ {
+ "gpt-4",
+ "gpt-3.5-turbo",
+ "gpt-3.5",
+ "gpt-35-turbo",
+ "davinci-002",
+ "babbage-002",
+ }
+)
+_O200K_MODELS = frozenset({"o1", "o3", "o4-mini", "gpt-5", "gpt-4.1", "gpt-4o"})
class TokenCountingStrategy(Protocol):
- """Contract for anything that can count tokens for a chunk of text."""
+ """Contract for an authoritative raw-text token counter."""
def count_text(self, text: str, model: str) -> int:
- """Return the token count for ``text`` under ``model``."""
+ """Return the exact token count for ``text`` under ``model``."""
...
-class HeuristicTokenCounter:
- """Deterministic, dependency-free token estimator.
+class TokenCountUnavailable(RuntimeError):
+ """An authoritative tokenizer or provider count is unavailable."""
- Counts word-ish units (words and standalone punctuation) and applies a
- fixed BPE expansion factor. Not exact, but stable and monotonic — good
- enough for cost attribution when ``pg_tiktoken`` is not reachable, and it
- never varies between runs so tests can assert on it.
- """
- def __init__(self, tokens_per_word: float = _TOKENS_PER_WORD) -> None:
- self.tokens_per_word = tokens_per_word
-
- def count_text(self, text: str, model: str = "") -> int:
- """Estimate the number of tokens in ``text``."""
- if not text:
- return 0
- units = _WORD_RE.findall(text)
- if not units:
- return 0
- return max(1, math.ceil(len(units) * self.tokens_per_word))
-
- def count_messages(self, messages: List[dict], model: str = "") -> int:
- """Estimate prompt tokens across a list of chat messages."""
- total = 0
- for message in messages:
- content = message.get("content", "") if isinstance(message, dict) else ""
- total += self.count_text(str(content), model)
- # Per-message framing overhead (role tags, delimiters).
- total += 3
- return total
+def _validated_count(value: Any) -> int:
+ """Normalize a tokenizer count and reject invalid numeric evidence."""
+ if isinstance(value, bool):
+ raise TokenCountUnavailable("the tokenizer returned an invalid count")
+ try:
+ count = operator.index(value)
+ except TypeError as exc:
+ raise TokenCountUnavailable("the tokenizer returned an invalid count") from exc
+ if count < 0:
+ raise TokenCountUnavailable("the tokenizer returned an invalid count")
+ return count
class PgTiktokenAdapter:
- """Adapter delegating to ``pg_llm_batch.TokenCounter`` (pg_tiktoken)."""
+ """Adapter delegating raw-text counts to ``pg_llm_batch.TokenCounter``."""
def __init__(self, pg_counter: Any) -> None:
self._counter = pg_counter
def count_text(self, text: str, model: str = "") -> int:
- """Count tokens via the Postgres ``pg_tiktoken`` extension."""
- # pg_llm_batch.TokenCounter exposes count_tokens(text, model).
- return int(self._counter.count_tokens(text, model))
+ """Count raw text through the configured PostgreSQL tokenizer."""
+ try:
+ return _validated_count(self._counter.count_tokens(text, model))
+ except TokenCountUnavailable:
+ raise
+ except Exception as exc: # noqa: BLE001 - external tokenizer boundary.
+ raise TokenCountUnavailable("the PostgreSQL tokenizer is unavailable") from exc
- def count_messages(self, messages: List[dict], model: str = "") -> int:
- """Count prompt tokens across chat messages via pg_tiktoken."""
- total = 0
- for message in messages:
- content = message.get("content", "") if isinstance(message, dict) else ""
- total += self.count_text(str(content), model)
- return total
+ def count_messages(self, messages: list[dict], model: str = "") -> int:
+ """Reject chat prompts whose provider framing is unreconstructible."""
+ raise TokenCountUnavailable("provider chat framing is unavailable")
-def build_token_counter(
- postgres_dsn: Optional[str] = None,
- *,
- config: Any = None,
-) -> HeuristicTokenCounter | PgTiktokenAdapter:
- """Return the best available token counter.
+class NativeExactTokenCounter:
+ """Dispatch exact raw-text counts for explicitly mapped model identifiers."""
+
+ def __init__(self, native_module: Any) -> None:
+ self._native_module = native_module
+
+ def count_text(self, text: str, model: str = "") -> int:
+ """Count raw text for a declared model or fail closed."""
+ if model in _CL100K_MODELS:
+ function_name = "count_cl100k"
+ elif model in _O200K_MODELS:
+ function_name = "count_o200k"
+ else:
+ raise TokenCountUnavailable(
+ f"no authoritative tokenizer is declared for {model!r}"
+ )
+ try:
+ function = getattr(self._native_module, function_name)
+ return _validated_count(function(text))
+ except Exception as exc: # noqa: BLE001 - optional native boundary.
+ raise TokenCountUnavailable("the native tokenizer is unavailable") from exc
+
+ def count_messages(self, messages: list[dict], model: str = "") -> int:
+ """Reject chat prompts whose framing/tools cannot be counted as raw text."""
+ raise TokenCountUnavailable("provider chat framing is unavailable")
+
+ def pack_text(self, text: str, model: str, max_tokens: int) -> list[tuple[str, int]]:
+ """Split one declared cl100k input at exact native token boundaries."""
+ if model not in _CL100K_EMBEDDING_MODELS:
+ raise TokenCountUnavailable(f"no authoritative tokenizer is declared for {model!r}")
+ try:
+ parts, _shards = self._native_module.pack_cl100k(
+ [text], max_tokens, 1, max_tokens
+ )
+ return [(part.text, _validated_count(part.token_count)) for part in parts]
+ except Exception as exc: # noqa: BLE001 - optional native boundary.
+ raise TokenCountUnavailable("the native cl100k packer is unavailable") from exc
+
+
+# Compatibility name retained for embedding callers; behavior remains exact.
+NativeCl100kTokenCounter = NativeExactTokenCounter
+
+
+class UnavailableTokenCounter:
+ """Represent absence of an authoritative tokenizer."""
+
+ def count_text(self, text: str, model: str = "") -> int:
+ """Fail closed instead of fabricating a raw-text token count."""
+ raise TokenCountUnavailable(f"no authoritative tokenizer is available for {model!r}")
+
+ def count_messages(self, messages: list[dict], model: str = "") -> int:
+ """Fail closed instead of fabricating a chat prompt count."""
+ raise TokenCountUnavailable("provider chat framing is unavailable")
- Prefers ``pg_tiktoken`` (via ``pg_llm_batch``) when a DSN is supplied and the
- dependency is importable; otherwise returns the heuristic estimator. Never
- reads the environment.
- """
+
+UnavailableEmbeddingTokenCounter = UnavailableTokenCounter
+
+
+def _native_token_counter() -> NativeExactTokenCounter | None:
+ """Load the optional extension without making startup depend on it."""
+ try:
+ module = importlib.import_module("contextual_orchestrator._token_packer")
+ except Exception: # noqa: BLE001 - incompatible wheel equals absence.
+ return None
+ functions = ("count_cl100k", "count_o200k", "pack_cl100k")
+ if not all(callable(getattr(module, name, None)) for name in functions):
+ return None
+ return NativeExactTokenCounter(module)
+
+
+def _build_counter(postgres_dsn: Optional[str], config: Any) -> Any:
+ """Build the configured authoritative counter or unavailable seam."""
if postgres_dsn:
try: # pragma: no cover - needs Postgres + pg_tiktoken extension
from pg_llm_batch import TokenCounter as PgTokenCounter # type: ignore
return PgTiktokenAdapter(PgTokenCounter(postgres_dsn, config=config))
- except Exception: # pragma: no cover - degrade to heuristic
- return HeuristicTokenCounter()
- return HeuristicTokenCounter()
+ except Exception: # pragma: no cover - optional authoritative boundary
+ pass
+ return _native_token_counter() or UnavailableTokenCounter()
+
+
+def build_embedding_token_counter(
+ postgres_dsn: Optional[str] = None,
+ *,
+ config: Any = None,
+) -> PgTiktokenAdapter | NativeExactTokenCounter | UnavailableTokenCounter:
+ """Return an authoritative embedding counter or explicit unavailable seam."""
+ return _build_counter(postgres_dsn, config)
+
+
+def build_token_counter(
+ postgres_dsn: Optional[str] = None,
+ *,
+ config: Any = None,
+) -> PgTiktokenAdapter | NativeExactTokenCounter | UnavailableTokenCounter:
+ """Return an authoritative raw-text counter or explicit unavailable seam."""
+ return _build_counter(postgres_dsn, config)
diff --git a/docs/adr/0002-control-plane-orchestrator.md b/docs/adr/0002-control-plane-orchestrator.md
index ac317f697..d70f12410 100644
--- a/docs/adr/0002-control-plane-orchestrator.md
+++ b/docs/adr/0002-control-plane-orchestrator.md
@@ -47,7 +47,12 @@ trained Fugu, TRINITY, or Conductor clone.
3. **Access lists.** Each `WorkflowStep` carries an access list so a worker
sees only the prior outputs deliberately exposed to it (Conductor-style
visibility, implemented as data on the step, not as a trained topology
- policy).
+ policy). The worker preserves the caller message array exactly once, then
+ receives only its subtask and deliberately exposed prior outputs in the
+ added user envelope. That envelope does not repeat the current task or
+ source attachments. Caller system instructions are reasserted in the
+ stage-role system message so their authority survives provider translation;
+ they are not copied into the added user envelope.
4. **Deterministic policy.** Worker and role selection uses a deterministic
capability-hint heuristic so the lab runs without training data, GPUs, or
vendor credentials. The heuristic is never an answer-quality,
diff --git a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md
new file mode 100644
index 000000000..043ac4ea4
--- /dev/null
+++ b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md
@@ -0,0 +1,105 @@
+# ADR 0005: Provider-embedding lease and token-accounting boundary
+
+- Status: Accepted
+- Date: 2026-08-31
+- Decision owners: ContextualWisdomLab
+- Series: `docs/adr` only. This is not a planning-ADR number.
+
+## Context
+
+Provider embedding jobs may outlive one Valkey execution-claim lease. A
+worker whose renewal fails can no longer prove that it owns the job, while a
+different worker may acquire the same claim. Publishing the first worker's
+result after that point would make usage, result, and terminal state depend
+on an expired lease.
+
+The embedding subsystem also builds a PyO3 extension backed by
+`tiktoken-rs`, but the Python runtime did not call it. Its historical local
+fallback estimated tokens from word units. Such an estimate cannot enforce a
+provider token limit or support billed cost accounting. OpenAI's published
+`tiktoken` mapping declares `cl100k_base` for
+`text-embedding-ada-002`, `text-embedding-3-small`, and
+`text-embedding-3-large`; it does not authorize guessing a tokenizer for an
+unknown model identifier.
+
+Redis's distributed-lock guidance requires a client to act only while it
+still owns the lock and recommends fencing when correctness depends on
+exclusive work. PyO3's module interface supports an in-package extension, so
+the existing Rust library can be loaded without adding a provider SDK or a
+second tokenization implementation.
+
+Chubby's production experience likewise separates coarse-grained advisory
+locking from the application-specific checks needed before publishing state;
+that operational boundary grounds the explicit terminal-publication fence here.
+
+## Decision
+
+1. **Lease loss is observable.** A durable execution claim exposes an
+ ownership check. A failed renewal, an elapsed renewal deadline, an
+ ownership-check error, or a negative ownership result marks the claim
+ lost and fails closed.
+2. **Terminal publication is fenced and atomic.** One Valkey Lua transaction
+ verifies the execution-claim token, accepts only `queued` or `running`, and
+ writes result/usage/error plus the terminal state together. The same fence
+ protects success and failure. A stale worker leaves recoverable `running`
+ state and retries claim acquisition while the backend remains live; it does
+ not overwrite a successor or require a process restart.
+3. **The provider call is not declared exactly-once.** A synchronous
+ provider call cannot be cancelled retroactively when its lease is lost.
+ Duplicate upstream execution remains possible during a partition. This
+ decision fences stale local publication; provider-side idempotency needs a
+ separate supported contract.
+4. **Embedding counts are authoritative or unavailable.** An explicitly
+ configured `pg_tiktoken` counter remains first. Otherwise the packaged
+ Rust counter is used only for the three published cl100k embedding model
+ identifiers above. Missing or failing native code and an unknown tokenizer
+ produce an explicit unavailable outcome. Embedding splitting, provider
+ dispatch, usage publication, and cost publication do not substitute a
+ word-count or BPE heuristic.
+5. **No routing/model inference is added.** The mapping is an exact tokenizer
+ contract, not evidence for model quality, provider selection, or
+ equivalence. The native extension does not count chat framing.
+6. **Legacy chat estimation was a known gap.** This ADR made no global
+ token-accounting compliance claim. ADR 0006 subsequently replaced that
+ legacy chat estimate with authoritative-or-unavailable accounting.
+
+## Consequences
+
+### Positive
+
+- An expired worker cannot publish a terminal provider-embedding result.
+- Installed production wheels execute the existing Rust cl100k counter.
+- Missing tokenizer authority is visible before provider dispatch and before
+ a cost record can be fabricated.
+- The PostgreSQL tokenizer boundary and dependency-injected test seams remain
+ available.
+
+### Negative
+
+- Embedding requests for an undeclared tokenizer reject unless an
+ authoritative PostgreSQL counter is configured.
+- Lease fencing prevents stale publication but does not eliminate duplicate
+ provider work.
+- Chat accounting is governed separately by ADR 0006.
+
+## References
+
+Burrows, M. (2006). *The Chubby lock service for loosely-coupled distributed
+systems*. 7th USENIX Symposium on Operating Systems Design and Implementation.
+https://research.google/pubs/the-chubby-lock-service-for-loosely-coupled-distributed-systems/
+
+The publisher-hosted paper is linked rather than copied because repository
+redistribution permission was not established.
+
+Redis Ltd. (n.d.). *Distributed locks with Redis*.
+https://redis.io/docs/latest/develop/clients/patterns/distributed-locks/
+
+PyO3 Project. (n.d.). *Python modules*.
+https://pyo3.rs/main/module
+
+OpenAI. (n.d.). *OpenAI public encodings*.
+https://github.com/openai/tiktoken/blob/main/tiktoken_ext/openai_public.py
+
+ContextualWisdomLab. (2026). *Cost-aware sync-versus-batch routing*
+(ADR 0003).
+https://github.com/ContextualWisdomLab/contextual-orchestrator/blob/main/docs/adr/0003-cost-aware-sync-batch-routing.md
diff --git a/docs/adr/0006-authoritative-chat-token-accounting.md b/docs/adr/0006-authoritative-chat-token-accounting.md
new file mode 100644
index 000000000..9f72d1001
--- /dev/null
+++ b/docs/adr/0006-authoritative-chat-token-accounting.md
@@ -0,0 +1,100 @@
+# ADR 0006: Authoritative chat token accounting
+
+- Status: Accepted
+- Date: 2026-08-31
+- Decision owners: ContextualWisdomLab
+- Series: `docs/adr` only. This is not a planning-ADR number.
+
+## Context
+
+Chat routing, run budgets, cost records, analytics, and streamed responses
+historically filled missing provider usage with deterministic word or
+character estimates. Those values were reproducible but not authoritative:
+chat framing, tools, and multimodal parts are provider protocol inputs rather
+than raw text, and their serialization is not reconstructible from a generic
+OpenAI-compatible request.
+
+OpenAI's public `tiktoken` table maps exact model identifiers to encodings and
+separately exposes prefix mappings. A prefix match can also accept a model
+identifier that does not exist, so it is not evidence that an arbitrary model
+uses an encoding. The packaged PyO3 extension already provides exact raw-text
+tokenization without adding a provider SDK.
+
+OpenAI-compatible Chat Completions responses expose provider usage when the
+provider measured it. Streamed responses may omit the terminal usage frame,
+including when a stream is interrupted. Missing usage is therefore an
+unavailable measurement, not zero and not permission to estimate.
+
+## Decision
+
+1. **Provider usage is authoritative.** Valid non-negative integral prompt
+ and completion counts reported by the provider are the usage and billing
+ evidence. The orchestrator does not replace or reconcile them with a local
+ estimate.
+2. **Local tokenization is raw-output-only.** The packaged Rust extension may
+ count a textual model output only when the complete model identifier has an
+ explicit encoding mapping in this decision. `gpt-4`, `gpt-3.5-turbo`,
+ `gpt-3.5`, `gpt-35-turbo`, `davinci-002`, and `babbage-002` use
+ `cl100k_base`; `o1`, `o3`, `o4-mini`, `gpt-5`, `gpt-4.1`, and `gpt-4o` use
+ `o200k_base`. Prefix matching and provider/model-name inference are
+ prohibited.
+3. **Chat prompts are not locally reconstructed.** Raw tokenizers do not count
+ provider chat framing, tool schemas, or multimodal serialization. When a
+ provider does not report prompt usage, prompt usage and any cost that
+ depends on it are explicitly unavailable. A PostgreSQL raw-text tokenizer
+ follows the same boundary.
+4. **Routing is conservative when usage is unavailable.** Explicit batch,
+ bulk, and latency declarations continue to decide routing. When only the
+ token threshold could select batch and prompt usage is unavailable, the
+ request stays synchronous.
+5. **Enabled budgets fail closed.** An enabled output-token or cost budget
+ cannot be evaluated from an unavailable required count. Dispatch is blocked
+ with an explicit unavailable measurement status; unavailable is never
+ interpreted as zero remaining spend.
+6. **Wire and ledger truth remain distinguishable.** API usage and cost fields
+ are nullable and include an explicit measurement status. Storage schemas
+ that require numeric token/cost columns may store zero only as a sentinel
+ beside `measurement_status=unavailable`; readers must return null rather
+ than treating that sentinel as measured free usage.
+7. **Streams do not synthesize usage.** A terminal provider usage frame is
+ forwarded as measured. Its absence produces `usage=null` with explicit
+ unavailable status. The request and SSE event protocol otherwise remains
+ unchanged.
+8. **Heuristic fields are removed from this subsystem.** Chat accounting does
+ not publish `estimated_*` usage or cost aliases. Historical psychometric
+ simulation estimates outside chat accounting are unaffected by this
+ decision.
+
+## Consequences
+
+### Positive
+
+- Routing, budgets, accounting, analytics, and SSE agree on one evidence
+ boundary.
+- A missing provider usage frame is visible instead of being converted to a
+ plausible-looking charge.
+- Known raw textual outputs can still support exact output-token budgets
+ through the existing native dependency.
+
+### Negative
+
+- Some providers and model identifiers now expose unavailable usage/cost where
+ the previous implementation returned an estimate.
+- Token-triggered batch routing stays synchronous without authoritative prompt
+ usage.
+- Cost budgets block when required usage evidence is absent.
+
+## References
+
+OpenAI. (n.d.). *Tiktoken model mappings*.
+https://github.com/openai/tiktoken/blob/main/tiktoken/model.py
+
+OpenAI. (n.d.). *Chat Completions API reference*.
+https://platform.openai.com/docs/api-reference/chat/create
+
+PyO3 Project. (n.d.). *Python modules*.
+https://pyo3.rs/main/module
+
+ContextualWisdomLab. (2026). *Cost-aware sync-versus-batch routing*
+(ADR 0003).
+https://github.com/ContextualWisdomLab/contextual-orchestrator/blob/main/docs/adr/0003-cost-aware-sync-batch-routing.md
diff --git a/docs/adr/README.md b/docs/adr/README.md
index c3b50506d..a11c4342b 100644
--- a/docs/adr/README.md
+++ b/docs/adr/README.md
@@ -13,6 +13,8 @@ They do not share numbering with `docs/planning/adrs/`.
| [0003](0003-cost-aware-sync-batch-routing.md) | Cost-aware sync-versus-batch routing | Accepted | Chen et al. (2023) FrugalGPT arXiv:2305.05176; Ong et al. (2024) RouteLLM arXiv:2406.18665; Ding et al. (2024) Hybrid LLM arXiv:2404.14618 |
| [0004](0004-msa-leaf-composition.md) | MSA leaf — standalone and callable | Accepted | NIST SP 800-204 independent deployability; planning ADR 0001 fail-closed judge composition |
| [0005](0005-verbose-debug-logging.md) | Verbose/debug logging with a redaction safety net | Accepted | OWASP Logging Cheat Sheet; NIST SP 800-92 log management; Python `logging` HOWTO |
+| [0005](0005-provider-embedding-lease-and-token-accounting.md) | Provider-embedding lease and token-accounting boundary | Accepted | Redis distributed-lock ownership/fencing guidance; PyO3 modules; OpenAI public cl100k mappings |
+| [0006](0006-authoritative-chat-token-accounting.md) | Authoritative chat token accounting | Accepted | OpenAI Chat usage contract and exact tiktoken model mappings; PyO3 modules |
Each record uses Context / Decision / Consequences plus an APA 7th
**References** section. Cite only verified DOI or official URLs. arXiv
diff --git a/docs/library_research.md b/docs/library_research.md
index 66ff1ca65..e2ac34480 100644
--- a/docs/library_research.md
+++ b/docs/library_research.md
@@ -21,6 +21,9 @@ primitives use maintained libraries when the enterprise target requires them.
| SSE usage capture | Python stdlib streaming parser already used by `ModelClient._stream_send` | Reuse the existing line-delimited SSE parser and capture only provider-declared usage frames; do not add an SSE or provider SDK dependency. | OpenAI's Responses and Chat Completions references define terminal usage fields, while interrupted streams may omit the final usage frame. |
| Verbose/debug logging | Python stdlib `logging` (researched: `structlog`, `loguru`) | Use stdlib `logging` exclusively -- `logging.basicConfig(..., force=True)` for one configuration entrypoint, a `logging.Filter` on the installed handler for redaction, `%`-style lazy formatting for cost-free DEBUG below its threshold. `structlog`/`loguru` add structured/prettier output this repo's existing `print(json.dumps(...))` CLI-report convention and `_LOGGER = logging.getLogger(__name__)` precedent (3 modules) do not need yet. | Python's own `logging` HOWTO documents `basicConfig`'s one-shot-unless-`force` behavior, handler-level `Filter`s, and that `isEnabledFor` gates expensive argument construction, not just formatting -- covering every requirement (level control, lazy evaluation, a redaction hook) with zero new dependency surface. |
| Distributed tracing | [OpenTelemetry Python](https://github.com/open-telemetry/opentelemetry-python) (already a runtime dependency since ADR 0122; recorded here for completeness -- this row was missing when that ADR shipped) | Keep as the request-correlation/span backend for cross-provider tracing (`telemetry.py`); it stays a separate system from stdlib logging (verbose/debug logging row above) -- two systems, not one, because OTel's span/attribute model and stdlib `logging`'s line-oriented model solve different problems and merging them would require a third abstraction neither currently needs. | OpenTelemetry's Python SDK and OTLP HTTP exporter are the maintained reference implementation for the vendor-neutral tracing API this repo's GenAI span conventions already target (see ADR 0122's References). |
+| Provider-embedding claim ownership | Existing `redis-py` lock token plus one Valkey Lua transaction | Propagate renewal loss, compare the live execution token, and atomically write terminal state with result/usage/error. A live worker retries claim acquisition until terminal state or deadline; provider-side exactly-once execution is not claimed. | Redis's official distributed-lock guidance requires ownership-safe release/extension and recommends fencing when correctness depends on exclusive work. Reused the existing registry and skipped a new coordination dependency, forced cancellation of synchronous provider I/O, and an unsupported provider-idempotency claim. |
+| Provider-embedding token accounting | Existing PyO3 + `tiktoken-rs` extension, with configured `pg_tiktoken` first | Load the packaged Rust extension in the production embedding path for the exact OpenAI-published cl100k embedding model IDs. Missing/failing native code and unknown tokenizers are explicitly unavailable; splitting, provider dispatch, usage, and cost fail closed instead of estimating. | OpenAI's public encoding table maps `text-embedding-ada-002`, `text-embedding-3-small`, and `text-embedding-3-large` to cl100k; PyO3 publishes the existing module in-package. Skipped tokenizer-name inference, a second tokenizer implementation, a provider SDK, and heuristic fallback. ADR 0006 now governs chat accounting separately. |
+| Chat token accounting | Existing provider usage fields plus the packaged PyO3 + `tiktoken-rs` extension | Treat valid provider usage as authoritative. Use Rust only for raw textual output from exact model IDs declared by ADR 0006; prompt framing, tools, multimodal input, unknown models, missing native code, and missing stream usage are explicitly unavailable. Enabled budgets fail closed and token-threshold routing remains synchronous when the required count is unavailable. | OpenAI's Chat Completions contract carries provider usage and notes streamed usage can be absent; OpenAI's public tiktoken model table separates exact mappings from unsafe prefix matching. Reused the existing extension and storage status seam. Skipped a provider SDK, prompt-serialization reimplementation, prefix/name inference, heuristic estimates, and fabricated zero-cost reporting. |
## Ponytail Decision
@@ -115,6 +118,7 @@ missing or invalid.
| What may an observation claim? | Council of Europe CEFR Companion Volume and linking manual | Preserve criterion-level evidence and transparent linking inputs; do not emit a CEFR level or placement decision. | A local CEFR scale or standard-setting algorithm. |
| What makes an assessment result defensible? | AERA, APA, and NCME Standards for Educational and Psychological Testing | Keep task, rubric, criterion, anchor, evidence, rater, prompt, workflow, parse, verifier, and replay provenance explicit. | Calling provider success or a model label validity evidence. |
| How should provider calls be governed? | Existing `TaskOrchestrator`, KV credential registry, model discovery, and structured-output passthrough | Reuse the existing gateway; require exact external contract compatibility and fail closed on missing capability or provider failure. | Direct provider SDK calls or a second credential path. |
+| What proves a discovered configured-gateway chat row is usable? | The existing provider-error boundary distinguishes catalog metadata from runtime success. | Require one bounded synthetic structured-output probe before activation and share missing-model exclusions across the full virtual request. | Treating list membership as readiness or retrying the same missing model in later workflow roles. |
Official references:
diff --git a/docs/planning/adrs/0011-provider-error-boundary.md b/docs/planning/adrs/0011-provider-error-boundary.md
index 4c6885412..155e1f561 100644
--- a/docs/planning/adrs/0011-provider-error-boundary.md
+++ b/docs/planning/adrs/0011-provider-error-boundary.md
@@ -29,6 +29,14 @@ available through a public gateway error or an exception cause.
request failed` / `provider batch request failed`) because a stream may
already have emitted bytes (no retry, no failover), and raw connection
resets outside ``URLError`` map to the same stable ``transport_error`` code.
+7. A virtual structured request keeps one request-scoped set of models proven
+ missing. Evidence roles and final synthesis share it, so each missing model
+ is attempted at most once and complete exhaustion terminates with the typed
+ `model_not_found` surface rather than restarting the same catalog sequence.
+8. Configured-gateway readiness probes use a distinct capability-probe
+ telemetry operation. They are not caller attempts. Conversely, an explicit
+ structured model pin constrains evidence, model judgment, and synthesis to
+ that same agent; only virtual selectors may replace a missing model.
## Consequences
@@ -41,6 +49,7 @@ responses have a wider audience than provider credentials and request data.
## Verification
`tests/test_model_discovery.py`, `tests/test_provider_reliability.py`,
+`tests/test_auto_discovery_server.py`, `tests/test_chat_response_format_http_honesty.py`,
`tests/test_true_streaming.py`, and `tests/test_model_judge.py` assert that
provider response text is absent from public messages and causes, while the
full suite must remain green before merge.
diff --git a/docs/planning/adrs/0015-durable-provider-catalog.md b/docs/planning/adrs/0015-durable-provider-catalog.md
index 16a3e5cbd..a5dc4b33f 100644
--- a/docs/planning/adrs/0015-durable-provider-catalog.md
+++ b/docs/planning/adrs/0015-durable-provider-catalog.md
@@ -64,10 +64,15 @@ realtime transports. They are never used to infer reasoning, verification,
coding, vision, or provider-native effort capabilities. Those require explicit
catalog or measured evidence under the gateway-owned policy.
Provider-declared `supported_parameters=response_format` is retained as the
-`response_format` serving tag. Virtual structured synthesis requires that tag
-when selecting from an automatically discovered pool; unknown support does not
-become presumed support. An explicitly requested operator-managed model remains
-the operator's transport contract.
+`response_format` serving tag and preferred for virtual structured synthesis.
+Missing catalog metadata does not itself prove support; configured-gateway rows
+must instead pass the runtime probe below. An explicitly requested
+operator-managed model remains the operator's transport contract.
+Configured-gateway runtime activation additionally performs one bounded,
+synthetic `json_object` probe per discovered chat row. Only rows that return a
+valid structured object remain chat-serving candidates; a listing alone is not
+readiness evidence. Embedding rows retain their separate capability route and
+are never subjected to a chat probe.
When live discovery activates a real chat model, only agents explicitly tagged
`bootstrap_seed` are retired. A `mock://` transport alone is not proof that an
agent is disposable; operator-configured mock agents remain in the declared
@@ -98,7 +103,8 @@ equivalent machine-readable balance contract.
The merge gate covers normalized DDL, secret-column absence, provider-account
isolation, parameterized PostgreSQL statements, last-known-good retention,
withdrawal after authoritative success, non-chat filtering, secret-free
-evidence, and end-to-end recovery when one provider fails.
+evidence, bounded configured-gateway capability probes, and end-to-end recovery
+when one provider fails.
## References
diff --git a/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md
new file mode 100644
index 000000000..d447a2a1c
--- /dev/null
+++ b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md
@@ -0,0 +1,62 @@
+# ADR 0039: Request-scoped configured endpoint routing
+
+- Status: Accepted
+- Date: 2026-09-01
+
+## Context
+
+Chat Completions and Responses may route across agents discovered from several
+configured provider endpoints. A caller can have an authorization or data
+residency requirement to remain on one of those endpoints. A caller-supplied
+URL must never become a transport destination or alter global discovery.
+
+## Decision
+
+The optional `routing.endpoint` field is accepted only by Chat Completions and
+Responses. It selects an absolute HTTP(S) endpoint already present in the
+enabled agent configuration; it is never used as a destination. Endpoint
+identity canonicalizes scheme, host, default port, and path, treating one
+terminal `/v1` as a transport suffix. Credentials, query strings, fragments,
+and non-HTTP(S) schemes are rejected.
+
+The matched agent identifiers form a request-local candidate set for planning,
+worker, verifier, judge, synthesizer, retry, failover, proxy, and streaming
+paths. Generated plans are checked against the same set, and routing caches are
+partitioned by a non-reversible digest of the endpoint identity. The constraint
+uses context-local state and never mutates the discovered agent pool.
+
+An unmatched selector or a conflict with an explicit concrete model fails with
+OpenAI-shaped HTTP 400 code `endpoint_unavailable`. Endpoint-bound work remains
+synchronous because a deferred job cannot retain this request-local boundary.
+The `routing` extension is removed before provider transport. Omitting the
+field preserves automatic discovery and routing; unsupported surfaces continue
+to reject the key.
+
+## Consequences
+
+Deployments can constrain one request without creating arbitrary outbound
+request capability or losing provider discovery. Tests use synthetic endpoint
+names; runtime endpoint selectors and credentials do not enter repository
+artifacts.
+
+## References
+
+NIST AI 600-1 treats third-party model and data-flow controls as explicit risk
+management boundaries; that supports selecting only an operator-configured
+endpoint, not converting caller input into a new destination. The Conductor and
+TRINITY papers describe orchestration across specialized agents, which motivates
+applying the same request constraint to every role and retry path. Redistribution
+permission for these publications has not been established, so this ADR links
+and summarizes them instead of vendoring PDFs.
+
+National Institute of Standards and Technology. (2024). *Artificial
+intelligence risk management framework: Generative artificial intelligence
+profile* (NIST AI 600-1). https://doi.org/10.6028/NIST.AI.600-1
+
+Nielsen, I., Motwani, S., Guan, Y., et al. (2025). *Learning to orchestrate
+agents in natural language with the conductor* [Preprint]. arXiv.
+https://arxiv.org/abs/2512.04388
+
+Xu, Z., Zhou, K., Shek, T., et al. (2025). *TRINITY: An evolved LLM
+coordinator for complex real-world workflows* [Preprint]. arXiv.
+https://arxiv.org/abs/2512.04695
diff --git a/docs/product-technical-gap-baseline.md b/docs/product-technical-gap-baseline.md
index 76c4346c0..38f45756b 100644
--- a/docs/product-technical-gap-baseline.md
+++ b/docs/product-technical-gap-baseline.md
@@ -2368,10 +2368,9 @@ that exact combination unconditionally, before any upstream call.
codebase has never sent `tools` + `stream=true` to a real provider in the
first place; `proxy_completion()` always forces `upstream["stream"] = False`
for tool-calling requests. `_chat_response_sse_chunks` (the SSE framing
-function this passthrough already calls) already had full, independently
-tested support for both tool_calls delta framing and honest usage-chunk
-emission (`usage_source: "reported"` vs `"estimated"`) — the only thing
-missing was reachability.
+function this passthrough already calls) had independently tested tool_calls
+delta framing. ADR 0006 subsequently removed its estimated-usage fallback:
+valid provider usage is reported and missing usage is explicitly unavailable.
**Fix**: [PR #925](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/925)
narrowed the rejection from *all* `tools`/`response_format` structured
@@ -2382,9 +2381,8 @@ fix, PR #924, took that broader approach and was closed in favor of #925):
traced `cost_router.py:456-507` — a conduct-mode multi-step workflow's
`result["usage"]` dict is built by summing per-step counts and is *always*
populated, but carries no `usage_source`/`measurement_status` tag of its own;
-when a step's provider response omits usage, the sum silently includes a
-local token-count estimate (the `measurement_status="estimated"` case, which
-only survives on the sibling `cost` key, never on `usage` itself).
+when a step's provider response omitted usage, the historical sum silently
+included a local token-count estimate.
`_chat_response_sse_chunks` labels any populated `usage` dict
`"usage_source": "reported"` unconditionally — it does not check
`payload["cost"]["measurement_status"]`, unlike the sibling
@@ -2397,10 +2395,9 @@ of fabricated-precision the project's Honest metrics convention exists to
prevent. `tools` passthrough doesn't have this exposure (always one
non-streaming upstream call; `payload["usage"]` there is always the raw
provider JSON's own field), which is also exactly the shape Strix needs.
-**Follow-up (not yet scheduled)**: teach `_chat_response_sse_chunks` the same
-`measurement_status == "measured"` gate `chat_completion_chunks` already has,
-so conduct-mode streamed usage can be exposed honestly too, instead of kept
-closed.
+**Follow-up delivered by ADR 0006**: `_chat_response_sse_chunks` now emits
+measured usage only for valid provider counts and otherwise emits null usage
+with an unavailable status rather than synthesizing the historical estimate.
**Duplicate-work consolidation** (concurrent autonomous sessions independently
converged on the same bug): closed contextual-orchestrator#924 (superseded by
diff --git a/fuzz/fuzz_endpoint_selector.py b/fuzz/fuzz_endpoint_selector.py
new file mode 100644
index 000000000..8beefeb0c
--- /dev/null
+++ b/fuzz/fuzz_endpoint_selector.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+"""Atheris harness for the request endpoint selector parser."""
+
+import sys
+from pathlib import Path
+
+import atheris
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+with atheris.instrument_imports():
+ from fuzz.targets import exercise_endpoint_selector
+
+
+def one_input(data: bytes) -> None:
+ """Feed arbitrary Unicode endpoint selectors through the shared invariant."""
+ fdp = atheris.FuzzedDataProvider(data)
+ exercise_endpoint_selector(fdp.ConsumeUnicodeNoSurrogates(fdp.remaining_bytes()))
+
+
+def main() -> None:
+ """Run the coverage-guided endpoint parser target."""
+ atheris.Setup(sys.argv, one_input)
+ atheris.Fuzz()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/fuzz/targets.py b/fuzz/targets.py
index 3cc3641e3..c8010a2f8 100644
--- a/fuzz/targets.py
+++ b/fuzz/targets.py
@@ -55,6 +55,7 @@
_parse_openai_compatible,
)
from contextual_orchestrator.orchestrator import (
+ EndpointUnavailableError,
ModelAgent,
TaskOrchestrator,
_parse_model_judge_reply,
@@ -62,6 +63,7 @@
chat_completion_chunks,
redact_text,
redact_value,
+ normalize_endpoint_selector,
sse_stream_body,
)
from contextual_orchestrator.pii_protection import PiiProtectionError, _decode_secret
@@ -78,6 +80,16 @@
# raise; everything else below is a legitimate stdlib decode/parse failure.
RequestError = server.RequestError
+
+def exercise_endpoint_selector(value: str) -> None:
+ """Normalize arbitrary endpoint text or reject it with the stable error."""
+ try:
+ normalized = normalize_endpoint_selector(value)
+ except EndpointUnavailableError:
+ return
+ assert normalized == normalize_endpoint_selector(normalized)
+ assert normalized.startswith(("http://", "https://"))
+
# Malformed bytes/JSON must surface only as these.
_EXPECTED_BODY_EXC = (
RequestError,
diff --git a/rust/.gitignore b/rust/.gitignore
new file mode 100644
index 000000000..b83d22266
--- /dev/null
+++ b/rust/.gitignore
@@ -0,0 +1 @@
+/target/
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
new file mode 100644
index 000000000..96d3c6a1d
--- /dev/null
+++ b/rust/Cargo.lock
@@ -0,0 +1,373 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "aho-corasick"
+version = "1.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "anyhow"
+version = "1.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470"
+
+[[package]]
+name = "base64"
+version = "0.22.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
+
+[[package]]
+name = "bit-set"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1"
+dependencies = [
+ "bit-vec",
+]
+
+[[package]]
+name = "bit-vec"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb"
+
+[[package]]
+name = "bstr"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f"
+dependencies = [
+ "memchr",
+ "regex-automata",
+ "serde_core",
+]
+
+[[package]]
+name = "contextual-token-packer"
+version = "0.1.0"
+dependencies = [
+ "pyo3",
+ "rayon",
+ "serde",
+ "serde_json",
+ "tiktoken-rs",
+]
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
+
+[[package]]
+name = "either"
+version = "1.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34"
+
+[[package]]
+name = "fancy-regex"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2"
+dependencies = [
+ "bit-set",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "itoa"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
+
+[[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+
+[[package]]
+name = "libc"
+version = "0.2.189"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2"
+
+[[package]]
+name = "memchr"
+version = "2.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98"
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.29.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.29.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.29.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.29.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.29.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "rayon"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "regex"
+version = "1.13.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
+
+[[package]]
+name = "rustc-hash"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
+
+[[package]]
+name = "serde"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba"
+dependencies = [
+ "serde_core",
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_core"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.229"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.4",
+]
+
+[[package]]
+name = "serde_json"
+version = "1.0.151"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14"
+dependencies = [
+ "itoa",
+ "memchr",
+ "serde",
+ "serde_core",
+ "zmij",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.119"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
+[[package]]
+name = "tiktoken-rs"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625"
+dependencies = [
+ "anyhow",
+ "base64",
+ "bstr",
+ "fancy-regex",
+ "lazy_static",
+ "regex",
+ "rustc-hash",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
+
+[[package]]
+name = "zmij"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b"
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
new file mode 100644
index 000000000..b7d83efbb
--- /dev/null
+++ b/rust/Cargo.toml
@@ -0,0 +1,3 @@
+[workspace]
+members = ["token_counter"]
+resolver = "2"
diff --git a/rust/token_counter/.gitignore b/rust/token_counter/.gitignore
new file mode 100644
index 000000000..b83d22266
--- /dev/null
+++ b/rust/token_counter/.gitignore
@@ -0,0 +1 @@
+/target/
diff --git a/rust/token_counter/Cargo.toml b/rust/token_counter/Cargo.toml
new file mode 100644
index 000000000..b31d4b976
--- /dev/null
+++ b/rust/token_counter/Cargo.toml
@@ -0,0 +1,15 @@
+[package]
+name = "contextual-token-packer"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+pyo3 = { version = "0.29", features = ["abi3-py310", "auto-initialize"] }
+rayon = "1.10"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+tiktoken-rs = "0.7"
+
+[lib]
+name = "_token_packer"
+crate-type = ["cdylib"]
diff --git a/rust/token_counter/pyproject.toml b/rust/token_counter/pyproject.toml
new file mode 100644
index 000000000..981ca701c
--- /dev/null
+++ b/rust/token_counter/pyproject.toml
@@ -0,0 +1,13 @@
+[build-system]
+requires = ["maturin>=1.8,<2"]
+build-backend = "maturin"
+
+[project]
+name = "contextual-token-packer"
+version = "0.1.0"
+requires-python = ">=3.10"
+
+[tool.maturin]
+module-name = "contextual_orchestrator._token_packer"
+python-source = "../.."
+features = ["pyo3/extension-module"]
diff --git a/rust/token_counter/src/lib.rs b/rust/token_counter/src/lib.rs
new file mode 100644
index 000000000..8ebf18875
--- /dev/null
+++ b/rust/token_counter/src/lib.rs
@@ -0,0 +1,441 @@
+//! Exact token counting, cl100k child chunking, and provider-request packing.
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::*;
+use rayon::prelude::*;
+use std::sync::OnceLock;
+use tiktoken_rs::{cl100k_base, o200k_base};
+
+static CL100K: OnceLock = OnceLock::new();
+static O200K: OnceLock = OnceLock::new();
+
+fn exact_tokenizer() -> PyResult<&'static tiktoken_rs::CoreBPE> {
+ if let Some(tokenizer) = CL100K.get() {
+ return Ok(tokenizer);
+ }
+ let tokenizer = cl100k_base().map_err(|_| PyValueError::new_err("cl100k unavailable"))?;
+ let _ = CL100K.set(tokenizer);
+ CL100K
+ .get()
+ .ok_or_else(|| PyValueError::new_err("cl100k initialization failed"))
+}
+
+fn o200k_tokenizer() -> PyResult<&'static tiktoken_rs::CoreBPE> {
+ if let Some(tokenizer) = O200K.get() {
+ return Ok(tokenizer);
+ }
+ let tokenizer = o200k_base().map_err(|_| PyValueError::new_err("o200k unavailable"))?;
+ let _ = O200K.set(tokenizer);
+ O200K
+ .get()
+ .ok_or_else(|| PyValueError::new_err("o200k initialization failed"))
+}
+
+#[pyclass(get_all, skip_from_py_object)]
+#[derive(Clone)]
+struct PackedPart {
+ source_index: usize,
+ part_index: usize,
+ part_count: usize,
+ token_start: usize,
+ token_end: usize,
+ token_count: usize,
+ text: String,
+}
+
+fn checked_token_total(current: usize, additional: usize) -> PyResult {
+ current
+ .checked_add(additional)
+ .ok_or_else(|| PyValueError::new_err("request token total overflow"))
+}
+
+fn utf8_token_ranges(
+ tokenizer: &tiktoken_rs::CoreBPE,
+ tokens: &[u32],
+ max_tokens_per_input: usize,
+) -> PyResult> {
+ let mut ranges = Vec::new();
+ let mut start = 0usize;
+ while start < tokens.len() {
+ let ceiling = start
+ .checked_add(max_tokens_per_input)
+ .map(|value| value.min(tokens.len()))
+ .ok_or_else(|| PyValueError::new_err("token range overflow"))?;
+ let mut end = ceiling;
+ let decoded = loop {
+ if end == start {
+ return Err(PyValueError::new_err(
+ "token limit cannot preserve a complete UTF-8 scalar",
+ ));
+ }
+ match tokenizer.decode(tokens[start..end].to_vec()) {
+ Ok(text) => break text,
+ Err(_) => end -= 1,
+ }
+ };
+ ranges.push((start, end, decoded));
+ start = end;
+ }
+ Ok(ranges)
+}
+
+#[pyfunction]
+fn sum_token_counts(values: Vec) -> PyResult {
+ values.into_iter().try_fold(0usize, checked_token_total)
+}
+
+#[pyfunction]
+fn count_cl100k(text: &str) -> PyResult {
+ let tokenizer = exact_tokenizer()?;
+ Ok(tokenizer.encode_ordinary(text).len())
+}
+
+#[pyfunction]
+fn count_o200k(text: &str) -> PyResult {
+ let tokenizer = o200k_tokenizer()?;
+ Ok(tokenizer.encode_ordinary(text).len())
+}
+
+#[pyfunction]
+fn cosine_similarity(vector_a: Vec, vector_b: Vec) -> PyResult> {
+ if vector_a.len() != vector_b.len() || vector_a.is_empty() {
+ return Ok(None);
+ }
+ if vector_a
+ .iter()
+ .chain(vector_b.iter())
+ .any(|value| !value.is_finite())
+ {
+ return Err(PyValueError::new_err("cosine inputs must be finite"));
+ }
+ let dot = vector_a
+ .iter()
+ .zip(vector_b.iter())
+ .map(|(a, b)| a * b)
+ .sum::();
+ let norm_a = vector_a
+ .iter()
+ .map(|value| value * value)
+ .sum::()
+ .sqrt();
+ let norm_b = vector_b
+ .iter()
+ .map(|value| value * value)
+ .sum::()
+ .sqrt();
+ if norm_a == 0.0 || norm_b == 0.0 {
+ return Ok(None);
+ }
+ let value = dot / (norm_a * norm_b);
+ if value.is_finite() {
+ Ok(Some(value))
+ } else {
+ Err(PyValueError::new_err("cosine result is not finite"))
+ }
+}
+
+#[pyfunction]
+fn root_mean_square_error(estimates: Vec, truths: Vec) -> PyResult {
+ if estimates.len() != truths.len() || estimates.is_empty() {
+ return Err(PyValueError::new_err(
+ "RMSE inputs must have one shared positive length",
+ ));
+ }
+ if estimates
+ .iter()
+ .chain(truths.iter())
+ .any(|value| !value.is_finite())
+ {
+ return Err(PyValueError::new_err("RMSE inputs must be finite"));
+ }
+ let mean_square = estimates
+ .iter()
+ .zip(truths.iter())
+ .map(|(estimate, truth)| (estimate - truth).powi(2))
+ .sum::()
+ / estimates.len() as f64;
+ let value = mean_square.sqrt();
+ if value.is_finite() {
+ Ok(value)
+ } else {
+ Err(PyValueError::new_err("RMSE result is not finite"))
+ }
+}
+
+#[pyfunction]
+fn weighted_average_embeddings(parts: Vec<(Vec, usize)>) -> PyResult> {
+ if parts.is_empty() {
+ return Ok(Vec::new());
+ }
+ let dimension = parts[0].0.len();
+ if dimension == 0 || parts.iter().any(|(vector, _)| vector.len() != dimension) {
+ return Err(PyValueError::new_err(
+ "embedding vectors must have one shared positive dimension",
+ ));
+ }
+ if parts.iter().any(|(_, weight)| *weight == 0) {
+ return Err(PyValueError::new_err(
+ "embedding token weights must be positive",
+ ));
+ }
+ let weights: Vec = parts.iter().map(|(_, weight)| *weight).collect();
+ let total_weight = sum_token_counts(weights.clone())?;
+ (0..dimension)
+ .map(|offset| {
+ let weighted_sum = parts
+ .iter()
+ .zip(weights.iter())
+ .map(|((vector, _weight), weight)| vector[offset] * (*weight as f64))
+ .sum::();
+ let reduced = weighted_sum / (total_weight as f64);
+ if reduced.is_finite() {
+ Ok((reduced * 100_000_000.0).round() / 100_000_000.0)
+ } else {
+ Err(PyValueError::new_err("embedding reduction is not finite"))
+ }
+ })
+ .collect()
+}
+
+#[pyfunction]
+fn pack_cl100k(
+ texts: Vec,
+ max_tokens_per_input: usize,
+ max_inputs: usize,
+ max_total_tokens: usize,
+) -> PyResult<(Vec, Vec>)> {
+ if max_tokens_per_input == 0 || max_inputs == 0 || max_total_tokens == 0 {
+ return Err(PyValueError::new_err("limits must be positive"));
+ }
+ if texts.iter().any(String::is_empty) {
+ return Err(PyValueError::new_err("embedding input must be non-empty"));
+ }
+ let tokenizer = exact_tokenizer()?;
+ let encoded: Vec> = texts
+ .par_iter()
+ .map(|text| tokenizer.encode_ordinary(text))
+ .collect();
+ let mut parts = Vec::new();
+ for (source_index, tokens) in encoded.iter().enumerate() {
+ let ranges = utf8_token_ranges(
+ &tokenizer,
+ tokens,
+ max_tokens_per_input.min(max_total_tokens),
+ )?;
+ let part_count = ranges.len();
+ for (part_index, (token_start, token_end, text)) in ranges.into_iter().enumerate() {
+ parts.push(PackedPart {
+ source_index,
+ part_index,
+ part_count,
+ token_start,
+ token_end,
+ token_count: token_end - token_start,
+ text,
+ });
+ }
+ }
+ let mut shards = Vec::new();
+ let mut current = Vec::new();
+ let mut current_tokens = 0usize;
+ for (index, part) in parts.iter().enumerate() {
+ let next = checked_token_total(current_tokens, part.token_count)?;
+ if !current.is_empty() && (current.len() >= max_inputs || next > max_total_tokens) {
+ shards.push(std::mem::take(&mut current));
+ current_tokens = 0;
+ }
+ current_tokens = checked_token_total(current_tokens, part.token_count)?;
+ current.push(index);
+ }
+ if !current.is_empty() {
+ shards.push(current);
+ }
+ Ok((parts, shards))
+}
+
+#[pymodule]
+fn _token_packer(module: &Bound<'_, PyModule>) -> PyResult<()> {
+ module.add_class::()?;
+ module.add_function(wrap_pyfunction!(pack_cl100k, module)?)?;
+ module.add_function(wrap_pyfunction!(sum_token_counts, module)?)?;
+ module.add_function(wrap_pyfunction!(count_cl100k, module)?)?;
+ module.add_function(wrap_pyfunction!(count_o200k, module)?)?;
+ module.add_function(wrap_pyfunction!(cosine_similarity, module)?)?;
+ module.add_function(wrap_pyfunction!(root_mean_square_error, module)?)?;
+ module.add_function(wrap_pyfunction!(weighted_average_embeddings, module)?)?;
+ Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn exact_token_text(token_count: usize) -> String {
+ let tokenizer = cl100k_base().unwrap();
+ let unit = tokenizer.encode_ordinary(" x");
+ assert_eq!(unit.len(), 1);
+ tokenizer.decode(vec![unit[0]; token_count]).unwrap()
+ }
+
+ #[test]
+ fn utf8_korean_emoji_combining_and_order_are_preserved() {
+ Python::attach(|_| {
+ let inputs = vec!["한국어🙂e\u{301}".into(), "두 번째🙂".into()];
+ let (parts, shards) = pack_cl100k(inputs.clone(), 8192, 2048, 300_000).unwrap();
+ assert!(parts.iter().all(|part| part.token_count <= 8192));
+ assert_eq!(parts.last().unwrap().source_index, 1);
+ assert_eq!(
+ parts
+ .iter()
+ .map(|part| part.text.as_str())
+ .collect::>(),
+ inputs.iter().map(String::as_str).collect::>()
+ );
+ assert_eq!(shards.iter().flatten().count(), parts.len());
+ });
+ }
+
+ #[test]
+ fn utf8_boundary_retreats_to_a_complete_scalar_or_fails_closed() {
+ Python::attach(|_| {
+ let input = format!("{}🙂", exact_token_text(8191));
+ let tokenizer = cl100k_base().unwrap();
+ assert_eq!(tokenizer.encode_ordinary(&input).len(), 8193);
+ let (parts, _) = pack_cl100k(vec![input.clone()], 8192, 2048, 300_000).unwrap();
+ assert_eq!(
+ parts
+ .iter()
+ .map(|part| part.text.as_str())
+ .collect::(),
+ input
+ );
+ assert_eq!(
+ parts
+ .iter()
+ .map(|part| part.token_count)
+ .collect::>(),
+ vec![8191, 2]
+ );
+ assert!(pack_cl100k(vec!["🙂".into()], 1, 2048, 300_000).is_err());
+ });
+ }
+
+ #[test]
+ fn rejects_empty_input_and_nonpositive_limits() {
+ Python::attach(|_| {
+ assert!(pack_cl100k(vec![String::new()], 8192, 2048, 300_000).is_err());
+ assert!(pack_cl100k(vec!["x".into()], 0, 2048, 300_000).is_err());
+ });
+ }
+
+ #[test]
+ fn per_input_boundary_is_exact_at_8192_and_8193() {
+ Python::attach(|_| {
+ let (exact, _) =
+ pack_cl100k(vec![exact_token_text(8192)], 8192, 2048, 300_000).unwrap();
+ assert_eq!(
+ exact
+ .iter()
+ .map(|part| part.token_count)
+ .collect::>(),
+ vec![8192]
+ );
+ let (over, _) = pack_cl100k(vec![exact_token_text(8193)], 8192, 2048, 300_000).unwrap();
+ assert_eq!(
+ over.iter().map(|part| part.token_count).collect::>(),
+ vec![8192, 1]
+ );
+ assert_eq!(
+ (
+ over[0].token_start,
+ over[0].token_end,
+ over[1].token_start,
+ over[1].token_end
+ ),
+ (0, 8192, 8192, 8193)
+ );
+ });
+ }
+
+ #[test]
+ fn total_token_boundary_is_exact_at_300000_and_300001() {
+ Python::attach(|_| {
+ let inputs = vec![exact_token_text(7500); 40];
+ let (_, exact) = pack_cl100k(inputs, 8192, 2048, 300_000).unwrap();
+ assert_eq!(exact.len(), 1);
+ assert_eq!(exact[0].len(), 40);
+ let mut over_inputs = vec![exact_token_text(7500); 40];
+ over_inputs.push(exact_token_text(1));
+ let (_, over) = pack_cl100k(over_inputs, 8192, 2048, 300_000).unwrap();
+ assert_eq!(over.iter().map(Vec::len).collect::>(), vec![40, 1]);
+ });
+ }
+
+ #[test]
+ fn total_token_limit_also_bounds_each_part() {
+ Python::attach(|_| {
+ let (parts, shards) = pack_cl100k(vec![exact_token_text(9)], 10, 2, 8).unwrap();
+ assert_eq!(
+ parts
+ .iter()
+ .map(|part| part.token_count)
+ .collect::>(),
+ vec![8, 1]
+ );
+ assert_eq!(shards.iter().map(Vec::len).collect::>(), vec![1, 1]);
+ });
+ }
+
+ #[test]
+ fn input_count_boundary_is_exact_at_2048_and_2049() {
+ Python::attach(|_| {
+ let (_, exact) = pack_cl100k(vec!["x".into(); 2048], 8192, 2048, 300_000).unwrap();
+ assert_eq!(exact.iter().map(Vec::len).collect::>(), vec![2048]);
+ let (_, over) = pack_cl100k(vec!["x".into(); 2049], 8192, 2048, 300_000).unwrap();
+ assert_eq!(over.iter().map(Vec::len).collect::>(), vec![2048, 1]);
+ });
+ }
+
+ #[test]
+ fn checked_total_fails_closed_on_integer_overflow() {
+ Python::attach(|_| {
+ assert!(checked_token_total(usize::MAX, 1).is_err());
+ assert!(sum_token_counts(vec![usize::MAX, 1]).is_err());
+ });
+ }
+
+ #[test]
+ fn vector_reduction_and_token_sum_are_rust_owned() {
+ Python::attach(|_| {
+ assert_eq!(sum_token_counts(vec![2, 3, 5]).unwrap(), 10);
+ assert_eq!(
+ weighted_average_embeddings(vec![(vec![1.0, 3.0], 1), (vec![3.0, 5.0], 3)])
+ .unwrap(),
+ vec![2.5, 4.5]
+ );
+ assert!(weighted_average_embeddings(vec![(vec![f64::INFINITY], 1)]).is_err());
+ assert!(weighted_average_embeddings(vec![(vec![1.0], 0)]).is_err());
+ assert!(
+ weighted_average_embeddings(vec![(vec![1.0], 1), (vec![1.0, 2.0], 1)]).is_err()
+ );
+ assert!(weighted_average_embeddings(vec![(Vec::new(), 1)]).is_err());
+ });
+ }
+
+ #[test]
+ fn exact_count_cosine_and_rmse_are_rust_owned() {
+ Python::attach(|_| {
+ assert_eq!(count_cl100k("hello world").unwrap(), 2);
+ assert_eq!(count_o200k("hello world").unwrap(), 2);
+ assert_eq!(
+ cosine_similarity(vec![1.0, 0.0], vec![1.0, 0.0]).unwrap(),
+ Some(1.0)
+ );
+ assert_eq!(cosine_similarity(vec![0.0], vec![1.0]).unwrap(), None);
+ assert_eq!(
+ root_mean_square_error(vec![1.0, 3.0], vec![1.0, 1.0]).unwrap(),
+ 2.0_f64.sqrt()
+ );
+ assert!(root_mean_square_error(Vec::new(), Vec::new()).is_err());
+ });
+ }
+}
diff --git a/tests/fuzz/test_fuzz_properties.py b/tests/fuzz/test_fuzz_properties.py
index e58858b58..2758c42b2 100644
--- a/tests/fuzz/test_fuzz_properties.py
+++ b/tests/fuzz/test_fuzz_properties.py
@@ -18,6 +18,7 @@
from fuzz.targets import (
exercise_agent_config,
+ exercise_endpoint_selector,
exercise_model_judge_reply,
exercise_models_dev_cost,
exercise_orchestration,
@@ -85,6 +86,12 @@ def test_agent_config_parser(value: object) -> None:
exercise_agent_config(value)
+@_SETTINGS
+@given(st.text(max_size=4096))
+def test_endpoint_selector_normalization_is_stable(value: str) -> None:
+ exercise_endpoint_selector(value)
+
+
@_SETTINGS
@given(
st.builds(
diff --git a/tests/test_admin_spend_view.py b/tests/test_admin_spend_view.py
index a5e28d615..30123412e 100644
--- a/tests/test_admin_spend_view.py
+++ b/tests/test_admin_spend_view.py
@@ -22,7 +22,7 @@ def test_admin_state_includes_spend_block() -> None:
assert "spend" in state
spend = state["spend"]
- assert spend["measurement_status"] == "local_runtime_estimate"
+ assert spend["measurement_status"] == "unavailable"
assert spend["totals"]["run_count"] == 1
assert isinstance(spend["by_model"], list) and spend["by_model"]
assert spend["by_model"][0]["model"] == "priced-model"
diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py
index 4bacc5072..7beb3d698 100644
--- a/tests/test_api_contract.py
+++ b/tests/test_api_contract.py
@@ -58,6 +58,22 @@ def test_openapi_documents_compatibility_front_door() -> None:
"content"
]["application/json"]["schema"]
assert chat_schema["properties"]["include_orchestration_trace"]["type"] == "boolean"
+ chat_response = OPENAPI_SPEC["components"]["schemas"]["ChatCompletionResponse"]
+ assert chat_response["properties"]["usage"]["$ref"].endswith("AuthoritativeUsage")
+ assert chat_response["properties"]["usage_measurement_status"]["enum"] == [
+ "measured",
+ "unavailable",
+ ]
+ measured, unavailable = chat_response["oneOf"]
+ assert measured["properties"]["usage_measurement_status"] == {"const": "measured"}
+ assert measured["properties"]["usage"]["required"] == [
+ "prompt_tokens",
+ "completion_tokens",
+ ]
+ assert unavailable["properties"]["usage_measurement_status"] == {
+ "const": "unavailable"
+ }
+ assert unavailable["properties"]["usage"] == {"type": "null"}
assert OPENAPI_SPEC["paths"]["/api/v1/access_reports/{workflow_run_id}"]["get"][
"security"
] == [{"admin_bearer_auth": [], "trace_bearer_auth": []}]
diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py
index af96ddb1d..031f4bb06 100644
--- a/tests/test_auto_discovery_server.py
+++ b/tests/test_auto_discovery_server.py
@@ -4,26 +4,36 @@
import os
from unittest.mock import patch
-from contextual_orchestrator.__main__ import _auto_discover_runtime_agents
-from contextual_orchestrator.model_discovery import DiscoveredModel
+import pytest
+
+from contextual_orchestrator.__main__ import (
+ _auto_discover_runtime_agents,
+ _probe_configured_gateway_structured_chat,
+ main,
+)
+from contextual_orchestrator.model_discovery import (
+ DiscoveredModel,
+ agent_from_discovered,
+ agent_id_for,
+)
from contextual_orchestrator.orchestrator import ModelAgent, TaskOrchestrator
-def test_auto_discovery_activates_only_chat_capable_agents(monkeypatch) -> None:
- """Startup routing excludes discovered deployments without chat evidence."""
+def test_auto_discovery_activates_chat_and_embedding_capable_agents(monkeypatch) -> None:
+ """Startup retains a provider-declared embedding route beside chat."""
chat = DiscoveredModel(
- provider_name="openai",
+ provider_name="configured_gateway",
model_id="chat-capable-model",
- credential_name="OPENAI_API_KEY",
- chat_base_url="https://api.openai.com/v1",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
auth_scheme="Bearer",
capabilities=("chat",),
)
embedding = DiscoveredModel(
- provider_name="openai",
+ provider_name="configured_gateway",
model_id="embedding-capable-model",
- credential_name="OPENAI_API_KEY",
- chat_base_url="https://api.openai.com/v1",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
auth_scheme="Bearer",
capabilities=("embedding",),
)
@@ -38,16 +48,351 @@ def test_auto_discovery_activates_only_chat_capable_agents(monkeypatch) -> None:
result = _auto_discover_runtime_agents(orchestrator)
- assert len(result["added"]) == 1
- agent = next(agent for agent in orchestrator.agents if agent.id == result["added"][0])
- assert agent.model == chat.model_id
- assert agent.disabled is False
- assert "chat" in agent.tags
- assert all(candidate.model != embedding.model_id for candidate in orchestrator.agents)
+ assert len(result["added"]) == 2
+ chat_agent = next(agent for agent in orchestrator.agents if agent.model == chat.model_id)
+ embedding_agent = next(agent for agent in orchestrator.agents if agent.model == embedding.model_id)
+ assert chat_agent.disabled is False
+ assert "chat" in chat_agent.tags
+ assert embedding_agent.disabled is False
+ assert "embedding" in embedding_agent.tags
+ assert orchestrator.select_capability_agent("embedding").id == embedding_agent.id
assert all(not candidate.base_url.startswith("mock://") for candidate in orchestrator.agents)
assert "bootstrap_agent" in result["updated"]
+def test_embedding_only_discovery_keeps_chat_fallbacks(monkeypatch) -> None:
+ embedding = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="text-embedding-only",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("embedding",),
+ spend_admitted=True,
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([embedding], []),
+ )
+ orchestrator = TaskOrchestrator(
+ [
+ ModelAgent(
+ "configured_gateway_placeholder",
+ "",
+ provider_name="configured_gateway",
+ base_url="https://gateway.synthetic.example/v1",
+ ),
+ ModelAgent(
+ "bootstrap_chat_agent",
+ "bootstrap-chat-model",
+ tags=("bootstrap_seed",),
+ ),
+ ]
+ )
+
+ _auto_discover_runtime_agents(orchestrator)
+
+ by_id = {agent.id: agent for agent in orchestrator.candidates}
+ assert "configured_gateway_placeholder" in by_id
+ assert by_id["bootstrap_chat_agent"].disabled is False
+
+
+def test_configured_gateway_discovery_retains_only_structured_probe_successes(
+ monkeypatch,
+) -> None:
+ """Configured-gateway chat rows require a successful bounded structured probe."""
+ models = [
+ DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id=model_id,
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ for model_id in ("stale-model", "live-model")
+ ]
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential",
+ lambda name: "present" if name == "LLM_GATEWAY_API_KEY" else None,
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: (models, []),
+ )
+ probes: list[str] = []
+
+ def probe(_orchestrator, model):
+ probes.append(model.model_id)
+ return model.model_id == "live-model"
+
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ probe,
+ )
+ orchestrator = TaskOrchestrator(
+ [ModelAgent("bootstrap_agent", "bootstrap-model", tags=("bootstrap_seed",))]
+ )
+
+ result = _auto_discover_runtime_agents(orchestrator)
+
+ assert probes == ["stale-model", "live-model"]
+ assert result["added"] == ["configured_gateway_live_model"]
+ assert all(agent.model != "stale-model" for agent in orchestrator.agents)
+
+
+def test_failed_gateway_catalog_probes_remove_unprobed_blank_seed(monkeypatch) -> None:
+ """A catalog-wide auth failure cannot fall through to the blank seed row."""
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="auth-failing-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ live_model = DiscoveredModel(
+ provider_name="openrouter",
+ model_id="synthetic-live-model",
+ credential_name="OPENROUTER_API_KEY",
+ chat_base_url="https://openrouter.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ seed = ModelAgent(
+ "configured_gateway_bootstrap",
+ "",
+ base_url="https://gateway.synthetic.example/v1",
+ provider_name="configured_gateway",
+ credential_key="LLM_GATEWAY_API_KEY",
+ tags=("bootstrap_seed",),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential", lambda _name: "present"
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([model, live_model], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: False,
+ )
+ orchestrator = TaskOrchestrator([seed])
+
+ _auto_discover_runtime_agents(orchestrator)
+
+ assert all(agent.id != seed.id for agent in orchestrator.candidates)
+ selected = orchestrator._select_agent("task", "synthesizer")
+ assert selected.provider_name == "openrouter"
+ assert selected.model == live_model.model_id
+
+
+def test_failed_gateway_catalog_probe_transiently_retires_the_only_blank_seed(
+ monkeypatch, tmp_path
+) -> None:
+ """A failed sole seed is uncallable now and can be reprobed after restart."""
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="auth-failing-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ seed = ModelAgent(
+ "configured_gateway_bootstrap",
+ "",
+ base_url="https://gateway.synthetic.example/v1",
+ provider_name="configured_gateway",
+ credential_key="LLM_GATEWAY_API_KEY",
+ tags=("bootstrap_seed",),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential", lambda _name: "present"
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([model], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: False,
+ )
+ agents_db = str(tmp_path / "agents.db")
+ orchestrator = TaskOrchestrator([seed], agents_db=agents_db)
+
+ result = _auto_discover_runtime_agents(orchestrator)
+
+ assert result["updated"] == []
+ assert orchestrator.agents == []
+ assert orchestrator.candidates == []
+ with pytest.raises(RuntimeError, match="no chat-compatible agent available"):
+ orchestrator._select_agent("task", "synthesizer")
+
+ restarted = TaskOrchestrator([seed], agents_db=agents_db)
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: True,
+ )
+
+ recovered = _auto_discover_runtime_agents(restarted)
+
+ assert recovered["added"] == [agent_id_for(model)]
+ assert restarted.agents[0].model == model.model_id
+
+
+def test_auto_discovery_restart_recovers_a_persisted_disabled_gateway_seed(
+ monkeypatch, tmp_path
+) -> None:
+ """A legacy failure tombstone cannot stop a later healthy discovery pass."""
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="recovered-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ seed = ModelAgent(
+ "configured_gateway_bootstrap",
+ "",
+ base_url="https://gateway.synthetic.example/v1",
+ provider_name="configured_gateway",
+ credential_key="LLM_GATEWAY_API_KEY",
+ tags=("bootstrap_seed",),
+ )
+ agents_db = str(tmp_path / "agents.db")
+ failed = TaskOrchestrator([seed], agents_db=agents_db)
+ failed.sync_discovered_agents([replace(seed, disabled=True)])
+ failed.close()
+
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.load_agents", lambda _path: [seed]
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential", lambda _name: "present"
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([model], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: True,
+ )
+ active_models = []
+
+ def complete(orchestrator, *_args, **_kwargs):
+ active_models.extend(agent.model for agent in orchestrator.agents)
+ return {"answer": "ok"}
+
+ monkeypatch.setattr(TaskOrchestrator, "complete", complete)
+
+ main([
+ "recover",
+ "--agents-db", agents_db,
+ "--auto-discover-model-agents",
+ ])
+
+ assert active_models == [model.model_id]
+
+
+def test_failed_gateway_probe_disables_persisted_discovered_agent_after_restart(
+ monkeypatch, tmp_path
+) -> None:
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="stale-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ existing = replace(agent_from_discovered(model), disabled=False)
+ agents_db = str(tmp_path / "agents.db")
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential", lambda _name: "present"
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([model], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: False,
+ )
+ orchestrator = TaskOrchestrator([existing], agents_db=agents_db)
+
+ _auto_discover_runtime_agents(orchestrator)
+ restarted = TaskOrchestrator([], agents_db=agents_db, allow_empty_agents=True)
+
+ assert restarted.candidates[0].disabled is True
+ assert "structured:blocked" in restarted.candidates[0].tags
+
+
+def test_failed_gateway_probe_keeps_persisted_embedding_capability(
+ monkeypatch, tmp_path
+) -> None:
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="mixed-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat", "embedding"),
+ )
+ mixed = replace(agent_from_discovered(model), disabled=False, priority=100)
+ fallback = ModelAgent("fallback_agent", "fallback-chat-model")
+ agents_db = str(tmp_path / "agents.db")
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.get_credential", lambda _name: "present"
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([model], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat",
+ lambda *_args: False,
+ )
+ orchestrator = TaskOrchestrator([fallback, mixed], agents_db=agents_db)
+
+ _auto_discover_runtime_agents(orchestrator)
+ restarted = TaskOrchestrator([fallback], agents_db=agents_db)
+
+ persisted = next(agent for agent in restarted.candidates if agent.id == mixed.id)
+ assert persisted.disabled is False
+ assert "structured:blocked" in persisted.tags
+ assert restarted.select_capability_agent("embedding").id == mixed.id
+ assert restarted._select_agent("task", "synthesizer").id == fallback.id
+
+
+def test_configured_gateway_structured_probe_is_bounded_and_validates_output() -> None:
+ """The startup probe proves JSON object service with a bounded synthetic call."""
+ model = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="candidate-model",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("chat",),
+ )
+ orchestrator = TaskOrchestrator([], allow_empty_agents=True)
+ observed = {}
+
+ def send(agent, payload):
+ observed.update(agent=agent, payload=payload)
+ return {"choices": [{"message": {"content": '{"status":"ok"}'}}]}
+
+ orchestrator.client.probe_structured_chat = send
+
+ assert _probe_configured_gateway_structured_chat(orchestrator, model) is True
+ assert observed["payload"]["max_tokens"] == 8
+ assert observed["payload"]["response_format"] == {"type": "json_object"}
+ assert "json" in observed["payload"]["messages"][0]["content"].casefold()
+
+
def test_auto_discovery_activates_a_free_vision_model_but_free_pool_excludes_it(
monkeypatch,
) -> None:
@@ -193,6 +538,30 @@ def test_auto_discovery_never_activates_openrouter_evidence_rows(monkeypatch) ->
assert [agent.model for agent in orchestrator.agents] == ["bootstrap-model"]
+def test_auto_discovery_never_activates_evidence_only_embedding_rows(monkeypatch) -> None:
+ """Embedding capability cannot bypass the evidence-only serving boundary."""
+ evidence = DiscoveredModel(
+ provider_name="configured_gateway",
+ model_id="provider/evidence-embedding",
+ credential_name="LLM_GATEWAY_API_KEY",
+ chat_base_url="https://gateway.synthetic.example/v1",
+ auth_scheme="Bearer",
+ capabilities=("embedding",),
+ evidence_only=True,
+ spend_admitted=True,
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.__main__.discover_all_models",
+ lambda *_args, **_kwargs: ([evidence], []),
+ )
+ orchestrator = TaskOrchestrator(
+ [ModelAgent("bootstrap_agent", "bootstrap-model", tags=("bootstrap_seed",))]
+ )
+
+ assert _auto_discover_runtime_agents(orchestrator) == {"added": [], "updated": []}
+ assert [agent.model for agent in orchestrator.agents] == ["bootstrap-model"]
+
+
def test_auto_discovery_keeps_metadata_free_general_chat_models(monkeypatch) -> None:
"""OpenAI-style model rows without capability metadata remain discoverable."""
discovered = DiscoveredModel(
@@ -307,8 +676,8 @@ def test_auto_discovery_keeps_last_enabled_placeholder_for_disabled_gateway_mode
assert orchestrator.agents == [placeholder]
-def test_auto_discovery_leaves_pool_unchanged_without_chat_capability_evidence(monkeypatch) -> None:
- """Startup fails closed without taking down an explicitly configured pool."""
+def test_auto_discovery_adds_embedding_without_disabling_configured_chat_pool(monkeypatch) -> None:
+ """A capability route coexists with an explicitly configured chat pool."""
embedding = DiscoveredModel(
provider_name="openai",
model_id="embedding-capable-model",
@@ -324,8 +693,11 @@ def test_auto_discovery_leaves_pool_unchanged_without_chat_capability_evidence(m
orchestrator = TaskOrchestrator([ModelAgent("bootstrap_agent", "bootstrap-model")])
- assert _auto_discover_runtime_agents(orchestrator) == {"added": [], "updated": []}
- assert [agent.id for agent in orchestrator.agents] == ["bootstrap_agent"]
+ result = _auto_discover_runtime_agents(orchestrator)
+ assert result == {"added": ["openai_embedding_capable_model"], "updated": []}
+ assert {agent.id for agent in orchestrator.agents} == {
+ "bootstrap_agent", "openai_embedding_capable_model"
+ }
def test_unrelated_discovery_keeps_configured_gateway_placeholder(monkeypatch) -> None:
@@ -371,7 +743,11 @@ def test_auto_discovery_uses_explicit_capabilities_before_model_id_heuristics(mo
orchestrator = TaskOrchestrator([ModelAgent("bootstrap_agent", "bootstrap-model")])
- assert _auto_discover_runtime_agents(orchestrator) == {"added": [], "updated": []}
+ result = _auto_discover_runtime_agents(orchestrator)
+ assert result["added"] == ["openai_generic_deployment"]
+ agent = orchestrator.select_capability_agent("embedding")
+ assert agent.model == "generic-deployment"
+ assert "chat" not in agent.tags
def test_auto_discovery_preserves_sole_real_bootstrap_seed(monkeypatch) -> None:
diff --git a/tests/test_batch_embeddings.py b/tests/test_batch_embeddings.py
index 36d6f7b38..153581df0 100644
--- a/tests/test_batch_embeddings.py
+++ b/tests/test_batch_embeddings.py
@@ -39,7 +39,16 @@
EmbeddingBatchResultItem,
)
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
-from contextual_orchestrator.token_counting import HeuristicTokenCounter # noqa: E402
+
+
+class _ExactTestCounter:
+ """Deterministic injected counter for synthetic embedding fixtures."""
+
+ def __init__(self, tokens_per_word: float = 1.0) -> None:
+ self.tokens_per_word = tokens_per_word
+
+ def count_text(self, text: str, model: str = "") -> int:
+ return int(len(text.split()) * self.tokens_per_word)
CONTRACT = json.loads(
@@ -75,7 +84,12 @@ def _serve():
price_book.set_price(
PriceEntry("acme-provider", "text-embedding-test", prompt_price_per_1k=0.13, completion_price_per_1k=0.0)
)
- coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_token_counter=_ExactTestCounter(),
+ )
token = "cost_token"
server = build_server(
orchestrator, port=0, security=SecurityConfig(auth_token=token), coordinator=coordinator
@@ -138,6 +152,8 @@ def retrieve(self, job):
class _PendingEmbeddingBackend(_RecordingEmbeddingBackend):
+ poll_after_ms = 250
+
def submit(self, requests, metadata=None):
super().submit(requests, metadata)
return BatchJob("pending-embeddings", self.name, "in_progress", len(requests))
@@ -166,7 +182,11 @@ def test_http_embeddings_try_cheapest_eligible_member_first(path: str, input_key
price_book.set_price(PriceEntry("cheap-provider", cheap.model, 0.01, 0.0))
backend = _RecordingEmbeddingBackend()
coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, embedding_batch_backend=backend
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_batch_backend=backend,
+ embedding_token_counter=_ExactTestCounter(),
)
token = "cheapest_http_token"
server = build_server(
@@ -222,7 +242,11 @@ def submit(self, requests, metadata=None):
backend = _FailCheapBackend()
coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, embedding_batch_backend=backend
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_batch_backend=backend,
+ embedding_token_counter=_ExactTestCounter(),
)
token = "health_order_http_token"
server = build_server(
@@ -290,7 +314,11 @@ def test_http_embeddings_omitted_model_reports_the_actually_served_model(
price_book.set_price(PriceEntry("cheap-served-provider", cheap.model, 0.01, 0.0))
backend = _RecordingEmbeddingBackend()
coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, embedding_batch_backend=backend
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_batch_backend=backend,
+ embedding_token_counter=_ExactTestCounter(),
)
token = "served_model_identity_http_token"
server = build_server(
@@ -443,7 +471,11 @@ def test_batch_embeddings_zdr_only_omitted_model_selects_zdr_capable_embedding_a
),
]
)
- coordinator = CostRoutingCoordinator(orchestrator, InMemoryConfigStore())
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ InMemoryConfigStore(),
+ embedding_token_counter=_ExactTestCounter(),
+ )
token = "zdr_batch_token"
server = build_server(
orchestrator, port=0, security=SecurityConfig(auth_token=token), coordinator=coordinator
@@ -468,6 +500,7 @@ def test_pending_batch_preserves_resolved_model_identity() -> None:
coordinator = CostRoutingCoordinator(
orchestrator,
InMemoryConfigStore(),
+ embedding_token_counter=_ExactTestCounter(),
embedding_batch_backend=_PendingEmbeddingBackend(),
)
@@ -476,6 +509,42 @@ def test_pending_batch_preserves_resolved_model_identity() -> None:
assert created["model"] == "resolved-embedding"
assert polled["model"] == "resolved-embedding"
+ assert created["poll_after_ms"] == _PendingEmbeddingBackend.poll_after_ms
+ assert created["job_retention_ms"] == coordinator.job_registry.retention_seconds * 1000
+
+
+def test_http_queued_embedding_admission_declares_owned_poll_and_retention() -> None:
+ """The public queued carrier exposes backend and registry lifecycle values."""
+ orchestrator = TaskOrchestrator(
+ [ModelAgent("embedding_worker", "resolved-embedding", tags=("embedding",))]
+ )
+ backend = _PendingEmbeddingBackend()
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ InMemoryConfigStore(),
+ embedding_token_counter=_ExactTestCounter(),
+ embedding_batch_backend=backend,
+ )
+ token = "synthetic_batch_token"
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token=token),
+ coordinator=coordinator,
+ )
+ threading.Thread(target=server.serve_forever, daemon=True).start()
+ try:
+ status, document = _request(
+ "POST",
+ f"http://127.0.0.1:{server.server_address[1]}/v1/batch/embeddings",
+ token,
+ {"inputs": ["synthetic"], "model": "resolved-embedding"},
+ )
+ assert status == 202
+ assert document["poll_after_ms"] == backend.poll_after_ms
+ assert document["job_retention_ms"] == coordinator.job_registry.retention_seconds * 1000
+ finally:
+ server.shutdown()
def test_empty_batch_preserves_resolved_model_identity() -> None:
@@ -510,7 +579,7 @@ def test_batch_embeddings_split_oversized_inputs_before_backend() -> None:
orchestrator,
config,
price_book=price_book,
- token_counter=HeuristicTokenCounter(tokens_per_word=1.0),
+ token_counter=_ExactTestCounter(tokens_per_word=1.0),
embedding_batch_backend=backend,
)
@@ -563,7 +632,7 @@ def test_batch_embeddings_char_guard_splits_no_whitespace_input() -> None:
coordinator = CostRoutingCoordinator(
orchestrator,
config,
- token_counter=HeuristicTokenCounter(tokens_per_word=1.0),
+ token_counter=_ExactTestCounter(tokens_per_word=1.0),
embedding_batch_backend=backend,
)
diff --git a/tests/test_batch_embeddings_encoding_dimensions_http_honesty.py b/tests/test_batch_embeddings_encoding_dimensions_http_honesty.py
index c6ef0e50d..e4cff750b 100644
--- a/tests/test_batch_embeddings_encoding_dimensions_http_honesty.py
+++ b/tests/test_batch_embeddings_encoding_dimensions_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "batch_embeddings_encoding_dimensions_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_batch_embeddings_endpoint_http_honesty.py b/tests/test_batch_embeddings_endpoint_http_honesty.py
index 41bbad5aa..eb174fb0a 100644
--- a/tests/test_batch_embeddings_endpoint_http_honesty.py
+++ b/tests/test_batch_embeddings_endpoint_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "batch_embeddings_endpoint_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_batch_embeddings_routing_http_honesty.py b/tests/test_batch_embeddings_routing_http_honesty.py
index ae041449a..ca23a8aaf 100644
--- a/tests/test_batch_embeddings_routing_http_honesty.py
+++ b/tests/test_batch_embeddings_routing_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "batch_embeddings_routing_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_batch_embeddings_user_http_honesty.py b/tests/test_batch_embeddings_user_http_honesty.py
index 92e475bc8..2881d174c 100644
--- a/tests/test_batch_embeddings_user_http_honesty.py
+++ b/tests/test_batch_embeddings_user_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "batch_embeddings_user_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py
index 6ef029ec3..39f745229 100644
--- a/tests/test_batch_job_registry.py
+++ b/tests/test_batch_job_registry.py
@@ -11,12 +11,18 @@
from __future__ import annotations
import sys
+import threading
+import time
+from types import SimpleNamespace
from pathlib import Path
from typing import Any, Dict
+import pytest
+
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from contextual_orchestrator.batch_job_registry import (
+ ClaimNotAcquired,
DEFAULT_RETENTION_SECONDS,
JobRegistryFactory,
ValkeyJsonMapping,
@@ -26,7 +32,9 @@
BatchJob,
BatchRequest,
BatchResultItem,
+ EmbeddingBatchRequest,
LocalBatchBackend,
+ ProviderEmbeddingBatchBackend,
)
from contextual_orchestrator.kv_config import InMemoryConfigStore
@@ -37,6 +45,54 @@ class FakeValkeyClient:
def __init__(self) -> None:
self.hashes: Dict[str, Dict[str, str]] = {}
self.expirations: Dict[str, int] = {}
+ self.strings: Dict[str, Any] = {}
+ self.execution_extension_attempted = threading.Event()
+ self.lose_execution_extension = True
+ self.execution_acquire_failures = 0
+
+ class LockNotOwnedError(RuntimeError):
+ pass
+
+ class _Lock:
+ def __init__(self, client: "FakeValkeyClient", name: str) -> None:
+ self._client = client
+ self.name = name
+ self.local = SimpleNamespace(token=f"token-{id(self)}".encode())
+ self._lose_on_extend = (
+ "provider_embedding_job_execution" in name
+ and client.lose_execution_extension
+ )
+ self._owned = False
+
+ def acquire(self) -> bool:
+ if (
+ "provider_embedding_job_execution" in self.name
+ and self._client.execution_acquire_failures > 0
+ ):
+ self._client.execution_acquire_failures -= 1
+ return False
+ self._owned = True
+ self._client.strings[self.name] = self.local.token
+ return True
+
+ def extend(self, _seconds: float, *, replace_ttl: bool) -> bool:
+ assert replace_ttl is True
+ if self._lose_on_extend:
+ self._owned = False
+ self._client.lose_execution_extension = False
+ self._client.strings.pop(self.name, None)
+ self._client.execution_extension_attempted.set()
+ return False
+ return self._owned
+
+ def owned(self) -> bool:
+ return self._owned
+
+ def release(self) -> None:
+ if not self._owned:
+ raise self._client.LockNotOwnedError("claim no longer owned")
+ self._owned = False
+ self._client.strings.pop(self.name, None)
def hget(self, key: str, field: str) -> Any:
return self.hashes.get(key, {}).get(field)
@@ -69,6 +125,53 @@ def expire(self, key: str, seconds: int) -> bool:
self.expirations[key] = seconds
return True
+ def lock(self, name: str, **_kwargs: Any) -> "FakeValkeyClient._Lock":
+ return self._Lock(self, name)
+
+ def eval(self, _script: str, key_count: int, *values: Any) -> int:
+ keys = values[:key_count]
+ args = values[key_count:]
+ if key_count == 2:
+ if len(args) == 5:
+ lock_key, states_key = keys
+ token, job_id, queued, running, retention = args
+ if self.strings.get(lock_key) != token:
+ return 0
+ current = self.hashes.get(states_key, {}).get(job_id)
+ if current not in {queued, running}:
+ return 0
+ self.hset(states_key, job_id, running)
+ self.expire(states_key, int(retention))
+ return 1
+ states_key, cancellations_key = keys
+ job_id, reserved, queued, running, cancellation, cancelled, retention = args
+ current = self.hashes.get(states_key, {}).get(job_id)
+ if current not in {reserved, queued, running}:
+ return 0
+ self.hset(cancellations_key, job_id, cancellation)
+ self.hset(states_key, job_id, cancelled)
+ for key in keys:
+ self.expire(key, int(retention))
+ return 1
+ lock_key, states_key, results_key, usage_key, errors_key = keys
+ token, job_id, running, queued, terminal, results, usage, error, retention = args
+ if self.strings.get(lock_key) != token:
+ return 0
+ current = self.hashes.get(states_key, {}).get(job_id)
+ if current not in {running, queued}:
+ return 0
+ for key, value in (
+ (results_key, results),
+ (usage_key, usage),
+ (errors_key, error),
+ ):
+ if value != "":
+ self.hset(key, job_id, value)
+ self.hset(states_key, job_id, terminal)
+ for key in keys[1:]:
+ self.expire(key, int(retention))
+ return 1
+
def test_mapping_round_trips_dataclasses_and_plain_values() -> None:
"""Dataclasses, dataclass lists, and JSON scalars all survive the trip."""
@@ -158,6 +261,291 @@ def test_default_retention_is_a_week() -> None:
assert DEFAULT_RETENTION_SECONDS == 7 * 24 * 3600
+def test_renewal_loss_is_visible_to_the_claim_holder() -> None:
+ """A failed CAS renewal fences the worker instead of becoming background noise."""
+ client = FakeValkeyClient()
+ factory = JobRegistryFactory(client)
+ with factory.lock(
+ "provider_embedding_job_execution",
+ "job",
+ lease_seconds=0.15,
+ renew_until_epoch=time.time() + 1,
+ ) as claim:
+ assert client.execution_extension_attempted.wait(timeout=1)
+ with pytest.raises(ClaimNotAcquired, match="ownership was lost"):
+ claim.ensure_owned()
+
+
+def test_provider_job_recovers_after_claim_renewal_loss_without_restart() -> None:
+ """A stale attempt cannot publish; the live worker reclaims and completes."""
+ client = FakeValkeyClient()
+ registry = JobRegistryFactory(client, retention_seconds=2)
+
+ calls = 0
+
+ def runner(_requests):
+ nonlocal calls
+ calls += 1
+ assert client.execution_extension_attempted.wait(timeout=1)
+ return [[float(calls)]], calls
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=registry,
+ claim_lease_seconds=0.15,
+ )
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ assert calls == 2
+ assert backend.retrieve(job)[0].embedding == [2.0]
+ assert backend.usage(job) == {"prompt_tokens": 2}
+ backend.close()
+
+
+def test_stale_provider_failure_is_fenced_before_live_recovery() -> None:
+ """A claim-losing failure cannot overwrite the succeeding attempt."""
+ client = FakeValkeyClient()
+ registry = JobRegistryFactory(client, retention_seconds=2)
+ calls = 0
+
+ def runner(_requests):
+ nonlocal calls
+ calls += 1
+ if calls == 1:
+ assert client.execution_extension_attempted.wait(timeout=1)
+ raise RuntimeError("stale provider failure")
+ return [[2.0]], 2
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner, job_registry=registry, claim_lease_seconds=0.15
+ )
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ assert backend.retrieve(job)[0].embedding == [2.0]
+ assert "provider_embedding_errors" not in client.hashes
+ backend.close()
+
+
+def test_terminal_transaction_rejects_a_transferred_claim_without_partial_writes() -> None:
+ """Claim transfer before EVAL leaves every terminal hash unchanged."""
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ registry = JobRegistryFactory(client)
+ states = registry.mapping("provider_embedding_states")
+ states["job"] = "running"
+ with registry.lock(
+ "provider_embedding_job_execution", "job", lease_seconds=1
+ ) as claim:
+ lock_name, _token = claim.atomic_identity()
+ client.strings[lock_name] = b"successor-token"
+ with pytest.raises(ClaimNotAcquired, match="before publication"):
+ registry.publish_provider_embedding_terminal(
+ claim,
+ "job",
+ status="completed",
+ results=[{"embedding": [1.0]}],
+ usage={"prompt_tokens": 1},
+ )
+
+ assert states["job"] == "running"
+ assert "batch_job_registry:provider_embedding_results" not in client.hashes
+ assert "batch_job_registry:provider_embedding_usage" not in client.hashes
+
+
+def test_durable_cancellation_wins_atomically_over_terminal_publication() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ release = threading.Event()
+ started = threading.Event()
+
+ def runner(_requests):
+ started.set()
+ assert release.wait(timeout=1)
+ return [[1.0]], 1
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=JobRegistryFactory(client),
+ claim_lease_seconds=1,
+ )
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+ assert started.wait(timeout=1)
+ assert backend.cancel(job, reason="caller cancelled")["status"] == "cancelled"
+ release.set()
+
+ assert backend.wait(job, timeout=1)["status"] == "cancelled"
+ assert backend.retrieve(job) == []
+ assert backend.usage(job) == {}
+ backend.close()
+
+
+def test_durable_job_past_deadline_becomes_failed_atomically() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ backend = ProviderEmbeddingBatchBackend(
+ lambda _requests: pytest.fail("expired work must not reach the provider"),
+ job_registry=JobRegistryFactory(client),
+ claim_lease_seconds=1,
+ )
+ job = backend.reserve(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+ backend._deadlines[job.job_id] = time.time() - 1
+ backend.start(job)
+
+ document = backend.wait(job, timeout=1)
+ assert document["status"] == "failed"
+ assert document["failure"] == {
+ "error_type": "TimeoutError",
+ "http_status": None,
+ "provider_code": "provider_embedding_deadline_exceeded",
+ "retryable": True,
+ "failed_shard_index": None,
+ }
+ assert backend.retrieve(job) == []
+ assert backend.usage(job) == {}
+ backend.close()
+
+
+def test_local_job_returning_after_deadline_fails() -> None:
+ def runner(_requests):
+ time.sleep(0.02)
+ return [[1.0]], 1
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=JobRegistryFactory(),
+ execution_timeout_seconds=0.01,
+ )
+ job = backend.submit([EmbeddingBatchRequest(input_text="synthetic")])
+
+ document = backend.wait(job, timeout=1)
+ assert document["status"] == "failed"
+ assert document["failure"]["provider_code"] == "provider_embedding_deadline_exceeded"
+ assert backend.retrieve(job) == []
+ backend.close()
+
+
+def test_execution_deadline_is_separate_from_result_retention(monkeypatch) -> None:
+ client = FakeValkeyClient()
+ registry = JobRegistryFactory(client, retention_seconds=123)
+ monkeypatch.setattr("contextual_orchestrator.batch_routing.time.time", lambda: 1000.0)
+ backend = ProviderEmbeddingBatchBackend(
+ lambda _requests: ([], 0),
+ job_registry=registry,
+ claim_lease_seconds=1,
+ execution_timeout_seconds=5,
+ )
+
+ job = backend.submit([])
+
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ assert backend._deadlines[job.job_id] == 1005.0
+ assert client.expirations["batch_job_registry:provider_embedding_deadlines"] == 123
+ backend.close()
+
+
+def test_queued_job_gets_its_lifetime_only_after_worker_claim() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ first_started = threading.Event()
+ release_first = threading.Event()
+ calls = 0
+
+ def runner(_requests):
+ nonlocal calls
+ calls += 1
+ if calls == 1:
+ first_started.set()
+ assert release_first.wait(timeout=1)
+ return [[float(calls)]], 1
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=JobRegistryFactory(client),
+ max_concurrency=1,
+ claim_lease_seconds=0.05,
+ execution_timeout_seconds=0.05,
+ )
+ first = backend.submit([EmbeddingBatchRequest(input_text="first")])
+ assert first_started.wait(timeout=1)
+ second = backend.submit([EmbeddingBatchRequest(input_text="second")])
+ assert second.job_id not in backend._deadlines
+ assert threading.Event().wait(0.1) is False
+ backend.cancel(first, reason="release queued worker")
+ release_first.set()
+
+ assert backend.wait(second, timeout=1)["status"] == "completed"
+ backend.close()
+
+
+def test_batch_lifetime_allows_one_client_timeout_per_request() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+
+ def runner(requests):
+ # Model the coordinator's worst case: each request becomes one
+ # sequential provider shard and consumes most of its client timeout.
+ for _request in requests:
+ assert threading.Event().wait(0.04) is False
+ return [[1.0] for _request in requests], len(requests)
+
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=JobRegistryFactory(client),
+ claim_lease_seconds=0.05,
+ execution_timeout_seconds=0.1,
+ )
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="one"), EmbeddingBatchRequest(input_text="two")]
+ )
+
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ backend.close()
+
+
+def test_batch_retries_initial_execution_deadline_claim() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ client.execution_acquire_failures = 1
+ backend = ProviderEmbeddingBatchBackend(
+ lambda requests: ([[1.0] for _request in requests], len(requests)),
+ job_registry=JobRegistryFactory(client),
+ claim_lease_seconds=0.01,
+ execution_timeout_seconds=1,
+ )
+
+ job = backend.submit([EmbeddingBatchRequest(input_text="one")])
+
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ backend.close()
+
+
+def test_durable_cancellation_cannot_be_overwritten_by_running_transition() -> None:
+ client = FakeValkeyClient()
+ client.lose_execution_extension = False
+ registry = JobRegistryFactory(client)
+ states = registry.mapping("provider_embedding_states")
+ states["job"] = "queued"
+ assert registry.cancel_provider_embedding("job", reason="caller cancelled")
+
+ with registry.lock(
+ "provider_embedding_job_execution", "job", lease_seconds=1
+ ) as claim:
+ with pytest.raises(ClaimNotAcquired, match="cancelled"):
+ registry.mark_provider_embedding_running(claim, "job")
+
+ assert states["job"] == "cancelled"
+
+
if __name__ == "__main__":
for name, value in sorted(globals().items()):
if name.startswith("test_") and callable(value):
diff --git a/tests/test_batch_optimizer.py b/tests/test_batch_optimizer.py
index 15db5f83d..350384a66 100644
--- a/tests/test_batch_optimizer.py
+++ b/tests/test_batch_optimizer.py
@@ -122,7 +122,7 @@ def test_mock_default_batch_route_works_without_usage() -> None:
with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None):
records = orchestrator.batch_route(["hello there"])
assert records[0]["answer"].startswith("[general_agent:")
- assert orchestrator.spend_analytics()["by_model"][0]["usage_source"] == "estimated"
+ assert orchestrator.spend_analytics()["by_model"][0]["usage_source"] == "unavailable"
@pytest.mark.parametrize("kind", ["missing", "content"])
@@ -473,9 +473,9 @@ def judge(self, **kwargs):
assert orchestrator.budget_status()["spent_output_tokens"] == 8
spend = orchestrator.spend_analytics()
- assert spend["totals"]["estimated_output_tokens"] == 8
+ assert spend["totals"]["output_tokens"] == 8
assert spend["by_model"][0]["output_tokens"] == 8
- assert spend["by_model"][0]["estimated_cost_usd"] == pytest.approx(0.00008)
+ assert spend["by_model"][0]["cost_usd"] == pytest.approx(0.00008)
def test_batch_route_judge_usage_survives_agent_pool_model_change() -> None:
@@ -513,11 +513,11 @@ def judge(self, **kwargs):
assert orchestrator.budget_status()["spent_output_tokens"] == 8
spend = orchestrator.spend_analytics()
- assert spend["totals"]["estimated_output_tokens"] == 8
+ assert spend["totals"]["output_tokens"] == 8
by_model = {row["model"]: row for row in spend["by_model"]}
assert "model-y" not in by_model # no historical spend leaks onto the new model
assert by_model["model-x"]["output_tokens"] == 8 # worker (6) + judge (2), pinned
- assert by_model["model-x"]["estimated_cost_usd"] == pytest.approx(0.00008)
+ assert by_model["model-x"]["cost_usd"] == pytest.approx(0.00008)
def test_batch_route_persists_earlier_group_spend_when_a_later_group_fails() -> None:
diff --git a/tests/test_batch_routing.py b/tests/test_batch_routing.py
index ccb564c86..2b84dd018 100644
--- a/tests/test_batch_routing.py
+++ b/tests/test_batch_routing.py
@@ -67,6 +67,14 @@ def test_batch_min_tokens_threshold_from_config() -> None:
assert policy.decide(RoutingHints(), prompt_tokens=800).channel == "batch"
+def test_batch_token_threshold_stays_sync_when_prompt_count_unavailable() -> None:
+ config = InMemoryConfigStore()
+ config.set("routing", "batch_min_tokens", 500)
+ decision = RoutingPolicy(config).decide(RoutingHints(), prompt_tokens=None)
+ assert decision.channel == "sync"
+ assert "unavailable" in decision.reason
+
+
def test_batch_disabled_config_forces_sync() -> None:
config = InMemoryConfigStore()
config.set("routing", "batch_enabled", False)
diff --git a/tests/test_batch_routing_boundaries.py b/tests/test_batch_routing_boundaries.py
index 08a90b54a..4b7f69171 100644
--- a/tests/test_batch_routing_boundaries.py
+++ b/tests/test_batch_routing_boundaries.py
@@ -126,6 +126,7 @@ class _FakeEmbeddingClient:
"""Async pg-llm-batch double for the embeddings endpoint."""
def __init__(self) -> None:
+ self.poll_interval_seconds = 1.0
self.uploaded: list[tuple[str, str]] = []
self.created: list[Dict[str, Any]] = []
self.status_calls: list[str] = []
@@ -243,6 +244,7 @@ def test_pg_embedding_backend_full_lifecycle_with_assembler() -> None:
def test_pg_embedding_backend_rejects_untracked_response_ids() -> None:
client = _FakeEmbeddingClient()
backend = PgLlmBatchEmbeddingBackend(client, endpoint_alias="nim-east")
+ assert backend.poll_after_ms == int(client.poll_interval_seconds * 1000)
requests = [EmbeddingBatchRequest(input_text="solo input", custom_id="emb_solo")]
job = backend.submit(requests)
assert job.status == "validating"
diff --git a/tests/test_batch_routing_boundaries_extra.py b/tests/test_batch_routing_boundaries_extra.py
index 859354012..02d848243 100644
--- a/tests/test_batch_routing_boundaries_extra.py
+++ b/tests/test_batch_routing_boundaries_extra.py
@@ -10,20 +10,16 @@
)
-def test_local_backend_without_token_counter_counts_word_units() -> None:
- """With no injected counter, token accounting falls back to word count."""
+def test_local_backend_without_token_counter_fails_closed() -> None:
+ """Missing authoritative accounting must not become a word-count estimate."""
backend = LocalEmbeddingBatchBackend()
request = EmbeddingBatchRequest(
custom_id=None,
model="local-embedding-model",
input_text="one two three four\nfive",
)
- job = backend.submit([request])
-
- assert job.request_count == 1
- results = backend.retrieve(job)
- assert len(results) == 1
- assert results[0].prompt_tokens == 5
+ with pytest.raises(RuntimeError, match="authoritative embedding tokenizer"):
+ backend.submit([request])
def test_local_backend_injected_counter_still_takes_precedence() -> None:
diff --git a/tests/test_budget_enforcement.py b/tests/test_budget_enforcement.py
index e6d571afc..9c1018e05 100644
--- a/tests/test_budget_enforcement.py
+++ b/tests/test_budget_enforcement.py
@@ -15,26 +15,40 @@
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
from contextual_orchestrator.orchestrator import BudgetExceededError # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
+from contextual_orchestrator.token_counting import UnavailableTokenCounter # noqa: E402
def _agent() -> ModelAgent:
return ModelAgent("general_agent", "test-model", tags=("reasoning",))
+class _ExactCounter:
+ """Exact synthetic raw-output counter for budget tests."""
+
+ def count_text(self, text: str, model: str) -> int:
+ return len(text.encode("utf-8"))
+
+
+def _orchestrator(agents: list[ModelAgent], **kwargs) -> TaskOrchestrator:
+ return TaskOrchestrator(agents, token_counter=_ExactCounter(), **kwargs)
+
+
def test_default_no_budget_is_unchanged() -> None:
- orchestrator = TaskOrchestrator([_agent()])
+ orchestrator = _orchestrator([_agent()])
orchestrator.run([{"role": "user", "content": "one"}])
orchestrator.run([{"role": "user", "content": "two"}]) # no cap, both allowed
assert orchestrator.spend_analytics()["budget"]["enabled"] is False
def test_token_budget_allows_then_blocks() -> None:
- orchestrator = TaskOrchestrator([_agent()], budget_max_output_tokens=1)
+ orchestrator = _orchestrator([_agent()], budget_max_output_tokens=1)
orchestrator.run([{"role": "user", "content": "first run is allowed"}]) # spent was 0 at entry
raised = False
@@ -47,8 +61,56 @@ def test_token_budget_allows_then_blocks() -> None:
assert raised
+def test_enabled_budget_fails_closed_after_unavailable_usage() -> None:
+ orchestrator = TaskOrchestrator(
+ [_agent()],
+ token_counter=UnavailableTokenCounter(),
+ budget_max_output_tokens=100,
+ )
+ orchestrator.run([{"role": "user", "content": "first call has unavailable usage"}])
+ budget = orchestrator.budget_status()
+ assert budget["measurement_status"] == "unavailable"
+ assert budget["enforcement_status"] == "blocked_unavailable"
+ assert budget["spent_output_tokens"] is None
+ with pytest.raises(BudgetExceededError, match="measurement unavailable"):
+ orchestrator.run([{"role": "user", "content": "must not dispatch"}])
+
+
+def test_cost_budget_fails_closed_for_an_unpriced_served_model() -> None:
+ orchestrator = _orchestrator(
+ [_agent()],
+ budget_max_cost_usd=10.0,
+ )
+ # With no price book, a cost budget cannot safely admit even the first call.
+ budget = orchestrator.budget_status()
+ assert budget["enforcement_status"] == "blocked_unavailable"
+ with pytest.raises(BudgetExceededError, match="measurement unavailable"):
+ orchestrator.run([{"role": "user", "content": "must not dispatch"}])
+
+
+def test_cost_budget_ignores_unpriced_embedding_only_candidate() -> None:
+ orchestrator = _orchestrator(
+ [
+ ModelAgent("chat_agent", "priced-chat", tags=("capability:chat",)),
+ ModelAgent(
+ "embedding_agent",
+ "unpriced-embedding",
+ tags=("capability:embedding",),
+ ),
+ ],
+ price_per_million={"priced-chat": 1.0},
+ budget_max_cost_usd=10.0,
+ )
+
+ orchestrator.run([{"role": "user", "content": "priced chat workflow"}])
+
+ assert orchestrator.budget_status()["enforcement_status"] == "within_budget"
+ assert orchestrator.budget_status()["spent_cost_usd"] > 0
+ assert orchestrator.spend_analytics()["budget"] == orchestrator.budget_status()
+
+
def test_budget_block_reports_spent_and_remaining() -> None:
- orchestrator = TaskOrchestrator([_agent()], budget_max_output_tokens=1000)
+ orchestrator = _orchestrator([_agent()], budget_max_output_tokens=1000)
orchestrator.run([{"role": "user", "content": "measure the budget"}])
budget = orchestrator.spend_analytics()["budget"]
@@ -60,7 +122,7 @@ def test_budget_block_reports_spent_and_remaining() -> None:
def test_cost_budget_blocks() -> None:
- orchestrator = TaskOrchestrator(
+ orchestrator = _orchestrator(
[ModelAgent("general_agent", "priced-model", tags=("reasoning",))],
price_per_million={"priced-model": 1_000_000.0}, # $1 per token, so any run exceeds a tiny cap
budget_max_cost_usd=0.001,
@@ -79,7 +141,7 @@ def test_cost_budget_blocks() -> None:
def test_budget_meter_matches_randomized_recorded_run_analytics() -> None:
"""Incremental token/cost state equals the authoritative full aggregation."""
agents = [ModelAgent("agent_one", "model-one"), ModelAgent("agent_two", "model-two")]
- orchestrator = TaskOrchestrator(
+ orchestrator = _orchestrator(
agents,
price_per_million={"model-one": 0.75, "model-two": 3.25},
budget_max_output_tokens=10_000,
@@ -107,7 +169,7 @@ def test_budget_meter_matches_randomized_recorded_run_analytics() -> None:
def test_budget_status_does_not_scan_recorded_runs() -> None:
"""The per-request gate remains independent of workflow-run cardinality."""
- orchestrator = TaskOrchestrator([_agent()], budget_max_output_tokens=1)
+ orchestrator = _orchestrator([_agent()], budget_max_output_tokens=1)
orchestrator.run([{"role": "user", "content": "record one run"}])
orchestrator.spend_analytics = lambda: (_ for _ in ()).throw(
AssertionError("budget_status scanned workflow runs")
@@ -118,7 +180,7 @@ def test_budget_status_does_not_scan_recorded_runs() -> None:
def test_budget_meter_rebuilds_from_persisted_runs(tmp_path: Path) -> None:
"""Restarted gates recover the same meter from durable workflow records."""
state_db = str(tmp_path / "budget_state.db")
- first = TaskOrchestrator(
+ first = _orchestrator(
[_agent()],
state_db=state_db,
price_per_million={"test-model": 2.5},
@@ -128,7 +190,7 @@ def test_budget_meter_rebuilds_from_persisted_runs(tmp_path: Path) -> None:
expected = first.spend_analytics()["budget"]
first.close()
- restored = TaskOrchestrator(
+ restored = _orchestrator(
[_agent()],
state_db=state_db,
price_per_million={"test-model": 2.5},
@@ -144,7 +206,7 @@ def test_budget_meter_reconciles_after_agent_status_change() -> None:
"""Pool mutations preserve historical spend and analytics parity."""
priced = ModelAgent("priced_agent", "priced-model")
fallback = ModelAgent("fallback_agent", "fallback-model")
- orchestrator = TaskOrchestrator(
+ orchestrator = _orchestrator(
[priced, fallback],
price_per_million={"priced-model": 10.0},
budget_max_cost_usd=1.0,
@@ -167,7 +229,7 @@ def test_budget_meter_reconciles_after_agent_status_change() -> None:
def test_replacing_a_run_does_not_accumulate_fractional_cost_drift() -> None:
"""Repeated replacements retain exact decimal cost accounting."""
agents = [ModelAgent("agent_one", "model-one"), ModelAgent("agent_two", "model-two")]
- orchestrator = TaskOrchestrator(
+ orchestrator = _orchestrator(
agents,
price_per_million={"model-one": 0.1, "model-two": 0.2},
budget_max_cost_usd=1.0,
@@ -194,7 +256,7 @@ def test_replacing_a_run_does_not_accumulate_fractional_cost_drift() -> None:
def test_http_over_budget_returns_429() -> None:
token = "budget_token"
- orchestrator = TaskOrchestrator([_agent()], budget_max_output_tokens=1)
+ orchestrator = _orchestrator([_agent()], budget_max_output_tokens=1)
orchestrator.run([{"role": "user", "content": "prime the budget over the cap"}]) # now exceeded
server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=token))
diff --git a/tests/test_chat_parallel_tool_calls_http_honesty.py b/tests/test_chat_parallel_tool_calls_http_honesty.py
index 8ad66c481..98b520369 100644
--- a/tests/test_chat_parallel_tool_calls_http_honesty.py
+++ b/tests/test_chat_parallel_tool_calls_http_honesty.py
@@ -9,9 +9,12 @@
from pathlib import Path
import sys
+from jsonschema import validate
+
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator.api_contract import OPENAPI_SPEC # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "chat_parallel_tool_calls_honesty_token" # noqa: S105
@@ -132,7 +135,14 @@ def test_http_chat_parallel_tool_calls_true_with_tools_passthrough() -> None:
)
# Mock passthrough returns chat-shaped body
assert status == 200, body
- assert "choices" in body or "id" in body
+ validate(
+ body,
+ {
+ "$ref": "#/components/schemas/ChatCompletionResponse",
+ "components": OPENAPI_SPEC["components"],
+ },
+ )
+ assert body["usage_measurement_status"] == "measured"
finally:
server.shutdown()
thread.join(timeout=5)
diff --git a/tests/test_chat_reasoning_effort_http_honesty.py b/tests/test_chat_reasoning_effort_http_honesty.py
index a22b8a244..20b00a369 100644
--- a/tests/test_chat_reasoning_effort_http_honesty.py
+++ b/tests/test_chat_reasoning_effort_http_honesty.py
@@ -105,6 +105,26 @@ def test_http_chat_accepts_reasoning_effort_none_as_omit() -> None:
thread.join(timeout=5)
+def test_http_chat_accepts_orchestrator_auto_without_provider_forwarding() -> None:
+ """Consumer default ``auto`` stays at the orchestration boundary."""
+ server, thread, port = _server()
+ try:
+ status, body = _post(
+ port,
+ {
+ "model": "mock-planner",
+ "messages": [{"role": "user", "content": "synthetic auto effort"}],
+ "reasoning_effort": "auto",
+ },
+ )
+ assert status == 200, body
+ assert "choices" in body
+ assert "reasoning_effort" not in body.get("echo", {})
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
def test_http_chat_rejects_reasoning_effort_bool() -> None:
server, thread, port = _server()
try:
diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py
index 154f9579b..43dc8ada3 100644
--- a/tests/test_chat_response_format_http_honesty.py
+++ b/tests/test_chat_response_format_http_honesty.py
@@ -12,6 +12,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator.provider_errors import ProviderUpstreamError # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "chat_response_format_http_honesty_token" # noqa: S105
@@ -83,6 +84,353 @@ def test_http_chat_accepts_response_format_json_object() -> None:
thread.join(timeout=5)
+def test_http_structured_synthesis_classifies_upstream_404() -> None:
+ """Final synthesis provider rejection is typed and never a raw 500."""
+ orchestrator = build()
+
+ def reject_synthesis(*_args, **_kwargs):
+ raise urllib.error.HTTPError(
+ "https://provider.synthetic.invalid/v1/chat/completions",
+ 404,
+ "not found",
+ {},
+ None,
+ )
+
+ orchestrator.client.proxy_send = reject_synthesis
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN),
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": "mock-planner",
+ "messages": [{"role": "user", "content": "json object mode"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 404, body
+ assert body["error"]["code"] == "model_not_found"
+ assert body["error"]["detail"]["transport"] == "structured_synthesis"
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_virtual_structured_synthesis_replaces_stale_model_on_same_endpoint() -> None:
+ """Virtual routing retries a missing catalog model without crossing endpoints."""
+ agents = [
+ ModelAgent("stale_agent", "stale-model", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("live_agent", "live-model", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("other_agent", "other-model", "mock://other", tags=("reasoning", "writing")),
+ ]
+ orchestrator = TaskOrchestrator(agents)
+ original_select_agent = orchestrator._select_agent
+ orchestrator._select_agent = lambda task, role, **kwargs: (
+ agents[0] if role == "synthesizer" else original_select_agent(task, role, **kwargs)
+ )
+ calls = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ if agent.id == "stale_agent":
+ raise urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None)
+ return {"choices": [{"message": {"content": '{"status":"ok"}'}}]}
+
+ orchestrator.client.proxy_send_once = send
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ "session_id": "synthetic-session",
+ },
+ )
+ assert status == 200, body
+ assert calls == ["stale_agent", "live_agent"]
+ assert "other_agent" not in calls
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_virtual_structured_schema_exhaustion_is_typed_and_non_repeating() -> None:
+ """Schema-invalid synthesis and repair exhaust each same-endpoint model once."""
+ agents = [
+ ModelAgent("first_agent", "first-model", "mock://catalog"),
+ ModelAgent("second_agent", "second-model", "mock://catalog"),
+ ModelAgent("other_agent", "other-model", "mock://other"),
+ ]
+ orchestrator = TaskOrchestrator(agents)
+ calls: list[str] = []
+ orchestrator.conduct = lambda *_args, **_kwargs: {"trace": []} # type: ignore[method-assign]
+ orchestrator._select_agent = lambda *_args, **_kwargs: agents[0] # type: ignore[method-assign]
+ orchestrator._failover_candidates = lambda *_args, **_kwargs: list(agents) # type: ignore[method-assign]
+
+ def invalid(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ return {"choices": [{"message": {"content": '{"milestones":[{"extra":true}]}'}}]}
+
+ orchestrator.client.proxy_send_once = invalid
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "synthetic structured request"}],
+ "response_format": {
+ "type": "json_schema",
+ "json_schema": {
+ "name": "milestone_list",
+ "strict": True,
+ "schema": {
+ "type": "object",
+ "properties": {
+ "milestones": {
+ "type": "array",
+ "items": {
+ "type": "object",
+ "properties": {"label": {"type": "string"}},
+ "required": ["label"],
+ "additionalProperties": False,
+ },
+ }
+ },
+ "required": ["milestones"],
+ "additionalProperties": False,
+ },
+ },
+ },
+ "session_id": "synthetic-session",
+ },
+ )
+ assert status == 502, body
+ assert body["error"]["code"] == "invalid_structured_output"
+ assert calls == ["first_agent", "first_agent", "second_agent", "second_agent"]
+ assert "other_agent" not in calls
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_virtual_structured_workflow_never_reuses_request_scoped_missing_model() -> None:
+ """A model missing in evidence work is excluded from every later role and synthesis."""
+ agents = [
+ ModelAgent("stale_agent", "stale-model", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("live_agent", "live-model", "mock://catalog", tags=("reasoning", "writing")),
+ ]
+ orchestrator = TaskOrchestrator(agents)
+ original_ranked = orchestrator._ranked_agents
+
+ def stale_first(*args, **kwargs):
+ ranked = original_ranked(*args, **kwargs)
+ return sorted(ranked, key=lambda candidate: candidate.id != "stale_agent")
+
+ orchestrator._ranked_agents = stale_first
+ chat_calls: list[str] = []
+ synthesis_calls: list[str] = []
+
+ def chat(agent, _messages, **_kwargs):
+ chat_calls.append(agent.id)
+ if agent.id == "stale_agent":
+ raise ProviderUpstreamError(
+ agent_id=agent.id,
+ model=agent.model,
+ error_code="model_not_found",
+ message="synthetic missing model",
+ client_status=404,
+ provider_status=404,
+ retryable=False,
+ )
+ return '{"status":"evidence"}'
+
+ def synthesize(agent, _endpoint, _payload):
+ synthesis_calls.append(agent.id)
+ return {"choices": [{"message": {"content": '{"status":"ok"}'}}]}
+
+ orchestrator.client.chat = chat
+ orchestrator.client.proxy_send_once = synthesize
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 200, body
+ assert chat_calls.count("stale_agent") == 1
+ assert synthesis_calls == ["live_agent"]
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_virtual_structured_workflow_exhausts_each_missing_model_once() -> None:
+ """A fully stale virtual pool terminates typed after one attempt per model."""
+ agents = [
+ ModelAgent("stale_a", "stale-a", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("stale_b", "stale-b", "mock://catalog", tags=("reasoning", "writing")),
+ ]
+ orchestrator = TaskOrchestrator(agents)
+ calls: list[str] = []
+
+ def missing(agent, _messages, **_kwargs):
+ calls.append(agent.id)
+ raise ProviderUpstreamError(
+ agent_id=agent.id,
+ model=agent.model,
+ error_code="model_not_found",
+ message="synthetic missing model",
+ client_status=404,
+ provider_status=404,
+ retryable=False,
+ )
+
+ orchestrator.client.chat = missing
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 404, body
+ assert body["error"]["code"] == "model_not_found"
+ assert sorted(calls) == ["stale_a", "stale_b"]
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_explicit_structured_model_preserves_model_not_found() -> None:
+ """An explicit model pin never switches models after a provider 404."""
+ orchestrator = TaskOrchestrator([
+ ModelAgent("stale_agent", "stale-model", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("live_agent", "live-model", "mock://catalog", tags=("reasoning", "writing")),
+ ])
+ calls = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ raise urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None)
+
+ orchestrator.client.proxy_send = send
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": "stale-model",
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 404
+ assert body["error"]["code"] == "model_not_found"
+ assert calls
+ assert set(calls) == {"stale_agent"}
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_explicit_structured_model_preserves_authentication_error() -> None:
+ """An explicit model pin returns its own 401 and never changes endpoints."""
+ orchestrator = TaskOrchestrator([
+ ModelAgent("auth_failing", "pinned-model", "mock://pinned", tags=("reasoning", "writing")),
+ ModelAgent("other_endpoint", "other-model", "mock://other", tags=("reasoning", "writing")),
+ ])
+ calls = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ raise urllib.error.HTTPError("https://synthetic.invalid", 401, "unauthorized", {}, None)
+
+ orchestrator.client.proxy_send = send
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": "pinned-model",
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 401
+ assert body["error"]["code"] == "authentication_error"
+ assert calls and set(calls) == {"auth_failing"}
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+def test_virtual_missing_models_record_each_circuit_failure_once() -> None:
+ """Same-endpoint model replacement does not double-count its final 404."""
+ agents = [
+ ModelAgent("stale_agent", "stale-model", "mock://catalog", tags=("reasoning", "writing")),
+ ModelAgent("live_agent", "live-model", "mock://catalog", tags=("reasoning", "writing")),
+ ]
+ orchestrator = TaskOrchestrator(agents)
+ original_select_agent = orchestrator._select_agent
+ orchestrator._select_agent = lambda task, role, **kwargs: (
+ agents[0] if role == "synthesizer" else original_select_agent(task, role, **kwargs)
+ )
+ calls = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ raise urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None)
+
+ orchestrator.client.proxy_send_once = send
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, body = _post(
+ server.server_address[1],
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ )
+ assert status == 404, body
+ assert calls == ["stale_agent", "live_agent"]
+ assert orchestrator._circuit["stale_agent"]["failures"] == 1
+ assert orchestrator._circuit["live_agent"]["failures"] == 1
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
def test_http_structured_chat_rejects_batch_routing() -> None:
"""Provider-native structured synthesis has no batch execution contract."""
server, thread, port = _server()
diff --git a/tests/test_compose_contract.py b/tests/test_compose_contract.py
index 6db9c518a..15de905e9 100644
--- a/tests/test_compose_contract.py
+++ b/tests/test_compose_contract.py
@@ -23,7 +23,9 @@ def test_compose_uses_postgres_kv_and_secret_bootstrap() -> None:
def test_gateway_image_installs_postgres_driver_and_ignores_secrets() -> None:
dockerfile = Path("Dockerfile").read_text()
assert "COPY pyproject.toml requirements.lock README.md LICENSE ./" in dockerfile
- assert "pip install --no-cache-dir --require-hashes -r requirements.lock" in dockerfile
+ assert "uv pip install --python 3.12 --require-hashes" in dockerfile
+ assert "COPY --from=dependency-builder /build/deps/" in dockerfile
+ assert "maturin build --locked --release" in dockerfile
assert "--production" in dockerfile
assert "--admin-token-key CONTEXTUAL_ORCHESTRATOR_ADMIN_TOKEN" in dockerfile
assert "--inference-token-key CONTEXTUAL_ORCHESTRATOR_INFERENCE_TOKEN" in dockerfile
diff --git a/tests/test_cost_review_server.py b/tests/test_cost_review_server.py
index b62105125..b32e69df2 100644
--- a/tests/test_cost_review_server.py
+++ b/tests/test_cost_review_server.py
@@ -200,7 +200,8 @@ def test_chat_completion_reports_real_usage_and_records_cost() -> None:
"messages": [{"role": "user", "content": "hello there world"}],
"attribution": {"team": "alpha", "company": "acme"}})
assert status == 200
- assert body["usage"]["total_tokens"] > 0
+ assert body["usage"] is None
+ assert body["usage_measurement_status"] == "unavailable"
assert body["orchestration"]["channel"] == "sync"
status, report = _request("GET", f"{base}/api/v1/cost_reports/rollup?dimension=team", token)
@@ -212,8 +213,8 @@ def test_chat_completion_reports_real_usage_and_records_cost() -> None:
server.shutdown()
-def test_chat_completion_fallback_path_labels_estimated_measurement_status() -> None:
- """Provider-unreported usage stays honestly labeled estimated end to end."""
+def test_chat_completion_missing_usage_is_unavailable_end_to_end() -> None:
+ """Provider-unreported usage stays unavailable end to end."""
server, port, token = _serve()
base = f"http://127.0.0.1:{port}"
try:
@@ -229,13 +230,27 @@ def test_chat_completion_fallback_path_labels_estimated_measurement_status() ->
assert status == 200, body
# The mock provider reports no usage, so both the completion cost
# payload and the analytics usage-ledger rows must carry the explicit
- # estimated measurement status instead of claiming provider-measured.
- assert body["orchestration"]["cost"]["measurement_status"] == "estimated"
+ # unavailable measurement status instead of claiming provider-measured.
+ assert body["orchestration"]["cost"]["measurement_status"] == "unavailable"
status, records = _request("GET", f"{base}/api/v1/llm_usage_records", token)
assert status == 200, records
assert records["total_count"] == 1
- assert records["items"][0]["measurement_status"] == "estimated"
+ row = records["items"][0]
+ assert row["measurement_status"] == "unavailable"
+ assert row["prompt_tokens"] is None
+ assert row["completion_tokens"] is None
+ assert row["total_tokens"] is None
+ assert row["cost_amount"] is None
+
+ status, report = _request(
+ "GET", f"{base}/api/v1/cost_reports/rollup?dimension=model_name", token
+ )
+ assert status == 200, report
+ assert report["grand_total"]["measurement_status"] == "unavailable"
+ assert report["grand_total"]["record_count_by_status"]["unavailable"] == 1
+ assert report["items"][0]["measurement_status"] == "unavailable"
+ assert report["items"][0]["record_count_by_status"]["unavailable"] == 1
finally:
server.shutdown()
diff --git a/tests/test_cost_router.py b/tests/test_cost_router.py
index 52c0c8410..82ca8c413 100644
--- a/tests/test_cost_router.py
+++ b/tests/test_cost_router.py
@@ -3,7 +3,7 @@
from __future__ import annotations
import sys
-import threading
+import time
from pathlib import Path
import pytest
@@ -22,11 +22,8 @@
TaskOrchestrator,
)
from contextual_orchestrator.batch_routing import ( # noqa: E402
- BatchDownloadError,
BatchJob,
BatchRequest,
- BatchResultItem,
- LocalBatchBackend,
PgLlmBatchBackend,
)
from contextual_orchestrator.cost_router import BatchModelSelectionError # noqa: E402
@@ -59,7 +56,8 @@ def test_sync_completion_records_usage_and_returns_costs() -> None:
attribution={"team": "alpha", "company": "acme"},
)
assert result["channel"] == "sync"
- assert result["usage"]["total_tokens"] > 0
+ assert result["usage"] is None
+ assert result["usage_measurement_status"] == "unavailable"
assert result["usage_record_id"].startswith("usage_")
records = coordinator.ledger.records()
assert len(records) == len(result["usage_record_ids"])
@@ -67,7 +65,7 @@ def test_sync_completion_records_usage_and_returns_costs() -> None:
assert all(record["provider_name"] == "mock" for record in records)
assert all(record["model_name"] == "mock-a" for record in records)
assert all(record["request_channel"] == "sync" for record in records)
- assert all(record["measurement_status"] == "estimated" for record in records)
+ assert all(record["measurement_status"] == "unavailable" for record in records)
def test_sync_completion_preserves_provider_reported_usage() -> None:
@@ -198,37 +196,6 @@ def test_completed_race_loser_usage_is_recorded_as_measured_provider_spend() ->
assert record["workflow_run_id"] == "run_race"
-def test_race_loser_with_unparseable_usage_is_recorded_as_unavailable() -> None:
- """A billable race-loser call with malformed usage still gets a ledger row."""
- coordinator = _coordinator()
- context = {
- "route_mode": "route",
- "attribution": {"team": "alpha"},
- "model_name": "contextual-orchestrator",
- "workflow_run_id": "run_race_unmeasurable",
- "workflow_ready": True,
- "records": [],
- "pending_usage": [],
- }
- token = coordinator._race_usage_context.set(context)
- try:
- coordinator._record_race_endpoint_usage(
- "mock_worker",
- ("duplicate", "mock_worker", None),
- )
- finally:
- coordinator._race_usage_context.reset(token)
- records = coordinator.ledger.records()
- assert len(records) == 1
- record = records[0]
- assert record["measurement_status"] == "unavailable"
- assert record["prompt_tokens"] == 0
- assert record["completion_tokens"] == 0
- assert record["provider_name"] == "mock"
- assert record["model_name"] == "mock-a"
- assert record["workflow_run_id"] == "run_race_unmeasurable"
-
-
def test_race_loser_derives_provider_from_base_url_when_name_is_absent() -> None:
"""Race-loser spend uses the same provider identity as winner accounting."""
agent = ModelAgent(
@@ -310,45 +277,6 @@ def run(*_args, **_kwargs):
assert result["cost"]["cost_amount"] == 0.022
-def test_sync_cost_reports_unavailable_when_a_race_loser_cannot_be_measured() -> None:
- """Devin review (#955): the plain orchestrator.run() sync path's own cost
- aggregation needs the same unavailable-outranks-estimated precedence as
- the provider_request path -- a measured winner plus an unavailable race
- loser must not report a confident "measured" total or sum a real cost
- over an unknown one.
- """
- coordinator = _coordinator()
-
- def run(*_args, **_kwargs):
- coordinator.orchestrator._race_usage_sink(
- "mock_worker",
- ("duplicate", "mock_worker", None), # unparseable usage
- )
- return {
- "workflow_run_id": "run_race_unavailable",
- "mode": "route",
- "answer": "winner",
- "trace": [
- {
- "agent_id": "mock_worker",
- "output": "winner",
- "usage": {"prompt_tokens": 7, "completion_tokens": 3},
- }
- ],
- }
-
- coordinator.orchestrator.run = run # type: ignore[method-assign]
- result = coordinator.complete([{"role": "user", "content": "race"}], mode="route")
-
- assert result["cost"]["measurement_status"] == "unavailable"
- assert result["cost"]["cost_amount"] is None
- assert "currency_components" not in result["cost"]
- assert {record["measurement_status"] for record in coordinator.ledger.records()} == {
- "measured",
- "unavailable",
- }
-
-
def test_ready_race_usage_without_workflow_id_is_not_discarded() -> None:
coordinator = _coordinator()
context = {
@@ -373,7 +301,7 @@ def test_ready_race_usage_without_workflow_id_is_not_discarded() -> None:
def test_conducted_plain_completion_records_every_step_usage() -> None:
- """A conducted run preserves measured and estimated evidence per provider call."""
+ """A conducted run preserves measured and unavailable evidence per call."""
coordinator = _coordinator()
coordinator.orchestrator.run = lambda *args, **kwargs: { # type: ignore[method-assign]
"workflow_run_id": "run_plain_conducted",
@@ -396,26 +324,20 @@ def test_conducted_plain_completion_records_every_step_usage() -> None:
assert len(records) == 2
assert result["usage_record_ids"] == [row["usage_record_id"] for row in records]
assert result["usage_record_id"] == records[-1]["usage_record_id"]
- assert [row["measurement_status"] for row in records] == ["measured", "estimated"]
- assert result["cost"]["measurement_status"] == "estimated"
- assert result["usage"] == {
- "prompt_tokens": sum(row["prompt_tokens"] for row in records),
- "completion_tokens": sum(row["completion_tokens"] for row in records),
- "total_tokens": sum(row["total_tokens"] for row in records),
- }
+ assert [row["measurement_status"] for row in records] == ["measured", "unavailable"]
+ assert result["cost"]["measurement_status"] == "unavailable"
+ assert result["usage"] is None
assert records[0]["prompt_tokens"] == 7
assert records[0]["completion_tokens"] == 3
- assert records[1]["prompt_tokens"] == coordinator.token_counter.count_messages(
- messages, "mock-a"
- )
+ assert records[1]["prompt_tokens"] is None
+ assert records[1]["completion_tokens"] is None
def test_sync_records_derive_provider_and_model_from_served_agent() -> None:
coordinator = _coordinator()
coordinator.complete([{"role": "user", "content": "do a thing"}])
row = coordinator.ledger.records()[0]
- # cost = prompt/1k * 1 + completion/1k * 2, both > 0 given the mock echo answer
- assert row["cost_amount"] >= 0.0
+ assert row["cost_amount"] is None
assert row["upstream_api"] == "mock"
@@ -481,106 +403,6 @@ def proxy_completion(*_args, **_kwargs):
assert len(result["usage_record_ids"]) == 2
-def test_structured_cost_reports_unavailable_when_a_race_loser_cannot_be_measured() -> None:
- """Devin review (#955): a measured winner plus an unavailable race loser
- must not roll up into a confident "measured" total -- the aggregate cost
- status and amount need the same honesty precedence record_stream_usage
- already applies, not just the raw per-record ledger label.
- """
- coordinator = _coordinator()
-
- def proxy_completion(*_args, **_kwargs):
- coordinator.orchestrator._race_usage_sink(
- "mock_worker",
- ("duplicate", "mock_worker", None), # unparseable usage
- )
- return {
- "model": "mock-a",
- "usage": {"prompt_tokens": 7, "completion_tokens": 3},
- "orchestration": {"workflow_run_id": "run_unavailable_loser"},
- }
-
- coordinator.orchestrator.proxy_completion = proxy_completion # type: ignore[method-assign]
- coordinator.orchestrator.get_workflow_run = lambda _run_id: { # type: ignore[method-assign]
- "workflow_run_id": "run_unavailable_loser",
- "mode": "route",
- "answer": "winner",
- "trace": None,
- }
- result = coordinator.complete(
- [{"role": "user", "content": "race unavailable"}],
- provider_request={
- "model": "mock-a",
- "messages": [{"role": "user", "content": "race unavailable"}],
- "response_format": {"type": "json_object"},
- },
- )
-
- assert result["cost"]["measurement_status"] == "unavailable"
- assert result["cost"]["cost_amount"] is None
- assert "currency_components" not in result["cost"]
- records = coordinator.ledger.records()
- assert {record["measurement_status"] for record in records} == {"measured", "unavailable"}
-
-
-@pytest.mark.parametrize("structured", [False, True])
-def test_unavailable_mixed_currency_cost_suppresses_partial_components(
- structured: bool,
-) -> None:
- """Unknown loser spend must not expose complete-looking currency subtotals."""
- agents = [
- ModelAgent("winner_agent", "winner-model", provider_name="winner_provider"),
- ModelAgent("loser_agent", "loser-model", provider_name="loser_provider"),
- ]
- orchestrator = TaskOrchestrator(agents)
- config = InMemoryConfigStore()
- price_book = PriceBook(config)
- price_book.set_price(PriceEntry("winner_provider", "winner-model", 1, 1, "USD"))
- price_book.set_price(PriceEntry("loser_provider", "loser-model", 1, 1, "EUR"))
- coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
-
- def workflow():
- return {
- "workflow_run_id": "run_mixed_unavailable",
- "mode": "route",
- "answer": "winner",
- "trace": [{
- "agent_id": "winner_agent",
- "output": "winner",
- "usage": {"prompt_tokens": 7, "completion_tokens": 3},
- }],
- }
-
- def emit_loser():
- orchestrator._race_usage_sink(
- "loser_agent", ("duplicate", "loser_agent", None)
- )
-
- if structured:
- def proxy_completion(*_args, **_kwargs):
- emit_loser()
- return {"orchestration": {"workflow_run_id": "run_mixed_unavailable"}}
- orchestrator.proxy_completion = proxy_completion # type: ignore[method-assign]
- orchestrator.get_workflow_run = lambda _run_id: workflow() # type: ignore[method-assign]
- result = coordinator.complete(
- [{"role": "user", "content": "mixed"}],
- provider_request={"model": "winner-model", "messages": []},
- )
- else:
- def run(*_args, **_kwargs):
- emit_loser()
- return workflow()
- orchestrator.run = run # type: ignore[method-assign]
- result = coordinator.complete([{"role": "user", "content": "mixed"}])
-
- assert result["cost"] == {
- "cost_amount": None,
- "currency_code": "MIXED",
- "price_known": True,
- "measurement_status": "unavailable",
- }
-
-
def test_structured_provider_workflow_estimates_each_unreported_call() -> None:
"""Mixed usage bills each measured call plus one fallback prompt estimate."""
coordinator = _coordinator()
@@ -609,28 +431,29 @@ def test_structured_provider_workflow_estimates_each_unreported_call() -> None:
records = coordinator.ledger.records()
assert len(result["usage_record_ids"]) == len(records) == len(trace)
statuses = [record["measurement_status"] for record in records]
- assert statuses.count("estimated") == 2
- assert set(statuses) == {"measured", "estimated"}
- assert result["cost"]["measurement_status"] == "estimated"
- assert records[1]["total_tokens"] > 0
- assert records[3]["total_tokens"] > 0
- request_prompt = coordinator.token_counter.count_messages(
- [{"role": "user", "content": "return mixed usage JSON"}], "mock-a"
- )
- assert sum(record["prompt_tokens"] for record in records) == request_prompt + 5
+ assert statuses.count("unavailable") == 2
+ assert set(statuses) == {"measured", "unavailable"}
+ assert result["cost"]["measurement_status"] == "unavailable"
+ assert records[1]["total_tokens"] is None
+ assert records[3]["total_tokens"] is None
+ assert sum(
+ record["prompt_tokens"]
+ for record in records
+ if record["prompt_tokens"] is not None
+ ) == 5
assert [
record["prompt_tokens"]
for record in records
if record["measurement_status"] == "measured"
] == [2, 3, 0]
- estimated = [
- record for record in records if record["measurement_status"] == "estimated"
+ unavailable = [
+ record for record in records if record["measurement_status"] == "unavailable"
]
- assert [record["prompt_tokens"] for record in estimated] == [request_prompt, 0]
+ assert [record["prompt_tokens"] for record in unavailable] == [None, None]
-def test_unreported_provider_calls_bill_request_prompt_once() -> None:
- """One completion attributes its request prompt once, not once per unreported call."""
+def test_unreported_provider_calls_remain_unavailable() -> None:
+ """Missing provider usage never reconstructs prompt framing."""
coordinator = _coordinator()
coordinator.orchestrator.client.take_usage = lambda: None
@@ -646,15 +469,11 @@ def test_unreported_provider_calls_bill_request_prompt_once() -> None:
records = coordinator.ledger.records()
unreported = [
- record for record in records if record["measurement_status"] == "estimated"
+ record for record in records if record["measurement_status"] == "unavailable"
]
assert len(unreported) >= 2
- request_prompt = coordinator.token_counter.count_messages(messages, "mock-a")
- # The full request prompt lands on the first unreported step only; later
- # unreported steps estimate just their own output tokens.
- assert sum(record["prompt_tokens"] for record in unreported) == request_prompt
- assert unreported[0]["prompt_tokens"] == request_prompt
- assert all(record["prompt_tokens"] == 0 for record in unreported[1:])
+ assert all(record["prompt_tokens"] is None for record in unreported)
+ assert all(record["completion_tokens"] is None for record in unreported)
def test_structured_mixed_currency_costs_are_never_implicitly_converted() -> None:
@@ -724,7 +543,7 @@ def test_sync_completion_survives_usage_persistence_failure() -> None:
result = coordinator.complete([{"role": "user", "content": "hello without blocking"}])
assert result["channel"] == "sync"
assert result["usage_record_id"].startswith("usage_")
- assert result["usage"]["total_tokens"] > 0
+ assert result["usage"] is None
assert ledger.flush(timeout=1.0)
assert ledger.telemetry_health()["store_failures"] == len(result["usage_record_ids"])
@@ -736,9 +555,8 @@ def test_sync_completion_survives_usage_persistence_failure() -> None:
def test_batch_completion_records_on_retrieve() -> None:
coordinator = _coordinator()
- prompt = "bulk job please"
submitted = coordinator.complete(
- [{"role": "user", "content": prompt}],
+ [{"role": "user", "content": "bulk job please"}],
hints={"latency_tolerant": True},
attribution={"team": "beta", "company": "acme"},
)
@@ -753,359 +571,6 @@ def test_batch_completion_records_on_retrieve() -> None:
assert all(record["request_channel"] == "batch" for record in records)
assert all(record["team_name"] == "beta" for record in records)
- # The mock runner behind LocalBatchBackend reports no real per-step usage,
- # so this legitimately falls back to an estimate -- but it must be an
- # honest estimate of the *real* prompt threaded through
- # BatchResultItem.messages, not the old hardcoded blank placeholder
- # (which always computed exactly ``count_messages([{"content": ""}])``
- # tokens regardless of how long the actual prompt was).
- result = retrieved["results"][0]
- blank_prompt_tokens = coordinator.token_counter.count_messages(
- [{"role": "user", "content": ""}]
- )
- real_prompt_tokens = coordinator.token_counter.count_messages(
- [{"role": "user", "content": prompt}]
- )
- assert result["measurement_status"] == "estimated"
- assert result["prompt_tokens"] == real_prompt_tokens
- assert result["prompt_tokens"] != blank_prompt_tokens
- assert result["completion_tokens"] > 0
-
-
-def test_local_batch_records_each_served_provider_at_its_own_price() -> None:
- agents = [
- ModelAgent(id="worker_a", model="model-a", base_url="mock://a", provider_name="alpha"),
- ModelAgent(id="worker_b", model="model-b", base_url="mock://b", provider_name="beta"),
- ]
- orchestrator = TaskOrchestrator(agents)
- config = InMemoryConfigStore()
- price_book = PriceBook(config)
- price_book.set_price(PriceEntry("alpha", "model-a", 1.0, 2.0))
- price_book.set_price(PriceEntry("beta", "model-b", 3.0, 4.0))
- backend = LocalBatchBackend(
- lambda *_: {
- "answer": "done",
- "mode": "conduct",
- "trace": [
- {"agent_id": "worker_a", "usage": {"prompt_tokens": 10, "completion_tokens": 5}},
- {"agent_id": "worker_b", "usage": {"prompt_tokens": 7, "completion_tokens": 3}},
- ],
- }
- )
- coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, batch_backend=backend
- )
-
- submitted = coordinator.complete(
- [{"role": "user", "content": "meter both"}], hints={"channel": "batch"}
- )
- result = coordinator.retrieve_batch(submitted["job_id"])["results"][0]
- rows = coordinator.ledger.records()
-
- assert [(row["provider_name"], row["model_name"]) for row in rows] == [
- ("alpha", "model-a"),
- ("beta", "model-b"),
- ]
- assert result["usage_record_ids"] == [row["usage_record_id"] for row in rows]
- assert result["prompt_tokens"] == 17
- assert result["completion_tokens"] == 8
- assert result["cost_amount"] == 0.053
- assert result["measurement_status"] == "measured"
-
-
-def test_local_batch_malformed_usage_falls_back_without_coercion() -> None:
- coordinator = _coordinator()
- coordinator.batch_backend = LocalBatchBackend(
- lambda *_: {
- "answer": "done",
- "trace": [
- {
- "agent_id": "mock_worker",
- "output": "fallback",
- "usage": {"prompt_tokens": None, "completion_tokens": 3},
- },
- {
- "agent_id": "mock_worker",
- "output": "again",
- "usage": {"prompt_tokens": "7", "completion_tokens": "3"},
- },
- ],
- }
- )
-
- submitted = coordinator.complete(
- [{"role": "user", "content": "do not coerce"}], hints={"channel": "batch"}
- )
- result = coordinator.retrieve_batch(submitted["job_id"])["results"][0]
-
- assert result["measurement_status"] == "estimated"
- assert all(
- row["measurement_status"] == "estimated"
- for row in coordinator.ledger.records()
- )
-
-
-def test_custom_batch_backend_preserves_one_sided_zero_usage() -> None:
- class _Backend:
- name = "custom"
- def submit(self, requests, metadata=None):
- return BatchJob("custom-zero", self.name, request_count=len(requests))
- def retrieve(self, job):
- return [BatchResultItem("request-zero", "answer", 0, 5)]
- coordinator = _coordinator()
- coordinator.batch_backend = _Backend()
- job = coordinator.complete([{"role": "user", "content": "zero prompt"}],
- hints={"channel": "batch"})
- result = coordinator.retrieve_batch(job["job_id"])["results"][0]
- assert result["measurement_status"] == "measured"
- assert (result["prompt_tokens"], result["completion_tokens"]) == (0, 5)
-
-
-def test_custom_batch_backend_rejects_negative_reported_usage() -> None:
- class _Backend:
- name = "custom"
- def submit(self, requests, metadata=None):
- return BatchJob("custom-negative", self.name, request_count=len(requests))
- def retrieve(self, job):
- return [BatchResultItem(
- "request-negative", "answer", -1, 5, usage_valid=True,
- messages=[{"role": "user", "content": "real prompt"}],
- )]
- coordinator = _coordinator()
- coordinator.batch_backend = _Backend()
- job = coordinator.complete(
- [{"role": "user", "content": "real prompt"}], hints={"channel": "batch"}
- )
-
- result = coordinator.retrieve_batch(job["job_id"])["results"][0]
-
- assert result["measurement_status"] == "estimated"
- assert result["prompt_tokens"] >= 0
- assert result["completion_tokens"] >= 0
-
-
-def test_multi_item_batch_settlement_queries_ledger_once(monkeypatch) -> None:
- class _Backend:
- name = "custom"
- def submit(self, requests, metadata=None):
- return BatchJob("custom-multi", self.name, request_count=len(requests))
- def retrieve(self, job):
- return [
- BatchResultItem("request-one", "one", 1, 1),
- BatchResultItem("request-two", "two", 2, 2),
- ]
- coordinator = _coordinator()
- coordinator.batch_backend = _Backend()
- calls = 0
- records = coordinator.ledger.records
-
- def counted_records(*args, **kwargs):
- nonlocal calls
- calls += 1
- return records(*args, **kwargs)
-
- monkeypatch.setattr(coordinator.ledger, "records", counted_records)
- job = coordinator.submit_batch([
- BatchRequest("request-one", [{"role": "user", "content": "one"}]),
- BatchRequest("request-two", [{"role": "user", "content": "two"}]),
- ])
-
- document = coordinator.retrieve_batch(job.job_id)
-
- assert document["result_count"] == 2
- assert document["usage_persistence_status"] == "settled"
- assert calls == 1
-
-
-def test_batch_retrieval_does_not_wait_forever_for_ledger_storage(monkeypatch) -> None:
- release = threading.Event()
-
- class _BlockedStore:
- def __init__(self):
- self.rows = []
- def append(self, record):
- release.wait(timeout=2)
- self.rows.append(record.as_dict())
- def query(self, start=None, end=None):
- return list(self.rows)
-
- config = InMemoryConfigStore()
- price_book = PriceBook(config)
- ledger = CostLedger(
- price_book,
- store=NonBlockingLedgerStore(_BlockedStore()),
- )
- coordinator = _coordinator(ledger=ledger)
- monkeypatch.setattr(
- "contextual_orchestrator.cost_router._BATCH_LEDGER_SETTLEMENT_TIMEOUT_SECONDS",
- 0.01,
- )
- job = coordinator.complete(
- [{"role": "user", "content": "blocked ledger"}], hints={"channel": "batch"}
- )
- try:
- document = coordinator.retrieve_batch(job["job_id"])
- finally:
- release.set()
-
- assert document["usage_persistence_status"] == "pending"
- assert document["result_count"] == 1
- assert document["results"][0]["measurement_status"] == "estimated"
- assert ledger.flush(timeout=1.0)
- settled = coordinator.retrieve_batch(job["job_id"])
- assert settled["usage_persistence_status"] == "settled"
- assert settled["results"] == document["results"]
-
-
-def test_batch_settlement_ignores_unrelated_global_flush_work(monkeypatch) -> None:
- release = threading.Event()
- unrelated_started = threading.Event()
-
- class _Store:
- def __init__(self):
- self.rows = []
- def append(self, record):
- self.rows.append(record.as_dict())
- if record.usage_record_id == "usage_unrelated":
- unrelated_started.set()
- release.wait(timeout=2)
- def query(self, start=None, end=None):
- return list(self.rows)
- def existing_usage_record_ids(self, usage_record_ids):
- return {row["usage_record_id"] for row in self.rows
- if row["usage_record_id"] in usage_record_ids}
-
- ledger = CostLedger(PriceBook(InMemoryConfigStore()),
- store=NonBlockingLedgerStore(_Store()))
- coordinator = _coordinator(ledger=ledger)
- wait_for_ids = ledger.wait_for_usage_record_ids
- def wait_with_unrelated_work(usage_record_ids, *, timeout=None):
- ledger.record_usage(provider="mock", model="mock", prompt_tokens=1,
- completion_tokens=1, usage_record_id="usage_unrelated")
- assert unrelated_started.wait(timeout=1)
- return wait_for_ids(usage_record_ids, timeout=timeout)
- monkeypatch.setattr(ledger, "wait_for_usage_record_ids", wait_with_unrelated_work)
- job = coordinator.complete([{"role": "user", "content": "batch settlement"}],
- hints={"channel": "batch"})
- try:
- document = coordinator.retrieve_batch(job["job_id"])
- finally:
- release.set()
- assert document["usage_persistence_status"] == "settled"
-
-
-def test_rejected_async_batch_write_remains_pending_and_idempotent(monkeypatch) -> None:
- class _RejectingStore:
- def append(self, record):
- return False
- def query(self, start=None, end=None):
- return []
- def existing_usage_record_ids(self, usage_record_ids):
- return set()
-
- ledger = CostLedger(PriceBook(InMemoryConfigStore()),
- store=NonBlockingLedgerStore(_RejectingStore()))
- coordinator = _coordinator(ledger=ledger)
- monkeypatch.setattr(
- "contextual_orchestrator.cost_router._BATCH_LEDGER_SETTLEMENT_TIMEOUT_SECONDS", 0.01
- )
- job = coordinator.complete([{"role": "user", "content": "rejected write"}],
- hints={"channel": "batch"})
- first = coordinator.retrieve_batch(job["job_id"])
- second = coordinator.retrieve_batch(job["job_id"])
- assert first["usage_persistence_status"] == "pending"
- assert second["usage_persistence_status"] == "pending"
- assert second["results"][0]["usage_record_ids"] == first["results"][0]["usage_record_ids"]
- assert ledger.records() == []
-
-
-def test_local_batch_cache_hit_does_not_rebill_provider() -> None:
- coordinator = _coordinator()
- coordinator.batch_backend = LocalBatchBackend(lambda *_: {
- "answer": "cached", "cache_status": "hit",
- "trace": [{"agent_id": "mock_worker", "usage": {
- "prompt_tokens": 20, "completion_tokens": 10}}],
- })
- job = coordinator.complete([{"role": "user", "content": "cached"}],
- hints={"channel": "batch"})
- result = coordinator.retrieve_batch(job["job_id"])["results"][0]
- assert (result["prompt_tokens"], result["completion_tokens"], result["cost_amount"]) == (0, 0, 0)
- assert [(row["provider_name"], row["request_channel"])
- for row in coordinator.ledger.records()] == [("cache", "cache")]
-
-
-def test_local_batch_records_race_loser_once_across_retrievals() -> None:
- coordinator = _coordinator()
- def raced_complete(*_args, **_kwargs):
- coordinator.orchestrator._race_usage_sink(
- "mock_worker", ("loser", "mock_worker", {
- "prompt_tokens": 2, "completion_tokens": 3}))
- return {"answer": "winner", "trace": [{"agent_id": "mock_worker", "usage": {
- "prompt_tokens": 5, "completion_tokens": 7}}]}
- coordinator.orchestrator.complete = raced_complete
- job = coordinator.complete([{"role": "user", "content": "race"}],
- hints={"channel": "batch"})
- first = coordinator.retrieve_batch(job["job_id"])
- assert first == coordinator.retrieve_batch(job["job_id"])
- assert (first["results"][0]["prompt_tokens"],
- first["results"][0]["completion_tokens"]) == (7, 10)
- assert len(coordinator.ledger.records()) == 2
-
-
-def test_pg_batch_malformed_usage_estimates_real_prompt() -> None:
- captured = {}
- class _Assembler:
- def assemble(self, lines):
- captured["lines"] = lines
- return "memory://captured"
- class _Client:
- async def upload_jsonl(self, *_args):
- return {"id": "file-1"}
- async def create_batch_job(self, *_args, **_kwargs):
- return {"id": "batch-malformed", "status": "completed"}
- async def download_results(self, *_args):
- return {"success": True, "responses": [{
- "custom_id": captured["lines"][0]["custom_id"],
- "response": {"body": {
- "choices": [{"message": {"content": "answer"}}],
- "usage": {"prompt_tokens": "9", "completion_tokens": None},
- }},
- }]}
- coordinator = _coordinator()
- coordinator.batch_backend = PgLlmBatchBackend(_Client(), payload_assembler=_Assembler())
- messages = [{"role": "user", "content": "the real remote prompt"}]
- job = coordinator.complete(messages, hints={"channel": "batch"})
- result = coordinator.retrieve_batch(job["job_id"])["results"][0]
- assert result["measurement_status"] == "estimated"
- assert result["prompt_tokens"] == coordinator.token_counter.count_messages(
- messages, "contextual-orchestrator")
-
-
-def test_batch_mixed_currency_and_retrieval_idempotency() -> None:
- agents = [
- ModelAgent("worker_alpha", "model-a", "mock://a", provider_name="alpha"),
- ModelAgent("worker_beta", "model-b", "mock://b", provider_name="beta"),
- ]
- config = InMemoryConfigStore()
- prices = PriceBook(config)
- prices.set_price(PriceEntry("alpha", "model-a", 1, 1, currency_code="USD"))
- prices.set_price(PriceEntry("beta", "model-b", 1, 1, currency_code="EUR"))
- coordinator = CostRoutingCoordinator(TaskOrchestrator(agents), config,
- price_book=prices, batch_backend=LocalBatchBackend(lambda *_: {
- "answer": "done", "trace": [
- {"agent_id": "worker_alpha", "usage": {"prompt_tokens": 1, "completion_tokens": 1}},
- {"agent_id": "worker_beta", "usage": {"prompt_tokens": 2, "completion_tokens": 2}},
- ]}))
- job = coordinator.complete([{"role": "user", "content": "mixed"}],
- hints={"channel": "batch"})
- first = coordinator.retrieve_batch(job["job_id"])
- prices.set_price(PriceEntry("alpha", "model-a", 99, 99, currency_code="USD"))
- prices.set_price(PriceEntry("beta", "model-b", 99, 99, currency_code="EUR"))
- assert first == coordinator.retrieve_batch(job["job_id"])
- assert [part["currency_code"] for part in
- first["results"][0]["currency_components"]] == ["EUR", "USD"]
- assert len(coordinator.ledger.records()) == 2
-
def test_default_local_batch_backend_reuses_orchestrator_concurrency() -> None:
class _Client:
@@ -1129,12 +594,10 @@ def test_cost_report_rolls_up_across_sync_and_batch() -> None:
)
job = coordinator.complete([{"role": "user", "content": "batch one"}],
hints={"channel": "batch"}, attribution={"company": "acme"})
- batch = coordinator.retrieve_batch(job["job_id"])
+ coordinator.retrieve_batch(job["job_id"])
report = coordinator.cost_report("company")
- assert report["grand_total"]["record_count"] == (
- len(sync["usage_record_ids"]) + len(batch["results"][0]["usage_record_ids"])
- )
+ assert report["grand_total"]["record_count"] == len(coordinator.ledger.records())
assert report["items"][0]["dimension_value"] == "acme"
@@ -1196,47 +659,6 @@ async def download_results(self, batch_id, endpoint_alias):
assert row["request_channel"] == "batch"
-def test_retrieve_batch_propagates_download_failure_instead_of_fake_empty_success() -> None:
- """An explicit download failure must never be reported as a zero-result success.
-
- Regression for the bug where ``PgLlmBatchBackend.retrieve()`` mapped
- ``success: False`` to ``[]``, which ``retrieve_batch()`` then returned as
- an ordinary ``result_count: 0`` success -- indistinguishable from a batch
- that legitimately completed with nothing to report.
- """
-
- class _FailingClient:
- async def upload_jsonl(self, file_path, endpoint_alias, purpose="batch"):
- return {"id": "file-1"}
-
- async def create_batch_job(self, input_file_id, endpoint_alias, endpoint="/v1/chat/completions", metadata=None):
- return {"id": "batch-failed", "status": "validating"}
-
- async def get_batch_status(self, batch_id, endpoint_alias):
- return {"status": "completed", "is_complete": True}
-
- async def download_results(self, batch_id, endpoint_alias):
- return {"success": False, "reason": "Batch not complete"}
-
- agents = [ModelAgent(id="mock_worker", model="mock-a", base_url="mock://a", provider_name="mock",
- tags=("reasoning",), priority=1)]
- orchestrator = TaskOrchestrator(agents)
- config = InMemoryConfigStore()
- price_book = PriceBook(config)
- price_book.set_price(PriceEntry("mock", "*", 1.0, 2.0))
- backend = PgLlmBatchBackend(_FailingClient())
- coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book, batch_backend=backend)
-
- submitted = coordinator.complete([{"role": "user", "content": "route to pg-llm-batch"}],
- hints={"channel": "batch"}, attribution={"provider": "mock"})
- with pytest.raises(BatchDownloadError) as excinfo:
- coordinator.retrieve_batch(submitted["job_id"])
- assert excinfo.value.job_id == submitted["job_id"]
- assert excinfo.value.reason == "Batch not complete"
- # No usage was recorded for the failed retrieval.
- assert coordinator.ledger.records() == []
-
-
def test_zdr_batch_resolves_each_request_to_a_member_of_its_configured_pool() -> None:
non_zdr = ModelAgent(
"non_zdr_member",
@@ -1344,8 +766,20 @@ def retrieve(self, job):
def fail_reselection(*_args, **_kwargs):
raise AssertionError("the selected embedding member must not be re-resolved")
+ class _SyntheticExactCounter:
+ def count_text(self, text, model):
+ return 1
+
orchestrator.select_capability_agent = fail_reselection # type: ignore[method-assign]
- coordinator = CostRoutingCoordinator(orchestrator, embedding_batch_backend=backend)
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ embedding_batch_backend=backend,
+ embedding_token_counter=type(
+ "ExactSyntheticCounter",
+ (),
+ {"count_text": lambda self, text, model="": len(text)},
+ )(),
+ )
document = coordinator.complete_embeddings_batch(
["private"], model=second.model, zdr_only=True, agent_id=second.id
)
@@ -1355,16 +789,87 @@ def fail_reselection(*_args, **_kwargs):
assert backend.requests[0].agent_id == second.id
+def test_provider_embedding_runner_accepts_empty_direct_batch() -> None:
+ """The provider backend's direct contract returns an empty batch without indexing it."""
+ agent = ModelAgent(
+ "remote_embedding",
+ "synthetic-embedding",
+ "https://provider.synthetic.invalid/v1",
+ tags=("embedding",),
+ )
+ orchestrator = TaskOrchestrator([agent])
+ orchestrator.client.embed = lambda *_args, **_kwargs: (_ for _ in ()).throw(
+ AssertionError("an empty batch must not call the provider")
+ )
+ coordinator = CostRoutingCoordinator(orchestrator)
+ backend = coordinator.embedding_batch_backend
+ job = backend.submit([])
+ deadline = time.time() + 2
+ while backend.poll(job)["status"] not in {"completed", "failed"} and time.time() < deadline:
+ time.sleep(0.01)
+ assert backend.poll(job)["status"] == "completed"
+ assert backend.retrieve(job) == []
+
+
+def test_virtual_embedding_model_binds_concrete_tokenizer_before_accounting() -> None:
+ """Default embedding traffic counts against the selected provider model."""
+ agent = ModelAgent(
+ "remote_embedding",
+ "text-embedding-3-small",
+ "https://provider.synthetic.invalid/v1",
+ tags=("embedding",),
+ )
+
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent]),
+ embedding_token_counter=type(
+ "ExactSyntheticCounter",
+ (),
+ {"count_text": lambda self, text, model="": len(text)},
+ )(),
+ )
+
+ resolved = coordinator._resolve_embedding_target(
+ "contextual-orchestrator", False, None
+ )
+
+ assert resolved == ("text-embedding-3-small", "remote_embedding")
+
+
+def test_non_zdr_batch_preserves_an_explicit_model_outside_the_pool() -> None:
+ """The ZDR resolver must not change ordinary batch passthrough behavior."""
+ captured: list[BatchRequest] = []
+
+ class _CapturingBackend:
+ name = "capturing"
+
+ def submit(self, requests, metadata=None):
+ captured.extend(requests)
+ return BatchJob("batch-ordinary", self.name, status="submitted", request_count=len(requests))
+
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([ModelAgent("configured_agent", "configured-model", "mock://configured")]),
+ batch_backend=_CapturingBackend(),
+ )
+ request = BatchRequest(
+ messages=[{"role": "user", "content": "ordinary batch"}],
+ model="unconfigured-provider-model",
+ )
+
+ coordinator.submit_batch([request])
+
+ assert captured == [request]
+
+
+if __name__ == "__main__": # pragma: no cover
+ for _name, _fn in sorted(globals().items()):
+ if _name.startswith("test_") and callable(_fn):
+ _fn()
+ print(f"ok {_name}")
+ print("ok")
+
+
def test_embedding_batch_selects_cheapest_capability_candidate_when_unspecified() -> None:
- """Price-aware selection: an unspecified member picks price, not rank order.
-
- ``ranked_first`` outranks ``cheaper`` under the orchestrator's own
- priority-based ordering (verified below), so a price-blind ``candidates[0]``
- pick would return it. The coordinator must instead resolve to the cheaper
- member via ``_cheapest_capability_candidate``'s direct
- ``PriceBook.get_price()`` lookup and raw ``prompt_price_per_1k``
- comparison (it does not call ``cheapest_upstream``).
- """
from contextual_orchestrator.batch_routing import EmbeddingBatchResultItem
ranked_first = ModelAgent(
@@ -1415,7 +920,15 @@ def retrieve(self, job):
PriceEntry("cheap-provider", "cheap-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0)
)
coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, embedding_batch_backend=backend
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_batch_backend=backend,
+ embedding_token_counter=type(
+ "ExactSyntheticCounter",
+ (),
+ {"count_text": lambda self, text, model="": len(text)},
+ )(),
)
document = coordinator.complete_embeddings_batch(["private"], zdr_only=True)
@@ -1426,12 +939,6 @@ def retrieve(self, job):
def test_cheapest_capability_candidate_prefers_known_price_over_unpriced() -> None:
- """An unpriced member must not win as a false zero-cost candidate.
-
- ``ranked_first`` outranks ``priced`` under the orchestrator's own
- priority order but carries no price-book entry at all; it must lose to
- the known, paid ``priced`` member rather than being treated as free.
- """
ranked_first = ModelAgent(
"ranked_first_member", "unpriced-embedding", "mock://unpriced",
provider_name="unpriced-provider", tags=("embedding",), priority=10,
@@ -1448,7 +955,7 @@ def test_cheapest_capability_candidate_prefers_known_price_over_unpriced() -> No
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(
- PriceEntry("paid-provider", "paid-embedding", prompt_price_per_1k=5.0, completion_price_per_1k=0.0)
+ PriceEntry("paid-provider", "paid-embedding", prompt_price_per_1k=1.0, completion_price_per_1k=0.0)
)
coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
@@ -1457,165 +964,96 @@ def test_cheapest_capability_candidate_prefers_known_price_over_unpriced() -> No
assert chosen.id == priced.id
-def test_cheapest_capability_candidate_ignores_unpriced_and_keeps_ranked_order() -> None:
- """An all-unpriced pool must preserve the pre-existing ranked order, not always pick index 0."""
- ranked_first = ModelAgent(
- "ranked_first_member", "unpriced-one", "mock://one",
- provider_name="unpriced-provider-one", tags=("embedding",), priority=10,
- )
- ranked_second = ModelAgent(
- "ranked_second_member", "unpriced-two", "mock://two",
- provider_name="unpriced-provider-two", tags=("embedding",), priority=1,
- )
- orchestrator = TaskOrchestrator([ranked_first, ranked_second])
- with orchestrator.request_policy(False):
- candidates = orchestrator._capability_agents("embedding", None)
- assert [agent.id for agent in candidates] == [ranked_first.id, ranked_second.id]
-
- coordinator = CostRoutingCoordinator(orchestrator, InMemoryConfigStore())
-
- chosen = coordinator._cheapest_capability_candidate(candidates)
-
- assert chosen.id == ranked_first.id
-
-
-def test_cheapest_capability_candidate_ignores_mismatched_currency() -> None:
- """A different-currency price must not be compared to a default-currency price by face value.
-
- ``foreign_member`` has a numerically smaller price but in a currency this
- repo has no exchange rate for; it must lose to the known, comparable
- ``domestic_member`` price rather than win on raw face value.
- """
- foreign_member = ModelAgent(
- "foreign_priced_member", "foreign-embedding", "mock://foreign",
- provider_name="foreign-provider", tags=("embedding",), priority=10,
+def test_cheapest_capability_candidate_respects_health_filter() -> None:
+ sick_cheaper = ModelAgent(
+ "sick_cheaper_member", "cheap-embedding", "mock://cheap",
+ provider_name="cheap-provider", tags=("embedding",), priority=1,
)
- domestic_member = ModelAgent(
- "domestic_priced_member", "domestic-embedding", "mock://domestic",
- provider_name="domestic-provider", tags=("embedding",), priority=1,
+ healthy_expensive = ModelAgent(
+ "healthy_expensive_member", "expensive-embedding", "mock://expensive",
+ provider_name="expensive-provider", tags=("embedding",), priority=10,
)
- orchestrator = TaskOrchestrator([foreign_member, domestic_member])
- with orchestrator.request_policy(False):
- candidates = orchestrator._capability_agents("embedding", None)
-
+ orchestrator = TaskOrchestrator([sick_cheaper, healthy_expensive])
config = InMemoryConfigStore()
- price_book = PriceBook(config) # default_currency == "USD"
+ price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "foreign-provider", "foreign-embedding",
- prompt_price_per_1k=1.0, completion_price_per_1k=0.0, currency_code="JPY",
- )
+ PriceEntry("cheap-provider", "cheap-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0)
)
price_book.set_price(
- PriceEntry(
- "domestic-provider", "domestic-embedding",
- prompt_price_per_1k=5.0, completion_price_per_1k=0.0, currency_code="USD",
- )
+ PriceEntry("expensive-provider", "expensive-embedding", prompt_price_per_1k=5.0, completion_price_per_1k=0.0)
)
coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
- chosen = coordinator._cheapest_capability_candidate(candidates)
+ with orchestrator.request_policy(False):
+ candidates = orchestrator._capability_agents("embedding", None)
+
+ for _ in range(5):
+ orchestrator._group_router.observe_failure(sick_cheaper.id)
- assert chosen.id == domestic_member.id
+ ordered = coordinator._cost_ordered_capability_candidates(candidates)
+ assert [agent.id for agent in ordered] == [healthy_expensive.id, sick_cheaper.id]
-def test_cheapest_capability_candidate_ignores_completion_price_for_embeddings() -> None:
- """Embedding routing must not price nonexistent completion tokens.
- ``cheap_input_member`` has the lowest true (input-only) cost but a huge
- completion price that would dominate under a nonzero
- ``assumed_completion_tokens``; it must still win because embeddings never
- consume completion tokens.
- """
- cheap_input_member = ModelAgent(
- "cheap_input_member", "cheap-input-embedding", "mock://cheap-input",
- provider_name="cheap-input-provider", tags=("embedding",), priority=10,
+def test_cheapest_capability_candidate_ignores_uncomparable_currency() -> None:
+ ranked_first = ModelAgent(
+ "ranked_first_member", "eur-embedding", "mock://eur",
+ provider_name="eur-provider", tags=("embedding",), priority=10,
)
- expensive_input_member = ModelAgent(
- "expensive_input_member", "expensive-input-embedding", "mock://expensive-input",
- provider_name="expensive-input-provider", tags=("embedding",), priority=1,
+ cheaper_mismatched = ModelAgent(
+ "cheaper_mismatched_member", "gbp-embedding", "mock://gbp",
+ provider_name="gbp-provider", tags=("embedding",), priority=1,
)
- orchestrator = TaskOrchestrator([cheap_input_member, expensive_input_member])
+ orchestrator = TaskOrchestrator([ranked_first, cheaper_mismatched])
with orchestrator.request_policy(False):
candidates = orchestrator._capability_agents("embedding", None)
+ assert [agent.id for agent in candidates] == [ranked_first.id, cheaper_mismatched.id]
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "cheap-input-provider", "cheap-input-embedding",
- prompt_price_per_1k=0.01, completion_price_per_1k=100.0,
- )
+ PriceEntry("eur-provider", "eur-embedding", prompt_price_per_1k=1.0, completion_price_per_1k=0.0, currency_code="EUR")
)
price_book.set_price(
- PriceEntry(
- "expensive-input-provider", "expensive-input-embedding",
- prompt_price_per_1k=1.0, completion_price_per_1k=0.0,
- )
+ PriceEntry("gbp-provider", "gbp-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0, currency_code="GBP")
)
coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
chosen = coordinator._cheapest_capability_candidate(candidates)
- assert chosen.id == cheap_input_member.id
-
+ assert chosen.id == ranked_first.id
-def test_cheapest_capability_candidate_compares_equivalent_currency_spellings() -> None:
- """Lowercase/padded currency codes must still compare as the same currency.
- ``padded_lowercase_member`` has the true lowest price, but its price
- entry stores the currency as ``" usd "`` (lowercase, padded) rather
- than the price book's canonical ``"USD"``. An exact-string currency
- comparison would wrongly treat that as a different, non-comparable
- currency and let the costlier ``canonical_member`` win instead. This
- mirrors ``model_discovery._currency_is_comparable``'s own
- non-empty/trimmed/case-insensitive normalization.
- """
- padded_lowercase_member = ModelAgent(
- "padded_lowercase_member", "padded-embedding", "mock://padded",
- provider_name="padded-provider", tags=("embedding",), priority=10,
+def test_cheapest_capability_candidate_compares_same_currency_case_insensitively() -> None:
+ ranked_first = ModelAgent(
+ "ranked_first_member", "upper-embedding", "mock://upper",
+ provider_name="upper-provider", tags=("embedding",), priority=10,
)
- canonical_member = ModelAgent(
- "canonical_member", "canonical-embedding", "mock://canonical",
- provider_name="canonical-provider", tags=("embedding",), priority=1,
+ cheaper = ModelAgent(
+ "cheaper_member", "lower-embedding", "mock://lower",
+ provider_name="lower-provider", tags=("embedding",), priority=1,
)
- orchestrator = TaskOrchestrator([padded_lowercase_member, canonical_member])
+ orchestrator = TaskOrchestrator([ranked_first, cheaper])
with orchestrator.request_policy(False):
candidates = orchestrator._capability_agents("embedding", None)
+ assert [agent.id for agent in candidates] == [ranked_first.id, cheaper.id]
config = InMemoryConfigStore()
- price_book = PriceBook(config) # default_currency == "USD"
+ price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "padded-provider", "padded-embedding",
- prompt_price_per_1k=0.01, completion_price_per_1k=0.0, currency_code=" usd ",
- )
+ PriceEntry("upper-provider", "upper-embedding", prompt_price_per_1k=5.0, completion_price_per_1k=0.0, currency_code="USD")
)
price_book.set_price(
- PriceEntry(
- "canonical-provider", "canonical-embedding",
- prompt_price_per_1k=5.0, completion_price_per_1k=0.0, currency_code="USD",
- )
+ PriceEntry("lower-provider", "lower-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0, currency_code=" usd ")
)
coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
chosen = coordinator._cheapest_capability_candidate(candidates)
- assert chosen.id == padded_lowercase_member.id
-
+ assert chosen.id == cheaper.id
-def test_cheapest_capability_candidate_breaks_ledger_rounding_ties_on_raw_price() -> None:
- """Devin round-3 bug: tiny embedding prices must not collapse into ties.
- ``ranked_first`` (0.00000049 per 1K) and ``cheaper`` (0.00000001 per 1K)
- are genuinely different prices, but ``PriceBook.compute_cost`` quantizes
- to six decimal places for ledger reporting, so both round to the same
- ``0.0`` cost for the assumed 1,000-token request that ``cheapest_upstream``
- prices candidates against. A comparison that goes through that rounded
- cost sees a tie and keeps the ranked-first (here, pricier) candidate;
- the fix must instead compare the raw, unrounded ``prompt_price_per_1k``
- and pick the actually-cheaper member.
- """
+def test_cheapest_capability_candidate_differentiates_sub_cent_prices() -> None:
ranked_first = ModelAgent(
"ranked_first_member", "tiny-expensive-embedding", "mock://tiny-expensive",
provider_name="tiny-expensive-provider", tags=("embedding",), priority=10,
@@ -1632,27 +1070,11 @@ def test_cheapest_capability_candidate_breaks_ledger_rounding_ties_on_raw_price(
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "tiny-expensive-provider", "tiny-expensive-embedding",
- prompt_price_per_1k=0.00000049, completion_price_per_1k=0.0,
- )
+ PriceEntry("tiny-expensive-provider", "tiny-expensive-embedding", prompt_price_per_1k=0.00000049, completion_price_per_1k=0.0)
)
price_book.set_price(
- PriceEntry(
- "tiny-cheap-provider", "tiny-cheap-embedding",
- prompt_price_per_1k=0.00000001, completion_price_per_1k=0.0,
- )
- )
- # Confirm both really do collapse to the same rounded ledger cost, so
- # this test would have failed against the pre-fix rounded comparison.
- expensive_cost, *_ = price_book.compute_cost(
- "tiny-expensive-provider", "tiny-expensive-embedding", 1000, 0
+ PriceEntry("tiny-cheap-provider", "tiny-cheap-embedding", prompt_price_per_1k=0.00000001, completion_price_per_1k=0.0)
)
- cheap_cost, *_ = price_book.compute_cost(
- "tiny-cheap-provider", "tiny-cheap-embedding", 1000, 0
- )
- assert expensive_cost == cheap_cost == 0.0
-
coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
chosen = coordinator._cheapest_capability_candidate(candidates)
@@ -1661,15 +1083,6 @@ def test_cheapest_capability_candidate_breaks_ledger_rounding_ties_on_raw_price(
def test_non_zdr_embedding_batch_selects_cheapest_capability_candidate_when_unspecified() -> None:
- """Devin bug: ordinary (non-ZDR) unspecified embedding batches must also price-route.
-
- Previously ``_resolve_embedding_target`` returned before
- ``_cheapest_capability_candidate`` ever ran whenever ``zdr_only=False``,
- so only ZDR embedding batches were cost-aware; ordinary batches always
- kept the orchestrator's ranked (not price) order. An unspecified
- (auto/group) model with multiple differently priced non-ZDR group
- members must now resolve to the cheaper one.
- """
from contextual_orchestrator.batch_routing import EmbeddingBatchResultItem
ranked_first = ModelAgent(
@@ -1691,9 +1104,7 @@ def __init__(self) -> None:
def submit(self, requests, metadata=None):
self.requests.extend(requests)
- return BatchJob(
- "non-zdr-cheapest", self.name, status="completed", request_count=len(requests)
- )
+ return BatchJob("non-zdr-cheapest", self.name, status="completed", request_count=len(requests))
def poll(self, job):
return {"is_complete": True, "status": "completed"}
@@ -1713,22 +1124,23 @@ def retrieve(self, job):
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "expensive-plain-provider", "expensive-plain-embedding",
- prompt_price_per_1k=5.0, completion_price_per_1k=0.0,
- )
+ PriceEntry("expensive-plain-provider", "expensive-plain-embedding", prompt_price_per_1k=5.0, completion_price_per_1k=0.0)
)
price_book.set_price(
- PriceEntry(
- "cheap-plain-provider", "cheap-plain-embedding",
- prompt_price_per_1k=0.01, completion_price_per_1k=0.0,
- )
+ PriceEntry("cheap-plain-provider", "cheap-plain-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0)
)
coordinator = CostRoutingCoordinator(
- orchestrator, config, price_book=price_book, embedding_batch_backend=backend
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_batch_backend=backend,
+ embedding_token_counter=type(
+ "ExactSyntheticCounter",
+ (),
+ {"count_text": lambda self, text, model="": len(text)},
+ )(),
)
- # zdr_only defaults False and agent_id defaults None: the ordinary path.
coordinator.submit_embeddings_batch(["ordinary input"])
assert backend.requests[0].model == cheaper.model
@@ -1736,9 +1148,6 @@ def retrieve(self, job):
def test_non_zdr_complete_embeddings_batch_selects_cheapest_capability_candidate() -> None:
- """Same fix, exercised through ``complete_embeddings_batch`` — the public,
- synchronous coordinator entry point production callers (e.g. naruon's
- batch embedding service) call directly."""
ranked_first = ModelAgent(
"ranked_first_public_member", "expensive-public-embedding", "mock://expensive-public",
provider_name="expensive-public-provider", tags=("embedding",), priority=10,
@@ -1752,18 +1161,21 @@ def test_non_zdr_complete_embeddings_batch_selects_cheapest_capability_candidate
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(
- PriceEntry(
- "expensive-public-provider", "expensive-public-embedding",
- prompt_price_per_1k=5.0, completion_price_per_1k=0.0,
- )
+ PriceEntry("expensive-public-provider", "expensive-public-embedding", prompt_price_per_1k=5.0, completion_price_per_1k=0.0)
)
price_book.set_price(
- PriceEntry(
- "cheap-public-provider", "cheap-public-embedding",
- prompt_price_per_1k=0.01, completion_price_per_1k=0.0,
- )
+ PriceEntry("cheap-public-provider", "cheap-public-embedding", prompt_price_per_1k=0.01, completion_price_per_1k=0.0)
+ )
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_token_counter=type(
+ "ExactSyntheticCounter",
+ (),
+ {"count_text": lambda self, text, model="": len(text)},
+ )(),
)
- coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
document = coordinator.complete_embeddings_batch(["ordinary input"])
@@ -1771,14 +1183,6 @@ def test_non_zdr_complete_embeddings_batch_selects_cheapest_capability_candidate
def test_non_zdr_embedding_batch_preserves_explicit_model_outside_the_pool() -> None:
- """An explicit model absent from the pool must still pass through unresolved.
-
- Widening ``_resolve_embedding_target`` to cost-route unspecified
- non-ZDR requests must not force *every* ordinary embedding request
- through ``_capability_agents`` — only the unspecified (auto/group
- placeholder) case should. An explicit model outside the configured pool
- has no matching agent to look up and must not raise or be rewritten.
- """
orchestrator = TaskOrchestrator(
[ModelAgent("configured_agent", "configured-embedding", tags=("embedding",))]
)
@@ -1793,8 +1197,7 @@ def test_non_zdr_embedding_batch_preserves_explicit_model_outside_the_pool() ->
def test_non_zdr_batch_preserves_an_explicit_model_outside_the_pool() -> None:
- """The ZDR resolver must not change ordinary batch passthrough behavior."""
- captured: list[BatchRequest] = []
+ captured = []
class _CapturingBackend:
name = "capturing"
@@ -1814,12 +1217,5 @@ def submit(self, requests, metadata=None):
coordinator.submit_batch([request])
- assert captured == [request]
-
-
-if __name__ == "__main__": # pragma: no cover
- for _name, _fn in sorted(globals().items()):
- if _name.startswith("test_") and callable(_fn):
- _fn()
- print(f"ok {_name}")
- print("ok")
+ assert len(captured) == 1
+ assert captured[0].model == "unconfigured-provider-model"
diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py
index 4ced59bc6..535777d94 100644
--- a/tests/test_cost_router_boundaries.py
+++ b/tests/test_cost_router_boundaries.py
@@ -2,6 +2,7 @@
from __future__ import annotations
+from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List
import pytest
@@ -11,6 +12,7 @@
InMemoryConfigStore,
ModelAgent,
PriceBook,
+ PriceEntry,
TaskOrchestrator,
)
from contextual_orchestrator.batch_routing import (
@@ -20,7 +22,10 @@
BatchResultItem,
EmbeddingBatchResultItem,
)
-from contextual_orchestrator.batch_job_registry import JobRegistryFactory
+from contextual_orchestrator.batch_job_registry import (
+ ClaimNotAcquired,
+ JobRegistryFactory,
+)
from contextual_orchestrator.cost_router import (
CostRoutingCoordinator as Coordinator,
)
@@ -29,6 +34,20 @@
_provider_from_base_url,
_weighted_average_embedding,
)
+from contextual_orchestrator.token_counting import (
+ TokenCountUnavailable,
+ UnavailableEmbeddingTokenCounter,
+)
+
+
+class _ExactTestCounter:
+ """Deterministic injected counter for synthetic embedding fixtures."""
+
+ def count_text(self, text: str, model: str = "") -> int:
+ return len(text.split())
+
+ def count_messages(self, messages: list[dict], model: str = "") -> int:
+ return sum(len(str(m.get("content", "")).split()) for m in messages)
def _coordinator(**kwargs: Any) -> Coordinator:
@@ -46,7 +65,9 @@ def _coordinator(**kwargs: Any) -> Coordinator:
]
orchestrator = TaskOrchestrator(agents)
config = InMemoryConfigStore()
- price_book = PriceBook(config)
+ price_book = kwargs.pop("price_book", None) or PriceBook(config)
+ if "token_counter" not in kwargs and "embedding_token_counter" not in kwargs:
+ kwargs["embedding_token_counter"] = _ExactTestCounter()
return Coordinator(orchestrator, config, price_book=price_book, **kwargs)
@@ -195,7 +216,7 @@ def retrieve(self, job): # type: ignore[no-untyped-def]
model="mock-a", usage_valid=False,
)]
- coordinator = _coordinator(batch_backend=InvalidUsageBackend())
+ coordinator = _coordinator(batch_backend=InvalidUsageBackend(), token_counter=_ExactTestCounter())
large_prompt = "word " * 500
job = coordinator.submit_batch([
BatchRequest(messages=[{"role": "user", "content": large_prompt}], model="mock-a")
@@ -237,7 +258,9 @@ def retrieve(self, job): # type: ignore[no-untyped-def]
)]
registry = LegacyRegistry()
- coordinator = _coordinator(batch_backend=InvalidUsageBackend(), job_registry=registry)
+ coordinator = _coordinator(
+ batch_backend=InvalidUsageBackend(), job_registry=registry, token_counter=_ExactTestCounter()
+ )
coordinator._batch_jobs["legacy-job"] = BatchJob( # noqa: SLF001
"legacy-job", "invalid-usage", request_count=1
)
@@ -294,7 +317,9 @@ def retrieve(self, job): # type: ignore[no-untyped-def]
registry = LegacyRegistry()
backend = TwoPassInvalidUsageBackend()
- coordinator = _coordinator(batch_backend=backend, job_registry=registry)
+ coordinator = _coordinator(
+ batch_backend=backend, job_registry=registry, token_counter=_ExactTestCounter()
+ )
coordinator._batch_jobs["legacy-job"] = BatchJob( # noqa: SLF001
"legacy-job", "invalid-usage", request_count=2
)
@@ -430,11 +455,10 @@ def count_messages(self, messages: Any, model: str = "") -> int:
return 3
-def test_embedding_token_count_tolerates_counter_failure_and_clamps() -> None:
+def test_embedding_token_count_propagates_counter_failure() -> None:
coordinator = _coordinator(token_counter=_ExplodingCounter())
- # Adapter failure falls back to whitespace units.
- assert coordinator._count_embedding_tokens("alpha beta gamma", "mock-e") == 3
- assert coordinator._count_embedding_tokens("", "mock-e") == 0
+ with pytest.raises(RuntimeError, match="counter backend offline"):
+ coordinator._count_embedding_tokens("alpha beta gamma", "mock-e")
class _ZeroCounter:
@@ -445,9 +469,10 @@ def count_messages(self, messages: Any, model: str = "") -> int:
return 1
-def test_embedding_token_count_clamps_positive_text_to_one() -> None:
+def test_embedding_token_count_rejects_non_positive_authoritative_result() -> None:
coordinator = _coordinator(token_counter=_ZeroCounter())
- assert coordinator._count_embedding_tokens("nonempty", "mock-e") == 1
+ with pytest.raises(RuntimeError, match="non-positive"):
+ coordinator._count_embedding_tokens("nonempty", "mock-e")
def test_split_empty_input_yields_single_empty_part() -> None:
@@ -456,6 +481,19 @@ def test_split_empty_input_yields_single_empty_part() -> None:
assert coordinator._force_token_safe_chunks("", model="m", max_tokens=4, max_chars=10) == [("", 0)]
+def test_split_uses_native_packer_when_available() -> None:
+ class Counter:
+ def pack_text(self, text: str, model: str, max_tokens: int):
+ assert (text, model, max_tokens) == ("alpha beta", "text-embedding-3-small", 4)
+ return [("alpha ", 2), ("beta", 1)]
+
+ coordinator = _coordinator(token_counter=Counter())
+
+ assert coordinator._split_embedding_input(
+ "alpha beta", model="text-embedding-3-small", max_tokens=4, max_chars=100
+ ) == [("alpha ", 2), ("beta", 1)]
+
+
class _WholeStringOnlyCounter:
"""Pathological counter: only the full original text exceeds the budget."""
@@ -526,6 +564,7 @@ class _DroppingEmbeddingBackend:
"""Local-shaped embedding backend that loses one requested vector."""
name = "dropping"
+ poll_after_ms = 250
def __init__(self) -> None:
self.jobs: Dict[str, BatchJob] = {}
@@ -558,6 +597,21 @@ def retrieve(self, job: BatchJob) -> List[EmbeddingBatchResultItem]:
]
+def test_embedding_submission_stops_before_backend_when_count_is_unavailable() -> None:
+ """No provider work or cost record may follow an unavailable exact count."""
+ backend = _DroppingEmbeddingBackend()
+ coordinator = _coordinator(
+ embedding_batch_backend=backend,
+ embedding_token_counter=UnavailableEmbeddingTokenCounter(),
+ )
+
+ with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"):
+ coordinator.submit_embeddings_batch(["synthetic input"], model="unknown-embedding")
+
+ assert backend.jobs == {}
+ assert coordinator.ledger.records() == []
+
+
def test_embeddings_document_reports_placeholder_for_missing_parts() -> None:
backend = _DroppingEmbeddingBackend()
coordinator = _coordinator(embedding_batch_backend=backend)
@@ -580,6 +634,53 @@ def test_embeddings_document_is_idempotent_after_completion() -> None:
assert backend.polled.count(job.job_id) == 1
+def test_concurrent_embedding_polls_record_usage_once() -> None:
+ backend = _DroppingEmbeddingBackend()
+ coordinator = _coordinator(embedding_batch_backend=backend)
+ job = coordinator.submit_embeddings_batch(["only one"])
+
+ with ThreadPoolExecutor(max_workers=2) as executor:
+ documents = list(
+ executor.map(
+ lambda _index: coordinator.embeddings_batch_document(job.job_id),
+ range(2),
+ )
+ )
+
+ assert documents[0] == documents[1]
+ assert len(coordinator.ledger.records()) == 1
+
+
+def test_contended_embedding_document_claim_returns_cache_or_pending() -> None:
+ backend = _DroppingEmbeddingBackend()
+ coordinator = _coordinator(embedding_batch_backend=backend)
+ job = coordinator.submit_embeddings_batch(["only one"])
+
+ class _ContendedClaim:
+ def __enter__(self):
+ raise ClaimNotAcquired("owned by another poller")
+
+ def __exit__(self, *_args):
+ return False
+
+ coordinator.job_registry.lock = lambda *_args, **_kwargs: _ContendedClaim()
+
+ pending = coordinator.embeddings_batch_document(job.job_id)
+ assert pending == {
+ "batch_id": job.job_id,
+ "status": "in_progress",
+ "backend": job.backend,
+ "model": "mock-a",
+ "embeddings": None,
+ "poll_after_ms": 250,
+ "job_retention_ms": coordinator.job_registry.retention_seconds * 1000,
+ }
+
+ cached = {**pending, "status": "completed", "embeddings": []}
+ coordinator._embedding_documents[job.job_id] = cached
+ assert coordinator.embeddings_batch_document(job.job_id) == cached
+
+
def test_unpriced_embeddings_document_omits_cost() -> None:
document = _coordinator().complete_embeddings_batch(["unpriced embedding"])
assert document["price_known"] is False
@@ -593,6 +694,19 @@ def test_embeddings_document_requires_known_batch() -> None:
coordinator.embeddings_batch_document("no_such_batch")
+def test_embeddings_document_requires_the_bound_owner() -> None:
+ coordinator = _coordinator(embedding_batch_backend=_DroppingEmbeddingBackend())
+ job = coordinator.submit_embeddings_batch(["private"], owner_id="principal-a")
+
+ with pytest.raises(KeyError, match="embeddings batch job"):
+ coordinator.embeddings_batch_document(job.job_id, owner_id="principal-b")
+
+ assert (
+ coordinator.embeddings_batch_document(job.job_id, owner_id="principal-a")["status"]
+ == "completed"
+ )
+
+
def test_embeddings_document_incomplete_status_has_no_vectors() -> None:
class _PendingBackend(_DroppingEmbeddingBackend):
def poll(self, job: BatchJob) -> Dict[str, Any]:
@@ -605,6 +719,62 @@ def poll(self, job: BatchJob) -> Dict[str, Any]:
assert document["status"] == "in_progress"
+def test_embeddings_document_preserves_failed_terminal_state_without_cost() -> None:
+ class _FailedBackend(_DroppingEmbeddingBackend):
+ def poll(self, job: BatchJob) -> Dict[str, Any]:
+ return {
+ "job_id": job.job_id,
+ "status": "failed",
+ "is_complete": True,
+ "failure": {"error_type": "ProviderError", "retryable": False},
+ }
+
+ def retrieve(self, job: BatchJob) -> List[EmbeddingBatchResultItem]:
+ raise AssertionError("failed jobs have no result payload")
+
+ coordinator = _coordinator(embedding_batch_backend=_FailedBackend())
+ job = coordinator.submit_embeddings_batch(["never billed"])
+
+ document = coordinator.embeddings_batch_document(job.job_id)
+
+ assert document["status"] == "failed"
+ assert document["embeddings"] is None
+ assert document["failure"]["error_type"] == "ProviderError"
+ assert coordinator.ledger.records() == []
+
+
+def test_embeddings_document_bills_the_selected_agent_not_caller_attribution() -> None:
+ agent = ModelAgent(
+ id="embedding_worker",
+ model="embedding-v1",
+ base_url="mock://embed",
+ provider_name="trusted-provider",
+ tags=("embedding",),
+ )
+ orchestrator = TaskOrchestrator([agent])
+ config = InMemoryConfigStore()
+ price_book = PriceBook(config)
+ price_book.set_price(PriceEntry("trusted-provider", "embedding-v1", 2.0, 0.0))
+ coordinator = Coordinator(
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_token_counter=_ExactTestCounter(),
+ embedding_batch_backend=_DroppingEmbeddingBackend(),
+ )
+
+ coordinator.complete_embeddings_batch(
+ ["bill selected route"],
+ model="embedding-v1",
+ agent_id=agent.id,
+ attribution={"provider": "spoofed-provider"},
+ )
+
+ row = coordinator.ledger.records()[0]
+ assert row["provider_name"] == "trusted-provider"
+ assert row["model_name"] == "embedding-v1"
+
+
class _FlakyEmbeddingBackend(_DroppingEmbeddingBackend):
"""Fails the first retrieve() with a download error, then succeeds."""
@@ -730,8 +900,8 @@ def test_document_recounts_tokens_when_backend_reports_zero_usage() -> None:
coordinator = _coordinator(embedding_batch_backend=_SilentEmbeddingBackend())
job = coordinator.submit_embeddings_batch(["alpha beta gamma"])
document = coordinator.embeddings_batch_document(job.job_id)
- # HeuristicTokenCounter counts word units with the BPE expansion factor.
- assert document["token_counts"] == [4]
+ # The injected exact test seam counts whitespace-delimited fixture units.
+ assert document["token_counts"] == [3]
def test_complete_embeddings_batch_round_trips_locally() -> None:
diff --git a/tests/test_embeddings_blank_input_http_honesty.py b/tests/test_embeddings_blank_input_http_honesty.py
index db3a89856..de12623fd 100644
--- a/tests/test_embeddings_blank_input_http_honesty.py
+++ b/tests/test_embeddings_blank_input_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_blank_input_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_encoding_format_base64_http_honesty.py b/tests/test_embeddings_encoding_format_base64_http_honesty.py
index 0bb7e04af..f85d2bd31 100644
--- a/tests/test_embeddings_encoding_format_base64_http_honesty.py
+++ b/tests/test_embeddings_encoding_format_base64_http_honesty.py
@@ -13,7 +13,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_encoding_format_base64_http_honesty_token" # noqa: S105
@@ -50,10 +50,13 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_embeddings_encoding_format_http_honesty.py b/tests/test_embeddings_encoding_format_http_honesty.py
index 1f55cc420..ea5eaf729 100644
--- a/tests/test_embeddings_encoding_format_http_honesty.py
+++ b/tests/test_embeddings_encoding_format_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import SecurityConfig, build_server
_TEST_AUTH_TOKEN = "embeddings_encoding_format_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_metadata_http_honesty.py b/tests/test_embeddings_metadata_http_honesty.py
index 404b8d698..16771816c 100644
--- a/tests/test_embeddings_metadata_http_honesty.py
+++ b/tests/test_embeddings_metadata_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_metadata_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_model_pool_http_honesty.py b/tests/test_embeddings_model_pool_http_honesty.py
index 4d655ed14..80e12fd02 100644
--- a/tests/test_embeddings_model_pool_http_honesty.py
+++ b/tests/test_embeddings_model_pool_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import SecurityConfig, build_server
_TEST_AUTH_TOKEN = "embeddings_model_pool_http_honesty_token"
@@ -42,7 +42,9 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_null_optional_noop_http_honesty.py b/tests/test_embeddings_null_optional_noop_http_honesty.py
index d4b72fdb0..34e9e02e3 100644
--- a/tests/test_embeddings_null_optional_noop_http_honesty.py
+++ b/tests/test_embeddings_null_optional_noop_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_null_optional_noop_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_routing_http_honesty.py b/tests/test_embeddings_routing_http_honesty.py
index 039b92a0e..59864d5d0 100644
--- a/tests/test_embeddings_routing_http_honesty.py
+++ b/tests/test_embeddings_routing_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_routing_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_embeddings_token_array_input_http_honesty.py b/tests/test_embeddings_token_array_input_http_honesty.py
index dd2dd1db4..9cc59f5fb 100644
--- a/tests/test_embeddings_token_array_input_http_honesty.py
+++ b/tests/test_embeddings_token_array_input_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_token_array_input_http_honesty_token" # noqa: S105
@@ -48,10 +48,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_embeddings_user_field_http_honesty.py b/tests/test_embeddings_user_field_http_honesty.py
index 0482cc6d2..54a35f9ed 100644
--- a/tests/test_embeddings_user_field_http_honesty.py
+++ b/tests/test_embeddings_user_field_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "embeddings_user_field_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_empty_string_encoding_tool_choice_endpoint_noop_http_honesty.py b/tests/test_empty_string_encoding_tool_choice_endpoint_noop_http_honesty.py
index 7237e9f9f..70e873227 100644
--- a/tests/test_empty_string_encoding_tool_choice_endpoint_noop_http_honesty.py
+++ b/tests/test_empty_string_encoding_tool_choice_endpoint_noop_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "empty_string_encoding_tool_choice_endpoint_noop_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_empty_string_numeric_controls_noop_http_honesty.py b/tests/test_empty_string_numeric_controls_noop_http_honesty.py
index 5a18c3a26..c68ccf982 100644
--- a/tests/test_empty_string_numeric_controls_noop_http_honesty.py
+++ b/tests/test_empty_string_numeric_controls_noop_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "empty_string_numeric_controls_noop_http_honesty_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_encoding_stream_logprobs_http_honesty.py b/tests/test_encoding_stream_logprobs_http_honesty.py
index 8d746ef01..e0696160d 100644
--- a/tests/test_encoding_stream_logprobs_http_honesty.py
+++ b/tests/test_encoding_stream_logprobs_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import SecurityConfig, build_server
_TEST_AUTH_TOKEN = "encoding_stream_logprobs_http_honesty_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_evolve_optimizer.py b/tests/test_evolve_optimizer.py
index 04ec0d209..c1ab2a303 100644
--- a/tests/test_evolve_optimizer.py
+++ b/tests/test_evolve_optimizer.py
@@ -16,6 +16,13 @@
from contextual_orchestrator.orchestrator import evolve_orchestration, _space_size # noqa: E402
+class _ExactCounter:
+ """Exact synthetic output counter for optimizer accounting."""
+
+ def count_text(self, text: str, model: str) -> int:
+ return len(text.encode("utf-8"))
+
+
# Config space: worker "tier" controls (deterministically) both quality and price.
TIERS = {
"small_worker": {"price": 1.0, "quality": 0.5},
@@ -34,6 +41,7 @@ def _build(config: dict) -> TaskOrchestrator:
return TaskOrchestrator(
[ModelAgent(config["tier"], "model-x", tags=("reasoning", "writing"))],
price_per_million={"model-x": tier["price"]},
+ token_counter=_ExactCounter(),
)
diff --git a/tests/test_generated_workflow.py b/tests/test_generated_workflow.py
index 47f015731..5ee9a22c4 100644
--- a/tests/test_generated_workflow.py
+++ b/tests/test_generated_workflow.py
@@ -85,11 +85,53 @@ def test_access_lists_actually_isolate_context() -> None:
step3_prompt = client.calls[4][-1]["content"]
assert "step-output(1)" not in step1_prompt # access [] -> sees no prior output
- assert "step-output(1)" in step2_prompt and "step-output(2)" in step2_prompt # access [0,1]
+ assert step2_prompt.count("step-output(1)") == 1 # access [0,1], copied once
+ assert step2_prompt.count("step-output(2)") == 1
assert "step-output(1)" not in step3_prompt # access [1,2] excludes step 0
assert "step-output(2)" in step3_prompt and "step-output(3)" in step3_prompt
+def test_workflow_step_carries_current_task_and_source_context_once() -> None:
+ """Workers receive one canonical copy of task and source-bearing history."""
+ orchestrator, client = _orch(json.dumps(PLAN))
+ orchestrator.conduct(
+ [
+ {"role": "system", "content": "CALLER_POLICY_SENTINEL"},
+ {"role": "user", "content": "EARLIER_TURN_SENTINEL"},
+ {"role": "assistant", "content": "EARLIER_ANSWER_SENTINEL"},
+ {"role": "user", "content": "CURRENT_TASK_SENTINEL"},
+ ]
+ )
+
+ step_messages = client.calls[1]
+ serialized = json.dumps(step_messages)
+ assert serialized.count("CURRENT_TASK_SENTINEL") == 1
+ assert serialized.count("EARLIER_TURN_SENTINEL") == 1
+ assert serialized.count("EARLIER_ANSWER_SENTINEL") == 1
+ assert "Caller instructions:\nCALLER_POLICY_SENTINEL" in step_messages[0]["content"]
+ assert sum(
+ message.get("role") == "user" and message.get("content") == "CURRENT_TASK_SENTINEL"
+ for message in step_messages
+ ) == 1
+ assert "CURRENT_TASK_SENTINEL" not in step_messages[-1]["content"]
+
+
+def test_workflow_step_reasserts_multipart_caller_instructions() -> None:
+ """Text content parts retain system authority in every worker stage."""
+ orchestrator, client = _orch(json.dumps(PLAN))
+ orchestrator.conduct(
+ [
+ {
+ "role": "system",
+ "content": [{"type": "text", "text": "MULTIPART_POLICY_SENTINEL"}],
+ },
+ {"role": "user", "content": "solve it"},
+ ]
+ )
+
+ assert "Caller instructions:\nMULTIPART_POLICY_SENTINEL" in client.calls[1][0]["content"]
+
+
def test_planner_prompt_lists_the_agent_pool() -> None:
orchestrator, client = _orch(json.dumps(PLAN))
orchestrator.conduct([{"role": "user", "content": "solve it"}])
diff --git a/tests/test_ledger_execution_identity_http_honesty.py b/tests/test_ledger_execution_identity_http_honesty.py
index 5ce178446..8313b70d2 100644
--- a/tests/test_ledger_execution_identity_http_honesty.py
+++ b/tests/test_ledger_execution_identity_http_honesty.py
@@ -40,7 +40,13 @@ def _serve():
config = InMemoryConfigStore()
price_book = PriceBook(config)
price_book.set_price(PriceEntry("mock", "mock-a", prompt_price_per_1k=1.0, completion_price_per_1k=2.0))
- coordinator = CostRoutingCoordinator(orchestrator, config, price_book=price_book)
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_token_counter=counter,
+ )
server = build_server(
orchestrator,
port=0,
diff --git a/tests/test_mode_casefold_http_honesty.py b/tests/test_mode_casefold_http_honesty.py
index da13fb768..9d3b6f22c 100644
--- a/tests/test_mode_casefold_http_honesty.py
+++ b/tests/test_mode_casefold_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import SecurityConfig, build_server
_TEST_AUTH_TOKEN = "mode_casefold_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py
index 0a5271f77..cfa249cea 100644
--- a/tests/test_model_judge.py
+++ b/tests/test_model_judge.py
@@ -23,6 +23,7 @@
from contextual_orchestrator.orchestrator import ( # noqa: E402
BudgetExceededError,
ModelClient,
+ ProviderRequestTooLargeError,
ProviderResponseError,
_parse_model_judge_reply,
_structured_output_error,
@@ -172,7 +173,15 @@ def judge(self, *, task: str, answer: str, criteria: tuple) -> object:
criterion_cls=_ScriptedCriterion,
format_error=ValueError,
)
+ class _ExactTestCounter:
+ def count_text(self, text: str, model: str = "") -> int:
+ return len(text.split())
+
+ def count_messages(self, messages: list[dict], model: str = "") -> int:
+ return sum(len(m.get("content", "").split()) for m in messages)
+
orchestrator, _ = _orch('{"decision":"ACCEPT","reason":"The report supports the answer."}')
+ orchestrator.token_counter = _ExactTestCounter()
with patch.object(
orchestrator_module,
"_resolve_fast_mlsirm_components",
@@ -190,12 +199,12 @@ def judge(self, *, task: str, answer: str, criteria: tuple) -> object:
assert verification["verifier_output"] != verification["judge_output_text"]
isolated_record = {"trace": [], "verification": verification}
- budget_contribution = orchestrator._run_budget_output_by_model(isolated_record)
- assert budget_contribution["model-x"] == orchestrator_module.estimate_tokens(
- verification["judge_output_text"]
+ budget_contribution, _ = orchestrator._run_budget_output_by_model(isolated_record)
+ assert budget_contribution["model-x"] == len(
+ verification["judge_output_text"].split()
)
- assert budget_contribution["model-x"] != orchestrator_module.estimate_tokens(
- verification["verifier_output"]
+ assert budget_contribution["model-x"] != len(
+ verification["verifier_output"].split()
)
assert budget_contribution["model-x"] > 0
@@ -281,6 +290,50 @@ def chat(self, agent: ModelAgent, messages: list, **kwargs: object) -> str: # t
assert set(client.agent_ids) == {"group_member"}
+def test_explicit_structured_group_model_pins_every_provider_call() -> None:
+ class _RecordingClient(_ScriptedClient):
+ def __init__(self) -> None:
+ super().__init__('{"decision":"ACCEPT","reason":"Exact judge passed."}')
+ self.calls_by_kind: list[tuple[str, str]] = []
+
+ def chat(self, agent: ModelAgent, messages: list, **kwargs: object) -> str: # type: ignore[override]
+ self.calls_by_kind.append(("evidence_or_judge", agent.id))
+ return super().chat(agent, messages, **kwargs)
+
+ def proxy_send(self, agent: ModelAgent, endpoint: str, body: dict) -> dict: # type: ignore[override]
+ del endpoint, body
+ self.calls_by_kind.append(("synthesis", agent.id))
+ return {"choices": [{"message": {"content": '{"status":"ok"}'}}]}
+
+ client = _RecordingClient()
+ selected = ModelAgent(
+ "selected_member", "selected-model", group_name="shared_model_group"
+ )
+ sibling = ModelAgent(
+ "sibling_member", "sibling-model", group_name="shared_model_group", priority=100
+ )
+ orchestrator = TaskOrchestrator([selected, sibling], client=client)
+ with patch.object(
+ orchestrator_module,
+ "_resolve_fast_mlsirm_components",
+ return_value=_scripted_fast_components(),
+ ):
+ result = orchestrator.proxy_completion(
+ {
+ "model": selected.model,
+ "messages": [{"role": "user", "content": "structured"}],
+ "response_format": {"type": "json_object"},
+ },
+ single_agent=False,
+ )
+
+ assert result["choices"][0]["message"]["content"] == '{"status":"ok"}'
+ assert client.calls_by_kind == [
+ *(('evidence_or_judge', selected.id) for _ in range(5)),
+ ("synthesis", selected.id),
+ ]
+
+
def test_free_structured_judge_uses_exact_free_agent_with_duplicate_model_id() -> None:
class _ProxyClient(ModelClient):
def __init__(self) -> None:
@@ -687,6 +740,152 @@ def test_strict_schema_validation_and_repair_stay_in_the_conduct_trace() -> None
assert _structured_output_error('{"input_count":6}', response_format) == "schema_violation"
+def test_structured_repair_does_not_retry_request_excluded_model() -> None:
+ stale = ModelAgent("stale_agent", "stale-model", "mock://catalog")
+ live = ModelAgent("live_agent", "live-model", "mock://catalog")
+ orchestrator = TaskOrchestrator([stale, live])
+ calls = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ if calls == [stale.id]:
+ raise urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None)
+ if calls == [stale.id, live.id]:
+ return {"choices": [{"message": {"content": '{"input_count":6}'}}]}
+ if agent.id == live.id:
+ raise urllib.error.HTTPError("https://synthetic.invalid", 413, "large", {}, None)
+ return {"choices": [{"message": {"content": '{"input_count":10}'}}]}
+
+ 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,
+ },
+ },
+ }
+ with (
+ patch.object(
+ orchestrator,
+ "conduct",
+ return_value={"mode": "conduct", "answer": "evidence", "trace": [], "verification": {}},
+ ),
+ patch.object(orchestrator, "_select_agent", return_value=stale),
+ patch.object(orchestrator, "_failover_candidates", return_value=[stale, live]),
+ patch.object(orchestrator.client, "proxy_send_once", side_effect=send),
+ pytest.raises(ProviderRequestTooLargeError),
+ ):
+ orchestrator.proxy_completion(
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "classify ten items"}],
+ "response_format": response_format,
+ },
+ single_agent=False,
+ )
+
+ assert calls == [stale.id, live.id, live.id]
+
+
+def test_virtual_structured_schema_failure_advances_on_same_endpoint() -> None:
+ """A virtual candidate that fails synthesis and repair is excluded once."""
+ first = ModelAgent("first_agent", "first-model", "mock://catalog")
+ second = ModelAgent("second_agent", "second-model", "mock://catalog")
+ other = ModelAgent("other_agent", "other-model", "mock://other")
+ orchestrator = TaskOrchestrator([first, second, other])
+ calls: list[str] = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ count = 10 if agent.id == second.id else 6
+ return {"choices": [{"message": {"content": f'{{"input_count": {count}}}'}}]}
+
+ 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,
+ },
+ },
+ }
+ with (
+ patch.object(orchestrator, "conduct", return_value={"trace": []}),
+ patch.object(orchestrator, "_select_agent", return_value=first),
+ patch.object(
+ orchestrator,
+ "_failover_candidates",
+ return_value=[first, second, other],
+ ),
+ patch.object(orchestrator.client, "proxy_send_once", side_effect=send),
+ ):
+ result = orchestrator.proxy_completion(
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "classify ten items"}],
+ "response_format": response_format,
+ },
+ single_agent=False,
+ )
+
+ assert result["choices"][0]["message"]["content"] == '{"input_count": 10}'
+ assert calls == [first.id, first.id, second.id]
+ assert other.id not in calls
+
+
+def test_explicit_structured_schema_failure_remains_pinned() -> None:
+ """An explicit model never switches after its synthesis and repair fail schema."""
+ first = ModelAgent("first_agent", "first-model", "mock://catalog")
+ second = ModelAgent("second_agent", "second-model", "mock://catalog")
+ orchestrator = TaskOrchestrator([first, second])
+ calls: list[str] = []
+
+ def send(agent, _endpoint, _payload):
+ calls.append(agent.id)
+ return {"choices": [{"message": {"content": '{"input_count": 6}'}}]}
+
+ 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,
+ },
+ },
+ }
+ with (
+ patch.object(orchestrator, "conduct", return_value={"trace": []}),
+ patch.object(orchestrator.client, "proxy_send", side_effect=send),
+ pytest.raises(
+ ProviderResponseError,
+ match="structured synthesis and repair violated response_format",
+ ),
+ ):
+ orchestrator.proxy_completion(
+ {
+ "model": first.model,
+ "messages": [{"role": "user", "content": "classify ten items"}],
+ "response_format": response_format,
+ },
+ single_agent=False,
+ )
+
+ assert calls == [first.id, first.id]
+
+
def test_structured_synthesis_failure_updates_provider_health() -> None:
"""A failed final provider is excluded by the existing circuit policy."""
orchestrator, _ = _orch("unused")
@@ -869,6 +1068,33 @@ def __init__(self, criterion_id: str, description: str, weight: float) -> None:
assert result["judge_orchestration_mode"] == "route"
+def test_model_judge_does_not_retry_request_excluded_agent() -> None:
+ stale = ModelAgent("stale_agent", "stale-model", tags=("reasoning", "writing"))
+ live = ModelAgent("live_agent", "live-model", tags=("reasoning", "writing"))
+ calls = []
+
+ class _Client(ModelClient):
+ def chat(self, agent, messages, temperature=None):
+ del messages, temperature
+ calls.append(agent.id)
+ return '{"decision":"ACCEPT","reason":"live judge"}'
+
+ orchestrator = TaskOrchestrator([stale, live], client=_Client())
+ with patch.object(
+ orchestrator_module,
+ "_resolve_fast_mlsirm_components",
+ return_value=_scripted_fast_components(),
+ ), patch.object(orchestrator, "_ranked_agents", return_value=[stale, live]):
+ result = orchestrator._model_judge_verification(
+ "task",
+ {"verifier_output": "report"},
+ excluded_agent_ids={stale.id},
+ )
+
+ assert result["accepted"] is True
+ assert calls == [live.id]
+
+
def test_fast_mlsirm_invalid_irt_projection_fails_closed() -> None:
class _Judge:
def __init__(self, _orchestrator, *, mode: str, accept_threshold: float) -> None:
diff --git a/tests/test_openai_passthrough.py b/tests/test_openai_passthrough.py
index a29a27951..9fc776560 100644
--- a/tests/test_openai_passthrough.py
+++ b/tests/test_openai_passthrough.py
@@ -32,6 +32,7 @@
build_server,
responses_sse_body,
)
+from contextual_orchestrator.telemetry import current_session_id # noqa: E402
def _build() -> TaskOrchestrator:
@@ -508,6 +509,61 @@ def test_http_chat_completions_accepts_response_format_and_passes_through() -> N
assert body["echo"]["response_format"] == {"type": "json_object"}
+def test_lineage_structured_payload_accepts_session_without_provider_forwarding() -> None:
+ """Lineage correlation is gateway metadata, not a provider request field."""
+ orchestrator = _build()
+ observed_sessions: list[str | None] = []
+ proxy_completion = orchestrator.proxy_completion
+
+ def capture_session(*args, **kwargs): # type: ignore[no-untyped-def]
+ observed_sessions.append(current_session_id())
+ return proxy_completion(*args, **kwargs)
+
+ orchestrator.proxy_completion = capture_session # type: ignore[method-assign]
+ token = "passthrough_token"
+ server = build_server(
+ orchestrator, port=0, security=SecurityConfig(auth_token=token)
+ )
+ threading.Thread(target=server.serve_forever, daemon=True).start()
+ port = server.server_address[1]
+ try:
+ status, body = _post(
+ f"http://127.0.0.1:{port}/v1/chat/completions",
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "Return synthetic JSON."}],
+ "response_format": {"type": "json_object"},
+ "session_id": "synthetic-lineage-session",
+ },
+ token,
+ )
+ finally:
+ server.shutdown()
+ assert status == 200
+ assert body["echo"]["response_format"] == {"type": "json_object"}
+ assert "session_id" not in body["echo"]
+ assert observed_sessions == ["synthetic-lineage-session"]
+
+
+@pytest.mark.parametrize("session_id", [7, "", "x" * 129, "line\nbreak"])
+def test_http_rejects_invalid_top_level_session_id(session_id: object) -> None:
+ server, port, token = _serve()
+ try:
+ status, body = _post(
+ f"http://127.0.0.1:{port}/v1/chat/completions",
+ {
+ "model": "mock-planner",
+ "messages": [{"role": "user", "content": "hello"}],
+ "session_id": session_id,
+ },
+ token,
+ )
+ finally:
+ server.shutdown()
+ assert status == 400
+ assert body["error"]["code"] == "invalid_session_id"
+
+
def test_http_gateway_default_response_format_resolves_concrete_agent() -> None:
"""The virtual gateway default remains valid on provider-native passthrough."""
server, port, token = _serve()
diff --git a/tests/test_openai_user_field_http_honesty.py b/tests/test_openai_user_field_http_honesty.py
index a5a331919..9071b9ea2 100644
--- a/tests/test_openai_user_field_http_honesty.py
+++ b/tests/test_openai_user_field_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import SecurityConfig, build_server
_TEST_AUTH_TOKEN = "openai_user_field_http_honesty_token" # noqa: S105
@@ -42,7 +42,9 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
- server = build_server(build(), port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN))
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
+ server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN), coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
return server, thread, server.server_address[1]
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 95c719796..d823049a4 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -17,11 +17,19 @@
from contextual_orchestrator.orchestrator import optimize_orchestration, _pareto_front # noqa: E402
+class _ExactCounter:
+ """Exact synthetic output counter for optimizer accounting."""
+
+ def count_text(self, text: str, model: str) -> int:
+ return len(text.encode("utf-8"))
+
+
def _candidate(name: str, agent_id: str, price: float) -> dict:
# Distinct agent id shows up in the mock answer, so quality_fn can differentiate configs.
orchestrator = TaskOrchestrator(
[ModelAgent(agent_id, "model-x", tags=("reasoning", "writing"))],
price_per_million={"model-x": price},
+ token_counter=_ExactCounter(),
)
return {"name": name, "orchestrator": orchestrator, "mode": "route"}
diff --git a/tests/test_orchestrated_responses_stream.py b/tests/test_orchestrated_responses_stream.py
index b958f20c5..3926ab52b 100644
--- a/tests/test_orchestrated_responses_stream.py
+++ b/tests/test_orchestrated_responses_stream.py
@@ -143,7 +143,10 @@ def test_streamed_responses_records_unavailable_usage_without_estimating_answer(
assert len({row["workflow_run_id"] for row in rows}) == 2
assert all(row["request_channel"] == "stream" for row in rows)
assert all(row["measurement_status"] == "unavailable" for row in rows)
- assert all(row["prompt_tokens"] == row["completion_tokens"] == 0 for row in rows)
+ assert all(
+ row["prompt_tokens"] is None and row["completion_tokens"] is None
+ for row in rows
+ )
def test_conduct_preserves_responses_instructions_for_every_stage() -> None:
diff --git a/tests/test_orchestrator_dispatch_boundaries.py b/tests/test_orchestrator_dispatch_boundaries.py
index 4de661f3c..9c3957108 100644
--- a/tests/test_orchestrator_dispatch_boundaries.py
+++ b/tests/test_orchestrator_dispatch_boundaries.py
@@ -14,6 +14,7 @@
)
from contextual_orchestrator.orchestrator import (
_COMMERCIAL_REPORT_CACHE,
+ BudgetExceededError,
ModelAgent,
TaskOrchestrator,
WorkflowStep,
@@ -143,9 +144,9 @@ def test_batch_route_enforces_budget_and_request_identifier_contract() -> None:
with pytest.raises(RuntimeError, match="spend budget exceeded"):
orch.batch_route(["prompt one"])
- # Without an exceeded budget the batch path answers and persists normally.
- records = orch.batch_route(["prompt two"])
- assert len(records) == 1
+ # Missing authoritative accounting blocks an enabled budget before dispatch.
+ with pytest.raises(BudgetExceededError, match="measurement unavailable"):
+ orch.batch_route(["prompt two"])
def test_batch_route_persists_runs_when_a_state_db_is_configured(tmp_path) -> None:
@@ -448,7 +449,7 @@ def test_openai_models_deduplicate_models_and_unknown_ids_raise() -> None:
# -- spend analytics usage-source classification ---------------------------------------
-def test_spend_analytics_marks_mixed_usage_sources_per_model() -> None:
+def test_spend_analytics_marks_partial_unknown_usage_unavailable() -> None:
orch = _orch(_agent(), _agent("builder_agent"))
run = {
"workflow_run_id": "run_mixed",
@@ -473,7 +474,8 @@ def test_spend_analytics_marks_mixed_usage_sources_per_model() -> None:
by_model = {row["model"]: row for row in report["by_model"]}
planner_row = by_model["mock-model"]
assert planner_row["step_count"] == 3
- assert planner_row["usage_source"] == "mixed"
+ assert planner_row["usage_source"] == "unavailable"
+ assert planner_row["output_tokens"] is None
# -- readiness criteria -----------------------------------------------------------------
diff --git a/tests/test_passthrough_provider_failover.py b/tests/test_passthrough_provider_failover.py
index 6b0fb876d..769311e43 100644
--- a/tests/test_passthrough_provider_failover.py
+++ b/tests/test_passthrough_provider_failover.py
@@ -21,6 +21,7 @@
from contextual_orchestrator.orchestrator import (
ModelClient,
ProviderRequestTooLargeError,
+ _structured_output_error,
)
from contextual_orchestrator.provider_errors import ProviderUpstreamError
@@ -641,6 +642,33 @@ def test_all_virtual_candidates_rejecting_size_preserves_request_too_large() ->
)
+def test_stale_model_then_size_failure_preserves_request_too_large() -> None:
+ client = SequencedProxyClient(
+ {
+ "primary_agent": _http_error(404),
+ "fallback_agent": _http_error(413),
+ }
+ )
+
+ orchestrator = _build(client)
+ orchestrator.conduct = lambda *args, **kwargs: { # type: ignore[method-assign]
+ "mode": "conduct",
+ "answer": "evidence",
+ "trace": [],
+ "verification": {"accepted": True, "reason": "test", "verifier_output": ""},
+ }
+
+ with pytest.raises(ProviderRequestTooLargeError, match="every eligible provider"):
+ orchestrator.proxy_completion(
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "large request"}],
+ "response_format": {"type": "json_object"},
+ },
+ single_agent=False,
+ )
+
+
def test_mixed_failures_surface_the_final_classified_provider_failure() -> None:
"""Mixed exhaustion keeps the final provider's actionable typed failure."""
client = SequencedProxyClient(
@@ -1339,3 +1367,59 @@ def test_default_mock_endpoint_represents_one_fixture_provider() -> None:
assert caught.value.agent_id == "first_mock"
assert [agent_id for agent_id, _ in client.calls] == ["first_mock"]
+
+
+def test_virtual_structured_synthesis_skips_reasoning_only_model_on_same_endpoint() -> None:
+ """A contentless virtual candidate cannot terminate same-endpoint synthesis."""
+ endpoint = "https://synthetic.invalid/v1"
+ client = SequencedProxyClient(
+ {
+ "reasoning_only": {
+ "choices": [{"message": {"content": None, "reasoning": "bounded"}}]
+ },
+ "structured_live": {
+ "model": "structured-model",
+ "choices": [{"message": {"content": '{"status":"synthetic_ok"}'}}]
+ },
+ }
+ )
+ first = ModelAgent(
+ "reasoning_only", "reasoning-model", endpoint, priority=10,
+ tags=("response_format",),
+ )
+ second = ModelAgent(
+ "structured_live", "structured-model", endpoint, priority=1,
+ tags=("response_format",),
+ )
+ orchestrator = TaskOrchestrator([first, second], client=client)
+ orchestrator.conduct = lambda *args, **kwargs: { # type: ignore[method-assign]
+ "mode": "conduct",
+ "answer": "evidence",
+ "trace": [],
+ "verification": {"accepted": True},
+ }
+
+ result = orchestrator.proxy_completion(
+ {
+ "model": TaskOrchestrator.AUTO_MODEL,
+ "messages": [{"role": "user", "content": "synthetic structured request"}],
+ "response_format": {"type": "json_object"},
+ },
+ single_agent=False,
+ )
+
+ assert result["model"] == "structured-model"
+ assert result["choices"][0]["message"]["content"] == '{"status":"synthetic_ok"}'
+ assert [agent_id for agent_id, _ in client.calls] == [
+ "reasoning_only",
+ "structured_live",
+ ]
+
+
+def test_json_object_contract_rejects_non_json_and_non_object_values() -> None:
+ """json_object validation cannot accept provider prose or JSON scalars."""
+ response_format = {"type": "json_object"}
+
+ assert _structured_output_error("not json", response_format) == "invalid_json"
+ assert _structured_output_error("[]", response_format) == "invalid_json_object"
+ assert _structured_output_error('{"status":"synthetic_ok"}', response_format) is None
diff --git a/tests/test_prediction_modalities_model_strip_http_honesty.py b/tests/test_prediction_modalities_model_strip_http_honesty.py
index 051eac00e..5c8bb5f48 100644
--- a/tests/test_prediction_modalities_model_strip_http_honesty.py
+++ b/tests/test_prediction_modalities_model_strip_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "prediction_modalities_model_strip_http_honesty_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py
new file mode 100644
index 000000000..0eb661fbb
--- /dev/null
+++ b/tests/test_provider_embedding_batch_backend.py
@@ -0,0 +1,459 @@
+"""Focused synthetic contracts for provider-backed embedding batches."""
+
+import threading
+import time
+
+import pytest
+
+from contextual_orchestrator.batch_routing import (
+ EmbeddingBatchRequest,
+ ProviderEmbeddingBatchBackend,
+)
+from contextual_orchestrator.batch_job_registry import JobRegistryFactory
+from contextual_orchestrator import (
+ CostRoutingCoordinator,
+ InMemoryConfigStore,
+ ModelAgent,
+ PriceBook,
+ PriceEntry,
+ TaskOrchestrator,
+)
+from contextual_orchestrator.orchestrator import ModelClient
+from contextual_orchestrator.provider_errors import ProviderUpstreamError
+from contextual_orchestrator.server import SecurityConfig, build_server
+from contextual_orchestrator.token_counting import (
+ TokenCountUnavailable,
+ UnavailableEmbeddingTokenCounter,
+)
+
+
+class _SyntheticProviderClient(ModelClient):
+ def __init__(self):
+ super().__init__()
+ self.embedding_calls = []
+
+ def embed(self, agent, texts):
+ self.embedding_calls.append(list(texts))
+ return [[float(len(text))] for text in texts]
+
+ def embed_with_usage(self, agent, texts):
+ return self.embed(agent, texts), sum(len(text.encode("utf-8")) for text in texts)
+
+
+class _SyntheticExactCounter:
+ def count_text(self, text, model):
+ """Return a deterministic synthetic authoritative count."""
+ return len(text.split())
+
+
+def test_unknown_tokenizer_uses_authoritative_provider_usage() -> None:
+ """A byte-safe request completes only after the provider supplies exact usage."""
+ agent = ModelAgent(
+ "provider_embedding", "provider-embedding-model", "https://provider.synthetic.invalid/v1", tags=("embedding",)
+ )
+ orchestrator = TaskOrchestrator([agent], client=_SyntheticProviderClient())
+ config = InMemoryConfigStore()
+ price_book = PriceBook(config)
+ price_book.set_price(
+ PriceEntry("provider.synthetic.invalid", "provider-embedding-model", 1.0, 0.0)
+ )
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ config,
+ price_book=price_book,
+ embedding_token_counter=UnavailableEmbeddingTokenCounter(),
+ )
+
+ document = coordinator.complete_embeddings_batch(["synthetic input"])
+
+ assert document["status"] == "completed"
+ assert document["total_tokens"] == len("synthetic input".encode("utf-8"))
+ assert document["cost_micro_usd"] > 0
+
+
+def test_unknown_tokenizer_rejects_missing_provider_usage() -> None:
+ """Vectors without authoritative provider usage never become a successful job."""
+ agent = ModelAgent(
+ "provider_embedding", "provider-embedding-model", "https://provider.synthetic.invalid/v1", tags=("embedding",)
+ )
+ client = _SyntheticProviderClient()
+ client.embed_with_usage = lambda agent, texts: (client.embed(agent, texts), None)
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent], client=client),
+ embedding_token_counter=UnavailableEmbeddingTokenCounter(),
+ )
+
+ job = coordinator.submit_embeddings_batch(["synthetic input"])
+ deadline = time.time() + 2
+ while coordinator.embedding_batch_backend.poll(job)["status"] not in {"completed", "failed"} and time.time() < deadline:
+ time.sleep(0.01)
+
+ assert coordinator.embedding_batch_backend.poll(job)["status"] == "failed"
+
+
+def test_unknown_tokenizer_fails_before_provider_when_byte_bound_exceeds_budget() -> None:
+ """An unprovable preflight token budget never reaches provider I/O."""
+ agent = ModelAgent(
+ "provider_embedding", "provider-embedding-model", "https://provider.synthetic.invalid/v1", tags=("embedding",)
+ )
+ client = _SyntheticProviderClient()
+ client.embed_with_usage = lambda *_args: (_ for _ in ()).throw(
+ AssertionError("provider must not be called")
+ )
+ config = InMemoryConfigStore()
+ config.set("routing", "embedding_max_tokens_per_request", 3)
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent], client=client),
+ config,
+ embedding_token_counter=UnavailableEmbeddingTokenCounter(),
+ )
+
+ with pytest.raises(TokenCountUnavailable):
+ coordinator.submit_embeddings_batch(["four"])
+
+
+@pytest.mark.parametrize("text", ["한글🙂e\u0301", "<|special|>", "\x00\U0010ffff"])
+def test_unknown_tokenizer_byte_bound_never_becomes_recorded_usage(text) -> None:
+ """Unicode byte proofs admit provider I/O but never masquerade as token usage."""
+ agent = ModelAgent(
+ "provider_embedding", "provider-embedding-model", "https://provider.synthetic.invalid/v1", tags=("embedding",)
+ )
+ client = _SyntheticProviderClient()
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent], client=client),
+ embedding_token_counter=UnavailableEmbeddingTokenCounter(),
+ )
+
+ document = coordinator.complete_embeddings_batch([text])
+
+ assert document["total_tokens"] == len(text.encode("utf-8"))
+
+
+def test_provider_batch_returns_before_terminal_result() -> None:
+ release = threading.Event()
+
+ def runner(requests):
+ release.wait(timeout=1)
+ return [[float(len(request.input_text))] for request in requests], 2
+
+ backend = ProviderEmbeddingBatchBackend(runner)
+ request = EmbeddingBatchRequest(input_text="synthetic input", model="synthetic-model")
+ started = time.monotonic()
+ job = backend.submit([request])
+ assert time.monotonic() - started < 0.1
+ release.set()
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ assert backend.retrieve(job)[0].embedding == [15.0]
+ assert backend.usage(job) == {"prompt_tokens": 2}
+ backend.close()
+
+
+def test_queued_document_exposes_backend_poll_and_registry_retention_contract() -> None:
+ """Queued HTTP documents carry owned cadence/retention, not caller guesses."""
+ release = threading.Event()
+
+ def runner(requests):
+ release.wait(timeout=1)
+ return [[1.0] for _request in requests], len(requests)
+
+ registry = JobRegistryFactory(retention_seconds=123)
+ backend = ProviderEmbeddingBatchBackend(
+ runner,
+ job_registry=registry,
+ claim_lease_seconds=None,
+ )
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([], allow_empty_agents=True),
+ embedding_batch_backend=backend,
+ embedding_token_counter=_SyntheticExactCounter(),
+ job_registry=registry,
+ )
+ job = coordinator.submit_embeddings_batch(["synthetic"], model="synthetic-model")
+
+ document = coordinator.embeddings_batch_document(job.job_id)
+
+ assert document["status"] in {"queued", "running"}
+ assert document["poll_after_ms"] == backend.poll_after_ms
+ assert document["job_retention_ms"] == 123_000
+ release.set()
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ backend.close()
+
+
+def test_provider_reservation_does_not_execute_before_public_registration() -> None:
+ called = threading.Event()
+
+ def runner(requests):
+ called.set()
+ return [[1.0] for _request in requests], len(requests)
+
+ backend = ProviderEmbeddingBatchBackend(runner)
+ job = backend.reserve(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+
+ assert backend.poll(job)["status"] == "reserved"
+ assert called.is_set() is False
+ backend.start(job)
+ assert backend.wait(job, timeout=1)["status"] == "completed"
+ assert called.is_set() is True
+ backend.close()
+
+
+def test_provider_batch_failure_is_terminal_without_payload_leak() -> None:
+ def runner(_requests):
+ raise RuntimeError("synthetic provider failure")
+
+ backend = ProviderEmbeddingBatchBackend(runner)
+ job = backend.submit([EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")])
+ assert backend.wait(job, timeout=1)["status"] == "failed"
+ assert backend.retrieve(job) == []
+ backend.close()
+
+
+def test_provider_batch_failure_preserves_classified_provider_details() -> None:
+ def runner(_requests):
+ raise ProviderUpstreamError(
+ agent_id="embedding_worker",
+ model="embedding-model",
+ error_code="rate_limit_exceeded",
+ message="provider request failed",
+ client_status=429,
+ provider_status=503,
+ retryable=True,
+ transport="embedding",
+ )
+
+ backend = ProviderEmbeddingBatchBackend(runner)
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="synthetic", model="embedding-model")]
+ )
+
+ failure = backend.wait(job, timeout=1)["failure"]
+ assert failure["http_status"] == 429
+ assert failure["provider_code"] == "rate_limit_exceeded"
+ assert failure["retryable"] is True
+ backend.close()
+
+
+def test_provider_batch_cancellation_preserves_the_reason() -> None:
+ release = threading.Event()
+
+ def runner(_requests):
+ release.wait(timeout=1)
+ return [[1.0]], 1
+
+ backend = ProviderEmbeddingBatchBackend(runner)
+ job = backend.submit(
+ [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")]
+ )
+ backend.cancel(job, reason="synchronous request deadline elapsed")
+
+ assert backend.poll(job)["cancellation"] == {
+ "reason": "synchronous request deadline elapsed"
+ }
+ release.set()
+ backend.close()
+
+
+def test_close_waits_for_start_to_submit_work() -> None:
+ submit_entered = threading.Event()
+ release_submit = threading.Event()
+ shutdown_called = threading.Event()
+
+ class Executor:
+ def submit(self, *_args):
+ submit_entered.set()
+ assert release_submit.wait(timeout=1)
+
+ def shutdown(self, **_kwargs):
+ shutdown_called.set()
+
+ backend = ProviderEmbeddingBatchBackend(lambda _requests: ([], 0))
+ job = backend.reserve([])
+ backend._executor = Executor()
+ starter = threading.Thread(target=backend.start, args=(job,))
+ starter.start()
+ assert submit_entered.wait(timeout=1)
+
+ closer = threading.Thread(target=backend.close)
+ closer.start()
+ assert not shutdown_called.wait(timeout=0.05)
+ release_submit.set()
+ starter.join(timeout=1)
+ closer.join(timeout=1)
+ assert shutdown_called.is_set()
+
+
+def test_server_shutdown_closes_embedding_workers() -> None:
+ class ClosingBackend:
+ name = "closing"
+ closed = False
+
+ def close(self):
+ self.closed = True
+
+ backend = ClosingBackend()
+ orchestrator = TaskOrchestrator([ModelAgent("embedding_worker", "embedding-model")])
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ embedding_token_counter=_SyntheticExactCounter(),
+ embedding_batch_backend=backend,
+ )
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="shutdown-token"),
+ coordinator=coordinator,
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+
+ server.shutdown()
+ thread.join(timeout=1)
+
+ assert backend.closed is True
+
+
+def test_server_close_closes_embedding_workers_after_abnormal_exit() -> None:
+ class ClosingBackend:
+ name = "closing"
+ closed = False
+
+ def close(self):
+ self.closed = True
+
+ backend = ClosingBackend()
+ orchestrator = TaskOrchestrator(
+ [ModelAgent("embedding_worker", "embedding-model")]
+ )
+ coordinator = CostRoutingCoordinator(
+ orchestrator,
+ embedding_token_counter=_SyntheticExactCounter(),
+ embedding_batch_backend=backend,
+ )
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="close-token"),
+ coordinator=coordinator,
+ )
+
+ server.server_close()
+
+ assert backend.closed is True
+
+
+def test_remote_embedding_member_selects_provider_backend() -> None:
+ agent = ModelAgent(
+ "synthetic_embedding",
+ "synthetic-embedding-model",
+ base_url="https://synthetic.invalid/v1",
+ tags=("embedding",),
+ )
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent], client=_SyntheticProviderClient()),
+ embedding_token_counter=_SyntheticExactCounter(),
+ )
+ job = coordinator.submit_embeddings_batch(
+ ["synthetic one", "synthetic two"], model=agent.model, agent_id=agent.id
+ )
+ for _attempt in range(100):
+ document = coordinator.embeddings_batch_document(job.job_id)
+ if document["status"] == "completed":
+ break
+ time.sleep(0.01)
+ assert document["status"] == "completed"
+ assert [item["embedding"] for item in document["embeddings"]] == [[13.0], [13.0]]
+
+
+def test_runtime_added_remote_embedding_member_uses_provider_backend() -> None:
+ client = _SyntheticProviderClient()
+ orchestrator = TaskOrchestrator([], client=client, allow_empty_agents=True)
+ coordinator = CostRoutingCoordinator(
+ orchestrator, embedding_token_counter=_SyntheticExactCounter()
+ )
+ orchestrator.add_agent(
+ "default",
+ ModelAgent(
+ "runtime_embedding",
+ "runtime-embedding-model",
+ base_url="https://synthetic.invalid/v1",
+ tags=("embedding",),
+ ).to_config(),
+ )
+
+ document = coordinator.complete_embeddings_batch(
+ ["provider input"], model="contextual-orchestrator", wait_timeout=1
+ )
+
+ assert document["status"] == "completed"
+ assert client.embedding_calls == [["provider input"]]
+
+ orchestrator.add_agent(
+ "default",
+ ModelAgent(
+ "runtime_mock", "runtime-mock-model", base_url="mock://local", tags=("embedding",)
+ ).to_config(),
+ )
+ local_document = coordinator.complete_embeddings_batch(
+ ["local input"], agent_id="runtime_mock"
+ )
+ assert local_document["status"] == "completed"
+ assert client.embedding_calls == [["provider input"]]
+
+
+def test_local_startup_registers_provider_backend_for_recovered_jobs() -> None:
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([], allow_empty_agents=True),
+ embedding_token_counter=_SyntheticExactCounter(),
+ )
+
+ assert isinstance(
+ coordinator._embedding_backends["provider"], ProviderEmbeddingBatchBackend
+ )
+
+
+def test_server_closes_provider_backend_added_after_startup() -> None:
+ orchestrator = TaskOrchestrator([], allow_empty_agents=True)
+ coordinator = CostRoutingCoordinator(
+ orchestrator, embedding_token_counter=_SyntheticExactCounter()
+ )
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="runtime_backend_close_token"),
+ coordinator=coordinator,
+ )
+ closed = threading.Event()
+ coordinator._embedding_backends["provider"].close = closed.set
+
+ server.server_close()
+
+ assert closed.is_set()
+
+
+def test_provider_embedding_requests_are_sharded_by_the_existing_token_limit() -> None:
+ agent = ModelAgent(
+ "synthetic_embedding",
+ "synthetic-embedding-model",
+ base_url="https://synthetic.invalid/v1",
+ tags=("embedding",),
+ )
+ client = _SyntheticProviderClient()
+ coordinator = CostRoutingCoordinator(
+ TaskOrchestrator([agent], client=client),
+ embedding_token_counter=_SyntheticExactCounter(),
+ )
+ coordinator.config.set("routing", "embedding_max_tokens_per_request", 3)
+ coordinator.config.set("routing", "embedding_max_inputs_per_request", 2)
+
+ document = coordinator.complete_embeddings_batch(
+ ["one two", "three four", "five"],
+ model=agent.model,
+ agent_id=agent.id,
+ wait_timeout=1,
+ )
+
+ assert document["status"] == "completed"
+ assert client.embedding_calls == [["one two"], ["three four", "five"]]
diff --git a/tests/test_provider_error_taxonomy.py b/tests/test_provider_error_taxonomy.py
index 9aba81f39..17f5ebd95 100644
--- a/tests/test_provider_error_taxonomy.py
+++ b/tests/test_provider_error_taxonomy.py
@@ -47,6 +47,22 @@ def _body_http_error(code: int, payload: dict) -> urllib.error.HTTPError:
# -- message redaction --------------------------------------------------------
+def test_reclassification_preserves_failure_and_updates_boundary_transport() -> None:
+ original = classify_provider_failure(
+ _http_error(404), agent_id="synthetic-agent", model="synthetic-model", transport="passthrough"
+ )
+ classified = classify_provider_failure(
+ original,
+ agent_id="synthetic-agent",
+ model="synthetic-model",
+ transport="structured_synthesis",
+ )
+ assert classified is not original
+ assert classified.error_code == original.error_code
+ assert classified.provider_status == original.provider_status
+ assert classified.transport == "structured_synthesis"
+
+
def test_safe_message_prefers_nested_provider_error_fields() -> None:
"""``error.message`` / ``error.code`` / top-level fields are the only pass-through."""
nested = safe_provider_message(_body_http_error(400, {"error": {"message": "max_tokens too large"}}))
@@ -61,6 +77,31 @@ def test_safe_message_prefers_nested_provider_error_fields() -> None:
assert detail == "validation failed"
+def test_safe_message_keeps_actionable_schema_diagnostics_without_payloads() -> None:
+ """Schema field names are useful; field values and request bodies remain private."""
+ actionable = "'messages' must contain the word 'json' to use json_object"
+ assert safe_provider_message(
+ _body_http_error(400, {"error": {"message": actionable}})
+ ) == "messages must mention json when response_format is json_object"
+ for diagnostic in (
+ "messages=[{'role':'user','content':'customer secret'}]",
+ '"messages": [{"role":"user","content":"customer secret"}]',
+ "'content': 'customer secret'",
+ "prompt=customer secret",
+ "input: customer secret",
+ ):
+ assert safe_provider_message(
+ _body_http_error(400, {"error": {"message": diagnostic}})
+ ) is None
+
+ assert safe_provider_message(
+ _body_http_error(
+ 400,
+ {"error": {"message": "messages rejected; customer-private-text"}},
+ )
+ ) is None
+
+
def test_safe_message_hides_unparseable_bodies_and_urls() -> None:
"""Non-JSON bodies return None so URLs/reasons never leak through fallback text."""
assert safe_provider_message(_http_error(500, b"upstream-secret http://10.0.0.9/internal")) is None
diff --git a/tests/test_provider_usage_capture.py b/tests/test_provider_usage_capture.py
index 18fc4665c..6e5f89b8d 100644
--- a/tests/test_provider_usage_capture.py
+++ b/tests/test_provider_usage_capture.py
@@ -1,8 +1,8 @@
-"""Real provider-reported usage capture — prefer reported tokens over the estimate.
+"""Provider-reported usage capture and explicit unavailable fallbacks.
A gateway that already sees provider `usage` should bill on it, not a char heuristic.
These assert reported completion_tokens flow into spend_analytics and are labeled,
-while the mock path stays estimated.
+while the mock path stays unavailable.
"""
from __future__ import annotations
@@ -54,8 +54,8 @@ def test_reported_usage_preferred_and_labeled() -> None:
row = next(r for r in orchestrator.spend_analytics()["by_model"] if r["model"] == "priced-model")
assert row["usage_source"] == "reported"
- assert row["output_tokens"] == 50 # reported completion tokens, not the char estimate
- assert row["estimated_cost_usd"] == round(50 / 1_000_000 * 10.0, 6) # cost from reported tokens
+ assert row["output_tokens"] == 50
+ assert row["cost_usd"] == round(50 / 1_000_000 * 10.0, 6)
def test_reported_prompt_tokens_surface_in_totals() -> None:
@@ -65,29 +65,28 @@ def test_reported_prompt_tokens_surface_in_totals() -> None:
orchestrator.run([{"role": "user", "content": "route once"}])
totals = orchestrator.spend_analytics()["totals"]
assert totals["prompt_tokens_source"] == "reported"
- assert totals["reported_prompt_tokens"] == 5 # provider-reported prompt tokens, one route step
+ assert totals["prompt_tokens"] == 5
-def test_mock_prompt_tokens_source_is_estimated() -> None:
+def test_mock_prompt_tokens_source_is_unavailable() -> None:
orchestrator = TaskOrchestrator([ModelAgent("general_agent", "free-model", tags=("reasoning",))])
# Isolate mock worker accounting from the optional model-judge extra,
# whose real invocation has its own reported prompt-token evidence.
with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None):
orchestrator.run([{"role": "user", "content": "no usage here"}])
totals = orchestrator.spend_analytics()["totals"]
- assert totals["prompt_tokens_source"] == "estimated"
- assert totals["reported_prompt_tokens"] == 0
- assert totals["estimated_prompt_tokens"] > 0
+ assert totals["prompt_tokens_source"] == "unavailable"
+ assert totals["prompt_tokens"] is None
-def test_mock_path_stays_estimated() -> None:
+def test_mock_path_stays_unavailable() -> None:
orchestrator = TaskOrchestrator([ModelAgent("general_agent", "free-model", tags=("reasoning",))])
with patch.object(orchestrator_module, "_resolve_fast_mlsirm_components", return_value=None):
orchestrator.run([{"role": "user", "content": "do the work"}])
row = next(r for r in orchestrator.spend_analytics()["by_model"] if r["model"] == "free-model")
- assert row["usage_source"] == "estimated"
- assert row["output_tokens"] == row["estimated_output_tokens"] # falls back to the estimate
+ assert row["usage_source"] == "unavailable"
+ assert row["output_tokens"] is None
def test_conduct_all_steps_reported() -> None:
diff --git a/tests/test_reasoning_none_text_empty_logprobs_zero_http_honesty.py b/tests/test_reasoning_none_text_empty_logprobs_zero_http_honesty.py
index 209924412..168a29621 100644
--- a/tests/test_reasoning_none_text_empty_logprobs_zero_http_honesty.py
+++ b/tests/test_reasoning_none_text_empty_logprobs_zero_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "reasoning_none_text_empty_logprobs_zero_http_honesty_token" # noqa: S105
@@ -48,10 +48,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_routing_endpoint_constraint.py b/tests/test_routing_endpoint_constraint.py
new file mode 100644
index 000000000..88410ddde
--- /dev/null
+++ b/tests/test_routing_endpoint_constraint.py
@@ -0,0 +1,349 @@
+"""Request-scoped endpoint routing is exact, isolated, and fail closed."""
+
+from __future__ import annotations
+
+import json
+import threading
+import urllib.error
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator.orchestrator import EndpointUnavailableError, ModelClient
+from contextual_orchestrator.server import (
+ RequestError,
+ SecurityConfig,
+ _validate_routing,
+ build_server,
+)
+
+
+class _RecordingClient(ModelClient):
+ """Record selected agents while returning deterministic synthetic text."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.agent_ids: list[str] = []
+
+ def chat(self, agent: ModelAgent, messages: list, temperature: float = 0.2, **_kwargs) -> str: # type: ignore[override]
+ self.agent_ids.append(agent.id)
+ return json.dumps({"workflow_required": False})
+
+ def proxy_send( # type: ignore[override]
+ self, agent: ModelAgent, endpoint: str, payload: dict
+ ) -> dict:
+ self.agent_ids.append(agent.id)
+ if endpoint == "responses":
+ return {
+ "object": "response",
+ "model": payload["model"],
+ "output": [
+ {
+ "type": "message",
+ "content": [{"type": "output_text", "text": "synthetic response"}],
+ }
+ ],
+ }
+ return {
+ "id": "chatcmpl-endpoint-test",
+ "object": "chat.completion",
+ "model": payload["model"],
+ "choices": [
+ {
+ "index": 0,
+ "message": {"role": "assistant", "content": "synthetic response"},
+ "finish_reason": "stop",
+ }
+ ],
+ }
+
+
+def _orchestrator() -> TaskOrchestrator:
+ return TaskOrchestrator(
+ [
+ ModelAgent(
+ "agent_a",
+ "model-a",
+ base_url="https://a.example/v1",
+ group_name="shared_reasoning_model",
+ ),
+ ModelAgent("agent_b", "model-b", base_url="https://b.example/v1"),
+ ],
+ client=_RecordingClient(),
+ cache_ttl=60,
+ )
+
+
+def test_endpoint_scope_filters_every_role_and_does_not_leak() -> None:
+ orchestrator = _orchestrator()
+ with orchestrator.routing_endpoint_scope("https://a.example", "orchestrator/auto"):
+ for role in ("thinker", "worker", "verifier", "judge", "synthesizer"):
+ assert [agent.id for agent in orchestrator._ranked_agents("task", role)] == [
+ "agent_a"
+ ]
+ assert {agent.id for agent in orchestrator._ranked_agents("task", "worker")} == {
+ "agent_a",
+ "agent_b",
+ }
+
+
+def test_endpoint_scope_is_concurrent_and_rejects_model_conflicts() -> None:
+ orchestrator = _orchestrator()
+ barrier = threading.Barrier(2)
+
+ def select(endpoint: str) -> str:
+ with orchestrator.routing_endpoint_scope(endpoint, "orchestrator/auto"):
+ barrier.wait()
+ return orchestrator._ranked_agents("task", "worker")[0].id
+
+ with ThreadPoolExecutor(max_workers=2) as executor:
+ assert set(executor.map(select, ("https://a.example", "https://b.example"))) == {
+ "agent_a",
+ "agent_b",
+ }
+ with (
+ pytest.raises(EndpointUnavailableError),
+ orchestrator.routing_endpoint_scope("https://a.example", "model-b"),
+ ):
+ pass
+
+
+def test_endpoint_scope_partitions_response_and_triage_caches() -> None:
+ orchestrator = _orchestrator()
+ messages = [{"role": "user", "content": "same synthetic prompt"}]
+ with orchestrator.routing_endpoint_scope("https://a.example", "orchestrator/auto"):
+ cache_a = orchestrator._cache_key(messages, "route")
+ orchestrator._triage_workflow_required("same synthetic prompt")
+ with orchestrator.routing_endpoint_scope("https://b.example", "orchestrator/auto"):
+ cache_b = orchestrator._cache_key(messages, "route")
+ orchestrator._triage_workflow_required("same synthetic prompt")
+ assert cache_a != cache_b
+ assert len(orchestrator._triage_cache) == 2
+
+
+@pytest.mark.parametrize(
+ "endpoint",
+ [
+ "https://user:secret@a.example",
+ "https://a.example?x=1",
+ "https://a.example#x",
+ "ftp://a.example",
+ "https://a.example:bad",
+ ],
+)
+def test_invalid_endpoint_selector_is_rejected(endpoint: str) -> None:
+ with pytest.raises(RequestError) as exc_info:
+ _validate_routing({"endpoint": endpoint}, allow_endpoint=True)
+ assert exc_info.value.code == "endpoint_unavailable"
+
+
+def test_malformed_ipv6_endpoint_selector_preserves_error_contract() -> None:
+ with pytest.raises(RequestError) as exc_info:
+ _validate_routing({"endpoint": "https://[::1"}, allow_endpoint=True)
+ assert exc_info.value.code == "endpoint_unavailable"
+
+
+def test_endpoint_capacity_check_performs_no_provider_io() -> None:
+ class _ProbeRecordingClient(_RecordingClient):
+ def __init__(self) -> None:
+ super().__init__()
+ self.embed_calls = 0
+
+ def embed(self, agent: ModelAgent, texts: list[str]) -> list[list[float]]:
+ self.embed_calls += 1
+ return [[0.0] for _text in texts]
+
+ client = _ProbeRecordingClient()
+ orchestrator = TaskOrchestrator(
+ [
+ ModelAgent(
+ "chat_agent",
+ "chat-model",
+ base_url="https://a.example/v1",
+ ),
+ ModelAgent(
+ "embedding_agent",
+ "embedding-model",
+ base_url="https://a.example/v1",
+ tags=("embedding",),
+ ),
+ ],
+ client=client,
+ )
+
+ with orchestrator.routing_endpoint_scope(
+ "https://a.example", TaskOrchestrator.AUTO_MODEL
+ ):
+ pass
+
+ assert client.agent_ids == []
+ assert client.embed_calls == 0
+
+
+def test_endpoint_is_limited_to_supported_surfaces_and_forces_sync() -> None:
+ with pytest.raises(RequestError) as exc_info:
+ _validate_routing({"endpoint": "https://a.example"})
+ assert exc_info.value.code == "invalid_routing"
+ assert _validate_routing(
+ {"endpoint": "https://a.example", "priority": "bulk"},
+ allow_endpoint=True,
+ ) == {
+ "endpoint": "https://a.example",
+ "priority": "bulk",
+ "channel": "sync",
+ }
+
+
+def _post_json(server: object, path: str, body: dict) -> tuple[int, dict]:
+ port = server.server_address[1] # type: ignore[attr-defined]
+ request = urllib.request.Request(
+ f"http://127.0.0.1:{port}{path}",
+ data=json.dumps(body).encode(),
+ headers={
+ "authorization": "Bearer endpoint-test-token",
+ "content-type": "application/json",
+ },
+ method="POST",
+ )
+ try:
+ with urllib.request.urlopen(request, timeout=10) as response:
+ return response.status, json.loads(response.read())
+ except urllib.error.HTTPError as exc:
+ return exc.code, json.loads(exc.read())
+
+
+@pytest.mark.parametrize(
+ ("path", "payload"),
+ [
+ (
+ "/v1/chat/completions",
+ {
+ "model": "orchestrator/auto",
+ "messages": [{"role": "user", "content": "synthetic answer"}],
+ },
+ ),
+ ("/v1/responses", {"model": "orchestrator/auto", "input": "synthetic answer"}),
+ ],
+)
+def test_http_surfaces_constrain_candidates_and_preserve_envelopes(path: str, payload: dict) -> None:
+ orchestrator = _orchestrator()
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="endpoint-test-token"),
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ payload["routing"] = {"endpoint": "https://a.example"}
+ status, document = _post_json(server, path, payload)
+ assert status == 200, document
+ assert document["object"] in {"chat.completion", "response"}
+ assert orchestrator.client.agent_ids
+ assert set(orchestrator.client.agent_ids) == {"agent_a"}
+
+ payload["routing"] = {"endpoint": "https://missing.example"}
+ status, document = _post_json(server, path, payload)
+ assert status == 400
+ assert document["error"]["code"] == "endpoint_unavailable"
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+@pytest.mark.parametrize(
+ ("path", "payload"),
+ [
+ (
+ "/v1/chat/completions",
+ {
+ "model": " model-a ",
+ "messages": [{"role": "user", "content": "synthetic answer"}],
+ },
+ ),
+ ("/v1/responses", {"model": " shared-reasoning-model ", "input": "synthetic answer"}),
+ ],
+)
+def test_http_endpoint_scope_accepts_normalized_models_and_group_aliases(
+ path: str, payload: dict
+) -> None:
+ orchestrator = _orchestrator()
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="endpoint-test-token"),
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, document = _post_json(
+ server,
+ path,
+ {**payload, "routing": {"endpoint": "https://a.example"}},
+ )
+ assert status == 200, document
+ assert set(orchestrator.client.agent_ids) == {"agent_a"}
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
+
+
+@pytest.mark.parametrize(
+ ("path", "payload"),
+ [
+ (
+ "/v1/chat/completions",
+ {
+ "messages": [{"role": "user", "content": "synthetic answer"}],
+ },
+ ),
+ (
+ "/v1/responses",
+ {
+ "model": TaskOrchestrator.FREE_MODEL,
+ "input": "synthetic answer",
+ },
+ ),
+ ],
+)
+def test_http_endpoint_scope_rejects_endpoint_without_local_virtual_capacity(
+ path: str, payload: dict
+) -> None:
+ orchestrator = TaskOrchestrator(
+ [
+ ModelAgent(
+ "embedding_only",
+ "embedding-model",
+ base_url="https://paid.example/v1",
+ tags=("embedding",),
+ ),
+ ModelAgent(
+ "free_elsewhere",
+ "free-model",
+ base_url="https://free.example/v1",
+ tags=("cost:free",),
+ ),
+ ],
+ client=_RecordingClient(),
+ cache_ttl=60,
+ )
+ server = build_server(
+ orchestrator,
+ port=0,
+ security=SecurityConfig(auth_token="endpoint-test-token"),
+ )
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ try:
+ status, document = _post_json(
+ server,
+ path,
+ {**payload, "routing": {"endpoint": "https://paid.example"}},
+ )
+ assert status == 400
+ assert document["error"]["code"] == "endpoint_unavailable"
+ finally:
+ server.shutdown()
+ thread.join(timeout=5)
diff --git a/tests/test_service_tier_encoding_format_strip_http_honesty.py b/tests/test_service_tier_encoding_format_strip_http_honesty.py
index d1fba4d2f..4854808e8 100644
--- a/tests/test_service_tier_encoding_format_strip_http_honesty.py
+++ b/tests/test_service_tier_encoding_format_strip_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "service_tier_encoding_format_strip_http_honesty_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_spend_analytics.py b/tests/test_spend_analytics.py
index e03a61414..7ef56568d 100644
--- a/tests/test_spend_analytics.py
+++ b/tests/test_spend_analytics.py
@@ -1,153 +1,74 @@
-"""Spend observability — estimated per-model token + cost analytics.
-
-The LLM-gateway category monetizes on spend tracking; this product discarded usage
-entirely. These assert the token estimate, per-model aggregation, cost math when a
-price is configured, honest nulls when it is not, and the read-only HTTP endpoint.
-"""
+"""Authoritative per-model token and cost analytics."""
from __future__ import annotations
import json
-from pathlib import Path
-import sys
import threading
-import urllib.error
import urllib.request
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
-from contextual_orchestrator.orchestrator import estimate_tokens # noqa: E402
-from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
+from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator.server import SecurityConfig, build_server
+from contextual_orchestrator.token_counting import TokenCountUnavailable
-def test_estimate_tokens_heuristic() -> None:
- assert estimate_tokens("") == 0
- assert estimate_tokens("abcd") == 1 # 4 chars ~ 1 token
- assert estimate_tokens("abcde") == 2 # (5 + 3) // 4
- assert estimate_tokens("a" * 400) == 100
+class _ExactCounter:
+ """Injected exact raw-output counter for synthetic fixtures."""
+ def count_text(self, text: str, model: str) -> int:
+ if model != "gpt-4":
+ raise TokenCountUnavailable("unknown synthetic model")
+ return len(text.encode("utf-8"))
-def test_spend_without_prices_reports_null_cost() -> None:
- orchestrator = TaskOrchestrator([ModelAgent("general_agent", "free-model", tags=("reasoning",))])
- orchestrator.run([{"role": "user", "content": "estimate my spend"}])
- report = orchestrator.spend_analytics()
- assert report["pricing_configured"] is False
- assert report["totals"]["run_count"] == 1
- assert report["totals"]["estimated_output_tokens"] > 0
- assert report["totals"]["estimated_cost_usd"] is None
- assert "free-model" in report["unpriced_models"]
- row = next(r for r in report["by_model"] if r["model"] == "free-model")
- assert row["estimated_cost_usd"] is None
+def _orchestrator(*, price: float | None = None) -> TaskOrchestrator:
+ return TaskOrchestrator(
+ [ModelAgent("general_agent", "gpt-4", tags=("reasoning",))],
+ price_per_million={"gpt-4": price} if price is not None else None,
+ token_counter=_ExactCounter(),
+ )
-def test_spend_with_price_computes_cost() -> None:
- orchestrator = TaskOrchestrator(
- [ModelAgent("general_agent", "priced-model", tags=("reasoning",))],
- price_per_million={"priced-model": 10.0},
- )
- orchestrator.run([{"role": "user", "content": "compute my cost please"}])
+def test_exact_output_without_prompt_usage_is_explicitly_unavailable() -> None:
+ orchestrator = _orchestrator()
+ orchestrator.run([{"role": "user", "content": "account for this"}])
report = orchestrator.spend_analytics()
+ row = report["by_model"][0]
- assert report["pricing_configured"] is True
- row = next(r for r in report["by_model"] if r["model"] == "priced-model")
- assert row["price_per_million_usd"] == 10.0
- # Cost is billed on output_tokens (provider-reported usage when available,
- # e.g. from a real realtime judge call, else the text estimate) -- not on
- # estimated_output_tokens, which stays a text-length estimate even when a
- # judge call reports real usage that diverges from it (see
- # test_spend_analytics_bills_reported_judge_usage_not_its_text_estimate).
- expected = round(row["output_tokens"] / 1_000_000 * 10.0, 6)
- assert row["estimated_cost_usd"] == expected
- assert report["totals"]["estimated_cost_usd"] == expected # single priced model
- assert report["unpriced_models"] == []
-
-
-def test_spend_analytics_bills_reported_judge_usage_not_its_text_estimate() -> None:
- """A realtime judge's reported usage can diverge from its own text-length estimate.
-
- ``spend_analytics`` folds a completed judge call's usage into the same
- per-model bucket as its worker steps (honest accounting: a real judge
- call is a real incurred cost). When the judge's provider-reported
- ``completion_tokens`` differs from ``estimate_tokens(judge_output_text)``
- -- as happens whenever a real fast-mlsirm judge call is exercised, only
- possible where that optional native dependency is actually installed --
- ``estimated_output_tokens`` (a text-length estimate, purely informational)
- and ``output_tokens`` (the actual billing basis) must diverge too, and
- cost must track ``output_tokens``, never the estimate. This is the exact
- field this repo's own CI hit and a machine without fast-mlsirm installed
- cannot reproduce locally.
- """
- orchestrator = TaskOrchestrator(
- [ModelAgent("general_agent", "priced-model", tags=("reasoning",))],
- price_per_million={"priced-model": 10.0},
- )
- worker_output = "worker answer"
- judge_text = "judge rationale, much longer than its reported token count"
- orchestrator._replace_workflow_run(
- {
- "workflow_run_id": "run_judge_usage_divergence",
- "prompt_text": "compute my cost please",
- "trace": [
- {
- "id": 0,
- "role": "worker",
- "agent_id": "general_agent",
- "output": worker_output,
- },
- ],
- "verification": {
- "judge_agent_id": "general_agent",
- "judge_model": "priced-model",
- "judge_usage": {"completion_tokens": 3},
- "judge_output_text": judge_text,
- },
- }
- )
- report = orchestrator.spend_analytics()
+ assert report["measurement_status"] == "unavailable"
+ assert report["totals"]["output_tokens"] > 0
+ assert report["totals"]["prompt_tokens"] is None
+ assert report["totals"]["cost_usd"] is None
+ assert row["usage_source"] == "mixed"
+ assert row["cost_usd"] is None
+ assert not any("estimated" in key for key in row | report["totals"])
- row = next(r for r in report["by_model"] if r["model"] == "priced-model")
- worker_tokens = estimate_tokens(worker_output)
- judge_text_estimate = estimate_tokens(judge_text)
- assert row["estimated_output_tokens"] == worker_tokens + judge_text_estimate
- assert row["output_tokens"] == worker_tokens + 3
- assert row["output_tokens"] != row["estimated_output_tokens"]
- expected = round(row["output_tokens"] / 1_000_000 * 10.0, 6)
- assert row["estimated_cost_usd"] == expected
- assert report["totals"]["estimated_cost_usd"] == expected
-
-
-def test_call_time_price_overrides_instance() -> None:
- orchestrator = TaskOrchestrator(
- [ModelAgent("general_agent", "priced-model", tags=("reasoning",))],
- price_per_million={"priced-model": 10.0},
- )
- orchestrator.run([{"role": "user", "content": "override the price"}])
- row = next(
- r for r in orchestrator.spend_analytics(price_per_million={"priced-model": 20.0})["by_model"]
- if r["model"] == "priced-model"
- )
- assert row["price_per_million_usd"] == 20.0
+def test_exact_output_cost_uses_operator_price() -> None:
+ orchestrator = _orchestrator(price=10.0)
+ orchestrator.run([{"role": "user", "content": "calculate exact output cost"}])
+ report = orchestrator.spend_analytics()
+ row = report["by_model"][0]
+ expected = row["output_tokens"] / 1_000_000 * 10.0
+ assert row["cost_usd"] == expected
+ assert report["totals"]["cost_usd"] == expected
-def test_spend_empty_when_no_runs() -> None:
- report = TaskOrchestrator([ModelAgent("general_agent", "some-model")]).spend_analytics()
+
+def test_empty_analytics_are_zero_not_estimated() -> None:
+ report = _orchestrator().spend_analytics()
assert report["totals"]["run_count"] == 0
+ assert report["totals"]["output_tokens"] == 0
assert report["by_model"] == []
- assert report["totals"]["estimated_output_tokens"] == 0
-def test_http_spend_endpoint_returns_report() -> None:
+def test_http_spend_endpoint_preserves_unavailable_status() -> None:
token = "spend_token"
- orchestrator = TaskOrchestrator([ModelAgent("general_agent", "priced-model", tags=("reasoning",))])
+ orchestrator = _orchestrator()
orchestrator.run([{"role": "user", "content": "seed a run"}])
server = build_server(orchestrator, port=0, security=SecurityConfig(auth_token=token))
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
- port = server.server_address[1]
request = urllib.request.Request(
- f"http://127.0.0.1:{port}/api/v1/spend_analytics/latest",
+ f"http://127.0.0.1:{server.server_address[1]}/api/v1/spend_analytics/latest",
headers={"authorization": f"Bearer {token}", "connection": "close"},
)
try:
@@ -156,13 +77,5 @@ def test_http_spend_endpoint_returns_report() -> None:
finally:
server.shutdown()
assert status == 200
- assert body["measurement_status"] == "local_runtime_estimate"
- assert body["totals"]["run_count"] == 1
-
-
-if __name__ == "__main__":
- for name, fn in sorted(globals().items()):
- if name.startswith("test_") and callable(fn):
- fn()
- print(f"ok {name}")
- print("ok")
+ assert body["measurement_status"] == "unavailable"
+ assert body["totals"]["prompt_tokens"] is None
diff --git a/tests/test_stream_options_null_flags_noop_http_honesty.py b/tests/test_stream_options_null_flags_noop_http_honesty.py
index 5101bd67a..51a773155 100644
--- a/tests/test_stream_options_null_flags_noop_http_honesty.py
+++ b/tests/test_stream_options_null_flags_noop_http_honesty.py
@@ -284,21 +284,17 @@ def base_url(self) -> str:
return f"http://127.0.0.1:{self._server.server_address[1]}"
-def test_http_chat_tools_streams_estimated_usage_when_provider_omits_it() -> None:
- """A tools-capable provider that omits ``usage`` gets an honest estimate, not a fake "reported" label.
+def test_http_chat_tools_streams_unavailable_usage_when_provider_omits_it() -> None:
+ """A tools-capable provider that omits ``usage`` reports unavailable.
Regression for a Devin finding on #925: the PR's own rationale for
narrowing the ``stream_options.include_usage`` rejection to exclude
``tools`` passthrough assumed the one upstream call "always" carries real
provider usage. That's the common case, not a guarantee --
``ModelClient.proxy_send`` returns the provider's raw JSON verbatim with
- no ``usage`` requirement. ``_chat_response_sse_chunks`` already has a
- fallback for exactly this (the same one already exercised for the
- non-tools streaming path): estimate from the visible content/tool_calls
- and label it ``usage_source: "estimated"`` rather than fabricate
- ``"reported"``. This proves that fallback over a real (if fake) HTTP
- provider response, since ``mock://`` agents always inject usage and can
- never exercise it.
+ no ``usage`` requirement. This proves that the gateway returns an explicit
+ unavailable outcome over a real (if fake) HTTP provider response, since
+ ``mock://`` agents always inject usage and cannot exercise this boundary.
"""
with _NoUsageToolProvider() as provider:
# The "local://" scheme is this codebase's sanctioned way to point an
@@ -354,29 +350,22 @@ def test_http_chat_tools_streams_estimated_usage_when_provider_omits_it() -> Non
]
usage_frames = [frame for frame in frames if frame.get("choices") == []]
assert len(usage_frames) == 1, frames
- usage = usage_frames[0]["usage"]
- assert usage["usage_source"] == "estimated", usage
- assert usage["prompt_tokens"] > 0
- assert usage["completion_tokens"] > 0
+ assert usage_frames[0]["usage"] is None
+ assert usage_frames[0]["usage_measurement_status"] == "unavailable"
finally:
server.shutdown()
thread.join(timeout=5)
-def test_http_chat_tools_estimated_usage_counts_the_tool_schema() -> None:
- """A larger ``tools`` schema must widen the estimated prompt token count.
+def test_http_chat_tools_do_not_reconstruct_tool_schema_usage() -> None:
+ """Tool schemas remain unavailable instead of being reconstructed.
- Regression for a Devin finding on #925: the estimated-usage fallback in
- ``_chat_response_sse_chunks`` used to build its prompt-token estimate
- from ``body["messages"]`` alone, excluding ``body["tools"]`` entirely.
- OpenAI-compatible providers count tool definitions (names, descriptions,
- JSON schemas) toward prompt usage, so a request with a large tool schema
- would get a materially understated estimate. Proves the fix by sending
- the identical messages with a tiny tools list vs. a large one and
- asserting the estimated prompt_tokens strictly increases.
+ A generic gateway cannot reproduce provider serialization for names,
+ descriptions, and JSON schemas. Send identical messages with small and
+ large tool sets and require the same unavailable outcome for both.
"""
- def estimated_prompt_tokens(tools: list[dict]) -> int:
+ def usage_frame(tools: list[dict]) -> dict:
with _NoUsageToolProvider() as provider:
local_base_url = provider.base_url.replace("http://", "local://", 1)
orchestrator = TaskOrchestrator(
@@ -417,7 +406,7 @@ def estimated_prompt_tokens(tools: list[dict]) -> int:
]
usage_frames = [frame for frame in frames if frame.get("choices") == []]
assert len(usage_frames) == 1, frames
- return usage_frames[0]["usage"]["prompt_tokens"]
+ return usage_frames[0]
finally:
server.shutdown()
thread.join(timeout=5)
@@ -446,9 +435,9 @@ def estimated_prompt_tokens(tools: list[dict]) -> int:
for index in range(5)
]
- small_estimate = estimated_prompt_tokens(small_tools)
- large_estimate = estimated_prompt_tokens(large_tools)
- assert large_estimate > small_estimate, (small_estimate, large_estimate)
+ for frame in (usage_frame(small_tools), usage_frame(large_tools)):
+ assert frame["usage"] is None
+ assert frame["usage_measurement_status"] == "unavailable"
def test_http_chat_response_format_only_streams_still_reject_include_usage() -> None:
@@ -504,8 +493,8 @@ def test_http_chat_rejects_non_boolean_non_null_flag() -> None:
test_http_responses_accepts_stream_options_null_flags()
test_http_chat_accepts_include_usage_true()
test_http_chat_tools_streams_include_reported_usage()
- test_http_chat_tools_streams_estimated_usage_when_provider_omits_it()
- test_http_chat_tools_estimated_usage_counts_the_tool_schema()
+ test_http_chat_tools_streams_unavailable_usage_when_provider_omits_it()
+ test_http_chat_tools_do_not_reconstruct_tool_schema_usage()
test_http_chat_response_format_only_streams_still_reject_include_usage()
test_http_chat_rejects_non_boolean_non_null_flag()
print("ok")
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
index 8994b93c6..1fd4518b6 100644
--- a/tests/test_streaming.py
+++ b/tests/test_streaming.py
@@ -135,7 +135,6 @@ def test_structured_sse_tool_call_deltas_include_indices() -> None:
},
model="tool-model",
include_usage=False,
- prompt_text="tool call",
)
tool_deltas = [
@@ -173,7 +172,6 @@ def test_structured_sse_tool_call_deltas_coexist_with_reported_usage() -> None:
},
model="tool-model",
include_usage=True,
- prompt_text="tool call",
)
tool_deltas = [
@@ -221,7 +219,6 @@ def test_structured_sse_normal_chunks_carry_null_usage_when_include_usage() -> N
},
model="tool-model",
include_usage=True,
- prompt_text="hi",
)
normal_chunks = [chunk for chunk in chunks if chunk["choices"] != []]
diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py
index 0c71b34ec..542cf2a23 100644
--- a/tests/test_telemetry.py
+++ b/tests/test_telemetry.py
@@ -83,6 +83,8 @@ def test_session_and_attribute_boundaries_reject_unsafe_values():
assert telemetry_module._normalize_session_id(value) is None
assert session_id_from_metadata(None) is None
assert telemetry_module._safe_attributes({"server.port": object()}) == {}
+ assert telemetry_module._safe_attributes({"server.address": "10.0.0.9"}) == {}
+ assert telemetry_module._safe_attributes({"server.address": "fd00::9"}) == {}
token = set_session_id("session-safe")
try:
@@ -847,6 +849,37 @@ def capture(name, attributes):
]
+def test_readiness_probe_uses_separate_operation_telemetry(monkeypatch):
+ """A startup readiness probe cannot masquerade as a caller request attempt."""
+ captured = []
+
+ @contextmanager
+ def capture(name, attributes):
+ captured.append((name, attributes))
+ yield None
+
+ client = ModelClient()
+ agent = ModelAgent(
+ "provider_agent",
+ "model-x",
+ base_url="https://provider.example/v1",
+ credential_key="",
+ provider_name="openai",
+ group_name="model-family-x",
+ )
+ monkeypatch.setattr(orchestrator_module, "traced", capture)
+ monkeypatch.setattr(client, "_validate_provider", lambda unused_agent: None)
+ monkeypatch.setattr(
+ client, "_send_raw_with_retry", lambda *_args, **_kwargs: {"ok": True}
+ )
+
+ assert client.probe_structured_chat(agent, {"messages": []}) == {"ok": True}
+ assert captured[0][0] == "capability_probe.chat model-x"
+ assert captured[0][1]["contextual_orchestrator.operation_kind"] == "capability_probe"
+ assert captured[0][1]["contextual_orchestrator.model_group"] == "model_family_x"
+ assert captured[0][1]["contextual_orchestrator.fallback_outcome"] == "not_attempted"
+
+
def test_traced_starts_safe_client_span_with_error_type_and_no_raw_exception(monkeypatch, caplog):
"""Failures remain classifiable without recording raw exception or session data."""
tracer = MagicMock()
@@ -878,7 +911,11 @@ def test_traced_starts_safe_client_span_with_error_type_and_no_raw_exception(mon
span.record_exception.assert_not_called()
# Failures record the CLASSIFIED cause family (network timeout here), not
# the Python exception class, and never the exception text.
- span.set_attribute.assert_called_once_with("error.type", "provider_connection_error")
+ span.set_attribute.assert_any_call("error.type", "provider_connection_error")
+ span.set_attribute.assert_any_call(
+ "contextual_orchestrator.error_summary",
+ "provider_connection_error",
+ )
assert "provider-response-secret" not in caplog.text
assert "session-secret" not in caplog.text
@@ -903,6 +940,129 @@ def test_traced_records_upstream_status_for_http_failures(monkeypatch):
)
+def test_traced_logs_actionable_bounded_failure_evidence(monkeypatch, caplog):
+ """Operators get a safe cause, model group, status, and fallback outcome."""
+ import urllib.error
+
+ tracer = MagicMock()
+ span = tracer.start_as_current_span.return_value.__enter__.return_value
+ monkeypatch.setattr(telemetry_module.trace, "get_tracer", lambda unused_name: tracer)
+ message = (
+ "'messages' must contain the word 'json' in some form, to use "
+ "'response_format' of type 'json_object'.No fallback model group found; "
+ "customer-private-text"
+ )
+ body = io.BytesIO(json.dumps({"error": {"message": message}}).encode())
+
+ with pytest.raises(urllib.error.HTTPError):
+ with traced(
+ "capability_probe.chat gpt-4.1",
+ {
+ "contextual_orchestrator.model_group": "gpt-4.1",
+ "contextual_orchestrator.fallback_outcome": "not_attempted",
+ },
+ ):
+ raise urllib.error.HTTPError("https://private.example", 400, "bad", None, body)
+
+ span.set_attribute.assert_any_call("error.type", "invalid_request_error")
+ span.set_attribute.assert_any_call("contextual_orchestrator.provider_status_code", 400)
+ span.set_attribute.assert_any_call(
+ "contextual_orchestrator.error_summary",
+ "messages must mention json when response_format is json_object",
+ )
+ assert "provider_status=400" in caplog.text
+ assert "model_group=gpt-4.1" in caplog.text
+ assert "fallback_outcome=not_attempted" in caplog.text
+ assert "messages must mention json when response_format is json_object" in caplog.text
+ assert "No fallback model group" not in caplog.text
+ assert "customer-private-text" not in caplog.text
+ assert "private.example" not in caplog.text
+
+
+def test_traced_recognizes_litellm_prefixed_json_object_diagnostic(
+ monkeypatch, caplog
+):
+ """A gateway prefix cannot hide Azure's actionable JSON-object contract."""
+ import urllib.error
+
+ tracer = MagicMock()
+ span = tracer.start_as_current_span.return_value.__enter__.return_value
+ monkeypatch.setattr(telemetry_module.trace, "get_tracer", lambda unused_name: tracer)
+ message = (
+ "AzureException BadRequestError - 'messages' must contain the word 'json' "
+ "in some form, to use 'response_format' of type 'json_object'."
+ "No fallback model group found; customer-private-text"
+ )
+ body = io.BytesIO(json.dumps({"error": {"message": message}}).encode())
+
+ with pytest.raises(urllib.error.HTTPError):
+ with traced(
+ "capability_probe.chat gpt-4.1",
+ {"contextual_orchestrator.model_group": "gpt-4.1"},
+ ):
+ raise urllib.error.HTTPError("https://private.example", 400, "bad", None, body)
+
+ span.set_attribute.assert_any_call(
+ "contextual_orchestrator.error_summary",
+ "messages must mention json when response_format is json_object",
+ )
+ assert "provider_status=400" in caplog.text
+ assert "model_group=gpt-4.1" in caplog.text
+ assert "AzureException" not in caplog.text
+ assert "No fallback model group" not in caplog.text
+ assert "customer-private-text" not in caplog.text
+ assert "private.example" not in caplog.text
+
+
+def test_traced_does_not_export_natural_language_provider_echo(monkeypatch, caplog):
+ """Unstructured provider prose is never evidence that request text is absent."""
+ import urllib.error
+
+ tracer = MagicMock()
+ span = tracer.start_as_current_span.return_value.__enter__.return_value
+ monkeypatch.setattr(telemetry_module.trace, "get_tracer", lambda unused_name: tracer)
+ echoed = "The supplied phrase customer-private-text is not valid JSON"
+ body = io.BytesIO(json.dumps({"error": {"message": echoed}}).encode())
+
+ with pytest.raises(urllib.error.HTTPError):
+ with traced("capability_probe.chat model-x"):
+ raise urllib.error.HTTPError("https://private.example", 400, "bad", None, body)
+
+ span.set_attribute.assert_any_call(
+ "contextual_orchestrator.error_summary", "invalid_request_error"
+ )
+ assert echoed not in caplog.text
+ assert "customer-private-text" not in caplog.text
+
+
+def test_traced_does_not_export_classified_provider_prose(monkeypatch, caplog):
+ """A previously classified provider error cannot bypass the summary allowlist."""
+ from contextual_orchestrator.provider_errors import ProviderUpstreamError
+
+ tracer = MagicMock()
+ span = tracer.start_as_current_span.return_value.__enter__.return_value
+ monkeypatch.setattr(telemetry_module.trace, "get_tracer", lambda unused_name: tracer)
+ echoed = "The supplied phrase customer-private-text is invalid"
+ error = ProviderUpstreamError(
+ agent_id="provider_agent",
+ model="model-x",
+ error_code="invalid_request_error",
+ message=echoed,
+ client_status=400,
+ provider_status=400,
+ )
+
+ with pytest.raises(ProviderUpstreamError):
+ with traced("chat model-x"):
+ raise error
+
+ span.set_attribute.assert_any_call(
+ "contextual_orchestrator.error_summary", "invalid_request_error"
+ )
+ assert echoed not in caplog.text
+ assert "customer-private-text" not in caplog.text
+
+
def test_annotate_and_usage_helpers_filter_to_allowed_genai_attributes():
"""Span annotation keeps approved scalars only; prompts never enter spans."""
span = MagicMock()
diff --git a/tests/test_token_counting_boundaries.py b/tests/test_token_counting_boundaries.py
index b56db6a65..76cda90e5 100644
--- a/tests/test_token_counting_boundaries.py
+++ b/tests/test_token_counting_boundaries.py
@@ -1,4 +1,4 @@
-"""Boundary tests for the token counting seam (statement + branch coverage)."""
+"""Boundary tests for authoritative token-count selection."""
from __future__ import annotations
@@ -8,42 +8,14 @@
import pytest
from contextual_orchestrator.token_counting import (
- HeuristicTokenCounter,
+ NativeExactTokenCounter,
PgTiktokenAdapter,
+ TokenCountUnavailable,
+ UnavailableTokenCounter,
build_token_counter,
)
-def test_heuristic_empty_and_whitespace_only_text_count_zero() -> None:
- counter = HeuristicTokenCounter()
- assert counter.count_text("") == 0
- # Whitespace-only text matches no word units and must not round up to 1.
- assert counter.count_text(" \t\n ") == 0
-
-
-def test_heuristic_punctuation_only_counts_standalone_symbols() -> None:
- counter = HeuristicTokenCounter()
- # Five standalone punctuation units, expanded by the BPE factor: ceil(5*1.3)=7.
- assert counter.count_text("!?...") == 7
-
-
-def test_heuristic_custom_tokens_per_word_scales_monotonically() -> None:
- text = "alpha beta gamma"
- low = HeuristicTokenCounter(tokens_per_word=1.0).count_text(text)
- high = HeuristicTokenCounter(tokens_per_word=2.5).count_text(text)
- assert low == 3
- assert high == 8
- assert high > low
-
-
-@pytest.mark.parametrize("bad_message", ["plain string", None, 42])
-def test_heuristic_non_dict_messages_contribute_framing_only(bad_message: object) -> None:
- counter = HeuristicTokenCounter()
- total = counter.count_messages([{"content": "hello"}, bad_message]) # type: ignore[list-item]
- # "hello" -> ceil(1*1.3)=2 tokens plus 3 framing for each of two messages.
- assert total == 2 + 3 + 0 + 3
-
-
class _StubPgCounter:
"""Minimal pg_llm_batch.TokenCounter double recording calls."""
@@ -52,56 +24,98 @@ def __init__(self, dsn: str, config: object = None) -> None:
self.config = config
self.calls: list[tuple[str, str]] = []
- def count_tokens(self, text: str, model: str) -> float:
+ def count_tokens(self, text: str, model: str) -> int:
self.calls.append((text, model))
- # Return a float to prove the adapter normalizes to int.
- return 7.9
+ return 7
-def _install_stub_pg_module(monkeypatch: pytest.MonkeyPatch) -> types.ModuleType:
- module = types.ModuleType("pg_llm_batch")
- module.TokenCounter = _StubPgCounter # type: ignore[attr-defined]
- monkeypatch.setitem(sys.modules, "pg_llm_batch", module)
- return module
+def test_postgres_counts_raw_text_but_not_chat_framing() -> None:
+ stub = _StubPgCounter("postgresql://x")
+ adapter = PgTiktokenAdapter(stub)
+ assert adapter.count_text("one", "gpt-4") == 7
+ with pytest.raises(TokenCountUnavailable, match="chat framing"):
+ adapter.count_messages([{"content": "one"}], "gpt-4")
-def test_build_prefers_pg_tiktoken_when_dsn_and_dependency_available(
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- _install_stub_pg_module(monkeypatch)
- config = {"model": "demo_model"}
- counter = build_token_counter("postgresql://ledger_user@localhost/usage_db", config=config)
- assert isinstance(counter, PgTiktokenAdapter)
- # Exact-count delegation returns an int even when the backend yields a float,
- # and forwards both text and model verbatim.
- assert counter.count_text("route me", "mock-generalist") == 7
- stub = getattr(counter, "_counter")
- assert stub.calls == [("route me", "mock-generalist")]
- assert stub.dsn == "postgresql://ledger_user@localhost/usage_db"
- assert stub.config is config
+def test_postgres_runtime_failure_is_unavailable() -> None:
+ class _FailingPgCounter:
+ def count_tokens(self, text: str, model: str) -> int:
+ raise ConnectionError("synthetic database loss")
+ with pytest.raises(TokenCountUnavailable, match="PostgreSQL tokenizer"):
+ PgTiktokenAdapter(_FailingPgCounter()).count_text("one", "gpt-4")
-def test_pg_adapter_counts_messages_and_normalizes_non_dicts() -> None:
- stub = _StubPgCounter("postgresql://x")
- adapter = PgTiktokenAdapter(stub)
- total = adapter.count_messages([{"content": "one"}, {"content": "two"}, "junk"])
- # Non-dict messages are normalized to empty-string content and still counted
- # (the exact backend bills framing), so three calls x 7 tokens each.
- assert total == 21
- assert [call[0] for call in stub.calls] == ["one", "two", ""]
-
-def test_build_degrades_to_heuristic_when_pg_dependency_fails_to_import(
- monkeypatch: pytest.MonkeyPatch,
-) -> None:
- broken = types.ModuleType("pg_llm_batch")
- monkeypatch.setitem(sys.modules, "pg_llm_batch", broken)
- # Attribute access on a bare module raises ImportError -> degrade cleanly.
- counter = build_token_counter("postgresql://ledger_user@localhost/usage_db")
- assert isinstance(counter, HeuristicTokenCounter)
+@pytest.mark.parametrize("invalid_count", [True, -1, 7.5, "7"])
+def test_postgres_rejects_non_integral_or_negative_counts(invalid_count: object) -> None:
+ counter = types.SimpleNamespace(
+ count_tokens=lambda _text, _model: invalid_count,
+ )
+ with pytest.raises(TokenCountUnavailable, match="invalid count"):
+ PgTiktokenAdapter(counter).count_text("one", "gpt-4")
-def test_build_defaults_to_heuristic_without_dsn() -> None:
+def test_build_prefers_configured_postgres(monkeypatch: pytest.MonkeyPatch) -> None:
+ module = types.ModuleType("pg_llm_batch")
+ module.TokenCounter = _StubPgCounter # type: ignore[attr-defined]
+ monkeypatch.setitem(sys.modules, "pg_llm_batch", module)
+ config = {"model": "demo_model"}
+ counter = build_token_counter("postgresql://ledger/usage", config=config)
+ assert isinstance(counter, PgTiktokenAdapter)
+ assert counter.count_text("route me", "gpt-4") == 7
+ assert counter._counter.config is config
+
+
+def test_native_exact_dispatches_only_full_declared_ids(monkeypatch: pytest.MonkeyPatch) -> None:
+ calls: list[tuple[str, str]] = []
+ module = types.SimpleNamespace(
+ count_cl100k=lambda text: calls.append(("cl100k", text)) or 2,
+ count_o200k=lambda text: calls.append(("o200k", text)) or 3,
+ pack_cl100k=lambda *_args: ([], []),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.token_counting.importlib.import_module",
+ lambda _name: module,
+ )
counter = build_token_counter()
- assert isinstance(counter, HeuristicTokenCounter)
- assert counter.count_text("hello world") >= 1
+ assert isinstance(counter, NativeExactTokenCounter)
+ assert counter.count_text("hello world", "gpt-4") == 2
+ assert counter.count_text("hello world", "gpt-4o") == 3
+ with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"):
+ counter.count_text("hello world", "gpt-4-2099-nonexistent")
+ with pytest.raises(TokenCountUnavailable, match="chat framing"):
+ counter.count_messages([{"role": "user", "content": "hello"}], "gpt-4")
+ assert calls == [("cl100k", "hello world"), ("o200k", "hello world")]
+
+
+def test_native_counter_does_not_flatten_multimodal_chat_prompts() -> None:
+ module = types.SimpleNamespace(count_cl100k=lambda _text: 1, count_o200k=lambda _text: 1)
+ counter = NativeExactTokenCounter(module)
+ messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "inspect the image"},
+ {
+ "type": "image_url",
+ "image_url": {"url": "https://example.invalid/synthetic.png"},
+ },
+ ],
+ }
+ ]
+
+ with pytest.raises(TokenCountUnavailable, match="chat framing"):
+ counter.count_messages(messages, "gpt-4o")
+
+
+def test_factory_is_unavailable_when_backends_fail(monkeypatch: pytest.MonkeyPatch) -> None:
+ broken = types.ModuleType("pg_llm_batch")
+ monkeypatch.setitem(sys.modules, "pg_llm_batch", broken)
+ monkeypatch.setattr(
+ "contextual_orchestrator.token_counting.importlib.import_module",
+ lambda _name: (_ for _ in ()).throw(ImportError("missing native")),
+ )
+ counter = build_token_counter("postgresql://ledger/usage")
+ assert isinstance(counter, UnavailableTokenCounter)
+ with pytest.raises(TokenCountUnavailable):
+ counter.count_text("hello", "gpt-4")
diff --git a/tests/test_token_counting_strategies.py b/tests/test_token_counting_strategies.py
index 67aa9e03e..f713da82a 100644
--- a/tests/test_token_counting_strategies.py
+++ b/tests/test_token_counting_strategies.py
@@ -1,105 +1,82 @@
-"""Behavioral coverage for deterministic and PostgreSQL token counters."""
+"""Parity and failure tests for exact native token counting and packing."""
from __future__ import annotations
-import sys
import types
-from typing import Any
+
+import pytest
from contextual_orchestrator.token_counting import (
- HeuristicTokenCounter,
- PgTiktokenAdapter,
- build_token_counter,
+ NativeExactTokenCounter,
+ TokenCountUnavailable,
+ UnavailableTokenCounter,
+ build_embedding_token_counter,
)
-class _PgCounter:
- """Small pg_llm_batch-compatible counter with constructor evidence."""
-
- def __init__(self, postgres_dsn: str, *, config: Any = None) -> None:
- self.postgres_dsn = postgres_dsn
- self.config = config
- self.calls: list[tuple[str, str]] = []
-
- def count_tokens(self, text: str, model: str) -> str:
- """Return a string count so the adapter's integer normalization is exercised."""
- self.calls.append((text, model))
- return str(len(text))
-
-
-def test_heuristic_counter_handles_empty_text_punctuation_and_calibration() -> None:
- """Keep dependency-free estimates deterministic across realistic text shapes."""
- default_counter = HeuristicTokenCounter()
- calibrated_counter = HeuristicTokenCounter(tokens_per_word=0.5)
-
- assert default_counter.count_text("") == 0
- assert default_counter.count_text(" \t\n") == 0
- assert default_counter.count_text("hello, world", model="ignored-model") == 4
- assert calibrated_counter.count_text("hello, world") == 2
-
-
-def test_heuristic_message_count_includes_framing_for_every_input_item() -> None:
- """Count message content plus framing even when a caller supplies a non-mapping item."""
- counter = HeuristicTokenCounter(tokens_per_word=1.0)
-
- assert counter.count_messages(
- [
- {"role": "user", "content": "hello world"},
- {"role": "assistant"},
- "malformed-message",
- ],
- model="ignored-model",
- ) == 11
-
+def test_native_factory_counts_and_packs_declared_models(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ calls: list[tuple[str, str]] = []
+ packed = types.SimpleNamespace(text="hello world", token_count=2)
+ module = types.SimpleNamespace(
+ count_cl100k=lambda text: calls.append(("cl100k", text)) or 2,
+ count_o200k=lambda text: calls.append(("o200k", text)) or 2,
+ pack_cl100k=lambda _texts, _per_input, _inputs, _total: ([packed], [[0]]),
+ )
+ monkeypatch.setattr(
+ "contextual_orchestrator.token_counting.importlib.import_module",
+ lambda _name: module,
+ )
-def test_postgres_adapter_delegates_text_and_message_counts() -> None:
- """Preserve exact database counts without adding heuristic framing overhead."""
- pg_counter = _PgCounter("postgresql://example/tokens")
- adapter = PgTiktokenAdapter(pg_counter)
+ counter = build_embedding_token_counter()
- assert adapter.count_text("four", model="gpt-example") == 4
- assert adapter.count_messages(
- [{"content": "abc"}, {"content": "de"}, "malformed-message"],
- model="gpt-example",
- ) == 5
- assert pg_counter.calls == [
- ("four", "gpt-example"),
- ("abc", "gpt-example"),
- ("de", "gpt-example"),
- ("", "gpt-example"),
+ assert isinstance(counter, NativeExactTokenCounter)
+ assert counter.count_text("hello world", "text-embedding-3-small") == 2
+ assert counter.count_text("hello world", "gpt-4o") == 2
+ assert counter.pack_text("hello world", "text-embedding-3-small", 8192) == [
+ ("hello world", 2)
]
+ with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"):
+ counter.count_text("hello world", "provider-unknown")
+ assert calls == [("cl100k", "hello world"), ("o200k", "hello world")]
-def test_counter_factory_uses_heuristic_without_a_database() -> None:
- """Keep standalone execution dependency-free when no DSN is requested."""
- assert isinstance(build_token_counter(), HeuristicTokenCounter)
-
-
-def test_counter_factory_builds_postgres_adapter_with_explicit_config(monkeypatch) -> None:
- """Pass the caller DSN and tokenizer configuration to pg_llm_batch."""
- module = types.ModuleType("pg_llm_batch")
- module.TokenCounter = _PgCounter
- monkeypatch.setitem(sys.modules, "pg_llm_batch", module)
- config = {"encoding_name": "cl100k_base"}
-
- counter = build_token_counter("postgresql://example/tokens", config=config)
+def test_native_failure_is_unavailable() -> None:
+ def fail(_text: str) -> int:
+ raise RuntimeError("synthetic native failure")
- assert isinstance(counter, PgTiktokenAdapter)
- assert counter._counter.postgres_dsn == "postgresql://example/tokens"
- assert counter._counter.config is config
-
-
-def test_counter_factory_falls_back_when_postgres_counter_cannot_start(monkeypatch) -> None:
- """Retain deterministic counting when the optional database boundary is unavailable."""
- class _UnavailableCounter:
- def __init__(self, _postgres_dsn: str, *, config: Any = None) -> None:
- raise ConnectionError("database unavailable")
+ module = types.SimpleNamespace(
+ count_cl100k=fail,
+ count_o200k=lambda _text: 1,
+ pack_cl100k=lambda *_args: ([], []),
+ )
+ counter = NativeExactTokenCounter(module)
+ with pytest.raises(TokenCountUnavailable, match="native tokenizer"):
+ counter.count_text("hello", "text-embedding-3-large")
+
+
+def test_installed_native_counter_matches_declared_encoding_parity() -> None:
+ module = pytest.importorskip("contextual_orchestrator._token_packer")
+ counter = NativeExactTokenCounter(module)
+ assert counter.count_text("hello world", "gpt-4") == 2
+ assert counter.count_text("hello world", "gpt-4o") == 2
+ assert counter.pack_text("hello world", "text-embedding-3-small", 8192) == [
+ ("hello world", 2)
+ ]
- module = types.ModuleType("pg_llm_batch")
- module.TokenCounter = _UnavailableCounter
- monkeypatch.setitem(sys.modules, "pg_llm_batch", module)
- assert isinstance(
- build_token_counter("postgresql://example/tokens"),
- HeuristicTokenCounter,
+def test_embedding_factory_missing_or_incomplete_native_is_unavailable(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(
+ "contextual_orchestrator.token_counting.importlib.import_module",
+ lambda _name: types.SimpleNamespace(
+ count_cl100k=lambda _text: 1,
+ count_o200k=lambda _text: 1,
+ ),
)
+ counter = build_embedding_token_counter()
+ assert isinstance(counter, UnavailableTokenCounter)
+ with pytest.raises(TokenCountUnavailable):
+ counter.count_text("hello", "text-embedding-3-small")
diff --git a/tests/test_token_id_whole_float_coerce_http_honesty.py b/tests/test_token_id_whole_float_coerce_http_honesty.py
index b85fe930b..86a065e1f 100644
--- a/tests/test_token_id_whole_float_coerce_http_honesty.py
+++ b/tests/test_token_id_whole_float_coerce_http_honesty.py
@@ -16,7 +16,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator
from contextual_orchestrator.server import (
SecurityConfig,
_coerce_embedding_token_sequence,
@@ -55,10 +55,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_user_null_omit_noop_http_honesty.py b/tests/test_user_null_omit_noop_http_honesty.py
index f3b232849..558503e36 100644
--- a/tests/test_user_null_omit_noop_http_honesty.py
+++ b/tests/test_user_null_omit_noop_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
_TEST_AUTH_TOKEN = "user_null_omit_noop_http_honesty_token" # noqa: S105
@@ -42,10 +42,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
diff --git a/tests/test_user_scalar_coerce_http_honesty.py b/tests/test_user_scalar_coerce_http_honesty.py
index 81166082c..afdeafd4e 100644
--- a/tests/test_user_scalar_coerce_http_honesty.py
+++ b/tests/test_user_scalar_coerce_http_honesty.py
@@ -11,7 +11,7 @@
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
-from contextual_orchestrator import ModelAgent, TaskOrchestrator # noqa: E402
+from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator # noqa: E402
from contextual_orchestrator.server import SecurityConfig, build_server # noqa: E402
from contextual_orchestrator.server import _validate_completions_user # noqa: E402
@@ -43,10 +43,13 @@ def _post(port: int, path: str, payload: dict) -> tuple[int, dict]:
def _server():
+ orchestrator = build()
+ counter = type("ExactSyntheticCounter", (), {"count_text": lambda self, text, model="": len(text)})()
server = build_server(
- build(),
+ orchestrator,
port=0,
security=SecurityConfig(auth_token=_TEST_AUTH_TOKEN, rate_limit_requests=10_000),
+ coordinator=CostRoutingCoordinator(orchestrator, embedding_token_counter=counter),
)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()