From 0a383c7be97ea81e493c63597bc55e6c6102f503 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 18:22:25 +0900 Subject: [PATCH 01/63] fix(embeddings): restore provider-backed batches Signed-off-by: Seongho Bae --- contextual_orchestrator/__init__.py | 1 + contextual_orchestrator/batch_job_registry.py | 91 +++++++ contextual_orchestrator/batch_routing.py | 244 +++++++++++++++++- contextual_orchestrator/cost_router.py | 55 +++- .../test_provider_embedding_batch_backend.py | 68 +++++ 5 files changed, 452 insertions(+), 7 deletions(-) create mode 100644 tests/test_provider_embedding_batch_backend.py diff --git a/contextual_orchestrator/__init__.py b/contextual_orchestrator/__init__.py index 06e135ea9..557bd5726 100644 --- a/contextual_orchestrator/__init__.py +++ b/contextual_orchestrator/__init__.py @@ -10,6 +10,7 @@ LocalEmbeddingBatchBackend, PgLlmBatchBackend, PgLlmBatchEmbeddingBackend, + ProviderEmbeddingBatchBackend, RoutingDecision, RoutingHints, RoutingPolicy, diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py index a7bae84fb..dcc5d6edd 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,10 @@ DEFAULT_RETENTION_SECONDS = 7 * 24 * 3600 +class ClaimNotAcquired(RuntimeError): + """Another worker owns a non-blocking durable job claim.""" + + 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 +138,95 @@ 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() + renewal_thread = None + if renew_until_epoch is not None: + + def renew_claim() -> None: + interval = max(0.05, min(lease_seconds / 3, 1.0)) + while not stop_renewal.wait(interval): + remaining = renew_until_epoch - time.time() + if remaining <= 0: + 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: + return + except Exception as exc: # noqa: BLE001 - redis is optional. + if type(exc).__name__ in {"LockError", "LockNotOwnedError"}: + return + + renewal_thread = threading.Thread( + target=renew_claim, + name="job-claim-renewal", + daemon=True, + ) + renewal_thread.start() + try: + yield claim + 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 + return lock @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 82ce36600..23fa731a1 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -24,6 +24,7 @@ import dataclasses import hashlib import json +import threading import time import uuid from contextlib import nullcontext @@ -32,6 +33,8 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Protocol +from .batch_job_registry import ClaimNotAcquired, JobRegistryFactory + _ROUTING_CATEGORY = "routing" _PROVIDER_CUSTOM_ID_MAX_LENGTH = 64 @@ -454,10 +457,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 @@ -572,7 +580,6 @@ 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()) def submit( @@ -604,6 +611,241 @@ 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, + ) -> 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._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 + 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.""" + 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_id = f"providerembed_{uuid.uuid4().hex}" + self._requests[job_id] = list(requests) + self._deadlines[job_id] = time.time() + self._registry.retention_seconds + self._states[job_id] = "queued" + self._terminal_events[job_id] = threading.Event() + with self._executor_lock: + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=self._max_concurrency) + executor = self._executor + executor.submit(copy_context().run, self._run_job, job_id) + return BatchJob(job_id=job_id, backend=self.name, status="queued", request_count=len(requests)) + + def _run_job(self, job_id: str) -> None: + """Execute one persisted job inside the bounded provider worker pool.""" + try: + deadline_epoch = float( + self._deadlines.get( + job_id, time.time() + self._registry.retention_seconds + ) + ) + self._deadlines[job_id] = deadline_epoch + with self._registry.lock( + "provider_embedding_job_execution", job_id, + lease_seconds=self._claim_lease_seconds, + renew_until_epoch=deadline_epoch, + ): + 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]) + vectors, prompt_tokens = self._runner(requests) + if self._states.get(job_id) == "cancelled": + 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 = [] + for index, (request, vector) in enumerate(zip(requests, vectors, strict=True)): + items.append( + EmbeddingBatchResultItem( + custom_id=request.custom_id, + index=index, + embedding=vector, + prompt_tokens=0, + model=request.model, + ) + ) + with self._registry.lock( + "provider_embedding_job_states", job_id, + lease_seconds=self._claim_lease_seconds, + ): + if self._states.get(job_id) == "cancelled": + return + self._usage[job_id] = {"prompt_tokens": int(prompt_tokens)} + self._results[job_id] = items + self._states[job_id] = "completed" + except ClaimNotAcquired: + return + except Exception as exc: # noqa: BLE001 - polling exposes a bounded terminal state + with self._registry.lock( + "provider_embedding_job_states", job_id, + lease_seconds=self._claim_lease_seconds, + ): + if self._states.get(job_id) not in {"cancelled", "completed"}: + self._errors[job_id] = { + "error_type": type(exc).__name__, + "http_status": getattr(exc, "status_code", None), + "provider_code": getattr(exc, "provider_code", None), + "retryable": bool(getattr(exc, "retryable", False)), + "failed_shard_index": getattr(exc, "failed_shard_index", None), + } + self._states[job_id] = "failed" + finally: + event = self._terminal_events.pop(job_id, None) + if event is not None: + event.set() + + 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, {})) + return document + + def cancel(self, job: BatchJob, *, reason: str) -> Dict[str, Any]: + """Mark queued/running work cancelled and discard any late provider result.""" + 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" + 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. diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 438a26fde..1b896fa89 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -34,6 +34,7 @@ EmbeddingBatchResultItem, LocalBatchBackend, LocalEmbeddingBatchBackend, + ProviderEmbeddingBatchBackend, RoutingHints, RoutingPolicy, ) @@ -109,12 +110,54 @@ 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: + def run_provider_embeddings( + requests: List[EmbeddingBatchRequest], + ) -> tuple[List[List[float]], int]: + first = requests[0] + agent = ( + orchestrator._agent(first.agent_id) + if first.agent_id is not None + else 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") + vectors = orchestrator.client.embed( + agent, [request.input_text for request in requests] + ) + prompt_tokens = sum( + int(self.token_counter.count_text(request.input_text, request.model)) + for request in requests + ) + return vectors, prompt_tokens + + self.embedding_batch_backend = ProviderEmbeddingBatchBackend( + run_provider_embeddings, + job_registry=registry, + max_concurrency=getattr(orchestrator.client, "local_concurrency", 1), + claim_lease_seconds=( + float(orchestrator.client.timeout) + if registry.durable and float(getattr(orchestrator.client, "timeout", 0)) > 0 + else None + ), + ) + else: + self.embedding_batch_backend = LocalEmbeddingBatchBackend( + token_counter=self.token_counter, 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)) # embeddings batch state: job handle + submitted requests + cached doc, diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py new file mode 100644 index 000000000..3294ae2ab --- /dev/null +++ b/tests/test_provider_embedding_batch_backend.py @@ -0,0 +1,68 @@ +"""Focused synthetic contracts for provider-backed embedding batches.""" + +import threading +import time + +from contextual_orchestrator.batch_routing import ( + EmbeddingBatchRequest, + ProviderEmbeddingBatchBackend, +) +from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator +from contextual_orchestrator.orchestrator import ModelClient + + +class _SyntheticProviderClient(ModelClient): + def embed(self, agent, texts): + return [[float(len(text))] for text in texts] + + +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_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_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()) + ) + 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]] From f252a55f8ce74a1772da36973d8183bf68eefb92 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 18:33:49 +0900 Subject: [PATCH 02/63] build(rust): restore native embedding packer Signed-off-by: Seongho Bae --- rust/.gitignore | 1 + rust/Cargo.lock | 373 ++++++++++++++++++++++++++ rust/Cargo.toml | 3 + rust/token_counter/.gitignore | 1 + rust/token_counter/Cargo.toml | 15 ++ rust/token_counter/pyproject.toml | 13 + rust/token_counter/src/lib.rs | 421 ++++++++++++++++++++++++++++++ 7 files changed, 827 insertions(+) create mode 100644 rust/.gitignore create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/token_counter/.gitignore create mode 100644 rust/token_counter/Cargo.toml create mode 100644 rust/token_counter/pyproject.toml create mode 100644 rust/token_counter/src/lib.rs 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..04a17a4c5 --- /dev/null +++ b/rust/token_counter/src/lib.rs @@ -0,0 +1,421 @@ +//! Exact cl100k child chunking and provider-request packing for Python. +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use rayon::prelude::*; +use std::sync::OnceLock; +use tiktoken_rs::cl100k_base; + +static CL100K: 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")) +} + +#[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 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!(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!( + 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()); + }); + } +} From 18c108be6b884b051198f33c903b1ee6eeed405f Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 18:44:44 +0900 Subject: [PATCH 03/63] build(docker): compile locked native runtime dependencies Signed-off-by: Seongho Bae --- Dockerfile | 19 ++++++++++++++++++- tests/test_compose_contract.py | 4 +++- 2 files changed, 21 insertions(+), 2 deletions(-) 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/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 From 315f1d7689708db0426c1d5eb3a452ed08851b7c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 19:16:21 +0900 Subject: [PATCH 04/63] fix(runtime): preserve Lineage session and embedding routes Signed-off-by: Seongho Bae --- contextual_orchestrator/__main__.py | 13 ++++++-- contextual_orchestrator/server.py | 3 ++ tests/test_auto_discovery_server.py | 47 +++++++++++++++++------------ tests/test_openai_passthrough.py | 21 +++++++++++++ 4 files changed, 62 insertions(+), 22 deletions(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index ffd82536d..23c0e0f2c 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -418,11 +418,18 @@ 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) ] + runtime_models = [ + model + for model in discovered + if model in chat_models or "embedding" in model.capabilities + ] existing_by_id = {agent.id: agent for agent in orchestrator.candidates} 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) + routable = is_routable_discovered_model(model) or ( + "embedding" in model.capabilities and model.spend_admitted + ) if existing is None: agents.append(replace(agent_from_discovered(model), disabled=not routable)) elif "discovered" not in existing.tags: @@ -455,7 +462,7 @@ 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): + if any(model.provider_name == "configured_gateway" for model in runtime_models): for agent in tuple(orchestrator.candidates): if ( agent.provider_name == "configured_gateway" diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index 4ea2f1d8b..a12be9a1b 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -183,6 +183,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 @@ -6325,6 +6326,8 @@ def do_POST(self) -> None: # noqa: N802 for key in ("metadata", "client_metadata") if isinstance((value := body.get(key)), dict) ] + if "session_id" in body: + metadata_values.append({"session_id": body["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) diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 86102624d..a0950d9ca 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -9,21 +9,21 @@ 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,12 +38,14 @@ 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"] @@ -307,8 +309,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 +326,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 +376,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_openai_passthrough.py b/tests/test_openai_passthrough.py index e5f06ccdb..ef6fe4c6c 100644 --- a/tests/test_openai_passthrough.py +++ b/tests/test_openai_passthrough.py @@ -402,6 +402,27 @@ 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.""" + server, port, token = _serve() + 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"] + + 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() From 36783268d1fe4254fb4484ed52b0f5937ec5736b Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 19:42:54 +0900 Subject: [PATCH 05/63] fix(runtime): classify structured synthesis failures Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 7 ++++ .../test_chat_response_format_http_honesty.py | 38 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 1f99f5fac..17b78d7c5 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4126,6 +4126,13 @@ def send_synthesis( "request body exceeds provider limit" ) from exc if not request_too_large: + if isinstance(exc, (urllib.error.HTTPError, ProviderUpstreamError)): + raise classify_provider_failure( + exc, + agent_id=candidate.id, + model=candidate.model, + transport="structured_synthesis", + ) from None raise raise ProviderRequestTooLargeError( "request body exceeds every eligible provider limit" diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index 154f9579b..a4de6ea3e 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -83,6 +83,44 @@ 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_http_structured_chat_rejects_batch_routing() -> None: """Provider-native structured synthesis has no batch execution contract.""" server, thread, port = _server() From 5ba91397ef1887517b355ea241382572e038d90c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 19:53:42 +0900 Subject: [PATCH 06/63] fix(runtime): enforce shared routing boundaries Signed-off-by: Seongho Bae --- contextual_orchestrator/__main__.py | 3 ++- contextual_orchestrator/cost_router.py | 2 ++ contextual_orchestrator/orchestrator.py | 1 + contextual_orchestrator/provider_errors.py | 13 +++++++++++- tests/test_auto_discovery_server.py | 24 ++++++++++++++++++++++ tests/test_cost_router.py | 23 +++++++++++++++++++++ tests/test_openai_passthrough.py | 1 + tests/test_provider_error_taxonomy.py | 16 +++++++++++++++ 8 files changed, 81 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 23c0e0f2c..7c81cc1a4 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -421,7 +421,8 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l runtime_models = [ model for model in discovered - if model in chat_models or "embedding" in model.capabilities + if not model.evidence_only + and (model in chat_models or "embedding" in model.capabilities) ] existing_by_id = {agent.id: agent for agent in orchestrator.candidates} agents = [] diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 1b896fa89..7265b76b6 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -124,6 +124,8 @@ def __init__( def run_provider_embeddings( requests: List[EmbeddingBatchRequest], ) -> tuple[List[List[float]], int]: + if not requests: + return [], 0 first = requests[0] agent = ( orchestrator._agent(first.agent_id) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 17b78d7c5..c4c0259f6 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -3559,6 +3559,7 @@ def _reload_state(self) -> None: "stream_options", "_required_agent_id", "_file_replicas", + "session_id", } ) diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py index ecc763d3e..ea47cee99 100644 --- a/contextual_orchestrator/provider_errors.py +++ b/contextual_orchestrator/provider_errors.py @@ -209,7 +209,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/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index a0950d9ca..21e725c36 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -195,6 +195,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( diff --git a/tests/test_cost_router.py b/tests/test_cost_router.py index 1b88f26c6..60bfbaca2 100644 --- a/tests/test_cost_router.py +++ b/tests/test_cost_router.py @@ -3,6 +3,7 @@ from __future__ import annotations import sys +import time from pathlib import Path import pytest @@ -784,6 +785,28 @@ 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_non_zdr_batch_preserves_an_explicit_model_outside_the_pool() -> None: """The ZDR resolver must not change ordinary batch passthrough behavior.""" captured: list[BatchRequest] = [] diff --git a/tests/test_openai_passthrough.py b/tests/test_openai_passthrough.py index ef6fe4c6c..821b41a79 100644 --- a/tests/test_openai_passthrough.py +++ b/tests/test_openai_passthrough.py @@ -421,6 +421,7 @@ def test_lineage_structured_payload_accepts_session_without_provider_forwarding( assert status == 200 assert body["echo"]["response_format"] == {"type": "json_object"} assert "session_id" not in body["echo"] + assert "session_id" in TaskOrchestrator._ORCHESTRATION_ONLY_KEYS def test_http_gateway_default_response_format_resolves_concrete_agent() -> None: diff --git a/tests/test_provider_error_taxonomy.py b/tests/test_provider_error_taxonomy.py index 9aba81f39..af4a32ad8 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"}})) From 6c77d9b0ec9e72939eb844b6dc97fc88db436694 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 20:15:52 +0900 Subject: [PATCH 07/63] fix(embeddings): fence leases and exact token counts Signed-off-by: Seongho Bae --- CHANGELOG.md | 5 + contextual_orchestrator/__init__.py | 11 ++- contextual_orchestrator/batch_job_registry.py | 64 ++++++++++++- contextual_orchestrator/batch_routing.py | 14 ++- contextual_orchestrator/cost_router.py | 47 +++++---- contextual_orchestrator/token_counting.py | 96 +++++++++++++++++-- ...er-embedding-lease-and-token-accounting.md | 95 ++++++++++++++++++ docs/adr/README.md | 1 + docs/library_research.md | 2 + tests/test_batch_embeddings.py | 14 ++- tests/test_batch_job_registry.py | 79 +++++++++++++++ tests/test_batch_routing_boundaries_extra.py | 12 +-- tests/test_cost_router.py | 10 +- tests/test_cost_router_boundaries.py | 34 +++++-- .../test_provider_embedding_batch_backend.py | 9 +- tests/test_token_counting_strategies.py | 85 ++++++++++++++++ 16 files changed, 526 insertions(+), 52 deletions(-) create mode 100644 docs/adr/0005-provider-embedding-lease-and-token-accounting.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 6b38e0770..8ddb03138 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,11 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ### Fixed +- 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). Legacy chat estimation remains a documented gap. - OpenRouter discovery no longer marks the entire credential account evidence-only. Authenticated catalog rows may serve ordinary requests, while ZDR-only requests still require explicit route-level ZDR evidence. diff --git a/contextual_orchestrator/__init__.py b/contextual_orchestrator/__init__.py index 557bd5726..2f23e0ccf 100644 --- a/contextual_orchestrator/__init__.py +++ b/contextual_orchestrator/__init__.py @@ -61,7 +61,13 @@ parse_reasoning_effort_profile, snapshot_role_effort_catalog, ) -from .token_counting import HeuristicTokenCounter, build_token_counter +from .token_counting import ( + HeuristicTokenCounter, + NativeCl100kTokenCounter, + TokenCountUnavailable, + build_embedding_token_counter, + build_token_counter, +) from .response_cache import ( RedisResponseCacheProvider, ResponseCacheProvider, @@ -120,6 +126,9 @@ "InMemoryConfigStore", "get_config_store", "HeuristicTokenCounter", + "NativeCl100kTokenCounter", + "TokenCountUnavailable", + "build_embedding_token_counter", "build_token_counter", "ResponseCacheProvider", "RedisResponseCacheProvider", diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py index dcc5d6edd..b52f1beb8 100644 --- a/contextual_orchestrator/batch_job_registry.py +++ b/contextual_orchestrator/batch_job_registry.py @@ -52,6 +52,46 @@ 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 _encode(value: Any) -> str: """Serialize one registry value (dataclasses included) to JSON.""" if dataclasses.is_dataclass(value) and not isinstance(value, type): @@ -169,6 +209,12 @@ 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: @@ -177,6 +223,7 @@ def renew_claim() -> None: 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 @@ -186,10 +233,11 @@ def renew_claim() -> None: replace_ttl=True, ) if not renewed: + lease.mark_lost() return - except Exception as exc: # noqa: BLE001 - redis is optional. - if type(exc).__name__ in {"LockError", "LockNotOwnedError"}: - return + except Exception: # noqa: BLE001 - redis is optional. + lease.mark_lost() + return renewal_thread = threading.Thread( target=renew_claim, @@ -198,7 +246,7 @@ def renew_claim() -> None: ) renewal_thread.start() try: - yield claim + yield lease finally: stop_renewal.set() if renewal_thread is not None: @@ -215,7 +263,13 @@ def renew_claim() -> None: if lock is None: lock = threading.Lock() self._local_locks[lock_name] = lock - return lock + + @contextmanager + def acquired_local_claim(): + with lock: + yield _ClaimLease() + + return acquired_local_claim() @property def durable(self) -> bool: diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 23fa731a1..8d16fc720 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -578,9 +578,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)) - 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 @@ -736,7 +736,8 @@ def _run_job(self, job_id: str) -> None: "provider_embedding_job_execution", job_id, lease_seconds=self._claim_lease_seconds, renew_until_epoch=deadline_epoch, - ): + ) as execution_claim: + execution_claim.ensure_owned() with self._registry.lock( "provider_embedding_job_states", job_id, lease_seconds=self._claim_lease_seconds, @@ -746,6 +747,7 @@ def _run_job(self, job_id: str) -> None: self._states[job_id] = "running" requests = list(self._requests[job_id]) vectors, prompt_tokens = self._runner(requests) + execution_claim.ensure_owned() if self._states.get(job_id) == "cancelled": return if len(vectors) != len(requests): @@ -768,6 +770,10 @@ def _run_job(self, job_id: str) -> None: "provider_embedding_job_states", job_id, lease_seconds=self._claim_lease_seconds, ): + # Refresh while holding the terminal-state lock. A stale + # worker must never publish after another worker can claim + # the same provider job. + execution_claim.ensure_owned(refresh=True) if self._states.get(job_id) == "cancelled": return self._usage[job_id] = {"prompt_tokens": int(prompt_tokens)} diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 7265b76b6..62bfe6757 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -41,7 +41,11 @@ from .batch_job_registry import JobRegistryFactory, build_job_registry from .cost_ledger import CostLedger, PriceBook from .kv_config import InMemoryConfigStore -from .token_counting import HeuristicTokenCounter, build_token_counter +from .token_counting import ( + HeuristicTokenCounter, + build_embedding_token_counter, + build_token_counter, +) _RACE_USAGE_CONTEXT: ContextVar[dict[str, Any] | None] = ContextVar( @@ -73,6 +77,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, @@ -89,6 +94,12 @@ def __init__( self.token_counter = token_counter or ( build_token_counter(postgres_dsn) if postgres_dsn else HeuristicTokenCounter() ) + 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) # Job registries live in Valkey when the credential registry carries # batch_job_registry_valkey_url, so submitted jobs survive a process @@ -137,13 +148,17 @@ def run_provider_embeddings( for request in requests ): raise RuntimeError("provider embedding batch must retain one selected route") - vectors = orchestrator.client.embed( - agent, [request.input_text for request in requests] - ) prompt_tokens = sum( - int(self.token_counter.count_text(request.input_text, request.model)) + int( + self.embedding_token_counter.count_text( + request.input_text, request.model + ) + ) for request in requests ) + vectors = orchestrator.client.embed( + agent, [request.input_text for request in requests] + ) return vectors, prompt_tokens self.embedding_batch_backend = ProviderEmbeddingBatchBackend( @@ -158,7 +173,7 @@ def run_provider_embeddings( ) else: self.embedding_batch_backend = LocalEmbeddingBatchBackend( - token_counter=self.token_counter, job_registry=registry + token_counter=self.embedding_token_counter, 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)) @@ -907,8 +922,7 @@ def _embedding_request_limits(self) -> tuple[int, int]: 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( @@ -1023,13 +1037,10 @@ 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]: @@ -1070,9 +1081,11 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: 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: + prompt_tokens = int( + self.embedding_token_counter.count_text(request.input_text, item.model) + ) parts_by_source.setdefault(source_index, []).append( { "part_index": request.part_index if request else 0, diff --git a/contextual_orchestrator/token_counting.py b/contextual_orchestrator/token_counting.py index 0779cb1d1..a7cca010b 100644 --- a/contextual_orchestrator/token_counting.py +++ b/contextual_orchestrator/token_counting.py @@ -1,7 +1,7 @@ """Token counting seam for usage/cost accounting. -The cost ledger needs prompt/completion token counts on every completion. Two -strategies are provided behind one :class:`TokenCounter`-compatible surface: +The cost ledger needs prompt/completion token counts on every completion. +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 @@ -9,13 +9,19 @@ * :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. +* :class:`NativeCl100kTokenCounter` β€” delegates declared cl100k embedding + models to the packaged Rust extension and reports all other cases as + unavailable. + +Legacy chat selection remains in :func:`build_token_counter`. Embedding +selection is isolated in :func:`build_embedding_token_counter` and never +returns a heuristic. Neither factory reads the environment: the DSN is passed +in by the caller. """ from __future__ import annotations +import importlib import math import re from typing import Any, List, Optional, Protocol @@ -25,6 +31,17 @@ # Rough BPE expansion: sub-word models emit slightly more tokens than words. _TOKENS_PER_WORD = 1.3 +# OpenAI's published tiktoken mapping assigns these embedding deployments to +# cl100k_base. Other model identifiers remain unavailable because a tokenizer +# must never be guessed from a provider/model name. +_CL100K_EMBEDDING_MODELS = frozenset( + { + "text-embedding-ada-002", + "text-embedding-3-small", + "text-embedding-3-large", + } +) + class TokenCountingStrategy(Protocol): """Contract for anything that can count tokens for a chunk of text.""" @@ -34,13 +51,18 @@ def count_text(self, text: str, model: str) -> int: ... +class TokenCountUnavailable(RuntimeError): + """An authoritative tokenizer is unavailable for the requested model.""" + + class HeuristicTokenCounter: """Deterministic, dependency-free token estimator. 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. + enough for legacy best-effort chat attribution when ``pg_tiktoken`` is not + reachable, and it never varies between runs so tests can assert on it. It + is not authoritative and must not be used for embedding limits or cost. """ def __init__(self, tokens_per_word: float = _TOKENS_PER_WORD) -> None: @@ -86,6 +108,66 @@ def count_messages(self, messages: List[dict], model: str = "") -> int: return total +class NativeCl100kTokenCounter: + """Use the bundled Rust cl100k counter only for explicitly mapped models.""" + + def __init__(self, native_module: Any) -> None: + self._native_module = native_module + + def count_text(self, text: str, model: str = "") -> int: + """Count a declared cl100k embedding model or fail closed.""" + if model not in _CL100K_EMBEDDING_MODELS: + raise TokenCountUnavailable(f"no authoritative tokenizer is declared for {model!r}") + try: + return int(self._native_module.count_cl100k(text)) + except Exception as exc: # noqa: BLE001 - optional native boundary. + raise TokenCountUnavailable("the native cl100k tokenizer is unavailable") from exc + + def count_messages(self, messages: List[dict], model: str = "") -> int: + """Reject chat counting because this counter is embedding-only.""" + raise TokenCountUnavailable("the native cl100k counter does not count chat framing") + + +class UnavailableEmbeddingTokenCounter: + """Represent the absence of an authoritative embedding tokenizer.""" + + def count_text(self, text: str, model: str = "") -> int: + """Fail closed instead of fabricating an embedding token count.""" + raise TokenCountUnavailable(f"no authoritative tokenizer is available for {model!r}") + + +def _native_token_counter() -> NativeCl100kTokenCounter | None: + """Load the optional in-package extension without making startup depend on it.""" + try: + module = importlib.import_module("contextual_orchestrator._token_packer") + except Exception: # noqa: BLE001 - an incompatible wheel is equivalent to absence. + return None + if not callable(getattr(module, "count_cl100k", None)): + return None + return NativeCl100kTokenCounter(module) + + +def build_embedding_token_counter( + postgres_dsn: Optional[str] = None, + *, + config: Any = None, +) -> PgTiktokenAdapter | NativeCl100kTokenCounter | UnavailableEmbeddingTokenCounter: + """Return an authoritative embedding counter or an explicit unavailable seam. + + PostgreSQL remains authoritative when explicitly configured. The bundled + Rust counter is the fallback only for model identifiers whose published + tokenizer mapping is cl100k. No heuristic estimate is returned here. + """ + 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 - optional authoritative boundary + pass + return _native_token_counter() or UnavailableEmbeddingTokenCounter() + + def build_token_counter( postgres_dsn: Optional[str] = None, *, 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..11ec8c9fd --- /dev/null +++ b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md @@ -0,0 +1,95 @@ +# 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. + +## 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.** A provider embedding worker checks + ownership before provider work, after provider work, and again by + refreshing its execution claim while holding the terminal-state lock. + Only then may it publish usage, results, and `completed`. A stale worker + leaves the job in recoverable `running` state; it does not overwrite a + successor's state or manufacture `failed`. +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 remains a known gap.** This ADR makes no global + token-accounting compliance claim. Existing chat routing and missing-usage + cost paths still use `HeuristicTokenCounter`; replacing that estimate with + authoritative provider/tokenizer evidence is required follow-up work. + +## 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 estimates remain outside this narrow compliance boundary. + +## References + +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/README.md b/docs/adr/README.md index d4347e352..822a3a1d6 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -12,6 +12,7 @@ They do not share numbering with `docs/planning/adrs/`. | [0002](0002-control-plane-orchestrator.md) | Control-plane orchestrator, not a trained coordinator | Accepted | Xu et al. (2025) TRINITY arXiv:2512.04695; Nielsen et al. (2025) Conductor arXiv:2512.04388; Sakana Fugu (2026) live pages | | [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-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 | 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 998004420..7640fa7b7 100644 --- a/docs/library_research.md +++ b/docs/library_research.md @@ -19,6 +19,8 @@ primitives use maintained libraries when the enterprise target requires them. | Rendered policy browser | [MCP Python SDK](https://github.com/modelcontextprotocol/python-sdk) | Keep as a pinned deployment-provided optional package for the existing Camoufox MCP transport; do not claim a repository `policy-browser` extra until this project owns that lock and publish contract. | Reuses the protocol client and Streamable HTTP lifecycle instead of implementing a second transport; static policy analysis does not install it. | | Structured-output validation | `jsonschema` | Use the maintained validator for provider-returned JSON against caller-supplied JSON Schema; keep parsing and the single repair policy in the existing orchestrator. | Reusing `validator_for`, schema checks, and bounded validation avoids an incomplete custom JSON Schema implementation. Provider output and schemas remain untrusted and fail closed. | | 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. | +| Provider-embedding claim ownership | Existing `redis-py` `Lock` acquire/extend/owned/release scripts | Propagate renewal loss to the claim holder and refresh ownership inside the terminal-state critical section before publishing usage, results, or completion. This fences stale local publication but does not claim provider-side exactly-once execution. | Redis's official distributed-lock guidance requires ownership-safe release/extension and recommends fencing when correctness depends on exclusive work. Skipped a custom lock protocol, 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. Legacy chat missing-usage accounting still uses `HeuristicTokenCounter` and remains a known noncompliant follow-up gap; this narrow embedding change is not global token-accounting compliance. | ## Ponytail Decision diff --git a/tests/test_batch_embeddings.py b/tests/test_batch_embeddings.py index a67881f2c..d8a8a985b 100644 --- a/tests/test_batch_embeddings.py +++ b/tests/test_batch_embeddings.py @@ -75,7 +75,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=HeuristicTokenCounter(), + ) token = "cost_token" server = build_server( orchestrator, port=0, security=SecurityConfig(auth_token=token), coordinator=coordinator @@ -271,7 +276,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=HeuristicTokenCounter(), + ) token = "zdr_batch_token" server = build_server( orchestrator, port=0, security=SecurityConfig(auth_token=token), coordinator=coordinator @@ -296,6 +305,7 @@ def test_pending_batch_preserves_resolved_model_identity() -> None: coordinator = CostRoutingCoordinator( orchestrator, InMemoryConfigStore(), + embedding_token_counter=HeuristicTokenCounter(), embedding_batch_backend=_PendingEmbeddingBackend(), ) diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 6ef029ec3..73f386474 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -11,12 +11,17 @@ from __future__ import annotations import sys +import threading +import time 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 +31,9 @@ BatchJob, BatchRequest, BatchResultItem, + EmbeddingBatchRequest, LocalBatchBackend, + ProviderEmbeddingBatchBackend, ) from contextual_orchestrator.kv_config import InMemoryConfigStore @@ -37,6 +44,36 @@ class FakeValkeyClient: def __init__(self) -> None: self.hashes: Dict[str, Dict[str, str]] = {} self.expirations: Dict[str, int] = {} + self.execution_extension_attempted = threading.Event() + + class LockNotOwnedError(RuntimeError): + pass + + class _Lock: + def __init__(self, client: "FakeValkeyClient", name: str) -> None: + self._client = client + self._lose_on_extend = "provider_embedding_job_execution" in name + self._owned = False + + def acquire(self) -> bool: + self._owned = True + 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.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 def hget(self, key: str, field: str) -> Any: return self.hashes.get(key, {}).get(field) @@ -69,6 +106,9 @@ 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 test_mapping_round_trips_dataclasses_and_plain_values() -> None: """Dataclasses, dataclass lists, and JSON scalars all survive the trip.""" @@ -158,6 +198,45 @@ 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_result_is_not_published_after_claim_renewal_loss() -> None: + """A stale provider worker leaves recovery state for the succeeding claimant.""" + client = FakeValkeyClient() + registry = JobRegistryFactory(client, retention_seconds=2) + + def runner(_requests): + assert client.execution_extension_attempted.wait(timeout=1) + return [[1.0]], 1 + + 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"] == "running" + assert backend.retrieve(job) == [] + assert backend.usage(job) == {} + backend.close() + + if __name__ == "__main__": for name, value in sorted(globals().items()): if name.startswith("test_") and callable(value): 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_cost_router.py b/tests/test_cost_router.py index 60bfbaca2..b2d280ef4 100644 --- a/tests/test_cost_router.py +++ b/tests/test_cost_router.py @@ -774,8 +774,16 @@ 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=_SyntheticExactCounter(), + ) document = coordinator.complete_embeddings_batch( ["private"], model=second.model, zdr_only=True, agent_id=second.id ) diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index d09d75ab5..e3a08f2fb 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -26,6 +26,11 @@ _provider_from_base_url, _weighted_average_embedding, ) +from contextual_orchestrator.token_counting import ( + HeuristicTokenCounter, + TokenCountUnavailable, + UnavailableEmbeddingTokenCounter, +) def _coordinator(**kwargs: Any) -> Coordinator: @@ -41,6 +46,8 @@ def _coordinator(**kwargs: Any) -> Coordinator: orchestrator = TaskOrchestrator(agents) config = InMemoryConfigStore() price_book = PriceBook(config) + if "token_counter" not in kwargs and "embedding_token_counter" not in kwargs: + kwargs["embedding_token_counter"] = HeuristicTokenCounter() return Coordinator(orchestrator, config, price_book=price_book, **kwargs) @@ -177,11 +184,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: @@ -192,9 +198,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: @@ -305,6 +312,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) diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 3294ae2ab..154a9ba51 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -16,6 +16,12 @@ def embed(self, agent, texts): return [[float(len(text))] for text in texts] +class _SyntheticExactCounter: + def count_text(self, text, model): + """Return a deterministic synthetic authoritative count.""" + return len(text.split()) + + def test_provider_batch_returns_before_terminal_result() -> None: release = threading.Event() @@ -54,7 +60,8 @@ def test_remote_embedding_member_selects_provider_backend() -> None: tags=("embedding",), ) coordinator = CostRoutingCoordinator( - TaskOrchestrator([agent], client=_SyntheticProviderClient()) + TaskOrchestrator([agent], client=_SyntheticProviderClient()), + embedding_token_counter=_SyntheticExactCounter(), ) job = coordinator.submit_embeddings_batch( ["synthetic one", "synthetic two"], model=agent.model, agent_id=agent.id diff --git a/tests/test_token_counting_strategies.py b/tests/test_token_counting_strategies.py index 67aa9e03e..9f9137704 100644 --- a/tests/test_token_counting_strategies.py +++ b/tests/test_token_counting_strategies.py @@ -6,9 +6,15 @@ import types from typing import Any +import pytest + from contextual_orchestrator.token_counting import ( HeuristicTokenCounter, + NativeCl100kTokenCounter, PgTiktokenAdapter, + TokenCountUnavailable, + UnavailableEmbeddingTokenCounter, + build_embedding_token_counter, build_token_counter, ) @@ -75,6 +81,69 @@ def test_counter_factory_uses_heuristic_without_a_database() -> None: assert isinstance(build_token_counter(), HeuristicTokenCounter) +def test_counter_factory_uses_native_cl100k_only_for_declared_embedding_models( + monkeypatch, +) -> None: + """The installed wheel is a real runtime path without guessing tokenizers.""" + calls: list[str] = [] + module = types.SimpleNamespace( + count_cl100k=lambda text: calls.append(text) or 2, + ) + monkeypatch.setattr( + "contextual_orchestrator.token_counting.importlib.import_module", + lambda _name: module, + ) + + counter = build_embedding_token_counter() + + assert isinstance(counter, NativeCl100kTokenCounter) + assert counter.count_text("hello world", "text-embedding-3-small") == 2 + assert calls == ["hello world"] + with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"): + counter.count_text("hello world", "provider-unknown") + assert calls == ["hello world"] + + +def test_native_counter_reports_unavailable_when_extension_call_fails(monkeypatch) -> None: + """One native failure must not fabricate embedding usage.""" + + def fail(_text: str) -> int: + raise RuntimeError("synthetic native failure") + + monkeypatch.setattr( + "contextual_orchestrator.token_counting.importlib.import_module", + lambda _name: types.SimpleNamespace(count_cl100k=fail), + ) + + counter = build_embedding_token_counter() + + assert isinstance(counter, NativeCl100kTokenCounter) + with pytest.raises(TokenCountUnavailable, match="native cl100k"): + counter.count_text("hello world", "text-embedding-3-large") + + +def test_installed_native_counter_matches_cl100k_reference_count() -> None: + """An installed wheel preserves the Rust cl100k parity boundary.""" + module = pytest.importorskip("contextual_orchestrator._token_packer") + counter = NativeCl100kTokenCounter(module) + + assert counter.count_text("hello world", "text-embedding-3-small") == 2 + + +def test_embedding_counter_without_authoritative_backend_is_unavailable(monkeypatch) -> None: + """Missing optional native code is an explicit unavailable result.""" + monkeypatch.setattr( + "contextual_orchestrator.token_counting.importlib.import_module", + lambda _name: (_ for _ in ()).throw(ImportError("synthetic missing wheel")), + ) + + counter = build_embedding_token_counter() + + assert isinstance(counter, UnavailableEmbeddingTokenCounter) + with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"): + counter.count_text("hello world", "text-embedding-3-small") + + 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") @@ -103,3 +172,19 @@ def __init__(self, _postgres_dsn: str, *, config: Any = None) -> None: build_token_counter("postgresql://example/tokens"), HeuristicTokenCounter, ) + + +def test_embedding_counter_prefers_postgres_over_native(monkeypatch) -> None: + """An explicitly configured authoritative PostgreSQL tokenizer stays first.""" + module = types.ModuleType("pg_llm_batch") + module.TokenCounter = _PgCounter + monkeypatch.setitem(sys.modules, "pg_llm_batch", module) + monkeypatch.setattr( + "contextual_orchestrator.token_counting.importlib.import_module", + lambda _name: pytest.fail("native fallback must not load when PostgreSQL starts"), + ) + + counter = build_embedding_token_counter("postgresql://example/tokens") + + assert isinstance(counter, PgTiktokenAdapter) + assert counter.count_text("four", "provider-specific") == 4 From 3c4daac62a61746f69a80dc1c4b8ffced695f489 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 20:19:30 +0900 Subject: [PATCH 08/63] test(embeddings): inject exact accounting boundary Signed-off-by: Seongho Bae --- ...est_batch_embeddings_encoding_dimensions_http_honesty.py | 6 ++++-- tests/test_batch_embeddings_endpoint_http_honesty.py | 6 ++++-- tests/test_batch_embeddings_routing_http_honesty.py | 6 ++++-- tests/test_batch_embeddings_user_http_honesty.py | 6 ++++-- 4 files changed, 16 insertions(+), 8 deletions(-) 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] From 6b5836a50cfa3c5a9f0aa1a91d78c72479293efb Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 20:26:01 +0900 Subject: [PATCH 09/63] test(embeddings): supply exact sync accounting Signed-off-by: Seongho Bae --- tests/test_embeddings_blank_input_http_honesty.py | 6 ++++-- .../test_embeddings_encoding_format_base64_http_honesty.py | 7 +++++-- tests/test_embeddings_encoding_format_http_honesty.py | 6 ++++-- tests/test_embeddings_metadata_http_honesty.py | 6 ++++-- tests/test_embeddings_model_pool_http_honesty.py | 6 ++++-- 5 files changed, 21 insertions(+), 10 deletions(-) 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] From 47f59963a559aba6032c1b1be81310dcc2883227 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 20:32:30 +0900 Subject: [PATCH 10/63] test(embeddings): cover exact optional accounting Signed-off-by: Seongho Bae --- tests/test_embeddings_null_optional_noop_http_honesty.py | 6 ++++-- tests/test_embeddings_routing_http_honesty.py | 6 ++++-- tests/test_embeddings_token_array_input_http_honesty.py | 7 +++++-- tests/test_embeddings_user_field_http_honesty.py | 6 ++++-- ...ring_encoding_tool_choice_endpoint_noop_http_honesty.py | 7 +++++-- 5 files changed, 22 insertions(+), 10 deletions(-) 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() From 994b75174de19897bb31d106ea8148456912a649 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 20:38:06 +0900 Subject: [PATCH 11/63] test(embeddings): inject authoritative counters comprehensively Signed-off-by: Seongho Bae --- ...est_empty_string_numeric_controls_noop_http_honesty.py | 7 +++++-- tests/test_encoding_stream_logprobs_http_honesty.py | 7 +++++-- tests/test_ledger_execution_identity_http_honesty.py | 8 +++++++- tests/test_mode_casefold_http_honesty.py | 6 ++++-- tests/test_openai_user_field_http_honesty.py | 6 ++++-- ...test_prediction_modalities_model_strip_http_honesty.py | 7 +++++-- ...easoning_none_text_empty_logprobs_zero_http_honesty.py | 7 +++++-- ...est_service_tier_encoding_format_strip_http_honesty.py | 7 +++++-- tests/test_token_id_whole_float_coerce_http_honesty.py | 7 +++++-- tests/test_user_null_omit_noop_http_honesty.py | 7 +++++-- tests/test_user_scalar_coerce_http_honesty.py | 7 +++++-- 11 files changed, 55 insertions(+), 21 deletions(-) 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_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_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_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_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_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_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() From 4cf3aff4d5d1b82a6fb3ec9ae1c559cd8125ff36 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:00:33 +0900 Subject: [PATCH 12/63] fix(embeddings): bind concrete tokenizer model Signed-off-by: Seongho Bae --- contextual_orchestrator/cost_router.py | 13 +++++++++-- tests/test_cost_router.py | 31 +++++++++++++++++++++++++- 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 62bfe6757..8722e9239 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -101,6 +101,7 @@ def __init__( else: self.embedding_token_counter = build_embedding_token_counter(postgres_dsn) self.policy = routing_policy or RoutingPolicy(self.config) + self._resolve_virtual_embedding_target = embedding_batch_backend is None # 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 @@ -862,11 +863,19 @@ def _resolve_embedding_target( self, model: str, zdr_only: bool, agent_id: Optional[str] ) -> tuple[str, Optional[str]]: """Resolve one embedding member without losing a caller's member choice.""" - if agent_id is None and not zdr_only: + virtual_models = { + "contextual-orchestrator", + getattr(self.orchestrator, "AUTO_MODEL", ""), + } + if ( + agent_id is None + and not zdr_only + and (model not in virtual_models or not self._resolve_virtual_embedding_target) + ): return model, None selection_model = ( None - if model in {"contextual-orchestrator", getattr(self.orchestrator, "AUTO_MODEL", "")} + if model in virtual_models else model ) with self.orchestrator.request_policy(zdr_only): diff --git a/tests/test_cost_router.py b/tests/test_cost_router.py index b2d280ef4..ddbe53898 100644 --- a/tests/test_cost_router.py +++ b/tests/test_cost_router.py @@ -782,7 +782,11 @@ def count_text(self, text, model): coordinator = CostRoutingCoordinator( orchestrator, embedding_batch_backend=backend, - embedding_token_counter=_SyntheticExactCounter(), + 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 @@ -815,6 +819,31 @@ def test_provider_embedding_runner_accepts_empty_direct_batch() -> None: 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] = [] From a640887c4d878f2f17298af226385deaf89dd700 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:09:23 +0900 Subject: [PATCH 13/63] fix(embeddings): require authoritative provider usage Signed-off-by: Seongho Bae --- contextual_orchestrator/batch_routing.py | 2 + contextual_orchestrator/cost_router.py | 76 ++++++++++++---- contextual_orchestrator/orchestrator.py | 17 +++- .../test_provider_embedding_batch_backend.py | 90 ++++++++++++++++++- 4 files changed, 167 insertions(+), 18 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 8d16fc720..d0e0f8610 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -817,6 +817,8 @@ def poll(self, job: BatchJob) -> Dict[str, Any]: } if status == "failed": document["failure"] = dict(self._errors.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]: diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 8722e9239..bdd18c64d 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -43,6 +43,7 @@ from .kv_config import InMemoryConfigStore from .token_counting import ( HeuristicTokenCounter, + TokenCountUnavailable, build_embedding_token_counter, build_token_counter, ) @@ -149,14 +150,24 @@ def run_provider_embeddings( for request in requests ): raise RuntimeError("provider embedding batch must retain one selected route") - prompt_tokens = sum( - int( - self.embedding_token_counter.count_text( - request.input_text, request.model + try: + prompt_tokens = sum( + int( + self.embedding_token_counter.count_text( + request.input_text, request.model + ) ) + for request in requests ) - for request in requests - ) + except TokenCountUnavailable: + vectors, provider_tokens = orchestrator.client.embed_with_usage( + agent, [request.input_text for request in requests] + ) + if provider_tokens is None: + raise TokenCountUnavailable( + "provider embedding response omitted authoritative usage" + ) + return vectors, provider_tokens vectors = orchestrator.client.embed( agent, [request.input_text for request in requests] ) @@ -962,9 +973,18 @@ 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: + 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( @@ -1084,6 +1104,10 @@ 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) @@ -1092,21 +1116,28 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: if prompt_tokens <= 0 and request is not None: prompt_tokens = request.token_count if request.input_text and prompt_tokens <= 0: - prompt_tokens = int( - self.embedding_token_counter.count_text(request.input_text, item.model) - ) + 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 {}, } ) embeddings: List[Dict[str, Any]] = [] - token_counts: List[int] = [] + token_counts: List[int | None] = [] total_cost_amount = 0.0 currency_code = "USD" for source_index in range(input_count): @@ -1116,6 +1147,17 @@ 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" + ) + 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( @@ -1150,7 +1192,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": { diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index c4c0259f6..b559aa1b8 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -1392,6 +1392,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: @@ -1404,7 +1411,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 @@ -1423,7 +1430,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, diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 154a9ba51..ae6d28250 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -3,18 +3,32 @@ import threading import time +import pytest + from contextual_orchestrator.batch_routing import ( EmbeddingBatchRequest, ProviderEmbeddingBatchBackend, ) -from contextual_orchestrator import CostRoutingCoordinator, ModelAgent, TaskOrchestrator +from contextual_orchestrator import ( + CostRoutingCoordinator, + InMemoryConfigStore, + ModelAgent, + TaskOrchestrator, +) from contextual_orchestrator.orchestrator import ModelClient +from contextual_orchestrator.token_counting import ( + TokenCountUnavailable, + UnavailableEmbeddingTokenCounter, +) class _SyntheticProviderClient(ModelClient): def embed(self, agent, 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): @@ -22,6 +36,80 @@ def count_text(self, text, model): 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()) + coordinator = CostRoutingCoordinator( + orchestrator, 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")) + + +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() From 1aa8c1c679964e4ba9b07d619430b36d982ff791 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:11:50 +0900 Subject: [PATCH 14/63] fix(embeddings): preserve provider batch truth --- contextual_orchestrator/batch_routing.py | 2 + contextual_orchestrator/cost_router.py | 67 ++++++++++++++----- contextual_orchestrator/server.py | 1 + ...er-embedding-lease-and-token-accounting.md | 11 +++ tests/test_cost_router_boundaries.py | 57 ++++++++++++++++ .../test_provider_embedding_batch_backend.py | 50 ++++++++++++++ 6 files changed, 172 insertions(+), 16 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 8d16fc720..3f480a655 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -817,6 +817,8 @@ def poll(self, job: BatchJob) -> Dict[str, Any]: } if status == "failed": document["failure"] = dict(self._errors.get(job.job_id, {})) + elif status == "cancelled": + document["cancellation"] = dict(self._cancellations.get(job.job_id, {})) return document def cancel(self, job: BatchJob, *, reason: str) -> Dict[str, Any]: diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 62bfe6757..692acaa28 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -148,17 +148,27 @@ def run_provider_embeddings( for request in requests ): raise RuntimeError("provider embedding batch must retain one selected route") - prompt_tokens = sum( - int( - self.embedding_token_counter.count_text( - request.input_text, request.model + max_tokens, _max_chars = self._embedding_request_limits() + vectors: List[List[float]] = [] + prompt_tokens = 0 + shard: List[EmbeddingBatchRequest] = [] + shard_tokens = 0 + for request in requests: + if shard and shard_tokens + request.token_count > max_tokens: + vectors.extend( + orchestrator.client.embed( + agent, [item.input_text for item in shard] + ) ) + shard = [] + shard_tokens = 0 + shard.append(request) + shard_tokens += request.token_count + prompt_tokens += request.token_count + if shard: + vectors.extend( + orchestrator.client.embed(agent, [item.input_text for item in shard]) ) - for request in requests - ) - vectors = orchestrator.client.embed( - agent, [request.input_text for request in requests] - ) return vectors, prompt_tokens self.embedding_batch_backend = ProviderEmbeddingBatchBackend( @@ -1068,6 +1078,20 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: "model": model_name, "embeddings": None, } + 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 items: List[EmbeddingBatchResultItem] = self.embedding_batch_backend.retrieve(job) request_by_custom_id = {request.custom_id: request for request in requests} @@ -1093,6 +1117,7 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: "prompt_tokens": max(0, prompt_tokens), "model": item.model, "attribution": dict(request.attribution) if request else {}, + "agent_id": request.agent_id if request else None, } ) @@ -1109,9 +1134,15 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: attribution = dict(parts[0]["attribution"]) 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, @@ -1164,13 +1195,13 @@ def complete_embeddings_batch( metadata: Optional[Dict[str, Any]] = None, zdr_only: bool = False, agent_id: Optional[str] = None, + wait_timeout: Optional[float] = 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, @@ -1180,6 +1211,10 @@ def complete_embeddings_batch( zdr_only=zdr_only, agent_id=agent_id, ) + if wait_timeout is not None and hasattr(self.embedding_batch_backend, "wait"): + status = self.embedding_batch_backend.wait(job, timeout=wait_timeout) + if not status.get("is_complete") and hasattr(self.embedding_batch_backend, "cancel"): + self.embedding_batch_backend.cancel(job, reason="synchronous request deadline elapsed") return self.embeddings_batch_document(job.job_id) def _require_embedding_job(self, batch_id: str) -> BatchJob: diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index a12be9a1b..0f43577e1 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -7023,6 +7023,7 @@ 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=float(orchestrator.client.timeout), )) except Exception as exc: # noqa: BLE001 - measured member failover last_embedding_error = exc diff --git a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md index 11ec8c9fd..3a0876902 100644 --- a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md +++ b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md @@ -28,6 +28,10 @@ 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 @@ -81,6 +85,13 @@ second tokenization implementation. ## 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/ diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index e3a08f2fb..fd0a1c8e5 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -11,6 +11,7 @@ InMemoryConfigStore, ModelAgent, PriceBook, + PriceEntry, TaskOrchestrator, ) from contextual_orchestrator.batch_routing import ( @@ -367,6 +368,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=HeuristicTokenCounter(), + 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" + + def test_cost_report_delegates_to_ledger_window() -> None: ledger = CostLedger(PriceBook(InMemoryConfigStore())) coordinator = _coordinator(ledger=ledger) diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 154a9ba51..83dd06e4e 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -12,7 +12,12 @@ 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] @@ -52,6 +57,26 @@ def runner(_requests): 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_remote_embedding_member_selects_provider_backend() -> None: agent = ModelAgent( "synthetic_embedding", @@ -73,3 +98,28 @@ def test_remote_embedding_member_selects_provider_backend() -> None: time.sleep(0.01) assert document["status"] == "completed" assert [item["embedding"] for item in document["embeddings"]] == [[13.0], [13.0]] + + +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) + + 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"]] From 0a2e1b79b1c4935702f34cb163d1e4f7e330dff4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:15:42 +0900 Subject: [PATCH 15/63] fix(embeddings): bill aggregate provider usage --- contextual_orchestrator/cost_router.py | 26 ++++++++++++++++++- .../test_provider_embedding_batch_backend.py | 13 +++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index e427d04dd..864d1aa71 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -102,7 +102,7 @@ def __init__( else: self.embedding_token_counter = build_embedding_token_counter(postgres_dsn) self.policy = routing_policy or RoutingPolicy(self.config) - self._resolve_virtual_embedding_target = embedding_batch_backend is None + self._resolve_virtual_embedding_target = False # 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 @@ -134,6 +134,7 @@ def __init__( agent for agent in embedding_agents if not agent.base_url.startswith("mock://") ] if remote_embedding_agents: + self._resolve_virtual_embedding_target = True def run_provider_embeddings( requests: List[EmbeddingBatchRequest], ) -> tuple[List[List[float]], int]: @@ -1174,6 +1175,7 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: token_counts: List[int | None] = [] total_cost_amount = 0.0 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: @@ -1187,6 +1189,28 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: 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"]), + ) + 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"]} diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 01cc51a5e..4b3a58e65 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -13,6 +13,8 @@ CostRoutingCoordinator, InMemoryConfigStore, ModelAgent, + PriceBook, + PriceEntry, TaskOrchestrator, ) from contextual_orchestrator.orchestrator import ModelClient @@ -47,14 +49,23 @@ def test_unknown_tokenizer_uses_authoritative_provider_usage() -> None: "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, embedding_token_counter=UnavailableEmbeddingTokenCounter() + 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: From 99c05a0f58972c7d1b9e19a3e0f9ef7b4cf113af Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:21:47 +0900 Subject: [PATCH 16/63] fix(routing): replace stale virtual synthesis models Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 24 +++++-- .../test_chat_response_format_http_honesty.py | 66 +++++++++++++++++++ 2 files changed, 86 insertions(+), 4 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index b559aa1b8..841e6b56d 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4090,10 +4090,12 @@ def _orchestrated_provider_completion( def send_synthesis( payload: dict[str, Any], ) -> tuple[dict[str, Any], ModelAgent]: - """Send once per eligible provider, advancing only after a proven 413 rejection.""" + """Retry 413 broadly and stale virtual models only within one endpoint.""" nonlocal final_agent seen_providers: set[str] = set() preferred = final_agent + preferred_endpoint = preferred.base_url.rstrip("/").casefold() + last_model_not_found: ProviderUpstreamError | None = None ordered_candidates = [ preferred, *( @@ -4103,12 +4105,15 @@ def send_synthesis( ), ] 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 @@ -4141,13 +4146,24 @@ def send_synthesis( ) from exc if not request_too_large: if isinstance(exc, (urllib.error.HTTPError, ProviderUpstreamError)): - raise classify_provider_failure( + classified = classify_provider_failure( exc, agent_id=candidate.id, model=candidate.model, transport="structured_synthesis", - ) from None + ) + if ( + virtual_model + and classified.error_code == "model_not_found" + and candidate_endpoint == preferred_endpoint + ): + last_model_not_found = classified + self._record_failure(candidate.id) + continue + raise classified from None raise + if last_model_not_found is not None: + raise last_model_not_found raise ProviderRequestTooLargeError( "request body exceeds every eligible provider limit" ) diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index a4de6ea3e..c3ef428e1 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -121,6 +121,72 @@ def reject_synthesis(*_args, **_kwargs): 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) + calls = [] + + def send(agent, _endpoint, _payload): + calls.append(agent.id) + if len(calls) == 1: + 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 len(calls) == 2 + assert set(calls) == {"stale_agent", "live_agent"} + assert "other_agent" not in calls + 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"))] + ) + orchestrator.client.proxy_send = lambda *_args, **_kwargs: (_ for _ in ()).throw( + urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None) + ) + 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" + 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() From b927ef3e27b52ad1731cba82c12e68a810695b78 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:22:02 +0900 Subject: [PATCH 17/63] test(telemetry): verify lineage session binding --- tests/test_openai_passthrough.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/tests/test_openai_passthrough.py b/tests/test_openai_passthrough.py index 821b41a79..eccae6394 100644 --- a/tests/test_openai_passthrough.py +++ b/tests/test_openai_passthrough.py @@ -28,6 +28,7 @@ build_server, responses_sse_body, ) +from contextual_orchestrator.telemetry import current_session_id # noqa: E402 def _build() -> TaskOrchestrator: @@ -404,7 +405,21 @@ def test_http_chat_completions_accepts_response_format_and_passes_through() -> N def test_lineage_structured_payload_accepts_session_without_provider_forwarding() -> None: """Lineage correlation is gateway metadata, not a provider request field.""" - server, port, token = _serve() + 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", @@ -421,7 +436,7 @@ def test_lineage_structured_payload_accepts_session_without_provider_forwarding( assert status == 200 assert body["echo"]["response_format"] == {"type": "json_object"} assert "session_id" not in body["echo"] - assert "session_id" in TaskOrchestrator._ORCHESTRATION_ONLY_KEYS + assert observed_sessions == ["synthetic-lineage-session"] def test_http_gateway_default_response_format_resolves_concrete_agent() -> None: From 3016b8ed87ee0af6b7c8310b9c3e79bc9963c32a Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:25:51 +0900 Subject: [PATCH 18/63] fix(server): close embedding workers on shutdown --- contextual_orchestrator/server.py | 19 ++++++++++-- .../test_provider_embedding_batch_backend.py | 31 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index 0f43577e1..33f4fecec 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -158,6 +158,16 @@ class ResponsiveThreadingHTTPServer(ThreadingHTTPServer): """ request_queue_size = socket.SOMAXCONN + embedding_batch_backend: Any = None + + def shutdown(self) -> None: + """Stop accepting requests and release embedding worker threads.""" + try: + super().shutdown() + finally: + close = getattr(self.embedding_batch_backend, "close", None) + if callable(close): + close() # OpenAI request params forwarded verbatim to the provider on passthrough. OPENAI_PASSTHROUGH_PARAM_KEYS = { @@ -8328,7 +8338,9 @@ 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 + return server def serve( @@ -8351,4 +8363,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/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 4b3a58e65..9e4b575df 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -18,6 +18,7 @@ TaskOrchestrator, ) from contextual_orchestrator.orchestrator import ModelClient +from contextual_orchestrator.server import SecurityConfig, build_server from contextual_orchestrator.token_counting import ( TokenCountUnavailable, UnavailableEmbeddingTokenCounter, @@ -176,6 +177,36 @@ def runner(_requests): backend.close() +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_remote_embedding_member_selects_provider_backend() -> None: agent = ModelAgent( "synthetic_embedding", From 4be6d0f5d663c2f4350806e5bcc1543583194b03 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:28:46 +0900 Subject: [PATCH 19/63] fix(embeddings): bound provider input batches --- contextual_orchestrator/cost_router.py | 23 +++++++++++++++---- .../test_provider_embedding_batch_backend.py | 1 + 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 864d1aa71..891106c97 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -56,6 +56,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 _EMBEDDING_UNIT_RE = re.compile(r"\S+\s*|\s+", re.UNICODE) @@ -151,7 +152,7 @@ def run_provider_embeddings( for request in requests ): raise RuntimeError("provider embedding batch must retain one selected route") - max_tokens, _max_chars = self._embedding_request_limits() + max_tokens, _max_chars, max_inputs = self._embedding_request_limits() vectors: List[List[float]] = [] prompt_tokens = 0 shard: List[EmbeddingBatchRequest] = [] @@ -160,7 +161,10 @@ def run_provider_embeddings( request_tokens = request.token_count or len( request.input_text.encode("utf-8") ) - if shard and shard_tokens + request_tokens > max_tokens: + if shard and ( + len(shard) >= max_inputs + or shard_tokens + request_tokens > max_tokens + ): shard_vectors, shard_usage = self._run_embedding_shard( agent, shard ) @@ -928,7 +932,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): @@ -955,9 +959,10 @@ 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 @@ -980,7 +985,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, diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 9e4b575df..340022618 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -243,6 +243,7 @@ def test_provider_embedding_requests_are_sharded_by_the_existing_token_limit() - 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"], From c44885f4baf6d06dcb73b9eb7d51ccf233757ad2 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:41:43 +0900 Subject: [PATCH 20/63] fix(embeddings): fence durable terminal publication --- contextual_orchestrator/batch_job_registry.py | 62 +++++- contextual_orchestrator/batch_routing.py | 182 +++++++++++------- ...er-embedding-lease-and-token-accounting.md | 12 +- docs/library_research.md | 2 +- tests/test_batch_job_registry.py | 106 +++++++++- 5 files changed, 279 insertions(+), 85 deletions(-) diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py index b52f1beb8..1c31ef3b2 100644 --- a/contextual_orchestrator/batch_job_registry.py +++ b/contextual_orchestrator/batch_job_registry.py @@ -91,6 +91,17 @@ def ensure_owned(self, *, refresh: bool = False) -> None: 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.""" @@ -269,7 +280,56 @@ def acquired_local_claim(): with lock: yield _ClaimLease() - return acquired_local_claim() + 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") @property def durable(self) -> bool: diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 9a4338f6f..1ea382952 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -630,6 +630,7 @@ def __init__( 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 @@ -690,6 +691,7 @@ def __init__( 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: @@ -711,6 +713,8 @@ def submit( self, requests: List[EmbeddingBatchRequest], metadata: Optional[Dict[str, Any]] = None ) -> BatchJob: """Persist a queued job and return immediately with a pollable handle.""" + 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._deadlines[job_id] = time.time() + self._registry.retention_seconds @@ -724,82 +728,120 @@ def submit( return BatchJob(job_id=job_id, backend=self.name, status="queued", request_count=len(requests)) def _run_job(self, job_id: str) -> None: - """Execute one persisted job inside the bounded provider worker pool.""" - try: - deadline_epoch = float( - self._deadlines.get( - job_id, time.time() + self._registry.retention_seconds - ) - ) - self._deadlines[job_id] = deadline_epoch - with self._registry.lock( - "provider_embedding_job_execution", job_id, - lease_seconds=self._claim_lease_seconds, - renew_until_epoch=deadline_epoch, - ) as execution_claim: - execution_claim.ensure_owned() - 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]) - vectors, prompt_tokens = self._runner(requests) - execution_claim.ensure_owned() - if self._states.get(job_id) == "cancelled": - 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 = [] - for index, (request, vector) in enumerate(zip(requests, vectors, strict=True)): - items.append( - EmbeddingBatchResultItem( - custom_id=request.custom_id, - index=index, - embedding=vector, - prompt_tokens=0, - model=request.model, - ) - ) + """Execute or reclaim one persisted job until it becomes terminal.""" + deadline_epoch = float( + self._deadlines.get(job_id, time.time() + self._registry.retention_seconds) + ) + self._deadlines[job_id] = deadline_epoch + 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_states", job_id, + "provider_embedding_job_execution", + job_id, lease_seconds=self._claim_lease_seconds, - ): - # Refresh while holding the terminal-state lock. A stale - # worker must never publish after another worker can claim - # the same provider job. - execution_claim.ensure_owned(refresh=True) - if self._states.get(job_id) == "cancelled": - return - self._usage[job_id] = {"prompt_tokens": int(prompt_tokens)} - self._results[job_id] = items - self._states[job_id] = "completed" - except ClaimNotAcquired: - return - except Exception as exc: # noqa: BLE001 - polling exposes a bounded terminal state - with self._registry.lock( - "provider_embedding_job_states", job_id, - lease_seconds=self._claim_lease_seconds, - ): - if self._states.get(job_id) not in {"cancelled", "completed"}: - self._errors[job_id] = { - "error_type": type(exc).__name__, - "http_status": getattr(exc, "status_code", None), - "provider_code": getattr(exc, "provider_code", None), - "retryable": bool(getattr(exc, "retryable", False)), - "failed_shard_index": getattr(exc, "failed_shard_index", None), - } - self._states[job_id] = "failed" - finally: + 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 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 _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() + 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 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, "status_code", None), + "provider_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) diff --git a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md index 3a0876902..e859c6364 100644 --- a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md +++ b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md @@ -38,12 +38,12 @@ that operational boundary grounds the explicit terminal-publication fence here. 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.** A provider embedding worker checks - ownership before provider work, after provider work, and again by - refreshing its execution claim while holding the terminal-state lock. - Only then may it publish usage, results, and `completed`. A stale worker - leaves the job in recoverable `running` state; it does not overwrite a - successor's state or manufacture `failed`. +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 diff --git a/docs/library_research.md b/docs/library_research.md index 7640fa7b7..b667093e7 100644 --- a/docs/library_research.md +++ b/docs/library_research.md @@ -19,7 +19,7 @@ primitives use maintained libraries when the enterprise target requires them. | Rendered policy browser | [MCP Python SDK](https://github.com/modelcontextprotocol/python-sdk) | Keep as a pinned deployment-provided optional package for the existing Camoufox MCP transport; do not claim a repository `policy-browser` extra until this project owns that lock and publish contract. | Reuses the protocol client and Streamable HTTP lifecycle instead of implementing a second transport; static policy analysis does not install it. | | Structured-output validation | `jsonschema` | Use the maintained validator for provider-returned JSON against caller-supplied JSON Schema; keep parsing and the single repair policy in the existing orchestrator. | Reusing `validator_for`, schema checks, and bounded validation avoids an incomplete custom JSON Schema implementation. Provider output and schemas remain untrusted and fail closed. | | 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. | -| Provider-embedding claim ownership | Existing `redis-py` `Lock` acquire/extend/owned/release scripts | Propagate renewal loss to the claim holder and refresh ownership inside the terminal-state critical section before publishing usage, results, or completion. This fences stale local publication but does not claim provider-side exactly-once execution. | Redis's official distributed-lock guidance requires ownership-safe release/extension and recommends fencing when correctness depends on exclusive work. Skipped a custom lock protocol, a new coordination dependency, forced cancellation of synchronous provider I/O, and an unsupported provider-idempotency claim. | +| 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. Legacy chat missing-usage accounting still uses `HeuristicTokenCounter` and remains a known noncompliant follow-up gap; this narrow embedding change is not global token-accounting compliance. | ## Ponytail Decision diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 73f386474..4f717d8fb 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -13,6 +13,7 @@ import sys import threading import time +from types import SimpleNamespace from pathlib import Path from typing import Any, Dict @@ -44,7 +45,9 @@ 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 class LockNotOwnedError(RuntimeError): pass @@ -52,17 +55,25 @@ class LockNotOwnedError(RuntimeError): class _Lock: def __init__(self, client: "FakeValkeyClient", name: str) -> None: self._client = client - self._lose_on_extend = "provider_embedding_job_execution" in name + 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: 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 @@ -74,6 +85,7 @@ 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) @@ -109,6 +121,28 @@ def expire(self, key: str, seconds: int) -> bool: 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:] + 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.""" @@ -213,14 +247,18 @@ def test_renewal_loss_is_visible_to_the_claim_holder() -> None: claim.ensure_owned() -def test_provider_result_is_not_published_after_claim_renewal_loss() -> None: - """A stale provider worker leaves recovery state for the succeeding claimant.""" +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 [[1.0]], 1 + return [[float(calls)]], calls backend = ProviderEmbeddingBatchBackend( runner, @@ -231,12 +269,66 @@ def runner(_requests): [EmbeddingBatchRequest(input_text="synthetic", model="synthetic-model")] ) - assert backend.wait(job, timeout=1)["status"] == "running" - assert backend.retrieve(job) == [] - assert backend.usage(job) == {} + 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 + + if __name__ == "__main__": for name, value in sorted(globals().items()): if name.startswith("test_") and callable(value): From 0b0b9fc21a04f382ba5326fbf8dbdf618e33e811 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 21:55:08 +0900 Subject: [PATCH 21/63] fix(embeddings): bind durable jobs to owners --- contextual_orchestrator/batch_routing.py | 29 +++++++++-- contextual_orchestrator/cost_router.py | 50 ++++++++++++++++--- contextual_orchestrator/server.py | 35 +++++++++++-- tests/test_cost_router_boundaries.py | 31 ++++++++++++ tests/test_openai_passthrough.py | 19 +++++++ .../test_provider_embedding_batch_backend.py | 20 ++++++++ 6 files changed, 169 insertions(+), 15 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 1ea382952..b2a460d53 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -713,19 +713,40 @@ 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._deadlines[job_id] = time.time() + self._registry.retention_seconds - self._states[job_id] = "queued" - self._terminal_events[job_id] = threading.Event() + 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.""" + 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() with self._executor_lock: if self._executor is None: self._executor = ThreadPoolExecutor(max_workers=self._max_concurrency) executor = self._executor - executor.submit(copy_context().run, self._run_job, job_id) - return BatchJob(job_id=job_id, backend=self.name, status="queued", request_count=len(requests)) + 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.""" diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 891106c97..221e06fb6 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -202,6 +202,7 @@ def run_provider_embeddings( # 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) ) @@ -209,6 +210,10 @@ def run_provider_embeddings( 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") + start_embedding_job = getattr(self.embedding_batch_backend, "start", None) + if callable(start_embedding_job): + for recovered_job in list(self._embedding_jobs.values()): + start_embedding_job(recovered_job) # ------------------------------------------------------------------ # Provider / model resolution @@ -864,6 +869,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. @@ -885,13 +891,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(self.embedding_batch_backend, "reserve", None) + start = getattr(self.embedding_batch_backend, "start", None) + if callable(reserve) and callable(start): + job = reserve(requests, metadata=metadata) + else: + job = self.embedding_batch_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( @@ -1105,7 +1119,23 @@ def _count_embedding_tokens(self, text: str, model: str) -> int: 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.""" + 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)) + with self.job_registry.lock( + "embedding_document", batch_id, lease_seconds=lease_seconds + ): + return self._embeddings_batch_document_locked( + batch_id, owner_id=owner_id + ) + + 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 @@ -1118,7 +1148,7 @@ 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) @@ -1220,6 +1250,7 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: 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 @@ -1249,6 +1280,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) currency_code = record.currency_code @@ -1297,6 +1329,7 @@ def complete_embeddings_batch( 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). @@ -1311,16 +1344,19 @@ def complete_embeddings_batch( metadata=metadata, zdr_only=zdr_only, agent_id=agent_id, + owner_id=owner_id, ) if wait_timeout is not None and hasattr(self.embedding_batch_backend, "wait"): status = self.embedding_batch_backend.wait(job, timeout=wait_timeout) if not status.get("is_complete") and hasattr(self.embedding_batch_backend, "cancel"): self.embedding_batch_backend.cancel(job, reason="synchronous request deadline elapsed") - return self.embeddings_batch_document(job.job_id) + return self.embeddings_batch_document(job.job_id, owner_id=owner_id) - def _require_embedding_job(self, batch_id: str) -> BatchJob: + 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/server.py b/contextual_orchestrator/server.py index 33f4fecec..f586374f2 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -61,6 +61,7 @@ detach_trace_context, reset_session_id, session_id_from_headers, + session_id_from_metadata, session_id_from_request, set_session_id, ) @@ -5569,7 +5570,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 @@ -6337,7 +6342,16 @@ def do_POST(self) -> None: # noqa: N802 if isinstance((value := body.get(key)), dict) ] if "session_id" in body: - metadata_values.append({"session_id": body["session_id"]}) + 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) @@ -7021,9 +7035,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( @@ -7033,7 +7053,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=float(orchestrator.client.timeout), + wait_timeout=remaining_timeout, + owner_id=security.principal_id(self.headers), )) except Exception as exc: # noqa: BLE001 - measured member failover last_embedding_error = exc @@ -7044,7 +7065,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, @@ -7115,6 +7141,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 diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index fd0a1c8e5..93ee527a9 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 @@ -350,12 +351,42 @@ 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_embeddings_document_requires_known_batch() -> None: coordinator = _coordinator() with pytest.raises(KeyError, match="embeddings batch job"): 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]: diff --git a/tests/test_openai_passthrough.py b/tests/test_openai_passthrough.py index eccae6394..44446c370 100644 --- a/tests/test_openai_passthrough.py +++ b/tests/test_openai_passthrough.py @@ -439,6 +439,25 @@ def capture_session(*args, **kwargs): # type: ignore[no-untyped-def] 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_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 340022618..aec6fb7df 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -146,6 +146,26 @@ def runner(requests): 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") From d8313445c09f6fc8fa3775c4246544ce0f49238b Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:01:21 +0900 Subject: [PATCH 22/63] fix(embeddings): preserve failures and close workers --- contextual_orchestrator/batch_routing.py | 8 ++- contextual_orchestrator/server.py | 16 +++++- .../test_provider_embedding_batch_backend.py | 55 +++++++++++++++++++ 3 files changed, 74 insertions(+), 5 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index b2a460d53..a9a125ce7 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -817,8 +817,12 @@ def _run_claimed_job(self, job_id: str, execution_claim: Any) -> None: except Exception as exc: # noqa: BLE001 - polling exposes bounded failure metadata error = { "error_type": type(exc).__name__, - "http_status": getattr(exc, "status_code", None), - "provider_code": getattr(exc, "provider_code", None), + "http_status": getattr( + exc, "provider_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), } diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index f586374f2..80db80a96 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -161,14 +161,24 @@ class ResponsiveThreadingHTTPServer(ThreadingHTTPServer): request_queue_size = socket.SOMAXCONN embedding_batch_backend: Any = None + def _close_embedding_backend(self) -> None: + 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: - close = getattr(self.embedding_batch_backend, "close", None) - if callable(close): - close() + 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 = { diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index aec6fb7df..903dd7cf4 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -18,6 +18,7 @@ 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, @@ -177,6 +178,31 @@ def runner(_requests): 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=429, + 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() @@ -227,6 +253,35 @@ def close(self): 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", From c37b4b6d10e108ef04a13be1da8b8ce14624eea4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:08:00 +0900 Subject: [PATCH 23/63] fix(routing): preserve chat fallback error truth --- contextual_orchestrator/__main__.py | 7 +++- contextual_orchestrator/orchestrator.py | 4 ++- tests/test_auto_discovery_server.py | 37 +++++++++++++++++++++ tests/test_passthrough_provider_failover.py | 27 +++++++++++++++ 4 files changed, 73 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 7c81cc1a4..16b9dae79 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -9,6 +9,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 @@ -424,6 +425,7 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l if not model.evidence_only and (model in chat_models or "embedding" in model.capabilities) ] + discovered_chat_agent_ids = {agent_id_for(model) for model in chat_models} existing_by_id = {agent.id: agent for agent in orchestrator.candidates} agents = [] for model in runtime_models: @@ -469,7 +471,9 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l agent.provider_name == "configured_gateway" and not agent.model.strip() and any( - candidate.id != agent.id and not candidate.disabled + candidate.id != agent.id + and not candidate.disabled + and candidate.id in discovered_chat_agent_ids for candidate in orchestrator.candidates ) ): @@ -478,6 +482,7 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l 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): diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 841e6b56d..a5527b969 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4096,6 +4096,7 @@ def send_synthesis( 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, *( @@ -4140,6 +4141,7 @@ def send_synthesis( return send(candidate, endpoint, candidate_payload), 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" @@ -4162,7 +4164,7 @@ def send_synthesis( continue raise classified from None raise - if last_model_not_found is not 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" diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 21e725c36..e8625a692 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -50,6 +50,43 @@ def test_auto_discovery_activates_chat_and_embedding_capable_agents(monkeypatch) 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_auto_discovery_activates_a_free_vision_model_but_free_pool_excludes_it( monkeypatch, ) -> None: diff --git a/tests/test_passthrough_provider_failover.py b/tests/test_passthrough_provider_failover.py index 927002bf5..c27b7af62 100644 --- a/tests/test_passthrough_provider_failover.py +++ b/tests/test_passthrough_provider_failover.py @@ -637,6 +637,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( From 6269a04756939800116499773a62a7e91fd8d909 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:13:41 +0900 Subject: [PATCH 24/63] fix(routing): prove structured model readiness Signed-off-by: Seongho Bae --- CHANGELOG.md | 3 + contextual_orchestrator/__main__.py | 43 ++++++++ contextual_orchestrator/orchestrator.py | 44 ++++++++ docs/library_research.md | 1 + .../adrs/0011-provider-error-boundary.md | 5 + .../adrs/0015-durable-provider-catalog.md | 16 ++- tests/test_auto_discovery_server.py | 74 ++++++++++++- .../test_chat_response_format_http_honesty.py | 100 ++++++++++++++++++ 8 files changed, 280 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ddb03138..87bec12c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ### Fixed +- 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. - 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 diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 16b9dae79..19a451a80 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -403,6 +403,37 @@ 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 {"status":"ok"}.', + } + ], + "response_format": {"type": "json_object"}, + "max_tokens": 8, + "stream": False, + } + try: + response = orchestrator.client.proxy_send_once( + agent, + "chat/completions", + 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. @@ -419,6 +450,18 @@ 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 + ) + if configured_gateway_probe_required: + chat_models = [ + model + for model in chat_models + if model.provider_name != "configured_gateway" + or _probe_configured_gateway_structured_chat(orchestrator, model) + ] runtime_models = [ model for model in discovered diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index a5527b969..0d6dfb4c8 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -3959,9 +3959,11 @@ 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, + _excluded_agent_ids=request_exclusions, ) 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} @@ -4086,6 +4088,31 @@ def _orchestrated_provider_completion( if virtual_model else [final_agent] ) + 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 send_synthesis( payload: dict[str, Any], @@ -4160,6 +4187,7 @@ def send_synthesis( and candidate_endpoint == preferred_endpoint ): last_model_not_found = classified + request_exclusions.add(candidate.id) self._record_failure(candidate.id) continue raise classified from None @@ -5423,6 +5451,7 @@ def conduct( model_name: str = GATEWAY_DEFAULT_MODEL, progress: Any = None, workflow_run_id: str | None = None, + _excluded_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() @@ -5527,6 +5556,7 @@ def conduct( text=task, role=step.role, allowed_agent_ids=free_ids if model_name == self.FREE_MODEL else None, + excluded_agent_ids=_excluded_agent_ids, ) elapsed = (time.perf_counter() - start) * 1000 outputs[step.id] = output @@ -6518,6 +6548,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, dict[str, Any] | None]: """Call an agent with bounded, safety-aware tool retry and failover. @@ -6547,6 +6578,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") @@ -6652,6 +6689,13 @@ def call(agent: ModelAgent) -> tuple[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 diff --git a/docs/library_research.md b/docs/library_research.md index b667093e7..13e6b27fa 100644 --- a/docs/library_research.md +++ b/docs/library_research.md @@ -102,6 +102,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..bc4127af5 100644 --- a/docs/planning/adrs/0011-provider-error-boundary.md +++ b/docs/planning/adrs/0011-provider-error-boundary.md @@ -29,6 +29,10 @@ 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. ## Consequences @@ -41,6 +45,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/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index e8625a692..7d9750f48 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -4,7 +4,10 @@ import os from unittest.mock import patch -from contextual_orchestrator.__main__ import _auto_discover_runtime_agents +from contextual_orchestrator.__main__ import ( + _auto_discover_runtime_agents, + _probe_configured_gateway_structured_chat, +) from contextual_orchestrator.model_discovery import DiscoveredModel from contextual_orchestrator.orchestrator import ModelAgent, TaskOrchestrator @@ -87,6 +90,75 @@ def test_embedding_only_discovery_keeps_chat_fallbacks(monkeypatch) -> None: 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_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, endpoint, payload): + observed.update(agent=agent, endpoint=endpoint, payload=payload) + return {"choices": [{"message": {"content": '{"status":"ok"}'}}]} + + orchestrator.client.proxy_send_once = send + + assert _probe_configured_gateway_structured_chat(orchestrator, model) is True + assert observed["endpoint"] == "chat/completions" + assert observed["payload"]["max_tokens"] == 8 + assert observed["payload"]["response_format"] == {"type": "json_object"} + + def test_auto_discovery_activates_a_free_vision_model_but_free_pool_excludes_it( monkeypatch, ) -> None: diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index c3ef428e1..1ae88528c 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 @@ -160,6 +161,105 @@ def send(agent, _endpoint, _payload): 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( From 583c4b6be7e616cf484752927cf9b7f2644afbd9 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:23:26 +0900 Subject: [PATCH 25/63] fix(embeddings): follow live routes safely --- contextual_orchestrator/batch_routing.py | 15 +- contextual_orchestrator/cost_router.py | 190 ++++++++++++------ contextual_orchestrator/orchestrator.py | 6 +- contextual_orchestrator/token_counting.py | 17 +- .../test_chat_response_format_http_honesty.py | 66 +++++- tests/test_cost_router_boundaries.py | 13 ++ .../test_provider_embedding_batch_backend.py | 65 ++++++ tests/test_token_counting_strategies.py | 10 +- 8 files changed, 296 insertions(+), 86 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index a9a125ce7..1ee421852 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -736,17 +736,16 @@ def reserve( def start(self, job: BatchJob) -> None: """Make a fully registered reservation executable.""" - 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() 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) - executor = self._executor - executor.submit(copy_context().run, self._run_job, job.job_id) + 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.""" diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 221e06fb6..a86e5096d 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 ( @@ -104,6 +105,8 @@ def __init__( 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 @@ -136,66 +139,18 @@ def __init__( ] if remote_embedding_agents: self._resolve_virtual_embedding_target = True - def run_provider_embeddings( - requests: List[EmbeddingBatchRequest], - ) -> tuple[List[List[float]], int]: - if not requests: - return [], 0 - first = requests[0] - agent = ( - orchestrator._agent(first.agent_id) - if first.agent_id is not None - else 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 - - self.embedding_batch_backend = ProviderEmbeddingBatchBackend( - run_provider_embeddings, - job_registry=registry, - max_concurrency=getattr(orchestrator.client, "local_concurrency", 1), - claim_lease_seconds=( - float(orchestrator.client.timeout) - if registry.durable and float(getattr(orchestrator.client, "timeout", 0)) > 0 - else None - ), - ) + 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 + ) # job_id -> submitted BatchJob (so poll/retrieve can be driven by id) self._batch_jobs = registry.mapping("batch_jobs", decode=lambda raw: BatchJob(**raw)) # embeddings batch state: job handle + submitted requests + cached doc, @@ -239,6 +194,104 @@ def _run_embedding_shard( ) return vectors, provider_tokens + def _provider_embedding_backend(self) -> ProviderEmbeddingBatchBackend: + return ProviderEmbeddingBatchBackend( + self._run_provider_embeddings, + job_registry=self.job_registry, + max_concurrency=getattr(self.orchestrator.client, "local_concurrency", 1), + claim_lease_seconds=( + float(self.orchestrator.client.timeout) + if self.job_registry.durable + and float(getattr(self.orchestrator.client, "timeout", 0)) > 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._provider_embedding_backend() + self._embedding_backends[ + self.embedding_batch_backend.name + ] = self.embedding_batch_backend + + def _embedding_backend_for(self, job: BatchJob) -> EmbeddingBatchBackend: + """Keep already-submitted jobs bound to the backend that owns them.""" + return self._embedding_backends.get(job.backend, self.embedding_batch_backend) + + 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 _served_provider_model(self, result: Dict[str, Any], fallback_model: str) -> tuple[str, str]: """Derive ``(provider, model)`` from the served agent in the trace.""" trace = result.get("trace") or [] @@ -882,7 +935,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, @@ -891,12 +946,12 @@ def submit_embeddings_batch( zdr_only=zdr_only, agent_id=resolved_agent_id, ) - reserve = getattr(self.embedding_batch_backend, "reserve", None) - start = getattr(self.embedding_batch_backend, "start", None) + reserve = getattr(backend, "reserve", None) + start = getattr(backend, "start", None) if callable(reserve) and callable(start): job = reserve(requests, metadata=metadata) else: - job = self.embedding_batch_backend.submit(requests, metadata=metadata) + 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 @@ -1021,6 +1076,9 @@ def _split_embedding_input( if text == "": return [("", 0)] 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 ) @@ -1151,7 +1209,8 @@ def _embeddings_batch_document_locked( 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 { "batch_id": batch_id, @@ -1175,7 +1234,7 @@ def _embeddings_batch_document_locked( self._embedding_documents[batch_id] = document return document - items: List[EmbeddingBatchResultItem] = self.embedding_batch_backend.retrieve(job) + items: List[EmbeddingBatchResultItem] = backend.retrieve(job) request_by_custom_id = {request.custom_id: request for request in requests} input_count = self._embedding_input_counts.get(batch_id, len(requests)) part_counts = self._embedding_part_counts.get(batch_id, [1] * input_count) @@ -1346,10 +1405,11 @@ def complete_embeddings_batch( agent_id=agent_id, owner_id=owner_id, ) - if wait_timeout is not None and hasattr(self.embedding_batch_backend, "wait"): - status = self.embedding_batch_backend.wait(job, timeout=wait_timeout) - if not status.get("is_complete") and hasattr(self.embedding_batch_backend, "cancel"): - self.embedding_batch_backend.cancel(job, reason="synchronous request deadline elapsed") + 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( diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index a5527b969..adcfbc7a2 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4086,12 +4086,13 @@ def _orchestrated_provider_completion( if virtual_model else [final_agent] ) + synthesis_failure_recorded = False def send_synthesis( payload: dict[str, Any], ) -> tuple[dict[str, Any], ModelAgent]: """Retry 413 broadly and stale virtual models only within one endpoint.""" - nonlocal final_agent + nonlocal final_agent, synthesis_failure_recorded seen_providers: set[str] = set() preferred = final_agent preferred_endpoint = preferred.base_url.rstrip("/").casefold() @@ -4161,6 +4162,7 @@ def send_synthesis( ): last_model_not_found = classified self._record_failure(candidate.id) + synthesis_failure_recorded = True continue raise classified from None raise @@ -4174,7 +4176,7 @@ def send_synthesis( try: raw, final_agent = send_synthesis(upstream) except Exception as exc: - if not _is_request_too_large_error(exc): + if not _is_request_too_large_error(exc) and not synthesis_failure_recorded: 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) diff --git a/contextual_orchestrator/token_counting.py b/contextual_orchestrator/token_counting.py index a7cca010b..d17b7e261 100644 --- a/contextual_orchestrator/token_counting.py +++ b/contextual_orchestrator/token_counting.py @@ -127,6 +127,18 @@ def count_messages(self, messages: List[dict], model: str = "") -> int: """Reject chat counting because this counter is embedding-only.""" raise TokenCountUnavailable("the native cl100k counter does not count chat framing") + 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, int(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 + class UnavailableEmbeddingTokenCounter: """Represent the absence of an authoritative embedding tokenizer.""" @@ -142,7 +154,10 @@ def _native_token_counter() -> NativeCl100kTokenCounter | None: module = importlib.import_module("contextual_orchestrator._token_packer") except Exception: # noqa: BLE001 - an incompatible wheel is equivalent to absence. return None - if not callable(getattr(module, "count_cl100k", None)): + if not all( + callable(getattr(module, name, None)) + for name in ("count_cl100k", "pack_cl100k") + ): return None return NativeCl100kTokenCounter(module) diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index c3ef428e1..a90805068 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -129,11 +129,15 @@ def test_virtual_structured_synthesis_replaces_stale_model_on_same_endpoint() -> 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 len(calls) == 1: + if agent.id == "stale_agent": raise urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None) return {"choices": [{"message": {"content": '{"status":"ok"}'}}]} @@ -152,8 +156,7 @@ def send(agent, _endpoint, _payload): }, ) assert status == 200, body - assert len(calls) == 2 - assert set(calls) == {"stale_agent", "live_agent"} + assert calls == ["stale_agent", "live_agent"] assert "other_agent" not in calls finally: server.shutdown() @@ -162,12 +165,17 @@ def send(agent, _endpoint, _payload): 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"))] - ) - orchestrator.client.proxy_send = lambda *_args, **_kwargs: (_ for _ in ()).throw( - urllib.error.HTTPError("https://synthetic.invalid", 404, "missing", {}, None) - ) + 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() @@ -182,6 +190,46 @@ def test_explicit_structured_model_preserves_model_not_found() -> None: ) assert status == 404 assert body["error"]["code"] == "model_not_found" + assert calls == ["stale_agent"] + 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) diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index 93ee527a9..6c5231eeb 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -212,6 +212,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.""" diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 903dd7cf4..4d4e8145b 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -223,6 +223,35 @@ def runner(_requests): 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" @@ -305,6 +334,42 @@ def test_remote_embedding_member_selects_provider_backend() -> None: 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_provider_embedding_requests_are_sharded_by_the_existing_token_limit() -> None: agent = ModelAgent( "synthetic_embedding", diff --git a/tests/test_token_counting_strategies.py b/tests/test_token_counting_strategies.py index 9f9137704..5b9dbb6e8 100644 --- a/tests/test_token_counting_strategies.py +++ b/tests/test_token_counting_strategies.py @@ -86,8 +86,13 @@ def test_counter_factory_uses_native_cl100k_only_for_declared_embedding_models( ) -> None: """The installed wheel is a real runtime path without guessing tokenizers.""" calls: list[str] = [] + packed = types.SimpleNamespace(text="hello world", token_count=2) module = types.SimpleNamespace( count_cl100k=lambda text: calls.append(text) or 2, + pack_cl100k=lambda texts, per_input, inputs, total: ( + [packed], + [[0]], + ), ) monkeypatch.setattr( "contextual_orchestrator.token_counting.importlib.import_module", @@ -98,6 +103,9 @@ def test_counter_factory_uses_native_cl100k_only_for_declared_embedding_models( assert isinstance(counter, NativeCl100kTokenCounter) assert counter.count_text("hello world", "text-embedding-3-small") == 2 + assert counter.pack_text("hello world", "text-embedding-3-small", 8192) == [ + ("hello world", 2) + ] assert calls == ["hello world"] with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"): counter.count_text("hello world", "provider-unknown") @@ -112,7 +120,7 @@ def fail(_text: str) -> int: monkeypatch.setattr( "contextual_orchestrator.token_counting.importlib.import_module", - lambda _name: types.SimpleNamespace(count_cl100k=fail), + lambda _name: types.SimpleNamespace(count_cl100k=fail, pack_cl100k=fail), ) counter = build_embedding_token_counter() From 44c4147020c0be8533402c7cf7d80ebe6c91cb30 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:38:30 +0900 Subject: [PATCH 26/63] fix(routing): isolate readiness probe attempts Signed-off-by: Seongho Bae --- CHANGELOG.md | 2 + contextual_orchestrator/__main__.py | 6 +- contextual_orchestrator/orchestrator.py | 62 ++++++++++++++----- .../adrs/0011-provider-error-boundary.md | 4 ++ tests/test_auto_discovery_server.py | 7 +-- .../test_chat_response_format_http_honesty.py | 3 +- tests/test_telemetry.py | 28 +++++++++ 7 files changed, 85 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 87bec12c7..db901d69a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,8 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) - 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. - 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 diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index aa0944d6f..d27472f9d 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -423,11 +423,7 @@ def _probe_configured_gateway_structured_chat( "stream": False, } try: - response = orchestrator.client.proxy_send_once( - agent, - "chat/completions", - payload, - ) + 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 diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 27fc29426..4d1ddb094 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -1962,6 +1962,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, @@ -1969,6 +1981,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("/") @@ -1990,16 +2003,21 @@ 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, + "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" @@ -3877,6 +3895,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") @@ -3962,8 +3986,15 @@ def _orchestrated_provider_completion( request_exclusions: set[str] = set() workflow = self.conduct( messages, - model_name=self.FREE_MODEL if free_only else self.GATEWAY_DEFAULT_MODEL, + 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 = sum(_step_output_token_count(step) for step in workflow["trace"]) model_by_agent = {agent.id: agent.model for agent in self.agents} @@ -4047,12 +4078,6 @@ def _orchestrated_provider_completion( } ) active_profile = effort_profile or self._role_effort_profile("synthesizer") - virtual_model = requested_model in { - None, - "contextual-orchestrator", - self.AUTO_MODEL, - self.FREE_MODEL, - } allowed_agent_ids = ({final_agent.id} if isinstance(required_agent_id, str) else ( { candidate.id @@ -5454,6 +5479,7 @@ def conduct( 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() @@ -5557,7 +5583,9 @@ 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 diff --git a/docs/planning/adrs/0011-provider-error-boundary.md b/docs/planning/adrs/0011-provider-error-boundary.md index bc4127af5..155e1f561 100644 --- a/docs/planning/adrs/0011-provider-error-boundary.md +++ b/docs/planning/adrs/0011-provider-error-boundary.md @@ -33,6 +33,10 @@ available through a public gateway error or an exception cause. 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 diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 7d9750f48..a515093cb 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -147,14 +147,13 @@ def test_configured_gateway_structured_probe_is_bounded_and_validates_output() - orchestrator = TaskOrchestrator([], allow_empty_agents=True) observed = {} - def send(agent, endpoint, payload): - observed.update(agent=agent, endpoint=endpoint, payload=payload) + def send(agent, payload): + observed.update(agent=agent, payload=payload) return {"choices": [{"message": {"content": '{"status":"ok"}'}}]} - orchestrator.client.proxy_send_once = send + orchestrator.client.probe_structured_chat = send assert _probe_configured_gateway_structured_chat(orchestrator, model) is True - assert observed["endpoint"] == "chat/completions" assert observed["payload"]["max_tokens"] == 8 assert observed["payload"]["response_format"] == {"type": "json_object"} diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index 18c5904cb..ddb33315c 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -290,7 +290,8 @@ def send(agent, _endpoint, _payload): ) assert status == 404 assert body["error"]["code"] == "model_not_found" - assert calls == ["stale_agent"] + assert calls + assert set(calls) == {"stale_agent"} finally: server.shutdown() thread.join(timeout=5) diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 91e744318..4d68f9cc7 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -422,6 +422,34 @@ 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", + ) + 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" + + 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() From e2140a2e5d3879820771534981a58c899edf88d4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:47:36 +0900 Subject: [PATCH 27/63] fix(gateway): reconcile live runtime state --- contextual_orchestrator/__main__.py | 74 +++++++++++++++---- contextual_orchestrator/batch_job_registry.py | 31 ++++++++ contextual_orchestrator/batch_routing.py | 29 +++++--- contextual_orchestrator/cost_router.py | 33 +++++++-- contextual_orchestrator/orchestrator.py | 7 ++ contextual_orchestrator/server.py | 5 ++ tests/test_auto_discovery_server.py | 35 ++++++++- tests/test_batch_job_registry.py | 40 ++++++++++ tests/test_model_judge.py | 27 +++++++ .../test_provider_embedding_batch_backend.py | 30 ++++++++ 10 files changed, 275 insertions(+), 36 deletions(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index aa0944d6f..cc9dcd040 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -456,35 +456,67 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l 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: - chat_models = [ - model - for model in chat_models - if model.provider_name != "configured_gateway" - or _probe_configured_gateway_structured_chat(orchestrator, model) - ] + 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) + 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} - existing_by_id = {agent.id: agent for agent in orchestrator.candidates} agents = [] for model in runtime_models: existing = existing_by_id.get(agent_id_for(model)) - routable = is_routable_discovered_model(model) or ( + spend_routable = is_routable_discovered_model(model) or ( "embedding" in model.capabilities and model.spend_admitted ) + structured_routable = agent_id_for(model) not in failed_configured_gateway_probe_ids + routable = 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: continue elif not routable: - tags = (*existing.tags, "spend:blocked") - if existing.disabled and "spend:blocked" not in existing.tags: - tags = (*tags, "spend:blocked:preserve-disabled") + 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, @@ -492,15 +524,27 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l tags=tuple(dict.fromkeys(tags)), ) ) - elif "spend:blocked" in existing.tags: + elif 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", + } ), ) ) diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py index 1c31ef3b2..7bf50ae0c 100644 --- a/contextual_orchestrator/batch_job_registry.py +++ b/contextual_orchestrator/batch_job_registry.py @@ -331,6 +331,37 @@ def publish_provider_embedding_terminal( claim.mark_lost() raise ClaimNotAcquired("durable job claim ownership was lost before publication") + 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.""" diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 1ee421852..c6e6f80e6 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -891,18 +891,23 @@ def poll(self, job: BatchJob) -> Dict[str, Any]: def cancel(self, job: BatchJob, *, reason: str) -> Dict[str, Any]: """Mark queued/running work cancelled and discard any late provider result.""" - 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" - event = self._terminal_events.pop(job.job_id, None) - if event is not None: - event.set() + 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, diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index a86e5096d..09d35eac8 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -151,6 +151,8 @@ def __init__( 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() # job_id -> submitted BatchJob (so poll/retrieve can be driven by id) self._batch_jobs = registry.mapping("batch_jobs", decode=lambda raw: BatchJob(**raw)) # embeddings batch state: job handle + submitted requests + cached doc, @@ -165,9 +167,13 @@ 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") - start_embedding_job = getattr(self.embedding_batch_backend, "start", None) - if callable(start_embedding_job): - for recovered_job in list(self._embedding_jobs.values()): + 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) # ------------------------------------------------------------------ @@ -267,14 +273,25 @@ def _refresh_embedding_backend(self) -> None: ): return self._resolve_virtual_embedding_target = True - self.embedding_batch_backend = self._provider_embedding_backend() - self._embedding_backends[ - self.embedding_batch_backend.name - ] = self.embedding_batch_backend + 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.""" - return self._embedding_backends.get(job.backend, self.embedding_batch_backend) + 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] diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 27fc29426..45f6c0d7c 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -287,6 +287,7 @@ class _FastMLSIJudgeAdapter: served_agent_id: 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: @@ -309,6 +310,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, usage, self.mode if mode is None else mode) @@ -5592,6 +5594,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"): @@ -5604,6 +5607,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: @@ -6947,6 +6951,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", "") @@ -6978,6 +6983,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 @@ -6988,6 +6994,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, diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index 80db80a96..bc89148f5 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -160,8 +160,12 @@ 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() @@ -8377,6 +8381,7 @@ def _send_security_headers(self) -> None: server = ResponsiveThreadingHTTPServer((host, port), Handler) server.embedding_batch_backend = coordinator.embedding_batch_backend + server.embedding_backend_closer = coordinator.close_embedding_backends return server diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 7d9750f48..ecdbf5c31 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -8,7 +8,7 @@ _auto_discover_runtime_agents, _probe_configured_gateway_structured_chat, ) -from contextual_orchestrator.model_discovery import DiscoveredModel +from contextual_orchestrator.model_discovery import DiscoveredModel, agent_from_discovered from contextual_orchestrator.orchestrator import ModelAgent, TaskOrchestrator @@ -134,6 +134,39 @@ def probe(_orchestrator, model): assert all(agent.model != "stale-model" for agent in orchestrator.agents) +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_configured_gateway_structured_probe_is_bounded_and_validates_output() -> None: """The startup probe proves JSON object service with a bounded synthetic call.""" model = DiscoveredModel( diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 4f717d8fb..6229cd8d3 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -124,6 +124,17 @@ def lock(self, name: str, **_kwargs: Any) -> "FakeValkeyClient._Lock": def eval(self, _script: str, key_count: int, *values: Any) -> int: keys = values[:key_count] args = values[key_count:] + if key_count == 2: + 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: @@ -329,6 +340,35 @@ def test_terminal_transaction_rejects_a_transferred_claim_without_partial_writes 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() + + if __name__ == "__main__": for name, value in sorted(globals().items()): if name.startswith("test_") and callable(value): diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index bbf9caaf6..73a2af049 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -667,6 +667,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_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 4d4e8145b..c3b18c780 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -370,6 +370,36 @@ def test_runtime_added_remote_embedding_member_uses_provider_backend() -> None: 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", From 028bebc394b3ebf7c66b626c0efbe46190a0afe5 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 22:49:17 +0900 Subject: [PATCH 28/63] fix(embeddings): tolerate minimal server clients --- contextual_orchestrator/cost_router.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 09d35eac8..0c0b05de3 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -201,14 +201,15 @@ def _run_embedding_shard( return vectors, provider_tokens def _provider_embedding_backend(self) -> ProviderEmbeddingBatchBackend: + client = getattr(self.orchestrator, "client", None) return ProviderEmbeddingBatchBackend( self._run_provider_embeddings, job_registry=self.job_registry, - max_concurrency=getattr(self.orchestrator.client, "local_concurrency", 1), + max_concurrency=getattr(client, "local_concurrency", 1), claim_lease_seconds=( - float(self.orchestrator.client.timeout) + float(client.timeout) if self.job_registry.durable - and float(getattr(self.orchestrator.client, "timeout", 0)) > 0 + and float(getattr(client, "timeout", 0)) > 0 else None ), ) From b2a2607a3a45deeb1ccead374e398438837b06c8 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 23:07:27 +0900 Subject: [PATCH 29/63] fix(accounting): remove heuristic chat token usage (#975) Use provider-reported counts or exact native raw-text tokenizers for declared models. Keep unreconstructible chat usage nullable, route conservatively, and fail enabled budgets closed when measurement is unavailable. Commit-Message-Assisted-by: Claude (via Claude Code) Signed-off-by: Seongho Bae --- CHANGELOG.md | 8 +- README.md | 16 +- contextual_orchestrator/__init__.py | 6 +- contextual_orchestrator/admin.py | 28 +- contextual_orchestrator/api_contract.py | 47 ++- contextual_orchestrator/batch_routing.py | 27 +- contextual_orchestrator/cost_router.py | 127 ++++--- contextual_orchestrator/orchestrator.py | 355 ++++++++++++------ contextual_orchestrator/server.py | 57 +-- contextual_orchestrator/token_counting.py | 222 +++++------ ...er-embedding-lease-and-token-accounting.md | 9 +- ...006-authoritative-chat-token-accounting.md | 100 +++++ docs/adr/README.md | 1 + docs/library_research.md | 3 +- docs/product-technical-gap-baseline.md | 19 +- rust/token_counter/src/lib.rs | 24 +- tests/test_admin_spend_view.py | 2 +- tests/test_api_contract.py | 6 + tests/test_batch_embeddings.py | 21 +- tests/test_batch_optimizer.py | 2 +- tests/test_batch_routing.py | 8 + tests/test_budget_enforcement.py | 63 +++- tests/test_cost_review_server.py | 13 +- tests/test_cost_router.py | 59 ++- tests/test_cost_router_boundaries.py | 16 +- tests/test_evolve_optimizer.py | 8 + tests/test_optimizer.py | 8 + .../test_orchestrator_dispatch_boundaries.py | 12 +- tests/test_provider_usage_capture.py | 23 +- tests/test_spend_analytics.py | 114 +++--- ...am_options_null_flags_noop_http_honesty.py | 49 +-- tests/test_streaming.py | 3 - tests/test_token_counting_boundaries.py | 162 ++++---- tests/test_token_counting_strategies.py | 190 ++-------- 34 files changed, 1015 insertions(+), 793 deletions(-) create mode 100644 docs/adr/0006-authoritative-chat-token-accounting.md diff --git a/CHANGELOG.md b/CHANGELOG.md index db901d69a..25bf0f39e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,11 +17,17 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) 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). Legacy chat estimation remains a documented gap. + usage or cost (ADR 0005). Chat accounting is governed separately by ADR 0006. - OpenRouter discovery no longer marks the entire credential account evidence-only. Authenticated catalog rows may serve ordinary requests, while ZDR-only requests still require explicit route-level ZDR evidence. 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 9d44c0926..4e20ecb53 100644 --- a/contextual_orchestrator/__init__.py +++ b/contextual_orchestrator/__init__.py @@ -62,9 +62,10 @@ snapshot_role_effort_catalog, ) from .token_counting import ( - HeuristicTokenCounter, NativeCl100kTokenCounter, + NativeExactTokenCounter, TokenCountUnavailable, + UnavailableTokenCounter, build_embedding_token_counter, build_token_counter, ) @@ -125,9 +126,10 @@ # config / tokens "InMemoryConfigStore", "get_config_store", - "HeuristicTokenCounter", "NativeCl100kTokenCounter", + "NativeExactTokenCounter", "TokenCountUnavailable", + "UnavailableTokenCounter", "build_embedding_token_counter", "build_token_counter", "ResponseCacheProvider", diff --git a/contextual_orchestrator/admin.py b/contextual_orchestrator/admin.py index 9717eeeed..6d437b212 100644 --- a/contextual_orchestrator/admin.py +++ b/contextual_orchestrator/admin.py @@ -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", @@ -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": "판맀 쀀비도", @@ -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 2da5c8fea..fc535950f 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -20,6 +20,43 @@ }, }, "schemas": { + "AuthoritativeUsage": { + "type": ["object", "null"], + "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"], + "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"], @@ -167,7 +204,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"}, }, } diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index c6e6f80e6..e38c21681 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -94,7 +94,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 +114,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}" @@ -201,8 +205,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" @@ -405,15 +409,26 @@ async def _download() -> Dict[str, Any]: custom_id = entry.get("custom_id", "") body = (entry.get("response") or {}).get("body", {}) answer = _extract_answer(body) - usage = body.get("usage", {}) or {} + usage = body.get("usage") + usage = usage if isinstance(usage, dict) else {} + prompt_tokens = usage.get("prompt_tokens") + completion_tokens = usage.get("completion_tokens") raw_request = tracked.get(custom_id) request = BatchRequest(**raw_request) if raw_request else None items.append( BatchResultItem( custom_id=custom_id, answer=answer, - prompt_tokens=int(usage.get("prompt_tokens", 0)), - completion_tokens=int(usage.get("completion_tokens", 0)), + prompt_tokens=( + prompt_tokens + if type(prompt_tokens) is int and prompt_tokens >= 0 + else None + ), + completion_tokens=( + completion_tokens + if type(completion_tokens) is int and completion_tokens >= 0 + else None + ), attribution=dict(request.attribution) if request else {}, model=request.model if request else "contextual-orchestrator", mode=request.mode if request else "auto", diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 0c0b05de3..58aa705d0 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -43,7 +43,6 @@ from .cost_ledger import CostLedger, PriceBook from .kv_config import InMemoryConfigStore from .token_counting import ( - HeuristicTokenCounter, TokenCountUnavailable, build_embedding_token_counter, build_token_counter, @@ -94,9 +93,7 @@ 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: @@ -363,8 +360,6 @@ def _record_race_endpoint_usage(self, endpoint_id: str, value: Any) -> None: elif isinstance(value, dict): usage = value.get("usage") counts = self._provider_usage(usage) - if counts is None: - return agent = next( (item for item in self.orchestrator.candidates if item.id == endpoint_id), None, @@ -380,8 +375,8 @@ def _record_race_endpoint_usage(self, endpoint_id: str, value: Any) -> None: model_name=context["model_name"], provider_model=self._agent_provider_model(agent, context["model_name"]), 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) @@ -418,24 +413,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( @@ -546,19 +537,23 @@ def complete( provider_response["usage_record_ids"] = [ record.usage_record_id for record in records ] + measurement_status = ( + "unavailable" + if any(record.measurement_status == "unavailable" for record in records) + else "measured" + ) provider_response["cost"] = { "cost_amount": ( round(sum(record.cost_amount for record in records), 6) - if len(currencies) == 1 + if measurement_status == "measured" and len(currencies) == 1 else None ), "currency_code": next(iter(currencies)) if len(currencies) == 1 else "MIXED", - "measurement_status": ( - "estimated" - if any(record.measurement_status == "estimated" for record in records) - else "measured" - ), + "measurement_status": measurement_status, } + if measurement_status == "unavailable": + provider_response["usage"] = None + provider_response["usage_measurement_status"] = measurement_status if len(currencies) > 1: provider_response["cost"]["currency_components"] = [ { @@ -672,24 +667,35 @@ 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} + measurement_status = ( + "unavailable" + if any(item.measurement_status == "unavailable" for item in records) + else "measured" + ) result["cost"] = { "cost_amount": ( round(sum(item.cost_amount for item in records), 6) - if len(currencies) == 1 + if measurement_status == "measured" and len(currencies) == 1 else None ), "currency_code": next(iter(currencies)) if len(currencies) == 1 else "MIXED", - "measurement_status": ( - "estimated" - if any(item.measurement_status == "estimated" for item in records) - else "measured" - ), + "measurement_status": measurement_status, } if len(currencies) > 1: result["cost"]["currency_components"] = [ @@ -728,27 +734,16 @@ def _record_completion( ): """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) + measurement_status = "measured" if ( + prompt_tokens is not None and completion_tokens is not None + ) else "unavailable" + prompt_tokens = prompt_tokens if measurement_status == "measured" else 0 + completion_tokens = completion_tokens if measurement_status == "measured" else 0 return self.ledger.record_usage( provider=provider, model=model, @@ -796,9 +791,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} return { @@ -894,18 +887,24 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict model_name=item.model, provider_model=provider_model, workflow_run_id=job.job_id, - prompt_tokens=item.prompt_tokens or None, - completion_tokens=item.completion_tokens or None, + prompt_tokens=item.prompt_tokens, + completion_tokens=item.completion_tokens, ) recorded.append( { "custom_id": item.custom_id, "answer": item.answer, "usage_record_id": record.usage_record_id, - "cost_amount": record.cost_amount, + "cost_amount": ( + record.cost_amount if record.measurement_status == "measured" else None + ), "currency_code": record.currency_code, - "prompt_tokens": record.prompt_tokens, - "completion_tokens": record.completion_tokens, + "prompt_tokens": ( + record.prompt_tokens if record.measurement_status == "measured" else None + ), + "completion_tokens": ( + record.completion_tokens if record.measurement_status == "measured" else None + ), "measurement_status": record.measurement_status, } ) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index a6320211e..79e45ed09 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -84,6 +84,7 @@ apply_request_profile, snapshot_role_effort_catalog, ) +from .token_counting import TokenCountUnavailable, build_token_counter # content is usually str; multimodal vision messages use OpenAI content-parts lists. @@ -211,29 +212,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: @@ -3338,6 +3340,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). @@ -3382,6 +3385,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 @@ -3425,6 +3429,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) @@ -3998,23 +4003,10 @@ def _orchestrated_provider_completion( _excluded_agent_ids=request_exclusions, _allowed_agent_ids=None if virtual_model else {final_agent.id}, ) - 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 - ) + 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( @@ -4652,6 +4644,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( @@ -4717,8 +4711,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: @@ -4727,14 +4721,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 @@ -4743,19 +4752,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) @@ -4769,6 +4785,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] @@ -7374,17 +7392,23 @@ 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 - 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.""" @@ -7399,7 +7423,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) @@ -7418,10 +7452,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( ( @@ -7433,117 +7476,164 @@ 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 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) + ) 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]: @@ -7551,9 +7641,17 @@ 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) + ) + 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]: @@ -14125,14 +14223,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) @@ -14140,20 +14239,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} @@ -14196,20 +14296,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), } @@ -14262,19 +14371,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"]) @@ -14330,8 +14443,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"), @@ -14356,7 +14468,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}, } @@ -14380,7 +14493,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", } @@ -14448,13 +14562,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/server.py b/contextual_orchestrator/server.py index bc89148f5..d6e540d1e 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -40,7 +40,6 @@ ProviderResponseError, ModelAgent, TaskOrchestrator, - estimate_tokens, _new_chat_completion_id, _responses_to_chat_payload, chat_completion_chunks, @@ -5085,7 +5084,6 @@ def _chat_response_sse_chunks( *, 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") @@ -5167,21 +5165,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 @@ -6716,16 +6728,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, @@ -6855,13 +6863,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, - ), ) ) ) diff --git a/contextual_orchestrator/token_counting.py b/contextual_orchestrator/token_counting.py index d17b7e261..48c5de47f 100644 --- a/contextual_orchestrator/token_counting.py +++ b/contextual_orchestrator/token_counting.py @@ -1,39 +1,18 @@ -"""Token counting seam for usage/cost accounting. - -The cost ledger needs prompt/completion token counts on every completion. -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. -* :class:`NativeCl100kTokenCounter` β€” delegates declared cl100k embedding - models to the packaged Rust extension and reports all other cases as - unavailable. - -Legacy chat selection remains in :func:`build_token_counter`. Embedding -selection is isolated in :func:`build_embedding_token_counter` and never -returns a heuristic. Neither factory reads the environment: the DSN is passed -in by the caller. +"""Authoritative raw-text token counting for accounting boundaries. + +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 importlib -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 operator +from typing import Any, Optional, Protocol -# OpenAI's published tiktoken mapping assigns these embedding deployments to -# cl100k_base. Other model identifiers remain unavailable because a tokenizer -# must never be guessed from a provider/model name. _CL100K_EMBEDDING_MODELS = frozenset( { "text-embedding-ada-002", @@ -41,93 +20,91 @@ "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 TokenCountUnavailable(RuntimeError): - """An authoritative tokenizer is unavailable for the requested model.""" - - -class HeuristicTokenCounter: - """Deterministic, dependency-free token estimator. - - Counts word-ish units (words and standalone punctuation) and applies a - fixed BPE expansion factor. Not exact, but stable and monotonic β€” good - enough for legacy best-effort chat attribution when ``pg_tiktoken`` is not - reachable, and it never varies between runs so tests can assert on it. It - is not authoritative and must not be used for embedding limits or cost. - """ + """An authoritative tokenizer or provider count is unavailable.""" - 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") -class NativeCl100kTokenCounter: - """Use the bundled Rust cl100k counter only for explicitly mapped models.""" +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 a declared cl100k embedding model or fail closed.""" - if model not in _CL100K_EMBEDDING_MODELS: - raise TokenCountUnavailable(f"no authoritative tokenizer is declared for {model!r}") + """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: - return int(self._native_module.count_cl100k(text)) + 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 cl100k tokenizer is unavailable") from exc + raise TokenCountUnavailable("the native tokenizer is unavailable") from exc - def count_messages(self, messages: List[dict], model: str = "") -> int: - """Reject chat counting because this counter is embedding-only.""" - raise TokenCountUnavailable("the native cl100k counter does not count chat framing") + 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]]: + 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}") @@ -135,44 +112,44 @@ def pack_text(self, text: str, model: str, max_tokens: int) -> List[tuple[str, i parts, _shards = self._native_module.pack_cl100k( [text], max_tokens, 1, max_tokens ) - return [(part.text, int(part.token_count)) for part in parts] + 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 -class UnavailableEmbeddingTokenCounter: - """Represent the absence of an authoritative embedding tokenizer.""" +# 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 an embedding token count.""" + """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") + -def _native_token_counter() -> NativeCl100kTokenCounter | None: - """Load the optional in-package extension without making startup depend on it.""" +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 - an incompatible wheel is equivalent to absence. + except Exception: # noqa: BLE001 - incompatible wheel equals absence. return None - if not all( - callable(getattr(module, name, None)) - for name in ("count_cl100k", "pack_cl100k") - ): + functions = ("count_cl100k", "count_o200k", "pack_cl100k") + if not all(callable(getattr(module, name, None)) for name in functions): return None - return NativeCl100kTokenCounter(module) + return NativeExactTokenCounter(module) -def build_embedding_token_counter( - postgres_dsn: Optional[str] = None, - *, - config: Any = None, -) -> PgTiktokenAdapter | NativeCl100kTokenCounter | UnavailableEmbeddingTokenCounter: - """Return an authoritative embedding counter or an explicit unavailable seam. - - PostgreSQL remains authoritative when explicitly configured. The bundled - Rust counter is the fallback only for model identifiers whose published - tokenizer mapping is cl100k. No heuristic estimate is returned here. - """ +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 @@ -180,25 +157,22 @@ def build_embedding_token_counter( return PgTiktokenAdapter(PgTokenCounter(postgres_dsn, config=config)) except Exception: # pragma: no cover - optional authoritative boundary pass - return _native_token_counter() or UnavailableEmbeddingTokenCounter() + return _native_token_counter() or UnavailableTokenCounter() -def build_token_counter( +def build_embedding_token_counter( postgres_dsn: Optional[str] = None, *, config: Any = None, -) -> HeuristicTokenCounter | PgTiktokenAdapter: - """Return the best available token counter. +) -> PgTiktokenAdapter | NativeExactTokenCounter | UnavailableTokenCounter: + """Return an authoritative embedding counter or explicit unavailable seam.""" + return _build_counter(postgres_dsn, config) - 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. - """ - 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() +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/0005-provider-embedding-lease-and-token-accounting.md b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md index e859c6364..043ac4ea4 100644 --- a/docs/adr/0005-provider-embedding-lease-and-token-accounting.md +++ b/docs/adr/0005-provider-embedding-lease-and-token-accounting.md @@ -59,10 +59,9 @@ that operational boundary grounds the explicit terminal-publication fence here. 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 remains a known gap.** This ADR makes no global - token-accounting compliance claim. Existing chat routing and missing-usage - cost paths still use `HeuristicTokenCounter`; replacing that estimate with - authoritative provider/tokenizer evidence is required follow-up work. +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 @@ -81,7 +80,7 @@ that operational boundary grounds the explicit terminal-publication fence here. authoritative PostgreSQL counter is configured. - Lease fencing prevents stale publication but does not eliminate duplicate provider work. -- Chat estimates remain outside this narrow compliance boundary. +- Chat accounting is governed separately by ADR 0006. ## References 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 822a3a1d6..cd4b117e1 100644 --- a/docs/adr/README.md +++ b/docs/adr/README.md @@ -13,6 +13,7 @@ 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-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 13e6b27fa..0f6a82419 100644 --- a/docs/library_research.md +++ b/docs/library_research.md @@ -20,7 +20,8 @@ primitives use maintained libraries when the enterprise target requires them. | Structured-output validation | `jsonschema` | Use the maintained validator for provider-returned JSON against caller-supplied JSON Schema; keep parsing and the single repair policy in the existing orchestrator. | Reusing `validator_for`, schema checks, and bounded validation avoids an incomplete custom JSON Schema implementation. Provider output and schemas remain untrusted and fail closed. | | 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. | | 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. Legacy chat missing-usage accounting still uses `HeuristicTokenCounter` and remains a known noncompliant follow-up gap; this narrow embedding change is not global token-accounting compliance. | +| 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 diff --git a/docs/product-technical-gap-baseline.md b/docs/product-technical-gap-baseline.md index cbad42f7f..fb9f2de26 100644 --- a/docs/product-technical-gap-baseline.md +++ b/docs/product-technical-gap-baseline.md @@ -2271,10 +2271,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 @@ -2285,9 +2284,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 @@ -2300,10 +2298,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/rust/token_counter/src/lib.rs b/rust/token_counter/src/lib.rs index 04a17a4c5..8ebf18875 100644 --- a/rust/token_counter/src/lib.rs +++ b/rust/token_counter/src/lib.rs @@ -1,11 +1,12 @@ -//! Exact cl100k child chunking and provider-request packing for Python. +//! 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; +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() { @@ -18,6 +19,17 @@ fn exact_tokenizer() -> PyResult<&'static tiktoken_rs::CoreBPE> { .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 { @@ -77,6 +89,12 @@ fn count_cl100k(text: &str) -> PyResult { 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() { @@ -240,6 +258,7 @@ fn _token_packer(module: &Bound<'_, PyModule>) -> PyResult<()> { 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)?)?; @@ -406,6 +425,7 @@ mod tests { 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) 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 e67f5902d..0c7bac597 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -58,6 +58,12 @@ 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", + ] assert OPENAPI_SPEC["paths"]["/api/v1/access_reports/{workflow_run_id}"]["get"][ "security" ] == [{"admin_bearer_auth": [], "trace_bearer_auth": []}] diff --git a/tests/test_batch_embeddings.py b/tests/test_batch_embeddings.py index d8a8a985b..701426556 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( @@ -79,7 +88,7 @@ def _serve(): orchestrator, config, price_book=price_book, - embedding_token_counter=HeuristicTokenCounter(), + embedding_token_counter=_ExactTestCounter(), ) token = "cost_token" server = build_server( @@ -279,7 +288,7 @@ def test_batch_embeddings_zdr_only_omitted_model_selects_zdr_capable_embedding_a coordinator = CostRoutingCoordinator( orchestrator, InMemoryConfigStore(), - embedding_token_counter=HeuristicTokenCounter(), + embedding_token_counter=_ExactTestCounter(), ) token = "zdr_batch_token" server = build_server( @@ -305,7 +314,7 @@ def test_pending_batch_preserves_resolved_model_identity() -> None: coordinator = CostRoutingCoordinator( orchestrator, InMemoryConfigStore(), - embedding_token_counter=HeuristicTokenCounter(), + embedding_token_counter=_ExactTestCounter(), embedding_batch_backend=_PendingEmbeddingBackend(), ) @@ -348,7 +357,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, ) @@ -401,7 +410,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_optimizer.py b/tests/test_batch_optimizer.py index 210ee2628..a698016d5 100644 --- a/tests/test_batch_optimizer.py +++ b/tests/test_batch_optimizer.py @@ -115,7 +115,7 @@ def test_mock_default_batch_route_works_without_usage() -> None: orchestrator = _orch() # plain ModelClient: mock batch_chat is sync, usage 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"]) diff --git a/tests/test_batch_routing.py b/tests/test_batch_routing.py index 3b7cdc52f..b9ce2a379 100644 --- a/tests/test_batch_routing.py +++ b/tests/test_batch_routing.py @@ -65,6 +65,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_budget_enforcement.py b/tests/test_budget_enforcement.py index e6d571afc..e689b5728 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,35 @@ 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_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 +101,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 +120,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 +148,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 +159,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 +169,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 +185,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 +208,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 +235,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_cost_review_server.py b/tests/test_cost_review_server.py index 9efbb66af..da8fcd771 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,13 @@ 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" + assert records["items"][0]["measurement_status"] == "unavailable" finally: server.shutdown() diff --git a/tests/test_cost_router.py b/tests/test_cost_router.py index ddbe53898..4f30b2875 100644 --- a/tests/test_cost_router.py +++ b/tests/test_cost_router.py @@ -56,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"]) @@ -64,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: @@ -300,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", @@ -323,18 +324,13 @@ 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"] == 0 + assert records[1]["completion_tokens"] == 0 def test_sync_records_derive_provider_and_model_from_served_agent() -> None: @@ -436,28 +432,25 @@ 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"] == 0 + assert records[3]["total_tokens"] == 0 + assert sum(record["prompt_tokens"] for record in records) == 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] == [0, 0] -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 @@ -473,15 +466,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"] == 0 for record in unreported) + assert all(record["completion_tokens"] == 0 for record in unreported) def test_structured_mixed_currency_costs_are_never_implicitly_converted() -> None: @@ -551,7 +540,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"]) diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index 6c5231eeb..cf2853760 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -29,12 +29,18 @@ _weighted_average_embedding, ) from contextual_orchestrator.token_counting import ( - HeuristicTokenCounter, 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 _coordinator(**kwargs: Any) -> Coordinator: agents = [ ModelAgent( @@ -49,7 +55,7 @@ def _coordinator(**kwargs: Any) -> Coordinator: config = InMemoryConfigStore() price_book = PriceBook(config) if "token_counter" not in kwargs and "embedding_token_counter" not in kwargs: - kwargs["embedding_token_counter"] = HeuristicTokenCounter() + kwargs["embedding_token_counter"] = _ExactTestCounter() return Coordinator(orchestrator, config, price_book=price_book, **kwargs) @@ -452,7 +458,7 @@ def test_embeddings_document_bills_the_selected_agent_not_caller_attribution() - orchestrator, config, price_book=price_book, - embedding_token_counter=HeuristicTokenCounter(), + embedding_token_counter=_ExactTestCounter(), embedding_batch_backend=_DroppingEmbeddingBackend(), ) @@ -553,8 +559,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_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_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_orchestrator_dispatch_boundaries.py b/tests/test_orchestrator_dispatch_boundaries.py index 6777ed885..0e7d154eb 100644 --- a/tests/test_orchestrator_dispatch_boundaries.py +++ b/tests/test_orchestrator_dispatch_boundaries.py @@ -13,6 +13,7 @@ set_backend, ) from contextual_orchestrator.orchestrator import ( + BudgetExceededError, ModelAgent, TaskOrchestrator, WorkflowStep, @@ -142,9 +143,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: @@ -437,7 +438,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", @@ -462,7 +463,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_provider_usage_capture.py b/tests/test_provider_usage_capture.py index 3d57d5d25..50ba77202 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 @@ -52,8 +52,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: @@ -63,25 +63,24 @@ 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",))]) 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",))]) 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_spend_analytics.py b/tests/test_spend_analytics.py index f04e919e7..23a597f15 100644 --- a/tests/test_spend_analytics.py +++ b/tests/test_spend_analytics.py @@ -1,94 +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 - expected = round(row["estimated_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"] == [] + 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"] == "tokenizer" + assert row["cost_usd"] is None + assert not any("estimated" in key for key in row | report["totals"]) -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: @@ -97,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_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 5b9dbb6e8..f713da82a 100644 --- a/tests/test_token_counting_strategies.py +++ b/tests/test_token_counting_strategies.py @@ -1,98 +1,28 @@ -"""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, - NativeCl100kTokenCounter, - PgTiktokenAdapter, + NativeExactTokenCounter, TokenCountUnavailable, - UnavailableEmbeddingTokenCounter, + UnavailableTokenCounter, build_embedding_token_counter, - build_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_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) - - 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"), - ] - - -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_uses_native_cl100k_only_for_declared_embedding_models( - monkeypatch, +def test_native_factory_counts_and_packs_declared_models( + monkeypatch: pytest.MonkeyPatch, ) -> None: - """The installed wheel is a real runtime path without guessing tokenizers.""" - calls: list[str] = [] + calls: list[tuple[str, str]] = [] packed = types.SimpleNamespace(text="hello world", token_count=2) module = types.SimpleNamespace( - count_cl100k=lambda text: calls.append(text) or 2, - pack_cl100k=lambda texts, per_input, inputs, total: ( - [packed], - [[0]], - ), + 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", @@ -101,98 +31,52 @@ def test_counter_factory_uses_native_cl100k_only_for_declared_embedding_models( counter = build_embedding_token_counter() - assert isinstance(counter, NativeCl100kTokenCounter) + 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) ] - assert calls == ["hello world"] with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"): counter.count_text("hello world", "provider-unknown") - assert calls == ["hello world"] - + assert calls == [("cl100k", "hello world"), ("o200k", "hello world")] -def test_native_counter_reports_unavailable_when_extension_call_fails(monkeypatch) -> None: - """One native failure must not fabricate embedding usage.""" +def test_native_failure_is_unavailable() -> None: def fail(_text: str) -> int: raise RuntimeError("synthetic native failure") - monkeypatch.setattr( - "contextual_orchestrator.token_counting.importlib.import_module", - lambda _name: types.SimpleNamespace(count_cl100k=fail, pack_cl100k=fail), + module = types.SimpleNamespace( + count_cl100k=fail, + count_o200k=lambda _text: 1, + pack_cl100k=lambda *_args: ([], []), ) - - counter = build_embedding_token_counter() - - assert isinstance(counter, NativeCl100kTokenCounter) - with pytest.raises(TokenCountUnavailable, match="native cl100k"): - counter.count_text("hello world", "text-embedding-3-large") + counter = NativeExactTokenCounter(module) + with pytest.raises(TokenCountUnavailable, match="native tokenizer"): + counter.count_text("hello", "text-embedding-3-large") -def test_installed_native_counter_matches_cl100k_reference_count() -> None: - """An installed wheel preserves the Rust cl100k parity boundary.""" +def test_installed_native_counter_matches_declared_encoding_parity() -> None: module = pytest.importorskip("contextual_orchestrator._token_packer") - counter = NativeCl100kTokenCounter(module) - - assert counter.count_text("hello world", "text-embedding-3-small") == 2 + 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) + ] -def test_embedding_counter_without_authoritative_backend_is_unavailable(monkeypatch) -> None: - """Missing optional native code is an explicit unavailable result.""" +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: (_ for _ in ()).throw(ImportError("synthetic missing wheel")), + lambda _name: types.SimpleNamespace( + count_cl100k=lambda _text: 1, + count_o200k=lambda _text: 1, + ), ) - counter = build_embedding_token_counter() - - assert isinstance(counter, UnavailableEmbeddingTokenCounter) - with pytest.raises(TokenCountUnavailable, match="no authoritative tokenizer"): - counter.count_text("hello world", "text-embedding-3-small") - - -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) - - 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.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_counter_prefers_postgres_over_native(monkeypatch) -> None: - """An explicitly configured authoritative PostgreSQL tokenizer stays first.""" - module = types.ModuleType("pg_llm_batch") - module.TokenCounter = _PgCounter - monkeypatch.setitem(sys.modules, "pg_llm_batch", module) - monkeypatch.setattr( - "contextual_orchestrator.token_counting.importlib.import_module", - lambda _name: pytest.fail("native fallback must not load when PostgreSQL starts"), - ) - - counter = build_embedding_token_counter("postgresql://example/tokens") - - assert isinstance(counter, PgTiktokenAdapter) - assert counter.count_text("four", "provider-specific") == 4 + assert isinstance(counter, UnavailableTokenCounter) + with pytest.raises(TokenCountUnavailable): + counter.count_text("hello", "text-embedding-3-small") From cef760d05dca56134324cecf444252c1e1ef5809 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 23:07:42 +0900 Subject: [PATCH 30/63] fix(gateway): preserve capability and deadline fences --- contextual_orchestrator/__main__.py | 11 +++--- contextual_orchestrator/batch_routing.py | 31 +++++++++++++++++ contextual_orchestrator/orchestrator.py | 6 ++-- tests/test_auto_discovery_server.py | 37 ++++++++++++++++++++ tests/test_batch_job_registry.py | 28 +++++++++++++++ tests/test_model_judge.py | 44 ++++++++++++++++++++++++ 6 files changed, 149 insertions(+), 8 deletions(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 92f234aa8..22b2bcb9c 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -482,16 +482,15 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l agents = [] for model in runtime_models: existing = existing_by_id.get(agent_id_for(model)) - spend_routable = is_routable_discovered_model(model) or ( - "embedding" in model.capabilities and model.spend_admitted - ) + 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 = spend_routable and structured_routable + 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: continue - elif not routable: + elif not routable or not structured_routable: block_markers = { "spend:blocked", "spend:blocked:preserve-disabled", @@ -516,7 +515,7 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l agents.append( replace( existing, - disabled=True, + disabled=not routable or preserve_disabled, tags=tuple(dict.fromkeys(tags)), ) ) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index c6e6f80e6..6fd588310 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -769,11 +769,42 @@ def _run_job(self, job_id: str) -> None: 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 _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() diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index a6320211e..902b16ba8 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -567,7 +567,7 @@ def from_dict(cls, value: dict[str, Any]) -> "ModelAgent": # pragma: no cover 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] @@ -5516,7 +5516,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 diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index d14a1c46b..492a4f75a 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -167,6 +167,43 @@ def test_failed_gateway_probe_disables_persisted_discovered_agent_after_restart( 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( diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 6229cd8d3..9eeb54bd4 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -369,6 +369,34 @@ def runner(_requests): 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() + + if __name__ == "__main__": for name, value in sorted(globals().items()): if name.startswith("test_") and callable(value): diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index 73a2af049..af79f5f0d 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -208,6 +208,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: From 78ff520a50a43cea7c8cbada1b4fd750f999e6c4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 23:28:13 +0900 Subject: [PATCH 31/63] fix(embeddings): close remaining concurrency races --- contextual_orchestrator/batch_job_registry.py | 28 ++++++++++ contextual_orchestrator/batch_routing.py | 21 ++++---- contextual_orchestrator/cost_router.py | 30 ++++++++--- contextual_orchestrator/orchestrator.py | 3 +- tests/test_batch_job_registry.py | 28 ++++++++++ tests/test_cost_router_boundaries.py | 29 ++++++++++ tests/test_model_judge.py | 54 +++++++++++++++++++ .../test_provider_embedding_batch_backend.py | 2 +- 8 files changed, 176 insertions(+), 19 deletions(-) diff --git a/contextual_orchestrator/batch_job_registry.py b/contextual_orchestrator/batch_job_registry.py index 7bf50ae0c..1d01b4b19 100644 --- a/contextual_orchestrator/batch_job_registry.py +++ b/contextual_orchestrator/batch_job_registry.py @@ -331,6 +331,34 @@ def publish_provider_embedding_terminal( 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: diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 2447d919e..73a63f5bd 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -823,14 +823,17 @@ def _fail_expired_job(self, job_id: str) -> None: 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() - 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" + 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) @@ -863,7 +866,7 @@ def _run_claimed_job(self, job_id: str, execution_claim: Any) -> None: error = { "error_type": type(exc).__name__, "http_status": getattr( - exc, "provider_status", getattr(exc, "status_code", None) + exc, "client_status", getattr(exc, "status_code", None) ), "provider_code": getattr( exc, "error_code", getattr(exc, "provider_code", None) diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 58aa705d0..dfb17d589 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -39,7 +39,7 @@ 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 from .kv_config import InMemoryConfigStore from .token_counting import ( @@ -1198,15 +1198,29 @@ 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.""" - self._require_embedding_job(batch_id, owner_id=owner_id) + 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)) - with self.job_registry.lock( - "embedding_document", batch_id, lease_seconds=lease_seconds - ): - return self._embeddings_batch_document_locked( - batch_id, owner_id=owner_id - ) + try: + with self.job_registry.lock( + "embedding_document", batch_id, lease_seconds=lease_seconds + ): + return 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: + return cached + return { + "batch_id": batch_id, + "status": "in_progress", + "backend": job.backend, + "model": self._embedding_models.get( + batch_id, "contextual-orchestrator" + ), + "embeddings": None, + } def _embeddings_batch_document_locked( self, batch_id: str, *, owner_id: Optional[str] = None diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 63d1ffe59..afc1adc9c 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4145,11 +4145,12 @@ def send_synthesis( 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: diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 9eeb54bd4..d209fe06e 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -125,6 +125,17 @@ 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) @@ -397,6 +408,23 @@ def test_durable_job_past_deadline_becomes_failed_atomically() -> None: 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_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index cf2853760..641479fb5 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -20,6 +20,7 @@ BatchResultItem, EmbeddingBatchResultItem, ) +from contextual_orchestrator.batch_job_registry import ClaimNotAcquired from contextual_orchestrator.cost_router import ( CostRoutingCoordinator as Coordinator, ) @@ -387,6 +388,34 @@ def test_concurrent_embedding_polls_record_usage_once() -> None: 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": "contextual-orchestrator", + "embeddings": None, + } + + cached = {**pending, "status": "completed", "embeddings": []} + coordinator._embedding_documents[job.job_id] = cached + assert coordinator.embeddings_batch_document(job.job_id) == cached + + def test_embeddings_document_requires_known_batch() -> None: coordinator = _coordinator() with pytest.raises(KeyError, match="embeddings batch job"): diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index af79f5f0d..604b5a40d 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -12,6 +12,7 @@ import contextual_orchestrator.orchestrator as orchestrator_module import sys from types import SimpleNamespace +import urllib.error from unittest.mock import patch import pytest @@ -22,6 +23,7 @@ from contextual_orchestrator.orchestrator import ( # noqa: E402 BudgetExceededError, ModelClient, + ProviderRequestTooLargeError, ProviderResponseError, _parse_model_judge_reply, _structured_output_error, @@ -594,6 +596,58 @@ 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_structured_synthesis_failure_updates_provider_health() -> None: """A failed final provider is excluded by the existing circuit policy.""" orchestrator, _ = _orch("unused") diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index c3b18c780..8824d0b8b 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -186,7 +186,7 @@ def runner(_requests): error_code="rate_limit_exceeded", message="provider request failed", client_status=429, - provider_status=429, + provider_status=503, retryable=True, transport="embedding", ) From b9d54687ef1f7b703a7c71e591dc9bc221021adb Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 23:35:51 +0900 Subject: [PATCH 32/63] fix(gateway): scope budgets and tool usage honestly --- contextual_orchestrator/batch_routing.py | 10 +++++++- contextual_orchestrator/cost_router.py | 7 +++--- contextual_orchestrator/orchestrator.py | 13 ++++++++-- contextual_orchestrator/server.py | 24 ++++++++++++++++++- tests/test_batch_job_registry.py | 18 ++++++++++++++ tests/test_budget_enforcement.py | 21 ++++++++++++++++ ...t_chat_parallel_tool_calls_http_honesty.py | 12 +++++++++- 7 files changed, 97 insertions(+), 8 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 73a63f5bd..c241a8c72 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -638,6 +638,7 @@ def __init__( 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") @@ -652,6 +653,13 @@ def __init__( ): raise ValueError("durable provider backend claim lease must be positive") self._claim_lease_seconds = claim_lease_seconds + 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( @@ -740,7 +748,7 @@ def reserve( raise RuntimeError("provider embedding backend is closed") job_id = f"providerembed_{uuid.uuid4().hex}" self._requests[job_id] = list(requests) - self._deadlines[job_id] = time.time() + self._registry.retention_seconds + self._deadlines[job_id] = time.time() + self._execution_timeout_seconds self._states[job_id] = "reserved" return BatchJob( job_id=job_id, diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index dfb17d589..3ee86a592 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -199,16 +199,17 @@ def _run_embedding_shard( 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=( - float(client.timeout) - if self.job_registry.durable - and float(getattr(client, "timeout", 0)) > 0 + 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( diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index afc1adc9c..4c6acf9b3 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -2284,6 +2284,7 @@ def _mock_raw( return { "id": f"chatcmpl_mock_{agent.id}", "object": "chat.completion", + "created": int(time.time()), "model": agent.model, "choices": [ { @@ -7567,7 +7568,11 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> ) candidate_prices_available = ( self.budget_max_cost_usd is None - or all(agent.model in prices for agent in self.agents) + or all( + agent.model in prices + for agent in self.agents + if _is_general_chat_agent(agent) + ) ) return { "measurement_status": measurement_status, @@ -7646,7 +7651,11 @@ def budget_status(self) -> dict[str, Any]: 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) + 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 diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index d6e540d1e..c732f9516 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -5079,6 +5079,28 @@ 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], *, @@ -6852,7 +6874,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) ) diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index d209fe06e..de1471e7f 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -408,6 +408,24 @@ def test_durable_job_past_deadline_becomes_failed_atomically() -> None: 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.reserve([]) + + assert backend._deadlines[job.job_id] == 1005.0 + assert client.expirations["batch_job_registry:provider_embedding_deadlines"] == 123 + backend.close() + + def test_durable_cancellation_cannot_be_overwritten_by_running_transition() -> None: client = FakeValkeyClient() client.lose_execution_extension = False diff --git a/tests/test_budget_enforcement.py b/tests/test_budget_enforcement.py index e689b5728..9c1018e05 100644 --- a/tests/test_budget_enforcement.py +++ b/tests/test_budget_enforcement.py @@ -88,6 +88,27 @@ def test_cost_budget_fails_closed_for_an_unpriced_served_model() -> None: 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 = _orchestrator([_agent()], budget_max_output_tokens=1000) orchestrator.run([{"role": "user", "content": "measure the budget"}]) 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) From 9f8cd4b53b9c0d7df327ff2524f3873cbc0e2950 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 31 Aug 2026 23:56:29 +0900 Subject: [PATCH 33/63] fix(embeddings): bound claimed batch lifetime --- contextual_orchestrator/batch_routing.py | 28 +++++++++-- tests/test_batch_job_registry.py | 62 +++++++++++++++++++++++- 2 files changed, 84 insertions(+), 6 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index c241a8c72..b05eb8d16 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -748,7 +748,6 @@ def reserve( raise RuntimeError("provider embedding backend is closed") job_id = f"providerembed_{uuid.uuid4().hex}" self._requests[job_id] = list(requests) - self._deadlines[job_id] = time.time() + self._execution_timeout_seconds self._states[job_id] = "reserved" return BatchJob( job_id=job_id, @@ -772,10 +771,7 @@ def start(self, job: BatchJob) -> None: def _run_job(self, job_id: str) -> None: """Execute or reclaim one persisted job until it becomes terminal.""" - deadline_epoch = float( - self._deadlines.get(job_id, time.time() + self._registry.retention_seconds) - ) - self._deadlines[job_id] = deadline_epoch + deadline_epoch = self._execution_deadline(job_id) while ( not self._closed.is_set() and self._states.get(job_id) in {"queued", "running"} @@ -803,6 +799,28 @@ def _run_job(self, job_id: str) -> 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"}: diff --git a/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index de1471e7f..1cff3d79d 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -419,13 +419,73 @@ def test_execution_deadline_is_separate_from_result_retention(monkeypatch) -> No execution_timeout_seconds=5, ) - job = backend.reserve([]) + 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.05, + ) + job = backend.submit( + [EmbeddingBatchRequest(input_text="one"), EmbeddingBatchRequest(input_text="two")] + ) + + 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 From 6b3ccabc98725c222dcbe68195a915b3748ab91f Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 00:18:07 +0900 Subject: [PATCH 34/63] fix(gateway): align usage and batch contracts --- contextual_orchestrator/admin.py | 14 ++++++------- contextual_orchestrator/api_contract.py | 18 ++++++++++++++++ contextual_orchestrator/batch_routing.py | 9 +++++++- tests/test_api_contract.py | 10 +++++++++ tests/test_batch_job_registry.py | 26 +++++++++++++++++++++++- 5 files changed, 68 insertions(+), 9 deletions(-) diff --git a/contextual_orchestrator/admin.py b/contextual_orchestrator/admin.py index 6d437b212..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", @@ -316,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": "운영자 μ„Έμ…˜", @@ -1056,7 +1056,7 @@

Spend

estimate
-
ModelEst. output tokensSteps$/1MEst. cost
+
ModelOutput tokensSteps$/1MCost (USD)

diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index fc535950f..8b531832b 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -22,6 +22,7 @@ "schemas": { "AuthoritativeUsage": { "type": ["object", "null"], + "required": ["prompt_tokens", "completion_tokens"], "properties": { "prompt_tokens": {"type": "integer", "minimum": 0}, "completion_tokens": {"type": "integer", "minimum": 0}, @@ -43,6 +44,23 @@ "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"}, diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index b05eb8d16..2daa5b672 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -771,7 +771,14 @@ def start(self, job: BatchJob) -> None: def _run_job(self, job_id: str) -> None: """Execute or reclaim one persisted job until it becomes terminal.""" - deadline_epoch = self._execution_deadline(job_id) + 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"} diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index 0c7bac597..def7b677d 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -64,6 +64,16 @@ def test_openapi_documents_compatibility_front_door() -> None: "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_batch_job_registry.py b/tests/test_batch_job_registry.py index 1cff3d79d..7e6cfbaab 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -48,6 +48,7 @@ def __init__(self) -> None: 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 @@ -64,6 +65,12 @@ def __init__(self, client: "FakeValkeyClient", name: str) -> None: 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 @@ -476,7 +483,7 @@ def runner(requests): runner, job_registry=JobRegistryFactory(client), claim_lease_seconds=0.05, - execution_timeout_seconds=0.05, + execution_timeout_seconds=0.1, ) job = backend.submit( [EmbeddingBatchRequest(input_text="one"), EmbeddingBatchRequest(input_text="two")] @@ -486,6 +493,23 @@ def runner(requests): 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 From 49e1add9242f3c05ef45b0f6e838bbb79410a64a Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 00:32:33 +0900 Subject: [PATCH 35/63] fix(structured): fail over contentless virtual models Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 66 ++++++++++++++------- tests/test_passthrough_provider_failover.py | 57 ++++++++++++++++++ 2 files changed, 101 insertions(+), 22 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 4c6acf9b3..4c3a463cd 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -172,17 +172,24 @@ 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 not isinstance(instance, Mapping): + return "invalid_json_object" + if response_type == "json_object": + return None validator_type = validator_for(schema) try: validator_type.check_schema(schema) @@ -2247,7 +2254,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 = { @@ -4135,6 +4143,24 @@ def _orchestrated_provider_completion( 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]: @@ -4187,7 +4213,9 @@ 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 @@ -4196,6 +4224,15 @@ def send_synthesis( "request body exceeds provider limit" ) from exc if not request_too_large: + 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 if isinstance(exc, (urllib.error.HTTPError, ProviderUpstreamError)): classified = classify_provider_failure( exc, @@ -4230,22 +4267,7 @@ def send_synthesis( 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" - ) - - synthesis_output = provider_output(raw) + synthesis_output = provider_output(final_agent, raw) synthesis_step: dict[str, Any] = { "id": len(workflow["trace"]), "role": "synthesizer", @@ -4304,7 +4326,7 @@ 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) + repaired_output = provider_output(final_agent, repaired) if _structured_output_error(repaired_output, response_format) is not None: self._record_failure(final_agent.id) if final_agent.group_name: diff --git a/tests/test_passthrough_provider_failover.py b/tests/test_passthrough_provider_failover.py index c27b7af62..99e3bb43a 100644 --- a/tests/test_passthrough_provider_failover.py +++ b/tests/test_passthrough_provider_failover.py @@ -20,6 +20,7 @@ from contextual_orchestrator.orchestrator import ( ModelClient, ProviderRequestTooLargeError, + _structured_output_error, ) from contextual_orchestrator.provider_errors import ProviderUpstreamError @@ -1323,3 +1324,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 From 3cc20e6af581e590609f37bcd3cc787b0f6eeb05 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 00:34:06 +0900 Subject: [PATCH 36/63] fix(structured): preserve schema value validation Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 4c3a463cd..d8c9fb368 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -186,10 +186,8 @@ def _structured_output_error( instance = json.loads(content) except (TypeError, json.JSONDecodeError): return "invalid_json" - if not isinstance(instance, Mapping): - return "invalid_json_object" if response_type == "json_object": - return None + return None if isinstance(instance, Mapping) else "invalid_json_object" validator_type = validator_for(schema) try: validator_type.check_schema(schema) From e6329db1b9d0fb59b23cf63b4e4b056743b8a5da Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 00:55:23 +0900 Subject: [PATCH 37/63] fix(api): accept orchestrator auto reasoning effort Signed-off-by: Seongho Bae --- contextual_orchestrator/server.py | 7 ++++++- ...test_chat_reasoning_effort_http_honesty.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index c732f9516..ea90507fa 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -3778,13 +3778,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", ) 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: From 3db85065187ff276c573e2c849f23a94456c0710 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 02:23:26 +0900 Subject: [PATCH 38/63] fix(orchestration): fail over invalid structured candidates Signed-off-by: Seongho Bae --- CHANGELOG.md | 4 + contextual_orchestrator/orchestrator.py | 122 +++++++++++------- .../test_chat_response_format_http_honesty.py | 62 +++++++++ tests/test_model_judge.py | 94 ++++++++++++++ 4 files changed, 234 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 25bf0f39e..667341229 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ### Fixed +- 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. diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index d8c9fb368..fcc09823a 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -4256,37 +4256,41 @@ def send_synthesis( "request body exceeds every eligible provider limit" ) - synthesis_started = time.perf_counter() - try: - raw, final_agent = send_synthesis(upstream) - except Exception as exc: - if not _is_request_too_large_error(exc) and not synthesis_failure_recorded: - 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 - synthesis_output = provider_output(final_agent, 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: + synthesis_started = time.perf_counter() + 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 synthesis_failure_recorded: + 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 + 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 + in_flight_tokens, in_flight_cost = self._trace_budget_spend( [*workflow["trace"], synthesis_step] ) @@ -4325,28 +4329,50 @@ def send_synthesis( self._group_router.observe_failure(final_agent.id) raise repaired_output = provider_output(final_agent, 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) + 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( diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index ddb33315c..aea701b68 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -164,6 +164,68 @@ def send(agent, _endpoint, _payload): 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 = [ diff --git a/tests/test_model_judge.py b/tests/test_model_judge.py index 604b5a40d..7b14bbd23 100644 --- a/tests/test_model_judge.py +++ b/tests/test_model_judge.py @@ -648,6 +648,100 @@ def send(agent, _endpoint, _payload): 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") From c76b58ff5c87556c4bb92f8360694811b4b59d2c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 02:53:40 +0900 Subject: [PATCH 39/63] fix(accounting): hide unavailable usage sentinels --- contextual_orchestrator/batch_routing.py | 14 +++++++++ contextual_orchestrator/cost_ledger.py | 39 +++++++++++++++++++++--- tests/test_batch_job_registry.py | 19 ++++++++++++ tests/test_cost_review_server.py | 16 +++++++++- tests/test_cost_router.py | 23 ++++++++------ 5 files changed, 95 insertions(+), 16 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 2daa5b672..45446b50a 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -871,6 +871,20 @@ def _run_claimed_job(self, job_id: str, execution_claim: Any) -> None: 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} diff --git a/contextual_orchestrator/cost_ledger.py b/contextual_orchestrator/cost_ledger.py index 96be7772a..ef6195435 100644 --- a/contextual_orchestrator/cost_ledger.py +++ b/contextual_orchestrator/cost_ledger.py @@ -1407,9 +1407,18 @@ def rollup( "total_tokens": 0, "cost_amount": Decimal("0"), "currency_code": row.get("currency_code", "USD"), + "measurement_status": "measured", + "unavailable_record_count": 0, }, ) bucket["record_count"] += 1 + status = row.get("measurement_status", "unavailable") + if status == "unavailable": + bucket["measurement_status"] = "unavailable" + bucket["unavailable_record_count"] += 1 + continue + if status == "estimated" and bucket["measurement_status"] == "measured": + bucket["measurement_status"] = "estimated" 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)) @@ -1442,18 +1451,38 @@ def report( def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[str, Any]: """Return grand totals (cost + tokens + record count) over the window.""" rows = self.store.query(start, end) - cost = sum((Decimal(str(row.get("cost_amount", 0))) for row in rows), Decimal("0")) + available = [row for row in rows if row.get("measurement_status") != "unavailable"] + unavailable_count = len(rows) - len(available) + cost = sum((Decimal(str(row.get("cost_amount", 0))) for row in available), Decimal("0")) 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", 0)) for row in available), + "completion_tokens": sum(int(row.get("completion_tokens", 0)) for row in available), + "total_tokens": sum(int(row.get("total_tokens", 0)) for row in available), "cost_amount": float(cost.quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP)), + "measurement_status": ( + "unavailable" + if unavailable_count + else "estimated" + if any(row.get("measurement_status") == "estimated" for row in available) + else "measured" + ), + "unavailable_record_count": unavailable_count, } 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/tests/test_batch_job_registry.py b/tests/test_batch_job_registry.py index 7e6cfbaab..39f745229 100644 --- a/tests/test_batch_job_registry.py +++ b/tests/test_batch_job_registry.py @@ -415,6 +415,25 @@ def test_durable_job_past_deadline_becomes_failed_atomically() -> None: 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) diff --git a/tests/test_cost_review_server.py b/tests/test_cost_review_server.py index da8fcd771..e7a486018 100644 --- a/tests/test_cost_review_server.py +++ b/tests/test_cost_review_server.py @@ -236,7 +236,21 @@ def test_chat_completion_missing_usage_is_unavailable_end_to_end() -> None: 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"] == "unavailable" + 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"]["unavailable_record_count"] == 1 + assert report["items"][0]["measurement_status"] == "unavailable" + assert report["items"][0]["unavailable_record_count"] == 1 finally: server.shutdown() diff --git a/tests/test_cost_router.py b/tests/test_cost_router.py index 4f30b2875..8440e0d6b 100644 --- a/tests/test_cost_router.py +++ b/tests/test_cost_router.py @@ -329,16 +329,15 @@ def test_conducted_plain_completion_records_every_step_usage() -> None: assert result["usage"] is None assert records[0]["prompt_tokens"] == 7 assert records[0]["completion_tokens"] == 3 - assert records[1]["prompt_tokens"] == 0 - assert records[1]["completion_tokens"] == 0 + 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" @@ -435,9 +434,13 @@ def test_structured_provider_workflow_estimates_each_unreported_call() -> None: assert statuses.count("unavailable") == 2 assert set(statuses) == {"measured", "unavailable"} assert result["cost"]["measurement_status"] == "unavailable" - assert records[1]["total_tokens"] == 0 - assert records[3]["total_tokens"] == 0 - assert sum(record["prompt_tokens"] for record in records) == 5 + 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 @@ -446,7 +449,7 @@ def test_structured_provider_workflow_estimates_each_unreported_call() -> None: unavailable = [ record for record in records if record["measurement_status"] == "unavailable" ] - assert [record["prompt_tokens"] for record in unavailable] == [0, 0] + assert [record["prompt_tokens"] for record in unavailable] == [None, None] def test_unreported_provider_calls_remain_unavailable() -> None: @@ -469,8 +472,8 @@ def test_unreported_provider_calls_remain_unavailable() -> None: record for record in records if record["measurement_status"] == "unavailable" ] assert len(unreported) >= 2 - assert all(record["prompt_tokens"] == 0 for record in unreported) - assert all(record["completion_tokens"] == 0 for record in unreported) + 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: From 8d3188c1ae05d5144d91c49d6404b5cbaac1d096 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 02:54:56 +0900 Subject: [PATCH 40/63] fix: publish embedding batch lifecycle metadata Signed-off-by: Seongho Bae --- CHANGELOG.md | 4 ++ contextual_orchestrator/api_contract.py | 8 +++- contextual_orchestrator/batch_job_registry.py | 7 +++- contextual_orchestrator/batch_routing.py | 23 ++++++++++- contextual_orchestrator/cost_router.py | 39 +++++++++++++------ tests/test_batch_embeddings.py | 38 ++++++++++++++++++ tests/test_cost_router_boundaries.py | 3 ++ .../test_provider_embedding_batch_backend.py | 33 ++++++++++++++++ 8 files changed, 140 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 667341229..cec82ee9a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ### Fixed +- 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 diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index 8b531832b..259f76006 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -1026,7 +1026,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 1d01b4b19..21c7fc8a7 100644 --- a/contextual_orchestrator/batch_job_registry.py +++ b/contextual_orchestrator/batch_job_registry.py @@ -48,6 +48,11 @@ 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.""" @@ -230,7 +235,7 @@ def acquired_claim(): if renew_until_epoch is not None: def renew_claim() -> None: - interval = max(0.05, min(lease_seconds / 3, 1.0)) + interval = _claim_renewal_interval_seconds(lease_seconds) while not stop_renewal.wait(interval): remaining = renew_until_epoch - time.time() if remaining <= 0: diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 45446b50a..e9a269c68 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -33,7 +33,11 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional, Protocol -from .batch_job_registry import ClaimNotAcquired, JobRegistryFactory +from .batch_job_registry import ( + ClaimNotAcquired, + JobRegistryFactory, + _claim_renewal_interval_seconds, +) _ROUTING_CATEGORY = "routing" _PROVIDER_CUSTOM_ID_MAX_LENGTH = 64 @@ -526,6 +530,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 @@ -653,6 +658,14 @@ def __init__( ): 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 = ( @@ -1044,6 +1057,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_router.py b/contextual_orchestrator/cost_router.py index 3ee86a592..0dfcff0c4 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -1206,22 +1206,36 @@ def embeddings_batch_document( with self.job_registry.lock( "embedding_document", batch_id, lease_seconds=lease_seconds ): - return self._embeddings_batch_document_locked( + 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: - return cached - return { - "batch_id": batch_id, - "status": "in_progress", - "backend": job.backend, - "model": self._embedding_models.get( - batch_id, "contextual-orchestrator" - ), - "embeddings": 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 @@ -1244,13 +1258,14 @@ def _embeddings_batch_document_locked( 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 = { diff --git a/tests/test_batch_embeddings.py b/tests/test_batch_embeddings.py index 701426556..fa7808d37 100644 --- a/tests/test_batch_embeddings.py +++ b/tests/test_batch_embeddings.py @@ -152,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)) @@ -323,6 +325,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: diff --git a/tests/test_cost_router_boundaries.py b/tests/test_cost_router_boundaries.py index 641479fb5..1bbd94e5a 100644 --- a/tests/test_cost_router_boundaries.py +++ b/tests/test_cost_router_boundaries.py @@ -302,6 +302,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] = {} @@ -409,6 +410,8 @@ def __exit__(self, *_args): "backend": job.backend, "model": "contextual-orchestrator", "embeddings": None, + "poll_after_ms": 250, + "job_retention_ms": coordinator.job_registry.retention_seconds * 1000, } cached = {**pending, "status": "completed", "embeddings": []} diff --git a/tests/test_provider_embedding_batch_backend.py b/tests/test_provider_embedding_batch_backend.py index 8824d0b8b..0eb661fbb 100644 --- a/tests/test_provider_embedding_batch_backend.py +++ b/tests/test_provider_embedding_batch_backend.py @@ -9,6 +9,7 @@ EmbeddingBatchRequest, ProviderEmbeddingBatchBackend, ) +from contextual_orchestrator.batch_job_registry import JobRegistryFactory from contextual_orchestrator import ( CostRoutingCoordinator, InMemoryConfigStore, @@ -147,6 +148,38 @@ def runner(requests): 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() From 5153080b3d6b79df759e153efd87ddb66c65aef4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:06:22 +0900 Subject: [PATCH 41/63] test: expect hidden unavailable stream counts Signed-off-by: Seongho Bae --- tests/test_orchestrated_responses_stream.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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: From 8812f25410605f67c86de37b4238ed78962b41f9 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:06:39 +0900 Subject: [PATCH 42/63] test(accounting): assert unavailable usage is null --- tests/test_orchestrated_responses_stream.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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: From edb508b36834404bc18d7a4b593086cd74ff1f6f Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:11:54 +0900 Subject: [PATCH 43/63] fix(batch): default embedding poll cadence --- contextual_orchestrator/batch_routing.py | 2 +- tests/test_batch_routing_boundaries.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index e9a269c68..48b8dd2e3 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -1063,7 +1063,7 @@ def __init__( if isinstance(poll_interval, (int, float)) and not isinstance(poll_interval, bool) and poll_interval > 0 - else 0 + else 1000 ) self._endpoint_alias = endpoint_alias self._endpoint = endpoint diff --git a/tests/test_batch_routing_boundaries.py b/tests/test_batch_routing_boundaries.py index bb04dbbdf..5754b94f5 100644 --- a/tests/test_batch_routing_boundaries.py +++ b/tests/test_batch_routing_boundaries.py @@ -228,6 +228,7 @@ def test_pg_embedding_backend_full_lifecycle_with_assembler() -> None: def test_pg_embedding_backend_memory_fallback_and_defaults() -> None: client = _FakeEmbeddingClient() backend = PgLlmBatchEmbeddingBackend(client, endpoint_alias="nim-east") + assert backend.poll_after_ms == 1000 requests = [EmbeddingBatchRequest(input_text="solo input", custom_id="emb_solo")] job = backend.submit(requests) assert job.status == "validating" From e66caeca750b9dc9034205e827c2598376fc8502 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:17:28 +0900 Subject: [PATCH 44/63] fix: require backend-owned embedding poll cadence Signed-off-by: Seongho Bae --- contextual_orchestrator/batch_routing.py | 2 +- tests/test_batch_routing_boundaries.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 48b8dd2e3..e9a269c68 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -1063,7 +1063,7 @@ def __init__( if isinstance(poll_interval, (int, float)) and not isinstance(poll_interval, bool) and poll_interval > 0 - else 1000 + else 0 ) self._endpoint_alias = endpoint_alias self._endpoint = endpoint diff --git a/tests/test_batch_routing_boundaries.py b/tests/test_batch_routing_boundaries.py index 5754b94f5..f6c40b700 100644 --- a/tests/test_batch_routing_boundaries.py +++ b/tests/test_batch_routing_boundaries.py @@ -111,6 +111,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] = [] @@ -228,7 +229,7 @@ def test_pg_embedding_backend_full_lifecycle_with_assembler() -> None: def test_pg_embedding_backend_memory_fallback_and_defaults() -> None: client = _FakeEmbeddingClient() backend = PgLlmBatchEmbeddingBackend(client, endpoint_alias="nim-east") - assert backend.poll_after_ms == 1000 + 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" From f678ab9a0e7ffa40dfa023674adb18b06aaedb11 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:22:44 +0900 Subject: [PATCH 45/63] fix: remove auth-failing gateway bootstrap Signed-off-by: Seongho Bae --- CHANGELOG.md | 4 ++ contextual_orchestrator/__main__.py | 21 ++++++--- tests/test_auto_discovery_server.py | 47 +++++++++++++++++++ .../test_chat_response_format_http_honesty.py | 33 +++++++++++++ 4 files changed, 99 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cec82ee9a..062aae9ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) ### Fixed +- 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; + explicit model pins still return their own typed authentication 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 diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 22b2bcb9c..72e3c4989 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -548,16 +548,25 @@ 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 runtime_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 - and candidate.id in discovered_chat_agent_ids - for candidate in orchestrator.candidates + and ( + has_ready_discovered_chat + or bool(failed_configured_gateway_probe_ids) ) ): orchestrator.remove_agent("default", agent.id) diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 492a4f75a..478b7ce22 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -134,6 +134,53 @@ def probe(_orchestrator, 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_probe_disables_persisted_discovered_agent_after_restart( monkeypatch, tmp_path ) -> None: diff --git a/tests/test_chat_response_format_http_honesty.py b/tests/test_chat_response_format_http_honesty.py index aea701b68..43dc8ada3 100644 --- a/tests/test_chat_response_format_http_honesty.py +++ b/tests/test_chat_response_format_http_honesty.py @@ -359,6 +359,39 @@ def send(agent, _endpoint, _payload): 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 = [ From d326a0d4698d0f683e680be2d5a7064c1bf840ca Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:31:49 +0900 Subject: [PATCH 46/63] fix: keep startup alive when catalog probes fail --- contextual_orchestrator/__main__.py | 11 +++++++- tests/test_auto_discovery_server.py | 42 +++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 72e3c4989..2452876a6 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -569,7 +569,16 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l 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: + retired = orchestrator.sync_discovered_agents( + [replace(agent, disabled=True)] + ) + result["updated"].extend(retired["updated"]) has_real_runtime_agent = any( not candidate.disabled and not candidate.base_url.startswith("mock://") diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 478b7ce22..a5f142a3a 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -4,6 +4,8 @@ import os from unittest.mock import patch +import pytest + from contextual_orchestrator.__main__ import ( _auto_discover_runtime_agents, _probe_configured_gateway_structured_chat, @@ -181,6 +183,46 @@ def test_failed_gateway_catalog_probes_remove_unprobed_blank_seed(monkeypatch) - assert selected.model == live_model.model_id +def test_failed_gateway_catalog_probe_disables_the_only_blank_seed(monkeypatch) -> None: + """A failed sole catalog row leaves startup alive but no callable blank seed.""" + 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, + ) + orchestrator = TaskOrchestrator([seed]) + + result = _auto_discover_runtime_agents(orchestrator) + + assert result["updated"] == [seed.id] + assert orchestrator.agents == [] + assert orchestrator.candidates == [replace(seed, disabled=True)] + with pytest.raises(RuntimeError, match="no chat-compatible agent available"): + orchestrator._select_agent("task", "synthesizer") + + def test_failed_gateway_probe_disables_persisted_discovered_agent_after_restart( monkeypatch, tmp_path ) -> None: From 46ed14d8d32a5ba0144ae041712c9868fbef8f46 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 03:48:12 +0900 Subject: [PATCH 47/63] fix: reprobe retired gateway seeds after restart Signed-off-by: Seongho Bae --- CHANGELOG.md | 3 ++- contextual_orchestrator/__main__.py | 5 +---- contextual_orchestrator/orchestrator.py | 14 ++++++++++++ tests/test_auto_discovery_server.py | 30 ++++++++++++++++++++----- 4 files changed, 41 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 062aae9ef..1f11f18e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,7 +15,8 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) - 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; - explicit model pins still return their own typed authentication error. + 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 diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 2452876a6..eda87e459 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -575,10 +575,7 @@ def _auto_discover_runtime_agents(orchestrator: TaskOrchestrator) -> dict[str, l ): orchestrator.remove_agent("default", agent.id) else: - retired = orchestrator.sync_discovered_agents( - [replace(agent, disabled=True)] - ) - result["updated"].extend(retired["updated"]) + orchestrator._retire_runtime_agent(agent.id) has_real_runtime_agent = any( not candidate.disabled and not candidate.base_url.startswith("mock://") diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index fcc09823a..9221d43e3 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -5350,6 +5350,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], diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index a5f142a3a..786c7c9a9 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -10,7 +10,11 @@ _auto_discover_runtime_agents, _probe_configured_gateway_structured_chat, ) -from contextual_orchestrator.model_discovery import DiscoveredModel, agent_from_discovered +from contextual_orchestrator.model_discovery import ( + DiscoveredModel, + agent_from_discovered, + agent_id_for, +) from contextual_orchestrator.orchestrator import ModelAgent, TaskOrchestrator @@ -183,8 +187,10 @@ def test_failed_gateway_catalog_probes_remove_unprobed_blank_seed(monkeypatch) - assert selected.model == live_model.model_id -def test_failed_gateway_catalog_probe_disables_the_only_blank_seed(monkeypatch) -> None: - """A failed sole catalog row leaves startup alive but no callable blank seed.""" +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", @@ -212,16 +218,28 @@ def test_failed_gateway_catalog_probe_disables_the_only_blank_seed(monkeypatch) "contextual_orchestrator.__main__._probe_configured_gateway_structured_chat", lambda *_args: False, ) - orchestrator = TaskOrchestrator([seed]) + agents_db = str(tmp_path / "agents.db") + orchestrator = TaskOrchestrator([seed], agents_db=agents_db) result = _auto_discover_runtime_agents(orchestrator) - assert result["updated"] == [seed.id] + assert result["updated"] == [] assert orchestrator.agents == [] - assert orchestrator.candidates == [replace(seed, disabled=True)] + 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_failed_gateway_probe_disables_persisted_discovered_agent_after_restart( monkeypatch, tmp_path From 8c6787886cf452fc55eeeb19aea44a030dcabe5c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 04:06:17 +0900 Subject: [PATCH 48/63] fix(discovery): recover disabled seed on restart --- contextual_orchestrator/__main__.py | 1 + tests/test_auto_discovery_server.py | 57 +++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index eda87e459..3decbdb3e 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -725,6 +725,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, ) if args.auto_discover_model_agents: _auto_discover_runtime_agents(orchestrator) diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 786c7c9a9..069e6522b 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -9,6 +9,7 @@ from contextual_orchestrator.__main__ import ( _auto_discover_runtime_agents, _probe_configured_gateway_structured_chat, + main, ) from contextual_orchestrator.model_discovery import ( DiscoveredModel, @@ -241,6 +242,62 @@ def test_failed_gateway_catalog_probe_transiently_retires_the_only_blank_seed( 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: From 2ccdd94f42d7c594ef598690c1b0bf031ba24fce Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:08:46 +0900 Subject: [PATCH 49/63] fix(routing): restore request endpoint constraints Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 144 ++++++++++++++- contextual_orchestrator/server.py | 75 +++++++- ...uest-scoped-configured-endpoint-routing.md | 40 +++++ tests/test_routing_endpoint_constraint.py | 164 ++++++++++++++++++ 4 files changed, 409 insertions(+), 14 deletions(-) create mode 100644 docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md create mode 100644 tests/test_routing_endpoint_constraint.py diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 9221d43e3..c2a3249ed 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -87,6 +87,73 @@ 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 +) + + +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.""" + parsed = urlparse(value) + scheme = parsed.scheme.casefold() + try: + 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. ChatMessage = dict[str, Any] ProviderDestination = tuple[int, tuple[Any, ...]] @@ -4436,6 +4503,40 @@ def send_synthesis( } return raw + @contextmanager + def routing_endpoint_scope(self, endpoint: str | None, requested_model: Any): + """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") + if requested_model not in { + None, + self.GATEWAY_DEFAULT_MODEL, + self.AUTO_MODEL, + self.FREE_MODEL, + } and not any( + agent.id in matching and agent.model == requested_model + for agent in self.agents + ): + raise EndpointUnavailableError("endpoint_unavailable") + ids_token = _REQUEST_ENDPOINT_AGENT_IDS.set(matching) + identity_token = _REQUEST_ENDPOINT_IDENTITY.set(normalized) + try: + yield + finally: + _REQUEST_ENDPOINT_IDENTITY.reset(identity_token) + _REQUEST_ENDPOINT_AGENT_IDS.reset(ids_token) + 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 { @@ -4449,6 +4550,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) ) @@ -4670,6 +4772,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, @@ -5796,7 +5904,9 @@ def _plan_generated(self, task: str) -> list[WorkflowStep]: pool = "\n".join( f"- {agent.id}: model={agent.model}, tags={', '.join(agent.tags) or 'none'}" for agent in self.agents - if _is_general_chat_agent(agent) and self._zdr_agent_allowed(agent) + if _is_general_chat_agent(agent) + and self._zdr_agent_allowed(agent) + and _agent_matches_request_endpoint(agent) ) system = ( "You are the workflow conductor. Decompose the user's task into a short workflow.\n" @@ -5828,7 +5938,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: @@ -5942,6 +6056,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 @@ -6182,7 +6297,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: @@ -6202,7 +6319,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) @@ -6267,7 +6390,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) @@ -6284,7 +6412,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] @@ -6989,7 +7119,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 diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index ea90507fa..e47219538 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -35,11 +35,13 @@ from .batch_routing import BatchRequest from .orchestrator import ( BudgetExceededError, + EndpointUnavailableError, MAX_LOCAL_CONCURRENCY, ProviderRequestTooLargeError, ProviderResponseError, ModelAgent, TaskOrchestrator, + normalize_endpoint_selector, _new_chat_completion_id, _responses_to_chat_payload, chat_completion_chunks, @@ -3159,7 +3161,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 @@ -3170,7 +3174,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") @@ -3214,6 +3221,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 {} @@ -6250,6 +6281,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": @@ -6389,6 +6421,23 @@ 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"), + ) + 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") @@ -6812,7 +6861,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 @@ -6913,7 +6964,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 @@ -7386,7 +7439,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( @@ -7554,7 +7609,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 @@ -7650,7 +7707,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", @@ -7764,6 +7823,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) 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..14bcfeb19 --- /dev/null +++ b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md @@ -0,0 +1,40 @@ +# 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. diff --git a/tests/test_routing_endpoint_constraint.py b/tests/test_routing_endpoint_constraint.py new file mode 100644 index 000000000..ba6857de3 --- /dev/null +++ b/tests/test_routing_endpoint_constraint.py @@ -0,0 +1,164 @@ +"""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 _orchestrator() -> TaskOrchestrator: + return TaskOrchestrator( + [ + ModelAgent("agent_a", "model-a", base_url="https://a.example/v1"), + 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 + + +@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_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) From 070ac41e56bf8924a240f6db365b3f90a80eaf1e Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:17:40 +0900 Subject: [PATCH 50/63] test(routing): prove endpoint cache isolation Signed-off-by: Seongho Bae --- tests/test_routing_endpoint_constraint.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_routing_endpoint_constraint.py b/tests/test_routing_endpoint_constraint.py index ba6857de3..74a3c6513 100644 --- a/tests/test_routing_endpoint_constraint.py +++ b/tests/test_routing_endpoint_constraint.py @@ -77,6 +77,19 @@ def select(endpoint: str) -> str: 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", [ From 3175e6d9e04cbdb2cfa24e004cb01ab7c9b06156 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:32:00 +0900 Subject: [PATCH 51/63] fix(routing): validate endpoint-local model eligibility --- contextual_orchestrator/orchestrator.py | 63 +++++++-- contextual_orchestrator/server.py | 1 + ...uest-scoped-configured-endpoint-routing.md | 14 ++ tests/test_routing_endpoint_constraint.py | 131 +++++++++++++++++- 4 files changed, 197 insertions(+), 12 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index c2a3249ed..ec854c363 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -93,6 +93,7 @@ _REQUEST_ENDPOINT_IDENTITY: ContextVar[str | None] = ContextVar( "contextual_orchestrator_request_endpoint_identity", default=None ) +_INVALID_REQUESTED_MODEL = object() class EndpointUnavailableError(ValueError): @@ -4504,7 +4505,13 @@ def send_synthesis( return raw @contextmanager - def routing_endpoint_scope(self, endpoint: str | None, requested_model: Any): + 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 @@ -4519,24 +4526,58 @@ def routing_endpoint_scope(self, endpoint: str | None, requested_model: Any): ) if not matching: raise EndpointUnavailableError("endpoint_unavailable") - if requested_model not in { - None, - self.GATEWAY_DEFAULT_MODEL, - self.AUTO_MODEL, - self.FREE_MODEL, - } and not any( - agent.id in matching and agent.model == requested_model - for agent in self.agents - ): - 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: + """Fail closed only for real endpoint/model conflicts on the active request.""" + if requested_model is _INVALID_REQUESTED_MODEL: + return True + if requested_model in { + None, + self.GATEWAY_DEFAULT_MODEL, + self.AUTO_MODEL, + self.FREE_MODEL, + }: + try: + ranked = self._ranked_agents( + "request endpoint probe", + "worker", + free_only=requested_model == self.FREE_MODEL, + ) + except RuntimeError: + return False + return bool(ranked) + 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 { diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index e47219538..92bd1a88b 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -6428,6 +6428,7 @@ def do_POST(self) -> None: # noqa: N802 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__() diff --git a/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md index 14bcfeb19..e47be4479 100644 --- a/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md +++ b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md @@ -38,3 +38,17 @@ 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 + +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/tests/test_routing_endpoint_constraint.py b/tests/test_routing_endpoint_constraint.py index 74a3c6513..bce90fd6d 100644 --- a/tests/test_routing_endpoint_constraint.py +++ b/tests/test_routing_endpoint_constraint.py @@ -31,11 +31,44 @@ def chat(self, agent: ModelAgent, messages: list, temperature: float = 0.2, **_k 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"), + 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(), @@ -175,3 +208,99 @@ def test_http_surfaces_constrain_candidates_and_preserve_envelopes(path: str, pa 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) From cf3c4cdde8f937ae2aab160358a124be9ec15222 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:39:54 +0900 Subject: [PATCH 52/63] fix(discovery): request JSON in gateway probe Signed-off-by: Seongho Bae --- contextual_orchestrator/__main__.py | 2 +- tests/test_auto_discovery_server.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/__main__.py b/contextual_orchestrator/__main__.py index 3decbdb3e..7d8d86abf 100644 --- a/contextual_orchestrator/__main__.py +++ b/contextual_orchestrator/__main__.py @@ -415,7 +415,7 @@ def _probe_configured_gateway_structured_chat( "messages": [ { "role": "user", - "content": 'Return only {"status":"ok"}.', + "content": 'Return only this JSON object: {"status":"ok"}.', } ], "response_format": {"type": "json_object"}, diff --git a/tests/test_auto_discovery_server.py b/tests/test_auto_discovery_server.py index 069e6522b..260aa6c88 100644 --- a/tests/test_auto_discovery_server.py +++ b/tests/test_auto_discovery_server.py @@ -390,6 +390,7 @@ def send(agent, payload): 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( From c64ee4c0ace5fce224801b9692af4a870adf3929 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:42:52 +0900 Subject: [PATCH 53/63] fix(routing): keep endpoint validation side effect free --- contextual_orchestrator/orchestrator.py | 24 +++++------ ...uest-scoped-configured-endpoint-routing.md | 8 ++++ fuzz/fuzz_endpoint_selector.py | 28 ++++++++++++ fuzz/targets.py | 12 ++++++ tests/fuzz/test_fuzz_properties.py | 7 +++ tests/test_routing_endpoint_constraint.py | 43 +++++++++++++++++++ 6 files changed, 110 insertions(+), 12 deletions(-) create mode 100644 fuzz/fuzz_endpoint_selector.py diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index ec854c363..2bb885f79 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -102,9 +102,9 @@ class EndpointUnavailableError(ValueError): def normalize_endpoint_selector(value: str) -> str: """Normalize an endpoint selector without ever using it as transport input.""" - parsed = urlparse(value) - scheme = parsed.scheme.casefold() try: + parsed = urlparse(value) + scheme = parsed.scheme.casefold() hostname = parsed.hostname port = parsed.port except ValueError as exc: @@ -4554,7 +4554,7 @@ def _normalize_endpoint_requested_model( return normalized def _request_endpoint_supports_model(self, requested_model: Any) -> bool: - """Fail closed only for real endpoint/model conflicts on the active request.""" + """Check endpoint-local eligibility without ranking or provider I/O.""" if requested_model is _INVALID_REQUESTED_MODEL: return True if requested_model in { @@ -4563,15 +4563,15 @@ def _request_endpoint_supports_model(self, requested_model: Any) -> bool: self.AUTO_MODEL, self.FREE_MODEL, }: - try: - ranked = self._ranked_agents( - "request endpoint probe", - "worker", - free_only=requested_model == self.FREE_MODEL, - ) - except RuntimeError: - return False - return bool(ranked) + 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: diff --git a/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md index e47be4479..d447a2a1c 100644 --- a/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md +++ b/docs/planning/adrs/0039-request-scoped-configured-endpoint-routing.md @@ -41,6 +41,14 @@ 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 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/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_routing_endpoint_constraint.py b/tests/test_routing_endpoint_constraint.py index bce90fd6d..88410ddde 100644 --- a/tests/test_routing_endpoint_constraint.py +++ b/tests/test_routing_endpoint_constraint.py @@ -139,6 +139,49 @@ def test_invalid_endpoint_selector_is_rejected(endpoint: str) -> None: 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"}) From bb98f0b66d4f314addb281e77b03f63716a6aafc Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:45:03 +0900 Subject: [PATCH 54/63] fix(telemetry): retain actionable probe failures Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 4 +++ contextual_orchestrator/provider_errors.py | 3 +- contextual_orchestrator/telemetry.py | 21 ++++++++++-- tests/test_provider_error_taxonomy.py | 16 +++++++++ tests/test_telemetry.py | 39 +++++++++++++++++++++- 5 files changed, 79 insertions(+), 4 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 9221d43e3..39fdd6d5a 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -2018,6 +2018,10 @@ def _proxy_send( "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), } diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py index ea47cee99..f159d1012 100644 --- a/contextual_orchestrator/provider_errors.py +++ b/contextual_orchestrator/provider_errors.py @@ -49,7 +49,8 @@ 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)\b|" + r"\b(?:messages?|content)\s*(?:[:=]|\[|\{)" r")" ) diff --git a/contextual_orchestrator/telemetry.py b/contextual_orchestrator/telemetry.py index 3b9fdccb6..ce2f4c9d0 100644 --- a/contextual_orchestrator/telemetry.py +++ b/contextual_orchestrator/telemetry.py @@ -51,7 +51,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", @@ -306,21 +310,34 @@ 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 + error_summary = safe_provider_message(classified) or 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/tests/test_provider_error_taxonomy.py b/tests/test_provider_error_taxonomy.py index af4a32ad8..b8ec08194 100644 --- a/tests/test_provider_error_taxonomy.py +++ b/tests/test_provider_error_taxonomy.py @@ -77,6 +77,22 @@ 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}}) + ) == actionable + for diagnostic in ( + "messages=[{'role':'user','content':'customer secret'}]", + "prompt=customer secret", + "input: customer secret", + ): + assert safe_provider_message( + _body_http_error(400, {"error": {"message": diagnostic}}) + ) 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_telemetry.py b/tests/test_telemetry.py index 4d68f9cc7..f41016b2b 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -438,6 +438,7 @@ def capture(name, attributes): 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) @@ -448,6 +449,8 @@ def capture(name, attributes): 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): @@ -481,7 +484,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", + "the provider connection failed or did not finish in time", + ) assert "provider-response-secret" not in caplog.text assert "session-secret" not in caplog.text @@ -506,6 +513,36 @@ 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' to use json_object" + 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", message) + 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 message in caplog.text + assert "private.example" 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() From 50ea768e38d42eeacb8ac66d3b6842a8ddaacb69 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:48:12 +0900 Subject: [PATCH 55/63] Restrict provider telemetry summaries --- contextual_orchestrator/telemetry.py | 26 +++++++++++++++++++++++++- tests/test_telemetry.py | 21 +++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/telemetry.py b/contextual_orchestrator/telemetry.py index ce2f4c9d0..32b58731e 100644 --- a/contextual_orchestrator/telemetry.py +++ b/contextual_orchestrator/telemetry.py @@ -4,6 +4,8 @@ import hashlib import logging +import re +import urllib.error from collections.abc import Iterator, Mapping from contextlib import contextmanager from contextvars import ContextVar, Token @@ -39,6 +41,12 @@ # 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['\"]?" + r"(?: in some form,)? to use " + r"(?:['\"]?response_format['\"]? of type ['\"]?json_object['\"]?|json_object)\.?$", + re.IGNORECASE, +) _ALLOWED_ATTRIBUTE_KEYS = frozenset( { "gen_ai.operation.name", @@ -315,7 +323,23 @@ def traced( classified = classify_provider_failure(exc, agent_id="", model="") failure_code = classified.error_code provider_status = classified.provider_status - error_summary = safe_provider_message(classified) or failure_code + # Arbitrary provider prose can echo caller content even when it has + # no assignment-shaped marker. Only one fixed-vocabulary schema + # diagnostic is safe enough to export; every other HTTP/provider + # failure uses the package-owned stable code. + provider_summary = ( + safe_provider_message(exc) + if isinstance(exc, urllib.error.HTTPError) + else None + ) + error_summary = ( + provider_summary + if provider_summary is not None + and _SAFE_SCHEMA_DIAGNOSTIC.fullmatch(provider_summary) + else failure_code + if isinstance(exc, urllib.error.HTTPError) + else safe_provider_message(classified) or failure_code + ) model_group = safe.get("contextual_orchestrator.model_group", "ungrouped") fallback_outcome = safe.get( "contextual_orchestrator.fallback_outcome", "not_observed" diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index f41016b2b..ec9a72dc7 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -543,6 +543,27 @@ def test_traced_logs_actionable_bounded_failure_evidence(monkeypatch, caplog): 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_annotate_and_usage_helpers_filter_to_allowed_genai_attributes(): """Span annotation keeps approved scalars only; prompts never enter spans.""" span = MagicMock() From 4e548b65b0018c43a3dd19ba79e513f1d301437a Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:51:47 +0900 Subject: [PATCH 56/63] fix(security): reject quoted request fields in provider errors --- contextual_orchestrator/provider_errors.py | 2 +- tests/test_provider_error_taxonomy.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py index f159d1012..f142e2430 100644 --- a/contextual_orchestrator/provider_errors.py +++ b/contextual_orchestrator/provider_errors.py @@ -50,7 +50,7 @@ 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)\b|" - r"\b(?:messages?|content)\s*(?:[:=]|\[|\{)" + r"\b(?:messages?|content)[\"']?\s*(?:[:=]|\[|\{)" r")" ) diff --git a/tests/test_provider_error_taxonomy.py b/tests/test_provider_error_taxonomy.py index b8ec08194..1a3528923 100644 --- a/tests/test_provider_error_taxonomy.py +++ b/tests/test_provider_error_taxonomy.py @@ -85,6 +85,8 @@ def test_safe_message_keeps_actionable_schema_diagnostics_without_payloads() -> ) == actionable for diagnostic in ( "messages=[{'role':'user','content':'customer secret'}]", + '"messages": [{"role":"user","content":"customer secret"}]', + "'content': 'customer secret'", "prompt=customer secret", "input: customer secret", ): From ccbbaf132322433b7ebf3977739a49f99614cd0a Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:53:09 +0900 Subject: [PATCH 57/63] fix(telemetry): close provider summary leaks Signed-off-by: Seongho Bae --- contextual_orchestrator/provider_errors.py | 2 +- contextual_orchestrator/telemetry.py | 25 ++++++------ tests/test_telemetry.py | 45 ++++++++++++++++++++-- 3 files changed, 53 insertions(+), 19 deletions(-) diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py index f142e2430..40c95fbdd 100644 --- a/contextual_orchestrator/provider_errors.py +++ b/contextual_orchestrator/provider_errors.py @@ -50,7 +50,7 @@ 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)\b|" - r"\b(?:messages?|content)[\"']?\s*(?:[:=]|\[|\{)" + r"(? Date: Tue, 1 Sep 2026 12:18:40 +0900 Subject: [PATCH 58/63] fix(telemetry): recognize prefixed schema failures Signed-off-by: Seongho Bae --- contextual_orchestrator/telemetry.py | 12 +++++++-- tests/test_telemetry.py | 37 ++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/telemetry.py b/contextual_orchestrator/telemetry.py index eac87cdd1..45a4a2b2e 100644 --- a/contextual_orchestrator/telemetry.py +++ b/contextual_orchestrator/telemetry.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import ipaddress import logging import re from collections.abc import Iterator, Mapping @@ -41,7 +42,7 @@ # 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['\"]?" + r"['\"]?messages['\"]? must contain the word ['\"]?json['\"]?" r"(?: in some form,)? to use " r"(?:['\"]?response_format['\"]? of type ['\"]?json_object['\"]?|json_object)" r"(?:\.|$)", @@ -194,6 +195,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 @@ -334,7 +342,7 @@ def traced( error_summary = ( _SAFE_SCHEMA_ERROR_SUMMARY if provider_summary is not None - and _SAFE_SCHEMA_DIAGNOSTIC.match(provider_summary) + and _SAFE_SCHEMA_DIAGNOSTIC.search(provider_summary) else failure_code ) model_group = safe.get("contextual_orchestrator.model_group", "ungrouped") diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index adc351514..24c8740b0 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -59,6 +59,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: @@ -552,6 +554,41 @@ def test_traced_logs_actionable_bounded_failure_evidence(monkeypatch, caplog): 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 From 073b87b165b9621c92e2971875143810ff6f85e9 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 12:29:48 +0900 Subject: [PATCH 59/63] fix(security): canonicalize provider schema diagnostics --- contextual_orchestrator/provider_errors.py | 13 +++++++++++-- contextual_orchestrator/telemetry.py | 5 ++--- tests/test_provider_error_taxonomy.py | 9 ++++++++- 3 files changed, 21 insertions(+), 6 deletions(-) diff --git a/contextual_orchestrator/provider_errors.py b/contextual_orchestrator/provider_errors.py index 40c95fbdd..3b862792d 100644 --- a/contextual_orchestrator/provider_errors.py +++ b/contextual_orchestrator/provider_errors.py @@ -49,10 +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)\b|" - r"(? ``(client_status, error_code, retryable)`` surface. #: The client status is what this gateway returns; ``error_code`` follows the @@ -149,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 diff --git a/contextual_orchestrator/telemetry.py b/contextual_orchestrator/telemetry.py index 45a4a2b2e..8a3894313 100644 --- a/contextual_orchestrator/telemetry.py +++ b/contextual_orchestrator/telemetry.py @@ -42,10 +42,9 @@ # 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['\"]?" - r"(?: in some form,)? to use " + r"messages (?:must contain the word ['\"]?json['\"]?(?: in some form,)? to use " r"(?:['\"]?response_format['\"]? of type ['\"]?json_object['\"]?|json_object)" - r"(?:\.|$)", + r"(?:\.|$)|must mention json when response_format is json_object)", re.IGNORECASE, ) _SAFE_SCHEMA_ERROR_SUMMARY = ( diff --git a/tests/test_provider_error_taxonomy.py b/tests/test_provider_error_taxonomy.py index 1a3528923..17f5ebd95 100644 --- a/tests/test_provider_error_taxonomy.py +++ b/tests/test_provider_error_taxonomy.py @@ -82,7 +82,7 @@ def test_safe_message_keeps_actionable_schema_diagnostics_without_payloads() -> actionable = "'messages' must contain the word 'json' to use json_object" assert safe_provider_message( _body_http_error(400, {"error": {"message": actionable}}) - ) == 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"}]', @@ -94,6 +94,13 @@ def test_safe_message_keeps_actionable_schema_diagnostics_without_payloads() -> _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.""" From b36c2f145a3dcab6dc1fb13a678364a916bb6b1e Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:02:19 +0900 Subject: [PATCH 60/63] fix(workflow): avoid duplicating caller prompt Signed-off-by: Seongho Bae (cherry picked from commit 97f722de4f6b4b0c3b366dc48a51c9e2ca89147d) Signed-off-by: Seongho Bae --- CHANGELOG.md | 3 +++ contextual_orchestrator/orchestrator.py | 19 ++------------ docs/adr/0002-control-plane-orchestrator.md | 5 +++- tests/test_generated_workflow.py | 28 ++++++++++++++++++++- 4 files changed, 36 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f11f18e0..47414083f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ 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, instructions, or 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; diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index fd4fc4938..e8febfa91 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -5721,11 +5721,6 @@ def conduct( task = self._latest_user_text(messages) source_images = self._source_image_parts(messages) required_tags = ("vision",) if source_images else () - caller_instructions = "\n\n".join( - message["content"] - for message in messages - if message.get("role") == "system" and isinstance(message.get("content"), str) - ) plan_source = "template" if model_name not in {self.GATEWAY_DEFAULT_MODEL, self.AUTO_MODEL}: steps = self._plan(task, model_name=model_name) @@ -5788,16 +5783,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", @@ -5805,13 +5791,12 @@ def conduct( f"Role: {step.role}\n" "Use only the original task and the accessed prior steps. " "Return concise, directly useful work." - + (f"\n\nCaller instructions:\n{caller_instructions}" if caller_instructions else "") ), }, *copy.deepcopy(messages), { "role": "user", - "content": user_content, + "content": instruction, }, ] start = time.perf_counter() diff --git a/docs/adr/0002-control-plane-orchestrator.md b/docs/adr/0002-control-plane-orchestrator.md index ac317f697..06b2ee93a 100644 --- a/docs/adr/0002-control-plane-orchestrator.md +++ b/docs/adr/0002-control-plane-orchestrator.md @@ -47,7 +47,10 @@ 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 worker envelope. The envelope does not repeat the current task, + caller instructions, or source attachments. 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/tests/test_generated_workflow.py b/tests/test_generated_workflow.py index e3ef8b99f..8f7e59f01 100644 --- a/tests/test_generated_workflow.py +++ b/tests/test_generated_workflow.py @@ -81,11 +81,37 @@ 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_caller_instructions_once() -> None: + """Workers receive one canonical copy of current task and caller policy.""" + 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("CALLER_POLICY_SENTINEL") == 1 + assert serialized.count("EARLIER_TURN_SENTINEL") == 1 + assert serialized.count("EARLIER_ANSWER_SENTINEL") == 1 + 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_planner_prompt_lists_the_agent_pool() -> None: orchestrator, client = _orch(json.dumps(PLAN)) orchestrator.conduct([{"role": "user", "content": "solve it"}]) From c5a0ae93f0b2307383119890ce189004cb52f3fc Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:14:22 +0900 Subject: [PATCH 61/63] fix(workflow): preserve caller instruction authority Signed-off-by: Seongho Bae (cherry picked from commit 69054734c167c5b99be656707ecc66100defb9bf) Signed-off-by: Seongho Bae --- CHANGELOG.md | 2 +- contextual_orchestrator/orchestrator.py | 6 ++++++ docs/adr/0002-control-plane-orchestrator.md | 6 ++++-- tests/test_generated_workflow.py | 6 +++--- 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 47414083f..41f6914af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,7 @@ and this project uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html) - 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, instructions, or attachments. + 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; diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index e8febfa91..ddd5be599 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -5721,6 +5721,11 @@ def conduct( task = self._latest_user_text(messages) source_images = self._source_image_parts(messages) required_tags = ("vision",) if source_images else () + caller_instructions = "\n\n".join( + message["content"] + for message in messages + if message.get("role") == "system" and isinstance(message.get("content"), str) + ) plan_source = "template" if model_name not in {self.GATEWAY_DEFAULT_MODEL, self.AUTO_MODEL}: steps = self._plan(task, model_name=model_name) @@ -5791,6 +5796,7 @@ def conduct( f"Role: {step.role}\n" "Use only the original task and the accessed prior steps. " "Return concise, directly useful work." + + (f"\n\nCaller instructions:\n{caller_instructions}" if caller_instructions else "") ), }, *copy.deepcopy(messages), diff --git a/docs/adr/0002-control-plane-orchestrator.md b/docs/adr/0002-control-plane-orchestrator.md index 06b2ee93a..d70f12410 100644 --- a/docs/adr/0002-control-plane-orchestrator.md +++ b/docs/adr/0002-control-plane-orchestrator.md @@ -49,8 +49,10 @@ trained Fugu, TRINITY, or Conductor clone. visibility, implemented as data on the step, not as a trained topology policy). The worker preserves the caller message array exactly once, then receives only its subtask and deliberately exposed prior outputs in the - added worker envelope. The envelope does not repeat the current task, - caller instructions, or source attachments. + 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/tests/test_generated_workflow.py b/tests/test_generated_workflow.py index 8f7e59f01..fb7aa2c10 100644 --- a/tests/test_generated_workflow.py +++ b/tests/test_generated_workflow.py @@ -87,8 +87,8 @@ def test_access_lists_actually_isolate_context() -> None: assert "step-output(2)" in step3_prompt and "step-output(3)" in step3_prompt -def test_workflow_step_carries_current_task_and_caller_instructions_once() -> None: - """Workers receive one canonical copy of current task and caller policy.""" +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( [ @@ -102,9 +102,9 @@ def test_workflow_step_carries_current_task_and_caller_instructions_once() -> No step_messages = client.calls[1] serialized = json.dumps(step_messages) assert serialized.count("CURRENT_TASK_SENTINEL") == 1 - assert serialized.count("CALLER_POLICY_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 From 2ea70dc0386b2f2417da64c07b0151c8a6fc8466 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:25:32 +0900 Subject: [PATCH 62/63] Preserve multipart workflow instructions (cherry picked from commit f6b4591e78034dc1dea94429977305444a18af30) Signed-off-by: Seongho Bae --- contextual_orchestrator/orchestrator.py | 5 +++-- tests/test_generated_workflow.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index ddd5be599..ec7c1de1c 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -5722,9 +5722,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}: diff --git a/tests/test_generated_workflow.py b/tests/test_generated_workflow.py index fb7aa2c10..7417ef626 100644 --- a/tests/test_generated_workflow.py +++ b/tests/test_generated_workflow.py @@ -112,6 +112,22 @@ def test_workflow_step_carries_current_task_and_source_context_once() -> None: 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"}]) From 858a13ed00fd87155d0cc48ec30d761af603c78a Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 12:46:45 +0900 Subject: [PATCH 63/63] fix(accounting): hide incomplete aggregate totals --- contextual_orchestrator/cost_ledger.py | 35 ++++++++++++++++++++++---- tests/test_cost_ledger.py | 33 ++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/contextual_orchestrator/cost_ledger.py b/contextual_orchestrator/cost_ledger.py index ef6195435..ab9ba59fc 100644 --- a/contextual_orchestrator/cost_ledger.py +++ b/contextual_orchestrator/cost_ledger.py @@ -1424,9 +1424,20 @@ def rollup( bucket["total_tokens"] += int(row.get("total_tokens", 0)) bucket["cost_amount"] += Decimal(str(row.get("cost_amount", 0))) for bucket in buckets.values(): - bucket["cost_amount"] = float( - bucket["cost_amount"].quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP) - ) + if bucket["unavailable_record_count"]: + for field in ( + "prompt_tokens", + "completion_tokens", + "total_tokens", + "cost_amount", + ): + bucket[field] = None + else: + bucket["cost_amount"] = float( + bucket["cost_amount"].quantize( + Decimal("0.000001"), rounding=ROUND_HALF_UP + ) + ) return buckets def report( @@ -1438,7 +1449,12 @@ def report( """Return a report envelope: per-value rollup plus a grand total.""" buckets = self.rollup(dimension, start, end) items = sorted( - buckets.values(), key=lambda item: item["cost_amount"], reverse=True + buckets.values(), + key=lambda item: ( + item["cost_amount"] is not None, + item["cost_amount"] or 0, + ), + reverse=True, ) grand_total = self.total(start, end) return { @@ -1454,7 +1470,7 @@ def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[ available = [row for row in rows if row.get("measurement_status") != "unavailable"] unavailable_count = len(rows) - len(available) cost = sum((Decimal(str(row.get("cost_amount", 0))) for row in available), Decimal("0")) - return { + totals: Dict[str, Any] = { "record_count": len(rows), "prompt_tokens": sum(int(row.get("prompt_tokens", 0)) for row in available), "completion_tokens": sum(int(row.get("completion_tokens", 0)) for row in available), @@ -1469,6 +1485,15 @@ def total(self, start: Optional[int] = None, end: Optional[int] = None) -> Dict[ ), "unavailable_record_count": unavailable_count, } + if unavailable_count: + for field in ( + "prompt_tokens", + "completion_tokens", + "total_tokens", + "cost_amount", + ): + totals[field] = None + return totals def records(self, start: Optional[int] = None, end: Optional[int] = None) -> List[Dict[str, Any]]: """Return raw usage record rows in the optional window.""" diff --git a/tests/test_cost_ledger.py b/tests/test_cost_ledger.py index 709a30014..13b07d73e 100644 --- a/tests/test_cost_ledger.py +++ b/tests/test_cost_ledger.py @@ -355,6 +355,39 @@ def test_report_envelope_sorts_by_cost_desc_and_includes_grand_total() -> None: assert report["grand_total"]["cost_amount"] == 11.0 +@pytest.mark.parametrize("include_measured", [False, True]) +def test_unavailable_rollups_do_not_publish_partial_totals(include_measured: bool) -> None: + ledger = _priced_ledger() + if include_measured: + ledger.record_usage( + provider="openai", + model="gpt-x", + prompt_tokens=1000, + completion_tokens=1000, + attribution={"team": "alpha"}, + ) + ledger.record_usage( + provider="openai", + model="gpt-x", + prompt_tokens=0, + completion_tokens=0, + attribution={"team": "alpha"}, + measurement_status="unavailable", + ) + + report = ledger.report("team") + expected_count = 2 if include_measured else 1 + for aggregate in (report["items"][0], report["grand_total"]): + assert aggregate["record_count"] == expected_count + assert aggregate["unavailable_record_count"] == 1 + assert aggregate["measurement_status"] == "unavailable" + assert aggregate["prompt_tokens"] is None + assert aggregate["completion_tokens"] is None + assert aggregate["total_tokens"] is None + assert aggregate["cost_amount"] is None + assert report["items"][0]["currency_code"] == "USD" + + def test_sql_ledger_store_on_sqlite_creates_objects_and_rolls_up() -> None: conn = sqlite3.connect(":memory:") store = SqlLedgerStore(conn, paramstyle="qmark")