From a9f88586fdd3a873ecb1ae82146aa4fdc6bd76e4 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 06:32:52 +0900 Subject: [PATCH 1/5] docs(security): align issue 117 status with runtime --- .../adrs/0026-trace-purpose-authorization.md | 10 ++++++---- docs/product-technical-gap-baseline.md | 2 +- tests/test_issue117_traceability.py | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 5 deletions(-) create mode 100644 tests/test_issue117_traceability.py diff --git a/docs/planning/adrs/0026-trace-purpose-authorization.md b/docs/planning/adrs/0026-trace-purpose-authorization.md index dc3c812f8..7b75e1f17 100644 --- a/docs/planning/adrs/0026-trace-purpose-authorization.md +++ b/docs/planning/adrs/0026-trace-purpose-authorization.md @@ -50,10 +50,12 @@ closed for the trace purpose because it has no verified trace claim. contains no prompt, output, credential, or PII value; - audit failure returns a generic `503` instead of releasing the trace. -Issue #117 remains open: batch jobs need their protected-main integration, and -the authorization adapter still needs tenant/resource/lifetime context. The -legacy single-token production migration is now guarded in the CLI and -canonical Compose path; its explicit local escape hatch remains documented. +Issue #117 remains open only for richer tenant/workspace/resource/lifetime and +revocation context supplied by the external authorization adapter. Protected +main already owner-binds batch jobs and workflow/evaluation records, gates +batch-result trace release, and rejects legacy single-token mode in production +CLI and canonical Compose paths. The explicit local single-token escape hatch +remains documented and must not be presented as production trace authority. ## Research grounding diff --git a/docs/product-technical-gap-baseline.md b/docs/product-technical-gap-baseline.md index cbad42f7f..c7056c5ed 100644 --- a/docs/product-technical-gap-baseline.md +++ b/docs/product-technical-gap-baseline.md @@ -1885,7 +1885,7 @@ guarantees. | [#123](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/123) | A sole collaborator can be unable to satisfy last-push approval. | Add governance evidence/runbook or a protected-rule-compatible process; never bypass approval. | | [#119](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/119) | Ambiguous or unbounded inbound framing threatens request integrity. | The #776/#783 implementation stack is merged into non-main branches; protected-main integration still requires exact-head hosted evidence and independent approval. | | [#118](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/118) | Liveness and authenticated readiness are not yet fully separated. | PR [#780](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/780) implements the minimal `/healthz` and authenticated `/readyz` contract; merge only after exact-head Checks and independent approval. | -| [#117](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/117) | Trace access and inference access need separate authority. | PR [#780](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/780) merged the minimal liveness/readiness and trace-authority slice; the issue remains open for batch ownership, full purpose/tenant/resource/lifetime/revocation context, and the single-token migration gate. | +| [#117](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/117) | Trace access and inference access need separate authority. | Protected main now enforces separate trace purpose, owner-bound batch/workflow/evaluation reads, strict trace-flag validation, audit-before-release, and a production gate against legacy single-token mode. The issue remains open only for richer tenant/workspace/resource/lifetime/revocation claims supplied by the external identity adapter; core code must not invent those deployment policies. | | [#116](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/116) | Browser admin sessions need separation from long-lived bearer credentials. | PR [#788](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/788) implements opaque bounded sessions, Secure-by-default cookies, same-origin state-change checks, logout/revocation, and regression evidence. | | [#103](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/103) | Release readiness must fail closed on stale head, missing review, or missing Checks evidence. | PR [#784](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/784) merged the semantic split, but the issue remains open until a trusted `.github` producer artifact and consumer verification bind complete exact-head policy evidence; caller-supplied dictionaries remain insufficient. | | [#899](https://github.com/ContextualWisdomLab/contextual-orchestrator/issues/899) | CEFR writing/speaking observations need a governed, evidence-bound orchestration boundary. | PR [#903](https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/903) is the current implementation head; validate its exact protected checks, research traceability, and independent approval before delivery. | diff --git a/tests/test_issue117_traceability.py b/tests/test_issue117_traceability.py new file mode 100644 index 000000000..1ca442e58 --- /dev/null +++ b/tests/test_issue117_traceability.py @@ -0,0 +1,15 @@ +"""Issue #117 canonical status remains aligned with protected-main behavior.""" + +from pathlib import Path + + +def test_issue117_docs_do_not_reopen_integrated_security_slices() -> None: + """Batch ownership and production migration stay recorded as implemented.""" + baseline = Path("docs/product-technical-gap-baseline.md").read_text() + adr = Path("docs/planning/adrs/0026-trace-purpose-authorization.md").read_text() + + assert "owner-bound batch/workflow/evaluation reads" in baseline + assert "production gate against legacy single-token mode" in baseline + assert "remains open only for richer tenant/workspace/resource/lifetime" in adr + assert "Protected\nmain already owner-binds batch jobs" in adr + assert "batch jobs need their protected-main integration" not in adr From 47de3a70f6ba7138d6789c24c21e799d5c2cbb17 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 06:39:00 +0900 Subject: [PATCH 2/5] fix(security): make batch trace retrieval opt-in --- contextual_orchestrator/api_contract.py | 20 ++++- contextual_orchestrator/server.py | 21 +++-- docs/doctoring/trace-purpose-authorization.md | 11 +-- .../adrs/0026-trace-purpose-authorization.md | 10 ++- tests/test_api_contract.py | 16 +++- tests/test_cost_review_server.py | 82 +++++++++++++++++++ 6 files changed, 142 insertions(+), 18 deletions(-) diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index 2da5c8fea..2a329dba7 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -904,10 +904,28 @@ "post": { "operationId": "create_batch_routing_job_results", "summary": "Retrieve principal-owned batch results and record their usage + cost", - "security": [{"inference_bearer_auth": [], "trace_bearer_auth": []}], + "security": [{"inference_bearer_auth": []}], "parameters": [ {"name": "batch_routing_job_id", "in": "path", "required": True, "schema": {"type": "string"}} ], + "requestBody": { + "required": False, + "content": { + "application/json": { + "schema": { + "type": "object", + "additionalProperties": False, + "properties": { + "include_orchestration_trace": { + "type": "boolean", + "default": False, + "description": "When true, the same caller must also have the trace purpose.", + } + }, + } + } + }, + }, "responses": { "200": {"description": "Batch results with recorded usage"}, "404": { diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index 4ea2f1d8b..fc3e0f29c 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -202,6 +202,7 @@ class ResponsiveThreadingHTTPServer(ThreadingHTTPServer): "previous_response_id", "conversation", "truncation", "include", "text", } | OPENAI_PASSTHROUGH_PARAM_KEYS ALLOWED_BATCH_KEYS = {"requests", "attribution", "routing", "model", "zdr_only"} +ALLOWED_BATCH_RESULTS_KEYS = {"include_orchestration_trace"} ALLOWED_EMBEDDINGS_BATCH_KEYS = {"model", "input", "inputs", "endpoint", "metadata", "attribution", "user", "encoding_format", "dimensions", "routing", "zdr_only"} ALLOWED_EMBEDDINGS_KEYS = { "model", "input", "encoding_format", "dimensions", "user", "metadata", "attribution", "routing", "zdr_only", @@ -7170,7 +7171,10 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di return if path.startswith("/api/v1/batch_routing_jobs/") and path.endswith("/results"): job_id = path[len("/api/v1/batch_routing_jobs/"):-len("/results")] - self._authorize_trace_access() + _reject_unknown_keys(body, ALLOWED_BATCH_RESULTS_KEYS) + include_trace = self._validate_trace_request(body, default=False) + if include_trace: + self._authorize_trace_access() try: retrieved = self._run( lambda: coordinator.retrieve_batch( @@ -7180,8 +7184,11 @@ def register_video_job(agent: ModelAgent, provider_result: dict[str, Any]) -> di except KeyError: self._send_error(404, "batch_job_not_found", f"batch job {job_id} not found") return - self._audit_trace_disclosure("/api/v1/batch_routing_jobs/{job_id}/results") - self._send(_response_payload(retrieved, include_trace=True)) + if include_trace: + self._audit_trace_disclosure( + "/api/v1/batch_routing_jobs/{job_id}/results" + ) + self._send(_response_payload(retrieved, include_trace=include_trace)) return if path == "/v1/responses": # The Responses API has no chat-completions verifier equivalent, @@ -7791,10 +7798,14 @@ def _audit_trace_disclosure(self, endpoint_path: str) -> None: "trace access audit is unavailable", ) from exc - def _validate_trace_request(self, body: dict[str, Any]) -> bool: + def _validate_trace_request( + self, body: dict[str, Any], *, default: bool | None = None + ) -> bool: """Validate whether the caller requested trace disclosure.""" if "include_orchestration_trace" not in body: - include_trace = security.expose_trace_by_default + include_trace = ( + security.expose_trace_by_default if default is None else default + ) elif type(body["include_orchestration_trace"]) is not bool: raise RequestError( 400, diff --git a/docs/doctoring/trace-purpose-authorization.md b/docs/doctoring/trace-purpose-authorization.md index dfad0142c..f9a898359 100644 --- a/docs/doctoring/trace-purpose-authorization.md +++ b/docs/doctoring/trace-purpose-authorization.md @@ -1,11 +1,12 @@ # Trace-purpose authorization Trace-bearing responses are privileged data-access requests. Chat, admin -simulation, workflow/evaluation creation, and batch-result retrieval first -authenticate the caller, then require a separate verified `trace` purpose -scope before returning a trace. Owner-facing workflow/evaluation reads apply -the same gate when trace exposure is enabled. A denied purpose returns `401` -and no trace body. +simulation, workflow/evaluation creation, and an explicit batch-result trace +request first authenticate the caller, then require a separate verified +`trace` purpose scope before returning a trace. Plain inference owners may +retrieve batch answers and cost evidence with traces stripped. Owner-facing +workflow/evaluation reads apply the same gate when trace exposure is enabled. +A denied explicit trace purpose returns `401` and no trace body. Access reports always require the `trace` purpose because accessed outputs are trace evidence. The authorization check runs before resource lookup, so a diff --git a/docs/planning/adrs/0026-trace-purpose-authorization.md b/docs/planning/adrs/0026-trace-purpose-authorization.md index 7b75e1f17..d9a4f1d05 100644 --- a/docs/planning/adrs/0026-trace-purpose-authorization.md +++ b/docs/planning/adrs/0026-trace-purpose-authorization.md @@ -16,10 +16,12 @@ affected_components: ## Decision Inference or admin authentication alone does not authorize a trace-bearing -response. Chat, admin simulation, workflow/evaluation creation, batch-result -retrieval, and trace-enabled workflow/evaluation reads require the verified -`trace` purpose scope. The server records a metadata-only audit event before -releasing the response and never trusts a caller-supplied purpose header. +response. Chat, admin simulation, workflow/evaluation creation, explicit +batch-result trace requests, and trace-enabled workflow/evaluation reads +require the verified `trace` purpose scope. Plain inference owners retrieve +batch answers and cost evidence with traces stripped. The server records a +metadata-only audit event before releasing a trace response and never trusts a +caller-supplied purpose header. Chat validates the trace flag before selecting structured, tool, streaming, or ordinary execution, so no early-return path can weaken the request contract. diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index e67f5902d..dadb640dd 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -68,9 +68,19 @@ def test_openapi_documents_compatibility_front_door() -> None: assert OPENAPI_SPEC["components"]["securitySchemes"]["trace_bearer_auth"]["scheme"] == ( "bearer" ) - assert OPENAPI_SPEC["paths"]["/api/v1/batch_routing_jobs/{batch_routing_job_id}/results"]["post"][ - "security" - ] == [{"inference_bearer_auth": [], "trace_bearer_auth": []}] + batch_results = OPENAPI_SPEC["paths"][ + "/api/v1/batch_routing_jobs/{batch_routing_job_id}/results" + ]["post"] + assert batch_results["security"] == [{"inference_bearer_auth": []}] + batch_results_schema = batch_results["requestBody"]["content"]["application/json"][ + "schema" + ] + assert batch_results_schema["additionalProperties"] is False + assert batch_results_schema["properties"]["include_orchestration_trace"] == { + "type": "boolean", + "default": False, + "description": "When true, the same caller must also have the trace purpose.", + } def test_openapi_documents_orchestrator_owned_embedding_model_selection() -> None: diff --git a/tests/test_cost_review_server.py b/tests/test_cost_review_server.py index 9efbb66af..feaffc8d4 100644 --- a/tests/test_cost_review_server.py +++ b/tests/test_cost_review_server.py @@ -404,6 +404,88 @@ def test_batch_routing_jobs_are_principal_bound_for_poll_and_results() -> None: server.shutdown() +def test_batch_results_require_explicit_trace_authority(monkeypatch) -> None: + """Plain inference retrieves answers while trace opt-in stays privileged.""" + inference_token = "batch_inference_owner" + trace_token = "batch_trace_reader" + security = SecurityConfig( + bearer_verifier=lambda token, scope: ( + token == inference_token and scope == "inference" + ) + or (token == trace_token and scope in {"inference", "trace"}), + principal_resolver=lambda token: ( + "tenant-a/owner-a" if token in {inference_token, trace_token} else None + ), + ) + agent = ModelAgent("mock_worker", "mock-a", base_url="mock://a") + orchestrator = TaskOrchestrator([agent]) + coordinator = CostRoutingCoordinator(orchestrator, InMemoryConfigStore()) + original_retrieve = coordinator.retrieve_batch + + def retrieve_with_trace(job_id, *, owner_id=None): + result = original_retrieve(job_id, owner_id=owner_id) + return {**result, "trace": [{"output": "intermediate"}]} + + monkeypatch.setattr(coordinator, "retrieve_batch", retrieve_with_trace) + server = build_server( + orchestrator, + port=0, + security=security, + coordinator=coordinator, + ) + threading.Thread(target=server.serve_forever, daemon=True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + try: + status, submitted = _request( + "POST", + f"{base}/api/v1/batch_routing_jobs", + inference_token, + {"requests": [{"messages": [{"role": "user", "content": "owned"}]}]}, + ) + assert status == 201 + results_url = f"{base}/api/v1/batch_routing_jobs/{submitted['job_id']}/results" + + status, plain = _request("POST", results_url, inference_token, {}) + assert status == 200 + assert plain["results"][0]["answer"] + assert "cost_amount" in plain["results"][0] + assert "trace" not in json.dumps(plain) + + for invalid in ("false", 1, None, [], {}): + status, body = _request( + "POST", + results_url, + inference_token, + {"include_orchestration_trace": invalid}, + ) + assert status == 400 + assert body["error"]["code"] == "invalid_include_orchestration_trace" + + status, denied = _request( + "POST", + results_url, + inference_token, + {"include_orchestration_trace": True}, + ) + assert status == 401 + assert denied["error"]["code"] == "unauthorized" + + status, traced = _request( + "POST", + results_url, + trace_token, + {"include_orchestration_trace": True}, + ) + assert status == 200 + assert traced["trace"] == [{"output": "intermediate"}] + assert any( + event["event_type"] == "orchestration_trace_access_granted" + for event in orchestrator.list_recent_audit_events() + ) + finally: + server.shutdown() + + def test_batch_routing_jobs_endpoint_submits_multiple_requests() -> None: server, port, token = _serve() base = f"http://127.0.0.1:{port}" From 69f5a4fe0e5c90c7bc4bdd7faff025a34147991c Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 06:44:54 +0900 Subject: [PATCH 3/5] fix(batch): preserve authorized result traces --- contextual_orchestrator/batch_routing.py | 3 +++ contextual_orchestrator/cost_router.py | 7 ++++--- contextual_orchestrator/server.py | 19 +++++++++++++++---- tests/test_cost_review_server.py | 22 +++++++++++++--------- 4 files changed, 35 insertions(+), 16 deletions(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 82ce36600..f47ec42d1 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -203,6 +203,7 @@ class BatchResultItem: attribution: Dict[str, Any] = field(default_factory=dict) model: str = "contextual-orchestrator" mode: str = "auto" + trace: Any = None class BatchBackend(Protocol): @@ -274,6 +275,7 @@ def run(request: BatchRequest) -> BatchResultItem: attribution=dict(request.attribution), model=request.model, mode=result.get("mode", request.mode), + trace=result.get("trace"), ) if self.max_concurrency == 1 or len(requests) <= 1: items = [run(request) for request in requests] @@ -414,6 +416,7 @@ async def _download() -> Dict[str, Any]: attribution=dict(request.attribution) if request else {}, model=request.model if request else "contextual-orchestrator", mode=request.mode if request else "auto", + trace=body.get("trace"), ) ) return items diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index 438a26fde..d54ccd5c5 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -725,8 +725,7 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict prompt_tokens=item.prompt_tokens or None, completion_tokens=item.completion_tokens or None, ) - recorded.append( - { + result_payload = { "custom_id": item.custom_id, "answer": item.answer, "usage_record_id": record.usage_record_id, @@ -736,7 +735,9 @@ def retrieve_batch(self, job_id: str, *, owner_id: Optional[str] = None) -> Dict "completion_tokens": record.completion_tokens, "measurement_status": record.measurement_status, } - ) + if item.trace is not None: + result_payload["trace"] = item.trace + recorded.append(result_payload) return { "job_id": job_id, "backend": job.backend, diff --git a/contextual_orchestrator/server.py b/contextual_orchestrator/server.py index fc3e0f29c..0311031d2 100644 --- a/contextual_orchestrator/server.py +++ b/contextual_orchestrator/server.py @@ -6316,7 +6316,11 @@ def do_POST(self) -> None: # noqa: N802 min(security.max_body_bytes, MAX_MULTIMODAL_JSON_BODY_BYTES) if large_inference_json else security.max_body_bytes - ) + ), + allow_empty_without_content_type=( + path.startswith("/api/v1/batch_routing_jobs/") + and path.endswith("/results") + ), ) zdr_only = _validate_zdr_only(body) request_policy = orchestrator.request_policy(zdr_only) @@ -7847,9 +7851,12 @@ def _parse_optional_int(self, query: dict[str, list[str]], field_name: str) -> i return None return int(raw) - def _read_json(self, *, max_body_bytes: int | None = None) -> dict[str, Any]: - if self.headers.get("content-type", "").split(";", 1)[0].strip().lower() != "application/json": - raise RequestError(415, "unsupported_media_type", "content-type must be application/json") + def _read_json( + self, + *, + max_body_bytes: int | None = None, + allow_empty_without_content_type: bool = False, + ) -> dict[str, Any]: try: body_size = _request_body_size( self.headers, @@ -7859,6 +7866,10 @@ def _read_json(self, *, max_body_bytes: int | None = None) -> dict[str, Any]: # Do not let a peer reuse a connection after an ambiguous frame. self.close_connection = True raise + if ( + body_size > 0 or not allow_empty_without_content_type + ) and self.headers.get("content-type", "").split(";", 1)[0].strip().lower() != "application/json": + raise RequestError(415, "unsupported_media_type", "content-type must be application/json") raw = self.rfile.read(body_size) if len(raw) != body_size: self.close_connection = True diff --git a/tests/test_cost_review_server.py b/tests/test_cost_review_server.py index feaffc8d4..7b4643985 100644 --- a/tests/test_cost_review_server.py +++ b/tests/test_cost_review_server.py @@ -404,7 +404,7 @@ def test_batch_routing_jobs_are_principal_bound_for_poll_and_results() -> None: server.shutdown() -def test_batch_results_require_explicit_trace_authority(monkeypatch) -> None: +def test_batch_results_require_explicit_trace_authority() -> None: """Plain inference retrieves answers while trace opt-in stays privileged.""" inference_token = "batch_inference_owner" trace_token = "batch_trace_reader" @@ -420,13 +420,6 @@ def test_batch_results_require_explicit_trace_authority(monkeypatch) -> None: agent = ModelAgent("mock_worker", "mock-a", base_url="mock://a") orchestrator = TaskOrchestrator([agent]) coordinator = CostRoutingCoordinator(orchestrator, InMemoryConfigStore()) - original_retrieve = coordinator.retrieve_batch - - def retrieve_with_trace(job_id, *, owner_id=None): - result = original_retrieve(job_id, owner_id=owner_id) - return {**result, "trace": [{"output": "intermediate"}]} - - monkeypatch.setattr(coordinator, "retrieve_batch", retrieve_with_trace) server = build_server( orchestrator, port=0, @@ -445,6 +438,17 @@ def retrieve_with_trace(job_id, *, owner_id=None): assert status == 201 results_url = f"{base}/api/v1/batch_routing_jobs/{submitted['job_id']}/results" + bodyless_request = urllib.request.Request( + results_url, + method="POST", + headers={"authorization": f"Bearer {inference_token}"}, + ) + with urllib.request.urlopen(bodyless_request) as response: + bodyless = json.loads(response.read()) + assert response.status == 200 + assert bodyless["results"][0]["answer"] + assert "trace" not in json.dumps(bodyless) + status, plain = _request("POST", results_url, inference_token, {}) assert status == 200 assert plain["results"][0]["answer"] @@ -477,7 +481,7 @@ def retrieve_with_trace(job_id, *, owner_id=None): {"include_orchestration_trace": True}, ) assert status == 200 - assert traced["trace"] == [{"output": "intermediate"}] + assert traced["results"][0]["trace"] assert any( event["event_type"] == "orchestration_trace_access_granted" for event in orchestrator.list_recent_audit_events() From b6ca947aa7cc876e3bdd0cdb83c1fbf63ad6011f Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Tue, 1 Sep 2026 11:26:24 +0900 Subject: [PATCH 4/5] fix(batch): extract nested orchestration traces --- contextual_orchestrator/batch_routing.py | 6 +++++- tests/test_batch_routing.py | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index f47ec42d1..74771165d 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -405,6 +405,10 @@ async def _download() -> Dict[str, Any]: body = (entry.get("response") or {}).get("body", {}) answer = _extract_answer(body) usage = body.get("usage", {}) or {} + orchestration = body.get("orchestration") + trace = body.get("trace") + if trace is None and isinstance(orchestration, dict): + trace = orchestration.get("trace") raw_request = tracked.get(custom_id) request = BatchRequest(**raw_request) if raw_request else None items.append( @@ -416,7 +420,7 @@ async def _download() -> Dict[str, Any]: attribution=dict(request.attribution) if request else {}, model=request.model if request else "contextual-orchestrator", mode=request.mode if request else "auto", - trace=body.get("trace"), + trace=trace, ) ) return items diff --git a/tests/test_batch_routing.py b/tests/test_batch_routing.py index 3b7cdc52f..d009bc339 100644 --- a/tests/test_batch_routing.py +++ b/tests/test_batch_routing.py @@ -192,6 +192,7 @@ async def download_results(self, batch_id, endpoint_alias): "body": { "choices": [{"message": {"role": "assistant", "content": "batched answer"}}], "usage": {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20}, + "orchestration": {"trace": [{"output": "intermediate"}]}, } }, } @@ -222,6 +223,7 @@ def test_pg_llm_batch_backend_submits_and_retrieves() -> None: assert item.prompt_tokens == 12 assert item.completion_tokens == 8 assert item.model == "gpt-x" + assert item.trace == [{"output": "intermediate"}] assert "download_results" in client.calls From 80e31c87d5e5b924cfe0c31a4c7eee6a15492df7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 2 Sep 2026 06:59:51 +0000 Subject: [PATCH 5/5] fix(batch): drop non-dict trace elements before they reach cost accounting PgLlmBatchBackend.retrieve() validated that a downloaded batch response's trace was a list, but not that each element was a dict. Every downstream consumer (CostRoutingCoordinator.retrieve_batch()'s billable_steps loop) calls step.get(...) unconditionally, so a malformed element (None, a string, a number) from an untrusted batch-provider response would raise AttributeError instead of failing closed. Filter to dict elements only, matching the existing defensive handling for a non-list trace. Also add the missing requestBody.required assertion CodeRabbit flagged on the batch-results OpenAPI contract test, so a future accidental required:true regression is caught. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01BV96rXhqoR3tYZ9AeAVur4 --- contextual_orchestrator/batch_routing.py | 1 + tests/test_api_contract.py | 1 + tests/test_batch_routing.py | 37 ++++++++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index 0282e790f..6cd3fb0c8 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -562,6 +562,7 @@ async def _download() -> Dict[str, Any]: trace = orchestration.get("trace") if not isinstance(trace, list): trace = [] + trace = [step for step in trace if isinstance(step, dict)] prompt_tokens = usage.get("prompt_tokens") completion_tokens = usage.get("completion_tokens") usage_valid = ( diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index 2004b30df..f7fb0fabb 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -104,6 +104,7 @@ def test_openapi_documents_compatibility_front_door() -> None: "/api/v1/batch_routing_jobs/{batch_routing_job_id}/results" ]["post"] assert batch_results["security"] == [{"inference_bearer_auth": []}] + assert batch_results["requestBody"]["required"] is False batch_results_schema = batch_results["requestBody"]["content"]["application/json"][ "schema" ] diff --git a/tests/test_batch_routing.py b/tests/test_batch_routing.py index edd2648a5..5074dbbb9 100644 --- a/tests/test_batch_routing.py +++ b/tests/test_batch_routing.py @@ -237,6 +237,43 @@ def test_pg_llm_batch_backend_submits_and_retrieves() -> None: assert "download_results" in client.calls +class _MalformedTraceElementsBatchApiClient(_FakeBatchApiClient): + async def download_results(self, batch_id, endpoint_alias): + self.calls.append("download_results") + return { + "success": True, + "batch_id": batch_id, + "responses": [ + { + "custom_id": "a", + "response": { + "body": { + "choices": [{"message": {"role": "assistant", "content": "batched answer"}}], + "usage": {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20}, + "orchestration": { + "trace": [{"output": "intermediate"}, None, "not-a-step", 7] + }, + } + }, + } + ], + } + + +def test_pg_llm_batch_backend_drops_non_dict_trace_elements() -> None: + """A downloaded trace with non-dict elements must not crash retrieval.""" + client = _MalformedTraceElementsBatchApiClient() + backend = PgLlmBatchBackend(client, endpoint_alias="prod_gateway") + requests = [BatchRequest(messages=[{"role": "user", "content": "batch me"}], custom_id="a", model="gpt-x")] + + job = backend.submit(requests, metadata={"routing_reason": "latency-tolerant"}) + backend.poll(job) + + results = backend.retrieve(job) + assert len(results) == 1 + assert results[0].trace == [{"output": "intermediate"}] + + def test_pg_llm_batch_backend_incomplete_download_raises_download_error() -> None: class _IncompleteClient(_FakeBatchApiClient): async def download_results(self, batch_id, endpoint_alias):