# How to Implement Custom Guardrails for PII Masking and Prompt Injection Detection in OmniRoute

> Learn to implement custom guardrails in OmniRoute for PII masking and prompt injection detection. Extend BaseGuardrail and register your solutions easily. Secure your AI applications now.

- Repository: [Diego Rodrigues de Sa e Souza/OmniRoute](https://github.com/diegosouzapw/OmniRoute)
- Tags: how-to-guide
- Published: 2026-07-19

---

**OmniRoute provides a plugin-style guardrail framework that intercepts requests via `preCall` and `postCall` hooks, allowing you to create custom PII maskers and prompt injection detectors by extending the `BaseGuardrail` class in [`src/lib/guardrails/base.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/base.ts) and registering them in [`src/lib/guardrails/registry.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/registry.ts).**

OmniRoute is an open-source routing layer for AI models that ships with an extensible guardrail system to enforce security policies. Whether you need to redact sensitive personally identifiable information (PII) or block adversarial prompt injection attempts, you can implement custom guardrails by leveraging the framework located under `src/lib/guardrails/`. This guide walks through the architecture and implementation steps based on the OmniRoute source code.

## Understanding the Guardrail Architecture

The guardrail framework consists of four core components that process requests before they reach upstream providers and responses before they return to clients.

**`BaseGuardrail`** – Located in [`src/lib/guardrails/base.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/base.ts), this abstract class defines the contract that every guardrail must implement. It requires the properties `name` (string identifier), `priority` (numeric execution order), and the methods `preCall` and `postCall`. The class also provides utilities for context handling and logging.

**Guardrail Registry** – The [`src/lib/guardrails/registry.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/registry.ts) file maintains a global registry that holds all active guardrail instances. The registry executes guardrails in ascending priority order and respects the `x-omniroute-disabled-guardrails` request header to skip specific implementations on a per-request basis.

**Built-in Implementations** – OmniRoute ships with three production-ready guardrails in the same directory:

- [`src/lib/guardrails/promptInjection.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/promptInjection.ts) – Detects adversarial prompts using configurable heuristics.
- [`src/lib/guardrails/piiMasker.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/piiMasker.ts) – Redacts personally identifiable information from requests and responses.
- [`src/lib/guardrails/visionBridge.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/visionBridge.ts) – Handles vision-specific input validation.

**Middleware Integration** – The [`src/middleware/promptInjectionGuard.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/middleware/promptInjectionGuard.ts) file wires the registry into the request pipeline. This middleware invokes the registry for every API route, ensuring uniform policy enforcement across the application.

## Implementing a Custom Guardrail

To add a custom guardrail—such as a stricter PII filter or a new prompt-injection heuristic—follow these five steps:

### 1. Extend the BaseGuardrail Class

Create a new TypeScript file in `src/lib/guardrails/` that extends `BaseGuardrail`. Define a unique `name`, assign a `priority` (lower numbers execute earlier), and implement the `preCall` and `postCall` methods.

```typescript
// src/lib/guardrails/customInjectionFilter.ts
import { BaseGuardrail, GuardrailContext, GuardrailResult } from "./base";

export class CustomInjectionFilter extends BaseGuardrail {
  readonly name = "custom-injection";
  readonly priority = 3;  // Runs before the default injection guard (priority 5)

  async preCall(payload: unknown, ctx: GuardrailContext): Promise<GuardrailResult> {
    const content = (payload as any).messages?.[0]?.content ?? "";
    // Reject prompts containing URLs
    if (/(https?:\/\/|www\.)[\w./-]+/i.test(content)) {
      return { action: "reject", reason: "URLs are not allowed in prompts" };
    }
    return { action: "allow" };
  }

  async postCall(response: unknown, ctx: GuardrailContext): Promise<GuardrailResult> {
    return { action: "allow" };
  }
}

```

### 2. Export from the Index File

Add the export to [`src/lib/guardrails/index.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/index.ts) to make the class available for import:

```typescript
// src/lib/guardrails/index.ts
export * from "./customInjectionFilter";

```

### 3. Register in the Guardrail Registry

Instantiate and register your guardrail in [`src/lib/guardrails/registry.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/registry.ts). The file already registers the built-in implementations, so add your instantiation after those:

```typescript
// src/lib/guardrails/registry.ts
import { CustomInjectionFilter } from "./customInjectionFilter";

guardrailRegistry.register(new CustomInjectionFilter());

```

### 4. Configure via Environment Variables (Optional)

To make your guardrail opt-in by default, follow the `isEnabled` pattern used in [`promptInjection.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/promptInjection.ts). Check for an environment variable such as `CUSTOM_INJECTION_ENABLED` within your class constructor to control activation without code changes.

### 5. Write Unit Tests

Create a test file under `tests/unit/guardrails/` using the existing suite as a template. Import your guardrail, construct a mock `GuardrailContext`, and assert that `preCall` returns the expected `action` and `reason` values for both allowed and rejected inputs.

## Working with Built-in PII Masking

The **PII masking guardrail** ([`src/lib/guardrails/piiMasker.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/piiMasker.ts)) operates in two phases to ensure data privacy. By default, this guardrail is **disabled** and must be explicitly activated.

**Pre-call Redaction** – When the request header `x-omniroute-pii-redact` is set to `"true"` or the environment variable `PII_REDACTION_ENABLED` is truthy, the guardrail scans the request payload and replaces personally identifiable fields with `***REDACTED***`.

**Post-call Scrubbing** – After receiving the upstream response, the guardrail applies the same redaction rules to prevent leaked PII from reaching the client.

Enable PII masking for a single request by including the header:

```typescript
await fetch("https://omniroute.example.com/v1/chat/completions", {
  method: "POST",
  headers: {
    "Content-Type": "application/json",
    "x-omniroute-pii-redact": "true",
  },
  body: JSON.stringify({
    model: "openai/gpt-4o",
    messages: [{ role: "user", content: "My SSN is 123-45-6789." }],
  }),
});

```

## Configuring Prompt Injection Detection

The **prompt injection guardrail** ([`src/lib/guardrails/promptInjection.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/promptInjection.ts)) uses a configurable heuristic to detect adversarial inputs. It runs at **priority 5**, ensuring it executes early in the pipeline before PII masking or vision processing.

**Operating Modes** – Configure the guardrail to either `"block"` (reject the request with a `403` error) or `"pass-through"` (allow but annotate the request).

**Threshold Scoring** – The `evaluatePromptInjection` function computes a numeric risk score. If the score exceeds the configured threshold, the guardrail rejects the request.

**Disabling Per Request** – Clients can bypass this check by sending the header `x-omniroute-disabled-guardrails: prompt-injection`.

## Wiring Guardrails into the Request Pipeline

The middleware in [`src/middleware/promptInjectionGuard.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/middleware/promptInjectionGuard.ts) integrates the registry into every API route. When a request enters the pipeline, the middleware:

1. Creates a `GuardrailContext` from the incoming request.
2. Executes `guardrailRegistry.runPreCalls(payload, ctx)` to process all enabled guardrails.
3. If any guardrail returns `{ action: "reject" }`, the middleware immediately returns a `403` response with the provided reason.
4. Upon receiving the upstream response, it calls `guardrailRegistry.runPostCalls(response.body, ctx)` to apply post-processing rules.

This uniform integration ensures that custom PII maskers and injection detectors apply consistently across all chat completion and translation endpoints.

## Summary

- **Extend `BaseGuardrail`** in [`src/lib/guardrails/base.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/base.ts) to define custom `preCall` and `postCall` logic for PII masking or prompt injection detection.
- **Register implementations** in [`src/lib/guardrails/registry.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/registry.ts) and control execution order via the `priority` property (lower values run first).
- **Enable built-in PII masking** using the `x-omniroute-pii-redact` header or `PII_REDACTION_ENABLED` environment variable; the feature defaults to off for privacy safety.
- **Configure prompt injection** thresholds and modes in [`src/lib/guardrails/promptInjection.ts`](https://github.com/diegosouzapw/OmniRoute/blob/main/src/lib/guardrails/promptInjection.ts), which runs at priority 5.
- **Disable specific guardrails** per request using the `x-omniroute-disabled-guardrails` header with comma-separated guardrail names.
- **Test custom implementations** under `tests/unit/guardrails/` to validate rejection logic and payload mutations.

## Frequently Asked Questions

### How do I disable a specific guardrail for a single request?

Send the `x-omniroute-disabled-guardrails` header with a comma-separated list of guardrail names. For example, `"prompt-injection,pii-masker"` skips both the injection detector and PII masker for that call only. The registry checks this header before executing any guardrail's `preCall` or `postCall` methods.

### What is the difference between `preCall` and `postCall` in OmniRoute guardrails?

The `preCall` method intercepts the request payload before it reaches the upstream AI provider, allowing you to mutate data (such as redacting PII) or reject the request entirely. The `postCall` method processes the response after it returns from the provider, enabling you to scrub sensitive data from model outputs or append diagnostic metadata before sending the result to the client.

### How does guardrail priority work in the OmniRoute registry?

The registry sorts guardrails by their `priority` property in ascending order (lower numbers execute first). The built-in prompt injection guardrail uses priority `5`, while the PII masker typically runs later. When implementing custom guardrails, assign priorities below `5` to run before injection checks, or above `5` to run after them.

### Can I enable PII masking globally for all requests?

Yes, set the environment variable `PII_REDACTION_ENABLED` to a truthy value. This activates the PII masker for every request without requiring the `x-omniroute-pii-redact` header. However, the guardrail follows a hard rule of defaulting to off, so you must explicitly enable this variable; there is no "always on" default in the source code to prevent accidental data processing restrictions.