From 3b712fe281e0bd3e5f09e7fe88df1fbf4c236918 Mon Sep 17 00:00:00 2001 From: karthickmk2001 Date: Wed, 11 Nov 2026 20:12:35 +0000 Subject: [PATCH] feat(api): pre-flight prompt-size cap with --max-input-tokens (#32) The api agent pastes the brief plus target source into one prompt, so a large target produced a very large prompt whose cost was only known after the call. --max-cost is checked after the runs finish, so it can't stop that. Add --max-input-tokens on `secfoo run` (falls back to [agents.api].max_input_tokens). After the api agent assembles the prompt, it estimates its size with LiteLLM's tokenizer for the model (or ~4 chars per token when there is none). If the prompt is over the limit, the run fails before litellm.completion() is called, and the error says how large the prompt was and which limit applied. CLI agents read the target themselves, so they warn and ignore the setting instead of pretending to cap it. The other half of #32, a distinct exit code when --max-cost is exceeded, already shipped in #64 (exit 3). This also fixes `secfoo run --help`, which still said --max-cost exits 2. Closes #32 Co-Authored-By: Claude Opus 5.5 --- docs/REFERENCE.md | 5 +++ src/secfoo/agents/api.py | 34 ++++++++++++++- src/secfoo/agents/base.py | 18 ++++++++ src/secfoo/cli.py | 14 +++++- src/secfoo/config.example.toml | 2 + src/secfoo/runner.py | 5 +++ src/secfoo/settings.py | 16 ++++++- tests/test_agents/test_api.py | 78 ++++++++++++++++++++++++++++++++++ tests/test_agents/test_base.py | 11 +++++ tests/test_cli.py | 27 ++++++++++++ tests/test_settings.py | 14 ++++++ 11 files changed, 221 insertions(+), 3 deletions(-) diff --git a/docs/REFERENCE.md b/docs/REFERENCE.md index 387ae8e..1956fee 100644 --- a/docs/REFERENCE.md +++ b/docs/REFERENCE.md @@ -417,6 +417,7 @@ secfoo run --skill sast --skill secret-scanning --agent api --target . \ |---|---| | `--fail-on critical\|high\|medium\|low` | Exit `2` if any run reports a finding at that severity or worse. Counts come from the report contract's `### [SEVERITY] Fn:` headings, not from prose. | | `--max-cost ` | Exit `3` if the summed agent-reported spend for this invocation exceeds the cap. Evaluated after the runs complete. | +| `--max-input-tokens ` | `api` agent only. Checked **before** the model is called: if the assembled prompt (brief + target source) is larger than `n` input tokens, the run fails without sending anything, and the error says how large the prompt was and which limit applied. Token counts come from LiteLLM's tokenizer for the model, or ~4 characters per token when it has none. Falls back to `[agents.api].max_input_tokens`. Other agents read the target themselves, so they warn and ignore it. | | `--json` | Write one JSON document to stdout (per-skill status, severity counts, tokens, USD, gate verdicts, resolved exit code); progress and notes go to stderr. Also suppresses the interactive project-name prompt. | Exit codes: `0` clean, `1` a run failed or timed out (gate verdicts are @@ -489,6 +490,10 @@ max_cost_usd = 2.00 [agents.claude] model = "claude-sonnet-4-6" +[agents.api] +# Refuse to send a prompt larger than this (see --max-input-tokens). +max_input_tokens = 150000 + [[mcp_servers]] name = "Atlassian-Rovo-MCP" command = "npx" diff --git a/src/secfoo/agents/api.py b/src/secfoo/agents/api.py index aee86da..a6a9384 100644 --- a/src/secfoo/agents/api.py +++ b/src/secfoo/agents/api.py @@ -46,11 +46,30 @@ ) +class PromptTooLargeError(RuntimeError): + """The assembled prompt is over the configured max_input_tokens.""" + + +def estimate_input_tokens(litellm, model: str, messages: list[dict]) -> int: + """Input tokens for `messages`, from LiteLLM's tokenizer for `model`. + + Falls back to ~4 characters per token (a common rule of thumb for English + text and code) when LiteLLM can't count for this model, so the cap still + applies rather than silently switching off. + """ + try: + return int(litellm.token_counter(model=model, messages=messages)) + except Exception: # noqa: BLE001 -- unknown model / missing tokenizer -> estimate + chars = sum(len(m.get("content") or "") for m in messages) + return -(-chars // 4) # ceil + + class _State(TypedDict): prompt: str workdir: str model: str timeout: int + max_input_tokens: int | None # #32: refuse before sending if the prompt is larger prev_commit: str | None # Memory Bank: SHA of the last scanned commit; None = full scan. # Memory Bank: paths an incremental scan covered (files sent + files # deleted); None = full scan. The runner only closes findings in these. @@ -198,7 +217,18 @@ def collect(state: _State) -> dict: f"## Source files\nYou cannot open files. " f"The complete target source is included below.\n{code}" ) - return {"messages": [{"role": "user", "content": content}], "scan_scope": scan_scope} + messages = [{"role": "user", "content": content}] + limit = state.get("max_input_tokens") + if limit: + size = estimate_input_tokens(litellm, state["model"], messages) + if size > limit: + # Raised before litellm.completion(), so nothing is sent or billed. + raise PromptTooLargeError( + f"Prompt is about {size:,} input tokens, over the max_input_tokens limit of {limit:,} " + f"(--max-input-tokens or [agents.api].max_input_tokens). Nothing was sent to the model. " + "Narrow the target with --exclude, or raise the limit." + ) + return {"messages": messages, "scan_scope": scan_scope} def assess(state: _State) -> dict: response = litellm.completion( @@ -250,6 +280,7 @@ def build_command(self, prompt: str, *, workdir: Path) -> list[str]: return [] # required by the base class; there is no program to launch supports_model = True + supports_max_input_tokens = True def run(self, prompt: str, *, workdir: Path, timeout: int | None = None, prev_commit: str | None = None) -> AgentResult: # Memory Bank: prev_commit triggers incremental diff scan. started = time.monotonic() @@ -272,6 +303,7 @@ def run(self, prompt: str, *, workdir: Path, timeout: int | None = None, prev_co "workdir": str(workdir), "model": self.model or os.environ.get("SECFOO_API_MODEL", DEFAULT_MODEL), "timeout": timeout or self.default_timeout_seconds, + "max_input_tokens": self.max_input_tokens, "prev_commit": prev_commit, # Memory Bank: None = full scan, SHA = incremental diff. "scan_scope": None, "messages": [], diff --git a/src/secfoo/agents/base.py b/src/secfoo/agents/base.py index f85499a..6b29fbf 100644 --- a/src/secfoo/agents/base.py +++ b/src/secfoo/agents/base.py @@ -109,6 +109,24 @@ def configure_model(self, model: str | None) -> None: model = None self.model = model + # Pre-flight cap on prompt size (#32). Only an adapter that builds the + # full prompt itself (the api agent) can measure it before sending; CLI + # agents read the target's files on their own, out of secfoo's sight. + max_input_tokens: int | None = None + supports_max_input_tokens: ClassVar[bool] = False + + def configure_max_input_tokens(self, limit: int | None) -> None: + if limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit <= 0): + raise ValueError("max_input_tokens must be a positive integer") + if limit is not None and not self.supports_max_input_tokens: + warnings.warn( + f"Agent {self.name!r} reads the target itself, so its prompt size can't be checked " + "before it runs; ignoring max_input_tokens", + stacklevel=2, + ) + limit = None + self.max_input_tokens = limit + def model_arguments(self) -> list[str]: return [self.model_flag, self.model] if self.model_flag and self.model is not None else [] diff --git a/src/secfoo/cli.py b/src/secfoo/cli.py index 7262435..cb65c32 100644 --- a/src/secfoo/cli.py +++ b/src/secfoo/cli.py @@ -259,10 +259,19 @@ def run( None, "--max-cost", min=0.0, - help="Exit 2 if the summed reported AI spend (USD) for this invocation " + help="Exit 3 if the summed reported AI spend (USD) for this invocation " "exceeds this amount. Checked after the runs finish (an agent CLI " "only reports cost on completion). Falls back to [defaults].max_cost_usd.", ), + max_input_tokens: int | None = typer.Option( + None, + "--max-input-tokens", + min=1, + help="api agent only: refuse to send a prompt larger than this many input " + "tokens, checked BEFORE the model is called (unlike --max-cost). The run " + "fails and says how large the prompt was. Falls back to " + "[agents.api].max_input_tokens.", + ), json_output: bool = typer.Option( False, "--json", @@ -327,6 +336,8 @@ def run( raise typer.Exit(code=1) if model is None: model = file_config.agent_models.get(agent.value) + if max_input_tokens is None: + max_input_tokens = file_config.agent_max_input_tokens.get(agent.value) if depth is None: depth = DepthId(file_defaults.depth) if file_defaults.depth else DepthId.QUICK if timeout is None: @@ -390,6 +401,7 @@ def handle_complete(skill_name: str, result_status: str) -> None: confluence_urls=confluence, agent_id=agent.value, model=model, + max_input_tokens=max_input_tokens, depth=depth.value, timeout=timeout, on_skill_start=handle_start, diff --git a/src/secfoo/config.example.toml b/src/secfoo/config.example.toml index fbfa3ef..e3a1e1d 100644 --- a/src/secfoo/config.example.toml +++ b/src/secfoo/config.example.toml @@ -31,6 +31,8 @@ depth = "quick" # quick | standard # model = "sonnet" # [agents.agent] # model = "your-cursor-model-id" +# [agents.api] +# max_input_tokens = 150000 # refuse a larger prompt before it is sent (--max-input-tokens) [memory] # secfoo-memory: a hosted service of curated vulnerability patterns the diff --git a/src/secfoo/runner.py b/src/secfoo/runner.py index d6628c8..ba76a0b 100644 --- a/src/secfoo/runner.py +++ b/src/secfoo/runner.py @@ -418,6 +418,7 @@ def _run_single_skill( exclude_paths: list[str] | None = None, stack: Stack | None = None, model: str | None = None, + max_input_tokens: int | None = None, ) -> RunOutcome: skill = load_skill(skill_id) memory_enabled = memory.is_active_for_agent(agent_id) @@ -456,6 +457,8 @@ def _run_single_skill( adapter = get_adapter(agent_id) if model is not None: adapter.configure_model(model) + if max_input_tokens is not None: + adapter.configure_max_input_tokens(max_input_tokens) # Memory Bank: look up the last successful commit so the adapter can diff and send only changed files. prev_commit = repo.get_last_commit(project_id=project_id, skill_id=skill.id) result = adapter.run(prompt, workdir=target_ctx.local_path, timeout=timeout, prev_commit=prev_commit) @@ -547,6 +550,7 @@ def execute_runs( confluence_urls: list[str], agent_id: str, model: str | None = None, + max_input_tokens: int | None = None, depth: Literal["quick", "standard"] = "quick", timeout: int | None = None, repo: RunRepository | None = None, @@ -647,6 +651,7 @@ def execute_runs( project_id=project_id, agent_id=agent_id, model=model, + max_input_tokens=max_input_tokens, target_ctx=target_ctx, confluence_urls=confluence_urls, depth=depth, diff --git a/src/secfoo/settings.py b/src/secfoo/settings.py index e5bca37..095d4c7 100644 --- a/src/secfoo/settings.py +++ b/src/secfoo/settings.py @@ -70,6 +70,8 @@ class SecfooConfig: # without a key there is nothing to wire in either way. memory_enabled: bool = True agent_models: dict[str, str] = field(default_factory=dict) + # [agents.].max_input_tokens -- pre-flight prompt-size cap (#32). + agent_max_input_tokens: dict[str, int] = field(default_factory=dict) def _parse_mcp_server(name: str, raw: dict) -> MCPServerConfig: @@ -124,6 +126,7 @@ def load_config(path: Path | None = None) -> SecfooConfig: if not isinstance(agents_raw, dict): raise ConfigError("[agents] must be a table") agent_models = {} + agent_max_input_tokens = {} for agent_id, options in agents_raw.items(): if not isinstance(options, dict): raise ConfigError(f"agents.{agent_id} must be a table") @@ -132,6 +135,11 @@ def load_config(path: Path | None = None) -> SecfooConfig: if not isinstance(model, str) or not model.strip(): raise ConfigError(f"agents.{agent_id}.model must be a non-empty string") agent_models[agent_id] = model + max_input_tokens = options.get("max_input_tokens") + if max_input_tokens is not None: + if isinstance(max_input_tokens, bool) or not isinstance(max_input_tokens, int) or max_input_tokens <= 0: + raise ConfigError(f"agents.{agent_id}.max_input_tokens must be a positive integer") + agent_max_input_tokens[agent_id] = max_input_tokens servers = [] seen_names: set[str] = set() @@ -151,4 +159,10 @@ def load_config(path: Path | None = None) -> SecfooConfig: if not isinstance(memory_enabled, bool): raise ConfigError("memory.enabled must be true or false") - return SecfooConfig(defaults=defaults, mcp_servers=servers, memory_enabled=memory_enabled, agent_models=agent_models) + return SecfooConfig( + defaults=defaults, + mcp_servers=servers, + memory_enabled=memory_enabled, + agent_models=agent_models, + agent_max_input_tokens=agent_max_input_tokens, + ) diff --git a/tests/test_agents/test_api.py b/tests/test_agents/test_api.py index 1efe508..3a1e06b 100644 --- a/tests/test_agents/test_api.py +++ b/tests/test_agents/test_api.py @@ -105,6 +105,84 @@ def test_run_success_returns_report_and_usage(tmp_path, monkeypatch): assert "===== FILE: app.py =====" in content +# ---------- #32: pre-flight prompt-size cap ---------- + + +def _api_target(tmp_path, monkeypatch, *, source_chars=40_000): + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + (tmp_path / "big.py").write_text("x" * source_chars, encoding="utf-8") + + +def test_prompt_over_max_input_tokens_fails_before_anything_is_sent(tmp_path, monkeypatch): + _api_target(tmp_path, monkeypatch) # ~40k chars -> ~10k tokens by the char estimate + fake = _fake_litellm(monkeypatch) + adapter = ApiAdapter() + adapter.configure_max_input_tokens(1_000) + + result = adapter.run("Review this.", workdir=tmp_path) + + assert result.status == "failed" + assert fake.calls == [] # nothing sent, nothing billed + assert "over the max_input_tokens limit of 1,000" in result.stderr + assert "Prompt is about 10," in result.stderr # says how large it was + assert "--max-input-tokens" in result.stderr # and which knob applies + assert result.cost_usd is None and result.input_tokens is None + + +def test_prompt_under_max_input_tokens_runs_normally(tmp_path, monkeypatch): + _api_target(tmp_path, monkeypatch, source_chars=400) + fake = _fake_litellm(monkeypatch) + adapter = ApiAdapter() + adapter.configure_max_input_tokens(100_000) + + result = adapter.run("Review this.", workdir=tmp_path) + + assert result.status == "success" + assert len(fake.calls) == 1 + + +def test_no_max_input_tokens_means_no_check(tmp_path, monkeypatch): + _api_target(tmp_path, monkeypatch) + fake = _fake_litellm(monkeypatch) + fake.token_counter = lambda **kw: pytest.fail("must not count without a limit") + + assert ApiAdapter().run("Review this.", workdir=tmp_path).status == "success" + + +def test_max_input_tokens_uses_litellms_tokenizer_for_the_model(tmp_path, monkeypatch): + _api_target(tmp_path, monkeypatch, source_chars=400) + monkeypatch.setenv("SECFOO_API_MODEL", "openai/test-model") + fake = _fake_litellm(monkeypatch) + seen = {} + + def token_counter(model, messages): + seen["model"] = model + return 5_000 # tokenizer says far more than the char estimate would + + fake.token_counter = token_counter + adapter = ApiAdapter() + adapter.configure_max_input_tokens(4_000) + + result = adapter.run("Review this.", workdir=tmp_path) + + assert result.status == "failed" + assert seen["model"] == "openai/test-model" + assert "about 5,000 input tokens" in result.stderr + assert fake.calls == [] + + +def test_estimate_falls_back_to_four_chars_per_token_when_the_tokenizer_fails(): + broken = types.SimpleNamespace(token_counter=lambda **kw: (_ for _ in ()).throw(ValueError("unknown model"))) + messages = [{"role": "user", "content": "a" * 401}] + assert api.estimate_input_tokens(broken, "x/unknown", messages) == 101 # ceil(401 / 4) + + +@pytest.mark.parametrize("bad", [0, -5, True, 1.5, "10"]) +def test_configure_max_input_tokens_rejects_non_positive_integers(bad): + with pytest.raises(ValueError, match="positive integer"): + ApiAdapter().configure_max_input_tokens(bad) + + def test_run_success_with_unpriced_model_keeps_tokens_and_leaves_cost_unknown(tmp_path, monkeypatch): monkeypatch.setenv("OPENAI_API_KEY", "sk-test") _fake_litellm(monkeypatch, cost=RuntimeError("no price for model")) diff --git a/tests/test_agents/test_base.py b/tests/test_agents/test_base.py index 979c5af..2e5bb3c 100644 --- a/tests/test_agents/test_base.py +++ b/tests/test_agents/test_base.py @@ -178,3 +178,14 @@ def _must_not_be_called(self, stdout, stderr): monkeypatch.setattr(CodexAdapter, "describe_failure", _must_not_be_called) fake_popen(returncode=0, stdout="report", stderr="") assert CodexAdapter().run("hi", workdir=tmp_path).status == "success" + + +def test_cli_agents_warn_and_ignore_max_input_tokens(): + # #32: a CLI agent reads the target itself, so secfoo can't measure its + # prompt up front -- it must say so rather than silently pretend to cap it. + import pytest + + adapter = ClaudeAdapter() + with pytest.warns(UserWarning, match="can't be checked"): + adapter.configure_max_input_tokens(1000) + assert adapter.max_input_tokens is None diff --git a/tests/test_cli.py b/tests/test_cli.py index 5b03167..790823a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -609,6 +609,33 @@ def execute(**kwargs): assert config.agent_models["codex"] == "config-model" +@pytest.mark.parametrize("override,expected", [(None, 150000), ("2000", 2000)]) +def test_run_max_input_tokens_flag_overrides_config(monkeypatch, override, expected): + captured = {} + config = SecfooConfig( + defaults=Defaults(agent="api"), mcp_servers=[], agent_max_input_tokens={"api": 150000, "claude": 9} + ) + monkeypatch.setattr("secfoo.cli.load_config", lambda: config) + monkeypatch.setattr("secfoo.cli.execute_runs", lambda **kw: captured.update(kw) or [_outcome("success")]) + args = ["run", "--skill", "sast"] + if override is not None: + args += ["--max-input-tokens", override] + result = runner.invoke(app, args) + assert result.exit_code == 0, result.output + assert captured["max_input_tokens"] == expected + + +def test_run_max_input_tokens_must_be_positive(monkeypatch): + monkeypatch.setattr("secfoo.cli.execute_runs", lambda **kw: pytest.fail("must not run")) + result = runner.invoke(app, ["run", "--skill", "sast", "--max-input-tokens", "0"]) + assert result.exit_code == 2 # typer usage error + + +def test_run_help_documents_the_cost_exit_code_as_3(): + result = runner.invoke(app, ["run", "--help"], env={"COLUMNS": "200"}) + assert "Exit 3 if the summed reported AI spend" in result.output + + def test_run_empty_model_rejected(monkeypatch): monkeypatch.setattr("secfoo.cli.execute_runs", lambda **kw: pytest.fail("must not run")) result = runner.invoke(app, ["run", "--skill", "sast", "--model", " "]) diff --git a/tests/test_settings.py b/tests/test_settings.py index cc81b8a..7c3b56d 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -133,6 +133,20 @@ def test_per_agent_models(tmp_path): assert load_config(path).agent_models == {"claude": "claude-model", "agent": "cursor-model"} +def test_per_agent_max_input_tokens(tmp_path): + path = tmp_path / "config.toml" + path.write_text("[agents.api]\nmax_input_tokens = 150000\n") + assert load_config(path).agent_max_input_tokens == {"api": 150000} + + +@pytest.mark.parametrize("value", ["0", "-1", "1.5", '"150000"', "true"]) +def test_invalid_max_input_tokens_is_rejected(tmp_path, value): + path = tmp_path / "config.toml" + path.write_text(f"[agents.api]\nmax_input_tokens = {value}\n") + with pytest.raises(ConfigError, match="positive integer"): + load_config(path) + + @pytest.mark.parametrize("value", ['""', '" "', '42', 'false', '[]']) def test_invalid_agent_model_is_rejected(tmp_path, value): path = tmp_path / "config.toml"