Skip to content

Repository files navigation

Defender by StackOne — Indirect prompt injection protection for MCP tool calls

PyPI versionlatest GitHub releaseGitHub starsLicensePython 3.11+

Model size: 22MBLatency: ~10msCPU onlyF1 Score: 90.8%


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 — aligned with @stackone/defender on npm.

Installation

pip

pip install stackone-defender

uv

uv add stackone-defender

Tier 2 (ONNX) — add extras:

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:

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

The [sfe] extra installs 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

fromstackone_defenderimportcreate_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")
ifnotresult.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

Defender flow: poisoned tool output is sanitized and evaluated; high-risk content can be blocked before the LLM

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 markersSYSTEM:, 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):

BenchmarkF1Samples
Qualifire (in-distribution)0.8686~1.5k
xxz224 (out-of-distribution)0.8834~22.5k
jayavibhav (adversarial)0.9717~1k
Average0.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:

fromtypingimportAnyfromstackone_defenderimportcreate_prompt_defense, set_default_tier3_providerfromstackone_defender.typesimportTier3VerdictclassMyProvider:
defclassify(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=(ctxor {}).get("toolName"))
returnTier3Verdict(decision="block"ifverdict.blockelse"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.
LevelTypical trigger
lowNo strong signals
mediumLighter pattern / sanitization signals
high / criticalStrong injection patterns, encoding signals, or high Tier 2 score

API

create_prompt_defense(**kwargs)

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 1block_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 preprocessingconfig={
"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.

fromdataclassesimportdataclass, field@dataclassclassDefenseResult:
allowed: bool# gating decision (respects block_high_risk)risk_level: RiskLevel# diagnostic; max of Tier 1 / Tier 2sanitized: Any# sentence-cleaned copy (input verbatim when sanitize_content=False); dropped runs leave a [CONTENT SANITISED] marker; best-effort, still gate on alloweddetections: list[str] # Tier 1 pattern names detectedfields_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_fieldpatterns_by_field: dict[str, list[str]] # patterns detected per fielddetected_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=Nonetier2_raw_score: float|None=Nonetier2_aux_score: float|None=None# multi-head models onlytier2_multihead_blocked: bool|None=Nonetier2_skip_reason: str|None=Nonemax_sentence: str|None=Nonetier3: Tier3Result|None=None# present when Tier 3 ranfields_dropped: list[str] =field(default_factory=list)
truncated_at_depth: bool|None=Nonelatency_ms: float=0.0# Cost telemetry — present only when the batched Tier 2 classifier ranphase_timings: PhaseTimings|None=None# prepare / infer / aggregate mstier2_stats: Tier2Stats|None=None# string/chunk/unique counts, real/padded tokenstier1_ms: float|None=Nonecold_load: bool|None=None# Operational signalstier2_available: bool|None=None# False when Tier 2 enabled but failed to loadcoverage_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.

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"},
])
forrinresults:
ifnotr.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).

results=awaitdefense.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

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

Integration example

fromstackone_defenderimportcreate_prompt_defensedefense=create_prompt_defense(block_high_risk=True)
defense.warmup_tier2()
defrun_tool_and_defend(raw_result: dict, tool_name: str):
outcome=defense.defend_tool_result(raw_result, tool_name)
ifnotoutcome.allowed:
return {"error": "Content blocked by safety filter", "risk_level": outcome.risk_level}
returnoutcome.sanitized# Example agent loopsanitized=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 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 patternScanned 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

uv sync --group dev
uv run pytest

License

Apache-2.0 — see LICENSE.

About

No description, website, or topics provided.

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Used by

Contributors

Languages