TailraceTailrace
Guides

Protect PII in the AI SDK

Wrap models and tools, configure streaming block behavior, and restore tokenized values at egress.

Use @tailrace/ai-sdk to enforce data policy on Vercel AI SDK model calls and tool executions - in-process, with no proxy.

When to use this

  • You send user or ticket data to OpenAI, Anthropic, or other providers through the AI SDK.
  • Agent tools read CRM records, databases, or files that may contain PII or secrets.
  • You need the same email to become the same token across a multi-step agent run.

Prerequisites

  • Quickstart completed, or equivalent createTailrace + wrapModel setup
  • ai@^5 and a provider package (@ai-sdk/openai, etc.)

Step 1: Choose an API shape

Two equivalent forms ship in v0.1:

Fluent (recommended for app code):

import { createTailrace } from "@tailrace/core";
import { withAiSdk } from "@tailrace/ai-sdk";

const tailrace = withAiSdk(createTailrace());
const model = tailrace.model(openai("gpt-4o"));
const tools = tailrace.tools({ search: searchTool });

Standalone (explicit imports):

import { wrapModel, wrapTools } from "@tailrace/ai-sdk";

const model = wrapModel(tailrace, openai("gpt-4o"));
const tools = wrapTools(tailrace, { search: searchTool });

Step 2: Set agent and workflow ID

const model = tailrace.model(openai("gpt-4o"), {
  agent: "support-bot",
  workflowId: sessionId,
});
OptionPurpose
agentSelects identities overrides in your policy document. Defaults to "default".
workflowIdVault scope for tokens. Same ID + same value → same token every time.

Pass the same workflowId to model, tools, and restore within one conversation.

Step 3: Wrap tools

import { tool } from "ai";
import { z } from "zod";

const tools = tailrace.tools(
  {
    fetchOrder: tool({
      description: "Fetch order by id",
      inputSchema: z.object({ orderId: z.string() }),
      execute: async ({ orderId }) => db.orders.get(orderId),
    }),
  },
  { agent: "support-bot", workflowId: sessionId },
);

Tailrace checks:

  • Tool arguments at { kind: "tool", name, direction: "out" }
  • Tool results at { kind: "tool", name, direction: "in" }

On block, the model receives a readable error (not a stack trace):

Blocked by data policy: api_key may not be sent to tool:fetchOrder:out (rule: entities.api_key)

Tools without execute are not wrapped.

Step 4: Restore at egress

The model wrapper tokenizes PII in prompts and completions. Your HTTP handler (or trusted UI sink) restores tokens:

const result = await generateText({ model, prompt });

const { output } = await tailrace.restore(result.text, {
  boundary: { kind: "egress", sink: "ui" },
  identity: { agent: "support-bot" },
  workflowId: sessionId,
});

return Response.json({ text: output });

Pick a sink name that matches your policy (egress:ui, egress:*, etc.).

Step 5: Configure streaming (optional)

For streamText, choose how policy block on model output is handled:

const model = tailrace.model(openai("gpt-4o"), {
  streamBlockBehavior: "abort", // default - recommended
});
ModeBehavior
abortHold-back scan; throw PolicyViolationError; fail-closed
bufferBuffer full response; check at end; throw on block
redactReplace blocked spans with [ENTITY]; stream continues - not fail-closed

Input blocking always aborts before the provider runs. redact applies only to streaming output.

Non-streaming generateText always throws on output block.

Verify it works

Block

curl -s -X POST http://localhost:3000/api/chat \
  -H 'content-type: application/json' \
  -d '{"prompt":"Use sk_test_51FakeKeyForTailraceTests000FAKE"}' | jq .

Expect 422 and "entity": "api_key".

Tokenize + restore

curl -s -X POST http://localhost:3000/api/chat \
  -H 'content-type: application/json' \
  -H 'x-workflow-id: test-session' \
  -d '{"prompt":"Email customer@example.com"}' | jq .

Expect 200, modelSaw containing <EMAIL_…>, and text containing the real email.

Troubleshooting

Tokens change between requests

You are not passing the same workflowId to model and restore. Thread the session ID through both.

restore throws INVARIANT

You called restore at a model or tool boundary. Use { kind: "egress", sink: "…" } only.

Stream aborts on benign output

A secret-class pattern matched in model output. Default policy blocks echoed secrets. Adjust policy only with care - secrets cannot be overridden to allow without dangerouslyAllowSecrets: true.

On this page