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

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 and registering them in 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, 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 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 – Detects adversarial prompts using configurable heuristics.
• src/lib/guardrails/piiMasker.ts – Redacts personally identifiable information from requests and responses.
• src/lib/guardrails/visionBridge.ts – Handles vision-specific input validation.

Middleware Integration – The 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.

// 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 to make the class available for import:

// 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. The file already registers the built-in implementations, so add your instantiation after those:

// 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. 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) 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:

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) 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 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 to define custom preCall and postCall logic for PII masking or prompt injection detection.
• Register implementations in 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, 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}