fix: complete PR 765 review remediations #810
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -68,6 +68,8 @@ def _temperature_capability_rejection(exc: Exception) -> bool: | |
| return False | ||
| try: | ||
| body = exc.read() | ||
| except http.client.IncompleteRead as read_error: | ||
| body = read_error.partial | ||
| except (OSError, ValueError): | ||
| body = b"" | ||
| body_bytes = body if isinstance(body, bytes) else str(body).encode("utf-8") | ||
|
|
@@ -220,6 +222,7 @@ def complete_structured( | |
| response = self.orchestrator.client.proxy_send(agent, "chat/completions", payload) | ||
| output = ModelClient._response_content(agent, response) | ||
| usage = response.get("usage") if isinstance(response.get("usage"), dict) else None | ||
| self.orchestrator._record_in_flight_provider_usage(agent, usage, output) | ||
| return self._completion_payload(output, agent.id, usage, self.mode if mode is None else mode) | ||
|
|
||
| def _completion_payload( | ||
|
|
@@ -313,6 +316,8 @@ def __post_init__(self) -> None: | |
| raise TypeError("local_credential_key must be a string") | ||
| if scheme == "local" and not self.local_credential_key: | ||
| raise ValueError("local:// gateway URLs require local_credential_key") | ||
| if scheme == "local" and not _is_local_provider_url(self.base_url): | ||
| raise ValueError("local:// gateway URLs require a well-formed explicit loopback endpoint") | ||
| if self.local_credential_key and scheme != "local": | ||
| raise ValueError("local_credential_key requires a local:// gateway URL") | ||
| if not self.auth_scheme or type(self.auth_scheme) is not str: | ||
|
|
@@ -456,7 +461,7 @@ def as_dict(self) -> dict[str, Any]: | |
| # is a caller or configuration error and must not be retried. | ||
| TRANSIENT_HTTP_STATUS = frozenset({408, 409, 425, 429, 500, 502, 503, 504}) | ||
| LOCAL_PROVIDER_SCHEMES = frozenset({"local"}) | ||
| LOCAL_PROVIDER_HOSTS = frozenset({"127.0.0.1", "::1", "localhost"}) | ||
|
|
||
|
|
||
| def _is_local_provider_url(base_url: str) -> bool: | ||
|
|
@@ -2018,14 +2023,15 @@ def close(self) -> None: | |
| class _StateStore: | ||
| """Minimal write-through sqlite persistence for orchestrator runtime state. | ||
|
|
||
| ponytail: one generic table, no ORM. Keyed kinds (workflow_run, evaluation_run, | ||
| budget_meter) | ||
| upsert by key; stream kinds (analytics, audit) append. Stream rows grow unbounded | ||
| on disk while the in-memory deques stay capped — add pruning if db size matters. | ||
| Runtime values (kind, key, payload, limit) are always bound through SQLite | ||
| placeholders so persisted prompts and identifiers cannot become SQL syntax. | ||
| """ | ||
|
|
||
| _KEYED = {"workflow_run", "evaluation_run", "budget_meter"} | ||
| _CREATE_RECORDS_SQL = ( | ||
| "CREATE TABLE IF NOT EXISTS records (" | ||
| "seq INTEGER PRIMARY KEY AUTOINCREMENT, kind TEXT NOT NULL, key TEXT, payload TEXT NOT NULL)" | ||
|
|
@@ -2169,6 +2175,9 @@ def __init__( | |
| self._audit_events: deque[dict[str, Any]] = deque(maxlen=256) | ||
| self._run_order: deque[str] = deque(maxlen=128) | ||
| self._workflow_run_lock = threading.Lock() | ||
| self._budget_spend_lock = threading.Lock() | ||
| self._budget_spent_output_tokens = 0 | ||
| self._budget_spent_cost_usd = 0.0 | ||
| self._archived_spend: dict[str, Any] = { | ||
| "run_count": 0, | ||
| "estimated_prompt_tokens": 0, | ||
|
|
@@ -2192,6 +2201,15 @@ def __init__( | |
| self._commercial_report_cache_local = threading.local() | ||
| if self._store is not None: | ||
| self._reload_state() | ||
| restored = self.spend_analytics()["totals"] | ||
| self._budget_spent_output_tokens = max( | ||
| self._budget_spent_output_tokens, | ||
| restored["estimated_output_tokens"], | ||
| ) | ||
| self._budget_spent_cost_usd = max( | ||
| self._budget_spent_cost_usd, | ||
| restored["estimated_cost_usd"] or 0.0, | ||
| ) | ||
|
Comment on lines
+2204
to
+2212
Contributor
There was a problem hiding this comment. 📝 Info: Budget ledger and run analytics intentionally diverge
Was this helpful? React with 👍 or 👎 to provide feedback.
Contributor
Author
There was a problem hiding this comment. Reviewed against ADR 0014. The budget ledger intentionally includes in-flight provider calls that may not yet have persisted workflow rows; restart reconciliation uses the durable meter and avoids double counting. No source change is justified by this informational finding. |
||
|
|
||
| def close(self) -> None: | ||
| """Release optional durable resources owned by this orchestrator.""" | ||
|
|
@@ -2248,6 +2266,14 @@ def provider_readiness_report( | |
| } | ||
|
|
||
| def _reload_state(self) -> None: | ||
| budget_meter = self._store.load("budget_meter") | ||
| if budget_meter: | ||
| self._budget_spent_output_tokens = int( | ||
| budget_meter[-1].get("spent_output_tokens", 0) | ||
| ) | ||
| self._budget_spent_cost_usd = float( | ||
| budget_meter[-1].get("spent_cost_usd", 0.0) | ||
| ) | ||
| for record in self._store.load("workflow_run"): | ||
| self._remember_workflow_run(record) | ||
| for evaluation in self._store.load("evaluation_run"): | ||
|
|
@@ -2339,6 +2365,11 @@ def proxy_completion( | |
| usage, | ||
| responses=endpoint.strip("/") == "responses", | ||
| ) | ||
| self._record_in_flight_provider_usage( | ||
| agent, | ||
| passthrough_step.get("usage"), | ||
| passthrough_output, | ||
| ) | ||
| self._persist_workflow_run( | ||
| { | ||
| "workflow_run_id": f"run_{uuid.uuid4().hex}", | ||
|
|
@@ -2386,29 +2417,47 @@ def _orchestrated_provider_completion( | |
| disabled_model = requested_model if requested_model is not None else final_agent.model | ||
| raise RuntimeError(f"requested model {disabled_model!r} is disabled") | ||
|
|
||
| with self._budget_spend_lock: | ||
| budget_before = ( | ||
| self._budget_spent_output_tokens, | ||
| self._budget_spent_cost_usd, | ||
| ) | ||
| workflow = self.conduct(messages, preserve_messages=True, judge=True) | ||
| workflow_output_tokens = 0 | ||
| workflow_cost_usd = 0.0 | ||
| model_by_agent = {agent.id: agent.model for agent in self.agents} | ||
| metered_steps = list(workflow["trace"]) | ||
| verification = workflow.get("verification") | ||
| if isinstance(verification, dict) and isinstance( | ||
| verification.get("judge_usage"), dict | ||
| ): | ||
| metered_steps.append( | ||
| { | ||
| "agent_id": verification.get("judge_agent_id", "unknown"), | ||
| "usage": verification["judge_usage"], | ||
| "output": "", | ||
| } | ||
| ) | ||
| for row in metered_steps: | ||
| output_tokens, _reported = _step_output_token_count(row) | ||
| workflow_output_tokens += output_tokens | ||
| model = model_by_agent.get( | ||
| row.get("served_agent_id") or row.get("agent_id") | ||
| ) | ||
| price = self.price_per_million.get(model) if model is not None else None | ||
| if price is not None: | ||
| workflow_cost_usd += output_tokens / 1_000_000 * price | ||
| with self._budget_spend_lock: | ||
| recorded_output_tokens = self._budget_spent_output_tokens - budget_before[0] | ||
| recorded_cost_usd = self._budget_spent_cost_usd - budget_before[1] | ||
| self._budget_spent_output_tokens += max( | ||
| 0, | ||
| workflow_output_tokens - recorded_output_tokens, | ||
| ) | ||
| self._budget_spent_cost_usd += max( | ||
| 0.0, | ||
| workflow_cost_usd - recorded_cost_usd, | ||
| ) | ||
|
Comment on lines
+2450
to
2460
Contributor
There was a problem hiding this comment. 📝 Info: Judge counted exactly once by reconciliation The reconciliation adds
Was this helpful? React with 👍 or 👎 to provide feedback.
Contributor
Author
There was a problem hiding this comment. Reviewed against ADR 0014. The reconciliation is the required safety net for judge usage when adapter metering is absent, while preventing duplicate counting when it is present. No source change is needed. |
||
| evidence = "\n\n".join( | ||
| f"Workflow step {row['id']} ({row['role']}):\n{row['output']}" | ||
| for row in workflow["trace"] | ||
|
|
@@ -2438,10 +2487,7 @@ def _orchestrated_provider_completion( | |
| ) | ||
| upstream["model"] = final_agent.model | ||
| upstream["stream"] = False | ||
| self._raise_if_spend_budget_exceeded() | ||
| synthesis_started = time.perf_counter() | ||
| raw = self.client.proxy_send(final_agent, "responses", upstream) | ||
| synthesis_output = raw.get("output_text") | ||
|
|
@@ -2487,10 +2533,7 @@ def _orchestrated_provider_completion( | |
| upstream["model"] = final_agent.model | ||
| upstream["messages"] = synthesis_messages | ||
| upstream["stream"] = False | ||
| self._raise_if_spend_budget_exceeded() | ||
| synthesis_started = time.perf_counter() | ||
| raw = self.client.proxy_send(final_agent, "chat/completions", upstream) | ||
| synthesis_output = "" | ||
|
|
@@ -2515,6 +2558,11 @@ def _orchestrated_provider_completion( | |
| usage, | ||
| responses=response_request, | ||
| ) | ||
| self._record_in_flight_provider_usage( | ||
| final_agent, | ||
| synthesis_step.get("usage"), | ||
| synthesis_output, | ||
| ) | ||
| trace = [*workflow["trace"], synthesis_step] | ||
| workflow_run_id = f"run_{uuid.uuid4().hex}" | ||
| record = { | ||
|
|
@@ -2599,6 +2647,7 @@ def stream_route(self, messages: list[ChatMessage], workflow_run_id: str | None | |
| parts.append(delta) | ||
| yield delta | ||
| answer = "".join(parts) | ||
| self._record_in_flight_provider_usage(agent, None, answer) | ||
| record = { | ||
| "workflow_run_id": workflow_run_id or f"run_{uuid.uuid4().hex}", | ||
| "created_at": int(time.time()), | ||
|
|
@@ -2749,6 +2798,11 @@ def batch_route(self, prompts: list[str]) -> list[dict[str, Any]]: | |
| } | ||
| if result.get("usage") is not None: | ||
| row["usage"] = result["usage"] | ||
| self._record_in_flight_provider_usage( | ||
| agent, | ||
| row.get("usage"), | ||
| result["content"], | ||
| ) | ||
| record = { | ||
| "workflow_run_id": f"run_{uuid.uuid4().hex}", | ||
| "created_at": int(time.time()), | ||
|
|
@@ -3235,10 +3289,15 @@ def _plan_generated(self, task: str) -> list[WorkflowStep]: | |
| "verifier step when correctness matters.\n" | ||
| f"Available agents:\n{pool}" | ||
| ) | ||
| raw, _served_id, _usage = self._invoke( | ||
| planner, | ||
| [ | ||
| {"role": "system", "content": system}, | ||
| {"role": "user", "content": task}, | ||
| ], | ||
| text=task, | ||
| role="thinker", | ||
| ) | ||
|
Comment on lines
+3292
to
+3300
Contributor
There was a problem hiding this comment. 📝 Info: BudgetExceededError caught by broad handler in generated planning
Was this helpful? React with 👍 or 👎 to provide feedback.
Contributor
Author
There was a problem hiding this comment. Accepted. Local follow-up 7929f70 adds an explicit BudgetExceededError re-raise before the template fallback and a regression test. Focused verification passed 67 tests. The remote branch remains at 513a815 because the normal push was rejected by the active required-workflow ruleset; no bypass was used. |
||
| return self._parse_workflow_plan(raw) | ||
|
|
||
| def _parse_workflow_plan(self, raw: str) -> list[WorkflowStep]: | ||
|
|
@@ -3363,13 +3422,17 @@ def _invoke( | |
| last_error: Exception | None = None | ||
| for agent in candidates: | ||
| try: | ||
| self._raise_if_spend_budget_exceeded() | ||
| output = self.client.chat(agent, messages) | ||
| except BudgetExceededError: | ||
| raise | ||
| except Exception as exc: # noqa: BLE001 - one agent failing routes to the next | ||
| last_error = exc | ||
| self._record_failure(agent.id) | ||
| continue | ||
| self._record_success(agent.id) | ||
| usage = self.client.take_usage() if hasattr(self.client, "take_usage") else None | ||
| self._record_in_flight_provider_usage(agent, usage, output) | ||
| return output, agent.id, usage | ||
| if isinstance(last_error, ProviderResponseError): | ||
| raise last_error | ||
|
|
@@ -3772,6 +3835,10 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> | |
| "estimated_cost_usd": cost, | ||
| }) | ||
|
|
||
| with self._budget_spend_lock: | ||
| budget_output_tokens = self._budget_spent_output_tokens | ||
| budget_cost_usd = self._budget_spent_cost_usd | ||
|
|
||
| return { | ||
| "measurement_status": "local_runtime_estimate", | ||
| "source_note": ( | ||
|
|
@@ -3792,7 +3859,10 @@ def spend_analytics(self, price_per_million: dict[str, float] | None = None) -> | |
| }, | ||
| "by_model": rows, | ||
| "unpriced_models": unpriced, | ||
| "budget": self._budget_block( | ||
| budget_output_tokens, | ||
| round(budget_cost_usd, 6) if prices else None, | ||
| ), | ||
| } | ||
|
|
||
| def _budget_block(self, spent_tokens: int, spent_cost: float | None) -> dict[str, Any]: | ||
|
|
@@ -3820,6 +3890,31 @@ def budget_status(self) -> dict[str, Any]: | |
| """Current spend-budget state (limits, spent, remaining, exceeded).""" | ||
| return self.spend_analytics()["budget"] | ||
|
|
||
| def _record_in_flight_provider_usage( | ||
| self, | ||
| agent: ModelAgent, | ||
| usage: dict[str, Any] | None, | ||
| output: str, | ||
| ) -> None: | ||
| """Add one completed provider call to the process budget ledger.""" | ||
| output_tokens, _reported = _step_output_token_count( | ||
| {"usage": usage, "output": output} | ||
| ) | ||
| price = self.price_per_million.get(agent.model) | ||
| with self._budget_spend_lock: | ||
| self._budget_spent_output_tokens += output_tokens | ||
| if price is not None: | ||
| self._budget_spent_cost_usd += output_tokens / 1_000_000 * price | ||
| if self._store is not None: | ||
| self._store.save( | ||
| "budget_meter", | ||
| "process_budget", | ||
| { | ||
| "spent_output_tokens": self._budget_spent_output_tokens, | ||
| "spent_cost_usd": self._budget_spent_cost_usd, | ||
| }, | ||
| ) | ||
|
|
||
| def _raise_if_spend_budget_exceeded( | ||
| self, | ||
| *, | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🔍 Existing host.docker.internal agents break on load
Removing
host.docker.internalfromLOCAL_PROVIDER_HOSTSplus the new__post_init__check makes anylocal://host.docker.internalagent raiseValueErrorat construction.from_dictshares this path, so a stored agent-pool row or agents JSON using that host raises duringload_all/load_agentson startup. Intended hardening, but a breaking migration for Docker-gateway configs.Was this helpful? React with 👍 or 👎 to provide feedback.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Reviewed against ADR 0012 and the ADR 0004 decision recorded on 2026-08-21. local:// is intentionally authenticated and loopback-only; host.docker.internal is not a loopback guarantee and must remain rejected. The DNS resolution and rebinding check remains in the provider validator. No compatibility reintroduction is appropriate.