fix(security): make batch trace retrieval opt-in #982
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -556,6 +556,13 @@ async def _download() -> Dict[str, Any]: | |
| body = response.get("body", {}) or {} | ||
| 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") | ||
| if not isinstance(trace, list): | ||
| trace = [] | ||
|
seonghobae marked this conversation as resolved.
|
||
| trace = [step for step in trace if isinstance(step, dict)] | ||
| prompt_tokens = usage.get("prompt_tokens") | ||
| completion_tokens = usage.get("completion_tokens") | ||
| usage_valid = ( | ||
|
|
@@ -576,6 +583,7 @@ async def _download() -> Dict[str, Any]: | |
| model=request.model if request else "contextual-orchestrator", | ||
| mode=request.mode if request else "auto", | ||
| messages=list(request.messages) if request else [], | ||
| trace=trace, | ||
|
There was a problem hiding this comment. 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift 트레이스 추가 후 항목 usage를 보존하세요.
🤖 Prompt for AI Agents
Contributor
Author
There was a problem hiding this comment. Verified this is real: once Not pushing a fix for this one: the two remediations you propose (make item-level usage authoritative regardless of trace presence, vs. define and implement an explicit trace-step usage-aggregation rule) are genuinely different accounting semantics, and this loop is shared with Generated by Claude Code There was a problem hiding this comment.
The fix requires an explicit policy for precedence between item-level usage and Keeping this thread open is correct for this security-scoped PR. Do you want me to open a follow-up GitHub issue for the ✏️ Learnings added
You are interacting with an AI system.
Contributor
Author
There was a problem hiding this comment. Yes, please open the follow-up issue — that keeps the accounting-policy decision (item-level usage vs. Generated by Claude Code There was a problem hiding this comment. 🧩 Analysis chain🏁 Script executed: #!/bin/bash
set -euo pipefail
repo="ContextualWisdomLab/contextual-orchestrator"
body=$(cat <<'EOF'
## 요약
`CostRoutingCoordinator.retrieve_batch()`의 usage 우선순위 정책을 명시하고 구현해야 합니다.
현재 `BatchResultItem.trace` 또는 `BatchResultItem.race_usage`가 있으면 공통 `billable_steps` 경로가 사용됩니다. 이 단계에 `usage`가 없으면, 유효한 `BatchResultItem.prompt_tokens` 및 `BatchResultItem.completion_tokens`가 있어도 item-level fallback이 실행되지 않습니다. 그 결과 측정된 usage와 cost가 `unavailable`로 기록될 수 있습니다.
## 재현 조건
1. `BatchResultItem.prompt_tokens=12` 및 `BatchResultItem.completion_tokens=8`을 설정합니다.
2. `BatchResultItem.trace` 또는 `BatchResultItem.race_usage`에 `usage`가 없는 step을 하나 이상 설정합니다.
3. `CostRoutingCoordinator.retrieve_batch()`를 실행합니다.
4. 현재 동작에서는 item-level usage fallback이 건너뛰어질 수 있습니다.
## 필요한 변경
- `contextual_orchestrator/cost_router.py`의 `CostRoutingCoordinator.retrieve_batch()`에서 item-level usage와 `billable_steps` usage의 우선순위 정책을 결정합니다.
- 이 정책은 `BatchResultItem.trace`와 기존 `BatchResultItem.race_usage`에 동일하게 적용해야 합니다.
- 선택한 정책에 따라 usage 및 cost 계산을 구현합니다.
- 정책의 근거와 fallback 조건을 코드 또는 관련 문서에 명확히 기록합니다.
- 회귀 테스트를 추가합니다.
## 영향 범위
- `contextual_orchestrator/cost_router.py`
- `CostRoutingCoordinator.retrieve_batch()`
- `billable_steps` 기반 usage 및 cost 계산
- 관련 batch retrieval 회귀 테스트
## 완료 조건
- trace step에 `usage`가 없어도 authoritative item-level usage `{prompt_tokens: 12, completion_tokens: 8}`를 처리하는 정책이 명확합니다.
- 선택한 정책에 따라 measured usage와 cost가 일관되게 계산됩니다.
- 동일 정책이 `trace` 및 `race_usage` 경로에 적용됩니다.
- 위 재현 조건을 검증하는 회귀 테스트가 있습니다.
## 관련 링크
- PR: https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/982
- 검토 스레드: https://github.com/ContextualWisdomLab/contextual-orchestrator/pull/982#discussion_r3911523497
- 요청자: `@seonghobae`
EOF
)
gh issue create \
--repo "$repo" \
--title "fix(cost): batch usage 우선순위 정책 정의 및 회귀 방지" \
--body "$body" \
--assignee "seonghobae"Length of output: 2081
이슈는 🧠 Learnings usedYou are interacting with an AI system. |
||
| usage_valid=usage_valid, | ||
| ) | ||
| ) | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -236,6 +236,7 @@ def server_close(self) -> None: | |
| "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", | ||
|
|
@@ -6564,7 +6565,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) | ||
|
|
@@ -7478,7 +7483,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")] | ||
| _reject_unknown_keys(body, ALLOWED_BATCH_RESULTS_KEYS) | ||
| include_trace = self._validate_trace_request(body, default=False) | ||
| if include_trace: | ||
| self._authorize_trace_access() | ||
|
seonghobae marked this conversation as resolved.
Comment on lines
+7487
to
+7489
Contributor
There was a problem hiding this comment. 🟡 Provider batches cannot return traces For 
Was this helpful? React with 👍 or 👎 to provide feedback.
Contributor
Author
There was a problem hiding this comment. CodeRabbit's latest pass raised the same gap twice more (outside the diff range, so no separate inline threads: Confirmed real and pre-existing: Generated by Claude Code |
||
| try: | ||
| retrieved = self._run( | ||
| lambda: coordinator.retrieve_batch( | ||
|
|
@@ -7488,8 +7496,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 | ||
| if include_trace: | ||
| self._audit_trace_disclosure( | ||
| "/api/v1/batch_routing_jobs/{job_id}/results" | ||
| ) | ||
| self._send(_response_payload(retrieved, include_trace=include_trace)) | ||
|
seonghobae marked this conversation as resolved.
seonghobae marked this conversation as resolved.
|
||
| return | ||
| if path == "/v1/responses": | ||
| # The Responses API has no chat-completions verifier equivalent, | ||
|
|
@@ -8114,10 +8125,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], *, 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 if default is None else default | ||
| ) | ||
|
seonghobae marked this conversation as resolved.
|
||
| elif type(body["include_orchestration_trace"]) is not bool: | ||
| raise RequestError( | ||
| 400, | ||
|
|
@@ -8159,9 +8174,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, | ||
| allow_empty_without_content_type: bool = False, | ||
| ) -> dict[str, Any]: | ||
| try: | ||
| body_size = _request_body_size( | ||
| self.headers, | ||
|
|
@@ -8171,6 +8189,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") | ||
|
seonghobae marked this conversation as resolved.
|
||
| raw = self.rfile.read(body_size) | ||
| if len(raw) != body_size: | ||
| self.close_connection = True | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.