From 762b0b103c3b32fbd0f4f03f0bae52395859c4e8 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 4 Jul 2026 22:13:43 +0000 Subject: [PATCH 1/8] feat: add file de-identification and re-identification tools Expose Skyflow's Detect file endpoints as three MCP tools, with support for passing files by signed/public URL and a job-handle pattern for Skyflow's asynchronous file processing. Tools: - de-identify-file: redact PII/PHI in images, PDFs, Office docs, text, and audio. Accepts a signed/public URL (downloaded server-side and converted to base64) or inline base64. Supports the common Skyflow file options (entities, allow/restrict regex, token type, masking method, OCR/transcription output, pixel density/max resolution, date shifting, audio bleep). - get-file-run-status: poll an async de-identification run by runId, with optional bounded server-side long-polling. - re-identify-file: restore original values in a previously de-identified text-based file. Async handling: the de-identify call waits a bounded time; if the run is still processing it returns the Skyflow runId + IN_PROGRESS status and a note directing the agent to get-file-run-status. This keeps the server stateless/serverless-friendly while covering Skyflow's async file API. Implementation: - Upgrade skyflow-node ^2.0.0 -> ^2.1.2 (fixes the IN_PROGRESS timeout return path used by the async flow). - src/lib/files/fileSource.ts: URL download + input resolution with a size cap enforced while streaming, a request timeout that covers the body read, and an SSRF guard (scheme allowlist, private/loopback/ link-local IP blocking across encodings, per-hop redirect re-validation). - src/lib/detect/detectRest.ts: fetch client for the runs and reidentify-file endpoints the SDK doesn't expose; tolerant of both camelCase and snake_case responses. - src/lib/mappings/fileFormats.ts: format allowlists, extension<->MIME maps, and the set of formats that honor transformations. - Derive real MIME types for processed files from their extension (the SDK reports a category label, not a MIME type). - Warn when date shifting is requested for a format Skyflow ignores it on; omit zero-valued count fields the SDK defaults to 0. - File-tool responses keep full data in structuredContent but omit the large base64 blobs from the text channel to avoid double-serializing multi-MB payloads. - Raise the JSON body limit 5MB -> 25MB for inline base64 files. New re-identify-file UI app; de-identify-file UI shows an IN_PROGRESS state and is shared by get-file-run-status. Docs (README, CLAUDE.md, CHANGELOG, mcp-apps-ui) updated. 76 new unit tests (244 total). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Nv6ppGHwgQFFZSd32FBiZA --- CHANGELOG.md | 15 + CLAUDE.md | 79 ++++- README.md | 55 +++- docs/mcp-apps-ui.md | 9 +- package.json | 4 +- pnpm-lock.yaml | 80 +++-- scripts/generate-ui-imports.ts | 1 + src/generated/ui-html.d.ts | 1 + src/lib/detect/detectRest.ts | 244 ++++++++++++++ src/lib/files/fileSource.ts | 383 ++++++++++++++++++++++ src/lib/mappings/entityMaps.ts | 12 + src/lib/mappings/fileFormats.ts | 164 +++++++++ src/lib/tools/deIdentifyFile.ts | 208 ++++++++++-- src/lib/tools/getFileRunStatus.ts | 157 +++++++++ src/lib/tools/reIdentifyFile.ts | 143 ++++++++ src/lib/tools/types.ts | 89 ++++- src/server.ts | 375 ++++++++++++++++++++- tests/unit/files/fileSource.test.ts | 281 ++++++++++++++++ tests/unit/tools/deIdentifyFile.test.ts | 336 ++++++++++++++++++- tests/unit/tools/getFileRunStatus.test.ts | 248 ++++++++++++++ tests/unit/tools/reIdentifyFile.test.ts | 243 ++++++++++++++ ui/de-identify-file/main.ts | 52 ++- ui/re-identify-file/main.ts | 139 ++++++++ ui/re-identify-file/mcp-app.html | 12 + ui/shared/types.ts | 17 +- 25 files changed, 3220 insertions(+), 127 deletions(-) create mode 100644 src/lib/detect/detectRest.ts create mode 100644 src/lib/files/fileSource.ts create mode 100644 src/lib/mappings/fileFormats.ts create mode 100644 src/lib/tools/getFileRunStatus.ts create mode 100644 src/lib/tools/reIdentifyFile.ts create mode 100644 tests/unit/files/fileSource.test.ts create mode 100644 tests/unit/tools/getFileRunStatus.test.ts create mode 100644 tests/unit/tools/reIdentifyFile.test.ts create mode 100644 ui/re-identify-file/main.ts create mode 100644 ui/re-identify-file/mcp-app.html diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a4b8f2..5ccf265 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,21 @@ ### Added +- **File de-identification and re-identification tools** — Three new MCP tools expose Skyflow's file endpoints: + - **`de-identify-file`** — Detects and redacts sensitive data in images, PDFs, Office documents, txt/csv/json/xml, dcm, and audio. Accepts the file as a **signed/public URL** (`fileUrl`, downloaded server-side with a 25 MB cap and converted to the base64 the Skyflow API requires) or as inline base64 (`fileDataBase64` + `fileName`). Supports the common Skyflow file options: `entities`, `allowRegexList`/`restrictRegexList`, `tokenType`, `maskingMethod`, `outputProcessedFile`, `outputOcrText`, `outputTranscription`, `pixelDensity`/`maxResolution`, `dateShift`, and audio `bleep` settings. + - **`get-file-run-status`** — Polls the asynchronous de-identification run by `runId`, with optional bounded server-side long-polling (`waitSeconds`). File processing at Skyflow is async: `de-identify-file` waits up to `waitTimeSeconds` (default 25s, max 64s) and returns either the completed result or `runId` + `IN_PROGRESS` with instructions to poll — the standard MCP job-handle pattern for long-running work on a stateless server. + - **`re-identify-file`** — Restores original values in previously de-identified csv/doc/docx/json/txt/xls/xlsx/xml files via Skyflow's synchronous reidentify-file endpoint, with `redactedEntities`/`maskedEntities`/`plainTextEntities` routing. + - New `ui/re-identify-file/` MCP Apps UI; `get-file-run-status` shares the de-identify-file app. New shared helpers: `src/lib/files/fileSource.ts` (URL download with SSRF guard, size/timeout limits, filename inference) and `src/lib/detect/detectRest.ts` (direct REST client for the runs and reidentify-file endpoints, tolerant of both camelCase and snake_case response fields). + - All three tools require authenticated mode; in anonymous mode they return setup instructions. + +### Changed (file tools) + +- **`skyflow-node` upgraded to ^2.1.2** — required for the fixed `IN_PROGRESS` timeout path in `deidentifyFile` (2.0.0 crashed with a destructuring TypeError when a run outlived the SDK wait window). +- **JSON body limit raised from 5 MB to 25 MB** to accommodate inline base64 file payloads. +- **`de-identify_file` handler reworked** — the previously disabled handler was redesigned (URL input, format validation, new options, polling note) and registered as `de-identify-file`. + +### Added + - **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 bb0ccbf..f1e610a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -44,6 +44,20 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" -H "Accept: application/json, text/event-stream" \ -H "Authorization: Bearer {your_bearer_token}" \ -d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"de-identify","arguments":{"inputString":"My email is john.doe@example.com"}},"id":2}' + +# Call de-identify-file tool with a signed/public URL +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":"de-identify-file","arguments":{"fileUrl":"https://example-bucket.s3.amazonaws.com/scan.pdf?X-Amz-Signature=..."}},"id":3}' + +# Poll an in-progress file run (runId comes from the de-identify-file response) +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":"get-file-run-status","arguments":{"runId":"{run_id}","waitSeconds":30}},"id":4}' ``` ## Architecture @@ -54,17 +68,17 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" - Serves a single `/mcp` endpoint that handles all MCP protocol requests - Accepts query parameters: `vaultId`, `vaultUrl`, `apiKey` (optional) - Uses credentials extraction middleware to validate either Authorization header or apiKey query parameter -- Configured with 5MB JSON payload limit to support base64-encoded files +- Configured with 25MB JSON payload limit to support base64-encoded files passed inline **MCP Server Instance** -- Registers two active tools: `de-identify` and `re-identify` +- Registers five active tools: `de-identify`, `re-identify`, `de-identify-file`, `get-file-run-status`, and `re-identify-file` - Each tool is registered via `registerAppTool` from `@modelcontextprotocol/ext-apps/server`, linking tools to interactive UI resources - Each tool is defined with Zod schemas for input validation and output structure - Uses the official `@modelcontextprotocol/sdk` library **MCP Apps UI Layer** (`ui/`) -- Active tool UIs: `ui/de-identify/` and `ui/re-identify/` -- `ui/de-identify-file/` exists but its resource is not registered (tool is disabled) +- Tool UIs: `ui/de-identify/`, `ui/re-identify/`, `ui/de-identify-file/`, `ui/re-identify-file/` +- The `get-file-run-status` tool shares the `ui/de-identify-file/` app (both render de-identification run results) - Shared theme/styles in `ui/shared/` (theme.ts, styles.css) - Built with Vite + `vite-plugin-singlefile` → single HTML files in `dist/ui/` - Resources registered via `registerAppResource` with `ui://` URIs @@ -85,6 +99,19 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" - Credentials are forwarded to Skyflow API in the appropriate format: `{ token: string }` or `{ apiKey: string }` - Uses `AsyncLocalStorage` to make Skyflow instance available to tools during request handling - Tools access the current request's Skyflow instance via `getCurrentSkyflow()` and mode via `isAnonymousMode()` +- The request context also carries `vaultUrl` and the raw credential value for tools that call Detect REST endpoints directly (via `getDetectRestContext()`) + +**Detect REST Helper** (`src/lib/detect/detectRest.ts`) +- Minimal fetch-based client for Detect endpoints the skyflow-node SDK doesn't expose: + `GET /v1/detect/runs/{run_id}` (run status) and `POST /v1/detect/reidentify/file` +- Sends the same `Authorization: Bearer ` the SDK would send (JWT or API key) +- Defensively parses both camelCase and snake_case response fields — the live API returns camelCase for several fields even though the OpenAPI types say snake_case + +**File Source Helper** (`src/lib/files/fileSource.ts`) +- Resolves tool file inputs: either a signed/public `fileUrl` (downloaded server-side, 25MB cap, 30s timeout) or inline `fileDataBase64` + `fileName` +- Skyflow's file endpoints only accept base64, so URLs are always downloaded and converted before forwarding +- Blocks obvious SSRF targets (localhost, private/link-local IP literals, `.internal`/`.local` hosts); https or http only +- Infers the file name from the explicit arg → `Content-Disposition` header → URL path → `Content-Type` fallback; the lowercased extension selects the Skyflow endpoint/data_format ### Tool Implementations @@ -101,15 +128,31 @@ curl -X POST "http://localhost:3000/mcp?vaultId={vault_id}&vaultUrl={vault_url}" - Returns `inputText` and `processedText` - Returns error with `anonymousModeRestricted: true` in anonymous mode -**de-identify_file tool** (`src/lib/tools/deIdentifyFile.ts`) -- Currently disabled (not registered) — handler code preserved in `deIdentifyFile.ts` for future re-enablement -- Was used to process images, PDFs, audio, and documents +**de-identify-file tool** (`src/lib/tools/deIdentifyFile.ts`) +- Detects and redacts sensitive information in files: images (bmp, jpeg, jpg, png, tif, tiff), PDFs, Word/Excel/PowerPoint, txt, csv, json, xml, dcm, and audio (mp3, wav) +- Accepts the file as a signed/public `fileUrl` (downloaded server-side and converted to base64) or as `fileDataBase64` + `fileName` +- Options: `entities`, `allowRegexList`, `restrictRegexList`, `tokenType` (`entity_unique_counter` | `entity_only`), `maskingMethod` (images), `outputProcessedFile` (images/audio), `outputOcrText` (images), `outputTranscription` (audio), `pixelDensity`/`maxResolution` (PDFs), `dateShift`, `bleep` (audio) +- **Async handling**: Skyflow file de-identification is asynchronous. The tool waits up to `waitTimeSeconds` (default 25, max 64) for the run to finish; if it's still processing, the response carries `runId`, `status: "IN_PROGRESS"`, and a `note` instructing the agent to poll with `get-file-run-status` +- Uses the SDK's `deidentifyFile`, which routes to the type-specific Skyflow endpoint based on the file extension + +**get-file-run-status tool** (`src/lib/tools/getFileRunStatus.ts`) +- Polls an asynchronous de-identification run by `runId` (from `de-identify-file`) +- Optional `waitSeconds` (0-55) long-polls server-side with backoff before returning +- On `SUCCESS`, returns the same output shape as `de-identify-file` (processed file base64, detected entities, counts) +- On `FAILED`, returns `isError` with Skyflow's failure message +- Calls `GET /v1/detect/runs/{run_id}` directly via the Detect REST helper + +**re-identify-file tool** (`src/lib/tools/reIdentifyFile.ts`) +- Restores original sensitive data in a previously de-identified file (formats: csv, doc, docx, json, txt, xls, xlsx, xml) +- Accepts `fileUrl` or `fileDataBase64` + `fileName`, plus optional `redactedEntities`/`maskedEntities`/`plainTextEntities` lists controlling how each entity type is restored +- The Skyflow endpoint is synchronous — the processed file is returned directly +- Calls `POST /v1/detect/reidentify/file` via the Detect REST helper (the SDK has no high-level file re-identify) **Tool handler extraction pattern** - Core logic is in `src/lib/tools/*.ts` as pure functions with explicit parameters -- `src/server.ts` calls these functions, passing `getCurrentSkyflow()`, `isAnonymousMode()` +- `src/server.ts` calls these functions, passing `getCurrentSkyflow()`, `isAnonymousMode()`, and for REST-based tools `getDetectRestContext()` - Shared types live in `src/lib/tools/types.ts` -- Tool files: `deIdentify.ts`, `reIdentify.ts`, `deIdentifyFile.ts` +- Tool files: `deIdentify.ts`, `reIdentify.ts`, `deIdentifyFile.ts`, `getFileRunStatus.ts`, `reIdentifyFile.ts` - This separation enables unit testing without `AsyncLocalStorage` context ### Type Safety Approach @@ -173,6 +216,9 @@ When no credentials are provided in a request, the server can operate in "anonym |------|---------------|-------------------| | `de-identify` | Works with `ENTITY_UNIQUE_COUNTER` tokens | Works with `VAULT_TOKEN` tokens | | `re-identify` | Returns error with setup instructions | Works normally | +| `de-identify-file` | Returns error with setup instructions | Works normally | +| `get-file-run-status` | Returns error with setup instructions | Works normally | +| `re-identify-file` | Returns error with setup instructions | Works normally | **Token Format Difference**: - **Anonymous mode**: Uses `TokenType.ENTITY_UNIQUE_COUNTER` - generates tokens like `[EMAIL_ADDRESS_1]`, `[SSN_2]`. Data is NOT persisted to vault. @@ -280,7 +326,7 @@ The `isError` property is set to `true` when a tool returns an error condition ( - `@modelcontextprotocol/sdk`: Official MCP TypeScript SDK (v1.27.1+) - `@modelcontextprotocol/ext-apps`: MCP Apps SDK for interactive tool UIs -- `skyflow-node`: Skyflow SDK for deidentification (v2.0.0+) +- `skyflow-node`: Skyflow SDK for deidentification (v2.1.2+ — required for the fixed `IN_PROGRESS` return path in `deidentifyFile`) - `express`: Web framework (v5.1.0+) - `zod`: Schema validation for tool inputs/outputs - `dotenv`: Environment variable management @@ -311,12 +357,19 @@ All of the above, plus: - [ ] **New UI** `ui//main.ts` — implement `ontoolinput` (loading state) and `ontoolresult` (result rendering) - [ ] **Update MCP Server Instance count** in this file +### When adding a new UI app + +- [ ] Add the app directory `ui//` with `mcp-app.html` and `main.ts` +- [ ] Add an `INPUT=/mcp-app.html vite build` step to `build:ui` in `package.json` +- [ ] Add the `{ varName, dir }` entry to `scripts/generate-ui-imports.ts` +- [ ] Add the export to `src/generated/ui-html.d.ts` (hand-maintained declaration for the generated module) + ### When disabling a tool - [ ] Remove `registerAppTool` and `registerAppResource` calls from `src/server.ts` - [ ] Remove unused HTML import from `./generated/ui-html.js` - [ ] Remove unused handler import; keep the handler file itself for future re-enablement -- [ ] Remove any helper functions only used by that tool (e.g. `getCurrentVaultId` was removed when `de-identify_file` was disabled) +- [ ] Remove any helper functions only used by that tool - [ ] Update CLAUDE.md: tool count, Tool Implementations section, Anonymous Mode table, Common Pitfalls ## Common Pitfalls @@ -327,5 +380,7 @@ All of the above, plus: 4. **Vault configuration** - `clusterId` is automatically extracted from `vaultUrl`, don't set it separately 5. **Entity type validation** - Use exact strings from `ENTITY_MAP` keys, not the enum values 6. **AsyncLocalStorage context** - Tools must run within the request context to access Skyflow instance via `getCurrentSkyflow()` and `isAnonymousMode()` -7. **Anonymous mode limitations** - Only the `de-identify` tool works in anonymous mode; `re-identify` returns an error with setup instructions +7. **Anonymous mode limitations** - Only the `de-identify` tool works in anonymous mode; the other tools return an error with setup instructions 8. **Keep schemas in sync** - When modifying tool inputs or return values, always update the corresponding `inputSchema` and `outputSchema` in the tool registration. The schemas must match the actual implementation. +9. **File de-identification is asynchronous** - `de-identify-file` may return `runId` + `status: "IN_PROGRESS"` instead of the processed file; results are then fetched with `get-file-run-status`. Keep server-side waits bounded (SDK max 64s; watch serverless execution limits when deploying). +10. **Detect REST response casing** - The runs and reidentify-file endpoints return camelCase fields at runtime despite snake_case OpenAPI types; `detectRest.ts` parses both. Don't "simplify" it to a single casing. diff --git a/README.md b/README.md index eb8d957..bf93d01 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ A remote MCP server for connecting to a Skyflow Vault for sensitive PII data det - [Skyflow Runtime MCP](#skyflow-runtime-mcp) - [Try it out online](#try-it-out-online) + - [Available Tools](#available-tools) - [Connecting in Authenticated Mode](#connecting-in-authenticated-mode) - [Connection Contract](#connection-contract) - [Minimal Example (curl)](#minimal-example-curl) @@ -28,6 +29,8 @@ A remote MCP server for connecting to a Skyflow Vault for sensitive PII data det - [List Available Tools](#list-available-tools) - [Call the De-identify Tool](#call-the-de-identify-tool) - [Call the Re-identify Tool](#call-the-re-identify-tool) + - [De-identify a File](#de-identify-a-file) + - [Re-identify a File](#re-identify-a-file) - [Integration with Claude Desktop](#integration-with-claude-desktop) - [Local Development](#local-development) - [Remote Connection (Recommended)](#remote-connection-recommended) @@ -42,6 +45,18 @@ This remote MCP server is hosted at `https://pii-mcp.dev/mcp`. Connect using you For a concrete client example, see [Integration with Claude Desktop](#integration-with-claude-desktop). +## Available Tools + +| Tool | What it does | +|------|--------------| +| `de-identify` | Detects and replaces sensitive data in text with tokens. | +| `re-identify` | Restores original sensitive data from tokenized text. | +| `de-identify-file` | Detects and redacts sensitive data in files — images (jpg, png, bmp, tif), PDFs, Word/Excel/PowerPoint, txt, csv, json, xml, dcm, and audio (mp3, wav). Pass the file as a signed/public URL (`fileUrl`) or inline base64 (`fileDataBase64` + `fileName`). | +| `get-file-run-status` | Checks (and optionally waits for) an asynchronous file de-identification run and returns the processed file when it completes. | +| `re-identify-file` | Restores original sensitive data in a previously de-identified file (csv, doc, docx, json, txt, xls, xlsx, xml). | + +**Files and async processing**: Skyflow processes files asynchronously. `de-identify-file` waits up to `waitTimeSeconds` (default 25s) for the run to finish — small files usually complete inline. Larger files return a `runId` with `status: "IN_PROGRESS"`; call `get-file-run-status` with that `runId` (optionally with `waitSeconds`) until it reports `SUCCESS`. Because Skyflow's file endpoints only accept base64 content, files passed by URL are downloaded server-side (25 MB limit) and converted before being forwarded — signed URLs (S3/GCS presigned, etc.) work as long as they're reachable from the server. + ## Connecting in Authenticated Mode Authenticated mode forwards your own Skyflow credentials to your vault. Any MCP client that supports Streamable HTTP can connect — the server has no client-specific logic. @@ -177,7 +192,7 @@ curl -X POST "https://pii-mcp.dev/mcp" \ ### Limitations in Anonymous Mode -- **Only `de-identify` tool available** - `re-identify` returns a helpful error +- **Only `de-identify` tool available** - `re-identify` and the file tools (`de-identify-file`, `get-file-run-status`, `re-identify-file`) return a helpful error - **Tokens use entity counters** - e.g., `[EMAIL_ADDRESS_1]`, `[SSN_2]` instead of vault tokens - **Data is NOT persisted** - tokens cannot be re-identified later - **Rate limited** - 10 requests per minute per IP (configurable by server operator) @@ -255,6 +270,42 @@ 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}' ``` +### De-identify a File + +Pass a signed or public URL — the server downloads the file and forwards it to Skyflow as base64: + +```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":"de-identify-file","arguments":{"fileUrl":"https://example-bucket.s3.amazonaws.com/intake-form.pdf?X-Amz-Signature=...","entities":["name","ssn","email_address"]}},"id":4}' +``` + +Small files complete inline. If the response reports `"status": "IN_PROGRESS"`, poll with the returned `runId`: + +```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":"get-file-run-status","arguments":{"runId":"{run_id}","waitSeconds":30}},"id":5}' +``` + +Inline base64 works too (instead of `fileUrl`): pass `fileDataBase64` plus `fileName` (the extension determines the format). + +### Re-identify a File + +Restore the original values in a previously de-identified text-based file (csv, doc, docx, json, txt, xls, xlsx, xml): + +```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-file","arguments":{"fileUrl":"https://example-bucket.s3.amazonaws.com/deidentified-notes.txt?X-Amz-Signature=..."}},"id":6}' +``` + ## 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`: @@ -311,7 +362,7 @@ After updating the config: 1. Save the file 2. Restart Claude Desktop completely (quit and reopen) -3. The `de-identify` and `re-identify` tools should now be available in Claude Desktop +3. The `de-identify`, `re-identify`, `de-identify-file`, `get-file-run-status`, and `re-identify-file` tools should now be available in Claude Desktop ## Architecture diff --git a/docs/mcp-apps-ui.md b/docs/mcp-apps-ui.md index ec93ddb..4d1c3d7 100644 --- a/docs/mcp-apps-ui.md +++ b/docs/mcp-apps-ui.md @@ -11,7 +11,14 @@ together, and the shared components each app is built from. |------|--------|--------------|--------------|-------------| | `de-identify` | Registered | `ui/de-identify/` | `ui://de-identify/mcp-app.html` | `DeIdentifyResult` | | `re-identify` | Registered | `ui/re-identify/` | `ui://re-identify/mcp-app.html` | `ReIdentifyResult` | -| `de-identify_file` | UI exists, tool disabled | `ui/de-identify-file/` | _not registered_ | `DeIdentifyFileResult` | +| `de-identify-file` | Registered | `ui/de-identify-file/` | `ui://de-identify-file/mcp-app.html` | `DeIdentifyFileResult` | +| `get-file-run-status` | Registered (shares the de-identify-file app) | `ui/de-identify-file/` | `ui://de-identify-file/mcp-app.html` | `DeIdentifyFileResult` | +| `re-identify-file` | Registered | `ui/re-identify-file/` | `ui://re-identify-file/mcp-app.html` | `ReIdentifyFileResult` | + +`get-file-run-status` intentionally reuses the de-identify-file app: both +render the same run-result shape (status banner, processed file viewer, +entity gallery), so the app keys its loading state off `runId` when no +`fileName`/`fileUrl` argument is present. Registration lives in `src/server.ts` — each tool calls `registerAppResource` (to expose the HTML at a `ui://` URI) and `registerAppTool` (passing the diff --git a/package.json b/package.json index 90b34f5..3b9a9f5 100644 --- a/package.json +++ b/package.json @@ -9,7 +9,7 @@ "dev": "./dev.sh", "build": "pnpm build:ui && pnpm build:ui-imports && pnpm build:server", "build:ui-imports": "npx -y tsx scripts/generate-ui-imports.ts", - "build:ui": "cd ui && INPUT=de-identify/mcp-app.html vite build && INPUT=re-identify/mcp-app.html vite build && INPUT=de-identify-file/mcp-app.html vite build", + "build:ui": "cd ui && INPUT=de-identify/mcp-app.html vite build && INPUT=re-identify/mcp-app.html vite build && INPUT=de-identify-file/mcp-app.html vite build && INPUT=re-identify-file/mcp-app.html vite build", "build:server": "tsc", "start": "node dist/server.js", "server": "npx -y tsx src/server.ts", @@ -23,7 +23,7 @@ "@modelcontextprotocol/sdk": "^1.27.1", "dotenv": "^17.2.3", "express": "^5.1.0", - "skyflow-node": "^2.0.0", + "skyflow-node": "^2.1.2", "zod": "^3.25.76" }, "devDependencies": { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 0670dc5..157d18e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -21,8 +21,8 @@ importers: specifier: ^5.1.0 version: 5.1.0 skyflow-node: - specifier: ^2.0.0 - version: 2.0.0 + specifier: ^2.1.2 + version: 2.1.2 zod: specifier: ^3.25.76 version: 3.25.76 @@ -64,10 +64,6 @@ packages: engines: {node: '>=6.0.0'} hasBin: true - '@babel/runtime@7.28.4': - resolution: {integrity: sha512-Q/N6JNWvIvPnLDvjlE1OUBLPQHH6l3CltCEsHIujp45zQUSSh8K+gHnaEX45yAT1nyngnINhvWtzN+Nb9D8RAQ==} - engines: {node: '>=6.9.0'} - '@babel/types@7.29.0': resolution: {integrity: sha512-LwdZHpScM4Qz8Xw2iKSzS+cfglZzJGvofQICy7W7v4caru4EaAmyUuO6BGrbyQ2mYV11W0U8j5mBhd14dd3B0A==} engines: {node: '>=6.9.0'} @@ -833,8 +829,8 @@ packages: resolution: {integrity: sha512-G6NsmEW15s0Uw9XnCg+33H3ViYRyiM0hMrMhhqQOR8NFc5GhYrI+6I3u7OTw7b91J2g8rtvMBZJDbcGb2YUniw==} engines: {node: '>= 18'} - form-data@4.0.4: - resolution: {integrity: sha512-KrGhL9Q4zjj0kiUt5OO4Mr/A/jlI2jDYs5eHBpYHPcBEVSiipAvn2Ko2HnPe20rmcuuvMHNdZFp+4IlGTMF0Ow==} + form-data@4.0.6: + resolution: {integrity: sha512-vKatAh4SlVfgbv+YtmhiRjhEMJsYpsG1Y2rMQtR+SVSbytsSD1YGzDIcrAJmdFec88u/+VoGmxnl+80gL1tRCQ==} engines: {node: '>= 6'} formdata-node@6.0.3: @@ -885,6 +881,10 @@ packages: resolution: {integrity: sha512-0hJU9SCPvmMzIBdZFqNPXWa6dqh7WdH0cII9y+CyS8rG3nL48Bclra9HmKhVVUHyPWNH5Y7xDwAB7bfgSjkUMQ==} engines: {node: '>= 0.4'} + hasown@2.0.4: + resolution: {integrity: sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==} + engines: {node: '>= 0.4'} + hono@4.12.9: resolution: {integrity: sha512-wy3T8Zm2bsEvxKZM5w21VdHDDcwVS1yUFFY6i8UobSsKfFceT7TOwhbhfKsDyx7tYQlmRM5FLpIuYvNFyjctiA==} engines: {node: '>=16.9.0'} @@ -955,18 +955,19 @@ packages: json-schema-typed@8.0.2: resolution: {integrity: sha512-fQhoXdcvc3V28x7C7BMs4P5+kNlgUURe2jmUT1T//oBRMDrqy1QPelJimwZGo7Hg9VPV3EQV5Bnq4hbFy2vetA==} - jsonwebtoken@9.0.2: - resolution: {integrity: sha512-PRp66vJ865SSqOlgqS8hujT5U4AOgMfhrwYIuIhfKaoSCZcirrmASQr8CX7cUg+RMih+hgznrjp99o+W4pJLHQ==} + jsonwebtoken@9.0.3: + resolution: {integrity: sha512-MT/xP0CrubFRNLNKvxJ2BYfy53Zkm++5bX9dtuPbqAeQpTVe0MQTFhao8+Cp//EmJp244xt6Drw/GVEGCUj40g==} engines: {node: '>=12', npm: '>=6'} - jwa@1.4.2: - resolution: {integrity: sha512-eeH5JO+21J78qMvTIDdBXidBd6nG2kZjg5Ohz/1fpa28Z4CcsWUzJ1ZZyFq/3z3N17aZy+ZuBoHljASbL1WfOw==} + jwa@2.0.1: + resolution: {integrity: sha512-hRF04fqJIP8Abbkq5NKGN0Bbr3JxlQ+qhZufXVr0DvujKy93ZCbXZMHDL4EOtodSbCWxOqR8MS1tXA5hwqCXDg==} - jws@3.2.2: - resolution: {integrity: sha512-YHlZCB6lMTllWDtSPHz/ZXTsi8S00usEV6v1tjq8tOUZzw7DpSDWVXjXDre6ed1w/pd495ODpHZYSdkRTsa0HA==} + jws@4.0.1: + resolution: {integrity: sha512-EKI/M/yqPncGUUh44xz0PxSidXFr/+r0pA70+gIYhjv+et7yxM+s29Y+VGDkovRofQem0fs7Uvf4+YmAdyRduA==} - jwt-decode@2.2.0: - resolution: {integrity: sha512-86GgN2vzfUu7m9Wcj63iUkuDzFNYFVmjeDm2GzWpUk+opB0pEpMsw6ePCMrhYkumz2C1ihqtZzOMAg7FiXcNoQ==} + jwt-decode@4.0.0: + resolution: {integrity: sha512-+KJGIyHgkGuIq3IEBNftfhW/LfWhXUIY6OmyVWjliu5KH1y0fw7VQ8YndE2O4qZdMSd9SqbnC8GOcZEy0Om7sA==} + engines: {node: '>=18'} lightningcss-android-arm64@1.32.0: resolution: {integrity: sha512-YK7/ClTt4kAK0vo6w3X+Pnm0D2cf2vPHbhOXdoNti1Ga0al1P4TBZhwjATvjNwLEBCnKvjJc2jQgHXH0NEwlAg==} @@ -1197,6 +1198,10 @@ packages: resolution: {integrity: sha512-mAZTtNCeetKMH+pSjrb76NAM8V9a05I9aBZOHztWy/UqcJdQYNsf59vrRKWnojAT9Y+GbIvoTBC++CPHqpDBhQ==} engines: {node: '>=0.6'} + qs@6.15.2: + resolution: {integrity: sha512-Rzq0KEyX/w/tEybncDgdkZrJgVUsUMk3xjh3t5bv3S1HTAtg+uOYt72+ZfwiQwKdysThkTBdL/rTi6HDmX9Ddw==} + engines: {node: '>=0.6'} + range-parser@1.2.1: resolution: {integrity: sha512-Hrgsx+orqoygnmhFbKaHE6c296J+HTAQXoxEF6gNupROmmGJRoyzfG3ccAveqCBrwr/2yxQ5BVd/GTl5agOwSg==} engines: {node: '>= 0.6'} @@ -1276,8 +1281,8 @@ packages: siginfo@2.0.0: resolution: {integrity: sha512-ybx0WO1/8bSBLEWXZvEd7gMW3Sn3JFlW3TvX1nREbDLRNQNaeNN8WK0meBwPdAaOI7TtRRRJn/Es1zhrrCHu7g==} - skyflow-node@2.0.0: - resolution: {integrity: sha512-ZYG7pjEBkUuXRHioApmG8/UjJmR1FtlvzLs8FpKbePKbA8cTxv/QzznG4uK9XsnlwKgbl8ccwaKrteXfhT1Nkw==} + skyflow-node@2.1.2: + resolution: {integrity: sha512-cpwYzw2jbC8HzEKi8baGk0zpOl1VJ0b8rzVjPYd7sf4+EXanACvD//dV1iAfa0XiJsgkbvkFHRuS/KLrtPkhYA==} source-map-js@1.2.1: resolution: {integrity: sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==} @@ -1516,8 +1521,6 @@ snapshots: dependencies: '@babel/types': 7.29.0 - '@babel/runtime@7.28.4': {} - '@babel/types@7.29.0': dependencies: '@babel/helper-string-parser': 7.27.1 @@ -2063,7 +2066,7 @@ snapshots: es-errors: 1.3.0 get-intrinsic: 1.3.0 has-tostringtag: 1.0.2 - hasown: 2.0.2 + hasown: 2.0.4 esbuild@0.25.12: optionalDependencies: @@ -2209,12 +2212,12 @@ snapshots: form-data-encoder@4.1.0: {} - form-data@4.0.4: + form-data@4.0.6: dependencies: asynckit: 0.4.0 combined-stream: 1.0.8 es-set-tostringtag: 2.1.0 - hasown: 2.0.2 + hasown: 2.0.4 mime-types: 2.1.35 formdata-node@6.0.3: {} @@ -2260,6 +2263,10 @@ snapshots: dependencies: function-bind: 1.1.2 + hasown@2.0.4: + dependencies: + function-bind: 1.1.2 + hono@4.12.9: {} html-escaper@2.0.2: {} @@ -2317,9 +2324,9 @@ snapshots: json-schema-typed@8.0.2: {} - jsonwebtoken@9.0.2: + jsonwebtoken@9.0.3: dependencies: - jws: 3.2.2 + jws: 4.0.1 lodash.includes: 4.3.0 lodash.isboolean: 3.0.3 lodash.isinteger: 4.0.4 @@ -2330,18 +2337,18 @@ snapshots: ms: 2.1.3 semver: 7.7.3 - jwa@1.4.2: + jwa@2.0.1: dependencies: buffer-equal-constant-time: 1.0.1 ecdsa-sig-formatter: 1.0.11 safe-buffer: 5.2.1 - jws@3.2.2: + jws@4.0.1: dependencies: - jwa: 1.4.2 + jwa: 2.0.1 safe-buffer: 5.2.1 - jwt-decode@2.2.0: {} + jwt-decode@4.0.0: {} lightningcss-android-arm64@1.32.0: optional: true @@ -2510,6 +2517,10 @@ snapshots: dependencies: side-channel: 1.1.0 + qs@6.15.2: + dependencies: + side-channel: 1.1.0 + range-parser@1.2.1: {} raw-body@3.0.1: @@ -2657,18 +2668,17 @@ snapshots: siginfo@2.0.0: {} - skyflow-node@2.0.0: + skyflow-node@2.1.2: dependencies: - '@babel/runtime': 7.28.4 dotenv: 16.6.1 - form-data: 4.0.4 + form-data: 4.0.6 form-data-encoder: 4.1.0 formdata-node: 6.0.3 js-base64: 3.7.7 - jsonwebtoken: 9.0.2 - jwt-decode: 2.2.0 + jsonwebtoken: 9.0.3 + jwt-decode: 4.0.0 node-fetch: 2.7.0 - qs: 6.14.0 + qs: 6.15.2 readable-stream: 4.7.0 url-join: 4.0.1 transitivePeerDependencies: diff --git a/scripts/generate-ui-imports.ts b/scripts/generate-ui-imports.ts index 50ed01e..88cdba4 100644 --- a/scripts/generate-ui-imports.ts +++ b/scripts/generate-ui-imports.ts @@ -16,6 +16,7 @@ const tools = [ { varName: "deIdentifyHtml", dir: "de-identify" }, { varName: "reIdentifyHtml", dir: "re-identify" }, { varName: "deIdentifyFileHtml", dir: "de-identify-file" }, + { varName: "reIdentifyFileHtml", dir: "re-identify-file" }, ] as const; function escapeForTemplateLiteral(s: string): string { diff --git a/src/generated/ui-html.d.ts b/src/generated/ui-html.d.ts index 2770d19..d9829b3 100644 --- a/src/generated/ui-html.d.ts +++ b/src/generated/ui-html.d.ts @@ -3,3 +3,4 @@ export declare const deIdentifyHtml: string; export declare const reIdentifyHtml: string; export declare const deIdentifyFileHtml: string; +export declare const reIdentifyFileHtml: string; diff --git a/src/lib/detect/detectRest.ts b/src/lib/detect/detectRest.ts new file mode 100644 index 0000000..f44a31f --- /dev/null +++ b/src/lib/detect/detectRest.ts @@ -0,0 +1,244 @@ +/** + * Minimal REST client for the Skyflow Detect endpoints that skyflow-node does + * not expose through its high-level API: + * + * - GET /v1/detect/runs/{run_id} (async de-identify run status/result) + * - POST /v1/detect/reidentify/file (synchronous file re-identification) + * + * Requests use the same credential the SDK would send: the bearer value from + * the Authorization header or the API key, both forwarded as + * `Authorization: Bearer `. + * + * NOTE on response casing: Skyflow's generated OpenAPI types describe these + * responses in snake_case, but the live API returns camelCase for several + * fields (the official SDKs read both). Parsers here accept either casing. + */ + +/** Context needed to call the Detect REST API for the current request. */ +export interface DetectRestContext { + /** Vault base URL, e.g. https://abc123.vault.skyflowapis.com */ + vaultUrl: string; + vaultId: string; + /** Bearer value: the caller's JWT or API key. Never log this. */ + credentialKey: string; +} + +/** Error carrying HTTP status + response payload from the Detect API. */ +export class DetectRestError extends Error { + httpCode?: number; + details?: unknown; + + constructor(message: string, httpCode?: number, details?: unknown) { + super(message); + this.name = "DetectRestError"; + this.httpCode = httpCode; + this.details = details; + } +} + +/** One output artifact from a detect run (processed file or entity crop). */ +export interface DetectRunOutputItem { + processedFile?: string; + processedFileType?: string; + processedFileExtension?: string; +} + +/** Parsed response from GET /v1/detect/runs/{run_id}. */ +export interface DetectRunResult { + status: string; + message?: string; + output: DetectRunOutputItem[]; + outputType?: string; + wordCount?: number; + charCount?: number; + sizeInKb?: number; + durationInSeconds?: number; + pageCount?: number; + slideCount?: number; +} + +/** Parsed response from POST /v1/detect/reidentify/file. */ +export interface ReidentifyFileResult { + status: string; + outputType?: string; + processedFile?: string; + processedFileType?: string; + processedFileExtension?: string; +} + +const REQUEST_TIMEOUT_MS = 65_000; + +function pick(record: Record, ...keys: string[]): T | undefined { + for (const key of keys) { + const value = record[key]; + if (value !== undefined && value !== null) { + return value as T; + } + } + return undefined; +} + +async function detectFetch( + context: DetectRestContext, + path: string, + init: { method: string; body?: unknown } +): Promise> { + const baseUrl = context.vaultUrl.replace(/\/+$/, ""); + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS); + + let payload: unknown; + let response: Response; + try { + response = await fetch(`${baseUrl}${path}`, { + method: init.method, + headers: { + Authorization: `Bearer ${context.credentialKey}`, + "Content-Type": "application/json", + Accept: "application/json", + }, + body: init.body === undefined ? undefined : JSON.stringify(init.body), + signal: controller.signal, + }); + + // Read the body inside the try so the timeout still covers a stalled body. + const text = await response.text(); + try { + payload = text ? JSON.parse(text) : {}; + } catch { + payload = { raw: text }; + } + } catch (error) { + const reason = + error instanceof Error && error.name === "AbortError" + ? `request timed out after ${REQUEST_TIMEOUT_MS / 1000}s` + : error instanceof Error + ? error.message + : "unknown network error"; + throw new DetectRestError(`Skyflow API request failed: ${reason}`); + } finally { + clearTimeout(timeout); + } + + if (!response.ok) { + const errorBody = payload as { + error?: { message?: string; details?: unknown }; + message?: string; + }; + const message = + errorBody?.error?.message ?? + errorBody?.message ?? + `Skyflow API returned HTTP ${response.status}`; + throw new DetectRestError(message, response.status, errorBody?.error?.details ?? payload); + } + + return (payload ?? {}) as Record; +} + +function parseOutputItems(rawOutput: unknown): DetectRunOutputItem[] { + if (!Array.isArray(rawOutput)) return []; + return rawOutput.map((item) => { + const record = (item ?? {}) as Record; + return { + processedFile: pick(record, "processedFile", "processed_file"), + processedFileType: pick(record, "processedFileType", "processed_file_type"), + processedFileExtension: pick( + record, + "processedFileExtension", + "processed_file_extension" + ), + }; + }); +} + +/** Fetch the current status/result of an async de-identify file run. */ +export async function getDetectRunRest( + context: DetectRestContext, + runId: string +): Promise { + const query = new URLSearchParams({ vault_id: context.vaultId }); + const data = await detectFetch( + context, + `/v1/detect/runs/${encodeURIComponent(runId)}?${query.toString()}`, + { method: "GET" } + ); + + const wordCharacterCount = pick>( + data, + "wordCharacterCount", + "word_character_count" + ); + + return { + status: (pick(data, "status") ?? "UNKNOWN").toUpperCase(), + message: pick(data, "message"), + output: parseOutputItems(data.output), + outputType: pick(data, "outputType", "output_type"), + wordCount: + pick(data, "wordCount", "word_count") ?? + (wordCharacterCount + ? pick(wordCharacterCount, "wordCount", "word_count") + : undefined), + charCount: + pick(data, "characterCount", "character_count") ?? + (wordCharacterCount + ? pick(wordCharacterCount, "characterCount", "character_count") + : undefined), + sizeInKb: pick(data, "size"), + durationInSeconds: pick(data, "duration"), + pageCount: pick(data, "pages"), + slideCount: pick(data, "slides"), + }; +} + +/** Entity format routing for re-identification. */ +export interface ReidentifyFormat { + redacted?: string[]; + masked?: string[]; + plaintext?: string[]; +} + +/** Re-identify a previously de-identified file (synchronous endpoint). */ +export async function reidentifyFileRest( + context: DetectRestContext, + file: { base64: string; dataFormat: string }, + format?: ReidentifyFormat +): Promise { + const body: Record = { + vault_id: context.vaultId, + file: { + base64: file.base64, + data_format: file.dataFormat, + }, + }; + + if (format && (format.redacted?.length || format.masked?.length || format.plaintext?.length)) { + body.format = { + ...(format.redacted?.length ? { redacted: format.redacted } : {}), + ...(format.masked?.length ? { masked: format.masked } : {}), + ...(format.plaintext?.length ? { plaintext: format.plaintext } : {}), + }; + } + + const data = await detectFetch(context, "/v1/detect/reidentify/file", { + method: "POST", + body, + }); + + const output = (pick>(data, "output") ?? {}) as Record< + string, + unknown + >; + + return { + status: (pick(data, "status") ?? "UNKNOWN").toUpperCase(), + outputType: pick(data, "outputType", "output_type"), + processedFile: pick(output, "processedFile", "processed_file"), + processedFileType: pick(output, "processedFileType", "processed_file_type"), + processedFileExtension: pick( + output, + "processedFileExtension", + "processed_file_extension" + ), + }; +} diff --git a/src/lib/files/fileSource.ts b/src/lib/files/fileSource.ts new file mode 100644 index 0000000..4e7d277 --- /dev/null +++ b/src/lib/files/fileSource.ts @@ -0,0 +1,383 @@ +import { isIP } from "node:net"; +import { + extensionFromFileName, + extensionFromMimeType, +} from "../mappings/fileFormats.js"; + +/** + * Resolves tool file inputs (a signed/public URL or inline base64) into the + * raw bytes the Skyflow file endpoints require. + * + * Skyflow's Detect file APIs only accept base64-encoded content, so URLs are + * downloaded server-side; callers base64-encode {@link ResolvedFile.buffer} + * once, at the point they build the request, to avoid extra round-trips. + */ + +/** Maximum file size accepted from a URL download (raw, decoded bytes). */ +export const MAX_DOWNLOAD_BYTES = 25 * 1024 * 1024; // 25 MB + +/** Timeout for downloading a file from a URL (covers the body read too). */ +export const DOWNLOAD_TIMEOUT_MS = 30_000; + +/** Maximum number of HTTP redirects to follow (each hop is re-validated). */ +const MAX_REDIRECTS = 5; + +/** A resolved input file ready to send to Skyflow. */ +export interface ResolvedFile { + /** Raw (decoded) file bytes. Base64-encode when building the request. */ + buffer: Buffer; + fileName: string; + /** Lowercased extension, e.g. "pdf". */ + extension: string; + /** Content-Type reported by the remote server, when downloaded from a URL. */ + contentType?: string; +} + +/** Error thrown for invalid or unfetchable file inputs. */ +export class FileSourceError extends Error { + constructor(message: string) { + super(message); + this.name = "FileSourceError"; + } +} + +/** + * Decide whether an IP address is private, loopback, link-local, or otherwise + * not a legitimate public download target (SSRF guard). + */ +function isBlockedIp(address: string): boolean { + const family = isIP(address); + if (family === 4) { + const octets = address.split(".").map(Number); + const [a, b] = octets; + if (a === 0 || a === 10 || a === 127) return true; // this-net, private, loopback + if (a === 169 && b === 254) return true; // link-local / cloud metadata + if (a === 172 && b >= 16 && b <= 31) return true; // private + if (a === 192 && b === 168) return true; // private + if (a === 100 && b >= 64 && b <= 127) return true; // carrier-grade NAT + if (a >= 224) return true; // multicast / reserved + return false; + } + if (family === 6) { + const host = address.toLowerCase(); + if (host === "::1" || host === "::") return true; + if (host.startsWith("fe80") || host.startsWith("fc") || host.startsWith("fd")) return true; + // IPv4-mapped IPv6 in dotted form, e.g. ::ffff:169.254.169.254 + const dotted = host.match(/(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})$/); + if (dotted && isIP(dotted[1]) === 4) return isBlockedIp(dotted[1]); + // IPv4-mapped IPv6 in hex form, e.g. ::ffff:a9fe:a9fe (Node normalizes to this) + const hexMapped = host.match(/^::ffff:([0-9a-f]{1,4}):([0-9a-f]{1,4})$/); + if (hexMapped) { + const hi = Number.parseInt(hexMapped[1], 16); + const lo = Number.parseInt(hexMapped[2], 16); + const ipv4 = [(hi >> 8) & 0xff, hi & 0xff, (lo >> 8) & 0xff, lo & 0xff].join("."); + return isBlockedIp(ipv4); + } + return false; + } + return false; +} + +/** + * Normalize a URL hostname to a dotted/colon IP literal if it is one in any + * common encoding (dotted-decimal, 32-bit decimal, hex, IPv6). Returns null for + * real hostnames. This catches SSRF bypasses like http://2130706433/ (decimal + * for 127.0.0.1) that a plain string blocklist misses. + */ +function ipLiteralFromHostname(hostname: string): string | null { + const host = hostname.replace(/^\[|\]$/g, ""); + if (isIP(host)) return host; + + // 32-bit decimal, e.g. 2130706433 + if (/^\d+$/.test(host)) { + const n = Number(host); + if (Number.isInteger(n) && n >= 0 && n <= 0xffffffff) { + return [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff].join("."); + } + } + + // Hex, e.g. 0x7f000001 + if (/^0x[0-9a-f]+$/i.test(host)) { + const n = Number.parseInt(host, 16); + if (Number.isInteger(n) && n >= 0 && n <= 0xffffffff) { + return [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff].join("."); + } + } + + return null; +} + +/** + * Reject hosts that must not be downloaded from. This is a best-effort guard + * against obvious SSRF targets (loopback, private ranges, cloud metadata) in + * their common literal encodings. It does NOT resolve DNS, so a public + * hostname whose A record points at an internal address is not caught here — + * rely on network egress controls for that. + */ +function assertHostAllowed(hostname: string): void { + const host = hostname.toLowerCase().replace(/^\[|\]$/g, ""); + + if ( + host === "localhost" || + host.endsWith(".localhost") || + host.endsWith(".local") || + host.endsWith(".internal") + ) { + throw new FileSourceError( + `fileUrl host "${hostname}" is not allowed. URLs must point to a publicly reachable file (e.g. a signed S3/GCS URL).` + ); + } + + const ipLiteral = ipLiteralFromHostname(host); + if (ipLiteral && isBlockedIp(ipLiteral)) { + throw new FileSourceError( + `fileUrl host "${hostname}" is not allowed. URLs must point to a publicly reachable file (e.g. a signed S3/GCS URL).` + ); + } +} + +/** Validate the URL scheme, returning the parsed URL. */ +function parseDownloadUrl(fileUrl: string): URL { + let url: URL; + try { + url = new URL(fileUrl); + } catch { + throw new FileSourceError(`fileUrl is not a valid URL: ${fileUrl}`); + } + if (url.protocol !== "https:" && url.protocol !== "http:") { + throw new FileSourceError( + `fileUrl must use http(s); got protocol "${url.protocol}"` + ); + } + return url; +} + +/** Parse a filename out of a Content-Disposition header, if present. */ +function fileNameFromContentDisposition(header: string | null): string | undefined { + if (!header) return undefined; + // RFC 5987 filename*=UTF-8''name.ext takes precedence over filename="name.ext" + const extended = header.match(/filename\*\s*=\s*(?:UTF-8'[^']*')?([^;]+)/i); + if (extended) { + try { + const decoded = decodeURIComponent(extended[1].trim().replace(/^"|"$/g, "")); + if (decoded) return decoded; + } catch { + // fall through to the plain filename parameter + } + } + const plain = header.match(/filename\s*=\s*"?([^";]+)"?/i); + return plain ? plain[1].trim() : undefined; +} + +/** Derive a filename from the URL path, e.g. ".../reports/scan.pdf?sig=..." → "scan.pdf". */ +function fileNameFromUrl(url: URL): string | undefined { + const segments = url.pathname.split("/").filter(Boolean); + const last = segments[segments.length - 1]; + if (!last) return undefined; + try { + return decodeURIComponent(last); + } catch { + return last; + } +} + +/** Read a response body into a Buffer, aborting if it exceeds the size cap. */ +async function readBodyWithCap( + response: Response, + controller: AbortController +): Promise { + const contentLength = Number(response.headers.get("content-length")); + if (Number.isFinite(contentLength) && contentLength > MAX_DOWNLOAD_BYTES) { + controller.abort(); + throw new FileSourceError( + `File is too large: ${contentLength} bytes (limit ${MAX_DOWNLOAD_BYTES}).` + ); + } + + if (!response.body) { + const arrayBuffer = await response.arrayBuffer(); + if (arrayBuffer.byteLength > MAX_DOWNLOAD_BYTES) { + throw new FileSourceError( + `File is too large: exceeds ${MAX_DOWNLOAD_BYTES} bytes.` + ); + } + return Buffer.from(arrayBuffer); + } + + // Stream so an absent/understated Content-Length can't force us to buffer an + // unbounded body into memory before the size check. + const chunks: Buffer[] = []; + let received = 0; + for await (const chunk of response.body as AsyncIterable) { + received += chunk.byteLength; + if (received > MAX_DOWNLOAD_BYTES) { + controller.abort(); + throw new FileSourceError( + `File is too large: exceeds ${MAX_DOWNLOAD_BYTES} bytes.` + ); + } + chunks.push(Buffer.from(chunk)); + } + return Buffer.concat(chunks, received); +} + +/** + * Download a file from a signed or public URL, enforcing scheme, host, size, + * and timeout limits. Redirects are followed manually so each hop's host is + * re-validated (an initial-host-only check is bypassable via a redirect to an + * internal address). Returns the raw bytes plus name hints from the response. + */ +export async function downloadFileFromUrl(fileUrl: string): Promise<{ + buffer: Buffer; + fileName?: string; + contentType?: string; +}> { + let url = parseDownloadUrl(fileUrl); + + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), DOWNLOAD_TIMEOUT_MS); + + try { + for (let redirects = 0; ; redirects++) { + assertHostAllowed(url.hostname); + + let response: Response; + try { + response = await fetch(url, { + signal: controller.signal, + redirect: "manual", + }); + } catch (error) { + const reason = + error instanceof Error && error.name === "AbortError" + ? `download timed out after ${DOWNLOAD_TIMEOUT_MS / 1000}s` + : error instanceof Error + ? error.message + : "unknown network error"; + throw new FileSourceError(`Failed to download fileUrl: ${reason}`); + } + + // Manual redirect handling: re-validate the destination host each hop. + if (response.status >= 300 && response.status < 400) { + const location = response.headers.get("location"); + if (!location) { + throw new FileSourceError( + `Failed to download fileUrl: server responded with HTTP ${response.status} but no redirect location.` + ); + } + if (redirects >= MAX_REDIRECTS) { + throw new FileSourceError( + `Failed to download fileUrl: too many redirects (>${MAX_REDIRECTS}).` + ); + } + try { + url = new URL(location, url); + } catch { + throw new FileSourceError( + `Failed to download fileUrl: invalid redirect location "${location}".` + ); + } + if (url.protocol !== "https:" && url.protocol !== "http:") { + throw new FileSourceError( + `Failed to download fileUrl: redirect to unsupported protocol "${url.protocol}".` + ); + } + continue; + } + + if (!response.ok) { + throw new FileSourceError( + `Failed to download fileUrl: server responded with HTTP ${response.status}. ` + + `If this is a signed URL, it may have expired.` + ); + } + + const buffer = await readBodyWithCap(response, controller); + const contentType = response.headers.get("content-type") ?? undefined; + const fileName = + fileNameFromContentDisposition(response.headers.get("content-disposition")) ?? + fileNameFromUrl(url); + + return { buffer, fileName, contentType }; + } + } finally { + clearTimeout(timeout); + } +} + +/** Input accepted by {@link resolveFileInput}. */ +export interface FileInputArgs { + fileUrl?: string; + fileDataBase64?: string; + fileName?: string; +} + +/** + * Resolve a tool's file input (URL or inline base64) to raw bytes and a file + * name with a usable extension. + * + * Resolution order for the name: explicit `fileName` arg → Content-Disposition + * header → URL path → extension inferred from Content-Type. + */ +export async function resolveFileInput(args: FileInputArgs): Promise { + const { fileUrl, fileDataBase64, fileName } = args; + + if (fileUrl && fileDataBase64) { + throw new FileSourceError( + "Provide either fileUrl or fileDataBase64, not both." + ); + } + + if (fileUrl) { + const { buffer, fileName: remoteName, contentType } = + await downloadFileFromUrl(fileUrl); + + let resolvedName = fileName ?? remoteName; + let extension = resolvedName ? extensionFromFileName(resolvedName) : undefined; + + if (!extension && contentType) { + const inferred = extensionFromMimeType(contentType); + if (inferred) { + extension = inferred; + resolvedName = `${resolvedName ?? "file"}.${inferred}`; + } + } + + if (!resolvedName || !extension) { + throw new FileSourceError( + "Could not determine the file type from the URL or response headers. " + + "Pass fileName (e.g. \"report.pdf\") so the file format is known." + ); + } + + return { buffer, fileName: resolvedName, extension, contentType }; + } + + if (fileDataBase64) { + if (!fileName) { + throw new FileSourceError( + "fileName is required when passing fileDataBase64 (the extension determines the file format)." + ); + } + const extension = extensionFromFileName(fileName); + if (!extension) { + throw new FileSourceError( + `fileName "${fileName}" has no extension; the extension determines the file format.` + ); + } + const buffer = Buffer.from(fileDataBase64, "base64"); + if (buffer.byteLength === 0) { + throw new FileSourceError("fileDataBase64 decoded to an empty file."); + } + if (buffer.byteLength > MAX_DOWNLOAD_BYTES) { + throw new FileSourceError( + `File is too large: ${buffer.byteLength} bytes (limit ${MAX_DOWNLOAD_BYTES}).` + ); + } + return { buffer, fileName, extension }; + } + + throw new FileSourceError( + "A file is required: pass fileUrl (signed or public URL) or fileDataBase64." + ); +} diff --git a/src/lib/mappings/entityMaps.ts b/src/lib/mappings/entityMaps.ts index 9a11b35..1b75fbc 100644 --- a/src/lib/mappings/entityMaps.ts +++ b/src/lib/mappings/entityMaps.ts @@ -105,6 +105,18 @@ export const TRANSCRIPTION_MAP: Record = { */ export const ENTITY_KEYS = Object.keys(ENTITY_MAP) as [string, ...string[]]; +/** Tuple of valid masking method strings, derived from MASKING_METHOD_MAP. */ +export const MASKING_METHOD_KEYS = Object.keys(MASKING_METHOD_MAP) as [ + string, + ...string[], +]; + +/** Tuple of valid transcription type strings, derived from TRANSCRIPTION_MAP. */ +export const TRANSCRIPTION_KEYS = Object.keys(TRANSCRIPTION_MAP) as [ + string, + ...string[], +]; + /** * Check if an entity type is valid */ diff --git a/src/lib/mappings/fileFormats.ts b/src/lib/mappings/fileFormats.ts new file mode 100644 index 0000000..63556a2 --- /dev/null +++ b/src/lib/mappings/fileFormats.ts @@ -0,0 +1,164 @@ +/** + * File format support for Skyflow Detect file endpoints. + * + * The Skyflow API routes files to type-specific endpoints based on the file + * extension (data_format). These lists mirror the generated REST API enums in + * skyflow-node and are used for upfront validation so callers get a clear + * error instead of a Skyflow 400. + */ + +/** Formats accepted by the de-identify file endpoints (all types combined). */ +export const DEIDENTIFY_FILE_FORMATS = [ + "bmp", + "csv", + "dcm", + "doc", + "docx", + "jpeg", + "jpg", + "json", + "mp3", + "pdf", + "png", + "ppt", + "pptx", + "tif", + "tiff", + "txt", + "wav", + "xls", + "xlsx", + "xml", +] as const; + +/** Formats accepted by the re-identify file endpoint. */ +export const REIDENTIFY_FILE_FORMATS = [ + "csv", + "doc", + "docx", + "json", + "txt", + "xls", + "xlsx", + "xml", +] as const; + +export type DeidentifyFileFormat = (typeof DEIDENTIFY_FILE_FORMATS)[number]; +export type ReidentifyFileFormat = (typeof REIDENTIFY_FILE_FORMATS)[number]; + +/** Image formats — support maskingMethod, outputProcessedFile, outputOcrText. */ +export const IMAGE_FORMATS: ReadonlySet = new Set([ + "bmp", + "jpeg", + "jpg", + "png", + "tif", + "tiff", +]); + +/** Audio formats — support outputProcessedFile, outputTranscription, bleep. */ +export const AUDIO_FORMATS: ReadonlySet = new Set(["mp3", "wav"]); + +/** + * Common MIME type → file extension fallbacks, used when a downloaded file's + * URL and Content-Disposition don't reveal a usable file name. + */ +const MIME_EXTENSION_MAP: Record = { + "application/json": "json", + "application/msword": "doc", + "application/pdf": "pdf", + "application/vnd.ms-excel": "xls", + "application/vnd.ms-powerpoint": "ppt", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": + "pptx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": + "docx", + "application/xml": "xml", + "audio/mp3": "mp3", + "audio/mpeg": "mp3", + "audio/wav": "wav", + "audio/x-wav": "wav", + "image/bmp": "bmp", + "image/jpeg": "jpg", + "image/png": "png", + "image/tiff": "tiff", + "text/csv": "csv", + "text/plain": "txt", + "text/xml": "xml", +}; + +/** Look up a file extension for a MIME type (parameters stripped, lowercased). */ +export function extensionFromMimeType(mimeType: string): string | undefined { + const normalized = mimeType.split(";")[0].trim().toLowerCase(); + return MIME_EXTENSION_MAP[normalized]; +} + +/** + * File extension → canonical MIME type, used to give processed-file outputs a + * real MIME type (the Skyflow SDK reports a category label like "redacted_image" + * in its `type` field, not a MIME type). + */ +const EXTENSION_MIME_MAP: Record = { + bmp: "image/bmp", + csv: "text/csv", + dcm: "application/dicom", + doc: "application/msword", + docx: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + jpeg: "image/jpeg", + jpg: "image/jpeg", + json: "application/json", + mp3: "audio/mpeg", + pdf: "application/pdf", + png: "image/png", + ppt: "application/vnd.ms-powerpoint", + pptx: "application/vnd.openxmlformats-officedocument.presentationml.presentation", + tif: "image/tiff", + tiff: "image/tiff", + txt: "text/plain", + wav: "audio/wav", + xls: "application/vnd.ms-excel", + xlsx: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + xml: "application/xml", +}; + +/** Look up a canonical MIME type for a file extension. */ +export function mimeTypeFromExtension(extension: string): string | undefined { + return EXTENSION_MIME_MAP[extension.toLowerCase()]; +} + +/** + * Formats for which the Skyflow SDK actually serializes `transformations` + * (date shifting). The SDK's per-type request builders omit transformations + * for image, PDF, Word document, spreadsheet, and presentation requests, so + * dateShift is silently ignored there — callers should be warned instead. + */ +export const TRANSFORMATION_SUPPORTED_FORMATS: ReadonlySet = new Set([ + "txt", + "json", + "xml", + "mp3", + "wav", + "dcm", +]); + +/** Extract the lowercased extension from a file name, or undefined if none. */ +export function extensionFromFileName(fileName: string): string | undefined { + const dotIndex = fileName.lastIndexOf("."); + if (dotIndex <= 0 || dotIndex === fileName.length - 1) { + return undefined; + } + return fileName.slice(dotIndex + 1).toLowerCase(); +} + +export function isDeidentifyFileFormat( + extension: string +): extension is DeidentifyFileFormat { + return (DEIDENTIFY_FILE_FORMATS as readonly string[]).includes(extension); +} + +export function isReidentifyFileFormat( + extension: string +): extension is ReidentifyFileFormat { + return (REIDENTIFY_FILE_FORMATS as readonly string[]).includes(extension); +} diff --git a/src/lib/tools/deIdentifyFile.ts b/src/lib/tools/deIdentifyFile.ts index a93f032..a5817bd 100644 --- a/src/lib/tools/deIdentifyFile.ts +++ b/src/lib/tools/deIdentifyFile.ts @@ -1,7 +1,11 @@ import { + Bleep, DeidentifyFileOptions, DeidentifyFileRequest, SkyflowError, + TokenFormat, + TokenType, + Transformations, } from "skyflow-node"; import type { Skyflow, FileInput } from "skyflow-node"; import { @@ -9,6 +13,15 @@ import { getMaskingMethodEnum, getTranscriptionEnum, } from "../mappings/entityMaps.js"; +import { + AUDIO_FORMATS, + DEIDENTIFY_FILE_FORMATS, + IMAGE_FORMATS, + TRANSFORMATION_SUPPORTED_FORMATS, + isDeidentifyFileFormat, + mimeTypeFromExtension, +} from "../mappings/fileFormats.js"; +import { FileSourceError, resolveFileInput } from "../files/fileSource.js"; import type { DeIdentifyFileArgs, DeIdentifyFileOutput, @@ -18,12 +31,41 @@ import type { ToolResult, } from "./types.js"; -/** Default maximum wait time for file de-identification operations (in seconds) */ -export const DEFAULT_MAX_WAIT_TIME_SECONDS = 64; +/** + * Default bounded wait for the initial de-identification call. Small files + * usually finish within this window; larger ones return a runId to poll via + * the get-file-run-status tool. + */ +export const DEFAULT_WAIT_TIME_SECONDS = 25; + +/** Maximum wait the Skyflow SDK allows for a single de-identify file call. */ +export const MAX_WAIT_TIME_SECONDS = 64; + +/** Note attached to responses for runs that are still processing. */ +export function stillProcessingNote(runId: string): string { + return ( + `File de-identification is still processing (runId: ${runId}). ` + + `Call the get-file-run-status tool with this runId to check progress and retrieve the result. ` + + `Pass waitSeconds (e.g. 30) to wait server-side for completion.` + ); +} + +/** Token type strings supported by the file endpoints (vault tokens are text-only). */ +const FILE_TOKEN_TYPE_MAP: Record = { + entity_unique_counter: TokenType.ENTITY_UNIQUE_COUNTER, + entity_only: TokenType.ENTITY_ONLY, +}; + +export const FILE_TOKEN_TYPE_KEYS = Object.keys(FILE_TOKEN_TYPE_MAP) as [ + string, + ...string[], +]; /** - * Handle the de-identify_file tool logic. - * Processes files to detect and redact sensitive information. + * Handle the de-identify-file tool logic. + * Accepts a file as a signed/public URL or inline base64, forwards it to + * Skyflow Detect (which requires base64), and waits a bounded amount of time + * for the asynchronous run to finish before handing back a runId to poll. */ export async function handleDeIdentifyFile( args: DeIdentifyFileArgs, @@ -34,7 +76,7 @@ export async function handleDeIdentifyFile( if (anonymousMode) { return { output: { - error: "de-identify_file is not available in anonymous mode", + error: "de-identify-file is not available in anonymous mode", anonymousModeRestricted: true, message: "File deidentification requires authenticated access for secure processing. " + @@ -52,26 +94,51 @@ export async function handleDeIdentifyFile( try { const { - fileData, + fileUrl, + fileDataBase64, fileName, mimeType, entities, + allowRegexList, + restrictRegexList, + tokenType, maskingMethod, outputProcessedFile, outputOcrText, outputTranscription, pixelDensity, maxResolution, - waitTime, + dateShift, + bleep, + waitTimeSeconds, } = args; - // Decode base64 to buffer - const buffer = Buffer.from(fileData, "base64"); + // Resolve URL or base64 input into the base64 payload Skyflow requires + const resolved = await resolveFileInput({ fileUrl, fileDataBase64, fileName }); - // Create a File object from the buffer - const file = new File([buffer], fileName, { type: mimeType }); + if (!isDeidentifyFileFormat(resolved.extension)) { + return { + output: { + error: true, + message: + `Unsupported file format ".${resolved.extension}". ` + + `Supported formats: ${DEIDENTIFY_FILE_FORMATS.join(", ")}.`, + }, + isError: true, + }; + } + + // The SDK routes to type-specific endpoints by extension, so give the File + // a name ending in the resolved lowercase extension the router recognizes. + const baseName = resolved.fileName.includes(".") + ? resolved.fileName.slice(0, resolved.fileName.lastIndexOf(".")) + : resolved.fileName; + const normalizedName = `${baseName}.${resolved.extension}`; + const effectiveMimeType = mimeType ?? resolved.contentType; + const file = new File([resolved.buffer], normalizedName, { + type: effectiveMimeType, + }); - // Construct the file input const fileInput: FileInput = { file: file }; const fileReq = new DeidentifyFileRequest(fileInput); @@ -84,21 +151,47 @@ export async function handleDeIdentifyFile( options.setEntities(entityEnums); } + if (allowRegexList && allowRegexList.length > 0) { + options.setAllowRegexList(allowRegexList); + } + + if (restrictRegexList && restrictRegexList.length > 0) { + options.setRestrictRegexList(restrictRegexList); + } + + if (tokenType) { + const tokenTypeEnum = FILE_TOKEN_TYPE_MAP[tokenType]; + if (!tokenTypeEnum) { + return { + output: { + error: true, + message: `Invalid tokenType "${tokenType}". Supported values: ${FILE_TOKEN_TYPE_KEYS.join(", ")}.`, + }, + isError: true, + }; + } + const tokenFormat = new TokenFormat(); + tokenFormat.setDefault(tokenTypeEnum); + options.setTokenFormat(tokenFormat); + } + // Set masking method for images - use type-safe mapping if (maskingMethod) { options.setMaskingMethod(getMaskingMethodEnum(maskingMethod)); } - // Set output options + // Set output options; processed-file output is only supported for images and audio const warnings: string[] = []; + const isImage = IMAGE_FORMATS.has(resolved.extension); + const isAudio = AUDIO_FORMATS.has(resolved.extension); if (outputProcessedFile !== undefined) { - if (mimeType?.startsWith("image/")) { + if (isImage) { options.setOutputProcessedImage(outputProcessedFile); - } else if (mimeType?.startsWith("audio/")) { + } else if (isAudio) { options.setOutputProcessedAudio(outputProcessedFile); } else { warnings.push( - `outputProcessedFile is not yet supported for ${mimeType ?? "unknown"} files. It currently only applies to image/* and audio/* types.` + `outputProcessedFile is not yet supported for .${resolved.extension} files. It currently only applies to image and audio formats.` ); } } @@ -119,8 +212,38 @@ export async function handleDeIdentifyFile( options.setMaxResolution(maxResolution); } - // Set wait time (default to max, or use provided value) - options.setWaitTime(waitTime || DEFAULT_MAX_WAIT_TIME_SECONDS); + if (dateShift) { + if (TRANSFORMATION_SUPPORTED_FORMATS.has(resolved.extension)) { + const transformations = new Transformations(); + transformations.setShiftDays({ + min: dateShift.minDays, + max: dateShift.maxDays, + entities: dateShift.entities.map((e) => getEntityEnum(e)), + }); + options.setTransformations(transformations); + } else { + warnings.push( + `dateShift is ignored for .${resolved.extension} files. Date shifting currently applies only to ${[...TRANSFORMATION_SUPPORTED_FORMATS].join(", ")} formats; detected dates in this file will be tokenized instead.` + ); + } + } + + if (bleep) { + const bleepOptions = new Bleep(); + if (bleep.gain !== undefined) bleepOptions.setGain(bleep.gain); + if (bleep.frequency !== undefined) bleepOptions.setFrequency(bleep.frequency); + if (bleep.startPadding !== undefined) bleepOptions.setStartPadding(bleep.startPadding); + if (bleep.stopPadding !== undefined) bleepOptions.setStopPadding(bleep.stopPadding); + options.setBleep(bleepOptions); + } + + // Bounded wait: the run continues at Skyflow if it doesn't finish in time, + // and the response then carries only runId + status for later polling. + const waitTime = Math.min( + Math.max(waitTimeSeconds ?? DEFAULT_WAIT_TIME_SECONDS, 1), + MAX_WAIT_TIME_SECONDS + ); + options.setWaitTime(waitTime); const response = await skyflow .detect(vaultId) @@ -128,21 +251,26 @@ export async function handleDeIdentifyFile( // Prepare the output with proper typing const output: DeIdentifyFileOutput = { - inputFileName: fileName, - inputMimeType: mimeType, + inputFileName: resolved.fileName, + inputMimeType: effectiveMimeType, }; - if (response.fileBase64) { - output.processedFileData = response.fileBase64; + if (fileUrl) { + output.inputFileUrl = fileUrl; } - if (response.type) { - output.mimeType = response.type; + if (response.fileBase64) { + output.processedFileData = response.fileBase64; } - if (response.extension) { - output.extension = response.extension; + // response.type is a Skyflow category label (e.g. "redacted_image"), not a + // MIME type — derive a real MIME from the processed file's extension so UIs + // can render/download it correctly. + const processedExtension = response.extension || resolved.extension; + if (processedExtension) { + output.extension = processedExtension; } + output.mimeType = mimeTypeFromExtension(processedExtension) ?? effectiveMimeType; if (response.entities && response.entities.length > 0) { output.detectedEntities = response.entities.map( @@ -153,33 +281,39 @@ export async function handleDeIdentifyFile( ); } - if (response.wordCount !== undefined) { + // The SDK defaults these counts to 0 when absent, so only surface positive + // values (0 carries no useful signal and clutters format-inapplicable cards). + if (response.wordCount) { output.wordCount = response.wordCount; } - if (response.charCount !== undefined) { + if (response.charCount) { output.charCount = response.charCount; } - if (response.sizeInKb !== undefined) { + if (response.sizeInKb) { output.sizeInKb = response.sizeInKb; } - if (response.durationInSeconds !== undefined) { + if (response.durationInSeconds) { output.durationInSeconds = response.durationInSeconds; } - if (response.pageCount !== undefined) { + if (response.pageCount) { output.pageCount = response.pageCount; } - if (response.slideCount !== undefined) { + if (response.slideCount) { output.slideCount = response.slideCount; } if (response.runId) { output.runId = response.runId; output.status = response.status; + + if (response.status?.toUpperCase() === "IN_PROGRESS") { + output.note = stillProcessingNote(response.runId); + } } if (warnings.length > 0) { @@ -188,7 +322,15 @@ export async function handleDeIdentifyFile( return { output }; } catch (error) { - if (error instanceof SkyflowError) { + if (error instanceof FileSourceError) { + return { + output: { + error: true, + message: error.message, + }, + isError: true, + }; + } else if (error instanceof SkyflowError) { return { output: { error: true, diff --git a/src/lib/tools/getFileRunStatus.ts b/src/lib/tools/getFileRunStatus.ts new file mode 100644 index 0000000..89b749e --- /dev/null +++ b/src/lib/tools/getFileRunStatus.ts @@ -0,0 +1,157 @@ +import { + getDetectRunRest, + DetectRestError, + type DetectRestContext, + type DetectRunResult, +} from "../detect/detectRest.js"; +import { mimeTypeFromExtension } from "../mappings/fileFormats.js"; +import { stillProcessingNote } from "./deIdentifyFile.js"; +import type { + GetFileRunStatusArgs, + DeIdentifyFileOutput, + GetFileRunStatusErrorOutput, + AnonymousModeError, + ToolResult, +} from "./types.js"; + +/** Longest server-side wait for a single status call (stays under gateway timeouts). */ +export const MAX_STATUS_WAIT_SECONDS = 55; + +/** The processed-file artifact type that carries detected entity crops. */ +const ENTITIES_OUTPUT_TYPE = "entities"; + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +/** Map a detect run result onto the shared file-tool output shape. */ +function toOutput(runId: string, run: DetectRunResult): DeIdentifyFileOutput { + const output: DeIdentifyFileOutput = { + runId, + status: run.status, + }; + + const processed = run.output.find( + (item) => item.processedFile && item.processedFileType !== ENTITIES_OUTPUT_TYPE + ); + if (processed) { + output.processedFileData = processed.processedFile; + if (processed.processedFileExtension) { + output.extension = processed.processedFileExtension; + // processedFileType is a Skyflow category label, not a MIME type — derive + // a real MIME from the extension so the UI can render/download correctly. + const mime = mimeTypeFromExtension(processed.processedFileExtension); + if (mime) output.mimeType = mime; + } + } + + const entities = run.output.filter( + (item) => item.processedFileType === ENTITIES_OUTPUT_TYPE && item.processedFile + ); + if (entities.length > 0) { + output.detectedEntities = entities.map((item) => ({ + file: item.processedFile as string, + extension: item.processedFileExtension ?? "", + })); + } + + if (run.wordCount !== undefined) output.wordCount = run.wordCount; + if (run.charCount !== undefined) output.charCount = run.charCount; + if (run.sizeInKb !== undefined) output.sizeInKb = run.sizeInKb; + if (run.durationInSeconds !== undefined) output.durationInSeconds = run.durationInSeconds; + if (run.pageCount !== undefined) output.pageCount = run.pageCount; + if (run.slideCount !== undefined) output.slideCount = run.slideCount; + if (run.message) output.message = run.message; + + return output; +} + +/** + * Handle the get-file-run-status tool logic. + * Checks (and optionally waits for) the status of an asynchronous file + * de-identification run started by the de-identify-file tool. + */ +export async function handleGetFileRunStatus( + args: GetFileRunStatusArgs, + context: DetectRestContext, + anonymousMode: boolean +): Promise> { + if (anonymousMode) { + return { + output: { + error: "get-file-run-status is not available in anonymous mode", + anonymousModeRestricted: true, + message: + "File deidentification runs require authenticated access. " + + "Configure your Skyflow credentials via the Authorization header " + + "('Bearer ') or the apiKey query parameter.", + helpUrl: "https://docs.skyflow.com/", + alternativeTool: "de-identify", + }, + isError: true, + }; + } + + const { runId } = args; + const waitSeconds = Math.min( + Math.max(args.waitSeconds ?? 0, 0), + MAX_STATUS_WAIT_SECONDS + ); + + try { + const deadline = Date.now() + waitSeconds * 1000; + let pollDelayMs = 2000; + let run = await getDetectRunRest(context, runId); + + while (run.status === "IN_PROGRESS" && Date.now() < deadline) { + const remaining = deadline - Date.now(); + await sleep(Math.min(pollDelayMs, remaining)); + pollDelayMs = Math.min(pollDelayMs * 2, 8000); + run = await getDetectRunRest(context, runId); + } + + if (run.status === "FAILED") { + return { + output: { + error: true, + message: run.message + ? `File de-identification run ${runId} failed: ${run.message}` + : `File de-identification run ${runId} failed.`, + details: { runId, status: run.status }, + }, + isError: true, + }; + } + + const output = toOutput(runId, run); + + // Any non-terminal or unknown status (IN_PROGRESS, UNKNOWN, etc.) means the + // result isn't ready — attach the polling note rather than returning a + // success-shaped response with no processed file. + if (run.status !== "SUCCESS") { + output.note = stillProcessingNote(runId); + } + + return { output }; + } catch (error) { + if (error instanceof DetectRestError) { + return { + output: { + error: true, + code: error.httpCode, + message: error.message, + details: error.details, + }, + isError: true, + }; + } + return { + output: { + error: true, + message: + error instanceof Error ? error.message : "Unknown error occurred", + }, + isError: true, + }; + } +} diff --git a/src/lib/tools/reIdentifyFile.ts b/src/lib/tools/reIdentifyFile.ts new file mode 100644 index 0000000..ce0f4ad --- /dev/null +++ b/src/lib/tools/reIdentifyFile.ts @@ -0,0 +1,143 @@ +import { + reidentifyFileRest, + DetectRestError, + type DetectRestContext, +} from "../detect/detectRest.js"; +import { FileSourceError, resolveFileInput } from "../files/fileSource.js"; +import { getEntityEnum } from "../mappings/entityMaps.js"; +import { + REIDENTIFY_FILE_FORMATS, + isReidentifyFileFormat, +} from "../mappings/fileFormats.js"; +import type { + ReIdentifyFileArgs, + ReIdentifyFileOutput, + ReIdentifyFileErrorOutput, + AnonymousModeError, + ToolResult, +} from "./types.js"; + +/** + * Handle the re-identify-file tool logic. + * Restores original sensitive data in a previously de-identified file. + * Accepts a signed/public URL or inline base64; the Skyflow endpoint is + * synchronous and returns the processed file directly. + */ +export async function handleReIdentifyFile( + args: ReIdentifyFileArgs, + context: DetectRestContext, + anonymousMode: boolean +): Promise> { + if (anonymousMode) { + return { + output: { + error: "re-identify-file is not available in anonymous mode", + anonymousModeRestricted: true, + message: + "The re-identify-file tool requires authenticated access to restore sensitive data from vault tokens. " + + "To use this feature, configure your Skyflow credentials:\n\n" + + "1. Get your API key from the Skyflow dashboard\n" + + "2. Add via Authorization header: 'Bearer '\n" + + " Or via query parameter: '?apiKey='", + helpUrl: "https://docs.skyflow.com/", + }, + isError: true, + }; + } + + try { + const { + fileUrl, + fileDataBase64, + fileName, + redactedEntities, + maskedEntities, + plainTextEntities, + } = args; + + const resolved = await resolveFileInput({ fileUrl, fileDataBase64, fileName }); + + if (!isReidentifyFileFormat(resolved.extension)) { + return { + output: { + error: true, + message: + `Unsupported file format ".${resolved.extension}" for re-identification. ` + + `Supported formats: ${REIDENTIFY_FILE_FORMATS.join(", ")}.`, + }, + isError: true, + }; + } + + // Validate entity names up front (throws on unknown entities) + const format = { + redacted: redactedEntities?.map((e) => getEntityEnum(e) as string), + masked: maskedEntities?.map((e) => getEntityEnum(e) as string), + plaintext: plainTextEntities?.map((e) => getEntityEnum(e) as string), + }; + + const result = await reidentifyFileRest( + context, + { base64: resolved.buffer.toString("base64"), dataFormat: resolved.extension }, + format + ); + + // Treat anything other than a SUCCESS that actually returned a file as an + // error, so a SUCCESS-without-file response never looks like an empty win. + if (result.status !== "SUCCESS" || !result.processedFile) { + return { + output: { + error: true, + message: `File re-identification did not complete (status: ${result.status}).`, + details: { status: result.status }, + }, + isError: true, + }; + } + + const output: ReIdentifyFileOutput = { + inputFileName: resolved.fileName, + status: result.status, + }; + + if (fileUrl) { + output.inputFileUrl = fileUrl; + } + + if (result.processedFile) { + output.processedFileData = result.processedFile; + } + + output.extension = result.processedFileExtension ?? resolved.extension; + + return { output }; + } catch (error) { + if (error instanceof FileSourceError) { + return { + output: { + error: true, + message: error.message, + }, + isError: true, + }; + } else if (error instanceof DetectRestError) { + return { + output: { + error: true, + code: error.httpCode, + message: error.message, + details: error.details, + }, + isError: true, + }; + } + return { + output: { + error: true, + message: + error instanceof Error ? error.message : "Unknown error occurred", + }, + isError: true, + }; + } +} diff --git a/src/lib/tools/types.ts b/src/lib/tools/types.ts index 1d5c320..9c9bb3f 100644 --- a/src/lib/tools/types.ts +++ b/src/lib/tools/types.ts @@ -40,9 +40,10 @@ export interface DetectedEntityItem { extension: string; } -/** Output from the de-identify_file tool handler */ +/** Output from the de-identify-file and get-file-run-status tool handlers */ export interface DeIdentifyFileOutput { inputFileName?: string; + inputFileUrl?: string; inputMimeType?: string; processedFileData?: string; mimeType?: string; @@ -59,6 +60,8 @@ export interface DeIdentifyFileOutput { slideCount?: number; runId?: string; status?: string; + message?: string; + note?: string; warnings?: string[]; } @@ -73,20 +76,68 @@ export interface ToolErrorOutput { export type DeIdentifyErrorOutput = ToolErrorOutput; export type ReIdentifyErrorOutput = ToolErrorOutput; export type DeIdentifyFileErrorOutput = ToolErrorOutput; +export type GetFileRunStatusErrorOutput = ToolErrorOutput; +export type ReIdentifyFileErrorOutput = ToolErrorOutput; + +/** Date-shifting transformation options for de-identification */ +export interface DateShiftArgs { + minDays: number; + maxDays: number; + entities: string[]; +} -/** Arguments for the de-identify_file tool */ +/** Audio bleep options for de-identification of audio files */ +export interface BleepArgs { + gain?: number; + frequency?: number; + startPadding?: number; + stopPadding?: number; +} + +/** Arguments for the de-identify-file tool */ export interface DeIdentifyFileArgs { - fileData: string; - fileName: string; + fileUrl?: string; + fileDataBase64?: string; + fileName?: string; mimeType?: string; entities?: string[]; + allowRegexList?: string[]; + restrictRegexList?: string[]; + tokenType?: string; maskingMethod?: string; outputProcessedFile?: boolean; outputOcrText?: boolean; outputTranscription?: string; pixelDensity?: number; maxResolution?: number; - waitTime?: number; + dateShift?: DateShiftArgs; + bleep?: BleepArgs; + waitTimeSeconds?: number; +} + +/** Arguments for the get-file-run-status tool */ +export interface GetFileRunStatusArgs { + runId: string; + waitSeconds?: number; +} + +/** Arguments for the re-identify-file tool */ +export interface ReIdentifyFileArgs { + fileUrl?: string; + fileDataBase64?: string; + fileName?: string; + redactedEntities?: string[]; + maskedEntities?: string[]; + plainTextEntities?: string[]; +} + +/** Output from the re-identify-file tool handler */ +export interface ReIdentifyFileOutput { + inputFileName?: string; + inputFileUrl?: string; + processedFileData?: string; + extension?: string; + status?: string; } /** Result wrapper for tool handlers that can return errors */ @@ -102,3 +153,31 @@ export interface ToolResult { export function toStructuredContent(output: object): Record { return output as Record; } + +/** + * Build an MCP tool result for file tools whose output can contain multi-MB + * base64 payloads. The full output is returned as `structuredContent`; the + * text `content` channel gets a compact view with the large base64 blobs + * replaced by placeholders, so the payload isn't serialized and shipped twice. + */ +export function toFileToolResult(result: ToolResult): { + content: { type: "text"; text: string }[]; + structuredContent: Record; + isError?: boolean; +} { + const output = result.output as Record; + const textView: Record = { ...output }; + + if (typeof textView.processedFileData === "string") { + textView.processedFileData = `[${textView.processedFileData.length} base64 chars omitted; see structuredContent]`; + } + if (Array.isArray(textView.detectedEntities)) { + textView.detectedEntities = `[${textView.detectedEntities.length} detected entity artifact(s); see structuredContent]`; + } + + return { + content: [{ type: "text" as const, text: JSON.stringify(textView) }], + structuredContent: output, + ...(result.isError ? { isError: true } : {}), + }; +} diff --git a/src/server.ts b/src/server.ts index 81923ed..77c8ab0 100644 --- a/src/server.ts +++ b/src/server.ts @@ -11,14 +11,34 @@ import { } from "@modelcontextprotocol/ext-apps/server"; import express, { type Express } from "express"; import { z } from "zod"; -import { deIdentifyHtml, reIdentifyHtml } from "./generated/ui-html.js"; +import { + deIdentifyHtml, + reIdentifyHtml, + deIdentifyFileHtml, + reIdentifyFileHtml, +} from "./generated/ui-html.js"; import { Skyflow } from "skyflow-node"; import { AsyncLocalStorage } from "async_hooks"; import { validateVaultConfig, looksLikePlaceholder } from "./lib/validation/vaultConfig.js"; -import { ENTITY_KEYS } from "./lib/mappings/entityMaps.js"; +import { + ENTITY_KEYS, + MASKING_METHOD_KEYS, + TRANSCRIPTION_KEYS, +} from "./lib/mappings/entityMaps.js"; import { handleDeIdentify } from "./lib/tools/deIdentify.js"; import { handleReIdentify } from "./lib/tools/reIdentify.js"; -import { toStructuredContent } from "./lib/tools/types.js"; +import { + handleDeIdentifyFile, + FILE_TOKEN_TYPE_KEYS, + MAX_WAIT_TIME_SECONDS, +} from "./lib/tools/deIdentifyFile.js"; +import { + handleGetFileRunStatus, + MAX_STATUS_WAIT_SECONDS, +} from "./lib/tools/getFileRunStatus.js"; +import { handleReIdentifyFile } from "./lib/tools/reIdentifyFile.js"; +import type { DetectRestContext } from "./lib/detect/detectRest.js"; +import { toStructuredContent, toFileToolResult } from "./lib/tools/types.js"; import { authenticateBearer } from "./lib/middleware/authenticateBearer.js"; import { createAnonymousRateLimiter, @@ -32,31 +52,57 @@ import { interface RequestContext { skyflow: Skyflow; vaultId: string; + vaultUrl: string; + /** Raw bearer value (JWT or API key) forwarded to Skyflow. Never log this. */ + credentialKey: string; isAnonymousMode: boolean; } const requestContextStorage = new AsyncLocalStorage(); /** - * Get the Skyflow instance for the current request context + * Get the full context for the current request */ -function getCurrentSkyflow(): Skyflow { +function getRequestContext(): RequestContext { const context = requestContextStorage.getStore(); if (!context) { - throw new Error("No Skyflow instance available in current request context"); + throw new Error("No request context available"); } - return context.skyflow; + return context; +} + +/** + * Get the Skyflow instance for the current request context + */ +function getCurrentSkyflow(): Skyflow { + return getRequestContext().skyflow; +} + +/** + * Get the vault ID for the current request context + */ +function getCurrentVaultId(): string { + return getRequestContext().vaultId; } /** * Check if the current request is in anonymous mode */ function isAnonymousMode(): boolean { - const context = requestContextStorage.getStore(); - if (!context) { - throw new Error("No request context available"); - } - return context.isAnonymousMode; + return getRequestContext().isAnonymousMode; +} + +/** + * Build the context used for direct Detect REST calls (run status, file + * re-identification) from the current request. + */ +function getDetectRestContext(): DetectRestContext { + const context = getRequestContext(); + return { + vaultUrl: context.vaultUrl, + vaultId: context.vaultId, + credentialKey: context.credentialKey, + }; } // Create an MCP server @@ -68,6 +114,8 @@ const server = new McpServer({ // MCP Apps: Resource URIs const DE_IDENTIFY_RESOURCE_URI = "ui://de-identify/mcp-app.html"; const RE_IDENTIFY_RESOURCE_URI = "ui://re-identify/mcp-app.html"; +const DE_IDENTIFY_FILE_RESOURCE_URI = "ui://de-identify-file/mcp-app.html"; +const RE_IDENTIFY_FILE_RESOURCE_URI = "ui://re-identify-file/mcp-app.html"; // Register UI resources for each tool registerAppResource(server, "De-identify UI", DE_IDENTIFY_RESOURCE_URI, {}, async () => ({ @@ -78,6 +126,62 @@ registerAppResource(server, "Re-identify UI", RE_IDENTIFY_RESOURCE_URI, {}, asyn contents: [{ uri: RE_IDENTIFY_RESOURCE_URI, mimeType: RESOURCE_MIME_TYPE, text: reIdentifyHtml }], })); +// Shared by de-identify-file and get-file-run-status (both render run results) +registerAppResource(server, "De-identify File UI", DE_IDENTIFY_FILE_RESOURCE_URI, {}, async () => ({ + contents: [{ uri: DE_IDENTIFY_FILE_RESOURCE_URI, mimeType: RESOURCE_MIME_TYPE, text: deIdentifyFileHtml }], +})); + +registerAppResource(server, "Re-identify File UI", RE_IDENTIFY_FILE_RESOURCE_URI, {}, async () => ({ + contents: [{ uri: RE_IDENTIFY_FILE_RESOURCE_URI, mimeType: RESOURCE_MIME_TYPE, text: reIdentifyFileHtml }], +})); + +/** + * Output schema shared by the de-identify-file and get-file-run-status tools. + * All fields are optional because the shape differs between completed runs, + * in-progress runs (runId/status only), and error responses. + */ +const fileRunOutputSchema = { + inputFileName: z.string().optional().describe("Name of the input file"), + inputFileUrl: z.string().optional().describe("URL the input file was downloaded from"), + inputMimeType: z.string().optional().describe("MIME type of the input file"), + processedFileData: z + .string() + .optional() + .describe("Base64-encoded de-identified file (present when the run has completed)"), + mimeType: z.string().optional().describe("Type of the processed output"), + extension: z.string().optional().describe("File extension of the processed output"), + detectedEntities: z + .array(z.object({ file: z.string(), extension: z.string() })) + .optional() + .describe("Detected entity artifacts (e.g. image crops of detected regions)"), + wordCount: z.number().optional(), + charCount: z.number().optional(), + sizeInKb: z.number().optional(), + durationInSeconds: z.number().optional().describe("Audio duration, for audio files"), + pageCount: z.number().optional().describe("Page count, for documents"), + slideCount: z.number().optional().describe("Slide count, for presentations"), + runId: z + .string() + .optional() + .describe("Skyflow run identifier for the asynchronous de-identification job"), + status: z + .string() + .optional() + .describe("Run status: IN_PROGRESS, SUCCESS, or FAILED"), + message: z.string().optional().describe("Status message from Skyflow, when provided"), + note: z + .string() + .optional() + .describe("Follow-up instructions, e.g. how to poll an in-progress run"), + warnings: z.array(z.string()).optional().describe("Non-fatal warnings about ignored options"), + 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"), + helpUrl: z.string().optional().describe("URL for setup documentation"), + alternativeTool: z.string().optional().describe("Suggested tool to use instead"), + code: z.number().optional().describe("HTTP error code from Skyflow API"), + details: z.unknown().optional().describe("Additional error details from Skyflow API"), +}; + /** * Skyflow De-identify Tool * Replaces sensitive information in text with placeholder tokens @@ -168,8 +272,238 @@ registerAppTool( } ); +/** + * Skyflow De-identify File Tool + * Detects and redacts sensitive information in files (images, PDFs, audio, + * documents, spreadsheets, presentations). File processing is asynchronous: + * if the run doesn't finish within the bounded wait, the response carries a + * runId to poll with the get-file-run-status tool. + */ +registerAppTool( + server, + "de-identify-file", + { + title: "Skyflow De-identify File Tool", + description: + "De-identify sensitive information in a file using Skyflow. " + + "Pass the file either as a signed/public URL (fileUrl) — the server downloads it and forwards it to Skyflow — " + + "or as base64 content (fileDataBase64 + fileName). " + + "Supports images (jpg, png, bmp, tif), PDFs, Word/Excel/PowerPoint documents, txt, csv, json, xml, dcm, and audio (mp3, wav). " + + "File processing is asynchronous: small files usually complete within the default wait and return the processed file inline; " + + "larger files return a runId with status IN_PROGRESS — call the get-file-run-status tool with that runId to retrieve the result.", + inputSchema: { + fileUrl: z + .string() + .optional() + .describe( + "Signed or public URL of the file to de-identify (e.g. an S3/GCS signed URL). The server downloads it (25 MB max) and converts it to base64 for Skyflow. Provide either fileUrl or fileDataBase64." + ), + fileDataBase64: z + .string() + .optional() + .describe("Base64-encoded file content. Provide either fileUrl or fileDataBase64."), + fileName: z + .string() + .optional() + .describe( + "File name including extension (e.g. \"report.pdf\"). Required with fileDataBase64; optional with fileUrl (inferred from the URL or response headers when omitted). The extension determines how Skyflow processes the file." + ), + mimeType: z.string().optional().describe("MIME type of the file (optional hint, e.g. \"application/pdf\")"), + entities: z + .array(z.enum(ENTITY_KEYS)) + .optional() + .describe("Specific entity types to detect. Leave empty to detect all supported entities."), + allowRegexList: z + .array(z.string()) + .optional() + .describe("Regex patterns for values that should NOT be de-identified (allowlist)"), + restrictRegexList: z + .array(z.string()) + .optional() + .describe("Regex patterns for additional values that SHOULD be de-identified (denylist)"), + tokenType: z + .enum(FILE_TOKEN_TYPE_KEYS) + .optional() + .describe( + "Token format for detected entities: entity_unique_counter (e.g. [SSN_1], default) or entity_only (e.g. [SSN])" + ), + maskingMethod: z + .enum(MASKING_METHOD_KEYS) + .optional() + .describe("How to mask detected regions in images: BLACKBOX or BLUR"), + outputProcessedFile: z + .boolean() + .optional() + .describe("Return the processed (redacted) file. Applies to image and audio files."), + outputOcrText: z + .boolean() + .optional() + .describe("Return OCR-extracted text for images"), + outputTranscription: z + .enum(TRANSCRIPTION_KEYS) + .optional() + .describe("Return a transcription for audio files: PLAINTEXT_TRANSCRIPTION or DIARIZED_TRANSCRIPTION"), + pixelDensity: z.number().positive().optional().describe("Pixel density for PDF rasterization"), + maxResolution: z.number().positive().optional().describe("Maximum resolution for PDF processing"), + dateShift: z + .object({ + minDays: z.number().int().describe("Minimum number of days to shift dates"), + maxDays: z.number().int().describe("Maximum number of days to shift dates"), + entities: z.array(z.enum(ENTITY_KEYS)).describe("Date entity types to shift (e.g. dob, date)"), + }) + .optional() + .describe("Shift detected dates by a random offset instead of tokenizing them"), + bleep: z + .object({ + gain: z.number().optional().describe("Bleep tone gain"), + frequency: z.number().optional().describe("Bleep tone frequency in Hz"), + startPadding: z.number().optional().describe("Seconds of padding before each bleep"), + stopPadding: z.number().optional().describe("Seconds of padding after each bleep"), + }) + .optional() + .describe("Bleep tone settings for redacting audio files"), + waitTimeSeconds: z + .number() + .int() + .min(1) + .max(MAX_WAIT_TIME_SECONDS) + .optional() + .describe( + `Seconds to wait for the run to complete before returning a runId to poll (1-${MAX_WAIT_TIME_SECONDS}, default 25)` + ), + }, + outputSchema: fileRunOutputSchema, + _meta: { ui: { resourceUri: DE_IDENTIFY_FILE_RESOURCE_URI } }, + }, + async (args) => { + const result = await handleDeIdentifyFile( + args, + getCurrentSkyflow(), + getCurrentVaultId(), + isAnonymousMode() + ); + return toFileToolResult(result); + } +); + +/** + * Skyflow File Run Status Tool + * Polls the status of an asynchronous file de-identification run and returns + * the processed file once the run completes. + */ +registerAppTool( + server, + "get-file-run-status", + { + title: "Skyflow File Run Status Tool", + description: + "Check the status of an asynchronous file de-identification run started by the de-identify-file tool, " + + "and retrieve the processed file when the run completes. " + + "Pass the runId returned by de-identify-file. " + + "Optionally pass waitSeconds to wait server-side for completion instead of polling repeatedly; " + + "if the run is still IN_PROGRESS after the wait, call this tool again.", + inputSchema: { + runId: z + .string() + .min(1) + .describe("Run identifier returned by the de-identify-file tool"), + waitSeconds: z + .number() + .int() + .min(0) + .max(MAX_STATUS_WAIT_SECONDS) + .optional() + .describe( + `Seconds to wait server-side for the run to complete (0-${MAX_STATUS_WAIT_SECONDS}, default 0 = single status check)` + ), + }, + outputSchema: fileRunOutputSchema, + _meta: { ui: { resourceUri: DE_IDENTIFY_FILE_RESOURCE_URI } }, + }, + async ({ runId, waitSeconds }) => { + const result = await handleGetFileRunStatus( + { runId, waitSeconds }, + getDetectRestContext(), + isAnonymousMode() + ); + return toFileToolResult(result); + } +); + +/** + * Skyflow Re-identify File Tool + * Restores original sensitive data in a previously de-identified file. + */ +registerAppTool( + server, + "re-identify-file", + { + title: "Skyflow Re-identify File Tool", + description: + "Re-identify a previously de-identified file using Skyflow, replacing tokens (like [SSN_abc123]) with the original sensitive data. " + + "Pass the file either as a signed/public URL (fileUrl) or as base64 content (fileDataBase64 + fileName). " + + "Supported formats: csv, doc, docx, json, txt, xls, xlsx, xml. " + + "Optionally control how specific entity types are restored (redacted, masked, or plaintext).", + inputSchema: { + fileUrl: z + .string() + .optional() + .describe( + "Signed or public URL of the de-identified file (25 MB max). Provide either fileUrl or fileDataBase64." + ), + fileDataBase64: z + .string() + .optional() + .describe("Base64-encoded de-identified file content. Provide either fileUrl or fileDataBase64."), + fileName: z + .string() + .optional() + .describe( + "File name including extension (e.g. \"notes.txt\"). Required with fileDataBase64; optional with fileUrl." + ), + redactedEntities: z + .array(z.enum(ENTITY_KEYS)) + .optional() + .describe("Entity types to keep redacted in the output"), + maskedEntities: z + .array(z.enum(ENTITY_KEYS)) + .optional() + .describe("Entity types to return masked (partially visible) in the output"), + plainTextEntities: z + .array(z.enum(ENTITY_KEYS)) + .optional() + .describe("Entity types to restore as plaintext in the output"), + }, + outputSchema: { + inputFileName: z.string().optional().describe("Name of the input file"), + inputFileUrl: z.string().optional().describe("URL the input file was downloaded from"), + processedFileData: z + .string() + .optional() + .describe("Base64-encoded re-identified file"), + extension: z.string().optional().describe("File extension of the processed output"), + status: z.string().optional().describe("Processing status: SUCCESS or FAILED"), + 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"), + helpUrl: z.string().optional().describe("URL for setup documentation"), + code: z.number().optional().describe("HTTP error code from Skyflow API"), + details: z.unknown().optional().describe("Additional error details from Skyflow API"), + }, + _meta: { ui: { resourceUri: RE_IDENTIFY_FILE_RESOURCE_URI } }, + }, + async (args) => { + const result = await handleReIdentifyFile( + args, + getDetectRestContext(), + isAnonymousMode() + ); + return toFileToolResult(result); + } +); + const app: Express = express(); -app.use(express.json({ limit: "5mb" })); // Limit for base64-encoded files +app.use(express.json({ limit: "25mb" })); // Limit for base64-encoded files passed inline // Serve static files from the public directory app.use(express.static("public")); @@ -249,7 +583,18 @@ app.post("/mcp", authenticateBearer, anonymousRateLimiter, async (req, res) => { } // Use validated config - const { vaultId: validatedVaultId, clusterId } = validation.config!; + const { vaultId: validatedVaultId, vaultUrl: validatedVaultUrl, clusterId } = validation.config!; + + // Normalize the vault URL for direct REST calls (scheme may be omitted) + const normalizedVaultUrl = /^https?:\/\//.test(validatedVaultUrl) + ? validatedVaultUrl + : `https://${validatedVaultUrl}`; + + // Raw bearer value (JWT or API key) for direct Detect REST calls + const credentialKey = + "token" in req.skyflowCredentials + ? req.skyflowCredentials.token + : req.skyflowCredentials.apiKey; // Create per-request Skyflow instance with credentials (bearer token or API key) let skyflowInstance: Skyflow; @@ -286,6 +631,8 @@ app.post("/mcp", authenticateBearer, anonymousRateLimiter, async (req, res) => { { skyflow: skyflowInstance, vaultId: validatedVaultId, + vaultUrl: normalizedVaultUrl, + credentialKey, isAnonymousMode: useAnonymousMode, }, async () => { diff --git a/tests/unit/files/fileSource.test.ts b/tests/unit/files/fileSource.test.ts new file mode 100644 index 0000000..0bd38c7 --- /dev/null +++ b/tests/unit/files/fileSource.test.ts @@ -0,0 +1,281 @@ +import { describe, it, expect, vi, afterEach } from "vitest"; +import { + resolveFileInput, + downloadFileFromUrl, + FileSourceError, + MAX_DOWNLOAD_BYTES, +} from "../../../src/lib/files/fileSource"; + +function stubFetch(response: Response | (() => Promise)) { + const impl = typeof response === "function" ? response : async () => response; + const mock = vi.fn(impl); + vi.stubGlobal("fetch", mock); + return mock; +} + +describe("fileSource", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + describe("resolveFileInput with base64", () => { + it("resolves base64 input with a file name", async () => { + const content = Buffer.from("hello world"); + const resolved = await resolveFileInput({ + fileDataBase64: content.toString("base64"), + fileName: "Notes.TXT", + }); + + expect(resolved.buffer.equals(content)).toBe(true); + expect(resolved.fileName).toBe("Notes.TXT"); + expect(resolved.extension).toBe("txt"); + }); + + it("rejects base64 payloads over the size limit", async () => { + const huge = "A".repeat(MAX_DOWNLOAD_BYTES + 8); + const big = Buffer.from(huge).toString("base64"); + await expect( + resolveFileInput({ fileDataBase64: big, fileName: "big.txt" }) + ).rejects.toThrow(/too large/); + }); + + it("rejects base64 input without a file name", async () => { + await expect( + resolveFileInput({ fileDataBase64: "aGVsbG8=" }) + ).rejects.toThrow(FileSourceError); + }); + + it("rejects file names without an extension", async () => { + await expect( + resolveFileInput({ fileDataBase64: "aGVsbG8=", fileName: "noext" }) + ).rejects.toThrow(/extension/); + }); + + it("rejects empty base64 payloads", async () => { + await expect( + resolveFileInput({ fileDataBase64: "!!!", fileName: "a.txt" }) + ).rejects.toThrow(/empty/); + }); + + it("rejects providing both fileUrl and fileDataBase64", async () => { + await expect( + resolveFileInput({ + fileUrl: "https://example.com/a.txt", + fileDataBase64: "aGVsbG8=", + fileName: "a.txt", + }) + ).rejects.toThrow(/not both/); + }); + + it("rejects when neither input is provided", async () => { + await expect(resolveFileInput({})).rejects.toThrow(/fileUrl/); + }); + }); + + describe("downloadFileFromUrl", () => { + it("downloads a file and infers the name from the URL path", async () => { + const bytes = Buffer.from("pdf bytes"); + stubFetch(new Response(bytes, { status: 200, headers: { "content-type": "application/pdf" } })); + + const result = await downloadFileFromUrl( + "https://bucket.s3.amazonaws.com/folder/report.pdf?X-Amz-Signature=abc" + ); + + expect(result.buffer.equals(bytes)).toBe(true); + expect(result.fileName).toBe("report.pdf"); + expect(result.contentType).toBe("application/pdf"); + }); + + it("prefers the Content-Disposition filename over the URL path", async () => { + stubFetch( + new Response(Buffer.from("x"), { + status: 200, + headers: { "content-disposition": 'attachment; filename="actual.docx"' }, + }) + ); + + const result = await downloadFileFromUrl("https://example.com/download?id=42"); + expect(result.fileName).toBe("actual.docx"); + }); + + it("rejects non-http(s) URLs", async () => { + await expect(downloadFileFromUrl("ftp://example.com/a.txt")).rejects.toThrow(/http/); + }); + + it("rejects invalid URLs", async () => { + await expect(downloadFileFromUrl("not a url")).rejects.toThrow(/not a valid URL/); + }); + + it.each([ + "https://localhost/file.txt", + "https://127.0.0.1/file.txt", + "https://10.0.0.5/file.txt", + "https://192.168.1.10/file.txt", + "https://172.20.3.4/file.txt", + "https://169.254.169.254/latest/meta-data", + "https://metadata.internal/file.txt", + "https://[::1]/file.txt", + // Alternate IP encodings for internal hosts + "https://2130706433/file.txt", // decimal for 127.0.0.1 + "https://0x7f000001/file.txt", // hex for 127.0.0.1 + "https://[::ffff:169.254.169.254]/file.txt", // IPv4-mapped IPv6 metadata + "https://100.64.0.1/file.txt", // carrier-grade NAT + ])("blocks private/internal host %s", async (url) => { + await expect(downloadFileFromUrl(url)).rejects.toThrow(/not allowed/); + }); + + it("allows public hosts", async () => { + stubFetch(new Response(Buffer.from("ok"), { status: 200 })); + const result = await downloadFileFromUrl("https://storage.googleapis.com/b/file.txt"); + expect(result.buffer.toString()).toBe("ok"); + }); + + it.each([ + "https://fcbarcelona.com/file.txt", // hostname starting with "fc" is not an IPv6 range + "https://fdn.example.com/file.txt", + "https://[fd12:3456::1]/file.txt", // but IPv6 unique-local literals are blocked + ])("only treats IPv6 literals as IPv6 ranges: %s", async (url) => { + stubFetch(new Response(Buffer.from("ok"), { status: 200 })); + if (url.includes("[fd")) { + await expect(downloadFileFromUrl(url)).rejects.toThrow(/not allowed/); + } else { + await expect(downloadFileFromUrl(url)).resolves.toBeDefined(); + } + }); + + it("rejects oversized files based on Content-Length", async () => { + stubFetch( + new Response(Buffer.from("x"), { + status: 200, + headers: { "content-length": String(MAX_DOWNLOAD_BYTES + 1) }, + }) + ); + + await expect(downloadFileFromUrl("https://example.com/huge.bin")).rejects.toThrow(/too large/); + }); + + it("surfaces HTTP error statuses with a helpful message", async () => { + stubFetch(new Response("expired", { status: 403 })); + + await expect(downloadFileFromUrl("https://example.com/file.pdf")).rejects.toThrow(/403/); + }); + + it("wraps network errors", async () => { + stubFetch(() => Promise.reject(new Error("ECONNREFUSED"))); + + await expect(downloadFileFromUrl("https://example.com/file.pdf")).rejects.toThrow( + /ECONNREFUSED/ + ); + }); + + it("follows a redirect to another public host", async () => { + const bytes = Buffer.from("redirected content"); + const fetchMock = vi.fn(async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("start.example.com")) { + return new Response(null, { + status: 302, + headers: { location: "https://cdn.example.com/real.pdf" }, + }); + } + return new Response(bytes, { + status: 200, + headers: { "content-type": "application/pdf" }, + }); + }); + vi.stubGlobal("fetch", fetchMock); + + const result = await downloadFileFromUrl("https://start.example.com/redir"); + expect(result.buffer.equals(bytes)).toBe(true); + expect(result.fileName).toBe("real.pdf"); + }); + + it("re-validates the host on each redirect hop (blocks SSRF via redirect)", async () => { + const fetchMock = vi.fn(async () => + new Response(null, { + status: 302, + headers: { location: "http://169.254.169.254/latest/meta-data/" }, + }) + ); + vi.stubGlobal("fetch", fetchMock); + + await expect( + downloadFileFromUrl("https://public.example.com/redir") + ).rejects.toThrow(/not allowed/); + }); + + it("rejects a redirect loop that exceeds the hop limit", async () => { + const fetchMock = vi.fn(async () => + new Response(null, { + status: 302, + headers: { location: "https://loop.example.com/again" }, + }) + ); + vi.stubGlobal("fetch", fetchMock); + + await expect( + downloadFileFromUrl("https://loop.example.com/again") + ).rejects.toThrow(/too many redirects/); + }); + + it("enforces the size cap while streaming when Content-Length is absent", async () => { + // A body larger than the cap, delivered with no Content-Length header. + const oversized = new Uint8Array(MAX_DOWNLOAD_BYTES + 1024); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(oversized); + controller.close(); + }, + }); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(stream, { status: 200 })) + ); + + await expect( + downloadFileFromUrl("https://example.com/huge-nolength.bin") + ).rejects.toThrow(/too large/); + }); + }); + + describe("resolveFileInput with URL", () => { + it("uses the explicit fileName over the remote name", async () => { + stubFetch(new Response(Buffer.from("data"), { status: 200 })); + + const resolved = await resolveFileInput({ + fileUrl: "https://example.com/blob", + fileName: "override.csv", + }); + + expect(resolved.fileName).toBe("override.csv"); + expect(resolved.extension).toBe("csv"); + }); + + it("falls back to Content-Type when no name is available", async () => { + stubFetch( + new Response(Buffer.from("data"), { + status: 200, + headers: { "content-type": "application/pdf" }, + }) + ); + + const resolved = await resolveFileInput({ fileUrl: "https://example.com/download" }); + + expect(resolved.extension).toBe("pdf"); + expect(resolved.fileName).toBe("download.pdf"); + }); + + it("errors when the file type cannot be determined", async () => { + stubFetch( + new Response(Buffer.from("data"), { + status: 200, + headers: { "content-type": "application/octet-stream" }, + }) + ); + + await expect( + resolveFileInput({ fileUrl: "https://example.com/download" }) + ).rejects.toThrow(/fileName/); + }); + }); +}); diff --git a/tests/unit/tools/deIdentifyFile.test.ts b/tests/unit/tools/deIdentifyFile.test.ts index 0f83388..a926ca8 100644 --- a/tests/unit/tools/deIdentifyFile.test.ts +++ b/tests/unit/tools/deIdentifyFile.test.ts @@ -1,9 +1,21 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { handleDeIdentifyFile, DEFAULT_MAX_WAIT_TIME_SECONDS } from "../../../src/lib/tools/deIdentifyFile"; -import type { DeIdentifyFileArgs, DeIdentifyFileOutput, DeIdentifyFileErrorOutput, AnonymousModeError } from "../../../src/lib/tools/types"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { + handleDeIdentifyFile, + DEFAULT_WAIT_TIME_SECONDS, + MAX_WAIT_TIME_SECONDS, +} from "../../../src/lib/tools/deIdentifyFile"; +import type { + DeIdentifyFileArgs, + DeIdentifyFileOutput, + DeIdentifyFileErrorOutput, + AnonymousModeError, +} from "../../../src/lib/tools/types"; // Mock the skyflow-node SDK const mockSetEntities = vi.fn(); +const mockSetAllowRegexList = vi.fn(); +const mockSetRestrictRegexList = vi.fn(); +const mockSetTokenFormat = vi.fn(); const mockSetMaskingMethod = vi.fn(); const mockSetOutputProcessedImage = vi.fn(); const mockSetOutputProcessedAudio = vi.fn(); @@ -11,8 +23,16 @@ const mockSetOutputOcrText = vi.fn(); const mockSetOutputTranscription = vi.fn(); const mockSetPixelDensity = vi.fn(); const mockSetMaxResolution = vi.fn(); +const mockSetTransformations = vi.fn(); +const mockSetBleep = vi.fn(); const mockSetWaitTime = vi.fn(); const mockDeidentifyFile = vi.fn(); +const mockTokenFormatSetDefault = vi.fn(); +const mockTransformationsSetShiftDays = vi.fn(); +const mockBleepSetGain = vi.fn(); +const mockBleepSetFrequency = vi.fn(); +const mockBleepSetStartPadding = vi.fn(); +const mockBleepSetStopPadding = vi.fn(); vi.mock("skyflow-node", () => { class MockSkyflowError extends Error { @@ -26,6 +46,9 @@ vi.mock("skyflow-node", () => { return { DeidentifyFileOptions: vi.fn(function (this: any) { this.setEntities = mockSetEntities; + this.setAllowRegexList = mockSetAllowRegexList; + this.setRestrictRegexList = mockSetRestrictRegexList; + this.setTokenFormat = mockSetTokenFormat; this.setMaskingMethod = mockSetMaskingMethod; this.setOutputProcessedImage = mockSetOutputProcessedImage; this.setOutputProcessedAudio = mockSetOutputProcessedAudio; @@ -33,9 +56,28 @@ vi.mock("skyflow-node", () => { this.setOutputTranscription = mockSetOutputTranscription; this.setPixelDensity = mockSetPixelDensity; this.setMaxResolution = mockSetMaxResolution; + this.setTransformations = mockSetTransformations; + this.setBleep = mockSetBleep; this.setWaitTime = mockSetWaitTime; }), DeidentifyFileRequest: vi.fn(function (this: any, fileInput: any) { this.fileInput = fileInput; }), + TokenFormat: vi.fn(function (this: any) { + this.setDefault = mockTokenFormatSetDefault; + }), + Transformations: vi.fn(function (this: any) { + this.setShiftDays = mockTransformationsSetShiftDays; + }), + Bleep: vi.fn(function (this: any) { + this.setGain = mockBleepSetGain; + this.setFrequency = mockBleepSetFrequency; + this.setStartPadding = mockBleepSetStartPadding; + this.setStopPadding = mockBleepSetStopPadding; + }), + TokenType: { + ENTITY_UNIQUE_COUNTER: "entity_unq_counter", + ENTITY_ONLY: "entity_only", + VAULT_TOKEN: "vault_token", + }, SkyflowError: MockSkyflowError, }; }); @@ -56,7 +98,7 @@ function createMockSkyflow(response: Record = {}) { } const baseArgs: DeIdentifyFileArgs = { - fileData: Buffer.from("test file content").toString("base64"), + fileDataBase64: Buffer.from("test file content").toString("base64"), fileName: "test.png", mimeType: "image/png", }; @@ -66,6 +108,10 @@ describe("handleDeIdentifyFile", () => { vi.clearAllMocks(); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + describe("anonymous mode", () => { it("should return error with anonymousModeRestricted flag", async () => { const skyflow = createMockSkyflow(); @@ -91,6 +137,119 @@ describe("handleDeIdentifyFile", () => { }); }); + describe("input validation", () => { + it("should error when neither fileUrl nor fileDataBase64 is provided", async () => { + const skyflow = createMockSkyflow(); + const result = await handleDeIdentifyFile({}, skyflow as any, "vault123", false); + const output = result.output as DeIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("fileUrl"); + expect(output.message).toContain("fileDataBase64"); + }); + + it("should error when both fileUrl and fileDataBase64 are provided", async () => { + const skyflow = createMockSkyflow(); + const result = await handleDeIdentifyFile( + { ...baseArgs, fileUrl: "https://example.com/file.png" }, + skyflow as any, + "vault123", + false + ); + + expect(result.isError).toBe(true); + expect((result.output as DeIdentifyFileErrorOutput).message).toContain("not both"); + }); + + it("should error when fileDataBase64 is provided without fileName", async () => { + const skyflow = createMockSkyflow(); + const result = await handleDeIdentifyFile( + { fileDataBase64: baseArgs.fileDataBase64 }, + skyflow as any, + "vault123", + false + ); + + expect(result.isError).toBe(true); + expect((result.output as DeIdentifyFileErrorOutput).message).toContain("fileName"); + }); + + it("should error on unsupported file formats", async () => { + const skyflow = createMockSkyflow(); + const result = await handleDeIdentifyFile( + { ...baseArgs, fileName: "video.mp4" }, + skyflow as any, + "vault123", + false + ); + const output = result.output as DeIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain('".mp4"'); + expect((skyflow as any).detect).not.toHaveBeenCalled(); + }); + + it("should error on invalid tokenType", async () => { + const skyflow = createMockSkyflow(); + const result = await handleDeIdentifyFile( + { ...baseArgs, tokenType: "vault_token" }, + skyflow as any, + "vault123", + false + ); + + expect(result.isError).toBe(true); + expect((result.output as DeIdentifyFileErrorOutput).message).toContain("tokenType"); + }); + }); + + describe("URL input", () => { + it("should download the file and pass it to Skyflow", async () => { + const fileBytes = Buffer.from("fake image bytes"); + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response(fileBytes, { + status: 200, + headers: { "content-type": "image/png" }, + })) + ); + + const skyflow = createMockSkyflow({}); + const result = await handleDeIdentifyFile( + { fileUrl: "https://bucket.s3.amazonaws.com/scan.png?sig=abc" }, + skyflow as any, + "vault123", + false + ); + const output = result.output as DeIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.inputFileName).toBe("scan.png"); + expect(output.inputFileUrl).toBe("https://bucket.s3.amazonaws.com/scan.png?sig=abc"); + expect(mockDeidentifyFile).toHaveBeenCalled(); + }); + + it("should return a clear error when the download fails", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => new Response("denied", { status: 403 })) + ); + + const skyflow = createMockSkyflow({}); + const result = await handleDeIdentifyFile( + { fileUrl: "https://bucket.s3.amazonaws.com/scan.png" }, + skyflow as any, + "vault123", + false + ); + const output = result.output as DeIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("403"); + expect(mockDeidentifyFile).not.toHaveBeenCalled(); + }); + }); + describe("authenticated mode", () => { it("should include inputFileName and inputMimeType in output", async () => { const skyflow = createMockSkyflow({}); @@ -102,10 +261,10 @@ describe("handleDeIdentifyFile", () => { expect(output.inputMimeType).toBe("image/png"); }); - it("should set outputProcessedImage for image mime types", async () => { + it("should set outputProcessedImage for image extensions", async () => { const skyflow = createMockSkyflow({}); await handleDeIdentifyFile( - { ...baseArgs, mimeType: "image/jpeg", outputProcessedFile: true }, + { ...baseArgs, fileName: "photo.jpeg", outputProcessedFile: true }, skyflow as any, "vault123", false ); @@ -113,10 +272,10 @@ describe("handleDeIdentifyFile", () => { expect(mockSetOutputProcessedAudio).not.toHaveBeenCalled(); }); - it("should set outputProcessedAudio for audio mime types", async () => { + it("should set outputProcessedAudio for audio extensions", async () => { const skyflow = createMockSkyflow({}); await handleDeIdentifyFile( - { ...baseArgs, mimeType: "audio/mp3", outputProcessedFile: true }, + { ...baseArgs, fileName: "call.mp3", outputProcessedFile: true }, skyflow as any, "vault123", false ); @@ -124,6 +283,19 @@ describe("handleDeIdentifyFile", () => { expect(mockSetOutputProcessedImage).not.toHaveBeenCalled(); }); + it("should warn when outputProcessedFile is unsupported for the format", async () => { + const skyflow = createMockSkyflow({}); + const result = await handleDeIdentifyFile( + { ...baseArgs, fileName: "doc.pdf", outputProcessedFile: true }, + skyflow as any, "vault123", false + ); + const output = result.output as DeIdentifyFileOutput; + + expect(output.warnings?.[0]).toContain(".pdf"); + expect(mockSetOutputProcessedImage).not.toHaveBeenCalled(); + expect(mockSetOutputProcessedAudio).not.toHaveBeenCalled(); + }); + it("should map entity strings to enums", async () => { const skyflow = createMockSkyflow({}); await handleDeIdentifyFile( @@ -144,23 +316,84 @@ describe("handleDeIdentifyFile", () => { expect(mockSetMaskingMethod).toHaveBeenCalledWith("ENUM_BLUR"); }); - it("should use DEFAULT_MAX_WAIT_TIME_SECONDS when waitTime not specified", async () => { + it("should set allow and restrict regex lists", async () => { + const skyflow = createMockSkyflow({}); + await handleDeIdentifyFile( + { ...baseArgs, allowRegexList: ["foo.*"], restrictRegexList: ["bar.*"] }, + skyflow as any, "vault123", false + ); + + expect(mockSetAllowRegexList).toHaveBeenCalledWith(["foo.*"]); + expect(mockSetRestrictRegexList).toHaveBeenCalledWith(["bar.*"]); + }); + + it("should configure token format from tokenType", async () => { + const skyflow = createMockSkyflow({}); + await handleDeIdentifyFile( + { ...baseArgs, tokenType: "entity_only" }, + skyflow as any, "vault123", false + ); + + expect(mockTokenFormatSetDefault).toHaveBeenCalledWith("entity_only"); + expect(mockSetTokenFormat).toHaveBeenCalled(); + }); + + it("should configure date-shift transformations for a supported format", async () => { + const skyflow = createMockSkyflow({}); + await handleDeIdentifyFile( + { ...baseArgs, fileName: "notes.txt", mimeType: "text/plain", dateShift: { minDays: 1, maxDays: 30, entities: ["dob"] } }, + skyflow as any, "vault123", false + ); + + expect(mockTransformationsSetShiftDays).toHaveBeenCalledWith({ + min: 1, + max: 30, + entities: ["ENUM_DOB"], + }); + expect(mockSetTransformations).toHaveBeenCalled(); + }); + + it("should configure bleep options for audio", async () => { + const skyflow = createMockSkyflow({}); + await handleDeIdentifyFile( + { ...baseArgs, fileName: "call.wav", bleep: { gain: 0.5, frequency: 800, startPadding: 0.1, stopPadding: 0.2 } }, + skyflow as any, "vault123", false + ); + + expect(mockBleepSetGain).toHaveBeenCalledWith(0.5); + expect(mockBleepSetFrequency).toHaveBeenCalledWith(800); + expect(mockBleepSetStartPadding).toHaveBeenCalledWith(0.1); + expect(mockBleepSetStopPadding).toHaveBeenCalledWith(0.2); + expect(mockSetBleep).toHaveBeenCalled(); + }); + + it("should use DEFAULT_WAIT_TIME_SECONDS when waitTimeSeconds not specified", async () => { const skyflow = createMockSkyflow({}); await handleDeIdentifyFile(baseArgs, skyflow as any, "vault123", false); - expect(mockSetWaitTime).toHaveBeenCalledWith(DEFAULT_MAX_WAIT_TIME_SECONDS); + expect(mockSetWaitTime).toHaveBeenCalledWith(DEFAULT_WAIT_TIME_SECONDS); }); - it("should use provided waitTime when specified", async () => { + it("should use provided waitTimeSeconds when specified", async () => { const skyflow = createMockSkyflow({}); await handleDeIdentifyFile( - { ...baseArgs, waitTime: 30 }, + { ...baseArgs, waitTimeSeconds: 30 }, skyflow as any, "vault123", false ); expect(mockSetWaitTime).toHaveBeenCalledWith(30); }); + it("should clamp waitTimeSeconds to the SDK maximum", async () => { + const skyflow = createMockSkyflow({}); + await handleDeIdentifyFile( + { ...baseArgs, waitTimeSeconds: 500 }, + skyflow as any, "vault123", false + ); + + expect(mockSetWaitTime).toHaveBeenCalledWith(MAX_WAIT_TIME_SECONDS); + }); + it("should include all optional response fields when present", async () => { const skyflow = createMockSkyflow({ fileBase64: "base64data", @@ -174,7 +407,7 @@ describe("handleDeIdentifyFile", () => { durationInSeconds: 30, entities: [{ file: "entity_base64", extension: "png" }], runId: "run_123", - status: "completed", + status: "SUCCESS", }); const result = await handleDeIdentifyFile(baseArgs, skyflow as any, "vault123", false); const output = result.output as DeIdentifyFileOutput; @@ -190,7 +423,82 @@ describe("handleDeIdentifyFile", () => { expect(output.durationInSeconds).toBe(30); expect(output.detectedEntities).toEqual([{ file: "entity_base64", extension: "png" }]); expect(output.runId).toBe("run_123"); - expect(output.status).toBe("completed"); + expect(output.status).toBe("SUCCESS"); + expect(output.note).toBeUndefined(); + }); + + it("should derive a real mimeType from the extension, not the SDK category label", async () => { + const skyflow = createMockSkyflow({ + fileBase64: "base64data", + type: "redacted_image", // SDK category label, not a MIME type + extension: "png", + }); + const result = await handleDeIdentifyFile(baseArgs, skyflow as any, "vault123", false); + const output = result.output as DeIdentifyFileOutput; + + expect(output.mimeType).toBe("image/png"); + expect(output.extension).toBe("png"); + }); + + it("should omit zero-valued counts the SDK defaults to 0", async () => { + const skyflow = createMockSkyflow({ + fileBase64: "base64data", + extension: "png", + wordCount: 0, + charCount: 0, + sizeInKb: 0, + pageCount: 0, + slideCount: 0, + durationInSeconds: 0, + }); + const result = await handleDeIdentifyFile(baseArgs, skyflow as any, "vault123", false); + const output = result.output as DeIdentifyFileOutput; + + expect(output.wordCount).toBeUndefined(); + expect(output.charCount).toBeUndefined(); + expect(output.sizeInKb).toBeUndefined(); + expect(output.pageCount).toBeUndefined(); + expect(output.slideCount).toBeUndefined(); + expect(output.durationInSeconds).toBeUndefined(); + }); + + it("should apply dateShift for supported formats", async () => { + const skyflow = createMockSkyflow({}); + const result = await handleDeIdentifyFile( + { ...baseArgs, fileName: "notes.txt", mimeType: "text/plain", dateShift: { minDays: 1, maxDays: 30, entities: ["dob"] } }, + skyflow as any, "vault123", false + ); + const output = result.output as DeIdentifyFileOutput; + + expect(mockSetTransformations).toHaveBeenCalled(); + expect(output.warnings).toBeUndefined(); + }); + + it("should warn (and skip) dateShift for formats that don't support it", async () => { + const skyflow = createMockSkyflow({}); + const result = await handleDeIdentifyFile( + { ...baseArgs, fileName: "doc.pdf", mimeType: "application/pdf", dateShift: { minDays: 1, maxDays: 30, entities: ["dob"] } }, + skyflow as any, "vault123", false + ); + const output = result.output as DeIdentifyFileOutput; + + expect(mockSetTransformations).not.toHaveBeenCalled(); + expect(output.warnings?.some((w) => w.includes("dateShift"))).toBe(true); + }); + + it("should return runId with polling note for in-progress runs", async () => { + const skyflow = createMockSkyflow({ + runId: "run_async_456", + status: "IN_PROGRESS", + }); + const result = await handleDeIdentifyFile(baseArgs, skyflow as any, "vault123", false); + const output = result.output as DeIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.runId).toBe("run_async_456"); + expect(output.status).toBe("IN_PROGRESS"); + expect(output.note).toContain("get-file-run-status"); + expect(output.note).toContain("run_async_456"); }); it("should omit optional fields when not in response", async () => { diff --git a/tests/unit/tools/getFileRunStatus.test.ts b/tests/unit/tools/getFileRunStatus.test.ts new file mode 100644 index 0000000..f6e36bb --- /dev/null +++ b/tests/unit/tools/getFileRunStatus.test.ts @@ -0,0 +1,248 @@ +import { describe, it, expect, vi, afterEach } from "vitest"; +import { handleGetFileRunStatus } from "../../../src/lib/tools/getFileRunStatus"; +import type { DetectRestContext } from "../../../src/lib/detect/detectRest"; +import type { + DeIdentifyFileOutput, + GetFileRunStatusErrorOutput, +} from "../../../src/lib/tools/types"; + +const context: DetectRestContext = { + vaultUrl: "https://cluster123.vault.skyflowapis.com", + vaultId: "vault123", + credentialKey: "test-api-key", +}; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); +} + +describe("handleGetFileRunStatus", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("returns error with anonymousModeRestricted flag in anonymous mode", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, true); + + expect(result.isError).toBe(true); + expect(result.output).toHaveProperty("anonymousModeRestricted", true); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("calls the runs endpoint with the runId and vault_id", async () => { + const fetchMock = vi.fn(async () => + jsonResponse({ status: "SUCCESS", output: [] }) + ); + vi.stubGlobal("fetch", fetchMock); + + await handleGetFileRunStatus({ runId: "run abc" }, context, false); + + const url = fetchMock.mock.calls[0][0] as string; + expect(url).toBe( + "https://cluster123.vault.skyflowapis.com/v1/detect/runs/run%20abc?vault_id=vault123" + ); + const init = fetchMock.mock.calls[0][1] as RequestInit; + expect((init.headers as Record).Authorization).toBe("Bearer test-api-key"); + }); + + it("maps a successful camelCase run response", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ + status: "SUCCESS", + outputType: "BASE64", + output: [ + { + processedFile: "cHJvY2Vzc2Vk", + processedFileType: "redacted_file", + processedFileExtension: "pdf", + }, + { + processedFile: "ZW50aXR5", + processedFileType: "entities", + processedFileExtension: "png", + }, + ], + wordCharacterCount: { wordCount: 12, characterCount: 80 }, + size: 42, + pages: 3, + }) + ) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as DeIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.runId).toBe("run1"); + expect(output.status).toBe("SUCCESS"); + expect(output.processedFileData).toBe("cHJvY2Vzc2Vk"); + expect(output.extension).toBe("pdf"); + expect(output.detectedEntities).toEqual([{ file: "ZW50aXR5", extension: "png" }]); + expect(output.wordCount).toBe(12); + expect(output.charCount).toBe(80); + expect(output.sizeInKb).toBe(42); + expect(output.pageCount).toBe(3); + }); + + it("maps a successful snake_case run response", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ + status: "SUCCESS", + output_type: "BASE64", + output: [ + { + processed_file: "cHJvY2Vzc2Vk", + processed_file_type: "redacted_file", + processed_file_extension: "txt", + }, + ], + word_count: 5, + character_count: 25, + }) + ) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as DeIdentifyFileOutput; + + expect(output.processedFileData).toBe("cHJvY2Vzc2Vk"); + expect(output.extension).toBe("txt"); + expect(output.wordCount).toBe(5); + expect(output.charCount).toBe(25); + }); + + it("derives a real mimeType from the processed file extension", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ + status: "SUCCESS", + output: [ + { processedFile: "eA==", processedFileType: "redacted_file", processedFileExtension: "png" }, + ], + }) + ) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as DeIdentifyFileOutput; + + expect(output.extension).toBe("png"); + expect(output.mimeType).toBe("image/png"); // not the "redacted_file" category label + }); + + it("returns runId with polling note for in-progress runs", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => jsonResponse({ status: "IN_PROGRESS", output: [] })) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as DeIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.status).toBe("IN_PROGRESS"); + expect(output.note).toContain("get-file-run-status"); + }); + + it("attaches a polling note for non-terminal/unknown statuses", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => jsonResponse({ status: "UNKNOWN", output: [] })) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as DeIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.status).toBe("UNKNOWN"); + expect(output.note).toContain("get-file-run-status"); + }); + + it("polls until success when waitSeconds is set", async () => { + vi.useFakeTimers(); + try { + let call = 0; + vi.stubGlobal( + "fetch", + vi.fn(async () => { + call += 1; + return call < 3 + ? jsonResponse({ status: "IN_PROGRESS", output: [] }) + : jsonResponse({ + status: "SUCCESS", + output: [{ processedFile: "ZG9uZQ==", processedFileType: "redacted_file" }], + }); + }) + ); + + const promise = handleGetFileRunStatus( + { runId: "run1", waitSeconds: 30 }, + context, + false + ); + // Walk through the 2s and 4s backoff sleeps + await vi.advanceTimersByTimeAsync(10_000); + const result = await promise; + const output = result.output as DeIdentifyFileOutput; + + expect(call).toBe(3); + expect(output.status).toBe("SUCCESS"); + expect(output.processedFileData).toBe("ZG9uZQ=="); + } finally { + vi.useRealTimers(); + } + }); + + it("returns isError with the Skyflow message for failed runs", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ status: "FAILED", message: "Unsupported file contents", output: [] }) + ) + ); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as GetFileRunStatusErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("Unsupported file contents"); + expect(output.message).toContain("run1"); + }); + + it("maps HTTP errors from the API", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ error: { message: "run not found" } }, 404) + ) + ); + + const result = await handleGetFileRunStatus({ runId: "missing" }, context, false); + const output = result.output as GetFileRunStatusErrorOutput; + + expect(result.isError).toBe(true); + expect(output.code).toBe(404); + expect(output.message).toBe("run not found"); + }); + + it("wraps network failures", async () => { + vi.stubGlobal("fetch", vi.fn(() => Promise.reject(new Error("socket hang up")))); + + const result = await handleGetFileRunStatus({ runId: "run1" }, context, false); + const output = result.output as GetFileRunStatusErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("socket hang up"); + }); +}); diff --git a/tests/unit/tools/reIdentifyFile.test.ts b/tests/unit/tools/reIdentifyFile.test.ts new file mode 100644 index 0000000..f36913c --- /dev/null +++ b/tests/unit/tools/reIdentifyFile.test.ts @@ -0,0 +1,243 @@ +import { describe, it, expect, vi, afterEach } from "vitest"; +import { handleReIdentifyFile } from "../../../src/lib/tools/reIdentifyFile"; +import type { DetectRestContext } from "../../../src/lib/detect/detectRest"; +import type { + ReIdentifyFileOutput, + ReIdentifyFileErrorOutput, +} from "../../../src/lib/tools/types"; + +const context: DetectRestContext = { + vaultUrl: "https://cluster123.vault.skyflowapis.com", + vaultId: "vault123", + credentialKey: "test-api-key", +}; + +const baseArgs = { + fileDataBase64: Buffer.from("Hello [NAME_1]!").toString("base64"), + fileName: "notes.txt", +}; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); +} + +const successBody = { + status: "SUCCESS", + output_type: "BASE64", + output: { + processed_file: Buffer.from("Hello John!").toString("base64"), + processed_file_type: "reidentified_file", + processed_file_extension: "txt", + }, +}; + +describe("handleReIdentifyFile", () => { + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("returns error with anonymousModeRestricted flag in anonymous mode", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + const result = await handleReIdentifyFile(baseArgs, context, true); + + expect(result.isError).toBe(true); + expect(result.output).toHaveProperty("anonymousModeRestricted", true); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("posts the file to the reidentify endpoint with snake_case body", async () => { + const fetchMock = vi.fn(async () => jsonResponse(successBody)); + vi.stubGlobal("fetch", fetchMock); + + await handleReIdentifyFile(baseArgs, context, false); + + const url = fetchMock.mock.calls[0][0] as string; + expect(url).toBe("https://cluster123.vault.skyflowapis.com/v1/detect/reidentify/file"); + + const init = fetchMock.mock.calls[0][1] as RequestInit; + expect(init.method).toBe("POST"); + expect((init.headers as Record).Authorization).toBe("Bearer test-api-key"); + + const body = JSON.parse(init.body as string); + expect(body.vault_id).toBe("vault123"); + expect(body.file.base64).toBe(baseArgs.fileDataBase64); + expect(body.file.data_format).toBe("txt"); + expect(body.format).toBeUndefined(); + }); + + it("includes the entity format routing when provided", async () => { + const fetchMock = vi.fn(async () => jsonResponse(successBody)); + vi.stubGlobal("fetch", fetchMock); + + await handleReIdentifyFile( + { + ...baseArgs, + redactedEntities: ["ssn"], + maskedEntities: ["email_address"], + plainTextEntities: ["name"], + }, + context, + false + ); + + const body = JSON.parse((vi.mocked(fetchMock).mock.calls[0][1] as RequestInit).body as string); + expect(body.format.redacted).toEqual(["ssn"]); + expect(body.format.masked).toEqual(["email_address"]); + expect(body.format.plaintext).toEqual(["name"]); + }); + + it("returns the processed file on success", async () => { + vi.stubGlobal("fetch", vi.fn(async () => jsonResponse(successBody))); + + const result = await handleReIdentifyFile(baseArgs, context, false); + const output = result.output as ReIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.status).toBe("SUCCESS"); + expect(output.inputFileName).toBe("notes.txt"); + expect(output.extension).toBe("txt"); + expect(Buffer.from(output.processedFileData!, "base64").toString()).toBe("Hello John!"); + }); + + it("parses camelCase responses as well", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => + jsonResponse({ + status: "SUCCESS", + outputType: "BASE64", + output: { + processedFile: Buffer.from("hi").toString("base64"), + processedFileType: "reidentified_file", + processedFileExtension: "csv", + }, + }) + ) + ); + + const result = await handleReIdentifyFile( + { ...baseArgs, fileName: "data.csv" }, + context, + false + ); + const output = result.output as ReIdentifyFileOutput; + + expect(output.extension).toBe("csv"); + expect(output.processedFileData).toBe(Buffer.from("hi").toString("base64")); + }); + + it("downloads the file when a fileUrl is provided", async () => { + const fileContent = Buffer.from("Tokenized [SSN_1]"); + const fetchMock = vi.fn(async (input: RequestInfo | URL) => { + const url = String(input); + if (url.includes("storage.example.com")) { + return new Response(fileContent, { + status: 200, + headers: { "content-type": "text/plain" }, + }); + } + return jsonResponse(successBody); + }); + vi.stubGlobal("fetch", fetchMock); + + const result = await handleReIdentifyFile( + { fileUrl: "https://storage.example.com/docs/tokenized.txt?sig=1" }, + context, + false + ); + const output = result.output as ReIdentifyFileOutput; + + expect(result.isError).toBeUndefined(); + expect(output.inputFileName).toBe("tokenized.txt"); + expect(output.inputFileUrl).toBe("https://storage.example.com/docs/tokenized.txt?sig=1"); + + const reidentifyCall = fetchMock.mock.calls.find(([input]) => + String(input).includes("/v1/detect/reidentify/file") + ); + const body = JSON.parse((reidentifyCall![1] as RequestInit).body as string); + expect(body.file.base64).toBe(fileContent.toString("base64")); + }); + + it("rejects unsupported formats for re-identification", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + const result = await handleReIdentifyFile( + { ...baseArgs, fileName: "image.png" }, + context, + false + ); + const output = result.output as ReIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain('".png"'); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("errors when input is missing", async () => { + const result = await handleReIdentifyFile({}, context, false); + + expect(result.isError).toBe(true); + expect((result.output as ReIdentifyFileErrorOutput).message).toContain("fileUrl"); + }); + + it("returns isError for FAILED responses", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => jsonResponse({ status: "FAILED", output: {} })) + ); + + const result = await handleReIdentifyFile(baseArgs, context, false); + const output = result.output as ReIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("FAILED"); + }); + + it("returns isError for a SUCCESS response that omits the processed file", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => jsonResponse({ status: "SUCCESS", output: {} })) + ); + + const result = await handleReIdentifyFile(baseArgs, context, false); + const output = result.output as ReIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.message).toContain("did not complete"); + }); + + it("maps HTTP errors from the API", async () => { + vi.stubGlobal( + "fetch", + vi.fn(async () => jsonResponse({ error: { message: "invalid vault" } }, 400)) + ); + + const result = await handleReIdentifyFile(baseArgs, context, false); + const output = result.output as ReIdentifyFileErrorOutput; + + expect(result.isError).toBe(true); + expect(output.code).toBe(400); + expect(output.message).toBe("invalid vault"); + }); + + it("rejects invalid entity names before calling the API", async () => { + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + const result = await handleReIdentifyFile( + { ...baseArgs, plainTextEntities: ["not_a_real_entity"] }, + context, + false + ); + + expect(result.isError).toBe(true); + expect((result.output as ReIdentifyFileErrorOutput).message).toContain("Invalid entity type"); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/de-identify-file/main.ts b/ui/de-identify-file/main.ts index 24a2539..b465d9b 100644 --- a/ui/de-identify-file/main.ts +++ b/ui/de-identify-file/main.ts @@ -12,14 +12,16 @@ function escapeHtml(text: string): string { return div.innerHTML; } -function showLoading(fileName?: string, mimeType?: string): void { +function showLoading(fileName?: string, mimeType?: string, runId?: string): void { + const fileLabel = fileName || runId; + const kindLabel = fileName ? "File" : runId ? "Run ID" : ""; root.innerHTML = `
- ${fileName ? ` + ${fileLabel ? `
- ${escapeHtml(fileName)} - File + ${escapeHtml(fileLabel)} + ${escapeHtml(kindLabel)}
${mimeType ? `
@@ -31,7 +33,7 @@ function showLoading(fileName?: string, mimeType?: string): void { ` : ""}
- Processing file for sensitive data... + ${runId && !fileName ? "Checking de-identification run..." : "Processing file for sensitive data..."}
`; @@ -99,6 +101,24 @@ function renderResult(data: DeIdentifyFileResult): void { } } + // In-progress runs: show a prominent processing banner with polling guidance + if (data.status && data.status.toUpperCase() === "IN_PROGRESS" && data.runId) { + root.innerHTML = ` +
+ +
+ + The file is still being processed. ${escapeHtml(data.note || "Check the run status to retrieve the result.")} +
+
+ `; + return; + } + // Build metadata cards const metaItems: { label: string; value: string }[] = []; if (data.inputFileName) metaItems.push({ label: "File Name", value: data.inputFileName }); @@ -130,7 +150,13 @@ function renderResult(data: DeIdentifyFileResult): void { // Status badge for async operations let statusHtml = ""; if (data.runId) { - const statusColor = data.status === "completed" ? "var(--color-text-success, #27ae60)" : "var(--color-text-warning, #f39c12)"; + const statusUpper = (data.status || "").toUpperCase(); + const statusColor = + statusUpper === "SUCCESS" || statusUpper === "COMPLETED" + ? "var(--color-text-success, #27ae60)" + : statusUpper === "FAILED" + ? "var(--color-text-danger, #e74c3c)" + : "var(--color-text-warning, #f39c12)"; statusHtml = `