Metadata-Version: 2.5
Name: stackone-defender
Version: 0.8.2
Summary: Indirect prompt injection defense for AI agents using tool calls
Project-URL: Homepage, https://github.com/StackOneHQ/stackone-defender
Project-URL: Repository, https://github.com/StackOneHQ/stackone-defender
Author-email: StackOne <support@stackone.com>
License-Expression: Apache-2.0
License-File: LICENSE
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Security
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.11
Provides-Extra: onnx
Requires-Dist: numpy>=1.24.0; extra == 'onnx'
Requires-Dist: onnxruntime>=1.16.0; extra == 'onnx'
Requires-Dist: tokenizers>=0.15.0; extra == 'onnx'
Provides-Extra: sfe
Requires-Dist: fasttext-ng>=0.9.3; extra == 'sfe'
Description-Content-Type: text/markdown

<div align="center">

  <picture>
    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/banner-dark.svg" />
    <img src="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/banner-light.svg" alt="Defender by StackOne — Indirect prompt injection protection for MCP tool calls" width="800" />
  </picture>

  <p>
    <a href="https://pypi.org/project/stackone-defender/"><img src="https://img.shields.io/pypi/v/stackone-defender?style=flat-square&color=047B43&label=pypi" alt="PyPI version" /></a>
    <a href="https://github.com/StackOneHQ/stackone-defender/releases"><img src="https://img.shields.io/github/v/release/StackOneHQ/stackone-defender?style=flat-square&color=047B43&label=release" alt="latest GitHub release" /></a>
    <a href="https://github.com/StackOneHQ/stackone-defender/stargazers"><img src="https://img.shields.io/github/stars/StackOneHQ/stackone-defender?style=flat-square&color=047B43" alt="GitHub stars" /></a>
    <a href="./LICENSE"><img src="https://img.shields.io/pypi/l/stackone-defender?style=flat-square&color=047B43" alt="License" /></a>
    <img src="https://img.shields.io/badge/Python-3.11+-047B43?style=flat-square" alt="Python 3.11+" />
  </p>
  <p>
    <img src="https://img.shields.io/badge/model-22MB-047B43?style=flat-square" alt="Model size: 22MB" />
    <img src="https://img.shields.io/badge/latency-~10ms-047B43?style=flat-square" alt="Latency: ~10ms" />
    <img src="https://img.shields.io/badge/CPU--only-no%20GPU%20needed-047B43?style=flat-square" alt="CPU only" />
    <img src="https://img.shields.io/badge/F1%20Score-90.8%25-047B43?style=flat-square" alt="F1 Score: 90.8%" />
  </p>

</div>

---

Indirect prompt injection defense for AI agents using tool calls (MCP, CLI, or direct APIs). Detects and neutralizes attacks hidden in tool results (emails, documents, PRs, etc.) before they reach your LLM.

**Python package:** [`stackone-defender`](https://pypi.org/project/stackone-defender/) — aligned with [`@stackone/defender`](https://www.npmjs.com/package/@stackone/defender) on npm.

## Installation

**pip**

```bash
pip install stackone-defender
```

**uv**

```bash
uv add stackone-defender
```

**Tier 2 (ONNX)** — add extras:

```bash
pip install stackone-defender[onnx]
# or: uv add "stackone-defender[onnx]"
```

The ONNX model (~22MB) is bundled in the wheel — no extra downloads at runtime.

**SFE preprocessor (optional)** — add extras:

```bash
pip install stackone-defender[sfe]
# or: uv add "stackone-defender[sfe]"
```

The `[sfe]` extra installs [`fasttext-ng`](https://pypi.org/project/fasttext-ng/) (provides the `fasttext` module). It requires **NumPy 2.3+**. PyPI may ship a wheel only for some platforms; otherwise pip/uv builds from source (needs a C++ toolchain).

## Quick start

```python
from stackone_defender import create_prompt_defense

# Tier 1 + Tier 2 are on by default. block_high_risk=True enables allow/block.
defense = create_prompt_defense(block_high_risk=True)

# Optional: preload ONNX to avoid first-call latency (requires [onnx] extra)
defense.warmup_tier2()

result = defense.defend_tool_result(tool_output, "gmail_get_message")

if not result.allowed:
    print(f"Blocked: risk={result.risk_level}, score={result.tier2_score}")
    print(f"Detections: {', '.join(result.detections)}")
else:
    send_to_llm(result.sanitized)
```

## How it works

<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/demo-dark.svg" />
  <img src="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/demo-light.svg" alt="Defender flow: poisoned tool output is sanitized and evaluated; high-risk content can be blocked before the LLM" width="900" />
</picture>

`defend_tool_result()` runs two tiers:

### Tier 1 — Pattern detection (sync, ~1 ms)

Detects (never rewrites) injection signals and records them as evidence:

- **Unicode normalization** (for matching only) — homoglyph resistance (e.g. Cyrillic `а` → ASCII `a`)
- **Role markers** — `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]`, etc.
- **Instruction-override patterns** — phrases like “ignore previous instructions”
- **Encoding detection** — decode-then-detect: a decoded payload is escalated only if it trips a real attack pattern
- **Boundary annotation (opt-in)** — `[UD-{id}]…[/UD-{id}]` wrappers when `annotate_boundary=True` (npm: `annotateBoundary`). Use `generate_boundary_instructions` from the package root in prompts when you enable wrapping.

### Tier 2 — ML classification (ONNX)

Packed-chunk MiniLM classifier (int8 ONNX ~22 MB, bundled):

- Split text into sentences, pack to model-sized chunks, score chunks in batched ONNX calls
- Catches paraphrased or novel injections missed by regex
- Uses chunked batch inference to bound memory on large payloads

### Optional SFE preprocessor

- `use_sfe=True` runs a field-level FastText pass to build a **classifier-only** view of the payload
- **Tier 1** detects on the raw tool value; SFE drops are classifier-only and never remove fields from the returned `sanitized` payload
- **Tier 2** extracts strings from the SFE-filtered tree; `fields_dropped` lists paths omitted from that extraction (not removed from `sanitized`)
- Fails open if the runtime/model is unavailable: payload continues unfiltered

**Benchmarks** (F1 @ threshold 0.5):

| Benchmark | F1 | Samples |
|-----------|-----|--------|
| Qualifire (in-distribution) | 0.8686 | ~1.5k |
| xxz224 (out-of-distribution) | 0.8834 | ~22.5k |
| jayavibhav (adversarial) | 0.9717 | ~1k |
| **Average** | **0.9079** | ~25k |

### Optional Tier 3 — LLM adjudication (consumer-supplied)

Authoritative LLM-based classification for the cases Tier 2 finds ambiguous. The package ships **only** the orchestration and the `Tier3Provider` interface — the actual model endpoint (a hosted LLM, OpenAI, an internal inference service) lives in your code, keeping proprietary models and credentials out of the package.

Two modes, selected via `defender_mode`:
- **`"cascade"`** (default): Tier 1 → Tier 2 → Tier 3, with Tier 3 invoked only when the Tier 2 effective score falls in the gray band (default `[0.3, 0.85)`). The Tier 3 verdict overrides Tier 2 on the escalated chunk — a `block` forces a block, an `allow` rescues it. Outside the band defender skips the round trip.
- **`"tier3_only"`**: skip Tier 2; the block/allow decision is the Tier 3 verdict alone. Tier 1 still runs detection; the returned `sanitized` payload is the original content.

Register a provider once at startup, then opt in per instance:

```python
from typing import Any

from stackone_defender import create_prompt_defense, set_default_tier3_provider
from stackone_defender.types import Tier3Verdict

class MyProvider:
    def classify(self, text: str, *, ctx: dict[str, Any] | None = None) -> Tier3Verdict:
        # Call your LLM endpoint here (sync or awaitable). ctx["toolName"] is available.
        verdict = call_my_llm(text, tool_name=(ctx or {}).get("toolName"))
        return Tier3Verdict(decision="block" if verdict.block else "allow", score=verdict.p_block)

set_default_tier3_provider(MyProvider())

defense = create_prompt_defense(
    block_high_risk=True,
    enable_tier3=True,
    defender_mode="cascade",  # or "tier3_only"
    tier3={
        "provider": MyProvider(),                          # optional: overrides the registry for this instance
        "escalation_band": {"lower": 0.3, "upper": 0.85},  # cascade gray band; [lower, upper), defaults shown
        "max_text_length": 10000,                          # caps text passed to the provider
        "block_threshold": 0.622,                          # optional; decide on score, not the model's word
    },
)
```

**Choosing the operating point (`block_threshold`).** By default the model's `decision` word is authoritative — but that word is its argmax, an implicit 0.5 cut that moves on its own whenever the model is retrained. Set `tier3.block_threshold` to decide on `verdict.score` (P(block)) instead: the cut becomes an explicit config value — raise it to trade recall for fewer false positives, lower it for the reverse; `0.5` reproduces argmax. It requires a provider that reports `score` as P(block) (not "confidence in whichever decision I made" — those invert on allows). If `score` is missing or outside `[0, 1]` the verdict's `decision` is used instead and defender warns once, so a provider that cannot report a score degrades to the default rather than failing.

**Fail-open semantics.** A provider error or timeout records a `skip_reason` on `result.tier3`; in cascade defender falls back to the Tier 2 decision, in `tier3_only` it allows the request. `enable_tier3=True` with no registered provider falls back to the standard Tier 1 + Tier 2 cascade and logs one warning per instance — Tier 3 misconfiguration never silently disables defense. `result.tier3` carries the verdict (a `Tier3Verdict`, or a `Tier3Skip` when the provider ran but returned nothing usable) when Tier 3 runs, and is `None` when it doesn't.

### `allowed` vs `risk_level`

- **`DefenseResult.sanitized`** is a **sentence-level cleaned** copy of the tool result (high-scoring sentences dropped within high-risk fields, optionally `[UD-…]` boundary-wrapped). Cleaning is best-effort (capped by detection) — still gate on `allowed`. Set `sanitize_content=False` for pure detect-and-gate: `sanitized` is then the input verbatim.
- Use **`allowed`** for gating when `block_high_risk=True`: `False` means do not pass `sanitized` to the model as-is.
- **`risk_level`** is diagnostic: it starts at `default_risk_level` (default `"low"`) and is **escalated** by Tier 1 / Tier 2 signals — not reduced. Use it for logging, not as the sole block signal unless you implement your own policy.

| Level | Typical trigger |
|-------|------------------|
| `low` | No strong signals |
| `medium` | Lighter pattern / sanitization signals |
| `high` / `critical` | Strong injection patterns, encoding signals, or high Tier 2 score |

## API

### `create_prompt_defense(**kwargs)`

```python
defense = create_prompt_defense(
    enable_tier1=True,
    enable_tier2=True,
    require_tier2=False,  # True: raise if Tier 2 can't load (fail closed) instead of degrading to Tier 1
    block_high_risk=False,
    default_risk_level="low",
    annotate_boundary=False,  # True: wrap risky strings with [UD-…] tags (npm: annotateBoundary)
    tier2_fields=["subject", "body", "snippet"],  # optional: scope Tier 2 to these JSON keys (default: all strings)
    use_sfe=True,  # optional: enable semantic field extractor preprocessing
    config={
        "tier2": {
            "high_risk_threshold": 0.8,
            "tier2_fields": None,  # or list[str]; constructor tier2_fields wins if set
        },
    },
)
```

### `defense.defend_tool_result(value, tool_name)`

Runs Tier 1 sanitization on risky fields of the **original** payload, then Tier 2 on strings from the SFE-filtered view when SFE is on (otherwise the full value). Optional `tier2_fields` restricts Tier 2 extraction to specific keys; omit it to classify **all** strings (matches `@stackone/defender` 0.6.3). **Synchronous** — no `await`.

```python
from dataclasses import dataclass, field

@dataclass
class DefenseResult:
    allowed: bool                             # gating decision (respects block_high_risk)
    risk_level: RiskLevel                     # diagnostic; max of Tier 1 / Tier 2
    sanitized: Any                            # sentence-cleaned copy (input verbatim when sanitize_content=False); dropped runs leave a [CONTENT SANITISED] marker; best-effort, still gate on allowed
    detections: list[str]                     # Tier 1 pattern names detected
    fields_sanitized: list[str]               # fields whose content the cleaner changed in sanitized (empty when sanitize_content=False or no Tier 2); for detections read detections/patterns_by_field
    patterns_by_field: dict[str, list[str]]   # patterns detected per field
    detected_field_count: int                 # count of fields with a Tier-1 detection (keys of patterns_by_field); threat-count signal (fields_sanitized len no longer tracks this)
    tier2_score: float | None = None
    tier2_raw_score: float | None = None
    tier2_aux_score: float | None = None      # multi-head models only
    tier2_multihead_blocked: bool | None = None
    tier2_skip_reason: str | None = None
    max_sentence: str | None = None
    tier3: Tier3Result | None = None          # present when Tier 3 ran
    fields_dropped: list[str] = field(default_factory=list)
    truncated_at_depth: bool | None = None
    latency_ms: float = 0.0
    # Cost telemetry — present only when the batched Tier 2 classifier ran
    phase_timings: PhaseTimings | None = None       # prepare / infer / aggregate ms
    tier2_stats: Tier2Stats | None = None           # string/chunk/unique counts, real/padded tokens
    tier1_ms: float | None = None
    cold_load: bool | None = None
    # Operational signals
    tier2_available: bool | None = None       # False when Tier 2 enabled but failed to load
    coverage_degraded: bool | None = None      # True when Tier 1 detection coverage was capped
```

### `defense.defend_tool_results(items)`

Sync batch API. When `enable_tier3=True`, uses one `asyncio.run()` and defends items **concurrently** via `asyncio.gather` (same scheduling model as npm `defendToolResults`; blocking sync providers still run one at a time on the event-loop thread). From async code, prefer `defend_tool_results_async`.

```python
results = defense.defend_tool_results([
    {"value": email_data, "tool_name": "gmail_get_message"},
    {"value": doc_data, "tool_name": "documents_get"},
    {"value": pr_data, "tool_name": "github_get_pull_request"},
])
for r in results:
    if not r.allowed:
        print("Blocked:", ", ".join(r.detections))
```

### `await defense.defend_tool_results_async(items)`

Async batch API — runs `defend_tool_result_async` per item concurrently via `asyncio.gather`. Required when Tier 3 is enabled inside a running event loop (e.g. FastAPI).

```python
results = await defense.defend_tool_results_async([
    {"value": email_data, "tool_name": "gmail_get_message"},
    {"value": doc_data, "tool_name": "documents_get"},
])
```

### `defense.analyze(text)`

Tier 1 only — useful for debugging pattern hits without full tool-result traversal.

### Tier 2 warmup

```python
defense = create_prompt_defense()
defense.warmup_tier2()  # no-op if enable_tier2=False or ONNX extra missing
```

## Integration example

```python
from stackone_defender import create_prompt_defense

defense = create_prompt_defense(block_high_risk=True)
defense.warmup_tier2()

def run_tool_and_defend(raw_result: dict, tool_name: str):
    outcome = defense.defend_tool_result(raw_result, tool_name)
    if not outcome.allowed:
        return {"error": "Content blocked by safety filter", "risk_level": outcome.risk_level}
    return outcome.sanitized

# Example agent loop
sanitized = run_tool_and_defend(gmail_api.get_message(msg_id), "gmail_get_message")
```

## Risky field detection

Only **string** values under configured “risky” keys are Tier-1-scanned — including strings nested inside arrays/objects under those keys (e.g. `{"name": ["…"]}`). [`RiskyFieldConfig`](https://github.com/StackOneHQ/stackone-defender/blob/main/src/stackone_defender/types.py) provides global names/patterns plus **`tool_overrides`** (wildcard tool names → field list), same idea as the npm package. (Tier 2 scans all extracted strings regardless.)

| Tool pattern | Scanned fields |
|--------------|----------------|
| `gmail_*`, `email_*` | subject, body, snippet, content |
| `documents_*` | name, description, content, title |
| `github_*` | name, title, body, description, message |
| `hris_*` | name, notes, bio, description |
| `ats_*` | name, notes, description, summary |
| `crm_*` | name, description, notes, content |

Otherwise the default list applies: `name`, `description`, `content`, `title`, `notes`, `summary`, `bio`, `body`, `text`, `message`, `comment`, `subject`, plus suffix patterns like `*_body`, `*_description`, etc. Structural keys such as `id`, `url`, `created_at` are not treated as risky by default.

## Development

```bash
uv sync --group dev
uv run pytest
```

## License

Apache-2.0 — see [LICENSE](./LICENSE).
