# Protect PII in the AI SDK

URL: https://tailrace.dev/docs/guides/protect-pii-in-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 [#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 [#prerequisites]

* [Quickstart](/docs/get-started/quickstart) completed, or equivalent `createTailrace` + `wrapModel` setup
* `ai@^5` and a provider package (`@ai-sdk/openai`, etc.)

## Step 1: Choose an API shape [#step-1-choose-an-api-shape]

Two equivalent forms ship in v0.1:

**Fluent** (recommended for app code):

```ts
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):

```ts
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 [#step-2-set-agent-and-workflow-id]

```ts
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 [#step-3-wrap-tools]

```ts
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):

```text
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 [#step-4-restore-at-egress]

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

```ts
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) [#step-5-configure-streaming-optional]

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

```ts
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 [#verify-it-works]

### Block [#block]

```bash
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 [#tokenize--restore]

```bash
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 [#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 [#related]

* [Next.js integration](/docs/integrations/nextjs) - runnable Demo 1 app
* [wrapModel reference](/docs/reference/ai-sdk/wrap-model)
* [streamBlockBehavior reference](/docs/reference/ai-sdk/options#streamblockbehavior)