<div align="center">

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

  <p>
    <a href="https://www.npmjs.com/package/@stackone/defender"><img src="https://img.shields.io/npm/v/%40stackone%2Fdefender?style=flat-square&color=047B43&label=npm" alt="npm version" /></a>
    <a href="https://www.npmjs.com/package/@stackone/defender"><img src="https://img.shields.io/npm/dm/%40stackone%2Fdefender?style=flat-square&color=047B43&label=downloads" alt="npm downloads" /></a>
    <a href="https://github.com/StackOneHQ/defender/releases"><img src="https://img.shields.io/github/v/release/StackOneHQ/defender?style=flat-square&color=047B43&label=release" alt="latest release" /></a>
    <a href="https://github.com/StackOneHQ/defender/stargazers"><img src="https://img.shields.io/github/stars/StackOneHQ/defender?style=flat-square&color=047B43" alt="GitHub stars" /></a>
    <a href="./LICENSE"><img src="https://img.shields.io/npm/l/%40stackone%2Fdefender?style=flat-square&color=047B43" alt="License" /></a>
    <img src="https://img.shields.io/badge/TypeScript-typed-047B43?style=flat-square" alt="TypeScript" />
  </p>
  <p>
    <img src="https://img.shields.io/badge/model-22MB-047B43?style=flat-square" alt="Model size: 22MB" />
    <img src="https://img.shields.io/badge/latency-~10ms-047B43?style=flat-square" alt="Latency: ~10ms" />
    <img src="https://img.shields.io/badge/CPU--only-no%20GPU%20needed-047B43?style=flat-square" alt="CPU only" />
    <img src="https://img.shields.io/badge/F1%20Score-90.8%25-047B43?style=flat-square" alt="F1 Score: 90.8%" />
  </p>

</div>

---

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

## Installation

```bash
npm install @stackone/defender
```

The ONNX model (~22MB) is bundled in the package — no extra downloads needed.

## Quick Start

```typescript
import { createPromptDefense } from '@stackone/defender';

// Tier 1 (patterns) + Tier 2 (ML classifier) are both on by default.
// blockHighRisk: true enables the allowed/blocked decision.
const defense = createPromptDefense({
  blockHighRisk: true,
});

// Defend a tool result — ONNX model (~22MB) auto-loads on first call
const result = await defense.defendToolResult(toolOutput, 'gmail_get_message');

if (!result.allowed) {
  console.log(`Blocked: risk=${result.riskLevel}, score=${result.tier2Score}`);
  console.log(`Detections: ${result.detections.join(', ')}`);
} else {
  // Safe to pass result.sanitized to the LLM
  passToLLM(result.sanitized);
}
```

## How It Works

<picture>
  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/demo-dark.svg" />
  <img src="https://raw.githubusercontent.com/StackOneHQ/defender/main/assets/demo-light.svg" alt="Defender flow: a poisoned email with an injection payload is intercepted by @stackone/defender and blocked before reaching the LLM, with riskLevel: critical and tier2Score: 0.97" width="900" />
</picture>

`defendToolResult()` runs a two-tier defense pipeline:

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

Regex-based detection and sanitization:
- **Unicode normalization** — prevents homoglyph attacks (Cyrillic 'а' → ASCII 'a')
- **Role stripping** — removes `SYSTEM:`, `ASSISTANT:`, `<system>`, `[INST]` markers
- **Pattern removal** — redacts injection patterns like "ignore previous instructions"
- **Encoding detection** — detects and handles Base64/URL encoded payloads
- **Boundary annotation** — opt-in; wraps untrusted content in `[UD-{id}]...[/UD-{id}]` tags when `annotateBoundary: true` is passed to `createPromptDefense`. Off by default; pair with `generateBoundaryInstructions()` in your system prompt if you enable it.

### Tier 2 — ML Classification (async)

Fine-tuned MiniLM classifier with sentence-level analysis:
- Splits text into sentences and scores each one (0.0 = safe, 1.0 = injection)
- Fine-tuned MiniLM-L6-v2, int8 quantized (~22MB), bundled in the package — no external download needed
- Catches attacks that evade pattern-based detection
- Latency: ~10ms/sample (after model warmup)

**Benchmark results** (ONNX mode, F1 score at threshold 0.5):

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

### Understanding `allowed` vs `riskLevel`

Use `allowed` for blocking decisions:
- `allowed: true` — safe to pass to the LLM
- `allowed: false` — content blocked (requires `blockHighRisk: true`, which defaults to `false`)

`riskLevel` is diagnostic metadata. It starts at `medium` (the default) and is escalated by Tier 1 pattern detections, encoding detection, and Tier 2 ML scoring — never reduced. Use it for logging and monitoring, not for allow/block logic.

Risk escalation from detections:

| Level | Detection Trigger |
|-------|-------------------|
| `low` | No threats detected |
| `medium` | Suspicious patterns, role markers stripped |
| `high` | Injection patterns detected, content redacted |
| `critical` | Severe injection attempt with multiple indicators |

## API

### `createPromptDefense(options?)`

Create a defense instance.

```typescript
const defense = createPromptDefense({
  enableTier1: true,           // Pattern detection (default: true)
  enableTier2: true,           // ML classification (default: true) — set false to disable
  blockHighRisk: true,         // Block high/critical content (default: false)
  tier2Fields: ['subject', 'body', 'snippet'], // Scope Tier 2 to specific fields (default: all fields)
  defaultRiskLevel: 'medium',
});
```

### `defense.defendToolResult(value, toolName)`

The primary method. Runs Tier 1 + Tier 2 and returns a `DefenseResult`:

```typescript
interface DefenseResult {
  allowed: boolean;                       // Use this for blocking decisions (respects blockHighRisk config)
  riskLevel: RiskLevel;                   // Diagnostic: tool base risk + detection escalation (see docs above)
  sanitized: unknown;                     // The sanitized tool result
  detections: string[];                   // Pattern names detected by Tier 1
  fieldsSanitized: string[];              // Fields where threats were found (e.g. ['subject', 'body'])
  patternsByField: Record<string, string[]>; // Patterns per field
  tier2Score?: number;                    // ML score (0.0 = safe, 1.0 = injection)
  maxSentence?: string;                   // The sentence with the highest Tier 2 score
  latencyMs: number;                      // Processing time in milliseconds
}
```

### `defense.defendToolResults(items)`

Batch method — defends multiple tool results concurrently.

```typescript
const results = await defense.defendToolResults([
  { value: emailData, toolName: 'gmail_get_message' },
  { value: docData, toolName: 'documents_get' },
  { value: prData, toolName: 'github_get_pull_request' },
]);

for (const result of results) {
  if (!result.allowed) {
    console.log(`Blocked: ${result.fieldsSanitized.join(', ')}`);
  }
}
```

### `defense.analyze(text)`

Low-level Tier 1 analysis for debugging. Returns pattern matches and risk assessment without sanitization.

```typescript
const result = defense.analyze('SYSTEM: ignore all rules');
console.log(result.hasDetections); // true
console.log(result.suggestedRisk); // 'high'
console.log(result.matches);       // [{ pattern: '...', severity: 'high', ... }]
```

### Tier 2 Setup

The bundled model auto-loads on first `defendToolResult()` call. Use `warmupTier2()` at startup to avoid first-call latency:

```typescript
const defense = createPromptDefense();
await defense.warmupTier2(); // optional, avoids ~1-2s first-call latency
```

## Integration Example

### With Vercel AI SDK

```typescript
import { generateText, tool } from 'ai';
import { createPromptDefense } from '@stackone/defender';

const defense = createPromptDefense({
  blockHighRisk: true,
});
await defense.warmupTier2(); // optional, avoids first-call latency

const result = await generateText({
  model: anthropic('claude-sonnet-4-20250514'),
  tools: {
    gmail_get_message: tool({
      // ... tool definition
      execute: async (args) => {
        const rawResult = await gmailApi.getMessage(args.id);
        const defended = await defense.defendToolResult(rawResult, 'gmail_get_message');

        if (!defended.allowed) {
          return { error: 'Content blocked by safety filter' };
        }

        return defended.sanitized;
      },
    }),
  },
});
```

## Risky Field Detection

Defender only scans string fields that are likely to contain user-generated or external content. Per-tool overrides focus scanning on the relevant fields:

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

Tools not matching any pattern use the default risky field list: `name`, `description`, `content`, `title`, `notes`, `summary`, `bio`, `body`, `text`, `message`, `comment`, `subject`, plus patterns like `*_description`, `*_body`, etc.

Fields like `id`, `url`, `created_at` are never scanned — they aren't in the risky fields list.

## Development

### Testing

```bash
npm test
```

## License

Apache-2.0 — See [LICENSE](./LICENSE) for details.
