From f6df57a8458a33368b5ea9b1a65780ff332e02aa Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 5 Jun 2026 22:14:59 +0000 Subject: [PATCH 1/2] fix(pickled-iac): install BudgetGuard via build_default_client in CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CLI helper bypassed pickled_core.llm.bootstrap.build_default_client and called build_client directly. As a side effect, _maybe_install_budget never ran, so a user-configured budget.max_cost_usd cap (in pickled.config.yaml or the PICKLED_MAX_COST_USD env var) was silently ignored by 'pickled-iac draft'. IaCDrafter retries the LLM up to 3 times per invocation, so a configured $0.50 cap could be overrun several-fold before the user noticed. Route the CLI helper through build_default_client (factory_env= PICKLED_IAC_LLM_FACTORY), matching pickled-bdd / pickled-data / pickled-diff and the MCP CLI of this same package. This also enables the disk-backed response cache when configured. Adds three regression tests pinning that the guard is installed from both yaml and the env override, and that the test-only LLM factory shortcut still works. Co-authored-by: Bartłomiej Rosa --- packages/pickled-iac/src/pickled_iac/cli.py | 26 ++-- .../tests/test_cli_budget_bootstrap.py | 119 ++++++++++++++++++ 2 files changed, 127 insertions(+), 18 deletions(-) create mode 100644 packages/pickled-iac/tests/test_cli_budget_bootstrap.py diff --git a/packages/pickled-iac/src/pickled_iac/cli.py b/packages/pickled-iac/src/pickled_iac/cli.py index c0e1b13..1b6c827 100644 --- a/packages/pickled-iac/src/pickled_iac/cli.py +++ b/packages/pickled-iac/src/pickled_iac/cli.py @@ -2,15 +2,13 @@ from __future__ import annotations -import importlib import json -import os import tempfile from pathlib import Path -from typing import cast import click from pickled_core import LLMClient, Verdict +from pickled_core.llm.bootstrap import build_default_client from pickled_core.llm.config import ConfigError from pickled_iac.drafter import IaCDrafter @@ -20,22 +18,14 @@ def _build_llm_client() -> LLMClient: - factory = os.environ.get("PICKLED_IAC_LLM_FACTORY") - if factory: - module_name, sep, attr = factory.partition(":") - if not sep: - raise click.ClickException( - "PICKLED_IAC_LLM_FACTORY must be 'module:callable'" - ) - module = importlib.import_module(module_name) - return cast(LLMClient, getattr(module, attr)()) - - from pickled_core.llm.config import load_config - from pickled_core.llm.factory import build_client - - provider = os.environ.get("PICKLED_LLM_PROVIDER", "anthropic") + """Build an LLM client honoring config, env, cache, and budget. + + Uses :func:`build_default_client` so the user's ``budget.max_cost_usd`` + cap (and disk cache) installed via ``pickled.config.yaml`` / + ``PICKLED_MAX_COST_USD`` are actually enforced for ``pickled-iac draft``. + """ try: - return build_client(provider, config=load_config()) + return build_default_client(factory_env="PICKLED_IAC_LLM_FACTORY") except ConfigError as exc: raise click.ClickException(str(exc)) from exc diff --git a/packages/pickled-iac/tests/test_cli_budget_bootstrap.py b/packages/pickled-iac/tests/test_cli_budget_bootstrap.py new file mode 100644 index 0000000..87c4b9a --- /dev/null +++ b/packages/pickled-iac/tests/test_cli_budget_bootstrap.py @@ -0,0 +1,119 @@ +"""Regression: ``pickled-iac`` CLI must honor ``budget.max_cost_usd`` cap. + +The bug: an earlier ``_build_llm_client`` called ``build_client`` directly, +which silently skipped :func:`_maybe_install_budget` from +:mod:`pickled_core.llm.bootstrap`. A user setting +``budget.max_cost_usd: "0.50"`` in ``pickled.config.yaml`` (or +``PICKLED_MAX_COST_USD=0.50``) saw the cap apply to ``pickled-bdd`` / +``pickled-data`` / ``pickled-diff`` and every MCP server, but ``pickled-iac +draft`` (which retries the LLM up to 3 times) ignored it. + +These tests pin the fix: the CLI now routes through +:func:`build_default_client`, which always installs a :class:`BudgetGuard` +when a cap is configured. +""" + +from __future__ import annotations + +from collections.abc import Generator +from decimal import Decimal +from pathlib import Path + +import pytest + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key in ( + "PICKLED_IAC_LLM_FACTORY", + "PICKLED_CACHE_DIR", + "PICKLED_CACHE_MODE", + "PICKLED_MAX_COST_USD", + "PICKLED_LLM_PROVIDER", + ): + monkeypatch.delenv(key, raising=False) + + +@pytest.fixture(autouse=True) +def _reset_budget_guard() -> Generator[None, None, None]: + from pickled_core.llm.budget_context import set_budget_guard + + set_budget_guard(None) + yield + set_budget_guard(None) + + +def _write_config(tmp_path: Path, cap: str) -> Path: + cfg = tmp_path / "pickled.config.yaml" + cfg.write_text( + f"""providers: + anthropic: + type: anthropic + default_model: claude-test + api_key_env: TEST_API_KEY +budget: + max_cost_usd: \"{cap}\" +cache: + mode: off +""", + encoding="utf-8", + ) + return cfg + + +def test_iac_cli_install_budget_guard_from_yaml( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``_build_llm_client`` must install BudgetGuard from yaml cap.""" + pytest.importorskip("anthropic") + _write_config(tmp_path, "0.50") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("TEST_API_KEY", "fake") + + from pickled_core.llm.budget_context import active_budget_guard + from pickled_iac.cli import _build_llm_client + + assert active_budget_guard() is None + _build_llm_client() + guard = active_budget_guard() + assert guard is not None, ( + "pickled-iac CLI bypassed the budget guard; users' max_cost_usd cap " + "would not have been enforced (see the docstring of this module)." + ) + assert guard._budget.max_cost_usd == Decimal("0.50") + + +def test_iac_cli_install_budget_guard_from_env( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``PICKLED_MAX_COST_USD`` env var must override yaml and install guard.""" + pytest.importorskip("anthropic") + _write_config(tmp_path, "0.50") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("TEST_API_KEY", "fake") + monkeypatch.setenv("PICKLED_MAX_COST_USD", "2.00") + + from pickled_core.llm.budget_context import active_budget_guard + from pickled_iac.cli import _build_llm_client + + _build_llm_client() + guard = active_budget_guard() + assert guard is not None + assert guard._budget.max_cost_usd == Decimal("2.00") + + +def test_iac_cli_factory_env_shortcut_still_works( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The ``PICKLED_IAC_LLM_FACTORY`` shortcut must still work for tests.""" + monkeypatch.setenv( + "PICKLED_IAC_LLM_FACTORY", + "pickled_bdd.testing:build_fake_llm", + ) + from pickled_bdd.testing import CannedLLMClient + from pickled_iac.cli import _build_llm_client + + client = _build_llm_client() + assert isinstance(client, CannedLLMClient) From 737eb29b626f530d3921af1122c01a4be524f39f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 5 Jun 2026 22:15:06 +0000 Subject: [PATCH 2/2] fix(pickled-schema): install BudgetGuard via build_default_client in CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same bug pattern as the prior pickled-iac commit: the CLI _build_llm_client called pickled_core.llm.factory.build_client directly, which skipped _maybe_install_budget. A configured budget.max_cost_usd cap (yaml or PICKLED_MAX_COST_USD) was silently ignored by 'pickled-schema draft'. OpenAPIDrafter retries up to 3 times per invocation, so the bypass had the same blast radius. Route through build_default_client (factory_env=PICKLED_SCHEMA_LLM_FACTORY) to match pickled-bdd / pickled-data / pickled-diff and this package's own MCP CLI. Disk cache is now also honored when configured. Adds three regression tests pinning guard installation and the test-only factory shortcut. Co-authored-by: Bartłomiej Rosa --- .../pickled-schema/src/pickled_schema/cli.py | 27 ++--- .../tests/test_cli_budget_bootstrap.py | 111 ++++++++++++++++++ 2 files changed, 120 insertions(+), 18 deletions(-) create mode 100644 packages/pickled-schema/tests/test_cli_budget_bootstrap.py diff --git a/packages/pickled-schema/src/pickled_schema/cli.py b/packages/pickled-schema/src/pickled_schema/cli.py index 9f1667d..dbd838f 100644 --- a/packages/pickled-schema/src/pickled_schema/cli.py +++ b/packages/pickled-schema/src/pickled_schema/cli.py @@ -2,14 +2,12 @@ from __future__ import annotations -import importlib import json -import os from pathlib import Path -from typing import cast import click from pickled_core import LLMClient, Verdict +from pickled_core.llm.bootstrap import build_default_client from pickled_core.llm.config import ConfigError from pickled_schema.gates import SchemaCoverageFinding, SchemaCoverageGate @@ -56,22 +54,15 @@ def _load_artifact(path: Path, fmt: SchemaFormat) -> SchemaArtifact: def _build_llm_client() -> LLMClient: - factory = os.environ.get("PICKLED_SCHEMA_LLM_FACTORY") - if factory: - module_name, sep, attr = factory.partition(":") - if not sep: - raise click.ClickException( - "PICKLED_SCHEMA_LLM_FACTORY must be 'module:callable'" - ) - module = importlib.import_module(module_name) - return cast(LLMClient, getattr(module, attr)()) - - from pickled_core.llm.config import load_config - from pickled_core.llm.factory import build_client - - provider = os.environ.get("PICKLED_LLM_PROVIDER", "anthropic") + """Build an LLM client honoring config, env, cache, and budget. + + Uses :func:`build_default_client` so the user's ``budget.max_cost_usd`` + cap (and disk cache) installed via ``pickled.config.yaml`` / + ``PICKLED_MAX_COST_USD`` are actually enforced for ``pickled-schema + draft``. + """ try: - return build_client(provider, config=load_config()) + return build_default_client(factory_env="PICKLED_SCHEMA_LLM_FACTORY") except ConfigError as exc: raise click.ClickException(str(exc)) from exc diff --git a/packages/pickled-schema/tests/test_cli_budget_bootstrap.py b/packages/pickled-schema/tests/test_cli_budget_bootstrap.py new file mode 100644 index 0000000..6d35a28 --- /dev/null +++ b/packages/pickled-schema/tests/test_cli_budget_bootstrap.py @@ -0,0 +1,111 @@ +"""Regression: ``pickled-schema`` CLI must honor ``budget.max_cost_usd`` cap. + +Mirror of ``pickled-iac``'s budget-bootstrap regression. See that test +module's docstring for the full bug context: the older CLI shortcut +bypassed :func:`pickled_core.llm.bootstrap.build_default_client` and +therefore never installed the user-configured :class:`BudgetGuard`. +``pickled-schema draft`` (which retries the LLM up to 3 times) would +silently ignore ``budget.max_cost_usd`` until this fix. +""" + +from __future__ import annotations + +from collections.abc import Generator +from decimal import Decimal +from pathlib import Path + +import pytest + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch: pytest.MonkeyPatch) -> None: + for key in ( + "PICKLED_SCHEMA_LLM_FACTORY", + "PICKLED_CACHE_DIR", + "PICKLED_CACHE_MODE", + "PICKLED_MAX_COST_USD", + "PICKLED_LLM_PROVIDER", + ): + monkeypatch.delenv(key, raising=False) + + +@pytest.fixture(autouse=True) +def _reset_budget_guard() -> Generator[None, None, None]: + from pickled_core.llm.budget_context import set_budget_guard + + set_budget_guard(None) + yield + set_budget_guard(None) + + +def _write_config(tmp_path: Path, cap: str) -> Path: + cfg = tmp_path / "pickled.config.yaml" + cfg.write_text( + f"""providers: + anthropic: + type: anthropic + default_model: claude-test + api_key_env: TEST_API_KEY +budget: + max_cost_usd: \"{cap}\" +cache: + mode: off +""", + encoding="utf-8", + ) + return cfg + + +def test_schema_cli_install_budget_guard_from_yaml( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.importorskip("anthropic") + _write_config(tmp_path, "0.50") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("TEST_API_KEY", "fake") + + from pickled_core.llm.budget_context import active_budget_guard + from pickled_schema.cli import _build_llm_client + + assert active_budget_guard() is None + _build_llm_client() + guard = active_budget_guard() + assert guard is not None, ( + "pickled-schema CLI bypassed the budget guard; users' max_cost_usd " + "cap would not have been enforced (see this module's docstring)." + ) + assert guard._budget.max_cost_usd == Decimal("0.50") + + +def test_schema_cli_install_budget_guard_from_env( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.importorskip("anthropic") + _write_config(tmp_path, "0.50") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("TEST_API_KEY", "fake") + monkeypatch.setenv("PICKLED_MAX_COST_USD", "2.00") + + from pickled_core.llm.budget_context import active_budget_guard + from pickled_schema.cli import _build_llm_client + + _build_llm_client() + guard = active_budget_guard() + assert guard is not None + assert guard._budget.max_cost_usd == Decimal("2.00") + + +def test_schema_cli_factory_env_shortcut_still_works( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "PICKLED_SCHEMA_LLM_FACTORY", + "pickled_bdd.testing:build_fake_llm", + ) + from pickled_bdd.testing import CannedLLMClient + from pickled_schema.cli import _build_llm_client + + client = _build_llm_client() + assert isinstance(client, CannedLLMClient)