From 6da6c617e69d873ad7a06850a19ddda788f1a19a Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 16:51:09 +0000
Subject: [PATCH 1/7] Add optional format support to the re-identify tool

Adds an optional `format` input to the re-identify tool so the calling
agent can specify, per entity type, how tokens are rendered on the way
out: fully restored (plaintext), partially masked, or fully redacted.
This mirrors the Skyflow Detect API `format` object and maps to the
skyflow-node SDK's `ReidentifyTextOptions`. Entity types not listed
default to full plaintext restoration, preserving existing behavior.

- handler: build `ReidentifyTextOptions` from the format and echo the
  applied format back in the output
- server: add `format` to the re-identify input/output Zod schemas and
  thread it through to the handler
- types: add `ReIdentifyFormat`, extend `ReIdentifyOutput`
- UI: summarize the applied format treatment in the re-identify app
- tests: 8 new cases (entity routing, echo-back, empty/invalid buckets,
  backward compatibility)
- docs: CLAUDE.md, README, wrapping guide, CHANGELOG

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 CHANGELOG.md                            |   2 +
 CLAUDE.md                               |   3 +-
 README.md                               |  18 ++++
 docs/wrapping-mcp-tools-with-skyflow.md |  35 ++++++++
 src/lib/tools/reIdentify.ts             |  48 ++++++++++-
 src/lib/tools/types.ts                  |  17 ++++
 src/server.ts                           |  37 ++++++++-
 tests/unit/tools/reIdentify.test.ts     | 105 ++++++++++++++++++++++++
 ui/re-identify/main.ts                  |  35 ++++++++
 ui/shared/styles.css                    |  30 +++++++
 ui/shared/types.ts                      |   8 ++
 11 files changed, 329 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a4b8f2..25b1a31 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ### Added
 
+- **Re-identify output format control** — The `re-identify` tool now accepts an optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings) that lets the caller specify, per entity type, how tokens are rendered on the way out — fully restored (plaintext), partially masked, or fully redacted. Entity types not listed default to full plaintext restoration, preserving existing behavior. Maps to the `skyflow-node` SDK's `ReidentifyTextOptions` per the Detect API spec. The applied format is echoed back in the response and summarized in the re-identify UI.
+
 - **MCP Apps UI for all three tools** — Each tool (`dehydrate`, `rehydrate`, `dehydrate_file`) now has an interactive vanilla TypeScript UI that renders inline in MCP Apps-capable hosts. Text-only hosts continue to receive JSON responses as before.
   - **Dehydrate UI**: Side-by-side before/after text panels with color-coded entity highlights, confidence scores, and an entity breakdown table. Shows anonymous mode banner when applicable.
   - **Rehydrate UI**: Token-to-original mapping display with color-matched highlights across before/after panels.
diff --git a/CLAUDE.md b/CLAUDE.md
index 7278309..82eb2e4 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -98,7 +98,8 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}"
 
 **re-identify tool** (`src/lib/tools/reIdentify.ts`)
 - Reverses de-identification by replacing tokens with original sensitive data
-- Returns `inputText` and `processedText`
+- Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed default to full plaintext restoration
+- Returns `inputText` and `processedText`, and echoes back the applied `format` when one was provided
 - Returns error with `anonymousModeRestricted: true` in anonymous mode
 
 **de-identify_file tool** (`src/lib/tools/deIdentifyFile.ts`)
diff --git a/README.md b/README.md
index 924fa57..beaf8e6 100644
--- a/README.md
+++ b/README.md
@@ -255,6 +255,24 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}"
   -d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"re-identify","arguments":{"inputString":"[REDACTED_TEXT_WITH_TOKENS]"}},"id":3}'
 ```
 
+By default every token is fully restored to its original value. To control how
+individual entity types are rendered, pass an optional `format` object with any
+of `redacted`, `masked`, or `plaintext` — each a list of entity types (the same
+lowercase names as the `de-identify` tool, e.g. `ssn`, `email_address`). Entity
+types not listed default to full plaintext restoration:
+
+```bash
+curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" \
+  -H "Content-Type: application/json" \
+  -H "Accept: application/json, text/event-stream" \
+  -H "Authorization: Bearer {your_bearer_token}" \
+  -d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"re-identify","arguments":{"inputString":"[REDACTED_TEXT_WITH_TOKENS]","format":{"masked":["ssn"],"redacted":["email_address"],"plaintext":["name"]}}},"id":3}'
+```
+
+In this example, SSNs come back partially masked, email addresses are fully
+redacted, names are restored in full, and any other detected entity type is
+restored to plaintext.
+
 ## Integration with Claude Desktop
 
 Claude Desktop is one concrete client for the [Connection Contract](#connection-contract). It uses `mcp-remote` as a bridge to the Streamable HTTP endpoint. Add the following to your `claude_desktop_config.json`:
diff --git a/docs/wrapping-mcp-tools-with-skyflow.md b/docs/wrapping-mcp-tools-with-skyflow.md
index c08b853..0e548ea 100644
--- a/docs/wrapping-mcp-tools-with-skyflow.md
+++ b/docs/wrapping-mcp-tools-with-skyflow.md
@@ -220,6 +220,29 @@ lowercase string → enum map in `src/lib/mappings/entityMaps.ts` (`ENTITY_MAP`)
 reference list. Common values include `email_address`, `ssn`, `credit_card`, `name`,
 `phone_number`, `ip_address`, `location`, and `bank_account`.
 
+### Optional: control re-identify output format
+
+Re-identify fully restores every token by default. To render specific entity types differently on
+the way out — reveal some, partially mask others, fully redact the rest — pass a
+`ReidentifyTextOptions` as the second argument. Entity types you don't list default to full plaintext
+restoration:
+
+```ts
+import { ReidentifyTextOptions, DetectEntities } from "skyflow-node";
+
+const options = new ReidentifyTextOptions();
+options.setPlainTextEntities([DetectEntities.NAME]);      // restore in full
+options.setMaskedEntities([DetectEntities.SSN]);          // partially masked (e.g. ***-**-6789)
+options.setRedactedEntities([DetectEntities.EMAIL_ADDRESS]); // fully hidden
+
+const res = await skyflow
+  .detect()
+  .reidentifyText(new ReidentifyTextRequest(text), options);
+```
+
+This is the same capability the `re-identify` MCP tool exposes via its optional `format` argument
+(`{ redacted, masked, plaintext }`); see `src/lib/tools/reIdentify.ts`.
+
 ## Approach B — Detect REST API (any language)
 
 The SDK just wraps Skyflow's **Detect REST API**, so any language can run the same round-trip over
@@ -279,6 +302,18 @@ curl -X POST "https://$CLUSTER_ID.vault.skyflowapis.com/v1/detect/reidentify/str
 { "text": "email john.doe@example.com about order 12345" }
 ```
 
+`text` and `vault_id` are required. Add an optional `format` object to control how each entity type
+is rendered — `{ "redacted": [...], "masked": [...], "plaintext": [...] }`, each a list of the same
+lowercase entity names. Entity types you don't list default to full plaintext restoration:
+
+```json
+{
+  "vault_id": "...",
+  "text": "email [EMAIL_ADDRESS_a1b2], ssn [SSN_c3d4]",
+  "format": { "redacted": ["email_address"], "masked": ["ssn"] }
+}
+```
+
 Note the re-identify response field is `text`, while de-identify returns `processed_text` — an
 asymmetry in the raw v1 wire format. The `skyflow-node` SDK (Approach A) smooths this over: it
 surfaces both as `processedText`, and exposes each detected entity with `entity` / `scores` /
diff --git a/src/lib/tools/reIdentify.ts b/src/lib/tools/reIdentify.ts
index 349d324..141f56f 100644
--- a/src/lib/tools/reIdentify.ts
+++ b/src/lib/tools/reIdentify.ts
@@ -1,15 +1,52 @@
-import { ReidentifyTextRequest, SkyflowError } from "skyflow-node";
+import { ReidentifyTextOptions, ReidentifyTextRequest, SkyflowError } from "skyflow-node";
 import type { Skyflow } from "skyflow-node";
-import type { ReIdentifyOutput, ReIdentifyErrorOutput, AnonymousModeError, ToolResult } from "./types.js";
+import { getEntityEnum } from "../mappings/entityMaps.js";
+import type { ReIdentifyFormat, ReIdentifyOutput, ReIdentifyErrorOutput, AnonymousModeError, ToolResult } from "./types.js";
+
+/**
+ * Build Skyflow re-identify options from the tool's `format` argument.
+ * Only buckets containing at least one entity type are set, so an empty or
+ * absent bucket falls back to the API default (plaintext re-identification).
+ * Returns `undefined` when no entity types are specified, so the SDK is called
+ * without options and every token is fully re-identified (the default behavior).
+ */
+function buildReidentifyOptions(
+  format: ReIdentifyFormat | undefined
+): ReidentifyTextOptions | undefined {
+  if (!format) return undefined;
+
+  const options = new ReidentifyTextOptions();
+  let hasAny = false;
+
+  if (format.redacted && format.redacted.length > 0) {
+    options.setRedactedEntities(format.redacted.map(getEntityEnum));
+    hasAny = true;
+  }
+  if (format.masked && format.masked.length > 0) {
+    options.setMaskedEntities(format.masked.map(getEntityEnum));
+    hasAny = true;
+  }
+  if (format.plaintext && format.plaintext.length > 0) {
+    options.setPlainTextEntities(format.plaintext.map(getEntityEnum));
+    hasAny = true;
+  }
+
+  return hasAny ? options : undefined;
+}
 
 /**
  * Handle the re-identify tool logic.
  * Restores original sensitive data from de-identified placeholders.
+ *
+ * An optional `format` controls how each entity type is rendered on the way out
+ * (redacted / masked / plaintext), per the Skyflow Detect API spec. When omitted,
+ * every token is fully re-identified to its original plaintext value.
  */
 export async function handleReIdentify(
   inputString: string,
   skyflow: Skyflow,
-  anonymousMode: boolean
+  anonymousMode: boolean,
+  format?: ReIdentifyFormat
 ): Promise<ToolResult<ReIdentifyOutput | AnonymousModeError | ReIdentifyErrorOutput>> {
   if (anonymousMode) {
     return {
@@ -29,14 +66,17 @@ export async function handleReIdentify(
   }
 
   try {
+    const options = buildReidentifyOptions(format);
+
     const response = await skyflow
       .detect()
-      .reidentifyText(new ReidentifyTextRequest(inputString));
+      .reidentifyText(new ReidentifyTextRequest(inputString), options);
 
     return {
       output: {
         inputText: inputString,
         processedText: response.processedText,
+        ...(format && { format }),
       },
     };
   } catch (error) {
diff --git a/src/lib/tools/types.ts b/src/lib/tools/types.ts
index 1d5c320..ca89e1d 100644
--- a/src/lib/tools/types.ts
+++ b/src/lib/tools/types.ts
@@ -19,10 +19,27 @@ export interface DeIdentifyOutput {
   note?: string;
 }
 
+/**
+ * Per-entity-type formatting for re-identification.
+ * Mirrors the Skyflow Detect API `format` object: each entity type listed is
+ * rendered according to its bucket. Entity types not listed in any bucket
+ * default to `plaintext` (full re-identification).
+ */
+export interface ReIdentifyFormat {
+  /** Entity types to fully redact (original value completely hidden). */
+  redacted?: string[];
+  /** Entity types to partially mask (only part of the original value revealed). */
+  masked?: string[];
+  /** Entity types to fully restore to their original plaintext value. */
+  plaintext?: string[];
+}
+
 /** Output from the re-identify tool handler */
 export interface ReIdentifyOutput {
   inputText: string;
   processedText: string;
+  /** The re-identification format that was applied, echoed back when provided. */
+  format?: ReIdentifyFormat;
 }
 
 /** Error output for tools that don't support anonymous mode */
diff --git a/src/server.ts b/src/server.ts
index 81923ed..b0166d4 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -144,11 +144,40 @@ registerAppTool(
   {
     title: "Skyflow Re-identify Tool",
     description:
-      "Re-identify previously de-identified sensitive information in strings using Skyflow. This tool accepts a string with redacted placeholders (like [CREDIT_CARD_abc123]) and returns the original sensitive data.",
-    inputSchema: { inputString: z.string().min(1).describe("Original Text — paste the tokenized text you want to restore") },
+      "Re-identify previously de-identified sensitive information in strings using Skyflow. This tool accepts a string with redacted placeholders (like [CREDIT_CARD_abc123]) and returns the original sensitive data. Optionally pass a `format` object to control how each entity type is rendered on the way out — fully restore (plaintext), partially mask, or fully redact specific entity types. Entity types not listed default to full plaintext restoration.",
+    inputSchema: {
+      inputString: z.string().min(1).describe("Original Text — paste the tokenized text you want to restore"),
+      format: z
+        .object({
+          redacted: z
+            .array(z.enum(ENTITY_KEYS))
+            .optional()
+            .describe("Entity types to fully redact — the original value is completely hidden (e.g. '[REDACTED]')."),
+          masked: z
+            .array(z.enum(ENTITY_KEYS))
+            .optional()
+            .describe("Entity types to partially mask — only part of the original value is revealed (e.g. an SSN shown as '***-**-6789')."),
+          plaintext: z
+            .array(z.enum(ENTITY_KEYS))
+            .optional()
+            .describe("Entity types to fully restore to their original plaintext value."),
+        })
+        .optional()
+        .describe(
+          "Optional per-entity-type re-identification format. List entity types under 'redacted', 'masked', or 'plaintext' to control how each is rendered. Any entity type not listed defaults to full plaintext restoration."
+        ),
+    },
     outputSchema: {
       inputText: z.string().optional().describe("The original tokenized input text"),
       processedText: z.string().optional(),
+      format: z
+        .object({
+          redacted: z.array(z.string()).optional(),
+          masked: z.array(z.string()).optional(),
+          plaintext: z.array(z.string()).optional(),
+        })
+        .optional()
+        .describe("The re-identification format that was applied, echoed back when provided"),
       error: z.union([z.boolean(), z.string()]).optional().describe("Error indicator or message"),
       anonymousModeRestricted: z.boolean().optional().describe("True when blocked due to anonymous mode"),
       message: z.string().optional().describe("Detailed error or setup instructions"),
@@ -158,8 +187,8 @@ registerAppTool(
     },
     _meta: { ui: { resourceUri: RE_IDENTIFY_RESOURCE_URI } },
   },
-  async ({ inputString }) => {
-    const result = await handleReIdentify(inputString, getCurrentSkyflow(), isAnonymousMode());
+  async ({ inputString, format }) => {
+    const result = await handleReIdentify(inputString, getCurrentSkyflow(), isAnonymousMode(), format);
     return {
       content: [{ type: "text", text: JSON.stringify(result.output) }],
       structuredContent: toStructuredContent(result.output),
diff --git a/tests/unit/tools/reIdentify.test.ts b/tests/unit/tools/reIdentify.test.ts
index 0923be9..f3f7870 100644
--- a/tests/unit/tools/reIdentify.test.ts
+++ b/tests/unit/tools/reIdentify.test.ts
@@ -4,6 +4,9 @@ import type { ReIdentifyOutput, ReIdentifyErrorOutput, AnonymousModeError } from
 
 // Mock the skyflow-node SDK
 const mockReidentifyText = vi.fn();
+const mockSetRedactedEntities = vi.fn();
+const mockSetMaskedEntities = vi.fn();
+const mockSetPlainTextEntities = vi.fn();
 
 vi.mock("skyflow-node", () => {
   class MockSkyflowError extends Error {
@@ -16,6 +19,16 @@ vi.mock("skyflow-node", () => {
   }
   return {
     ReidentifyTextRequest: vi.fn(function (this: any, input: string) { this.input = input; }),
+    ReidentifyTextOptions: vi.fn(function (this: any) {
+      this.setRedactedEntities = mockSetRedactedEntities;
+      this.setMaskedEntities = mockSetMaskedEntities;
+      this.setPlainTextEntities = mockSetPlainTextEntities;
+    }),
+    // Proxy returns lowercase prop names to mirror the real DetectEntities enum
+    // values (e.g. DetectEntities.SSN === "ssn"), so getEntityEnum round-trips.
+    DetectEntities: new Proxy({}, { get: (_t, prop) => String(prop).toLowerCase() }),
+    MaskingMethod: new Proxy({}, { get: (_t, prop) => prop }),
+    DetectOutputTranscription: new Proxy({}, { get: (_t, prop) => prop }),
     SkyflowError: MockSkyflowError,
   };
 });
@@ -62,6 +75,98 @@ describe("handleReIdentify", () => {
     });
   });
 
+  describe("format handling", () => {
+    it("should call reidentifyText without options when no format is provided", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      await handleReIdentify("[SSN_abc123]", skyflow as any, false);
+
+      expect(mockReidentifyText).toHaveBeenCalledTimes(1);
+      // Second argument (options) should be undefined for backward compatibility
+      expect(mockReidentifyText.mock.calls[0][1]).toBeUndefined();
+      expect(mockSetRedactedEntities).not.toHaveBeenCalled();
+      expect(mockSetMaskedEntities).not.toHaveBeenCalled();
+      expect(mockSetPlainTextEntities).not.toHaveBeenCalled();
+    });
+
+    it("should not echo a format field when none is provided", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("[SSN_abc123]", skyflow as any, false);
+      const output = result.output as ReIdentifyOutput;
+
+      expect(output).not.toHaveProperty("format");
+    });
+
+    it("should route entity types to redacted / masked / plaintext setters", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      await handleReIdentify("input", skyflow as any, false, {
+        redacted: ["ssn"],
+        masked: ["credit_card"],
+        plaintext: ["email_address", "name"],
+      });
+
+      expect(mockSetRedactedEntities).toHaveBeenCalledWith(["ssn"]);
+      expect(mockSetMaskedEntities).toHaveBeenCalledWith(["credit_card"]);
+      expect(mockSetPlainTextEntities).toHaveBeenCalledWith(["email_address", "name"]);
+      // Options object should be forwarded to the SDK call
+      expect(mockReidentifyText.mock.calls[0][1]).toBeDefined();
+    });
+
+    it("should pass options even when only one bucket is provided", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      await handleReIdentify("input", skyflow as any, false, { masked: ["ssn"] });
+
+      expect(mockSetMaskedEntities).toHaveBeenCalledWith(["ssn"]);
+      expect(mockSetRedactedEntities).not.toHaveBeenCalled();
+      expect(mockSetPlainTextEntities).not.toHaveBeenCalled();
+      expect(mockReidentifyText.mock.calls[0][1]).toBeDefined();
+    });
+
+    it("should ignore empty buckets and skip options when format has no entities", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      await handleReIdentify("input", skyflow as any, false, {
+        redacted: [],
+        masked: [],
+        plaintext: [],
+      });
+
+      expect(mockSetRedactedEntities).not.toHaveBeenCalled();
+      expect(mockSetMaskedEntities).not.toHaveBeenCalled();
+      expect(mockSetPlainTextEntities).not.toHaveBeenCalled();
+      // No entities means no options object is forwarded
+      expect(mockReidentifyText.mock.calls[0][1]).toBeUndefined();
+    });
+
+    it("should echo the applied format back in the output", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const format = { masked: ["ssn"], plaintext: ["name"] };
+      const result = await handleReIdentify("input", skyflow as any, false, format);
+      const output = result.output as ReIdentifyOutput;
+
+      expect(output.format).toEqual(format);
+    });
+
+    it("should echo an empty format object back when provided with no entities", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {});
+      const output = result.output as ReIdentifyOutput;
+
+      expect(output.format).toEqual({});
+    });
+
+    it("should return an error for an invalid entity type in the format", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {
+        masked: ["not_a_real_entity"],
+      });
+      const output = result.output as ReIdentifyErrorOutput;
+
+      expect(result.isError).toBe(true);
+      expect(output.error).toBe(true);
+      expect(output.message).toContain("Invalid entity type");
+      expect(mockReidentifyText).not.toHaveBeenCalled();
+    });
+  });
+
   describe("anonymous mode", () => {
     it("should return error with anonymousModeRestricted flag", async () => {
       const skyflow = createMockSkyflow({ processedText: "" });
diff --git a/ui/re-identify/main.ts b/ui/re-identify/main.ts
index ec00b6b..45bc73c 100644
--- a/ui/re-identify/main.ts
+++ b/ui/re-identify/main.ts
@@ -160,6 +160,39 @@ function highlightValuesInOutput(text: string, mappings: TokenMapping[]): string
   return html;
 }
 
+function renderFormatSummary(format: NonNullable<ReIdentifyResult["format"]>): string {
+  const groups: { label: string; entities: string[] }[] = [
+    { label: "Plaintext", entities: format.plaintext || [] },
+    { label: "Masked", entities: format.masked || [] },
+    { label: "Redacted", entities: format.redacted || [] },
+  ].filter((g) => g.entities.length > 0);
+
+  if (groups.length === 0) return "";
+
+  const groupsHtml = groups
+    .map((g) => {
+      const chips = g.entities
+        .map((e) => {
+          const cls = getEntityClass(e.toUpperCase());
+          const label = e.replace(/_/g, " ");
+          return `<span class="badge ${cls}"><span class="badge-dot" style="background: var(--entity-color)"></span>${escapeHtml(label)}</span>`;
+        })
+        .join(" ");
+      return `
+        <div class="format-group">
+          <span class="format-group-label">${escapeHtml(g.label)}</span>
+          <span class="format-group-chips">${chips}</span>
+        </div>
+      `;
+    })
+    .join("");
+
+  return `
+    <div class="section-heading">Format Applied</div>
+    <div class="format-summary">${groupsHtml}</div>
+  `;
+}
+
 function renderResult(data: ReIdentifyResult): void {
   if (data.error || data.anonymousModeRestricted) {
     const isAnonymous = data.anonymousModeRestricted;
@@ -211,6 +244,8 @@ function renderResult(data: ReIdentifyResult): void {
         </div>
       </div>
 
+      ${data.format ? renderFormatSummary(data.format) : ""}
+
       <div class="panels">
         <div class="panel">
           <div class="panel-header">Tokenized Input</div>
diff --git a/ui/shared/styles.css b/ui/shared/styles.css
index bd940af..2897870 100644
--- a/ui/shared/styles.css
+++ b/ui/shared/styles.css
@@ -247,6 +247,36 @@ body {
   flex-shrink: 0;
 }
 
+/* Re-identify format summary */
+.format-summary {
+  display: flex;
+  flex-direction: column;
+  gap: 8px;
+  margin-bottom: 16px;
+}
+
+.format-group {
+  display: flex;
+  align-items: baseline;
+  flex-wrap: wrap;
+  gap: 6px 10px;
+}
+
+.format-group-label {
+  font-size: 11px;
+  font-weight: var(--font-weight-semibold, 600);
+  color: var(--color-text-secondary, #6b6b80);
+  text-transform: uppercase;
+  letter-spacing: 0.04em;
+  min-width: 68px;
+}
+
+.format-group-chips {
+  display: inline-flex;
+  flex-wrap: wrap;
+  gap: 6px;
+}
+
 /* Info banner */
 .banner {
   padding: 10px 14px;
diff --git a/ui/shared/types.ts b/ui/shared/types.ts
index cbe6e3a..4855603 100644
--- a/ui/shared/types.ts
+++ b/ui/shared/types.ts
@@ -19,10 +19,18 @@ export interface DeIdentifyResult {
   note?: string;
 }
 
+/** Per-entity-type formatting applied during re-identification */
+export interface ReIdentifyFormat {
+  redacted?: string[];
+  masked?: string[];
+  plaintext?: string[];
+}
+
 /** Result from the re-identify tool */
 export interface ReIdentifyResult {
   inputText?: string;
   processedText?: string;
+  format?: ReIdentifyFormat;
   error?: string;
   message?: string;
   anonymousModeRestricted?: boolean;

From b036192b686e9fa51516e9ac4a7a9307cf175088 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:00:30 +0000
Subject: [PATCH 2/7] Harden re-identify format handling per review feedback

- Reject a format where the same entity type appears in more than one
  bucket (redacted/masked/plaintext), which would otherwise forward the
  entity to multiple SDK setters with undefined last-wins behavior.
- Normalize the echoed-back format to omit empty buckets so the response
  reflects only what was actually applied.
- UI: drop the redundant `.toUpperCase()` before `getEntityClass` (which
  lowercases internally) and give the format-summary badge dot a neutral
  color fallback so entity types without a dedicated CSS class still show.
- Docs/tests: note the overlap rejection in CLAUDE.md; add unit tests for
  overlap rejection, empty-bucket normalization, and intra-bucket dups.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 CLAUDE.md                           |  4 +--
 src/lib/tools/reIdentify.ts         | 45 ++++++++++++++++++++++++++++-
 tests/unit/tools/reIdentify.test.ts | 38 ++++++++++++++++++++++++
 ui/re-identify/main.ts              |  4 +--
 4 files changed, 86 insertions(+), 5 deletions(-)

diff --git a/CLAUDE.md b/CLAUDE.md
index 82eb2e4..ab8b060 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -98,8 +98,8 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}"
 
 **re-identify tool** (`src/lib/tools/reIdentify.ts`)
 - Reverses de-identification by replacing tokens with original sensitive data
-- Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed default to full plaintext restoration
-- Returns `inputText` and `processedText`, and echoes back the applied `format` when one was provided
+- Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed default to full plaintext restoration (the same as re-identifying without a format). Rejects with an error if the same entity type appears in more than one bucket
+- Returns `inputText` and `processedText`, and echoes back the applied `format` (empty buckets omitted) when one was provided
 - Returns error with `anonymousModeRestricted: true` in anonymous mode
 
 **de-identify_file tool** (`src/lib/tools/deIdentifyFile.ts`)
diff --git a/src/lib/tools/reIdentify.ts b/src/lib/tools/reIdentify.ts
index 141f56f..dc4aea5 100644
--- a/src/lib/tools/reIdentify.ts
+++ b/src/lib/tools/reIdentify.ts
@@ -34,6 +34,34 @@ function buildReidentifyOptions(
   return hasAny ? options : undefined;
 }
 
+/**
+ * Find entity types that appear in more than one format bucket. Each entity type
+ * may only be rendered one way, so an entity listed under, say, both `redacted`
+ * and `masked` is ambiguous (the SDK would forward it to both setters with
+ * last-wins/undefined behavior). Duplicates within a single bucket are ignored.
+ */
+function findFormatOverlaps(format: ReIdentifyFormat): string[] {
+  const bucketCount = new Map<string, number>();
+  for (const bucket of [format.redacted, format.masked, format.plaintext]) {
+    for (const entity of new Set(bucket ?? [])) {
+      bucketCount.set(entity, (bucketCount.get(entity) ?? 0) + 1);
+    }
+  }
+  return [...bucketCount.entries()].filter(([, count]) => count > 1).map(([entity]) => entity);
+}
+
+/**
+ * Reduce a format to only the buckets that actually carry entity types, so the
+ * value echoed back in the response reflects what was applied (no empty arrays).
+ */
+function normalizeFormat(format: ReIdentifyFormat): ReIdentifyFormat {
+  const normalized: ReIdentifyFormat = {};
+  if (format.redacted && format.redacted.length > 0) normalized.redacted = format.redacted;
+  if (format.masked && format.masked.length > 0) normalized.masked = format.masked;
+  if (format.plaintext && format.plaintext.length > 0) normalized.plaintext = format.plaintext;
+  return normalized;
+}
+
 /**
  * Handle the re-identify tool logic.
  * Restores original sensitive data from de-identified placeholders.
@@ -65,6 +93,21 @@ export async function handleReIdentify(
     };
   }
 
+  if (format) {
+    const overlaps = findFormatOverlaps(format);
+    if (overlaps.length > 0) {
+      return {
+        output: {
+          error: true,
+          message:
+            "Each entity type may appear in only one format bucket (redacted, masked, or plaintext). " +
+            `The following appear in more than one: ${overlaps.join(", ")}.`,
+        },
+        isError: true,
+      };
+    }
+  }
+
   try {
     const options = buildReidentifyOptions(format);
 
@@ -76,7 +119,7 @@ export async function handleReIdentify(
       output: {
         inputText: inputString,
         processedText: response.processedText,
-        ...(format && { format }),
+        ...(format && { format: normalizeFormat(format) }),
       },
     };
   } catch (error) {
diff --git a/tests/unit/tools/reIdentify.test.ts b/tests/unit/tools/reIdentify.test.ts
index f3f7870..d60e605 100644
--- a/tests/unit/tools/reIdentify.test.ts
+++ b/tests/unit/tools/reIdentify.test.ts
@@ -153,6 +153,44 @@ describe("handleReIdentify", () => {
       expect(output.format).toEqual({});
     });
 
+    it("should drop empty buckets from the echoed format", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {
+        masked: ["ssn"],
+        redacted: [],
+        plaintext: [],
+      });
+      const output = result.output as ReIdentifyOutput;
+
+      expect(output.format).toEqual({ masked: ["ssn"] });
+    });
+
+    it("should reject an entity type that appears in more than one bucket", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {
+        redacted: ["ssn"],
+        masked: ["ssn"],
+      });
+      const output = result.output as ReIdentifyErrorOutput;
+
+      expect(result.isError).toBe(true);
+      expect(output.error).toBe(true);
+      expect(output.message).toContain("ssn");
+      expect(output.message).toContain("only one format bucket");
+      expect(mockReidentifyText).not.toHaveBeenCalled();
+    });
+
+    it("should allow the same entity listed twice within a single bucket", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {
+        masked: ["ssn", "ssn"],
+      });
+
+      // A duplicate within one bucket is not a cross-bucket conflict
+      expect(result.isError).toBeUndefined();
+      expect(mockSetMaskedEntities).toHaveBeenCalledWith(["ssn", "ssn"]);
+    });
+
     it("should return an error for an invalid entity type in the format", async () => {
       const skyflow = createMockSkyflow({ processedText: "restored" });
       const result = await handleReIdentify("input", skyflow as any, false, {
diff --git a/ui/re-identify/main.ts b/ui/re-identify/main.ts
index 45bc73c..85f7e4a 100644
--- a/ui/re-identify/main.ts
+++ b/ui/re-identify/main.ts
@@ -173,9 +173,9 @@ function renderFormatSummary(format: NonNullable<ReIdentifyResult["format"]>): s
     .map((g) => {
       const chips = g.entities
         .map((e) => {
-          const cls = getEntityClass(e.toUpperCase());
+          const cls = getEntityClass(e);
           const label = e.replace(/_/g, " ");
-          return `<span class="badge ${cls}"><span class="badge-dot" style="background: var(--entity-color)"></span>${escapeHtml(label)}</span>`;
+          return `<span class="badge ${cls}"><span class="badge-dot" style="background: var(--entity-color, var(--color-text-secondary, #9ca3af))"></span>${escapeHtml(label)}</span>`;
         })
         .join(" ");
       return `

From fadff0f7284e3744c733e061f97e3c4a782b42cd Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:10:23 +0000
Subject: [PATCH 3/7] Refine re-identify format: dedupe, normalize once,
 clarify echo wording

Follow-up to the second review pass:

- Refactor buildReidentifyOptions / normalizeFormat / findFormatOverlaps
  around a shared FORMAT_BUCKETS list so the SDK options and the echoed
  format derive from one normalization step (no duplicated per-bucket
  logic that could drift).
- normalizeFormat now de-duplicates entity types within a bucket, so a
  value listed twice (masked: ["ssn","ssn"]) reaches the SDK and the echo
  once.
- Omit `format` from the response entirely when the normalized result is
  empty, instead of echoing `{}`.
- Reword docs/schema/type comments from "applied" to "requested": the
  echo reflects the caller's requested (normalized) format, not what the
  SDK verified it rendered.
- Tests: dedupe now expects a single entry; empty-format case asserts the
  field is omitted; add a test that the overlap guard short-circuits
  before entity-name validation.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 CHANGELOG.md                        |  2 +-
 CLAUDE.md                           |  2 +-
 src/lib/tools/reIdentify.ts         | 75 ++++++++++++++---------------
 src/lib/tools/types.ts              |  2 +-
 src/server.ts                       |  2 +-
 tests/unit/tools/reIdentify.test.ts | 29 +++++++++--
 6 files changed, 65 insertions(+), 47 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 25b1a31..f409259 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@
 
 ### Added
 
-- **Re-identify output format control** — The `re-identify` tool now accepts an optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings) that lets the caller specify, per entity type, how tokens are rendered on the way out — fully restored (plaintext), partially masked, or fully redacted. Entity types not listed default to full plaintext restoration, preserving existing behavior. Maps to the `skyflow-node` SDK's `ReidentifyTextOptions` per the Detect API spec. The applied format is echoed back in the response and summarized in the re-identify UI.
+- **Re-identify output format control** — The `re-identify` tool now accepts an optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings) that lets the caller specify, per entity type, how tokens are rendered on the way out — fully restored (plaintext), partially masked, or fully redacted. Entity types not listed default to full plaintext restoration, preserving existing behavior. Maps to the `skyflow-node` SDK's `ReidentifyTextOptions` per the Detect API spec. The requested format is echoed back (normalized) in the response and summarized in the re-identify UI.
 
 - **MCP Apps UI for all three tools** — Each tool (`dehydrate`, `rehydrate`, `dehydrate_file`) now has an interactive vanilla TypeScript UI that renders inline in MCP Apps-capable hosts. Text-only hosts continue to receive JSON responses as before.
   - **Dehydrate UI**: Side-by-side before/after text panels with color-coded entity highlights, confidence scores, and an entity breakdown table. Shows anonymous mode banner when applicable.
diff --git a/CLAUDE.md b/CLAUDE.md
index ab8b060..6f2dfee 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -99,7 +99,7 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}"
 **re-identify tool** (`src/lib/tools/reIdentify.ts`)
 - Reverses de-identification by replacing tokens with original sensitive data
 - Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed default to full plaintext restoration (the same as re-identifying without a format). Rejects with an error if the same entity type appears in more than one bucket
-- Returns `inputText` and `processedText`, and echoes back the applied `format` (empty buckets omitted) when one was provided
+- Returns `inputText` and `processedText`, and echoes back the requested `format` (normalized — empty buckets omitted) when one was provided
 - Returns error with `anonymousModeRestricted: true` in anonymous mode
 
 **de-identify_file tool** (`src/lib/tools/deIdentifyFile.ts`)
diff --git a/src/lib/tools/reIdentify.ts b/src/lib/tools/reIdentify.ts
index dc4aea5..45ae72a 100644
--- a/src/lib/tools/reIdentify.ts
+++ b/src/lib/tools/reIdentify.ts
@@ -3,35 +3,24 @@ import type { Skyflow } from "skyflow-node";
 import { getEntityEnum } from "../mappings/entityMaps.js";
 import type { ReIdentifyFormat, ReIdentifyOutput, ReIdentifyErrorOutput, AnonymousModeError, ToolResult } from "./types.js";
 
+/** The three re-identification treatment buckets, in output order. */
+const FORMAT_BUCKETS = ["redacted", "masked", "plaintext"] as const;
+
 /**
- * Build Skyflow re-identify options from the tool's `format` argument.
- * Only buckets containing at least one entity type are set, so an empty or
- * absent bucket falls back to the API default (plaintext re-identification).
- * Returns `undefined` when no entity types are specified, so the SDK is called
- * without options and every token is fully re-identified (the default behavior).
+ * Reduce a format to only the buckets that carry entity types, de-duplicating
+ * within each bucket and dropping empty/absent buckets. The result is what gets
+ * forwarded to the SDK and echoed back, so both stay in lockstep and neither
+ * carries empty arrays or redundant entries.
  */
-function buildReidentifyOptions(
-  format: ReIdentifyFormat | undefined
-): ReidentifyTextOptions | undefined {
-  if (!format) return undefined;
-
-  const options = new ReidentifyTextOptions();
-  let hasAny = false;
-
-  if (format.redacted && format.redacted.length > 0) {
-    options.setRedactedEntities(format.redacted.map(getEntityEnum));
-    hasAny = true;
-  }
-  if (format.masked && format.masked.length > 0) {
-    options.setMaskedEntities(format.masked.map(getEntityEnum));
-    hasAny = true;
-  }
-  if (format.plaintext && format.plaintext.length > 0) {
-    options.setPlainTextEntities(format.plaintext.map(getEntityEnum));
-    hasAny = true;
+function normalizeFormat(format: ReIdentifyFormat): ReIdentifyFormat {
+  const normalized: ReIdentifyFormat = {};
+  for (const bucket of FORMAT_BUCKETS) {
+    const entities = format[bucket];
+    if (entities && entities.length > 0) {
+      normalized[bucket] = [...new Set(entities)];
+    }
   }
-
-  return hasAny ? options : undefined;
+  return normalized;
 }
 
 /**
@@ -42,8 +31,8 @@ function buildReidentifyOptions(
  */
 function findFormatOverlaps(format: ReIdentifyFormat): string[] {
   const bucketCount = new Map<string, number>();
-  for (const bucket of [format.redacted, format.masked, format.plaintext]) {
-    for (const entity of new Set(bucket ?? [])) {
+  for (const bucket of FORMAT_BUCKETS) {
+    for (const entity of new Set(format[bucket] ?? [])) {
       bucketCount.set(entity, (bucketCount.get(entity) ?? 0) + 1);
     }
   }
@@ -51,15 +40,20 @@ function findFormatOverlaps(format: ReIdentifyFormat): string[] {
 }
 
 /**
- * Reduce a format to only the buckets that actually carry entity types, so the
- * value echoed back in the response reflects what was applied (no empty arrays).
+ * Build Skyflow re-identify options from an already-normalized format.
+ * Returns `undefined` when no entity types are specified, so the SDK is called
+ * without options and every token is fully re-identified (the default behavior).
  */
-function normalizeFormat(format: ReIdentifyFormat): ReIdentifyFormat {
-  const normalized: ReIdentifyFormat = {};
-  if (format.redacted && format.redacted.length > 0) normalized.redacted = format.redacted;
-  if (format.masked && format.masked.length > 0) normalized.masked = format.masked;
-  if (format.plaintext && format.plaintext.length > 0) normalized.plaintext = format.plaintext;
-  return normalized;
+function buildReidentifyOptions(
+  normalized: ReIdentifyFormat
+): ReidentifyTextOptions | undefined {
+  if (Object.keys(normalized).length === 0) return undefined;
+
+  const options = new ReidentifyTextOptions();
+  if (normalized.redacted) options.setRedactedEntities(normalized.redacted.map(getEntityEnum));
+  if (normalized.masked) options.setMaskedEntities(normalized.masked.map(getEntityEnum));
+  if (normalized.plaintext) options.setPlainTextEntities(normalized.plaintext.map(getEntityEnum));
+  return options;
 }
 
 /**
@@ -108,8 +102,13 @@ export async function handleReIdentify(
     }
   }
 
+  // Normalize once: drops empty buckets and intra-bucket duplicates. The same
+  // value feeds the SDK options and the echoed-back format.
+  const normalizedFormat = format ? normalizeFormat(format) : undefined;
+  const hasFormat = normalizedFormat !== undefined && Object.keys(normalizedFormat).length > 0;
+
   try {
-    const options = buildReidentifyOptions(format);
+    const options = normalizedFormat ? buildReidentifyOptions(normalizedFormat) : undefined;
 
     const response = await skyflow
       .detect()
@@ -119,7 +118,7 @@ export async function handleReIdentify(
       output: {
         inputText: inputString,
         processedText: response.processedText,
-        ...(format && { format: normalizeFormat(format) }),
+        ...(hasFormat && { format: normalizedFormat }),
       },
     };
   } catch (error) {
diff --git a/src/lib/tools/types.ts b/src/lib/tools/types.ts
index ca89e1d..84a9d4e 100644
--- a/src/lib/tools/types.ts
+++ b/src/lib/tools/types.ts
@@ -38,7 +38,7 @@ export interface ReIdentifyFormat {
 export interface ReIdentifyOutput {
   inputText: string;
   processedText: string;
-  /** The re-identification format that was applied, echoed back when provided. */
+  /** The re-identification format the caller requested (normalized), echoed back when provided. */
   format?: ReIdentifyFormat;
 }
 
diff --git a/src/server.ts b/src/server.ts
index b0166d4..5e750a9 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -177,7 +177,7 @@ registerAppTool(
           plaintext: z.array(z.string()).optional(),
         })
         .optional()
-        .describe("The re-identification format that was applied, echoed back when provided"),
+        .describe("The re-identification format the caller requested, normalized (empty buckets dropped) and echoed back when provided"),
       error: z.union([z.boolean(), z.string()]).optional().describe("Error indicator or message"),
       anonymousModeRestricted: z.boolean().optional().describe("True when blocked due to anonymous mode"),
       message: z.string().optional().describe("Detailed error or setup instructions"),
diff --git a/tests/unit/tools/reIdentify.test.ts b/tests/unit/tools/reIdentify.test.ts
index d60e605..a106bb3 100644
--- a/tests/unit/tools/reIdentify.test.ts
+++ b/tests/unit/tools/reIdentify.test.ts
@@ -145,12 +145,12 @@ describe("handleReIdentify", () => {
       expect(output.format).toEqual(format);
     });
 
-    it("should echo an empty format object back when provided with no entities", async () => {
+    it("should omit format from the output when provided with no entities", async () => {
       const skyflow = createMockSkyflow({ processedText: "restored" });
       const result = await handleReIdentify("input", skyflow as any, false, {});
       const output = result.output as ReIdentifyOutput;
 
-      expect(output.format).toEqual({});
+      expect(output).not.toHaveProperty("format");
     });
 
     it("should drop empty buckets from the echoed format", async () => {
@@ -180,15 +180,34 @@ describe("handleReIdentify", () => {
       expect(mockReidentifyText).not.toHaveBeenCalled();
     });
 
-    it("should allow the same entity listed twice within a single bucket", async () => {
+    it("should de-duplicate an entity listed twice within a single bucket", async () => {
       const skyflow = createMockSkyflow({ processedText: "restored" });
       const result = await handleReIdentify("input", skyflow as any, false, {
         masked: ["ssn", "ssn"],
       });
+      const output = result.output as ReIdentifyOutput;
 
-      // A duplicate within one bucket is not a cross-bucket conflict
+      // A duplicate within one bucket is not a cross-bucket conflict, and it is
+      // collapsed before reaching the SDK and the echoed format.
       expect(result.isError).toBeUndefined();
-      expect(mockSetMaskedEntities).toHaveBeenCalledWith(["ssn", "ssn"]);
+      expect(mockSetMaskedEntities).toHaveBeenCalledWith(["ssn"]);
+      expect(output.format).toEqual({ masked: ["ssn"] });
+    });
+
+    it("should reject an overlap before validating entity names (overlap guard runs first)", async () => {
+      const skyflow = createMockSkyflow({ processedText: "restored" });
+      const result = await handleReIdentify("input", skyflow as any, false, {
+        redacted: ["not_a_real_entity"],
+        masked: ["not_a_real_entity"],
+      });
+      const output = result.output as ReIdentifyErrorOutput;
+
+      // The overlap guard runs before getEntityEnum, so the error is about the
+      // bucket conflict, not the invalid entity name.
+      expect(result.isError).toBe(true);
+      expect(output.message).toContain("only one format bucket");
+      expect(output.message).not.toContain("Invalid entity type");
+      expect(mockReidentifyText).not.toHaveBeenCalled();
     });
 
     it("should return an error for an invalid entity type in the format", async () => {

From 9190abdc446c7d3afd02863bd6aa45c8549a1f9a Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:16:48 +0000
Subject: [PATCH 4/7] Docs: attribute re-identify default to Detect API; soften
 mask example
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Third review pass (doc accuracy):

- Reword the "unlisted entities → plaintext" behavior across the tool
  description, output schema, README, CLAUDE.md, CHANGELOG, and wrapping
  guide to attribute it to the Detect API's default rather than stating it
  as a verified invariant. The unit tests mock the SDK and can't exercise
  the real defaulting semantics, so the docs no longer over-claim.
- Decouple the (verified) backward-compat claim — omitting `format`
  behaves exactly as before — from the API-governed default for unlisted
  entities.
- Soften the partial-mask example: drop the specific "***-**-6789" shape
  in favor of noting the masked form depends on the vault's masking
  configuration.

No behavior change — the handler still forwards the format to the SDK
unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 CHANGELOG.md                            | 2 +-
 CLAUDE.md                               | 2 +-
 README.md                               | 2 +-
 docs/wrapping-mcp-tools-with-skyflow.md | 8 ++++----
 src/server.ts                           | 6 +++---
 5 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f409259..e9fef23 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@
 
 ### Added
 
-- **Re-identify output format control** — The `re-identify` tool now accepts an optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings) that lets the caller specify, per entity type, how tokens are rendered on the way out — fully restored (plaintext), partially masked, or fully redacted. Entity types not listed default to full plaintext restoration, preserving existing behavior. Maps to the `skyflow-node` SDK's `ReidentifyTextOptions` per the Detect API spec. The requested format is echoed back (normalized) in the response and summarized in the re-identify UI.
+- **Re-identify output format control** — The `re-identify` tool now accepts an optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings) that lets the caller specify, per entity type, how tokens are rendered on the way out — fully restored (plaintext), partially masked, or fully redacted. Omitting `format` preserves existing behavior exactly; entity types not listed in a provided `format` fall back to the Detect API's default and are restored as full plaintext. Maps to the `skyflow-node` SDK's `ReidentifyTextOptions` per the Detect API spec. The requested format is echoed back (normalized) in the response and summarized in the re-identify UI.
 
 - **MCP Apps UI for all three tools** — Each tool (`dehydrate`, `rehydrate`, `dehydrate_file`) now has an interactive vanilla TypeScript UI that renders inline in MCP Apps-capable hosts. Text-only hosts continue to receive JSON responses as before.
   - **Dehydrate UI**: Side-by-side before/after text panels with color-coded entity highlights, confidence scores, and an entity breakdown table. Shows anonymous mode banner when applicable.
diff --git a/CLAUDE.md b/CLAUDE.md
index 6f2dfee..42d38f8 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -98,7 +98,7 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}"
 
 **re-identify tool** (`src/lib/tools/reIdentify.ts`)
 - Reverses de-identification by replacing tokens with original sensitive data
-- Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed default to full plaintext restoration (the same as re-identifying without a format). Rejects with an error if the same entity type appears in more than one bucket
+- Accepts optional `format` object (`{ redacted?, masked?, plaintext? }`, each a list of entity type strings from `ENTITY_MAP`) to control how each entity type is rendered on the way out, per the Detect API spec. Maps to the SDK's `ReidentifyTextOptions` (`setRedactedEntities` / `setMaskedEntities` / `setPlainTextEntities`). Entity types not listed fall back to the Detect API's default and are restored as full plaintext (the same as re-identifying without a format). Rejects with an error if the same entity type appears in more than one bucket
 - Returns `inputText` and `processedText`, and echoes back the requested `format` (normalized — empty buckets omitted) when one was provided
 - Returns error with `anonymousModeRestricted: true` in anonymous mode
 
diff --git a/README.md b/README.md
index beaf8e6..0388ba5 100644
--- a/README.md
+++ b/README.md
@@ -259,7 +259,7 @@ By default every token is fully restored to its original value. To control how
 individual entity types are rendered, pass an optional `format` object with any
 of `redacted`, `masked`, or `plaintext` — each a list of entity types (the same
 lowercase names as the `de-identify` tool, e.g. `ssn`, `email_address`). Entity
-types not listed default to full plaintext restoration:
+types not listed fall back to the Detect API's default and are restored as full plaintext:
 
 ```bash
 curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" \
diff --git a/docs/wrapping-mcp-tools-with-skyflow.md b/docs/wrapping-mcp-tools-with-skyflow.md
index 0e548ea..19f83ef 100644
--- a/docs/wrapping-mcp-tools-with-skyflow.md
+++ b/docs/wrapping-mcp-tools-with-skyflow.md
@@ -224,15 +224,15 @@ reference list. Common values include `email_address`, `ssn`, `credit_card`, `na
 
 Re-identify fully restores every token by default. To render specific entity types differently on
 the way out — reveal some, partially mask others, fully redact the rest — pass a
-`ReidentifyTextOptions` as the second argument. Entity types you don't list default to full plaintext
-restoration:
+`ReidentifyTextOptions` as the second argument. Entity types you don't list fall back to the Detect
+API's default and are restored as full plaintext:
 
 ```ts
 import { ReidentifyTextOptions, DetectEntities } from "skyflow-node";
 
 const options = new ReidentifyTextOptions();
 options.setPlainTextEntities([DetectEntities.NAME]);      // restore in full
-options.setMaskedEntities([DetectEntities.SSN]);          // partially masked (e.g. ***-**-6789)
+options.setMaskedEntities([DetectEntities.SSN]);          // partially masked; exact form depends on vault config
 options.setRedactedEntities([DetectEntities.EMAIL_ADDRESS]); // fully hidden
 
 const res = await skyflow
@@ -304,7 +304,7 @@ curl -X POST "https://$CLUSTER_ID.vault.skyflowapis.com/v1/detect/reidentify/str
 
 `text` and `vault_id` are required. Add an optional `format` object to control how each entity type
 is rendered — `{ "redacted": [...], "masked": [...], "plaintext": [...] }`, each a list of the same
-lowercase entity names. Entity types you don't list default to full plaintext restoration:
+lowercase entity names. Entity types you don't list fall back to the Detect API's default and are restored as full plaintext:
 
 ```json
 {
diff --git a/src/server.ts b/src/server.ts
index 5e750a9..6048e06 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -144,7 +144,7 @@ registerAppTool(
   {
     title: "Skyflow Re-identify Tool",
     description:
-      "Re-identify previously de-identified sensitive information in strings using Skyflow. This tool accepts a string with redacted placeholders (like [CREDIT_CARD_abc123]) and returns the original sensitive data. Optionally pass a `format` object to control how each entity type is rendered on the way out — fully restore (plaintext), partially mask, or fully redact specific entity types. Entity types not listed default to full plaintext restoration.",
+      "Re-identify previously de-identified sensitive information in strings using Skyflow. This tool accepts a string with redacted placeholders (like [CREDIT_CARD_abc123]) and returns the original sensitive data. Optionally pass a `format` object to control how each entity type is rendered on the way out — fully restore (plaintext), partially mask, or fully redact specific entity types. Entity types not listed fall back to the Detect API's default and are restored as full plaintext (the same as re-identifying without a format).",
     inputSchema: {
       inputString: z.string().min(1).describe("Original Text — paste the tokenized text you want to restore"),
       format: z
@@ -156,7 +156,7 @@ registerAppTool(
           masked: z
             .array(z.enum(ENTITY_KEYS))
             .optional()
-            .describe("Entity types to partially mask — only part of the original value is revealed (e.g. an SSN shown as '***-**-6789')."),
+            .describe("Entity types to partially mask — only part of the original value is revealed. The exact masked form depends on the vault's masking configuration."),
           plaintext: z
             .array(z.enum(ENTITY_KEYS))
             .optional()
@@ -164,7 +164,7 @@ registerAppTool(
         })
         .optional()
         .describe(
-          "Optional per-entity-type re-identification format. List entity types under 'redacted', 'masked', or 'plaintext' to control how each is rendered. Any entity type not listed defaults to full plaintext restoration."
+          "Optional per-entity-type re-identification format. List entity types under 'redacted', 'masked', or 'plaintext' to control how each is rendered. Any entity type not listed falls back to the Detect API's default and is restored as full plaintext."
         ),
     },
     outputSchema: {

From 58820e376dae857a89d42fd891e52c129d5f0d19 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:22:42 +0000
Subject: [PATCH 5/7] Align UI wording to "requested" and give invalid entities
 a clean error

Fourth review pass:

- UI: rename the re-identify summary heading "Format Applied" ->
  "Requested Format" and the shared-type comment to match the "requested"
  framing used in the schema, type docs, and CHANGELOG (the echo reflects
  what the caller asked for, not what the SDK verified it rendered).
- Handler: validate entity names before the try, alongside the overlap
  guard, so an invalid entity type returns a distinct client-side
  validation error (no HTTP `code`) instead of surfacing through the
  generic catch looking like a Skyflow API error. Overlap still checked
  first (ordering test preserved).
- Test: assert the invalid-entity error carries no `code`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 src/lib/tools/reIdentify.ts         | 30 ++++++++++++++++++++++++++++-
 tests/unit/tools/reIdentify.test.ts |  3 +++
 ui/re-identify/main.ts              |  2 +-
 ui/shared/types.ts                  |  2 +-
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/lib/tools/reIdentify.ts b/src/lib/tools/reIdentify.ts
index 45ae72a..5c9082b 100644
--- a/src/lib/tools/reIdentify.ts
+++ b/src/lib/tools/reIdentify.ts
@@ -1,6 +1,6 @@
 import { ReidentifyTextOptions, ReidentifyTextRequest, SkyflowError } from "skyflow-node";
 import type { Skyflow } from "skyflow-node";
-import { getEntityEnum } from "../mappings/entityMaps.js";
+import { getEntityEnum, isValidEntity } from "../mappings/entityMaps.js";
 import type { ReIdentifyFormat, ReIdentifyOutput, ReIdentifyErrorOutput, AnonymousModeError, ToolResult } from "./types.js";
 
 /** The three re-identification treatment buckets, in output order. */
@@ -23,6 +23,21 @@ function normalizeFormat(format: ReIdentifyFormat): ReIdentifyFormat {
   return normalized;
 }
 
+/**
+ * Collect any entity types in the format that are not recognized. Returned so
+ * an invalid name can be reported as a client-side validation error (before any
+ * Skyflow call), distinct in shape from a Skyflow API error.
+ */
+function findInvalidEntities(format: ReIdentifyFormat): string[] {
+  const invalid = new Set<string>();
+  for (const bucket of FORMAT_BUCKETS) {
+    for (const entity of format[bucket] ?? []) {
+      if (!isValidEntity(entity)) invalid.add(entity);
+    }
+  }
+  return [...invalid];
+}
+
 /**
  * Find entity types that appear in more than one format bucket. Each entity type
  * may only be rendered one way, so an entity listed under, say, both `redacted`
@@ -100,6 +115,19 @@ export async function handleReIdentify(
         isError: true,
       };
     }
+
+    const invalidEntities = findInvalidEntities(format);
+    if (invalidEntities.length > 0) {
+      return {
+        output: {
+          error: true,
+          message:
+            `Invalid entity type(s) in format: ${invalidEntities.join(", ")}. ` +
+            "Use the same lowercase entity names as the de-identify tool.",
+        },
+        isError: true,
+      };
+    }
   }
 
   // Normalize once: drops empty buckets and intra-bucket duplicates. The same
diff --git a/tests/unit/tools/reIdentify.test.ts b/tests/unit/tools/reIdentify.test.ts
index a106bb3..384afa7 100644
--- a/tests/unit/tools/reIdentify.test.ts
+++ b/tests/unit/tools/reIdentify.test.ts
@@ -220,6 +220,9 @@ describe("handleReIdentify", () => {
       expect(result.isError).toBe(true);
       expect(output.error).toBe(true);
       expect(output.message).toContain("Invalid entity type");
+      // Client-side validation error: no HTTP code, so it is distinguishable
+      // from a Skyflow API error (which carries `code`).
+      expect(output.code).toBeUndefined();
       expect(mockReidentifyText).not.toHaveBeenCalled();
     });
   });
diff --git a/ui/re-identify/main.ts b/ui/re-identify/main.ts
index 85f7e4a..06679da 100644
--- a/ui/re-identify/main.ts
+++ b/ui/re-identify/main.ts
@@ -188,7 +188,7 @@ function renderFormatSummary(format: NonNullable<ReIdentifyResult["format"]>): s
     .join("");
 
   return `
-    <div class="section-heading">Format Applied</div>
+    <div class="section-heading">Requested Format</div>
     <div class="format-summary">${groupsHtml}</div>
   `;
 }
diff --git a/ui/shared/types.ts b/ui/shared/types.ts
index 4855603..1cf1578 100644
--- a/ui/shared/types.ts
+++ b/ui/shared/types.ts
@@ -19,7 +19,7 @@ export interface DeIdentifyResult {
   note?: string;
 }
 
-/** Per-entity-type formatting applied during re-identification */
+/** Per-entity-type formatting requested during re-identification */
 export interface ReIdentifyFormat {
   redacted?: string[];
   masked?: string[];

From 468953ffac9737c14e97ebe88c51f6969a9ce11e Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:26:43 +0000
Subject: [PATCH 6/7] Clarify plaintext-bucket redundancy and UI ordering

Fifth review pass (clarity only, no behavior change):

- Note in the `plaintext` schema description that it is redundant with the
  default treatment for unlisted entities and exists only for explicitness,
  so callers don't assume they must list entities to get restoration.
- Comment renderFormatSummary's least-to-most-restrictive group ordering,
  which intentionally differs from the handler's FORMAT_BUCKETS order.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 src/server.ts          | 2 +-
 ui/re-identify/main.ts | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/server.ts b/src/server.ts
index 6048e06..b604c21 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -160,7 +160,7 @@ registerAppTool(
           plaintext: z
             .array(z.enum(ENTITY_KEYS))
             .optional()
-            .describe("Entity types to fully restore to their original plaintext value."),
+            .describe("Entity types to fully restore to their original plaintext value. This is already the default for any unlisted entity type, so listing here is only for explicitness."),
         })
         .optional()
         .describe(
diff --git a/ui/re-identify/main.ts b/ui/re-identify/main.ts
index 06679da..5560aab 100644
--- a/ui/re-identify/main.ts
+++ b/ui/re-identify/main.ts
@@ -161,6 +161,8 @@ function highlightValuesInOutput(text: string, mappings: TokenMapping[]): string
 }
 
 function renderFormatSummary(format: NonNullable<ReIdentifyResult["format"]>): string {
+  // Ordered least- to most-restrictive (revealed -> hidden) for readability;
+  // this intentionally differs from the handler's FORMAT_BUCKETS order.
   const groups: { label: string; entities: string[] }[] = [
     { label: "Plaintext", entities: format.plaintext || [] },
     { label: "Masked", entities: format.masked || [] },

From 06eab492582f2af594389f1f21bcf82453e2af41 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Tue, 21 Jul 2026 17:52:53 +0000
Subject: [PATCH 7/7] Comment the handler's intentional re-validation of entity
 names
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Sixth review pass (comment only): note that handleReIdentify re-validates
entity names even though the server inputSchema already enum-constrains
them, because the handler is a standalone reusable function — so a future
reader doesn't remove the guard as "redundant."

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AZZMCL2fjoR3FvvoM3vofZ
---
 src/lib/tools/reIdentify.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/lib/tools/reIdentify.ts b/src/lib/tools/reIdentify.ts
index 5c9082b..ba3d94a 100644
--- a/src/lib/tools/reIdentify.ts
+++ b/src/lib/tools/reIdentify.ts
@@ -116,6 +116,10 @@ export async function handleReIdentify(
       };
     }
 
+    // The server inputSchema already enum-constrains entity names, but this
+    // handler is a standalone, independently-testable function and must not
+    // assume the caller pre-validated — re-validate so a direct call (or a
+    // future caller that skips the schema) still gets a clean client-side error.
     const invalidEntities = findInvalidEntities(format);
     if (invalidEntities.length > 0) {
       return {