| 1 | # Infection Check |
| 2 | |
| 3 | Safety middleware that analyzes agent output for prompt injection and suspicious external influence before allowing tool execution. |
| 4 | |
| 5 | ## What It Does |
| 6 | |
| 7 | This plugin collects streamed reasoning and response text, analyzes that content with a configurable audit model, and blocks tool execution until the safety check either passes, requests clarification, or terminates the agent. |
| 8 | |
| 9 | ## How It Works |
| 10 | |
| 11 | 1. **Collection** — During streaming, the plugin collects the agent's reasoning and response text via `reasoning_stream_chunk` and `response_stream_chunk` extensions. |
| 12 | 2. **Analysis** — A security audit model analyzes the collected text against the configurable prompt. |
| 13 | 3. **Gate** — Before any tool executes, `gate()` blocks until the check completes. If the background task already finished (thoughts mode), the result is reused instantly. Otherwise the task is rebuilt with full tool context (name + arguments) for maximum accuracy. |
| 14 | 4. **Action** — Based on the verdict: |
| 15 | - `<ok/>` — Tool execution proceeds silently. |
| 16 | - `<terminate/>` — Agent is stopped immediately with a warning log and desktop notification. |
| 17 | - `<clarify>` — The agent is asked to explain; a back-and-forth loop runs until resolved. |
| 18 | |
| 19 | ## Modes |
| 20 | |
| 21 | | Mode | What is Analyzed | When Analysis Starts | Latency | |
| 22 | |---|---|---|---| |
| 23 | | **thoughts** (default) | Reasoning + thoughts | When `heading` or `tool_name` appear in response stream | Low — runs in parallel while tool args still stream | |
| 24 | | **complete** | Reasoning + full response | After the entire response stream ends | Higher — blocking wait before tool execution | |
| 25 | |
| 26 | ## Clarification Loop |
| 27 | |
| 28 | When the check model returns `<clarify>`: |
| 29 | |
| 30 | 1. A warning log item is displayed with the safety concern. |
| 31 | 2. The conversation history is cloned and the clarification questions appended. |
| 32 | 3. The **main** model responds (streamed into the UI log). |
| 33 | 4. The response is fed back to the check model for re-evaluation. |
| 34 | 5. Repeats up to **Max Clarifications** times; exceeding the limit triggers termination. |
| 35 | |
| 36 | ## Termination Behavior |
| 37 | |
| 38 | When the check results in `<terminate/>` (directly or after exhausting clarifications): |
| 39 | |
| 40 | 1. A warning is logged with the full chain-of-thought. |
| 41 | 2. The last AI message in history is replaced with `[BLOCKED]`. |
| 42 | 3. A desktop notification is sent. |
| 43 | 4. Queued messages are scheduled to resume after the current task stops (since the normal `process_chain_end` extension does not fire after `HandledException`). |
| 44 | 5. `HandledException` is raised to stop the agent. |
| 45 | |
| 46 | ## Configuration |
| 47 | |
| 48 | | Setting | Default | Description | |
| 49 | |---|---|---| |
| 50 | | Mode | `thoughts` | `thoughts` or `complete` | |
| 51 | | Model | `utility` | `utility` (faster/cheaper) or `main` (more capable) | |
| 52 | | Max Clarifications | `3` | Clarification rounds before auto-terminate | |
| 53 | | History Size | `10` | Recent messages included as context | |
| 54 | | Prompt | *(built-in)* | Fully customizable security audit system prompt | |
| 55 | |
| 56 | ## Key Files |
| 57 | |
| 58 | - **Checker logic** |
| 59 | - `helpers/checker.py` implements stream collection, background analysis, gating, clarification, and termination. |
| 60 | - **Extensions** |
| 61 | - `extensions/python/reasoning_stream_chunk/_50_infection_collect.py` |
| 62 | - `extensions/python/response_stream_chunk/_50_infection_collect.py` |
| 63 | - `extensions/python/response_stream/_50_infection_analyze.py` |
| 64 | - `extensions/python/response_stream_end/_50_infection_analyze.py` |
| 65 | - `extensions/python/tool_execute_before/_50_infection_check.py` |
| 66 | |
| 67 | ## Extension Points Used |
| 68 | |
| 69 | | Extension Point | File | Purpose | |
| 70 | |---|---|---| |
| 71 | | `reasoning_stream_chunk` | `_50_infection_collect.py` | Accumulate reasoning text | |
| 72 | | `response_stream_chunk` | `_50_infection_collect.py` | Accumulate response text | |
| 73 | | `response_stream` | `_50_infection_analyze.py` | Detect thoughts complete → start background analysis | |
| 74 | | `response_stream_end` | `_50_infection_analyze.py` | Start analysis (complete mode / fallback) | |
| 75 | | `tool_execute_before` | `_50_infection_check.py` | Await check result → gate tool execution | |
| 76 | |
| 77 | ## Configuration Scope |
| 78 | |
| 79 | - **Settings section**: `agent` |
| 80 | - **Per-project config**: `true` |
| 81 | - **Per-agent config**: `true` |
| 82 | |
| 83 | ## Plugin Metadata |
| 84 | |
| 85 | - **Name**: `_infection_check` |
| 86 | - **Title**: `Infection Check` |
| 87 | - **Description**: Safety check for prompt injection from external sources. |