Add optional format support to the re-identify tool

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. 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.
   - **Rehydrate UI**: Token-to-original mapping display with color-matched highlights across before/after panels.

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 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
 
 **de-identify_file tool** (`src/lib/tools/deIdentifyFile.ts`)

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 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}" \
+  -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`:

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 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; exact form depends on vault config
+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 fall back to the Detect API's default and are restored as full plaintext:
+
+```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` /

src/lib/tools/reIdentify.ts

@@ -1,15 +1,89 @@
-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, isValidEntity } 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;
+
+/**
+ * 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 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 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`
+ * 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_BUCKETS) {
+    for (const entity of new Set(format[bucket] ?? [])) {
+      bucketCount.set(entity, (bucketCount.get(entity) ?? 0) + 1);
+    }
+  }
+  return [...bucketCount.entries()].filter(([, count]) => count > 1).map(([entity]) => entity);
+}
+
+/**
+ * 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 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;
+}
 
 /**
  * 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 {
@@ -28,15 +102,55 @@ 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,
+      };
+    }
+
+    // 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 {
+        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
+  // 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 = normalizedFormat ? buildReidentifyOptions(normalizedFormat) : undefined;
+
     const response = await skyflow
       .detect()
-      .reidentifyText(new ReidentifyTextRequest(inputString));
+      .reidentifyText(new ReidentifyTextRequest(inputString), options);
 
     return {
       output: {
         inputText: inputString,
         processedText: response.processedText,
+        ...(hasFormat && { format: normalizedFormat }),
       },
     };
   } catch (error) {

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 the caller requested (normalized), echoed back when provided. */
+  format?: ReIdentifyFormat;
 }
 
 /** Error output for tools that don't support anonymous mode */

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 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
+        .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. The exact masked form depends on the vault's masking configuration."),
+          plaintext: z
+            .array(z.enum(ENTITY_KEYS))
+            .optional()
+            .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(
+          "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: {
       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 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"),
@@ -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),