feat(api): pre-flight prompt-size cap with --max-input-tokens (#32)

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 <usd>` | Exit `3` if the summed agent-reported spend for this invocation exceeds the cap. Evaluated after the runs complete. |
+| `--max-input-tokens <n>` | `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"

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": [],

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 []
 

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,

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

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,

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.<id>].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,
+    )

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"))

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

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", "   "])

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"