Implementing Guardrails for PII Masking and Prompt Injection in OmniRoute

OmniRoute provides a plug-in guardrail framework in src/lib/guardrails/ that intercepts LLM requests pre-call and post-call, enabling PII redaction via piiMasker.ts and prompt-injection blocking via promptInjection.ts with configurable priority levels and per-request opt-out headers.

OmniRoute is an open-source routing layer for AI models that ships with a built-in guardrail system to secure LLM interactions. The framework allows developers to implement PII masking and prompt injection detection by extending base classes and registering custom implementations. This guide walks through the architecture and implementation patterns found in the diegosouzapw/OmniRoute repository.

How the OmniRoute Guardrail Architecture Works

The guardrail system follows a registry pattern where each guardrail implements a common contract defined in src/lib/guardrails/base.ts. The GuardrailRegistry class in src/lib/guardrails/registry.ts orchestrates execution, running guardrails in ascending priority order—lower numbers execute first—for both incoming requests and outgoing responses.

Core Components

• BaseGuardrail (src/lib/guardrails/base.ts) – Abstract class defining the contract (preCall, postCall, priority, name) and providing utilities such as context handling and logging.
• Guardrail Registry (src/lib/guardrails/registry.ts) – Holds all active guardrails, runs them in priority order, and respects the x-omniroute-disabled-guardrails request header.
• Built-in Guardrails (src/lib/guardrails/promptInjection.ts, src/lib/guardrails/piiMasker.ts) – Ready-made implementations for prompt-injection blocking and PII redaction.
• Middleware Glue (src/middleware/promptInjectionGuard.ts) – Calls the registry for each incoming request so guardrails apply uniformly across all API routes.

Execution Flow

1. Request enters the API route → the route’s middleware invokes guardrailRegistry.runPreCalls(payload, ctx).
2. Each guardrail’s preCall can mutate the payload (e.g., redact PII) or reject it (e.g., block a suspicious prompt).
3. The modified request proceeds through the normal OmniRoute pipeline (translator → executor → provider).
4. After the upstream response, guardrailRegistry.runPostCalls(response, ctx) executes so guardrails may scrub the response or add diagnostics.

Built-in Guardrails for Security

OmniRoute ships with production-ready guardrails for two critical security concerns: data privacy and prompt safety.

PII Masking Guardrail

The built-in PII masking implementation lives in src/lib/guardrails/piiMasker.ts. It operates in two phases:

• Pre-call – Redacts personally identifiable fields in the request body when the header x-omniroute-pii-redact is present or when the environment variable PII_REDACTION_ENABLED is truthy.
• Post-call – Scrubs the upstream response of any values matching the same redaction rules, ensuring leaked data never reaches the client.

By design, this guardrail defaults to off and must be explicitly enabled per-request or globally.

Prompt Injection Guardrail

Located in src/lib/guardrails/promptInjection.ts, this guardrail employs configurable heuristics:

• Mode – Supports "block" (reject) or "pass-through" (allow but annotate) behaviors.
• Threshold – Numeric score computed by evaluatePromptInjection; exceeding the threshold triggers a 403 rejection.
• Disable mechanism – Send the header x-omniroute-disabled-guardrails: prompt-injection to bypass.

With a default priority of 5, it runs early in the chain, before PII masking or vision-bridge processing.

Implementing Custom Guardrails

To add a custom guardrail (e.g., a stricter PII filter or proprietary injection heuristic), extend the base class and register your implementation.

Step 1: Extend BaseGuardrail

Create a TypeScript file in src/lib/guardrails/ that implements the abstract 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 guardrail (priority 5)

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

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

Step 2: Register in the Registry

Export the class from src/lib/guardrails/index.ts and instantiate it in src/lib/guardrails/registry.ts:

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

guardrailRegistry.register(new CustomInjectionFilter());

Step 3: Configure Environment Variables

Follow the isEnabled pattern from src/lib/guardrails/promptInjection.ts to make your guardrail opt-in via environment variables (e.g., CUSTOM_INJECTION_ENABLED).

Wiring Guardrails into the Request Pipeline

The middleware in src/middleware/promptInjectionGuard.ts demonstrates the integration pattern. It creates a GuardrailContext from the request, runs guardrailRegistry.runPreCalls(), and handles rejections with HTTP 403:

import { guardrailRegistry } from "@/lib/guardrails/registry";

export async function middleware(req, res, next) {
  const ctx = createGuardrailContext(req);
  const guardResult = await guardrailRegistry.runPreCalls(req.body, ctx);
  
  if (guardResult.action === "reject") {
    return res.status(403).json({ error: guardResult.reason });
  }
  
  const response = await next();
  await guardrailRegistry.runPostCalls(response.body, ctx);
  return response;
}

This middleware attaches to every API route through Next.js route files (e.g., src/app/api/v1/chat/completions/route.ts).

Practical Implementation Examples

Enabling PII Masking for a Single Request

Send the x-omniroute-pii-redact header to activate redaction:

import fetch from "node-fetch";

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." }],
  }),
});

The payload is automatically redacted (***REDACTED***) before forwarding, and the response undergoes the same scrubbing.

Creating a URL-Based Injection Filter

Block prompts containing URLs by implementing a custom guardrail with priority 3 (running before the default injection guard at 5):

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

export class CustomInjectionFilter extends BaseGuardrail {
  readonly name = "custom-injection";
  readonly priority = 3;

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

Disabling Guardrails Selectively

Use the x-omniroute-disabled-guardrails header to skip specific checks for debugging or trusted workloads:

await fetch("/v1/chat/completions", {
  method: "POST",
  headers: {
    "Content-Type": "application/json",
    "x-omniroute-disabled-guardrails": "prompt-injection,pii-masker",
  },
  body: JSON.stringify({ /* ... */ }),
});

Summary

• OmniRoute's guardrail framework lives in src/lib/guardrails/ and uses a priority-based registry pattern.
• BaseGuardrail provides the contract: implement preCall for request modification/rejection and postCall for response scrubbing.
• Built-in implementations in piiMasker.ts and promptInjection.ts provide production-ready PII masking and prompt injection detection.
• Register custom guardrails in src/lib/guardrails/registry.ts after importing them through src/lib/guardrails/index.ts.
• Control guardrails via headers (x-omniroute-pii-redact, x-omniroute-disabled-guardrails) or environment variables.

Frequently Asked Questions

How do I change the execution order of guardrails?

Set the priority property in your BaseGuardrail subclass. Lower numbers execute first. The built-in prompt injection guardrail uses priority 5, while the PII masker runs later. Register a custom guardrail with priority 3 to intercept requests before the default injection checks.

Can guardrails modify the request payload or only reject it?

Guardrails can both mutate and reject. In preCall, modify the payload object directly to redact content, or return { action: "reject", reason: "..." } to block the request. In postCall, modify the response before it returns to the client.

Is PII masking enabled by default in OmniRoute?

No. According to the hard rule implemented in src/lib/guardrails/piiMasker.ts, PII redaction defaults to off. You must explicitly enable it via the x-omniroute-pii-redact request header or set the PII_REDACTION_ENABLED environment variable to truthy.

How do I test a custom guardrail without deploying?

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

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