diff --git a/doc/architecture/SPEC/kb_service.md b/doc/architecture/SPEC/kb_service.md index e9073758..48444b40 100644 --- a/doc/architecture/SPEC/kb_service.md +++ b/doc/architecture/SPEC/kb_service.md @@ -301,3 +301,13 @@ diff 时保留原摘录（证据不会因此变得更"好过"）。只扩展带 ## 2026-09-30 模型回复解析 `models.parse_json_object` 取裸 JSON 对象或最后一个 ```json 围栏块，用 `json.loads(strict=False)`：字符串里的原始控制字符（多行值里的换行、制表符，例如图示）按数据接受，不再把结构正确的回答判为 `ModelUnavailable`。测试：`test_kb_intake_gate.py`。 + +## 2026-09-30 单次调用花费阈值（kb init 预算硬上限的一块积木） +`ModelGateway.call_json(..., max_budget_usd=x)`：x 是**停止阈值**，不是硬上限 —— 花费达到 x 后不再发起新请求，但越过 x 的 +那个请求照样计费。硬上限由调用方负责：调用前预留 x + 单个请求的最坏情况。transport 的 `stops_at_spend` 为 False 时在发出调用 +**之前**抛 `ModelUnavailable(" cannot stop a call at a spend threshold")`（fail-closed，绝不用 `max_tokens` 近似）； +x 必须为正。因阈值停下的调用（`stop_reason="max_budget"`）→ `ModelUnavailable`，并作为失败记录（不会成为训练样本）。 +`ModelReply.cost_usd` 取自 `usage["cost_usd"]`（未知为 None），trace 记录同时写入 `cost_usd` 与 `max_budget_usd`。 +不传上限的现有调用方行为不变：该参数根本不会发给 transport。 +`runtime.trace_recorder` 把阈值与花费持久化到 `model_call` 记录的 `result.max_budget_usd` / `result.cost_usd` +（未设或未知为 null），包括 transport 抛错的失败记录。 diff --git a/doc/architecture/SPEC/providers/base.md b/doc/architecture/SPEC/providers/base.md index e48cd786..dea613f3 100644 --- a/doc/architecture/SPEC/providers/base.md +++ b/doc/architecture/SPEC/providers/base.md @@ -1,6 +1,6 @@ # providers/base.py —— 规范 - + `LOC ~161 · provider 层契约 + 子进程环境白名单 · refactor-status: ok` @@ -53,3 +53,10 @@ stdlib + `..scopes.ToolScope`。它是一个叶子契约模块。 ## 2026-09-28 `complete()` 接受 `effort`（推理强度）：codex 以 `-c model_reasoning_effort=""` 生效并校验取值；其他 transport 接受并忽略（模型 id 已决定推理预算）。 + +## 2026-09-30 单次调用花费阈值（不是硬上限） +`complete()` 接受 `max_budget_usd`，类属性 `stops_at_spend`（默认 False）声明 CLI 是否会在本次调用花费达到阈值后**不再发起** +新的 API 请求。越过阈值的那一个请求仍会完成并计费，所以超出量最多一个请求 —— 这是阈值，不是硬上限； +需要硬上限的调用方要预留"阈值 + 单个请求的最坏情况"。不支持的 transport 接受并忽略该参数，调用方必须先检查 +`stops_at_spend`（知识服务的 gateway 会拒绝）。不得出现声称能硬性限制单次调用花费的属性。 +`max_tokens` 对 CLI transport 不是花费上限。transport 报告的本次花费放在 `Reply.usage["cost_usd"]`（未知时不出现）。 diff --git a/doc/architecture/SPEC/providers/claude_code.md b/doc/architecture/SPEC/providers/claude_code.md index dabbc232..cab70c0e 100644 --- a/doc/architecture/SPEC/providers/claude_code.md +++ b/doc/architecture/SPEC/providers/claude_code.md @@ -1,6 +1,6 @@ # providers/claude_code.py —— 规范 - + `LOC ~205 · harness transport（Claude 订阅） · refactor-status: ok` @@ -43,3 +43,10 @@ stdlib + `.base` + `..agent_loop.AgentOutcome` + `..llm` 的类型。 ## 2026-09-28 `complete()` 接受并忽略 `effort`（与 base 契约一致）。 + +## 2026-09-30 单次调用花费阈值 +`stops_at_spend = True`：`complete(max_budget_usd=x)` 把 `--max-budget-usd x` 传给 CLI（仅 print 模式有效，`_run` 正是 print 模式）， +x 必须为正。CLI 在每个 API 请求完成**之后**检查预算：它阻止下一个请求，不能截断当前请求，所以花费可超出 x 一个请求。CLI 因预算停下（`subtype` 含 budget）时返回 `stop_reason="max_budget"` 且无文本 —— 被截断的回答永远不可用。 +`total_cost_usd` 以 `usage["cost_usd"]` 返回。 +阈值参数由 `budget_arg()` 精确生成：浮点数的最短往返十进制、普通（非科学）记数法；非正、非有限值 → `ValueError`。 +固定精度格式会移动边界（`.4f` 把 0.00001 变成 "0.0000" 即没有阈值，把 0.12345 变成 "0.1235" 即更高的阈值），所以不允许。 diff --git a/doc/architecture/SPEC/providers/codex.md b/doc/architecture/SPEC/providers/codex.md index fac54c8d..8108fc45 100644 --- a/doc/architecture/SPEC/providers/codex.md +++ b/doc/architecture/SPEC/providers/codex.md @@ -1,6 +1,6 @@ # providers/codex.py —— 规范 - + `LOC ~197 · harness transport（ChatGPT 订阅） · refactor-status: ok` @@ -42,3 +42,6 @@ OS 级控制可用的后端；把两者混同，会模糊掉"当前实际生效 ## 2026-09-28 `complete()` 接受 `effort`（推理强度）：codex 以 `-c model_reasoning_effort=""` 生效并校验取值；其他 transport 接受并忽略（模型 id 已决定推理预算）。 + +## 2026-09-30 +`complete()` 接受并忽略 `max_budget_usd`；`stops_at_spend` 为 False（CLI 无花费阈值）。 diff --git a/doc/architecture/SPEC/providers/cursor.md b/doc/architecture/SPEC/providers/cursor.md index 23e58cea..91109a4b 100644 --- a/doc/architecture/SPEC/providers/cursor.md +++ b/doc/architecture/SPEC/providers/cursor.md @@ -1,6 +1,6 @@ # providers/cursor.py —— 规范 - + `LOC ~207 · harness transport（Cursor 订阅） · refactor-status: ok` @@ -48,3 +48,6 @@ scope 约束，native 调用只被**记录**。 ## 2026-09-28 `complete()` 接受并忽略 `effort`（与 base 契约一致）。 + +## 2026-09-30 +`complete()` 接受并忽略 `max_budget_usd`；`stops_at_spend` 为 False（CLI 无花费阈值）。 diff --git a/doc/architecture/SPEC/providers/deepseek.md b/doc/architecture/SPEC/providers/deepseek.md index 213f73c8..67199eab 100644 --- a/doc/architecture/SPEC/providers/deepseek.md +++ b/doc/architecture/SPEC/providers/deepseek.md @@ -1,6 +1,6 @@ # providers/deepseek.py —— 规范 - + `LOC ~502 · harness transport（dsh，API-keyed） · refactor-status: oversized` @@ -59,3 +59,6 @@ stdlib + `.base` + `.registry` + `..agent_loop` + `..llm` 的类型 + dsh SDK（ ## 2026-09-28 `complete()` 接受并忽略 `effort`（与 base 契约一致）。 + +## 2026-09-30 +`complete()` 接受并忽略 `max_budget_usd`；`stops_at_spend` 为 False。 diff --git a/doc/architecture/SPEC/providers/zcode.md b/doc/architecture/SPEC/providers/zcode.md index a38fa55a..f9392455 100644 --- a/doc/architecture/SPEC/providers/zcode.md +++ b/doc/architecture/SPEC/providers/zcode.md @@ -49,3 +49,5 @@ stdlib + `.base` + `..agent_loop.AgentOutcome` + `..llm` 的类型与 `ModelMism ## 测试 `test_provider_zcode.py`（离线，事件形状取自 zcode 0.16.9 实跑）；2026-09-30 用真实 zcode 做过 一次端到端冒烟（工具桥 `read_file` 被调用，`bridge_trace.jsonl` 有记录）。 + +`complete()` 接受并忽略 `max_budget_usd`；`stops_at_spend` 为 False（CLI 无花费阈值，`ModelGateway` 请求阈值时在派发前拒绝）。 diff --git a/src/infermatrix_copilot/kb_service/models.py b/src/infermatrix_copilot/kb_service/models.py index df58387f..9ac13565 100644 --- a/src/infermatrix_copilot/kb_service/models.py +++ b/src/infermatrix_copilot/kb_service/models.py @@ -2,8 +2,8 @@ The generator drafts knowledge changes; the judge (a different model family) grades them. Both are pinned by (provider, model, reasoning effort) and -recorded on every call. An unavailable backend, a timeout or an unparseable -reply raises ``ModelUnavailable``; the caller leaves the work queued. It never +recorded on every call. An unavailable backend, a timeout, a call stopped +at its spend cap or an unparseable reply raises ``ModelUnavailable``; the caller leaves the work queued. It never switches to a weaker model: a verdict must always say which model made it. """ @@ -61,6 +61,7 @@ class ModelReply: served_model: str usage: dict seconds: float + cost_usd: float | None = None # the transport's reported spend (None: unknown) def parse_json_object(text: str) -> dict: @@ -104,28 +105,47 @@ def _transport(self, provider: str): return transport def call_json(self, role: ModelRole, *, system: str, prompt: str, - validate: Callable[[dict], None] | None = None) -> ModelReply: + validate: Callable[[dict], None] | None = None, + max_budget_usd: float | None = None) -> ModelReply: + """``max_budget_usd`` is a per-call STOP THRESHOLD, not a hard cap: + the transport starts no further API request once the call's spend + reaches it, but the request that crosses it is billed in full. A + caller enforcing a hard ceiling must reserve the threshold plus one + request's worst case before calling. Requesting a threshold from a + transport that cannot stop at one (``stops_at_spend`` False) is + refused before dispatch; it is never approximated with ``max_tokens``.""" transport = self._transport(role.provider) + cap: dict = {} + if max_budget_usd is not None: + if not max_budget_usd > 0: + raise ModelUnavailable(f"{role.label()}: max_budget_usd must be positive") + if not getattr(transport, "stops_at_spend", False): + raise ModelUnavailable(f"{role.provider} cannot stop a call at a spend threshold") + cap = {"max_budget_usd": max_budget_usd} started = time.time() identity = {"role": role.name, "requested": role.label(), "provider": role.provider, "model": role.model, "effort": role.effort} try: reply = transport.complete( system=system, messages=[{"role": "user", "content": prompt}], - model=role.model, effort=role.effort, role=role.name) + model=role.model, effort=role.effort, role=role.name, **cap) except Exception as exc: # the transport's own failure modes if self._recorder is not None: self._recorder({**identity, "served_model": "", "stop_reason": "", "usage": {}, + "cost_usd": None, "max_budget_usd": max_budget_usd, "seconds": round(time.time() - started, 3), "system": system, "prompt": prompt, "reply": "", "error": str(exc)[:2000]}) raise ModelUnavailable(f"{role.label()} failed: {exc}") from exc seconds = time.time() - started text = "".join(getattr(block, "text", "") or "" for block in getattr(reply, "blocks", [])) + usage = dict(getattr(reply, "usage", {}) or {}) + cost = usage.get("cost_usd") + cost_usd = float(cost) if isinstance(cost, (int, float)) and not isinstance(cost, bool) else None record = { **identity, "served_model": getattr(reply, "model", "") or "", "stop_reason": getattr(reply, "stop_reason", ""), - "usage": dict(getattr(reply, "usage", {}) or {}), + "usage": usage, "cost_usd": cost_usd, "max_budget_usd": max_budget_usd, "seconds": round(seconds, 3), "system": system, "prompt": prompt, "reply": text, } # the record carries the call's FINAL verdict: a truncated, empty, @@ -133,7 +153,9 @@ def call_json(self, role: ModelRole, *, system: str, prompt: str, # never becomes a training example failure: ModelUnavailable | None = None data: dict = {} - if getattr(reply, "stop_reason", "") == "max_tokens" or not text.strip(): + if getattr(reply, "stop_reason", "") == "max_budget": + failure = ModelUnavailable(f"{role.label()} stopped at its spend threshold ${max_budget_usd}") + elif getattr(reply, "stop_reason", "") == "max_tokens" or not text.strip(): failure = ModelUnavailable(f"{role.label()} timed out or returned nothing") else: try: @@ -149,4 +171,4 @@ def call_json(self, role: ModelRole, *, system: str, prompt: str, self._recorder({**record, "error": str(failure) if failure else ""}) if failure is not None: raise failure - return ModelReply(role, data, text, record["served_model"], record["usage"], seconds) + return ModelReply(role, data, text, record["served_model"], usage, seconds, cost_usd) diff --git a/src/infermatrix_copilot/kb_service/runtime.py b/src/infermatrix_copilot/kb_service/runtime.py index 1fc14a9e..0c1b2697 100644 --- a/src/infermatrix_copilot/kb_service/runtime.py +++ b/src/infermatrix_copilot/kb_service/runtime.py @@ -128,7 +128,12 @@ def record(entry: dict) -> None: outputs={"reply": entry["reply"]}, model={k: entry.get(k, "") for k in ("role", "provider", "model", "effort", "served_model")}, usage=entry.get("usage") or {}, seconds=entry.get("seconds"), - result={"stop_reason": entry.get("stop_reason", "")}, error=entry.get("error", "")) + # spend: the requested stop threshold and the reported cost + # (None when unset/unknown), on failures too + result={"stop_reason": entry.get("stop_reason", ""), + "max_budget_usd": entry.get("max_budget_usd"), + "cost_usd": entry.get("cost_usd")}, + error=entry.get("error", "")) except Exception: # noqa: BLE001 - tracing never breaks a model call pass diff --git a/src/infermatrix_copilot/providers/base.py b/src/infermatrix_copilot/providers/base.py index b912d845..943831cf 100644 --- a/src/infermatrix_copilot/providers/base.py +++ b/src/infermatrix_copilot/providers/base.py @@ -140,18 +140,35 @@ def auth_gap(self) -> str | None: return None # -- contract ------------------------------------------------------------ + # True only where the CLI itself stops a `complete(max_budget_usd=...)` + # call from STARTING another API request once its spend reaches the + # threshold. The request that crosses it still completes and is billed, + # so this bounds the overshoot to one request; it is not a hard ceiling. + # `max_tokens` is no spend bound at all for CLI transports. + stops_at_spend: bool = False + def run_session(self, req: AgentSessionRequest): """Run one delegated agent step; returns `agent_loop.AgentOutcome`.""" raise NotImplementedError def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = ""): + role: str = "", effort: str = "", + max_budget_usd: float | None = None): """One-shot tool-less completion; returns a normalized `llm.Reply`. `effort` is a reasoning-effort pin (`low`/`medium`/`high`/`xhigh`) for transports whose CLI takes one (codex); the others accept and ignore - it, because their model id already fixes the reasoning budget.""" + it, because their model id already fixes the reasoning budget. + + `max_budget_usd` is a stop threshold where the CLI honours one + (`stops_at_spend`): no further API request starts once the call's + spend reaches it, but the request that crosses it is still billed. A + caller that needs a hard ceiling reserves the threshold PLUS one + request's worst case. Transports without it accept and ignore the + argument, so callers must check `stops_at_spend` first (the knowledge + service's gateway refuses them). The call's cost, when the transport + reports one, is `Reply.usage["cost_usd"]`.""" raise NotImplementedError diff --git a/src/infermatrix_copilot/providers/claude_code.py b/src/infermatrix_copilot/providers/claude_code.py index ff382020..13fbe0f4 100644 --- a/src/infermatrix_copilot/providers/claude_code.py +++ b/src/infermatrix_copilot/providers/claude_code.py @@ -22,8 +22,10 @@ from __future__ import annotations import json +import math import subprocess import sys +from decimal import Decimal from pathlib import Path from ..agent_loop import AgentOutcome @@ -43,15 +45,31 @@ _BRIDGE_SERVER = "infermatrix-tools" +def budget_arg(value: float) -> str: + """The threshold as the CLI argument, EXACTLY: the float's shortest + round-trip decimal in plain (never scientific) notation. Fixed-precision + formatting would move the boundary: `.4f` turned 0.00001 into "0.0000" + (no threshold at all) and 0.12345 into "0.1235" (a higher one).""" + if isinstance(value, bool) or not isinstance(value, (int, float)) \ + or not math.isfinite(value) or not value > 0: + raise ValueError(f"max_budget_usd must be a positive finite number, got {value!r}") + return format(Decimal(repr(float(value))), "f") + + class ClaudeCodeTransport(HarnessTransport): """claude CLI (headless) as a Strict backend.""" spec = PROVIDERS["claude-code"] + # `--max-budget-usd` (print mode, which _run uses) is checked after each + # API request completes: it stops the next one, it cannot cut the current + # one short. A threshold, never a hard cap. + stops_at_spend = True # -- process plumbing ---------------------------------------------------- def _run(self, prompt_text: str, *, system: str, cwd: str, timeout_s: float, max_turns: int, model: str = "", - mcp_config: Path | None = None) -> tuple[dict, bool]: + mcp_config: Path | None = None, + max_budget_usd: float | None = None) -> tuple[dict, bool]: """One CLI invocation → (parsed result object, timed_out). Prompt on stdin (argv has a 128KiB per-arg limit; evidence packs exceed it).""" cmd = [self.require_cli(), "-p", "--output-format", "json", @@ -63,6 +81,8 @@ def _run(self, prompt_text: str, *, system: str, cwd: str, selected = model or self.settings.strict_backend_model if selected: cmd += ["--model", selected] + if max_budget_usd is not None: + cmd += ["--max-budget-usd", budget_arg(max_budget_usd)] if mcp_config is not None: cmd += ["--mcp-config", str(mcp_config), "--strict-mcp-config", "--allowedTools", f"mcp__{_BRIDGE_SERVER}"] @@ -177,23 +197,35 @@ def run_session(self, req: AgentSessionRequest) -> AgentOutcome: def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = "") -> Reply: + role: str = "", effort: str = "", + max_budget_usd: float | None = None) -> Reply: """Tool-less one-shot: built-ins denied, no MCP config, two turns (one to think, the cap as a backstop). Runs in the run-less scratch of the process cwd — with every tool denied there is nothing to - contain.""" + contain. + + `max_budget_usd` is passed as `--max-budget-usd`: once the call's + spend reaches it the CLI starts no further API request (the request + that crossed it is already billed, so spend can exceed the threshold + by one request). A call the CLI ended on its budget (a `subtype` + naming the budget) comes back with `stop_reason "max_budget"` and no + text. The CLI's `total_cost_usd` is reported as + `usage["cost_usd"]` (absent when the CLI gave none).""" import tempfile with tempfile.TemporaryDirectory(prefix="imc-claude-oneshot-") as td: data, timed_out = self._run( flatten_messages("", messages), system=system, cwd=td, timeout_s=self.settings.strict_backend_timeout_s, - max_turns=2, model=model) + max_turns=2, model=model, max_budget_usd=max_budget_usd) usage = self._usage(data) - text = str(data.get("result") or "") + over_budget = "budget" in str(data.get("subtype") or "") + text = "" if over_budget else str(data.get("result") or "") + extra = {} if usage.cost_usd is None else {"cost_usd": usage.cost_usd} return Reply( blocks=[Block(type="text", text=text)] if text else [], - stop_reason="max_tokens" if timed_out else "end_turn", + stop_reason=("max_tokens" if timed_out + else "max_budget" if over_budget else "end_turn"), usage={"input_tokens": usage.input_tokens, "output_tokens": usage.output_tokens, "cache_read_input_tokens": int( @@ -201,5 +233,6 @@ def complete(self, *, system: str, messages: list[dict], "cache_read_input_tokens") or 0), "cache_creation_input_tokens": int( (data.get("usage") or {}).get( - "cache_creation_input_tokens") or 0)}, + "cache_creation_input_tokens") or 0), + **extra}, model=usage.served_model) diff --git a/src/infermatrix_copilot/providers/codex.py b/src/infermatrix_copilot/providers/codex.py index f85a000d..ffa289ac 100644 --- a/src/infermatrix_copilot/providers/codex.py +++ b/src/infermatrix_copilot/providers/codex.py @@ -182,7 +182,8 @@ def run_session(self, req: AgentSessionRequest) -> AgentOutcome: def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = "") -> Reply: + role: str = "", effort: str = "", + max_budget_usd: float | None = None) -> Reply: """Tool-less one-shot in an empty scratch cwd (read-only sandbox + nothing to read = contained).""" import tempfile diff --git a/src/infermatrix_copilot/providers/cursor.py b/src/infermatrix_copilot/providers/cursor.py index 4113aa02..4835dc02 100644 --- a/src/infermatrix_copilot/providers/cursor.py +++ b/src/infermatrix_copilot/providers/cursor.py @@ -203,7 +203,8 @@ def run_session(self, req: AgentSessionRequest) -> AgentOutcome: def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = "") -> Reply: + role: str = "", effort: str = "", + max_budget_usd: float | None = None) -> Reply: """Tool-less one-shot. Runs in an EMPTY scratch cwd so cursor-agent's native tools have nothing to read — the containment for calls that need no repo at all (intent, reducer, repair).""" diff --git a/src/infermatrix_copilot/providers/deepseek.py b/src/infermatrix_copilot/providers/deepseek.py index 4f0ae018..ac3f1e89 100644 --- a/src/infermatrix_copilot/providers/deepseek.py +++ b/src/infermatrix_copilot/providers/deepseek.py @@ -462,7 +462,8 @@ def _on_notification(note) -> None: def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = "") -> Reply: + role: str = "", effort: str = "", + max_budget_usd: float | None = None) -> Reply: """Tool-less one-shot in an EMPTY scratch cwd — the containment for calls that need no repo at all (intent, reducer, repair). No bridge spec, so no MCP layer is written; bash and the editor see an empty diff --git a/src/infermatrix_copilot/providers/zcode.py b/src/infermatrix_copilot/providers/zcode.py index ba9cd5b7..a1dc3211 100644 --- a/src/infermatrix_copilot/providers/zcode.py +++ b/src/infermatrix_copilot/providers/zcode.py @@ -329,7 +329,8 @@ def run_session(self, req: AgentSessionRequest) -> AgentOutcome: def complete(self, *, system: str, messages: list[dict], model: str = "", max_tokens: int | None = None, - role: str = "", effort: str = "") -> Reply: + role: str = "", effort: str = "", + max_budget_usd: float | None = None) -> Reply: """Tool-less one-shot in an empty scratch cwd with no bridge. Every native read is removed, except `Read` when the prompt is too big for argv and rides an attachment — then any read outside the scratch dir diff --git a/test/test_kb_spend_cap.py b/test/test_kb_spend_cap.py new file mode 100644 index 00000000..d57e4302 --- /dev/null +++ b/test/test_kb_spend_cap.py @@ -0,0 +1,245 @@ +"""Per-call spend thresholds: `claude --max-budget-usd` on one-shots, and the +knowledge gateway refusing to dispatch a thresholded call to a transport that +cannot stop at one. The threshold is checked after each API request, so it +bounds the overshoot to one request; a hard ceiling is the caller's +reservation of threshold + one request's worst case. Fully offline.""" + +import json +import stat +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from infermatrix_copilot.config import Settings +from infermatrix_copilot.kb_service.models import ModelGateway, ModelRole, ModelUnavailable +from infermatrix_copilot.providers.base import HarnessTransport +from infermatrix_copilot.providers.claude_code import ClaudeCodeTransport +from infermatrix_copilot.providers.codex import CodexTransport +from infermatrix_copilot.providers.cursor import CursorTransport + +_FAKE_CLI = """#!/usr/bin/env python3 +import json, os, sys +here = os.path.dirname(os.path.abspath(__file__)) +sys.stdin.read() +with open(os.path.join(here, "capture.json"), "w") as f: + json.dump({"argv": sys.argv[1:]}, f) +over = os.path.exists(os.path.join(here, "over_budget")) +print(json.dumps({ + "type": "result", "subtype": "error_max_budget_usd" if over else "success", + "is_error": over, "result": "partial" if over else '{"ok": true}', + "num_turns": 1, "stop_reason": "end_turn", "total_cost_usd": 0.07, + "usage": {"input_tokens": 10, "output_tokens": 2}, + "modelUsage": {"claude-opus-5-5": {"costUSD": 0.07}}})) +""" + + +def _claude(tmp_path: Path) -> ClaudeCodeTransport: + cli = tmp_path / "bin" / "claude" + cli.parent.mkdir(exist_ok=True) + cli.write_text(_FAKE_CLI, encoding="utf-8") + cli.chmod(cli.stat().st_mode | stat.S_IXUSR) + return ClaudeCodeTransport(Settings(_env_file=None, strict_backend="claude-code", + strict_backend_cli=str(cli))) + + +def _argv(tmp_path: Path) -> list[str]: + return json.loads((tmp_path / "bin" / "capture.json").read_text(encoding="utf-8"))["argv"] + + +MESSAGES = [{"role": "user", "content": "hi"}] + + +def test_claude_passes_the_cap_only_when_given(tmp_path): + transport = _claude(tmp_path) + transport.complete(system="S", messages=MESSAGES) + assert "--max-budget-usd" not in _argv(tmp_path) + + reply = transport.complete(system="S", messages=MESSAGES, max_budget_usd=0.5) + argv = _argv(tmp_path) + assert argv[argv.index("--max-budget-usd") + 1] == "0.5" + assert reply.usage["cost_usd"] == pytest.approx(0.07) + assert reply.stop_reason == "end_turn" + + +def test_claude_reports_a_call_stopped_at_its_cap(tmp_path): + transport = _claude(tmp_path) + (tmp_path / "bin" / "over_budget").write_text("", encoding="utf-8") + reply = transport.complete(system="S", messages=MESSAGES, max_budget_usd=0.01) + assert reply.stop_reason == "max_budget" + assert reply.text == "" # a budget-truncated answer is never usable + + +def test_claude_rejects_a_non_positive_cap(tmp_path): + with pytest.raises(ValueError): + _claude(tmp_path).complete(system="S", messages=MESSAGES, max_budget_usd=0) + + +def test_only_claude_code_claims_to_stop_at_a_threshold(): + assert ClaudeCodeTransport.stops_at_spend is True + assert HarnessTransport.stops_at_spend is False + assert CodexTransport.stops_at_spend is False + assert CursorTransport.stops_at_spend is False + + +def test_no_transport_advertises_a_hard_cap(): + """Regression: the CLI checks its budget after a request completes, so no + transport may claim to bound a single call's spend. The old name + (`supports_spend_cap`) promised exactly that and must not come back.""" + for cls in (HarnessTransport, ClaudeCodeTransport, CodexTransport, CursorTransport): + assert not hasattr(cls, "supports_spend_cap") + + +class _Transport: + """A gateway-level fake: records the kwargs it was called with.""" + + def __init__(self, *, capped: bool, text: str = '{"ok": true}', stop: str = "end_turn", + usage: dict | None = None): + self.stops_at_spend = capped + self.calls: list[dict] = [] + self._reply = SimpleNamespace(blocks=[SimpleNamespace(text=text)], stop_reason=stop, + usage=usage if usage is not None else {"cost_usd": 0.12}, + model="served-model") + + def complete(self, **kwargs): + self.calls.append(kwargs) + return self._reply + + +ROLE = ModelRole("generator", "claude-code", "claude-opus-5-5") + + +def _gateway(transport, records: list | None = None) -> ModelGateway: + return ModelGateway(None, transport_factory=lambda provider: transport, + recorder=(records.append if records is not None else None)) + + +def test_gateway_refuses_an_uncapped_transport_before_dispatch(): + transport = _Transport(capped=False) + with pytest.raises(ModelUnavailable, match="cannot stop a call at a spend threshold"): + _gateway(transport).call_json(ROLE, system="S", prompt="P", max_budget_usd=1.0) + assert transport.calls == [] # fail-closed: nothing was sent + + +def test_gateway_passes_the_cap_and_surfaces_cost(): + transport = _Transport(capped=True) + records: list = [] + reply = _gateway(transport, records).call_json(ROLE, system="S", prompt="P", max_budget_usd=1.5) + assert transport.calls[0]["max_budget_usd"] == 1.5 + assert reply.cost_usd == pytest.approx(0.12) + assert records[-1]["cost_usd"] == pytest.approx(0.12) + assert records[-1]["max_budget_usd"] == 1.5 + + +def test_gateway_without_a_cap_is_unchanged(): + """Existing callers pass no cap: the kwarg is not sent at all, so fakes + and transports with the old signature keep working.""" + transport = _Transport(capped=False, usage={}) + reply = _gateway(transport).call_json(ROLE, system="S", prompt="P") + assert "max_budget_usd" not in transport.calls[0] + assert reply.cost_usd is None + assert reply.data == {"ok": True} + + +def test_gateway_treats_a_call_stopped_at_its_cap_as_unavailable(): + transport = _Transport(capped=True, text="", stop="max_budget") + records: list = [] + with pytest.raises(ModelUnavailable, match="spend threshold"): + _gateway(transport, records).call_json(ROLE, system="S", prompt="P", max_budget_usd=0.2) + assert records[-1]["error"] # recorded as a failure, never a training example + + +def test_gateway_rejects_a_non_positive_cap(): + with pytest.raises(ModelUnavailable, match="positive"): + _gateway(_Transport(capped=True)).call_json(ROLE, system="S", prompt="P", max_budget_usd=-1) + + +# -- exact threshold argument ---------------------------------------------------- + +@pytest.mark.parametrize("value, expected", [ + (0.00001, "0.00001"), # `.4f` made this "0.0000": no threshold at all + (0.12345, "0.12345"), # `.4f` made this "0.1235": a HIGHER threshold + (1e-9, "0.000000001"), # never scientific notation + (0.5, "0.5"), + (3, "3.0"), + (1234.5678, "1234.5678"), +]) +def test_budget_arg_is_exact(value, expected): + from decimal import Decimal + + from infermatrix_copilot.providers.claude_code import budget_arg + + arg = budget_arg(value) + assert arg == expected + assert Decimal(arg) > 0 and float(arg) == float(value) # never rounded to 0, never moved + + +@pytest.mark.parametrize("bad", [0, -0.1, float("nan"), float("inf"), True, "0.5"]) +def test_budget_arg_rejects_non_positive_or_non_finite(bad): + from infermatrix_copilot.providers.claude_code import budget_arg + + with pytest.raises(ValueError): + budget_arg(bad) + + +def test_claude_passes_a_tiny_threshold_unrounded(tmp_path): + _claude(tmp_path).complete(system="S", messages=MESSAGES, max_budget_usd=0.00001) + argv = _argv(tmp_path) + assert argv[argv.index("--max-budget-usd") + 1] == "0.00001" + + +# -- persisted through the REAL recorder ------------------------------------------- + +def _store_gateway(tmp_path, transport): + from infermatrix_copilot.kb_service.runtime import trace_recorder + from infermatrix_copilot.trace_store import TraceStore + + store = TraceStore(tmp_path / "traces", environ={}) + return store, ModelGateway(None, transport_factory=lambda provider: transport, + recorder=trace_recorder(store)) + + +def test_trace_persists_threshold_and_cost(tmp_path): + store, gateway = _store_gateway(tmp_path, _Transport(capped=True)) + gateway.call_json(ROLE, system="S", prompt="P", max_budget_usd=0.00001) + (record,) = store.query(kind="model_call") + assert record["result"]["max_budget_usd"] == 0.00001 + assert record["result"]["cost_usd"] == pytest.approx(0.12) + + +def test_trace_persists_threshold_on_transport_failure(tmp_path): + class Exploding(_Transport): + def complete(self, **kwargs): + raise RuntimeError("cli crashed") + + store, gateway = _store_gateway(tmp_path, Exploding(capped=True)) + with pytest.raises(ModelUnavailable): + gateway.call_json(ROLE, system="S", prompt="P", max_budget_usd=0.25) + (record,) = store.query(kind="model_call") + assert record["error"] + assert record["result"]["max_budget_usd"] == 0.25 + assert record["result"]["cost_usd"] is None + + +def test_trace_without_a_threshold_records_none(tmp_path): + store, gateway = _store_gateway(tmp_path, _Transport(capped=False, usage={})) + gateway.call_json(ROLE, system="S", prompt="P") + (record,) = store.query(kind="model_call") + assert record["result"]["max_budget_usd"] is None + assert record["result"]["cost_usd"] is None + + +def test_every_transport_accepts_a_spend_threshold(): + """Every one-shot transport takes ``max_budget_usd`` (most accept and + ignore it); only one that stops at a threshold may claim to.""" + import inspect + + from infermatrix_copilot.providers import deepseek, zcode + + transports = [HarnessTransport, ClaudeCodeTransport, CodexTransport, CursorTransport] + transports += [obj for module in (deepseek, zcode) for obj in vars(module).values() + if isinstance(obj, type) and issubclass(obj, HarnessTransport) + and obj is not HarnessTransport] + for cls in transports: + assert "max_budget_usd" in inspect.signature(cls.complete).parameters, cls.__name__ + assert cls.stops_at_spend is (cls is ClaudeCodeTransport), cls.__name__