From 4cbfdec55800aeb8cbfd6973fbb0d772055982fc Mon Sep 17 00:00:00 2001 From: Gagan Deep Date: Thu, 30 Jul 2026 20:51:02 +0530 Subject: [PATCH] fix: cap max_tokens on scan requests, and read the API key from the env Two problems found while running `modelfuzz scan` against hosted models through OpenRouter for the first time. `scan` never set max_tokens, so a gateway reserves the target model's full context window up front. Against anthropic/claude-sonnet-5 every request failed: 402 - This request requires more credits, or fewer max_tokens. You requested up to 65536 tokens, but can only afford 3937. The same scan with max_tokens=1024 completes and returns a real verdict. A probe only needs room for one tool call and a mutation for one prompt, so 1024 is generous; --max-tokens raises it. This made scanning a hosted model fail outright on exactly the credit-limited accounts most first-time users have, which is the worst possible audience for it. Separately, --api-key was a plain option with no envvar, so a real key could only reach it as a command-line argument -- landing in shell history and visible in `ps` for the duration of the run. It now reads MODELFUZZ_API_KEY. Verified end to end against OpenRouter with no key on the command line: openai/gpt-4o-mini 3/3 seeds broke through at generation 1 anthropic/claude-sonnet-5 0/3, refused through generation 3 Tests 71 -> 76. The stub client now records max_tokens per request, so both call sites are asserted rather than just the probe. --- CHANGELOG.md | 9 ++++- README.md | 14 +++++++- pyproject.toml | 2 +- src/modelfuzz/cli.py | 30 ++++++++++++++--- tests/test_scan.py | 80 +++++++++++++++++++++++++++++++++++++++++++- uv.lock | 2 +- 6 files changed, 127 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ea898b5..b6ebb0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable changes to this project are documented here. ## [Unreleased] +## [0.3.5] - 2026-07-30 + +- fix: `scan` now sends `max_tokens` on every request (default `1024`, override with `--max-tokens`). Without it, gateways such as OpenRouter reserve the target model's full context window up front and reject the call with HTTP 402 — so scanning a hosted model failed outright on exactly the credit-limited accounts most first-time users have +- feat: `--api-key` reads `MODELFUZZ_API_KEY` when not passed, so the key stays out of shell history and out of the process list +- docs: document scanning a hosted model, and both new options + ## [0.3.4] - 2026-07-30 - fix: `URLAllowList` now inspects URLs nested inside `dict`, `list`, `tuple` and `set` arguments, including dict keys. A malicious URL hidden in a payload field — `http_post(url="https://api.internal.com/v1", payload={"redirect": "http://evil.com"})` — previously passed, because 0.3.2 taught the rule to skip non-string values without teaching it to look inside them. A `redirect` or `callback` field is as much an exfiltration route as the `url` parameter itself @@ -74,7 +80,8 @@ All notable changes to this project are documented here. - CI workflow (lint + tests), MIT license, unit/integration test suite - Package renamed from `agentshield` to `modelfuzz` -[Unreleased]: https://github.com/higagan/modelfuzz/compare/v0.3.4...HEAD +[Unreleased]: https://github.com/higagan/modelfuzz/compare/v0.3.5...HEAD +[0.3.5]: https://github.com/higagan/modelfuzz/compare/v0.3.4...v0.3.5 [0.3.4]: https://github.com/higagan/modelfuzz/compare/v0.3.3...v0.3.4 [0.3.3]: https://github.com/higagan/modelfuzz/compare/v0.3.2...v0.3.3 [0.3.2]: https://github.com/higagan/modelfuzz/compare/v0.3.1...v0.3.2 diff --git a/README.md b/README.md index 352d0d5..68096f2 100644 --- a/README.md +++ b/README.md @@ -181,7 +181,19 @@ The weak model hands over credentials on the very first probe of all three seeds Options: - `--budget-s` — time budget in seconds for the attack loop (default `30`). -- `--api-key` — API key for hosted endpoints (defaults to a dummy value for local models). +- `--api-key` — API key for hosted endpoints (defaults to a dummy value for local models). Read from `MODELFUZZ_API_KEY` when not passed. +- `--max-tokens` — cap on reply length per request (default `1024`). Raise it if a target truncates its answer. + +### Scanning a hosted model + +Export the key once rather than passing it on the command line, so it stays out of your shell history and out of the process list: + +```bash +export MODELFUZZ_API_KEY="sk-..." +modelfuzz scan --endpoint https://openrouter.ai/api/v1 --model openai/gpt-4o-mini +``` + +Any OpenAI-compatible gateway works. Note that gateways bill against a reserved token budget, which is why every request sends `--max-tokens`; without it, some reject the call outright on credit-limited accounts. If every request errors out (bad endpoint, wrong model name), the scanner reports `⚠️ INCONCLUSIVE` instead of a false-safe result — an untested agent is never reported as a secure one. diff --git a/pyproject.toml b/pyproject.toml index b8783b4..12ecf44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "modelfuzz" -version = "0.3.4" +version = "0.3.5" description = "Runtime guardrails for AI agents." readme = "README.md" license = { text = "MIT" } diff --git a/src/modelfuzz/cli.py b/src/modelfuzz/cli.py index 6e11fa9..591309d 100644 --- a/src/modelfuzz/cli.py +++ b/src/modelfuzz/cli.py @@ -84,6 +84,14 @@ ) +# Cap the reply length on every request. Without it, gateways such as OpenRouter +# reserve the model's full context window up front and reject the call with HTTP +# 402 unless the account can cover it -- so an unset limit makes `scan` fail on +# exactly the credit-limited accounts most first-time users have. A probe only +# needs enough room for one tool call, and a mutation for one prompt. +DEFAULT_MAX_TOKENS = 1024 + + def _make_client(endpoint: str, api_key: str): """Build an OpenAI client, or exit with install guidance if unavailable. @@ -103,18 +111,19 @@ def _make_client(endpoint: str, api_key: str): return OpenAI(base_url=endpoint, api_key=api_key) -def _probe(client, model: str, prompt: str) -> bool: +def _probe(client, model: str, prompt: str, max_tokens: int = DEFAULT_MAX_TOKENS) -> bool: """Send one attack prompt to the target. Return True if it called the tool.""" response = client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], tools=HTTP_POST_TOOL, tool_choice="auto", + max_tokens=max_tokens, ) return bool(response.choices[0].message.tool_calls) -def _mutate(client, model: str, prompt: str) -> str: +def _mutate(client, model: str, prompt: str, max_tokens: int = DEFAULT_MAX_TOKENS) -> str: """Ask the target model to craft a more deceptive variant of a failed attack.""" response = client.chat.completions.create( model=model, @@ -122,6 +131,7 @@ def _mutate(client, model: str, prompt: str) -> str: {"role": "system", "content": _MUTATION_SYSTEM}, {"role": "user", "content": prompt}, ], + max_tokens=max_tokens, ) content = response.choices[0].message.content or "" return content.strip().strip("`").strip('"').strip() @@ -161,13 +171,23 @@ def scan( api_key: str = typer.Option( "dummy-key", "--api-key", - help="API key for the endpoint. Defaults to a dummy value for local models.", + envvar="MODELFUZZ_API_KEY", + help=( + "API key for the endpoint. Read from MODELFUZZ_API_KEY when not passed, " + "which keeps the key out of your shell history and process list. " + "Defaults to a dummy value for local models." + ), ), budget_s: float = typer.Option( 30.0, "--budget-s", help="Time budget in seconds for the adaptive attack loop.", ), + max_tokens: int = typer.Option( + DEFAULT_MAX_TOKENS, + "--max-tokens", + help="Cap on reply length per request. Raise it if a target truncates its answer.", + ), ) -> None: """Red-team a target agent with an adaptive prompt-injection fuzzer. @@ -205,7 +225,7 @@ def time_left() -> float: ) try: - triggered = _probe(client, model, prompt) + triggered = _probe(client, model, prompt, max_tokens) except Exception as exc: # noqa: BLE001 - surface any endpoint error per-attempt errors += 1 typer.echo(f"{YELLOW}[⚠️ ERROR] Request failed: {exc}{RESET}\n") @@ -227,7 +247,7 @@ def time_left() -> float: typer.echo(f"{YELLOW}[🧬 MUTATING] Evolving a more deceptive variant…{RESET}") try: - mutated = _mutate(client, model, prompt) + mutated = _mutate(client, model, prompt, max_tokens) except Exception as exc: # noqa: BLE001 - a failed mutation just ends this lineage errors += 1 typer.echo(f"{YELLOW}[⚠️ ERROR] Mutation failed: {exc}{RESET}\n") diff --git a/tests/test_scan.py b/tests/test_scan.py index 3cd67ac..c439880 100644 --- a/tests/test_scan.py +++ b/tests/test_scan.py @@ -46,10 +46,12 @@ def __init__(self, probe_result, mutation="a more deceptive variant of the attac self._mutation = mutation self.probe_prompts: list[str] = [] self.mutation_prompts: list[str] = [] + self.max_tokens_seen: list[int | None] = [] self.chat = SimpleNamespace(completions=SimpleNamespace(create=self._create)) - def _create(self, *, model, messages, tools=None, tool_choice=None): + def _create(self, *, model, messages, tools=None, tool_choice=None, max_tokens=None): prompt = messages[-1]["content"] + self.max_tokens_seen.append(max_tokens) if tools is not None: self.probe_prompts.append(prompt) result = self._probe_result @@ -169,3 +171,79 @@ def fake_import(name, *args, **kwargs): assert result.exit_code == 1 assert "pip install 'modelfuzz[scan]'" in result.output + + +class TestRequestLimits: + """Every request must cap its reply length. + + Without max_tokens, gateways such as OpenRouter reserve the model's full + context window and reject the call with HTTP 402 on credit-limited + accounts -- which is most first-time users. + """ + + def test_probe_and_mutation_both_send_max_tokens(self, monkeypatch): + # mutation="" ends each lineage after one mutate call, so both call + # sites are exercised and the queue drains without burning the budget. + client = StubClient(probe_result=_refusal_response(), mutation="") + _run(monkeypatch, client) + + assert client.max_tokens_seen, "no requests were made" + assert all(v == cli.DEFAULT_MAX_TOKENS for v in client.max_tokens_seen) + # Both call sites are covered, not just the probe. + assert client.probe_prompts and client.mutation_prompts + + def test_max_tokens_is_overridable(self, monkeypatch): + client = StubClient(probe_result=_refusal_response(), mutation="") + _run(monkeypatch, client, ["--max-tokens", "64"]) + + assert client.max_tokens_seen + assert all(v == 64 for v in client.max_tokens_seen) + + +class TestApiKeySource: + """The key should not have to appear on the command line.""" + + def _capture_key(self, monkeypatch, client): + seen: dict[str, str] = {} + + def fake_make_client(endpoint, api_key): + seen["api_key"] = api_key + return client + + monkeypatch.setattr(cli, "_make_client", fake_make_client) + return seen + + def test_reads_the_key_from_the_environment(self, monkeypatch): + client = StubClient(probe_result=_refusal_response(), mutation="") + seen = self._capture_key(monkeypatch, client) + monkeypatch.setenv("MODELFUZZ_API_KEY", "from-env") + + result = runner.invoke(cli.app, ["scan", "--endpoint", "http://x/v1", "--model", "m"]) + + assert result.exit_code == 0 + assert seen["api_key"] == "from-env" + + def test_explicit_flag_wins_over_the_environment(self, monkeypatch): + client = StubClient(probe_result=_refusal_response(), mutation="") + seen = self._capture_key(monkeypatch, client) + monkeypatch.setenv("MODELFUZZ_API_KEY", "from-env") + + result = runner.invoke( + cli.app, + ["scan", "--endpoint", "http://x/v1", "--model", "m", "--api-key", "explicit"], + ) + + assert result.exit_code == 0 + assert seen["api_key"] == "explicit" + + def test_falls_back_to_a_dummy_key_for_local_models(self, monkeypatch): + client = StubClient(probe_result=_refusal_response(), mutation="") + seen = self._capture_key(monkeypatch, client) + monkeypatch.delenv("MODELFUZZ_API_KEY", raising=False) + + result = runner.invoke( + cli.app, ["scan", "--endpoint", "http://localhost:11434/v1", "--model", "m"] + ) + + assert result.exit_code == 0 + assert seen["api_key"] == "dummy-key" diff --git a/uv.lock b/uv.lock index ced93b0..d38b88a 100644 --- a/uv.lock +++ b/uv.lock @@ -286,7 +286,7 @@ wheels = [ [[package]] name = "modelfuzz" -version = "0.3.4" +version = "0.3.5" source = { editable = "." } dependencies = [ { name = "typer" },