stackone-defender 0.7.4__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.4 → stackone_defender-0.8.2}/CHANGELOG.md +20 -0
  3. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/PKG-INFO +34 -17
  4. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/README.md +33 -16
  5. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/pyproject.toml +1 -1
  6. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/onnx_classifier.py +85 -2
  7. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/patterns.py +6 -2
  8. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/tier2_classifier.py +45 -5
  9. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/config.py +6 -1
  10. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/core/prompt_defense.py +139 -18
  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.4 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/encoding_detector.py +3 -1
  15. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/types.py +31 -3
  16. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_integration.py +205 -8
  17. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_sanitizers.py +0 -194
  18. stackone_defender-0.8.2/tests/test_sentence_cleaner.py +87 -0
  19. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_tier2_classifier.py +94 -0
  20. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/uv.lock +1 -1
  21. stackone_defender-0.7.4/.release-please-manifest.json +0 -1
  22. stackone_defender-0.7.4/src/stackone_defender/core/tool_result_sanitizer.py +0 -361
  23. stackone_defender-0.7.4/src/stackone_defender/sanitizers/__init__.py +0 -31
  24. stackone_defender-0.7.4/src/stackone_defender/sanitizers/pattern_remover.py +0 -113
  25. stackone_defender-0.7.4/src/stackone_defender/sanitizers/sanitizer.py +0 -243
  26. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/.gitignore +0 -0
  27. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/.python-version +0 -0
  28. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/.release-please-config.json +0 -0
  29. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/LICENSE +0 -0
  30. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/__init__.py +0 -0
  31. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/__init__.py +0 -0
  32. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/pattern_detector.py +0 -0
  33. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/classifiers/tier3_orchestrator.py +0 -0
  34. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/core/__init__.py +0 -0
  35. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/classifier_config.json +0 -0
  36. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/config.json +0 -0
  37. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/model_quantized.onnx +0 -0
  38. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/tokenizer.json +0 -0
  39. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/models/minilm-multihead-v5/tokenizer_config.json +0 -0
  40. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/py.typed +0 -0
  41. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/leet_normalizer.py +0 -0
  42. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/normalizer.py +0 -0
  43. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sanitizers/role_stripper.py +0 -0
  44. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sfe/__init__.py +0 -0
  45. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sfe/model.ftz +0 -0
  46. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/sfe/preprocess.py +0 -0
  47. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/utils/__init__.py +0 -0
  48. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/utils/boundary.py +0 -0
  49. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/utils/field_detection.py +0 -0
  50. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/src/stackone_defender/utils/structure.py +0 -0
  51. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/__init__.py +0 -0
  52. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_onnx_classifier.py +0 -0
  53. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_pattern_detector.py +0 -0
  54. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_sfe.py +0 -0
  55. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_tier3.py +0 -0
  56. {stackone_defender-0.7.4 → stackone_defender-0.8.2}/tests/test_utils.py +0 -0
@@ -0,0 +1 @@
1
+ {".":"0.8.2"}
@@ -1,5 +1,25 @@
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
+
3
23
  ## [0.7.4](https://github.com/StackOneHQ/defender-py/compare/stackone-defender-v0.7.3...stackone-defender-v0.7.4) (2026-08-13)
4
24
 
5
25
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: stackone-defender
3
- Version: 0.7.4
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.4"
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"
@@ -75,7 +75,23 @@ class OnnxClassifier:
75
75
  # function of the string alone, not its batch neighbours (deterministic).
76
76
  _PAD_BUCKETS = (32, 64, 128, 256)
77
77
 
78
- def __init__(self, model_path: str | None = None, temperature_t: float | None = None):
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
+ ):
79
95
  self._model_path = model_path or get_default_model_path()
80
96
  self._session = None
81
97
  self._tokenizer = None
@@ -83,6 +99,14 @@ class OnnxClassifier:
83
99
  self._count_tokenizer = None
84
100
  self._max_length = 256
85
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
86
110
  # Output mode is detected lazily from the logits shape on the first
87
111
  # inference call. ``None`` until then.
88
112
  self._output_mode: Literal["single", "multi"] | None = None
@@ -145,8 +169,10 @@ class OnnxClassifier:
145
169
  import onnxruntime as ort
146
170
  from tokenizers import Tokenizer
147
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.
148
175
  self._load_failed = True
149
- _logger.warning("[defender] ONNX model failed to load: %s", e)
150
176
  raise ImportError(
151
177
  "ONNX dependencies not installed. Install with: pip install stackone-defender[onnx]"
152
178
  ) from e
@@ -198,6 +224,10 @@ class OnnxClassifier:
198
224
  import numpy as np
199
225
 
200
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
201
231
  input_ids = np.array([encoding.ids], dtype=np.int64)
202
232
  attention_mask = np.array([encoding.attention_mask], dtype=np.int64)
203
233
 
@@ -264,8 +294,61 @@ class OnnxClassifier:
264
294
  for k, orig_idx in enumerate(idxs):
265
295
  pairs[orig_idx] = chunk_pairs[k]
266
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
+
267
303
  return cast(list[tuple[float, float | None]], pairs)
268
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
+
269
352
  def _classify_batch_chunk_pair(
270
353
  self, encodings: list, pad_to: int | None = None
271
354
  ) -> list[tuple[float, float | None]]:
@@ -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",
@@ -9,12 +9,14 @@ from __future__ import annotations
9
9
 
10
10
  import json
11
11
  import logging
12
+ import math
12
13
  import os
13
14
  import re
14
15
  import time
15
16
  from dataclasses import dataclass
16
17
  from typing import Any
17
18
 
19
+ from ..sanitizers.normalizer import normalize_unicode
18
20
  from ..types import MultiheadConfig, RiskLevel, Tier2Result
19
21
  from ..utils.boundary import strip_boundary_patterns
20
22
  from .onnx_classifier import BatchTokenStats, OnnxClassifier, get_default_model_path
@@ -26,6 +28,14 @@ DEFAULT_TIER2_CLASSIFIER_CONFIG: dict[str, Any] = {
26
28
  "medium_risk_threshold": 0.5,
27
29
  "min_text_length": 10,
28
30
  "max_text_length": 10000,
31
+ # Token-degeneracy (OOD) guard threshold. A chunk is dropped from Tier 2
32
+ # scoring (its mean-pooled score is arbitrary off-distribution) only when its
33
+ # most-frequent content token covers >= this share AND it uses few distinct
34
+ # tokens AND the dominant token is not [UNK] — repeated rule/box chars,
35
+ # base64/hex. The distinct-token floor keeps the share test from being padded
36
+ # around; the [UNK] exclusion keeps homoglyph/OOV rows from being suppressed.
37
+ # Tier 1 still catches literal and (via decode) encoded attacks. > 1 disables.
38
+ "degeneracy_max_token_share": 2 / 3,
29
39
  }
30
40
 
31
41
 
@@ -150,8 +160,11 @@ class Tier2Classifier:
150
160
  self._model_path: str = merged["onnx_model_path"]
151
161
  self._temperature_t: float | None = merged.get("temperature_t")
152
162
  self._multihead: MultiheadConfig | None = merged.get("multihead")
163
+ self._degeneracy_max_token_share: float = float(merged["degeneracy_max_token_share"])
153
164
 
154
- self._onnx = OnnxClassifier(self._model_path, self._temperature_t)
165
+ self._onnx = OnnxClassifier(
166
+ self._model_path, self._temperature_t, self._degeneracy_max_token_share
167
+ )
155
168
 
156
169
  # ------------------------------------------------------------------
157
170
  # Lifecycle
@@ -167,11 +180,25 @@ class Tier2Classifier:
167
180
  # Single-text classify
168
181
  # ------------------------------------------------------------------
169
182
 
183
+ # Collapse 4+ repeats of the same non-word char down to 3. Decorative runs
184
+ # (box-drawing ``─``, ``===``, ``---``, ``###``) tokenize one-token-per-char,
185
+ # so a rule line becomes ~85% one repeated token; under mean pooling the
186
+ # pooled vector lands off-distribution and the head returns an arbitrary
187
+ # (often high) score. Classifier input only — the returned payload is never
188
+ # modified. (\\w is Unicode-aware in Python, so accented letters are kept.)
189
+ _DECORATIVE_RUN = re.compile(r"([^\w\s])\1{3,}")
190
+
191
+ def _normalize_for_classification(self, text: str) -> str:
192
+ # NFKC-fold unicode (fullwidth/math-styled -> ASCII) so obfuscated attacks
193
+ # tokenize as real words instead of [UNK]; then strip boundary markers and
194
+ # collapse decorative runs. Classifier input only — payload never mutated.
195
+ return self._DECORATIVE_RUN.sub(r"\1\1\1", strip_boundary_patterns(normalize_unicode(text)))
196
+
170
197
  def classify(self, text: str) -> Tier2Result:
171
198
  start = time.perf_counter()
172
199
  # Strip defender's own boundary markers before tokenization so nested
173
200
  # tool-call chains and spoofed boundary patterns don't corrupt scores.
174
- text = strip_boundary_patterns(text)
201
+ text = self._normalize_for_classification(text)
175
202
  if len(text) < self._min_text_length:
176
203
  return Tier2Result(
177
204
  score=0,
@@ -185,6 +212,18 @@ class Tier2Classifier:
185
212
 
186
213
  try:
187
214
  main, aux = self._onnx.classify_pair(analysis_text)
215
+ # A non-finite score (NaN/Infinity) means the model produced no
216
+ # usable output. Report a SKIP, not score 0 -- score 0 yields
217
+ # confidence 1.0 (|0 - 0.5| * 2), making a broken inference look
218
+ # like a max-confidence benign classification.
219
+ if not math.isfinite(main):
220
+ return Tier2Result(
221
+ score=0,
222
+ confidence=0,
223
+ skipped=True,
224
+ skip_reason="Non-finite model output (NaN/Infinity)",
225
+ latency_ms=_ms(start),
226
+ )
188
227
  confidence = abs(main - 0.5) * 2
189
228
  return Tier2Result(
190
229
  score=main, confidence=confidence, skipped=False, latency_ms=_ms(start), aux=aux
@@ -208,7 +247,7 @@ class Tier2Classifier:
208
247
  def classify_by_sentence(self, text: str) -> dict[str, Any]:
209
248
  """Classify text by sentence and return max main score."""
210
249
  start = time.perf_counter()
211
- text = strip_boundary_patterns(text)
250
+ text = self._normalize_for_classification(text)
212
251
  sentences = _split_into_sentences(text)
213
252
  if not sentences:
214
253
  return _skipped(start, "No sentences found")
@@ -252,7 +291,7 @@ class Tier2Classifier:
252
291
 
253
292
  def classify_by_chunks(self, text: str) -> dict[str, Any]:
254
293
  start = time.perf_counter()
255
- text = strip_boundary_patterns(text)
294
+ text = self._normalize_for_classification(text)
256
295
  if len(text) < self._min_text_length:
257
296
  return _skipped(start, "Text below minTextLength")
258
297
 
@@ -316,7 +355,7 @@ class Tier2Classifier:
316
355
  }
317
356
 
318
357
  def prepare_chunks(self, text: str) -> dict[str, Any]:
319
- text = strip_boundary_patterns(text)
358
+ text = self._normalize_for_classification(text)
320
359
  if len(text) < self._min_text_length:
321
360
  return {"chunks": [], "skipped": True, "skip_reason": "Text below minTextLength"}
322
361
 
@@ -418,6 +457,7 @@ class Tier2Classifier:
418
457
  "onnx_model_path": self._model_path,
419
458
  "temperature_t": self.get_temperature(),
420
459
  "multihead": self._multihead,
460
+ "degeneracy_max_token_share": self._degeneracy_max_token_share,
421
461
  }
422
462
 
423
463
  def get_risk_level(self, score: float) -> RiskLevel:
@@ -54,7 +54,7 @@ DEFAULT_TRAVERSAL_CONFIG = TraversalConfig(
54
54
  max_depth=10,
55
55
  max_size=10 * 1024 * 1024,
56
56
  large_array_threshold=1000,
57
- skip_large_arrays=True,
57
+ skip_large_arrays=False, # deprecated legacy cap; detection now bounded by max_size
58
58
  )
59
59
 
60
60
  DEFAULT_CUMULATIVE_RISK_THRESHOLDS: dict[str, int | float] = {
@@ -65,6 +65,11 @@ DEFAULT_CUMULATIVE_RISK_THRESHOLDS: dict[str, int | float] = {
65
65
  "patterns_fraction": 0.25,
66
66
  }
67
67
 
68
+ # Per-field cap (characters) on the text Tier 1 runs heavy regex / encoding
69
+ # detection over. ReDoS guard — content past the cap is not analysed and
70
+ # metadata.analysis_truncated is set. The full content is still returned.
71
+ DEFAULT_MAX_FIELD_ANALYSIS_LENGTH = 50000
72
+
68
73
  DEFAULT_TIER2_CONFIG = Tier2Config(
69
74
  high_risk_threshold=0.8,
70
75
  medium_risk_threshold=0.5,