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+wrapModelsetup ai@^5and 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,
});| Option | Purpose |
|---|---|
agent | Selects identities overrides in your policy document. Defaults to "default". |
workflowId | Vault 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
});| Mode | Behavior |
|---|---|
abort | Hold-back scan; throw PolicyViolationError; fail-closed |
buffer | Buffer full response; check at end; throw on block |
redact | Replace 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.
Related
- Next.js integration - runnable Demo 1 app
- wrapModel reference
- streamBlockBehavior reference