stackone-defender 0.7.3__tar.gz → 0.8.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. stackone_defender-0.8.2/.release-please-manifest.json +1 -0
  2. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/CHANGELOG.md +33 -0
  3. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/PKG-INFO +35 -18
  4. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/README.md +33 -16
  5. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/pyproject.toml +1 -1
  6. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/onnx_classifier.py +163 -23
  7. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/patterns.py +6 -2
  8. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/tier2_classifier.py +52 -10
  9. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/config.py +6 -1
  10. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/core/prompt_defense.py +215 -32
  11. stackone_defender-0.8.2/src/stackone_defender/core/sentence_cleaner.py +82 -0
  12. stackone_defender-0.8.2/src/stackone_defender/core/tool_result_sanitizer.py +460 -0
  13. stackone_defender-0.8.2/src/stackone_defender/sanitizers/__init__.py +15 -0
  14. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/encoding_detector.py +3 -1
  15. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/types.py +60 -3
  16. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_integration.py +278 -12
  17. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_onnx_classifier.py +73 -4
  18. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_sanitizers.py +0 -194
  19. stackone_defender-0.8.2/tests/test_sentence_cleaner.py +87 -0
  20. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_tier2_classifier.py +96 -2
  21. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_tier3.py +1 -1
  22. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/uv.lock +10 -10
  23. stackone_defender-0.7.3/.release-please-manifest.json +0 -1
  24. stackone_defender-0.7.3/src/stackone_defender/core/tool_result_sanitizer.py +0 -361
  25. stackone_defender-0.7.3/src/stackone_defender/sanitizers/__init__.py +0 -31
  26. stackone_defender-0.7.3/src/stackone_defender/sanitizers/pattern_remover.py +0 -113
  27. stackone_defender-0.7.3/src/stackone_defender/sanitizers/sanitizer.py +0 -243
  28. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/.gitignore +0 -0
  29. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/.python-version +0 -0
  30. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/.release-please-config.json +0 -0
  31. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/LICENSE +0 -0
  32. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/__init__.py +0 -0
  33. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/__init__.py +0 -0
  34. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/pattern_detector.py +0 -0
  35. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/tier3_orchestrator.py +0 -0
  36. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/core/__init__.py +0 -0
  37. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/classifier_config.json +0 -0
  38. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/config.json +0 -0
  39. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/model_quantized.onnx +0 -0
  40. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/tokenizer.json +0 -0
  41. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/tokenizer_config.json +0 -0
  42. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/py.typed +0 -0
  43. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/leet_normalizer.py +0 -0
  44. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/normalizer.py +0 -0
  45. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/role_stripper.py +0 -0
  46. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sfe/__init__.py +0 -0
  47. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sfe/model.ftz +0 -0
  48. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/sfe/preprocess.py +0 -0
  49. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/utils/__init__.py +0 -0
  50. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/utils/boundary.py +0 -0
  51. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/utils/field_detection.py +0 -0
  52. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/src/stackone_defender/utils/structure.py +0 -0
  53. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/__init__.py +0 -0
  54. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_pattern_detector.py +0 -0
  55. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_sfe.py +0 -0
  56. {stackone_defender-0.7.3 → stackone_defender-0.8.2}/tests/test_utils.py +0 -0
@@ -0,0 +1 @@
1
+ {".":"0.8.2"}
@@ -1,5 +1,38 @@
1
1
  # Changelog
2
2
 
3
+ ## [0.8.2](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.4...stackone-defender-v0.8.2) (2026-08-19)
4
+
5
+
6
+ ### Features
7
+
8
+ * **ENG-1084:** detect-and-gate + sentence-level cleaning — Tier 1 detects without mutating content (gate on `allowed`); `sanitized` is a sentence-cleaned copy of high/critical fields ([60b01d0](https://github.com/StackOneHQ/defender-py/commit/60b01d0117c04106a3017dc4bf536cce6cc7a6c4))
9
+ * **ENG-1084:** mark dropped runs with `[CONTENT SANITISED]` in the cleaned copy ([fdcdb94](https://github.com/StackOneHQ/defender-py/commit/fdcdb94e27698ac71a8dd8933d54ad47455be376))
10
+ * **ENG-1084:** `detected_field_count` threat-count signal; object-key scanning; Tier 2 availability (`require_tier2` / `tier2_available`) ([6f1b52a](https://github.com/StackOneHQ/defender-py/commit/6f1b52afd9d2932aab9462776c102d05e5bd9626))
11
+ * **ENG-1084:** NFKC-fold unicode before Tier 2 classification ([50741ee](https://github.com/StackOneHQ/defender-py/commit/50741eeb7c9d7e435a2a504d2f83c3bd9b2b8c04))
12
+
13
+
14
+ ### Bug Fixes
15
+
16
+ * **ENG-1084:** return large arrays in full — 0.7.4 truncated arrays >1000 items to 100 + a sentinel, dropping the rest; Tier 1 detection is now bounded by the `max_size` budget (flagged via `coverage_degraded`), never drops data ([26c7eec](https://github.com/StackOneHQ/defender-py/commit/26c7eecb10f55be0a9e7cb614a67629fa485d199))
17
+ * **ENG-1084:** scan strings inside risky array fields (`{"name": [INJ]}`) ([6f1b52a](https://github.com/StackOneHQ/defender-py/commit/6f1b52afd9d2932aab9462776c102d05e5bd9626))
18
+ * **ENG-1084:** evidence-driven encoding escalation; token-degeneracy (OOD) guard; decorative-run normalization ([f450992](https://github.com/StackOneHQ/defender-py/commit/f45099241e36ee4210daaf39f2573c2a88d3071b))
19
+ * **ENG-1084:** gate sentence-cleaning on aggregate risk; verbatim when nothing is dropped ([9e25335](https://github.com/StackOneHQ/defender-py/commit/9e2533576f90574f8bab5476728d79c5275b9060))
20
+ * **ENG-1084:** count objects/arrays once in size metrics + iterate `obj.items()` directly ([26c7eec](https://github.com/StackOneHQ/defender-py/commit/26c7eecb10f55be0a9e7cb614a67629fa485d199))
21
+ * **ENG-1084:** warmup fail-open + dict-subclass stripping ([f49ffcf](https://github.com/StackOneHQ/defender-py/commit/f49ffcf6e93e322623447d217994cc8215e02410))
22
+
23
+ ## [0.7.4](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.3...stackone-defender-v0.7.4) (2026-08-13)
24
+
25
+
26
+ ### Features
27
+
28
+ * **ENG-1761:** port Tier 2 buckets, dedupe, and cost telemetry from TS 0.7.4 ([b92b9f6](https://github.com/StackOneHQ/defender-py/commit/b92b9f661ef927b27542a664d92477e039279758))
29
+ * **ENG-1761:** port Tier 2 fixed-width buckets, dedupe, cost telemetry from TS 0.7.4 ([1e63d09](https://github.com/StackOneHQ/defender-py/commit/1e63d0999d003348c17d6700bf6a03b65a7b2242))
30
+
31
+
32
+ ### Bug Fixes
33
+
34
+ * **ENG-1761:** address review — cold_load on failure paths + stronger tests ([e897786](https://github.com/StackOneHQ/defender-py/commit/e897786ff013b95b7b74eecd8df554ee978dd98b))
35
+
3
36
  ## [0.7.3](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.2...stackone-defender-v0.7.3) (2026-07-22)
4
37
 
5
38
 
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: stackone-defender
3
- Version: 0.7.3
3
+ Version: 0.8.2
4
4
  Summary: Indirect prompt injection defense for AI agents using tool calls
5
5
  Project-URL: Homepage, https://github.com/StackOneHQ/stackone-defender
6
6
  Project-URL: Repository, https://github.com/StackOneHQ/stackone-defender
@@ -116,10 +116,12 @@ else:
116
116
 
117
117
  ### Tier 1 — Pattern detection (sync, ~1 ms)
118
118
 
119
- - **Unicode normalization** — homoglyph resistance (e.g. Cyrillic `а` → ASCII `a`)
120
- - **Role stripping** — `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]`, etc.
121
- - **Pattern removal** — phrases like “ignore previous instructions”
122
- - **Encoding detection** — suspicious Base64/URL-shaped payloads
119
+ Detects (never rewrites) injection signals and records them as evidence:
120
+
121
+ - **Unicode normalization** (for matching only) — homoglyph resistance (e.g. Cyrillic `а` → ASCII `a`)
122
+ - **Role markers** — `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]`, etc.
123
+ - **Instruction-override patterns** — phrases like “ignore previous instructions”
124
+ - **Encoding detection** — decode-then-detect: a decoded payload is escalated only if it trips a real attack pattern
123
125
  - **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.
124
126
 
125
127
  ### Tier 2 — ML classification (ONNX)
@@ -133,7 +135,7 @@ Packed-chunk MiniLM classifier (int8 ONNX ~22 MB, bundled):
133
135
  ### Optional SFE preprocessor
134
136
 
135
137
  - `use_sfe=True` runs a field-level FastText pass to build a **classifier-only** view of the payload
136
- - **Tier 1** always sanitizes the **original** tool value; **`sanitized`** in `DefenseResult` is unchanged by SFE drops
138
+ - **Tier 1** detects on the raw tool value; SFE drops are classifier-only and never remove fields from the returned `sanitized` payload
137
139
  - **Tier 2** extracts strings from the SFE-filtered tree; `fields_dropped` lists paths omitted from that extraction (not removed from `sanitized`)
138
140
  - Fails open if the runtime/model is unavailable: payload continues unfiltered
139
141
 
@@ -152,7 +154,7 @@ Authoritative LLM-based classification for the cases Tier 2 finds ambiguous. The
152
154
 
153
155
  Two modes, selected via `defender_mode`:
154
156
  - **`"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.
155
- - **`"tier3_only"`**: skip Tier 2; the block/allow decision is the Tier 3 verdict alone. Tier 1 still sanitizes the returned `sanitized` payload.
157
+ - **`"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.
156
158
 
157
159
  Register a provider once at startup, then opt in per instance:
158
160
 
@@ -189,8 +191,9 @@ defense = create_prompt_defense(
189
191
 
190
192
  ### `allowed` vs `risk_level`
191
193
 
194
+ - **`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.
192
195
  - Use **`allowed`** for gating when `block_high_risk=True`: `False` means do not pass `sanitized` to the model as-is.
193
- - **`risk_level`** is diagnostic: it starts at `default_risk_level` (default `"medium"`) 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.
196
+ - **`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.
194
197
 
195
198
  | Level | Typical trigger |
196
199
  |-------|------------------|
@@ -206,8 +209,9 @@ defense = create_prompt_defense(
206
209
  defense = create_prompt_defense(
207
210
  enable_tier1=True,
208
211
  enable_tier2=True,
212
+ require_tier2=False, # True: raise if Tier 2 can't load (fail closed) instead of degrading to Tier 1
209
213
  block_high_risk=False,
210
- default_risk_level="medium",
214
+ default_risk_level="low",
211
215
  annotate_boundary=False, # True: wrap risky strings with [UD-…] tags (npm: annotateBoundary)
212
216
  tier2_fields=["subject", "body", "snippet"], # optional: scope Tier 2 to these JSON keys (default: all strings)
213
217
  use_sfe=True, # optional: enable semantic field extractor preprocessing
@@ -229,18 +233,31 @@ from dataclasses import dataclass, field
229
233
 
230
234
  @dataclass
231
235
  class DefenseResult:
232
- allowed: bool
233
- risk_level: RiskLevel
234
- sanitized: Any
235
- detections: list[str]
236
- fields_sanitized: list[str]
237
- patterns_by_field: dict[str, list[str]]
236
+ allowed: bool # gating decision (respects block_high_risk)
237
+ risk_level: RiskLevel # diagnostic; max of Tier 1 / Tier 2
238
+ sanitized: Any # sentence-cleaned copy (input verbatim when sanitize_content=False); dropped runs leave a [CONTENT SANITISED] marker; best-effort, still gate on allowed
239
+ detections: list[str] # Tier 1 pattern names detected
240
+ 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
241
+ patterns_by_field: dict[str, list[str]] # patterns detected per field
242
+ 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)
238
243
  tier2_score: float | None = None
244
+ tier2_raw_score: float | None = None
245
+ tier2_aux_score: float | None = None # multi-head models only
246
+ tier2_multihead_blocked: bool | None = None
239
247
  tier2_skip_reason: str | None = None
240
248
  max_sentence: str | None = None
249
+ tier3: Tier3Result | None = None # present when Tier 3 ran
241
250
  fields_dropped: list[str] = field(default_factory=list)
242
251
  truncated_at_depth: bool | None = None
243
252
  latency_ms: float = 0.0
253
+ # Cost telemetry — present only when the batched Tier 2 classifier ran
254
+ phase_timings: PhaseTimings | None = None # prepare / infer / aggregate ms
255
+ tier2_stats: Tier2Stats | None = None # string/chunk/unique counts, real/padded tokens
256
+ tier1_ms: float | None = None
257
+ cold_load: bool | None = None
258
+ # Operational signals
259
+ tier2_available: bool | None = None # False when Tier 2 enabled but failed to load
260
+ coverage_degraded: bool | None = None # True when Tier 1 detection coverage was capped
244
261
  ```
245
262
 
246
263
  ### `defense.defend_tool_results(items)`
@@ -255,7 +272,7 @@ results = defense.defend_tool_results([
255
272
  ])
256
273
  for r in results:
257
274
  if not r.allowed:
258
- print("Blocked:", ", ".join(r.fields_sanitized))
275
+ print("Blocked:", ", ".join(r.detections))
259
276
  ```
260
277
 
261
278
  ### `await defense.defend_tool_results_async(items)`
@@ -300,7 +317,7 @@ sanitized = run_tool_and_defend(gmail_api.get_message(msg_id), "gmail_get_messag
300
317
 
301
318
  ## Risky field detection
302
319
 
303
- Only **string** values under configured “risky” keys are scanned and sanitized. [`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.
320
+ 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.)
304
321
 
305
322
  | Tool pattern | Scanned fields |
306
323
  |--------------|----------------|
@@ -90,10 +90,12 @@ else:
90
90
 
91
91
  ### Tier 1 — Pattern detection (sync, ~1 ms)
92
92
 
93
- - **Unicode normalization** — homoglyph resistance (e.g. Cyrillic `а` → ASCII `a`)
94
- - **Role stripping** — `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]`, etc.
95
- - **Pattern removal** — phrases like “ignore previous instructions”
96
- - **Encoding detection** — suspicious Base64/URL-shaped payloads
93
+ Detects (never rewrites) injection signals and records them as evidence:
94
+
95
+ - **Unicode normalization** (for matching only) — homoglyph resistance (e.g. Cyrillic `а` → ASCII `a`)
96
+ - **Role markers** — `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]`, etc.
97
+ - **Instruction-override patterns** — phrases like “ignore previous instructions”
98
+ - **Encoding detection** — decode-then-detect: a decoded payload is escalated only if it trips a real attack pattern
97
99
  - **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.
98
100
 
99
101
  ### Tier 2 — ML classification (ONNX)
@@ -107,7 +109,7 @@ Packed-chunk MiniLM classifier (int8 ONNX ~22 MB, bundled):
107
109
  ### Optional SFE preprocessor
108
110
 
109
111
  - `use_sfe=True` runs a field-level FastText pass to build a **classifier-only** view of the payload
110
- - **Tier 1** always sanitizes the **original** tool value; **`sanitized`** in `DefenseResult` is unchanged by SFE drops
112
+ - **Tier 1** detects on the raw tool value; SFE drops are classifier-only and never remove fields from the returned `sanitized` payload
111
113
  - **Tier 2** extracts strings from the SFE-filtered tree; `fields_dropped` lists paths omitted from that extraction (not removed from `sanitized`)
112
114
  - Fails open if the runtime/model is unavailable: payload continues unfiltered
113
115
 
@@ -126,7 +128,7 @@ Authoritative LLM-based classification for the cases Tier 2 finds ambiguous. The
126
128
 
127
129
  Two modes, selected via `defender_mode`:
128
130
  - **`"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.
129
- - **`"tier3_only"`**: skip Tier 2; the block/allow decision is the Tier 3 verdict alone. Tier 1 still sanitizes the returned `sanitized` payload.
131
+ - **`"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.
130
132
 
131
133
  Register a provider once at startup, then opt in per instance:
132
134
 
@@ -163,8 +165,9 @@ defense = create_prompt_defense(
163
165
 
164
166
  ### `allowed` vs `risk_level`
165
167
 
168
+ - **`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.
166
169
  - Use **`allowed`** for gating when `block_high_risk=True`: `False` means do not pass `sanitized` to the model as-is.
167
- - **`risk_level`** is diagnostic: it starts at `default_risk_level` (default `"medium"`) 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.
170
+ - **`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.
168
171
 
169
172
  | Level | Typical trigger |
170
173
  |-------|------------------|
@@ -180,8 +183,9 @@ defense = create_prompt_defense(
180
183
  defense = create_prompt_defense(
181
184
  enable_tier1=True,
182
185
  enable_tier2=True,
186
+ require_tier2=False, # True: raise if Tier 2 can't load (fail closed) instead of degrading to Tier 1
183
187
  block_high_risk=False,
184
- default_risk_level="medium",
188
+ default_risk_level="low",
185
189
  annotate_boundary=False, # True: wrap risky strings with [UD-…] tags (npm: annotateBoundary)
186
190
  tier2_fields=["subject", "body", "snippet"], # optional: scope Tier 2 to these JSON keys (default: all strings)
187
191
  use_sfe=True, # optional: enable semantic field extractor preprocessing
@@ -203,18 +207,31 @@ from dataclasses import dataclass, field
203
207
 
204
208
  @dataclass
205
209
  class DefenseResult:
206
- allowed: bool
207
- risk_level: RiskLevel
208
- sanitized: Any
209
- detections: list[str]
210
- fields_sanitized: list[str]
211
- patterns_by_field: dict[str, list[str]]
210
+ allowed: bool # gating decision (respects block_high_risk)
211
+ risk_level: RiskLevel # diagnostic; max of Tier 1 / Tier 2
212
+ sanitized: Any # sentence-cleaned copy (input verbatim when sanitize_content=False); dropped runs leave a [CONTENT SANITISED] marker; best-effort, still gate on allowed
213
+ detections: list[str] # Tier 1 pattern names detected
214
+ 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
215
+ patterns_by_field: dict[str, list[str]] # patterns detected per field
216
+ 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)
212
217
  tier2_score: float | None = None
218
+ tier2_raw_score: float | None = None
219
+ tier2_aux_score: float | None = None # multi-head models only
220
+ tier2_multihead_blocked: bool | None = None
213
221
  tier2_skip_reason: str | None = None
214
222
  max_sentence: str | None = None
223
+ tier3: Tier3Result | None = None # present when Tier 3 ran
215
224
  fields_dropped: list[str] = field(default_factory=list)
216
225
  truncated_at_depth: bool | None = None
217
226
  latency_ms: float = 0.0
227
+ # Cost telemetry — present only when the batched Tier 2 classifier ran
228
+ phase_timings: PhaseTimings | None = None # prepare / infer / aggregate ms
229
+ tier2_stats: Tier2Stats | None = None # string/chunk/unique counts, real/padded tokens
230
+ tier1_ms: float | None = None
231
+ cold_load: bool | None = None
232
+ # Operational signals
233
+ tier2_available: bool | None = None # False when Tier 2 enabled but failed to load
234
+ coverage_degraded: bool | None = None # True when Tier 1 detection coverage was capped
218
235
  ```
219
236
 
220
237
  ### `defense.defend_tool_results(items)`
@@ -229,7 +246,7 @@ results = defense.defend_tool_results([
229
246
  ])
230
247
  for r in results:
231
248
  if not r.allowed:
232
- print("Blocked:", ", ".join(r.fields_sanitized))
249
+ print("Blocked:", ", ".join(r.detections))
233
250
  ```
234
251
 
235
252
  ### `await defense.defend_tool_results_async(items)`
@@ -274,7 +291,7 @@ sanitized = run_tool_and_defend(gmail_api.get_message(msg_id), "gmail_get_messag
274
291
 
275
292
  ## Risky field detection
276
293
 
277
- Only **string** values under configured “risky” keys are scanned and sanitized. [`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.
294
+ 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.)
278
295
 
279
296
  | Tool pattern | Scanned fields |
280
297
  |--------------|----------------|
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "stackone-defender"
3
- version = "0.7.3"
3
+ version = "0.8.2"
4
4
  description = "Indirect prompt injection defense for AI agents using tool calls"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.11"
@@ -11,11 +11,20 @@ from __future__ import annotations
11
11
  import logging
12
12
  import math
13
13
  import threading
14
+ from dataclasses import dataclass
14
15
  from pathlib import Path
15
- from typing import Literal
16
+ from typing import Literal, cast
16
17
 
17
18
  _logger = logging.getLogger(__name__)
18
19
 
20
+
21
+ @dataclass
22
+ class BatchTokenStats:
23
+ """Padding accounting accumulated across a batched classification (#6 telemetry)."""
24
+
25
+ real_tokens: int = 0 # sum of real (non-pad) tokens across all chunks
26
+ padded_tokens: int = 0 # tokens actually run through ONNX, including padding
27
+
19
28
  # Shared across all OnnxClassifier instances (keyed by resolved model dir path).
20
29
  # Tuple is (session, tokenizer, count_tokenizer); the count tokenizer is
21
30
  # non-truncating so count_tokens reports true length (see _load_model).
@@ -61,8 +70,28 @@ class OnnxClassifier:
61
70
  """
62
71
 
63
72
  _MAX_BATCH_CHUNK = 32
64
-
65
- def __init__(self, model_path: str | None = None, temperature_t: float | None = None):
73
+ # Fixed pad widths (ascending). A string pads to the smallest bucket >= its
74
+ # token length, so its padded length — and thus its quantized score — is a
75
+ # function of the string alone, not its batch neighbours (deterministic).
76
+ _PAD_BUCKETS = (32, 64, 128, 256)
77
+
78
+ # Token-degeneracy (OOD) guard. Below this many content tokens, token-share
79
+ # is too coarse to mean anything (a 1-2 token row is trivially "dominated");
80
+ # matches the shortest decorative run that still false-fires post-collapse
81
+ # (``─``x3 -> 3 content tokens). Damp only when the row draws on at most
82
+ # _DEGENERACY_MAX_DISTINCT_TOKENS distinct tokens — padding an attack with
83
+ # copies of one token can push the share past 2/3 but cannot remove the
84
+ # attack's own vocabulary, so the distinct-token floor is a structural, not
85
+ # tuned, backstop against that bypass.
86
+ _DEGENERACY_MIN_CONTENT_TOKENS = 3
87
+ _DEGENERACY_MAX_DISTINCT_TOKENS = 4
88
+
89
+ def __init__(
90
+ self,
91
+ model_path: str | None = None,
92
+ temperature_t: float | None = None,
93
+ degeneracy_max_token_share: float | None = None,
94
+ ):
66
95
  self._model_path = model_path or get_default_model_path()
67
96
  self._session = None
68
97
  self._tokenizer = None
@@ -70,6 +99,14 @@ class OnnxClassifier:
70
99
  self._count_tokenizer = None
71
100
  self._max_length = 256
72
101
  self._load_failed = False
102
+ # Token-degeneracy guard threshold; > 1 disables. See _is_degenerate.
103
+ self._degeneracy_max_token_share = 2 / 3
104
+ if degeneracy_max_token_share is not None and math.isfinite(degeneracy_max_token_share):
105
+ self._degeneracy_max_token_share = float(degeneracy_max_token_share)
106
+ # Cached [UNK] id, resolved lazily from the loaded tokenizer. ``None`` =
107
+ # not yet resolved or no [UNK] concept; see _get_unk_token_id.
108
+ self._unk_token_id: int | None = None
109
+ self._unk_resolved = False
73
110
  # Output mode is detected lazily from the logits shape on the first
74
111
  # inference call. ``None`` until then.
75
112
  self._output_mode: Literal["single", "multi"] | None = None
@@ -132,8 +169,10 @@ class OnnxClassifier:
132
169
  import onnxruntime as ort
133
170
  from tokenizers import Tokenizer
134
171
  except ImportError as e:
172
+ # No warning here -- the ImportError propagates to the caller,
173
+ # which owns user-facing messaging (PromptDefense warns once per
174
+ # instance). Warning here logged a line on every failed call.
135
175
  self._load_failed = True
136
- _logger.warning("[defender] ONNX model failed to load: %s", e)
137
176
  raise ImportError(
138
177
  "ONNX dependencies not installed. Install with: pip install stackone-defender[onnx]"
139
178
  ) from e
@@ -142,7 +181,10 @@ class OnnxClassifier:
142
181
  tokenizer_path = str(Path(self._model_path) / "tokenizer.json")
143
182
  self._tokenizer = Tokenizer.from_file(tokenizer_path)
144
183
  self._tokenizer.enable_truncation(max_length=self._max_length)
145
- self._tokenizer.enable_padding(length=self._max_length)
184
+ # No fixed padding: classify_batch_pair pads each batch per bucket
185
+ # width (see _PAD_BUCKETS), so a string's padded length depends only
186
+ # on itself. Single classify runs at the string's real length.
187
+ self._tokenizer.no_padding()
146
188
 
147
189
  # tokenizer.json bakes in truncation+padding, so disable both
148
190
  # here to count true length rather than a value capped at
@@ -182,6 +224,10 @@ class OnnxClassifier:
182
224
  import numpy as np
183
225
 
184
226
  encoding = self._tokenizer.encode(text)
227
+ # Fix 3: token-degeneracy guard — skip inference on off-distribution
228
+ # input; its mean-pooled score is arbitrary. Damp to a benign 0.
229
+ if self._is_degenerate(encoding.ids):
230
+ return 0.0, None
185
231
  input_ids = np.array([encoding.ids], dtype=np.int64)
186
232
  attention_mask = np.array([encoding.attention_mask], dtype=np.int64)
187
233
 
@@ -199,36 +245,130 @@ class OnnxClassifier:
199
245
  main_logit = float(row[0]) if hasattr(row, "__len__") and len(row) > 0 else float(row)
200
246
  return _sigmoid(main_logit / t), None
201
247
 
202
- def classify_batch(self, texts: list[str]) -> list[float]:
248
+ def classify_batch(self, texts: list[str], stats: BatchTokenStats | None = None) -> list[float]:
203
249
  """Classify multiple texts; returns main-head scores only.
204
250
 
205
251
  Back-compat wrapper around :meth:`classify_batch_pair`.
206
252
  """
207
- return [main for main, _ in self.classify_batch_pair(texts)]
253
+ return [main for main, _ in self.classify_batch_pair(texts, stats)]
208
254
 
209
- def classify_batch_pair(self, texts: list[str]) -> list[tuple[float, float | None]]:
255
+ def classify_batch_pair(
256
+ self, texts: list[str], stats: BatchTokenStats | None = None
257
+ ) -> list[tuple[float, float | None]]:
210
258
  """Classify multiple texts, returning ``(main, aux)`` per row.
211
259
 
212
- Aux is ``None`` per-row for single-head models. Chunks the input to
213
- bound native memory; the attention matrix is ``O(chunk * seq_len^2)``,
214
- and for MiniLM (``max_length=256``) a chunk of 32 keeps memory
215
- under ~50MB per call.
260
+ Aux is ``None`` per-row for single-head models. Groups strings into fixed
261
+ pad-width buckets by token length, batches within each bucket, and pads
262
+ every row to the BUCKET width (not the batch max). This bounds padding
263
+ waste and makes a string's padded length — hence its quantized score —
264
+ depend only on itself, not its batch neighbours. Chunks of 32 bound
265
+ native memory (attention is ``O(chunk * seq_len^2)``). Results scatter
266
+ back to input order. ``stats`` accumulates padding counts when provided.
216
267
  """
217
268
  if not texts:
218
269
  return []
219
270
  self._ensure_loaded()
220
- all_pairs: list[tuple[float, float | None]] = []
221
- for offset in range(0, len(texts), self._MAX_BATCH_CHUNK):
222
- chunk = texts[offset : offset + self._MAX_BATCH_CHUNK]
223
- all_pairs.extend(self._classify_batch_chunk_pair(chunk))
224
- return all_pairs
225
271
 
226
- def _classify_batch_chunk_pair(self, texts: list[str]) -> list[tuple[float, float | None]]:
272
+ encodings = self._tokenizer.encode_batch(texts)
273
+ last_bucket = self._PAD_BUCKETS[-1]
274
+ buckets: dict[int, list[int]] = {}
275
+ for i, enc in enumerate(encodings):
276
+ length = len(enc.ids)
277
+ width = next((b for b in self._PAD_BUCKETS if length <= b), last_bucket)
278
+ buckets.setdefault(width, []).append(i)
279
+
280
+ pairs: list[tuple[float, float | None] | None] = [None] * len(texts)
281
+ for width in self._PAD_BUCKETS:
282
+ bucket_idxs = buckets.get(width)
283
+ if not bucket_idxs:
284
+ continue
285
+ for offset in range(0, len(bucket_idxs), self._MAX_BATCH_CHUNK):
286
+ idxs = bucket_idxs[offset : offset + self._MAX_BATCH_CHUNK]
287
+ chunk = [encodings[i] for i in idxs]
288
+ if stats is not None:
289
+ # Every row pads to the bucket width; real_tokens is useful work.
290
+ stats.padded_tokens += len(chunk) * width
291
+ for enc in chunk:
292
+ stats.real_tokens += len(enc.ids)
293
+ chunk_pairs = self._classify_batch_chunk_pair(chunk, width)
294
+ for k, orig_idx in enumerate(idxs):
295
+ pairs[orig_idx] = chunk_pairs[k]
296
+
297
+ # Fix 3: token-degeneracy guard — damp off-distribution rows to a benign
298
+ # 0 so they drop out of any upstream max. Reuses the already-computed ids.
299
+ for i, enc in enumerate(encodings):
300
+ if self._is_degenerate(enc.ids):
301
+ pairs[i] = (0.0, None)
302
+
303
+ return cast(list[tuple[float, float | None]], pairs)
304
+
305
+ def _get_unk_token_id(self) -> int | None:
306
+ """Resolve the tokenizer's ``[UNK]`` id, cached. Returns ``None`` when the
307
+ tokenizer has no ``[UNK]`` concept (the guard then skips its factor-3 check)."""
308
+ if self._unk_resolved:
309
+ return self._unk_token_id
310
+ self._unk_resolved = True
311
+ try:
312
+ self._unk_token_id = self._tokenizer.token_to_id("[UNK]")
313
+ except Exception:
314
+ self._unk_token_id = None
315
+ return self._unk_token_id
316
+
317
+ def _is_degenerate(self, ids: list[int]) -> bool:
318
+ """Token-degeneracy (OOD) test over a tokenized row. Damps only when ALL
319
+ of these hold over the content tokens (excluding [CLS]/[SEP]):
320
+
321
+ 1. the single most-frequent token covers >= ``degeneracy_max_token_share``,
322
+ 2. the row draws on <= ``_DEGENERACY_MAX_DISTINCT_TOKENS`` distinct tokens, and
323
+ 3. the dominant token is NOT [UNK].
324
+
325
+ Factor 2 blocks a padding attack; factor 3 blocks a homoglyph attack —
326
+ fullwidth / zero-width / other OOV chars collapse to repeated [UNK], the
327
+ signature of encoding evasion (more suspicious, not less), so those rows
328
+ are left to score rather than suppressed. Reuses the ids the model runs on.
329
+ """
330
+ threshold = self._degeneracy_max_token_share
331
+ if not (0 < threshold <= 1): # disabled
332
+ return False
333
+ has_specials = len(ids) >= 2
334
+ content = ids[1:-1] if has_specials else ids
335
+ n = len(content)
336
+ if n < self._DEGENERACY_MIN_CONTENT_TOKENS:
337
+ return False
338
+ counts: dict[int, int] = {}
339
+ max_freq = 0
340
+ dominant_id = -1
341
+ for tok in content:
342
+ c = counts.get(tok, 0) + 1
343
+ counts[tok] = c
344
+ if c > max_freq:
345
+ max_freq = c
346
+ dominant_id = tok
347
+ if max_freq / n < threshold or len(counts) > self._DEGENERACY_MAX_DISTINCT_TOKENS:
348
+ return False
349
+ unk = self._get_unk_token_id()
350
+ return unk is None or dominant_id != unk
351
+
352
+ def _classify_batch_chunk_pair(
353
+ self, encodings: list, pad_to: int | None = None
354
+ ) -> list[tuple[float, float | None]]:
355
+ """Run one pre-tokenized chunk through ONNX, padding rows to ``pad_to``.
356
+
357
+ Rows are padded to at least the chunk's real max, so a caller can never
358
+ under-pad. Inputs are pre-tokenized so ``classify_batch_pair`` can
359
+ length-bucket without re-tokenizing.
360
+ """
227
361
  import numpy as np
228
362
 
229
- encodings = self._tokenizer.encode_batch(texts)
230
- input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
231
- attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
363
+ actual_max = max(len(e.ids) for e in encodings)
364
+ max_len = max(pad_to, actual_max) if pad_to is not None else actual_max
365
+ batch_size = len(encodings)
366
+ input_ids = np.zeros((batch_size, max_len), dtype=np.int64)
367
+ attention_mask = np.zeros((batch_size, max_len), dtype=np.int64)
368
+ for i, enc in enumerate(encodings):
369
+ n = len(enc.ids)
370
+ input_ids[i, :n] = enc.ids
371
+ attention_mask[i, :n] = enc.attention_mask
232
372
 
233
373
  results = self._session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
234
374
  logits = results[0]
@@ -237,12 +377,12 @@ class OnnxClassifier:
237
377
  t = self._temperature_t
238
378
  pairs: list[tuple[float, float | None]] = []
239
379
  if self._output_mode == "multi":
240
- for i in range(len(texts)):
380
+ for i in range(batch_size):
241
381
  main = _sigmoid(float(logits[i][0]) / t)
242
382
  aux = _sigmoid(float(logits[i][1]) / t)
243
383
  pairs.append((main, aux))
244
384
  else:
245
- for i in range(len(texts)):
385
+ for i in range(batch_size):
246
386
  row = logits[i]
247
387
  # ``row`` may be a scalar (shape ``[batch]``) or 1-vector.
248
388
  main_logit = float(row[0]) if hasattr(row, "__len__") and len(row) > 0 else float(row)
@@ -134,7 +134,7 @@ ENCODING_SUSPICIOUS_PATTERNS: list[PatternDefinition] = [
134
134
  PatternDefinition("html_entity_abuse", re.compile(r"(?:&#\d{2,4};){4,}|(?:&#x[0-9a-fA-F]{2,4};){4,}", re.I), "encoding_suspicious", "medium", "HTML entity encoding (potential obfuscation)"),
135
135
  PatternDefinition("rot13_mention", re.compile(r"rot13|caesar\s+cipher|decode\s+this", re.I), "encoding_suspicious", "medium", "Mention of ROT13 or similar encoding schemes"),
136
136
  PatternDefinition("binary_string_encoding", re.compile(r"\b[01]{8}(?:\s+[01]{8}){2,}\b"), "encoding_suspicious", "medium", "Binary-encoded string (potential obfuscation)"),
137
- PatternDefinition("morse_code_encoding", re.compile(r"(?:[.-]+\s){4,}[.-]+"), "encoding_suspicious", "low", "Morse code pattern (potential obfuscation)"),
137
+ PatternDefinition("morse_code_encoding", re.compile(r"(?:[.-]{1,8}\s){4,}[.-]{1,8}"), "encoding_suspicious", "low", "Morse code pattern (potential obfuscation)"),
138
138
  PatternDefinition("leetspeak_injection", re.compile(r"1gn0r3|f0rg3t|byp4ss|syst3m|4dm1n|h4ck", re.I), "encoding_suspicious", "medium", "Leetspeak obfuscation of injection keywords"),
139
139
  ]
140
140
 
@@ -170,10 +170,14 @@ INDIRECT_INJECTION_PATTERNS: list[PatternDefinition] = [
170
170
  # ``[config](https://.../system-setup)`` triggered. Real smuggled-
171
171
  # instruction attacks include the full "ignore (all|the|previous|prior)"
172
172
  # phrasing in the URL/anchor.
173
+ # Negated, bounded char classes ([^\]] / [^)]) instead of ``.*?`` so the
174
+ # regex is linear -- the lazy-dot form could backtrack on long unclosed
175
+ # link/URL spans (ReDoS-prone).
173
176
  PatternDefinition(
174
177
  "markdown_hidden_instruction",
175
178
  re.compile(
176
- r"\[.*?\]\(.*?(?:ignore|disregard|forget|override)\W+(?:all|the|previous|prior)\W+.*?\)",
179
+ r"\[[^\]\n]{0,200}\]\([^)\n]{0,300}(?:ignore|disregard|forget|override)\W{1,8}"
180
+ r"(?:all|the|previous|prior)\W{1,8}[^)\n]{0,300}\)",
177
181
  re.I,
178
182
  ),
179
183
  "structural",