From 8ace8f402fca1f3c39e398471717aecb47fa67c4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Wed, 2 Sep 2026 17:20:22 +0900 Subject: [PATCH 1/3] fix(batch): gate real Batch API on declared batch_endpoint_supported TaskOrchestrator.batch_route() selects a worker agent for each prompt via _select_agent(), which filters only on chat-capability tags -- it can return any configured chat-capable provider (Anthropic-shaped, NVIDIA NIM, OpenRouter, a self-hosted OpenAI-compatible gateway, ...), many of which do not implement OpenAI's real async Batch API (/files, /batches, /files/{id}/content) at all. ModelClient.batch_chat()'s only gate before routing into that real endpoint was is_chat_compatible_model_id(), a model-id shape heuristic that says nothing about provider Batch API support -- so an unsupported provider's real HTTP call 404d and failed the entire batch group, with no fallback. Add ModelAgent.batch_endpoint_supported: bool | None (default None), mirroring the existing reasoning_effort_supported fail-closed tri-state field exactly: same __post_init__ validation, to_config/from_dict round-trip, and agent_pool SQLite persistence/migration. batch_chat() now only takes the real Batch API path when batch_endpoint_supported is True; every other remote agent falls back to the same per-item emulation _local_batch_chat already performs for local providers, aggregating into the identical result shape instead of either misrouting or hard-failing. pg-llm-batch (the real Batch API client contextual-orchestrator vendors) only validates the URL *path* submitted inside an already-created job -- it has no per-provider capability registry and, per its own module docstring, expects the orchestrator to know a target supports the Batch API before calling it. This closes that gap the same way this repo already closes it for reasoning_effort support, rather than inventing a new mechanism. Tests extend tests/test_batch_api.py's real-HTTP fake Batch/chat provider: - a batch_endpoint_supported=True agent still takes the real Batch API path - an unproven agent (None or False) falls back to emulation and returns a correctly aggregated result, via an actual per-request HTTP round trip through /chat/completions (not a stubbed return value) -- proving the real Batch API was never touched (zero /batches posts, no /files upload) Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.d/batch-endpoint-capability-gate.md | 1 + contextual_orchestrator/orchestrator.py | 60 +++++++++++-- tests/test_batch_api.py | 84 ++++++++++++++++++- 3 files changed, 137 insertions(+), 8 deletions(-) create mode 100644 CHANGELOG.d/batch-endpoint-capability-gate.md diff --git a/CHANGELOG.d/batch-endpoint-capability-gate.md b/CHANGELOG.d/batch-endpoint-capability-gate.md new file mode 100644 index 000000000..ce924207b --- /dev/null +++ b/CHANGELOG.d/batch-endpoint-capability-gate.md @@ -0,0 +1 @@ +Fixed `ModelClient.batch_chat()` routing any worker agent selected by `TaskOrchestrator.batch_route()` straight into the real provider Batch API (`/files`, `/batches`, `/files/{id}/content`), even when that provider never implements it — a 404 there failed the whole batch group with no fallback. Added `ModelAgent.batch_endpoint_supported: bool | None` (mirroring `reasoning_effort_supported`'s fail-closed pattern, including agent-pool persistence); the real Batch API is now only addressed when an operator has explicitly declared `batch_endpoint_supported: true`. Every other remote agent (`None`/`False`, the default) falls back to the same per-item emulation `_local_batch_chat` already performs for local providers. diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index a09c4e8f0..b36b995f2 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -593,6 +593,11 @@ class ModelAgent: # ``None`` means provider support is unproven. Opt-in effort profiles then # fail closed unless the profile explicitly requests the safe ``omit`` fallback. reasoning_effort_supported: bool | None = None + # ``None``/``False`` mean the provider's real async Batch API (/files, + # /batches, /files/{id}/content) is not proven to exist for this agent. + # batch_chat() then falls closed to per-item emulation through the normal + # chat endpoint rather than guessing a remote provider supports it. + batch_endpoint_supported: bool | None = None # Explicit reviewed replica contract. A group never races when this is absent. endpoint_equivalence: dict[str, Any] | None = None # Provider-declared support for the Chat Completions terminal usage frame. @@ -620,6 +625,8 @@ def __post_init__(self) -> None: ) if self.reasoning_effort_supported not in (None, True, False): raise TypeError("reasoning_effort_supported must be true, false, or null") + if self.batch_endpoint_supported not in (None, True, False): + raise TypeError("batch_endpoint_supported must be true, false, or null") if type(self.stream_usage_supported) is not bool: raise TypeError("stream_usage_supported must be a boolean") if self.endpoint_equivalence is not None: @@ -645,6 +652,7 @@ def to_config(self) -> dict[str, Any]: "context_window": self.context_window, "group_name": self.group_name, "reasoning_effort_supported": self.reasoning_effort_supported, + "batch_endpoint_supported": self.batch_endpoint_supported, "endpoint_equivalence": self.endpoint_equivalence, "stream_usage_supported": self.stream_usage_supported, } @@ -680,6 +688,7 @@ def from_dict(cls, value: dict[str, Any]) -> "ModelAgent": # pragma: no cover context_window=value.get("context_window"), group_name=value.get("group_name", ""), reasoning_effort_supported=value.get("reasoning_effort_supported"), + batch_endpoint_supported=value.get("batch_endpoint_supported"), endpoint_equivalence=value.get("endpoint_equivalence"), stream_usage_supported=value.get("stream_usage_supported", False), ) @@ -2843,6 +2852,19 @@ def batch_chat( ``requests`` maps a caller custom_id to its messages. Suited to eval/benchmark workloads (24h completion window, ~half the price); real-time chat should keep using ``chat``. The mock path answers synchronously so tests and local runs work. + + The real async Batch API (``/files``, ``/batches``, ``/files/{id}/content``) + is only ever addressed when ``agent.batch_endpoint_supported is True`` -- + an explicit, operator-declared capability, mirroring how + ``reasoning_effort_supported`` gates ``reasoning_effort``. A worker agent + can be selected for any configured chat-capable provider (Anthropic-shaped, + OpenRouter, a self-hosted OpenAI-compatible gateway, ...) many of which do + not implement ``/v1/batches`` at all; guessing support from the model id or + transport shape alone previously sent those requests straight into a real + HTTP call that 404s and fails the whole batch. An unproven agent + (``None``/``False``) instead falls back to the same per-item emulation + ``_local_batch_chat`` already performs for local providers -- looping each + request through ``chat()`` and aggregating into the identical result shape. """ if not is_chat_compatible_model_id(agent.model): raise ValueError( @@ -2855,7 +2877,7 @@ def batch_chat( } elif _is_local_provider_url(agent.base_url): results = self._local_batch_chat(agent, requests, temperature, effort_profile) - else: + elif agent.batch_endpoint_supported is True: destination = self._validate_provider(agent) # pragma: no cover batch_error: ProviderUpstreamError | None = None try: @@ -2872,6 +2894,11 @@ def batch_chat( ) if batch_error is not None: raise batch_error + else: + # Not proven to support the real Batch API: emulate rather than + # either misroute to an endpoint that likely 404s, or hard-fail the + # whole group for a provider that was never asked to prove itself. + results = self._local_batch_chat(agent, requests, temperature, effort_profile) return _validate_batch_results(requests, results) def _local_batch_chat( @@ -2881,7 +2908,11 @@ def _local_batch_chat( temperature: float | None, effort_profile: ReasoningEffortProfile | None = None, ) -> dict[str, dict[str, Any]]: - """Run local OpenAI-compatible requests concurrently through mlx-lm.""" + """Emulate a batch by running each request concurrently through ``chat()``. + + Used for local (mlx-lm) providers, and reused by ``batch_chat`` as the + fallback for any remote provider not proven to support the real Batch API. + """ request_settings = self.request_settings_snapshot() def complete(custom_id: str, messages: list[ChatMessage]) -> tuple[str, dict[str, Any]]: @@ -3110,6 +3141,7 @@ class _AgentPoolStore: "max_output_tokens", "context_window", "reasoning_effort_supported", + "batch_endpoint_supported", "stream_usage_supported", } ) @@ -3149,6 +3181,7 @@ def _create_normalized_schema(cls, conn: sqlite3.Connection) -> None: max_output_tokens INTEGER, context_window INTEGER, reasoning_effort_supported INTEGER, + batch_endpoint_supported INTEGER, stream_usage_supported INTEGER NOT NULL DEFAULT 0, CONSTRAINT agent_pool_disabled_flag_check CHECK (disabled IN (0, 1)), CONSTRAINT agent_pool_max_output_tokens_check @@ -3169,6 +3202,8 @@ def _create_normalized_schema(cls, conn: sqlite3.Connection) -> None: ), CONSTRAINT agent_pool_reasoning_effort_flag_check CHECK (reasoning_effort_supported IS NULL OR reasoning_effort_supported IN (0, 1)), + CONSTRAINT agent_pool_batch_endpoint_flag_check + CHECK (batch_endpoint_supported IS NULL OR batch_endpoint_supported IN (0, 1)), CONSTRAINT agent_pool_stream_usage_flag_check CHECK (stream_usage_supported IN (0, 1)) ) @@ -3209,8 +3244,9 @@ def _insert_agent(cls, conn: sqlite3.Connection, agent: "ModelAgent") -> None: INSERT INTO agent_pool ( agent_id, model_name, base_url, api_key_env, credential_key, priority, disabled, provider_name, local_credential_key, auth_scheme, - max_output_tokens, context_window, reasoning_effort_supported, stream_usage_supported - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + max_output_tokens, context_window, reasoning_effort_supported, + batch_endpoint_supported, stream_usage_supported + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( config["id"], @@ -3226,6 +3262,7 @@ def _insert_agent(cls, conn: sqlite3.Connection, agent: "ModelAgent") -> None: config["max_output_tokens"], config["context_window"], config["reasoning_effort_supported"], + config["batch_endpoint_supported"], int(config["stream_usage_supported"]), ), ) @@ -3278,6 +3315,12 @@ def _initialize_schema(cls, conn: sqlite3.Connection) -> None: "CHECK (reasoning_effort_supported IS NULL OR reasoning_effort_supported IN (0, 1))" ) columns.add("reasoning_effort_supported") + if "batch_endpoint_supported" not in columns: + conn.execute( + "ALTER TABLE agent_pool ADD COLUMN batch_endpoint_supported INTEGER " + "CHECK (batch_endpoint_supported IS NULL OR batch_endpoint_supported IN (0, 1))" + ) + columns.add("batch_endpoint_supported") if "max_output_tokens" not in columns: conn.execute( "ALTER TABLE agent_pool ADD COLUMN max_output_tokens INTEGER " @@ -3393,7 +3436,8 @@ def save(self, agent: "ModelAgent") -> None: priority = ?, disabled = ?, provider_name = ?, local_credential_key = ?, auth_scheme = ?, max_output_tokens = ?, context_window = ?, - reasoning_effort_supported = ?, stream_usage_supported = ? + reasoning_effort_supported = ?, batch_endpoint_supported = ?, + stream_usage_supported = ? WHERE agent_id = ? """, ( @@ -3409,6 +3453,7 @@ def save(self, agent: "ModelAgent") -> None: config["max_output_tokens"], config["context_window"], config["reasoning_effort_supported"], + config["batch_endpoint_supported"], int(config["stream_usage_supported"]), agent.id, ), @@ -3506,7 +3551,7 @@ def load_all(self) -> list["ModelAgent"]: SELECT agent_id, model_name, base_url, api_key_env, credential_key, priority, disabled, provider_name, local_credential_key, auth_scheme, max_output_tokens, context_window, - reasoning_effort_supported, stream_usage_supported + reasoning_effort_supported, batch_endpoint_supported, stream_usage_supported FROM agent_pool ORDER BY agent_id """ ).fetchall() @@ -3571,7 +3616,8 @@ def load_all(self) -> list["ModelAgent"]: max_output_tokens=row[10], context_window=row[11], reasoning_effort_supported=(None if row[12] is None else bool(row[12])), - stream_usage_supported=bool(row[13]), + batch_endpoint_supported=(None if row[13] is None else bool(row[13])), + stream_usage_supported=bool(row[14]), group_name=group_by_agent.get(row[0], ""), endpoint_equivalence=contract_by_agent.get(row[0]), ) diff --git a/tests/test_batch_api.py b/tests/test_batch_api.py index 57187ce04..32f44cc78 100644 --- a/tests/test_batch_api.py +++ b/tests/test_batch_api.py @@ -4,6 +4,12 @@ workloads. These drive ModelClient._batch_run against a local fake Batch server: multipart JSONL upload, batch creation, in_progress -> completed polling, JSONL result parsing with usage, and terminal-failure handling. Mock path stays sync. + +Also covers batch_chat()'s own real-vs-emulated routing gate: only an agent +explicitly declaring ``batch_endpoint_supported=True`` may reach the real +Batch API above; every other remote agent (unproven -- ``None``/``False``) +must fall back to the same per-item emulation local providers already use, +never guess-route into a provider that may not implement /v1/batches at all. """ from __future__ import annotations @@ -11,8 +17,10 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json from pathlib import Path +import socket import sys import threading +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -21,12 +29,14 @@ class _FakeBatchProvider: - """Implements /files, /batches, /batches/{id}, /files/{id}/content.""" + """Implements /files, /batches, /batches/{id}, /files/{id}/content, /chat/completions.""" def __init__(self, fail_status: str | None = None, polls_before_done: int = 2) -> None: outer = self self.uploaded_jsonl: bytes = b"" self.poll_count = 0 + self.batches_post_count = 0 + self.chat_completions_count = 0 class Handler(BaseHTTPRequestHandler): def _json(self, payload: dict, status: int = 200) -> None: @@ -44,7 +54,21 @@ def do_POST(self) -> None: # noqa: N802 outer.uploaded_jsonl = body # multipart wrapper included; checked via 'in' self._json({"id": "file_input_1"}) elif self.path == "/batches": + outer.batches_post_count += 1 self._json({"id": "batch_1", "status": "validating"}) + elif self.path == "/chat/completions": + # The emulation fallback's target: one ordinary chat completion + # per batch item, echoing the prompt so per-item routing is provable. + outer.chat_completions_count += 1 + payload = json.loads(body) + last_user = next( + (m["content"] for m in reversed(payload["messages"]) if m.get("role") == "user"), + "", + ) + self._json({ + "choices": [{"message": {"role": "assistant", "content": f"emulated: {last_user}"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }) else: self._json({"error": "not found"}, 404) @@ -150,6 +174,64 @@ def test_mock_path_answers_synchronously() -> None: assert results["task_a"]["usage"] is None +def _remote_agent(base_url: str, **overrides: object) -> ModelAgent: + # credential_key="" keeps chat()'s NotConfigured pre-check out of the way + # (no KV credential is registered in tests); base_url is real HTTP so + # _validate_provider's https/public-IP checks are bypassed per-test below, + # exactly like the destination-pinning pattern already proven in + # test_provider_integration.py::test_open_provider_uses_validated_destination_without_dns_relookup. + fields: dict[object, object] = { + "id": "worker_agent", + "model": "gpt-x", + "base_url": base_url, + "credential_key": "", + } + fields.update(overrides) + return ModelAgent(**fields) # type: ignore[arg-type] + + +def _pinned_destination(provider: "_FakeBatchProvider") -> tuple[int, tuple[str, int]]: + return (socket.AF_INET, ("127.0.0.1", provider._server.server_address[1])) + + +def test_batch_chat_routes_batch_capable_agent_to_real_batch_endpoint() -> None: + """(a) batch_endpoint_supported=True still takes the real Batch API path.""" + with _FakeBatchProvider(polls_before_done=2) as provider: + agent = _remote_agent(provider.base_url, batch_endpoint_supported=True) + client = _client() + with patch.object(client, "_validate_provider", return_value=_pinned_destination(provider)): + results = client.batch_chat(agent, REQUESTS, poll_interval=0.01, poll_timeout=30) + + assert results["task_a"]["content"] == "answer A" + assert results["task_b"]["usage"]["completion_tokens"] == 2 + assert provider.batches_post_count == 1 # the real Batch API was used + assert provider.chat_completions_count == 0 # emulation never triggered + + +def test_batch_chat_falls_back_to_emulation_when_batch_endpoint_unproven() -> None: + """(b) an unproven agent (None/False) still returns a correct aggregated + result, via emulation -- and (c) that emulation path is real, not stubbed: + every request actually crosses the wire as one /chat/completions call. + """ + for unsupported_value in (None, False): + with _FakeBatchProvider() as provider: + agent = _remote_agent(provider.base_url, batch_endpoint_supported=unsupported_value) + client = _client() + with patch.object(client, "_validate_provider", return_value=_pinned_destination(provider)): + results = client.batch_chat(agent, REQUESTS) + + assert set(results) == {"task_a", "task_b"} + assert results["task_a"]["content"] == "emulated: question A" + assert results["task_b"]["content"] == "emulated: question B" + assert results["task_a"]["usage"]["completion_tokens"] == 1 + # One real chat completion per request -- the emulation path was + # actually exercised, not just asserted from a mocked return value. + assert provider.chat_completions_count == 2 + # The real Batch API was never touched: no misroute, no hard failure. + assert provider.batches_post_count == 0 + assert provider.uploaded_jsonl == b"" + + if __name__ == "__main__": for name, fn in sorted(globals().items()): if name.startswith("test_") and callable(fn): From ecac975d8bff2090027ca13b44f94c020548466b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 2 Sep 2026 18:35:30 +0000 Subject: [PATCH 2/3] fix(tests): declare batch_endpoint_supported for real-Batch-API test agents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit batch_chat() now falls closed to per-item emulation unless ModelAgent.batch_endpoint_supported is explicitly True. Two boundary tests constructed remote HTTPS agents to exercise the real Batch API path directly and broke under the new fail-closed default; declare batch_endpoint_supported=True on those agents to keep exercising the real-Batch-API code path they were written for. Verified: full suite green for this change — 3350 passed, 2 skipped, 3 failed (test_admin_contract.py::test_model_group_mutations_refresh_audit_events, test_psychometric_routing.py::test_fast_mlsirm_fit_uses_judge_acceptance_item_for_context_score, test_spend_analytics.py::test_exact_output_without_prompt_usage_is_explicitly_unavailable), all three pre-existing and identical on origin/main, untouched by this branch's diff. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01BV96rXhqoR3tYZ9AeAVur4 --- tests/test_orchestrator_client_boundaries.py | 2 ++ tests/test_provider_reliability.py | 7 ++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_orchestrator_client_boundaries.py b/tests/test_orchestrator_client_boundaries.py index 9bb489b41..33a1b6dde 100644 --- a/tests/test_orchestrator_client_boundaries.py +++ b/tests/test_orchestrator_client_boundaries.py @@ -556,6 +556,7 @@ def test_batch_chat_success_on_https_provider_returns_validated_results() -> Non model="remote-chat-model", base_url="https://remote.example/v1", credential_key="REMOTE_API_KEY", + batch_endpoint_supported=True, ) client = ModelClient() requests = {"task_0": [{"role": "user", "content": "hi"}]} @@ -575,6 +576,7 @@ def test_batch_chat_wraps_provider_failures_without_provider_text() -> None: model="remote-chat-model", base_url="https://remote.example/v1", credential_key="REMOTE_API_KEY", + batch_endpoint_supported=True, ) client = ModelClient() requests = {"task_0": [{"role": "user", "content": "hi"}]} diff --git a/tests/test_provider_reliability.py b/tests/test_provider_reliability.py index fc816b937..1f1813f66 100644 --- a/tests/test_provider_reliability.py +++ b/tests/test_provider_reliability.py @@ -612,7 +612,12 @@ def _batch_run(self, agent, requests, temperature, poll_interval, poll_timeout, raise RuntimeError("provider-secret-batch-body") client = RawBatchFailureClient() - agent = ModelAgent("batch_worker", "gpt", base_url="https://provider.example/v1") + agent = ModelAgent( + "batch_worker", + "gpt", + base_url="https://provider.example/v1", + batch_endpoint_supported=True, + ) try: client.batch_chat(agent, {"task_0": [{"role": "user", "content": "ping"}]}) except RuntimeError as error: From 4e36447431322a1aa7570dcdb37798177e7ded8c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Fri, 4 Sep 2026 16:54:27 +0900 Subject: [PATCH 3/3] fix(batch): expose endpoint capability in admin API --- CHANGELOG.d/batch-endpoint-capability-gate.md | 2 + contextual_orchestrator/api_contract.py | 6 +++ contextual_orchestrator/orchestrator.py | 7 +++- contextual_orchestrator/server.py | 3 +- tests/test_agent_pool_db.py | 37 +++++++++++++++++++ tests/test_api_contract.py | 3 ++ tests/test_orchestrator_client_boundaries.py | 7 ++++ 7 files changed, 63 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.d/batch-endpoint-capability-gate.md b/CHANGELOG.d/batch-endpoint-capability-gate.md index ce924207b..dc18aa5b6 100644 --- a/CHANGELOG.d/batch-endpoint-capability-gate.md +++ b/CHANGELOG.d/batch-endpoint-capability-gate.md @@ -1 +1,3 @@ Fixed `ModelClient.batch_chat()` routing any worker agent selected by `TaskOrchestrator.batch_route()` straight into the real provider Batch API (`/files`, `/batches`, `/files/{id}/content`), even when that provider never implements it — a 404 there failed the whole batch group with no fallback. Added `ModelAgent.batch_endpoint_supported: bool | None` (mirroring `reasoning_effort_supported`'s fail-closed pattern, including agent-pool persistence); the real Batch API is now only addressed when an operator has explicitly declared `batch_endpoint_supported: true`. Every other remote agent (`None`/`False`, the default) falls back to the same per-item emulation `_local_batch_chat` already performs for local providers. + +Evidence: OpenAI documents Batch as a separate asynchronous contract that uploads a purpose-`batch` JSONL file and creates a `/v1/batches` job for one declared endpoint and completion window; OpenAI-compatible chat alone therefore does not prove Batch support ([Create batch](https://developers.openai.com/api/reference/resources/batches/methods/create), [Upload file](https://developers.openai.com/api/reference/resources/files/methods/create)). Yu et al., “Orca: A Distributed Serving System for Transformer-Based Generative Models,” OSDI 2022, describes selective batching as an explicit serving-system mechanism rather than a universal model property ([paper](https://www.usenix.org/conference/osdi22/presentation/yu)). Together these support an explicit per-provider capability gate while preserving synchronous emulation for unproven providers. diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index 03f89ba3a..83de16fae 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -523,6 +523,12 @@ ] }, "stream_usage_supported": {"type": "boolean"}, + "batch_endpoint_supported": { + "anyOf": [ + {"type": "boolean"}, + {"type": "null"}, + ] + }, }, }, }, diff --git a/contextual_orchestrator/orchestrator.py b/contextual_orchestrator/orchestrator.py index 23db7ae8b..e1d23a9c3 100644 --- a/contextual_orchestrator/orchestrator.py +++ b/contextual_orchestrator/orchestrator.py @@ -625,7 +625,7 @@ def __post_init__(self) -> None: ) if self.reasoning_effort_supported not in (None, True, False): raise TypeError("reasoning_effort_supported must be true, false, or null") - if self.batch_endpoint_supported not in (None, True, False): + if self.batch_endpoint_supported is not None and type(self.batch_endpoint_supported) is not bool: raise TypeError("batch_endpoint_supported must be true, false, or null") if type(self.stream_usage_supported) is not bool: raise TypeError("stream_usage_supported must be a boolean") @@ -6035,6 +6035,10 @@ def patch_agent(self, agent_pool_id: str, worker_agent_id: str, patch: dict[str, patched = replace( patched, stream_usage_supported=patch["stream_usage_supported"] ) + if "batch_endpoint_supported" in patch: + patched = replace( + patched, batch_endpoint_supported=patch["batch_endpoint_supported"] + ) updated_candidates = [patched if agent.id == worker_agent_id else agent for agent in self.candidates] updated_agents = [agent for agent in updated_candidates if not agent.disabled] @@ -8417,6 +8421,7 @@ def _agent_to_admin_payload(self, agent: ModelAgent) -> dict[str, Any]: "max_output_tokens": agent.max_output_tokens, "context_window": agent.context_window, "stream_usage_supported": agent.stream_usage_supported, + "batch_endpoint_supported": agent.batch_endpoint_supported, "group_name": agent.group_name, "group_routing": self._group_router.member_report(agent.id) if agent.group_name else None, } diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index eb6a77519..01677387a 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -283,7 +283,7 @@ def server_close(self) -> None: ALLOWED_AGENT_PATCH_KEYS = { "status", "priority", "tags", "provider_exclusions", "group_name", "endpoint_equivalence", "stream_usage_supported", "max_output_tokens", - "context_window", + "context_window", "batch_endpoint_supported", } ALLOWED_AGENT_CREATE_KEYS = { "id", @@ -301,6 +301,7 @@ def server_close(self) -> None: "stream_usage_supported", "max_output_tokens", "context_window", + "batch_endpoint_supported", } ALLOWED_MODEL_GROUP_KEYS = {"group_name", "member_agent_ids"} ALLOWED_MODEL_GROUP_PATCH_KEYS = {"member_agent_ids"} diff --git a/tests/test_agent_pool_db.py b/tests/test_agent_pool_db.py index 78b556392..cab6a80a8 100644 --- a/tests/test_agent_pool_db.py +++ b/tests/test_agent_pool_db.py @@ -39,6 +39,7 @@ def _seed() -> list[ModelAgent]: "max_output_tokens": 4096, "context_window": 128000, "stream_usage_supported": True, + "batch_endpoint_supported": True, } @@ -155,6 +156,21 @@ def test_stream_usage_capability_patch_survives_restart() -> None: assert restored._agent(agent.id).stream_usage_supported is True +def test_batch_endpoint_capability_patch_survives_restart() -> None: + with tempfile.TemporaryDirectory() as directory: + db = os.path.join(directory, "pool.db") + agent = ModelAgent("persisted_agent", "model-x") + first = TaskOrchestrator([agent], agents_db=db) + + updated = first.patch_agent( + "default", "persisted_agent", {"batch_endpoint_supported": True} + ) + assert updated["batch_endpoint_supported"] is True + + restored = TaskOrchestrator([agent], agents_db=db) + assert restored._agent(agent.id).batch_endpoint_supported is True + + def test_agent_pool_persists_limit_metadata_across_restart() -> None: with tempfile.TemporaryDirectory() as directory: db = os.path.join(directory, "pool.db") @@ -461,12 +477,21 @@ def test_http_create_and_delete_worker_agents() -> None: thread.start() base = f"http://127.0.0.1:{server.server_address[1]}/api/v1/agent_pools/default/worker_agents" try: + status, invalid_create = _call( + base, + "POST", + token, + {**NEW_AGENT, "id": "numeric_batch_agent", "batch_endpoint_supported": 1}, + ) + assert status == 400 and invalid_create["error"]["code"] == "invalid_request" + status, created = _call(base, "POST", token, NEW_AGENT) assert ( status == 201 and created["id"] == "coding_agent" and created["status"] == "active" and created["stream_usage_supported"] is True + and created["batch_endpoint_supported"] is True ) status, patched = _call( @@ -474,6 +499,17 @@ def test_http_create_and_delete_worker_agents() -> None: ) assert status == 200 and patched["stream_usage_supported"] is True + status, patched = _call( + f"{base}/general_agent", "PATCH", token, {"batch_endpoint_supported": True} + ) + assert status == 200 and patched["batch_endpoint_supported"] is True + + status, invalid_batch_capability = _call( + f"{base}/general_agent", "PATCH", token, {"batch_endpoint_supported": 1} + ) + assert status == 400 + assert invalid_batch_capability["error"]["code"] == "invalid_request" + status, invalid_capability = _call( f"{base}/general_agent", "PATCH", token, {"stream_usage_supported": 1} ) @@ -481,6 +517,7 @@ def test_http_create_and_delete_worker_agents() -> None: status, read = _call(f"{base}/general_agent", "GET", token) assert status == 200 and read["stream_usage_supported"] is True + assert read["batch_endpoint_supported"] is True status, dup = _call(base, "POST", token, NEW_AGENT) assert status == 400 # duplicate rejected diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index 7beb3d698..3cebcf1ac 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -81,6 +81,9 @@ def test_openapi_documents_compatibility_front_door() -> None: "/api/v1/agent_pools/{agent_pool_id}/worker_agents/{worker_agent_id}" ]["patch"]["requestBody"]["content"]["application/json"]["schema"] assert patch_schema["properties"]["stream_usage_supported"]["type"] == "boolean" + assert patch_schema["properties"]["batch_endpoint_supported"] == { + "anyOf": [{"type": "boolean"}, {"type": "null"}] + } assert patch_schema["properties"]["max_output_tokens"]["anyOf"] == [ { "type": "integer", diff --git a/tests/test_orchestrator_client_boundaries.py b/tests/test_orchestrator_client_boundaries.py index 33a1b6dde..18d226322 100644 --- a/tests/test_orchestrator_client_boundaries.py +++ b/tests/test_orchestrator_client_boundaries.py @@ -193,6 +193,13 @@ def test_model_agent_rejects_bad_local_credential_key_and_effort_flag() -> None: ModelAgent(id="agent_two", model="m", local_credential_key=123) # type: ignore[arg-type] with pytest.raises(TypeError, match="reasoning_effort_supported must be"): ModelAgent(id="agent_two", model="m", reasoning_effort_supported="yes") # type: ignore[arg-type] + for invalid_batch_flag in (0, 1, "true"): + with pytest.raises(TypeError, match="batch_endpoint_supported must be"): + ModelAgent( + id="agent_two", + model="m", + batch_endpoint_supported=invalid_batch_flag, # type: ignore[arg-type] + ) with pytest.raises(TypeError, match="max_output_tokens must be"): ModelAgent(id="agent_two", model="m", max_output_tokens=0) with pytest.raises(TypeError, match="max_output_tokens must be"):