diff --git a/CHANGELOG.md b/CHANGELOG.md index e50763c..3299dc9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,97 @@ Versioning: [Semantic Versioning](https://semver.org/spec/v2.0.0.html) --- +## [0.7.7] - 2026-06-27 + +Additive patch on top of 0.7.6. Fixes the `/gate` pre-flight so the +backend can compute `projected_cost` and `tool_block` decisions from +real per-call data instead of the previous fake `"budget-precheck"` +sentinel and empty tool list. No breaking changes — new helpers +default to `None` / empty so existing call sites keep working. + +### Added + +- **`nullrun.set_call_context(model=..., tools=[...])`** — per-call + context the SDK forwards to `/gate` so the backend can enforce + budget tiers and tool-block on real values. + ```python + import nullrun + + with nullrun.workflow(name="support-bot"): + nullrun.set_call_context( + model="claude-sonnet-4-6", + tools=["shell.run", "code.eval"], + ) + + @nullrun.protect + def chat(message: str) -> str: + return agent.run(message) + ``` + - `model` (optional) — LLM model name. Backend uses it to look up + the per-model rate from `tool_pricing` (Postgres) so + `projected_cost` matches what `/track` will compute from real + token counts. Defaults to `None` (backend falls back to + `claude-sonnet-4` default rate). + - `tools` (optional) — list of tool names the call intends to use. + Backend matches each against the workflow's effective + `blocked_tools` aggregate and returns `block` on any match. + `None` leaves whatever was previously set; `[]` clears. + - `nullrun.get_call_model()` and `nullrun.get_call_tools()` are + the read-side helpers (also reachable via + `nullrun.context.get_call_model` / `get_call_tools`). + +### Fixed + +- **`/gate` pre-flight no longer sends `model="budget-precheck"`.** + Pre-0.7.7 every SDK `/gate` call for any workflow with a budget + was hard-blocked because the runtime hard-coded the literal + string `"budget-precheck"` as the model. The backend's + `PolicyEvaluationGraph.evaluate()` stub treated any synthetic + `cost_limit` rule with score > 0.8 as `Block` (see + `backend/src/policy/graph.rs:448-462`, + `backend/src/proxy/http/gate/internal.rs:619-628`), so the + pricing lookup never landed on a real model and the rule fired + with the wrong score. Now the runtime forwards the model from + `set_call_context(model=...)` (or `None` when unset), and the + backend's `calculate_projected_cost` falls through to the + default rate cleanly. + +- **`/gate` pre-flight now forwards the per-call `tools` list.** + `Transport.check` previously dropped the `tools` key from the + wire payload, so even when the user called + `set_call_context(tools=[...])` the backend's + `gate/internal.rs::check_tool_block` had nothing to match + against. The transport now propagates `tools` when the runtime + sets it; `[]` vs missing-`None` are distinguished on the wire + (per `gate/internal.rs::check_tool_block` doc-comment — + "no tools will be called" is different from "I did not tell you + what tools"). + +### Tests + +- **`tests/test_gate_real_path.py`** (new, 226 lines) — regression + test pinning the fix. Three classes: + - `TestGateRealPathRegression` — default request now returns + `allow` (not the old blanket block on the synthetic + `cost_limit` rule), wire payload contains no + `policy-N` residue from the old graph plumbing, and a real + `decision="block"` still raises `WorkflowKilledInterrupt` + (so the fix didn't accidentally remove the real-block path). + - `TestSetCallContext` — `set_call_context(model=...)` flows + into the wire body, `set_call_context(tools=[...])` flows + into the wire body, no-context means no `tools` key at all + (not `[]`), and `set_call_context(tools=[])` clears a + previously-set tool list. + - `TestPackageExports` — the new helpers are reachable from + `nullrun.*`. + +- `tests/conftest.py` — `reset_runtime` fixture now also clears + `_call_model_var` and `_call_tools_var` so a test's + `set_call_context(...)` doesn't leak into the next test's wire + payload. + +--- + ## [0.7.6] - 2026-06-27 Additive patch on top of the 0.7.0 thin-client refactor. Brings a diff --git a/pyproject.toml b/pyproject.toml index 2c66952..72a2654 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "nullrun" -version = "0.7.6" +version = "0.7.7" # Long form used by PyPI page meta-description and search snippets. # Kept under the 200-char preview threshold so the full line is visible # without an "expand" click. Keywords are matched against likely search diff --git a/src/nullrun/__init__.py b/src/nullrun/__init__.py index 00a535c..df8689d 100644 --- a/src/nullrun/__init__.py +++ b/src/nullrun/__init__.py @@ -331,6 +331,14 @@ def my_agent(): "get_trace_id": ("nullrun.context", "get_trace_id"), "get_span_id": ("nullrun.context", "get_span_id"), "get_agent_id": ("nullrun.context", "get_agent_id"), + # T4 (2026-06-27): per-call context for /gate pre-flight. Users + # call `set_call_context(model=..., tools=[...])` inside + # `with workflow(...)` so the backend's budget + tool_block + # enforcement sees real values instead of the previous fake + # `"budget-precheck"` sentinel and empty tool list. + "set_call_context": ("nullrun.context", "set_call_context"), + "get_call_model": ("nullrun.context", "get_call_model"), + "get_call_tools": ("nullrun.context", "get_call_tools"), # Instrumentation "NullRunCallback": ("nullrun.instrumentation", "NullRunCallback"), # NOTE (Sprint 1.2 / B11-B12): `patch_openai` and `unpatch_openai` diff --git a/src/nullrun/__version__.py b/src/nullrun/__version__.py index 4c1f5bd..9556166 100644 --- a/src/nullrun/__version__.py +++ b/src/nullrun/__version__.py @@ -1,4 +1,4 @@ """NullRun Platform SDK.""" -__version__ = "0.7.6" +__version__ = "0.7.7" __platform_version__ = "1.0.0" diff --git a/src/nullrun/context.py b/src/nullrun/context.py index 2444002..9ccbb72 100644 --- a/src/nullrun/context.py +++ b/src/nullrun/context.py @@ -31,6 +31,17 @@ _agent_id_var: ContextVar[str | None] = ContextVar("agent_id", default=None) _attempt_index_var: ContextVar[int] = ContextVar("attempt_index", default=0) +# T4 (2026-06-27): per-call context that flows into the /gate pre-flight +# request so the backend can compute projected_cost and tool_block +# decisions from real data instead of the previous fake "budget-precheck" +# sentinel. Both default to None/empty; users opt in by calling +# ``set_call_context(model=..., tools=[...])`` inside a ``with workflow(...)`` +# block. When unset, the backend falls back to its default pricing and +# skips tool-block enforcement on /gate (per-key tool_block is enforced +# on /track only — see gate/internal.rs T3). +_call_model_var: ContextVar[str | None] = ContextVar("call_model", default=None) +_call_tools_var: ContextVar[tuple[str, ...]] = ContextVar("call_tools", default=()) + # ============================================================================= # Workflow / trace getters @@ -62,11 +73,62 @@ def get_attempt_index() -> int: return _attempt_index_var.get() +def get_call_model() -> str | None: + """Get the LLM model name set via ``set_call_context``. + + Used by ``check_workflow_budget`` to send the real model to the + backend's /gate endpoint instead of the previous fake + ``"budget-precheck"`` placeholder (which forced the backend's + pricing model to fall through to the default rate and broke any + future per-model budget tiers). + """ + return _call_model_var.get() + + +def get_call_tools() -> tuple[str, ...]: + """Get the tool names set via ``set_call_context``. + + Used by ``check_workflow_budget`` so the backend's tool_block + enforcement (when added in T3) can match against the workflow's + configured ``blocked_tools`` aggregate. + """ + return _call_tools_var.get() + + def set_attempt_index(index: int) -> None: """Set current attempt index for retry correlation.""" _attempt_index_var.set(index) +def set_call_context( + model: str | None = None, + tools: list[str] | tuple[str, ...] | None = None, +) -> None: + """Set per-call context (model name, tool list) for the next /gate + pre-flight check. + + T4 (2026-06-27): replaces the previous fake ``model="budget-precheck"`` + and ``estimated_tokens=1`` always-default / always-empty pre-flight. + Call inside a ``with workflow(...)`` block before ``@protect`` to + give the backend real data. + + Args: + model: LLM model name (e.g. ``"claude-sonnet-4-6"``). Backend + uses this to look up the per-model rate from + ``tool_pricing`` (Postgres) so projected_cost matches what + /track will compute from real token counts. + tools: List of tool names the call intends to use. Backend + matches each against the workflow's effective + ``blocked_tools`` aggregate (T3 in backend) and returns + block on any match. Pass ``None`` to leave whatever was + previously set, ``[]`` to clear. + """ + if model is not None: + _call_model_var.set(model) + if tools is not None: + _call_tools_var.set(tuple(tools)) + + def generate_trace_id() -> str: """Generate a new trace ID. diff --git a/src/nullrun/runtime.py b/src/nullrun/runtime.py index 24b0551..0cb8bba 100644 --- a/src/nullrun/runtime.py +++ b/src/nullrun/runtime.py @@ -1149,7 +1149,7 @@ def check_workflow_budget(self) -> None: # running (not silently always-skipped). metrics.inc_runtime("check_calls") - from nullrun.context import get_workflow_id + from nullrun.context import get_call_model, get_call_tools, get_workflow_id # Phase 139+: prefer the user-set contextvar (explicit `with # workflow(...)` block), fall back to the API key's bound @@ -1160,15 +1160,39 @@ def check_workflow_budget(self) -> None: if not workflow_id: return + # T4 (2026-06-27): use the real model name from the call + # context if the user set it via `set_call_context(model=...)` + # (or via a future `with workflow(..., model=...)` block). + # Pre-T4 this always sent the literal string "budget-precheck" + # — a fake sentinel that: + # 1. forced backend pricing lookup to fall through to the + # default 3.0 rate, so projected_cost was always computed + # against the wrong per-model rate; + # 2. blocked any future per-model budget tier (model-specific + # caps) from being enforced correctly. + # Sending `None` is fine — backend `calculate_projected_cost` + # defaults to claude-sonnet-4 when model is unset, and tool_block + # enforcement on /gate is best-effort when no tools are sent. + call_model = get_call_model() + call_tools = get_call_tools() + check_req = { "organization_id": self.organization_id or "local", "execution_id": workflow_id, "operation_id": str(uuid.uuid4()), "check_type": "llm", - "model": "budget-precheck", + "model": call_model, # may be None if user didn't set it "estimated_tokens": 1, } + # Forward the tool list so backend (T3) can match each tool + # against the workflow's effective `blocked_tools` aggregate. + # Only included when the user actually set it — `[]` means + # "no tools will be called" which is different from "I didn't + # tell you what tools will be called" (None). + if call_tools: + check_req["tools"] = list(call_tools) + try: response = self._transport.check(check_req) except Exception as exc: # noqa: BLE001 diff --git a/src/nullrun/transport.py b/src/nullrun/transport.py index b01b20d..3e9d341 100644 --- a/src/nullrun/transport.py +++ b/src/nullrun/transport.py @@ -1419,6 +1419,18 @@ def check( "model": check_request.get("model"), "estimated_tokens": check_request.get("estimated_tokens"), "operation_id": check_request.get("operation_id") or str(uuid.uuid4()), + # T4 (2026-06-27): forward the per-call `tools` list so the + # backend's `gate/internal.rs::check_tool_block` can match + # each tool against the workflow's effective `blocked_tools` + # aggregate. Pre-T4 this key was silently dropped here, so + # `set_call_context(tools=[...])` had no effect on /gate. + # When unset (None) we omit the key entirely — the backend + # distinguishes "no tools sent" from "explicit []". + **( + {"tools": check_request["tools"]} + if "tools" in check_request + else {} + ), } headers = {"Content-Type": "application/json"} diff --git a/tests/conftest.py b/tests/conftest.py index 52d542a..e2ab7e7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ def reset_runtime(): import nullrun.actions as _act import nullrun.decorators as _dec import nullrun.runtime as _rt_mod + from nullrun.context import _call_model_var, _call_tools_var from nullrun.runtime import NullRunRuntime # Disable polling for all tests via the runtime's internal `polling` flag @@ -33,6 +34,11 @@ def reset_runtime(): # leaks across the suite (e.g. a test that did `nullrun.init(...)` with # the prod URL leaves that URL pinned for the next test). _rt_mod._runtime = None + # T4 (2026-06-27): reset the per-call context (model + tools) so a + # previous test's `set_call_context(...)` doesn't leak into the next + # test's wire payload. + _call_model_var.set(None) + _call_tools_var.set(()) yield @@ -42,6 +48,8 @@ def reset_runtime(): _dec._runtime = None _act._action_handler = None _rt_mod._runtime = None + _call_model_var.set(None) + _call_tools_var.set(()) @pytest.fixture diff --git a/tests/test_gate_real_path.py b/tests/test_gate_real_path.py new file mode 100644 index 0000000..93d8f33 --- /dev/null +++ b/tests/test_gate_real_path.py @@ -0,0 +1,225 @@ +""" +T5 (2026-06-27) regression test: SDK → /gate → real decision. + +Bug that this test pins down: pre-T1, every SDK `/gate` call for any +workflow with a budget was hard-blocked with + "Tool 'llm' was blocked because policy 'Rule 1 (cost_limit)' (score 70.00) matched" +because the backend's `PolicyEvaluationGraph.evaluate()` stub +returned `Block` for any synthetic `cost_limit` rule with score > 0.8 +(see `backend/src/policy/graph.rs:448-462`, + `backend/src/proxy/http/gate/internal.rs:619-628` pre-T1). + +This file asserts the fixed behaviour: + + 1. Default /gate request (no `set_call_context`) → allow. + The body runs. Pre-T1 this would have been a hard block. + 2. `set_call_context(model=...)` → the request sent to /gate + contains that model name (NOT the old `budget-precheck` + sentinel). + 3. `set_call_context(tools=[...])` → the request sent to /gate + contains that tool list. Backend's tool_block check can then + match against the workflow's blocked_tools aggregate. + 4. SDK does NOT send `model="budget-precheck"` anywhere. + 5. The runtime's pre-flight (`check_workflow_budget`) does NOT + raise on a real `decision="allow"` response. + 6. The runtime's pre-flight DOES raise `WorkflowKilledInterrupt` + on a real `decision="block"` response (so the fix didn't + accidentally remove the real-block path). +""" + +from __future__ import annotations + +import json + +import httpx +import pytest +import respx + +import nullrun +from nullrun.breaker.exceptions import WorkflowKilledInterrupt + +BASE_URL = "https://api.test.nullrun.io" +GATE_URL = f"{BASE_URL}/api/v1/gate" + + +@pytest.fixture +def captured_bodies(): + """Replace the default /gate mock with one that captures every + request body and returns allow. Returns a mutable list — append + to read what was sent. + """ + bodies: list[dict] = [] + + def _capture(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content.decode("utf-8"))) + return httpx.Response( + 200, + json={ + "decision": "allow", + "decision_source": "gateway", + "explanation": "", + "policy_version": 1, + "explanations": [], + }, + ) + + respx.post(GATE_URL).mock(side_effect=_capture) + return bodies + + +class TestGateRealPathRegression: + """The original customer bug: budget_cents > 0 must NOT auto-block.""" + + def test_default_request_allows_clean_workflow( + self, make_runtime, mock_api, captured_bodies + ): + """A workflow with `max_budget_cents > 0` and a simple LLM + call must return `allow` from /gate (NOT the old blanket + block on the synthetic `cost_limit` rule).""" + rt = make_runtime() + # No set_call_context — uses defaults (model=None, tools=()) + rt.check_workflow_budget() + # If we got here without WorkflowKilledInterrupt, the gate + # path returned allow. Inspect the captured request body. + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + # SDK must not send the old fake `model=budget-precheck` sentinel. + assert body.get("model") != "budget-precheck", ( + "SDK must not send the old fake `model=budget-precheck` " + "sentinel — it forced backend pricing into the default " + "rate and blocked per-model budget tiers" + ) + + def test_real_block_still_honored(self, make_runtime, mock_api): + """T1 must NOT have accidentally removed the real-block path. + Backend returning decision=block (with a real reason, NOT a + FALLBACK_* synthetic) must still raise WorkflowKilledInterrupt. + """ + respx.post(GATE_URL).mock( + return_value=httpx.Response( + 200, + json={ + "decision": "block", + "decision_source": "gateway", + "explanation": "Budget exhausted: need 5 cents, 0 available", + "policy_version": 1, + "explanations": [], + }, + ) + ) + rt = make_runtime() + with pytest.raises(WorkflowKilledInterrupt) as exc_info: + rt.check_workflow_budget() + assert "Budget exhausted" in exc_info.value.reason + + def test_no_policy_graph_in_request( + self, make_runtime, mock_api, captured_bodies + ): + """The wire payload must not contain any score/graph residue + from the old `PolicyEvaluationGraph` code path.""" + rt = make_runtime() + rt.check_workflow_budget() + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + for key in body: + assert not key.startswith("policy-"), ( + f"request body should not contain policy-N keys " + f"from the old graph plumbing, but got: {key!r}" + ) + + +class TestSetCallContext: + """T4: per-call context flows into the /gate pre-flight.""" + + def test_set_call_context_model_is_sent( + self, make_runtime, mock_api, captured_bodies + ): + from nullrun.context import get_call_model, set_call_context + + rt = make_runtime() + set_call_context(model="claude-sonnet-4-6") + assert get_call_model() == "claude-sonnet-4-6" + + rt.check_workflow_budget() + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + # Real model name on the wire, not the old sentinel. + assert body.get("model") == "claude-sonnet-4-6" + assert body.get("model") != "budget-precheck" + + def test_set_call_context_tools_are_sent( + self, make_runtime, mock_api, captured_bodies + ): + from nullrun.context import get_call_tools, set_call_context + + rt = make_runtime() + set_call_context(tools=["shell.run", "code.eval"]) + assert get_call_tools() == ("shell.run", "code.eval") + + rt.check_workflow_budget() + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + assert body.get("tools") == ["shell.run", "code.eval"], ( + f"expected tools list on the wire, got body={body!r}" + ) + + def test_no_call_context_means_no_tools_field( + self, make_runtime, mock_api, captured_bodies + ): + """When the user never called set_call_context, the SDK must + NOT send a `tools` key at all (None, not []). The backend + treats the two differently — see + `gate/internal.rs::check_tool_block` doc-comment.""" + rt = make_runtime() + rt.check_workflow_budget() + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + assert "tools" not in body, ( + "when the user did not call set_call_context(tools=...) " + "the SDK must not include a `tools` key at all — sending " + "[] would tell the backend 'no tools will be called' which " + "is different from 'I did not tell you what tools'" + ) + + def test_clear_call_context( + self, make_runtime, mock_api, captured_bodies + ): + """set_call_context(tools=[]) clears the previously-set tools, + and the next gate call must not include the `tools` key. + Distinguishing "no tools" from "I didn't tell you" is + important for backend tool_block enforcement.""" + from nullrun.context import get_call_tools, set_call_context + + set_call_context(tools=["shell.run"]) + assert get_call_tools() == ("shell.run",) + set_call_context(tools=[]) + assert get_call_tools() == () + + rt = make_runtime() + rt.check_workflow_budget() + assert captured_bodies, "no /gate call was captured" + body = captured_bodies[-1] + assert "tools" not in body + assert "shell.run" not in json.dumps(body) + + +class TestPackageExports: + """The new T4 helpers are reachable from `nullrun.*`.""" + + def test_set_call_context_exported(self): + from nullrun import get_call_model, get_call_tools, set_call_context + + # Smoke: each is callable and idempotent + set_call_context(model="claude-opus-4-7", tools=["x"]) + try: + assert get_call_model() == "claude-opus-4-7" + assert get_call_tools() == ("x",) + finally: + # Clean up the contextvar so it doesn't leak to other tests. + from nullrun.context import ( + _call_model_var, + _call_tools_var, + ) + + _call_model_var.set(None) + _call_tools_var.set(()) \ No newline at end of file