diff --git a/.changeset/byok-package.md b/.changeset/byok-package.md new file mode 100644 index 000000000..27bae2483 --- /dev/null +++ b/.changeset/byok-package.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai-byok': minor +--- + +Add `@tanstack/ai-byok`: a bring-your-own-key toolkit for TanStack AI. Keys live client-side and travel to the relay in a per-provider header (`x-byok-`), never the request body or message history. + +- **Client** (`@tanstack/ai-byok`): `byokHeaders`, `withByok`/`byokFetch` (attach BYOK headers to a connection and detect the relay's `byokMissing` 401 so the UI can prompt for or unlock the missing key), `byokFetcher` (the same for the `fetcher` transport — `useChat`/`useGeneration` — covering both a fetch call and a TanStack Start server function via call-site headers), a typed provider registry, pluggable storage (session-only memory by default, opt-in passkey-encrypted persistence — WebAuthn PRF → HKDF → AES-256-GCM in IndexedDB, no plaintext option), and `validateKey`. Storage may expose `peek()` to report `provider → last-4` presence without decrypting. +- **React** (`@tanstack/ai-byok/react`): ``, `useByok()` (with `locked`/`unlock` for encrypted storage), and a drop-in `` settings UI that only ever shows the last 4 characters of a saved key. After a refresh, saved keys from encrypted storage surface as a `locked` status (with last-4) before the biometric unlock. +- **Server** (`@tanstack/ai-byok/server`): `getByokKey` (header-only, never logged), `byokMissing` (typed error response), and `scrubSecrets`/`maskKey` for keeping key material out of logs and errors. Stateless pass-through — no persistence, no central endpoint. diff --git a/docs/advanced/byok.md b/docs/advanced/byok.md new file mode 100644 index 000000000..d208a9b86 --- /dev/null +++ b/docs/advanced/byok.md @@ -0,0 +1,435 @@ +--- +title: Bring Your Own Key (BYOK) +id: byok +order: 9 +description: "Let users supply their own provider API keys with @tanstack/ai-byok — client-side keyring, per-request headers, and stateless server relays that never persist or log keys." +keywords: + - tanstack ai + - byok + - bring your own key + - api key + - passkey + - client-side keys + - x-byok + - getByokKey + - withByok +--- + +`@tanstack/ai-byok` is a bring-your-own-key toolkit for TanStack AI. Users paste their own provider API keys in the browser; the library attaches them **per request in a header**, and your server relay reads them **for that one call** without ever persisting or logging them. + +This page covers the security model, a full React + server relay setup, storage tiers, the missing-key flow, and how BYOK works with both [connection adapters](../chat/connection-adapters) and the `fetcher` transport. + +## First principle: never be a custodian + +- **Keys live client-side.** The browser is the system of record. +- **The server is a stateless pass-through.** It reads the key from the incoming request header, hands it to the TanStack AI adapter for one call, and never writes it to a database, cache, log, or observability stream. +- **No central endpoint.** Server helpers are trivially self-hostable — there is no hardcoded relay URL baked into the package. + +Keys travel in the `x-byok-` header — **never** the request body or message history — so they stay out of persisted conversations and the AG-UI event stream. + +## Installation + +```bash +npm install @tanstack/ai-byok +``` + +React bindings ship in the same package: + +```bash +# peer dependency — already present in most TanStack AI React apps +npm install react +``` + +## Package layout + +| Entry | What it contains | +| --- | --- | +| `@tanstack/ai-byok` | Provider registry, `byokHeaders`, `withByok` / `byokFetch`, `byokFetcher`, storage (`memoryStorage`, `passkeyStorage`), `validateKey` | +| `@tanstack/ai-byok/server` | `getByokKey`, `byokMissing`, `preferByokAdapter`, `requireByokOrEnv`, `scrubSecrets` / `maskKey` | +| `@tanstack/ai-byok/react` | ``, `useByok()`, ``, `` | +| `@tanstack/ai-byok/openrouter` | OpenRouter OAuth PKCE helpers (optional) | + +See the [API reference](../api/ai-byok) for every export. + +## Supported providers + +The toolkit understands these provider ids (used in the keyring map and header names): + +| Provider id | Label | Client-side validation | +| --- | --- | --- | +| `openai` | OpenAI | Yes | +| `anthropic` | Anthropic | Yes | +| `gemini` | Google Gemini | Yes | +| `openrouter` | OpenRouter | Yes | +| `groq` | Groq | Yes | +| `grok` | xAI Grok | Yes | +| `mistral` | Mistral | Yes | +| `elevenlabs` | ElevenLabs | Yes | +| `fal` | fal.ai | No browser-reachable endpoint | +| `ollama` | Ollama | Local — no API key | + +Header name for a provider: `x-byok-` (for example `x-byok-openai`). + +## Quick start (React + SSE relay) + +### 1. Wrap the app with a keyring + +```tsx +import { useEffect, useState } from "react"; +import { + ByokProvider, + isPasskeyStorageSupported, + memoryStorage, + passkeyStorage, +} from "@tanstack/ai-byok/react"; + +function App({ children }: { children: React.ReactNode }) { + const [storage, setStorage] = useState | null>( + null, + ); + + useEffect(() => { + let active = true; + async function pick() { + const platformAuth = + isPasskeyStorageSupported() && + (await globalThis.PublicKeyCredential?.isUserVerifyingPlatformAuthenticatorAvailable().catch( + () => false, + )); + if (active) { + setStorage(platformAuth ? passkeyStorage() : memoryStorage()); + } + } + void pick(); + return () => { + active = false; + }; + }, []); + + if (!storage) return null; + return {children}; +} +``` + +`storage` is chosen once when the provider mounts and cannot be changed later. Default to `memoryStorage()` (session-only, nothing persisted). Opt into `passkeyStorage()` where a platform authenticator is available. + +### 2. Attach keys to the chat connection + +```tsx +import { useRef, useCallback } from "react"; +import { useChat, fetchServerSentEvents } from "@tanstack/ai-react"; +import { useByok, withByok } from "@tanstack/ai-byok/react"; +import type { ProviderId } from "@tanstack/ai-byok/react"; +import { openKeyDialog } from "./byok-ui"; + +function Chat() { + const { keys, status, unlock } = useByok(); + const keysRef = useRef(keys); + keysRef.current = keys; + + const handleMissingKey = useCallback( + (provider: ProviderId) => { + if (status[provider]?.state === "locked") { + void unlock(); + } else { + openKeyDialog(provider); + } + }, + [status, unlock], + ); + + const { messages, sendMessage } = useChat({ + connection: fetchServerSentEvents( + "/api/chat", + withByok(() => keysRef.current, { onMissingKey: handleMissingKey }), + ), + }); + + return null; +} +``` + +`withByok` returns a function that produces fresh connection options on every request: BYOK headers merged on top of any static headers, plus an optional missing-key-aware `fetchClient`. + +Use a ref for `keys` so the getter always reads the latest keyring without recreating the connection adapter. + +### 3. Build a stateless server relay + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { createOpenaiChat } from "@tanstack/ai-openai"; +import { getByokKey, byokMissing } from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const apiKey = getByokKey(request, "openai"); + if (!apiKey) return byokMissing("openai"); + + const stream = chat({ + adapter: createOpenaiChat("gpt-5.2", apiKey), + messages, + }); + return toServerSentEventsResponse(stream); +} +``` + +`getByokKey` reads the header only — never the body. It returns `null` when absent and must not be wrapped in anything that attaches the value to a logger context. + +When neither a server env key nor a BYOK header is present, return `byokMissing(provider)` so the client can prompt for the right key instead of surfacing a generic provider error. + +### 4. Optional: drop-in settings UI + +```tsx +import { ByokKeyManager } from "@tanstack/ai-byok/react"; + +function Settings() { + return ; +} +``` + +`ByokKeyManager` is write-only from the UI's perspective: once saved, only the last four characters are ever shown. + +## Missing-key flow + +When the relay has no env key and the request carried no BYOK header, return a typed 401: + +```typescript +import { byokMissing } from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + // ...no BYOK header and no server env key for this provider + return byokMissing("openai"); + // → 401 { error: { type: "byok_missing", provider: "openai", message: "..." } } +} +``` + +On the client, `withByok` / `byokFetch` detect this response and invoke `onMissingKey(provider)` as a side channel — the SSE error path only exposes the HTTP status, not the JSON body. + +Typical UI response: + +1. If the provider's key is **saved but locked** (passkey storage after a refresh) → call `unlock()`. +2. Otherwise → open a key-entry dialog for that provider. + +## Storage tiers + +Pass one `KeyringStorage` to ``. Built-in strategies: + +### `defaultByokStorage()` (recommended) + +Passkey-encrypted storage when supported; session memory otherwise. **All keys** — pasted or OpenRouter PKCE — use the same tier and survive refresh behind a biometric unlock. + +```tsx +import { ByokProvider, defaultByokStorage } from "@tanstack/ai-byok/react"; + +function Root({ children }: { children: React.ReactNode }) { + return ( + {children} + ); +} +``` + +OpenRouter enforces PKCE key expiry server-side based on what the user chose at sign-in. + +### `memoryStorage()` (default alone) + +Session-only. All keys live in React state and vanish on refresh. Zero at-rest liability. Used when `` has no `storage` prop, or as the `defaultByokStorage()` fallback when passkeys are unavailable. + +### `passkeyStorage()` + +Encrypted at rest in IndexedDB (AES-256-GCM), with the encryption key derived from a passkey's WebAuthn **PRF** output via HKDF. Decryption happens entirely client-side with a biometric or PIN tap — no server, no custodian. + +Feature-detect with `isPasskeyStorageSupported()` and fall back to `memoryStorage()`. PRF support is solid on Android and Apple platform authenticators; it can be patchy on Firefox and with roaming keys on Safari. + +**Honest scope:** passkey storage protects against at-rest theft (stolen device, storage-dumping extensions, backups). It does **not** defeat live in-page XSS — an attacker running JavaScript in the origin after the user unlocks can read decrypted keys from memory. `` surfaces this caveat while passkey storage is active. + +`passkeyStorage` is **unlockable**: `` does not decrypt on mount. After a refresh, `peek()` reads unencrypted `provider → last-4` metadata so the UI can show saved keys as **locked** before the user taps unlock. Call `unlock()` (or save a key, which runs the registration ceremony on first use) to decrypt. + +Supply your own `KeyringStorage` implementation for custom strategies (for example a team-specific HSM workflow). + +## Fetcher transport + +`withByok` targets [connection adapters](../chat/connection-adapters) (`fetchServerSentEvents`, `fetchHttpStream`, etc.). For the **`fetcher`** transport — used by `useGeneration` (image, audio, video, speech, transcribe) and `useChat({ fetcher })` — use `byokFetcher`: + +```tsx +import { useRef } from "react"; +import { useGenerateAudio } from "@tanstack/ai-react"; +import { byokFetcher } from "@tanstack/ai-byok"; +import { openKeyDialog } from "./byok-ui"; +import { generateAudioFn } from "./server-functions"; + +function AudioGenerator() { + const keysRef = useRef({ elevenlabs: "xi-key" }); + + // Fetch-based fetcher + useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers, fetch, signal }) => + fetch("/api/generate/audio", { + method: "POST", + headers: { "content-type": "application/json", ...headers }, + body: JSON.stringify(input), + signal, + }), + { onMissingKey: (provider) => openKeyDialog(provider) }, + ), + }); + + // TanStack Start server function — forward headers at the call site + useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers }) => generateAudioFn({ data: input, headers }), + ), + }); + + return null; +} +``` + +On the server, read the key the same way: + +```typescript +import { getByokKey, byokMissing } from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + const apiKey = getByokKey(request, "elevenlabs"); + if (!apiKey) return byokMissing("elevenlabs"); + // ... +} +``` + +`onMissingKey` fires only when `byokFetcher` owns the `fetch` call. Server-function fetchers surface a missing key as a thrown error instead. + +## OpenRouter PKCE sign-in + +OpenRouter supports one-click login via [OAuth PKCE](https://openrouter.ai/docs/guides/overview/auth/oauth). The user signs in at OpenRouter (choosing key expiry and limits on OpenRouter's screen); your app receives a user-controlled API key and stores it in the keyring like any other BYOK key. + +### React hook + +```tsx +import { useOpenRouterPkce } from "@tanstack/ai-byok/openrouter/react"; + +function App() { + const { login, completing, error } = useOpenRouterPkce(); + // On return from OpenRouter (?code=…), exchanges the code and calls + // setKey("openrouter", key) automatically. + + return ( + + ); +} +``` + +`callbackUrl` defaults to `origin + pathname` of the current page. Override it when your OAuth return route differs. + +Pass `openRouter={{ onLogin, completing, error }}` to `` or `` to show a **Sign in with OpenRouter** button on the OpenRouter row when no key is set. + +### Lower-level client API + +```tsx +import { useByok } from "@tanstack/ai-byok/react"; +import { + startOpenRouterPkceLogin, + completeOpenRouterPkceFromUrl, +} from "@tanstack/ai-byok/openrouter"; + +function OpenRouterLogin() { + const { setKey } = useByok(); + + async function login() { + await startOpenRouterPkceLogin({ + callbackUrl: "https://myapp.com/chat", + }); + } + + async function completeFromCallback() { + const key = await completeOpenRouterPkceFromUrl(); + if (key) await setKey("openrouter", key); + } + + return null; +} +``` + +S256 PKCE is used by default. Localhost callbacks are supported on any port for local dev. + +## Key validation + +`validateKey(provider, key)` pings the provider's cheapest authenticated endpoint (usually a models list) before the user hits a wall mid-stream. + +Returns: + +- `'valid'` — provider accepted the key +- `'invalid'` — 401/403 from the provider +- `'unsupported'` — no browser-reachable validation endpoint for this provider + +Throws on network/CORS failures or unexpected HTTP status rather than guessing. Some providers block browser origins entirely; a thrown `TypeError` ("Failed to fetch") is the honest signal. + +In React, `useByok().validateKey(provider)` records the outcome in `status[provider]` and never throws — failures become `{ state: 'error', message }`. + +## Server logging and errors + +The server helpers never log key material. Relay authors are responsible for keeping keys out of their own logs and error responses: + +- Use `scrubSecrets(input, [apiKey])` before logging strings that may have interpolated a key (provider SDK error messages, URLs, stack traces). +- Use `maskKey(key)` when you need a display-safe representation (last four characters only). +- Never echo a full key back to the client in an error payload. + +## Env keys vs BYOK keys + +A common pattern (see the `ts-react-chat` example) keeps **server env keys** as the default and lets BYOK override per request: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { createOpenaiChat, openaiText } from "@tanstack/ai-openai"; +import { + preferByokAdapter, + requireByokOrEnv, +} from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + const { messages, model } = await request.json(); + + const blocked = requireByokOrEnv(request, "openai", ["OPENAI_API_KEY"]); + if (blocked) return blocked; + + const adapter = preferByokAdapter(request, "openai", model, { + byok: createOpenaiChat, + env: openaiText, + }); + + return toServerSentEventsResponse(chat({ adapter, messages })); +} +``` + +On the client, a server function can report which providers have env keys (booleans only — never the values) so the UI can warn before the user picks a model they cannot run. + +## Lower-level client API + +When you do not need React bindings: + +```typescript +import { byokHeaders, byokFetch } from "@tanstack/ai-byok"; +import { openKeyDialog } from "./byok-ui"; + +const headers = byokHeaders({ openai: "sk-...", gemini: "..." }); +// → { "x-byok-openai": "sk-...", "x-byok-gemini": "..." } + +const fetchWithByok = byokFetch((provider) => openKeyDialog(provider)); +``` + +Empty or absent keys are skipped — only providers with a non-empty key get a header. + +## Example + +The `examples/ts-react-chat` app integrates BYOK into the main chat page: a key icon opens a per-provider dialog, `withByok` attaches headers on every SSE request, and `/api/tanchat` prefers the BYOK header over env-configured adapters. E2E coverage lives in `testing/e2e/tests/byok.spec.ts`. + +## Related + +- [Connection Adapters](../chat/connection-adapters) — transports `withByok` plugs into +- [@tanstack/ai-byok API](../api/ai-byok) — full export reference +- [Debug Logging](./debug-logging) — keep custom middleware from logging secrets \ No newline at end of file diff --git a/docs/api/ai-byok.md b/docs/api/ai-byok.md new file mode 100644 index 000000000..a890558d1 --- /dev/null +++ b/docs/api/ai-byok.md @@ -0,0 +1,339 @@ +--- +title: "@tanstack/ai-byok" +slug: /api/ai-byok +order: 9 +description: "API reference for @tanstack/ai-byok — bring-your-own-key client keyring, per-request headers, React bindings, and stateless server helpers." +keywords: + - tanstack ai + - "@tanstack/ai-byok" + - byok + - api key + - getByokKey + - withByok + - byokFetcher + - api reference +--- + +Bring-your-own-key toolkit for TanStack AI. See the [BYOK guide](../advanced/byok) for architecture, security model, and end-to-end setup. + +## Installation + +```bash +npm install @tanstack/ai-byok +``` + +Subpath exports: + +```typescript +// Stateless server helpers (no React dependency) +import { getByokKey, byokMissing } from "@tanstack/ai-byok/server"; + +// React bindings (requires react peer) +import { ByokProvider, useByok } from "@tanstack/ai-byok/react"; + +// OpenRouter OAuth PKCE (optional vendor add-on) +import { useOpenRouterPkce } from "@tanstack/ai-byok/openrouter/react"; +``` + +## `@tanstack/ai-byok` (client) + +Framework-agnostic client toolkit. No React or server dependencies. + +### `byokHeaders(keys)` + +Turns a keyring into per-provider request headers. Skips empty or absent keys. + +```typescript +import { byokHeaders } from "@tanstack/ai-byok"; + +const headers = byokHeaders({ openai: "sk-live", anthropic: "" }); +// → { "x-byok-openai": "sk-live" } +``` + +### `withByok(getKeys, options?)` + +Builds BYOK connection options for fetch-based [connection adapters](../chat/connection-adapters). Returns a **function** that produces fresh options on every request. + +```tsx +import { useRef } from "react"; +import { useChat, fetchServerSentEvents } from "@tanstack/ai-react"; +import { withByok } from "@tanstack/ai-byok"; +import { openKeyDialog } from "./byok-ui"; +import { customFetch } from "./fetch"; + +function Chat() { + const keysRef = useRef({ openai: "sk-live" }); + const buildOptions = withByok(() => keysRef.current, { + onMissingKey: (provider) => openKeyDialog(provider), + headers: { "x-custom": "value" }, + fetchClient: customFetch, + }); + + useChat({ + connection: fetchServerSentEvents("/api/chat", buildOptions), + }); + + return null; +} +``` + +**`WithByokOptions`** + +| Field | Type | Description | +| --- | --- | --- | +| `onMissingKey?` | `(provider: ProviderId) => void` | Called when the relay returns a `byokMissing` 401 | +| `headers?` | `Record` | Extra headers merged under BYOK headers | +| `fetchClient?` | `typeof fetch` | Underlying fetch (defaults to global `fetch`) | + +### `buildByokRequestContext(getKeys, options?, signal?)` + +Shared header + fetch wiring used by `withByok` and `byokFetcher`. Returns `{ headers, fetch, signal? }`. + +### `byokFetch(onMissingKey, fetchImpl?)` + +Wraps `fetch` so a `byokMissing` 401 invokes `onMissingKey` with the provider id. The response is passed through unchanged. + +### `byokFetcher(getKeys, handler, options?)` + +The `fetcher` transport counterpart to `withByok`. Wraps a fetcher body so it receives BYOK `headers`, a missing-key-aware `fetch`, and the transport `signal`, read fresh on every call. + +```tsx +import { useRef } from "react"; +import { useGenerateAudio } from "@tanstack/ai-react"; +import { byokFetcher } from "@tanstack/ai-byok"; +import { openKeyDialog } from "./byok-ui"; + +function AudioPage() { + const keysRef = useRef({ elevenlabs: "xi-key" }); + + useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers, fetch, signal }) => + fetch("/api/generate", { + method: "POST", + headers: { "content-type": "application/json", ...headers }, + body: JSON.stringify(input), + signal, + }), + { onMissingKey: (provider) => openKeyDialog(provider) }, + ), + }); + + return null; +} +``` + +**`ByokFetcherContext`** + +| Field | Type | Description | +| --- | --- | --- | +| `headers` | `Record` | Per-provider BYOK headers for this request | +| `fetch` | `typeof fetch` | Missing-key-aware fetch (identical to global when `onMissingKey` is unset) | +| `signal?` | `AbortSignal` | Abort signal forwarded from `stop()`, when provided | + +### `isByokMissingBody(value)` + +Type guard for a `byokMissing` response body parsed from JSON. + +### Storage + +#### `defaultByokStorage(options?)` + +Recommended: passkey-encrypted storage when supported, otherwise session memory. All keys (pasted and OpenRouter PKCE) use the same tier. + +#### `memoryStorage()` + +Session-only — keys vanish on refresh, nothing persisted. Default when `` has no `storage` prop. + +#### `passkeyStorage(options?)` + +Passkey-encrypted persistence (WebAuthn PRF → HKDF → AES-256-GCM in IndexedDB). Unlockable — requires `unlock()` or a save before keys are usable after refresh. + +**`PasskeyStorageOptions`** + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `rpName?` | `string` | `"BYOK"` | Relying-party name in the passkey prompt | +| `userName?` | `string` | `"byok-keyring"` | Username label on the created passkey | +| `rpId?` | `string` | current origin | Parent domain for cross-subdomain sharing | +| `dbName?` | `string` | `"byok"` | IndexedDB database name | + +#### `isPasskeyStorageSupported()` + +Returns whether WebAuthn is available. PRF support is confirmed only during registration — catch errors from `passkeyStorage()` and fall back to `memoryStorage()`. + +### OpenRouter PKCE (`@tanstack/ai-byok/openrouter`) + +| Export | Description | +| --- | --- | +| `generateCodeVerifier()` | Random URL-safe PKCE verifier | +| `createS256CodeChallenge(verifier)` | S256 code challenge (base64url SHA-256) | +| `buildOpenRouterAuthUrl(options)` | OpenRouter `/auth` redirect URL | +| `startOpenRouterPkceLogin(options)` | Store pending state and redirect | +| `exchangeOpenRouterCode(options)` | POST code → API key | +| `completeOpenRouterPkceFromUrl(options?)` | Read `?code=`, exchange, clean up | +| `defaultOpenRouterCallbackUrl()` | `origin + pathname` default | +| `storeOpenRouterPkcePending` / `load` / `clear` | Session-scoped PKCE state | + +### `useOpenRouterPkce(options?)` (`@tanstack/ai-byok/openrouter/react`) + +React hook (requires ``). Auto-completes when the URL contains `?code=` from OpenRouter (unless `autoComplete: false`). + +| Option | Default | Description | +| --- | --- | --- | +| `callbackUrl?` | `origin + pathname` | OpenRouter redirect target | +| `autoComplete?` | `true` | Exchange code on mount | +| `useS256?` | `true` | S256 PKCE challenge | + +Returns `{ login, completing, error, callbackUrl }`. + +### `validateKey(provider, key)` + +Pings the provider's validation endpoint. Returns `'valid' | 'invalid' | 'unsupported'`. Throws on network/CORS failure or unexpected HTTP status. + +### Provider registry + +| Export | Description | +| --- | --- | +| `BYOK_PROVIDERS` | Static metadata map (`id`, `label`, optional `validate` config) | +| `PROVIDER_IDS` | Runtime array of all provider ids | +| `BYOK_HEADER_PREFIX` | `"x-byok-"` | +| `byokHeaderName(provider)` | Full header name for a provider | +| `isProviderId(value)` | Type guard for `ProviderId` | + +### Types + +| Type | Description | +| --- | --- | +| `Keyring` | `Partial>` — in-memory key map | +| `ProviderId` | Union of registered provider ids | +| `KeyringStorage` | Pluggable persistence interface (`load`, `save`, `clear`, optional `peek`, `unlockable`) | +| `ValidationStatus` | `'valid' \| 'invalid' \| 'unsupported'` | +| `ByokMissingBody` | Typed JSON body from `byokMissing` | + +## `@tanstack/ai-byok/server` + +Stateless server helpers. No persistence, no logging of key values. + +### `getByokKey(request, provider)` + +Reads a provider's BYOK key from the incoming request header. Returns the key string or `null` when absent. + +Accepts any object with a `Headers`-like `.get()` — works across Fetch-API runtimes (Workers, Deno, Bun, Node/undici). + +```typescript +import { getByokKey } from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + const apiKey = getByokKey(request, "openai"); + // ... +} +``` + +### `byokMissing(provider, init?)` + +Returns a typed JSON 401 telling the client which provider key is missing. Carries no key material. + +```typescript +import { byokMissing, getByokKey } from "@tanstack/ai-byok/server"; + +export async function POST(request: Request) { + const apiKey = getByokKey(request, "anthropic"); + if (!apiKey) return byokMissing("anthropic"); + // ... +} +``` + +### `preferByokAdapter(request, provider, model, factories)` + +Prefer a per-request BYOK header key over a server env-configured adapter factory. + +### `requireByokOrEnv(request, provider, envVarNames)` + +Return a typed `byokMissing` response when neither a BYOK header nor any named env var is present. Returns `null` when the request may proceed. + +### `scrubSecrets(input, secrets)` + +Replaces every occurrence of each secret in `input` with its masked form. Use before logging or returning strings that may have interpolated a key. + +### `maskKey(key)` / `lastFour(key)` + +Display-safe key representations (last four characters only). + +## `@tanstack/ai-byok/react` + +React bindings for the client keyring. + +### `` + +Provides the keyring context. `storage` is chosen once at mount and cannot change. + +```tsx +import { ByokProvider, memoryStorage } from "@tanstack/ai-byok/react"; + +function Root({ children }: { children: React.ReactNode }) { + return ( + {children} + ); +} +``` + +### `useByok()` + +Access the keyring and controls. Must be called under `` — throws otherwise. + +```tsx +import { useByok } from "@tanstack/ai-byok/react"; + +function KeySettings() { + const { + keys, // live keyring — pass to byokHeaders / withByok + setKey, // (provider, key) => Promise + clearKey, // (provider) => Promise + clearAll, // () => Promise + validateKey, // (provider, key?) => Promise + status, // per-provider KeyStatus map + storage, // configured KeyringStorage + locked, // true when unlockable storage may hold encrypted keys + unlock, // () => Promise — decrypt / load + hasKey, // (provider) => boolean + } = useByok(); + + return null; +} +``` + +**`KeyStatus` union** + +| `state` | Meaning | +| --- | --- | +| `'empty'` | No key stored | +| `'set'` | Key present and usable; `masked` shows last 4 | +| `'locked'` | Saved in unlockable storage but not decrypted this session | +| `'validating'` | Validation in flight | +| `'valid'` / `'invalid'` / `'unsupported'` | Validation outcome | +| `'error'` | Validation failed (network, etc.); includes `message` | + +### `` + +Drop-in settings UI for entering, validating, and clearing keys. Only ever shows the last four characters of a saved key. Accepts optional `envStatus`, `highlightProvider`, `openRouter`, and `variant` (`'light' | 'dark'`). + +### `` + +Modal wrapper around `` with a trigger button. Supports custom `trigger`, `overlayClassName`, and `panelClassName` for app styling. + +## Header convention + +Every present provider key is sent as: + +``` +x-byok-: +``` + +Keys are **never** placed in the request body, `forwardedProps`, or message history. + +## Related + +- [Bring Your Own Key (BYOK)](../advanced/byok) — full guide +- [Connection Adapters](../chat/connection-adapters) — `withByok` integration \ No newline at end of file diff --git a/docs/chat/connection-adapters.md b/docs/chat/connection-adapters.md index a67e17e6f..bebe017c5 100644 --- a/docs/chat/connection-adapters.md +++ b/docs/chat/connection-adapters.md @@ -66,6 +66,8 @@ const { messages } = useChat({ }); ``` +> **Tip:** For [bring-your-own-key (BYOK)](../advanced/byok) flows, use `withByok` from `@tanstack/ai-byok` as the connection options factory. It merges `x-byok-` headers on every request and detects the relay's `byokMissing` 401 so the UI can prompt for (or unlock) the missing key. + **Static body.** Anything in `options.body` is merged into the AG-UI `forwardedProps` payload sent to your server. Per-message data passed to `sendMessage` wins over this: ```typescript diff --git a/docs/config.json b/docs/config.json index 7debc1ba4..5052abd20 100644 --- a/docs/config.json +++ b/docs/config.json @@ -443,6 +443,11 @@ "label": "Typed Pre-Configured Options", "to": "advanced/typed-options", "addedAt": "2026-05-25" + }, + { + "label": "Bring Your Own Key (BYOK)", + "to": "advanced/byok", + "addedAt": "2026-07-16" } ] }, @@ -522,6 +527,11 @@ "to": "api/ai-angular", "addedAt": "2026-06-15", "updatedAt": "2026-07-08" + }, + { + "label": "@tanstack/ai-byok", + "to": "api/ai-byok", + "addedAt": "2026-07-16" } ] }, diff --git a/docs/getting-started/overview.md b/docs/getting-started/overview.md index 22e3cf1ad..93735e1c8 100644 --- a/docs/getting-started/overview.md +++ b/docs/getting-started/overview.md @@ -101,6 +101,15 @@ Solid hooks for TanStack AI: - Tool approval flow support - Type-safe message handling with `InferChatMessages` +### `@tanstack/ai-byok` +Bring-your-own-key toolkit for apps where users supply their own provider API keys: +- Client-side keyring with session-only or passkey-encrypted storage +- Per-request `x-byok-` headers (never the message body) +- Stateless server helpers (`getByokKey`, `byokMissing`) that never persist or log keys +- React bindings (``, `useByok`, ``) + +See the [BYOK guide](../advanced/byok). + ## Adapters With the help of adapters, TanStack AI can connect to various LLM providers. Available adapters include: diff --git a/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 67b58ad43..7faef3832 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -18,6 +18,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", + "@tanstack/ai-byok": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", "@tanstack/ai-code-mode": "workspace:*", diff --git a/examples/ts-react-chat/src/lib/byok-config.ts b/examples/ts-react-chat/src/lib/byok-config.ts new file mode 100644 index 000000000..53c1974af --- /dev/null +++ b/examples/ts-react-chat/src/lib/byok-config.ts @@ -0,0 +1,41 @@ +import { createServerFn } from '@tanstack/react-start' +import type { ProviderId } from '@tanstack/ai-byok' +import type { Provider } from '@/lib/model-selection' + +/** + * Maps the example's providers to a BYOK provider id and the env var(s) the + * adapter reads server-side. Providers not listed (ollama = local, bedrock = + * AWS credentials) don't use a single bring-your-own key. + */ +export const BYOK_PROVIDER_MAP: Partial< + Record }> +> = { + openai: { byokId: 'openai', envVars: ['OPENAI_API_KEY'] }, + anthropic: { byokId: 'anthropic', envVars: ['ANTHROPIC_API_KEY'] }, + gemini: { byokId: 'gemini', envVars: ['GEMINI_API_KEY', 'GOOGLE_API_KEY'] }, + 'gemini-interactions': { + byokId: 'gemini', + envVars: ['GEMINI_API_KEY', 'GOOGLE_API_KEY'], + }, + grok: { byokId: 'grok', envVars: ['XAI_API_KEY'] }, + groq: { byokId: 'groq', envVars: ['GROQ_API_KEY'] }, + openrouter: { byokId: 'openrouter', envVars: ['OPENROUTER_API_KEY'] }, +} + +/** The BYOK provider id backing an example provider, if any. */ +export function byokIdForProvider(provider: Provider): ProviderId | null { + return BYOK_PROVIDER_MAP[provider]?.byokId ?? null +} + +/** + * Reports which BYOK-supported providers have their key set in the server env, + * so the client can warn before a user picks a model it can't run. Returns only + * booleans — never the key values. + */ +export const getEnvKeyStatus = createServerFn({ method: 'GET' }).handler(() => { + const status: Partial> = {} + for (const { byokId, envVars } of Object.values(BYOK_PROVIDER_MAP)) { + status[byokId] = envVars.some((name) => Boolean(process.env[name])) + } + return status +}) diff --git a/examples/ts-react-chat/src/routes/api.tanchat.ts b/examples/ts-react-chat/src/routes/api.tanchat.ts index 775723aaa..2fca62512 100644 --- a/examples/ts-react-chat/src/routes/api.tanchat.ts +++ b/examples/ts-react-chat/src/routes/api.tanchat.ts @@ -7,15 +7,20 @@ import { mergeAgentTools, toServerSentEventsResponse, } from '@tanstack/ai' -import { openaiText } from '@tanstack/ai-openai' +import { createOpenaiChat, openaiText } from '@tanstack/ai-openai' import { ollamaText } from '@tanstack/ai-ollama' -import { anthropicText } from '@tanstack/ai-anthropic' -import { geminiText } from '@tanstack/ai-gemini' -import { geminiTextInteractions } from '@tanstack/ai-gemini/experimental' -import { openRouterText } from '@tanstack/ai-openrouter' -import { grokText } from '@tanstack/ai-grok' -import { groqText } from '@tanstack/ai-groq' +import { anthropicText, createAnthropicChat } from '@tanstack/ai-anthropic' +import { createGeminiChat, geminiText } from '@tanstack/ai-gemini' +import { + createGeminiTextInteractions, + geminiTextInteractions, +} from '@tanstack/ai-gemini/experimental' +import { createOpenRouterText, openRouterText } from '@tanstack/ai-openrouter' +import { createGrokText, grokText } from '@tanstack/ai-grok' +import { createGroqText, groqText } from '@tanstack/ai-groq' import { bedrockText } from '@tanstack/ai-bedrock' +import { preferByokAdapter, requireByokOrEnv } from '@tanstack/ai-byok/server' +import { BYOK_PROVIDER_MAP, byokIdForProvider } from '@/lib/byok-config' import type { AnyTextAdapter, ChatMiddleware } from '@tanstack/ai' import { addToCartToolDef, @@ -256,14 +261,20 @@ export const Route = createFileRoute('/api/tanchat')({ > = { anthropic: () => createChatOptions({ - adapter: anthropicText( + adapter: preferByokAdapter( + request, + 'anthropic', (model || 'claude-sonnet-4-6') as 'claude-sonnet-4-6', + { byok: createAnthropicChat, env: anthropicText }, ), }), openrouter: () => createChatOptions({ - adapter: openRouterText( + adapter: preferByokAdapter( + request, + 'openrouter', (model || 'openai/gpt-5.1') as 'openai/gpt-5.1', + { byok: createOpenRouterText, env: openRouterText }, ), modelOptions: { reasoning: { @@ -273,8 +284,11 @@ export const Route = createFileRoute('/api/tanchat')({ }), gemini: () => createChatOptions({ - adapter: geminiText( + adapter: preferByokAdapter( + request, + 'gemini', (model || 'gemini-3.1-pro-preview') as 'gemini-3.1-pro-preview', + { byok: createGeminiChat, env: geminiText }, ), modelOptions: { thinkingConfig: { @@ -285,8 +299,14 @@ export const Route = createFileRoute('/api/tanchat')({ }), 'gemini-interactions': () => createChatOptions({ - adapter: geminiTextInteractions( + adapter: preferByokAdapter( + request, + 'gemini', (model || 'gemini-3.1-pro-preview') as 'gemini-3.1-pro-preview', + { + byok: createGeminiTextInteractions, + env: geminiTextInteractions, + }, ), modelOptions: { previous_interaction_id: previousInteractionId, @@ -295,15 +315,21 @@ export const Route = createFileRoute('/api/tanchat')({ }), grok: () => createChatOptions({ - adapter: grokText( + adapter: preferByokAdapter( + request, + 'grok', (model || 'grok-build-0.1') as 'grok-build-0.1', + { byok: createGrokText, env: grokText }, ), modelOptions: {}, }), groq: () => createChatOptions({ - adapter: groqText( + adapter: preferByokAdapter( + request, + 'groq', (model || 'openai/gpt-oss-120b') as 'openai/gpt-oss-120b', + { byok: createGroqText, env: groqText }, ), }), bedrock: () => @@ -330,7 +356,12 @@ export const Route = createFileRoute('/api/tanchat')({ }), openai: () => createChatOptions({ - adapter: openaiText((model || 'gpt-5.2') as 'gpt-5.2'), + adapter: preferByokAdapter( + request, + 'openai', + (model || 'gpt-5.2') as 'gpt-5.2', + { byok: createOpenaiChat, env: openaiText }, + ), modelOptions: { prompt_cache_key: 'user-session-12345', prompt_cache_retention: '24h', @@ -344,6 +375,20 @@ export const Route = createFileRoute('/api/tanchat')({ requestedProvider in adapterConfig ? (requestedProvider as Provider) : 'openai' + + // If this provider has no server env key and the request carried no + // BYOK key, tell the client which key to add — a typed 401 the client + // detects (via withByok) instead of a generic 500 from the SDK. + const byokId = byokIdForProvider(provider) + if (byokId) { + const blocked = requireByokOrEnv( + request, + byokId, + BYOK_PROVIDER_MAP[provider]?.envVars ?? [], + ) + if (blocked) return blocked + } + // Get typed adapter options using createChatOptions pattern const options = adapterConfig[provider]() diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index b3228dd67..926a80df9 100644 --- a/examples/ts-react-chat/src/routes/index.tsx +++ b/examples/ts-react-chat/src/routes/index.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo, useRef, useState } from 'react' +import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { Link, createFileRoute } from '@tanstack/react-router' import { Braces, @@ -8,6 +8,7 @@ import { Github, Image, ImagePlus, + KeyRound, Mic, Music, Send, @@ -26,7 +27,17 @@ import { useChat, useTranscription, } from '@tanstack/ai-react' +import { + ByokKeyDialog, + ByokProvider, + defaultByokStorage, + useByok, + withByok, +} from '@tanstack/ai-byok/react' +import { useOpenRouterPkce } from '@tanstack/ai-byok/openrouter/react' import { clientTools } from '@tanstack/ai-client' +import type { ProviderId } from '@tanstack/ai-byok/react' +import { byokIdForProvider, getEnvKeyStatus } from '@/lib/byok-config' import { ThinkingPart } from '@tanstack/ai-react-ui' import type { UIMessage } from '@tanstack/ai-react' import type { ContentPart } from '@tanstack/ai' @@ -379,6 +390,50 @@ function Messages({ function ChatPage() { const [selectedModel, setSelectedModel] = useState(DEFAULT_MODEL_OPTION) + + // BYOK: read the browser-held keyring and report which providers have a + // server-side env key, so we can warn before a keyless model is used. + const { keys, status, unlock } = useByok() + const openRouterPkce = useOpenRouterPkce() + const keysRef = useRef(keys) + keysRef.current = keys + const statusRef = useRef(status) + statusRef.current = status + const [envKeyStatus, setEnvKeyStatus] = useState< + Partial> + >({}) + useEffect(() => { + void getEnvKeyStatus().then(setEnvKeyStatus) + }, []) + + // Key dialog, opened either by the toolbar icon or reactively when the relay + // reports a missing key. + const [keyDialog, setKeyDialog] = useState<{ + open: boolean + provider: ProviderId | null + }>({ open: false, provider: null }) + + // The relay returned a byokMissing 401 (no server key, no BYOK key). If we + // already hold that key but it's locked, unlock it; otherwise prompt to add. + const handleMissingKey = useCallback( + (provider: ProviderId) => { + if (statusRef.current[provider]?.state === 'locked') { + void unlock() + } else { + setKeyDialog({ open: true, provider }) + } + }, + [unlock], + ) + + const activeByokId = byokIdForProvider(selectedModel.provider) + // The selected model can't run right now if its provider has no server key + // and no decrypted key in the browser. + const notUsable = + activeByokId != null && !envKeyStatus[activeByokId] && !keys[activeByokId] + // A saved-but-locked key just needs unlocking; distinguish it from "no key". + const activeLocked = + activeByokId != null && status[activeByokId]?.state === 'locked' const [attachedImages, setAttachedImages] = useState< Array<{ id: string; base64: string; mimeType: string; preview: string }> >([]) @@ -422,7 +477,10 @@ function ChatPage() { addToolApprovalResponse, stop, } = useChat({ - connection: fetchServerSentEvents('/api/tanchat'), + connection: fetchServerSentEvents( + '/api/tanchat', + withByok(() => keysRef.current, { onMissingKey: handleMissingKey }), + ), tools, body, onCustomEvent: (eventType, data, context) => { @@ -619,6 +677,44 @@ function ChatPage() { ))} + setKeyDialog((s) => ({ ...s, open }))} + highlightProvider={keyDialog.provider} + openRouter={{ + onLogin: () => void openRouterPkce.login(), + completing: openRouterPkce.completing, + error: openRouterPkce.error, + }} + trigger={ + + } + overlayClassName="fixed inset-0 z-50 flex items-center justify-center bg-black/60 p-4" + panelClassName="w-full max-w-md rounded-xl border border-gray-700 bg-gray-900 p-5 shadow-2xl" + variant="dark" + description="Keys stay in your browser and are sent per-request in a header — never stored on the server. Providers with a server key already work without one." + /> + {notUsable && activeByokId ? ( + activeLocked ? ( +
+ Your {activeByokId} key is saved but locked. + +
+ ) : ( +
+ No key for {activeByokId}. Add one with the key icon to use this + model. +
+ ) + ) : null} defaultByokStorage()) + return ( + + + + ) +} + export const Route = createFileRoute('/')({ - component: ChatPage, + component: ChatRoute, }) diff --git a/knip.json b/knip.json index c7995c1aa..19c41a25c 100644 --- a/knip.json +++ b/knip.json @@ -48,6 +48,9 @@ "packages/ai-react": { "ignoreDependencies": ["@mcp-ui/client"] }, + "packages/ai-byok": { + "ignoreDependencies": ["react", "@types/react"] + }, "packages/ai-preact": { "ignoreDependencies": ["@mcp-ui/client"] }, diff --git a/packages/ai-byok/README.md b/packages/ai-byok/README.md new file mode 100644 index 000000000..1fbb32ae9 --- /dev/null +++ b/packages/ai-byok/README.md @@ -0,0 +1,187 @@ +# @tanstack/ai-byok + +Bring-your-own-key toolkit for [TanStack AI](https://tanstack.com/ai). Users +supply their own provider API keys; the library collects them **client-side**, +attaches them **per-request in a header**, and uses them **server-side without +ever persisting or logging them**. + +## First principle: never be a custodian + +- **Keys live client-side.** The browser is the system of record. +- **The server piece is a stateless pass-through.** It reads the key off the + incoming request header, hands it to the TanStack AI adapter for one call, and + never writes it to a DB, cache, log, or observability stream. +- **No central endpoint.** The server helper is trivially self-hostable; there + is no hardcoded relay URL. + +Keys travel in the `x-byok-` header — never the request body +or message history — so they stay out of persisted conversations and the +event/observability stream. + +## Client (React) + +```tsx +import { useRef } from 'react' +import { ByokProvider, ByokKeyManager, useByok } from '@tanstack/ai-byok/react' +import { withByok, defaultByokStorage } from '@tanstack/ai-byok/react' +import { fetchServerSentEvents } from '@tanstack/ai-client' +import { useChat } from '@tanstack/ai-react' + +// Wrap your app. `storage` is chosen once. defaultByokStorage() uses passkey +// encryption when supported — all keys, including OpenRouter PKCE. +function App({ children }) { + return {children} +} + +// Drop-in settings UI — shows the last 4 chars of a saved key, never the whole key. +function Settings() { + return +} + +// Attach the keys to the connection. `withByok` adds the per-request BYOK +// headers and, via `onMissingKey`, detects the relay's `byokMissing` 401 so you +// can prompt for (or unlock) the key the server was missing. +function Chat() { + const { keys, status, unlock } = useByok() + const keysRef = useRef(keys) + keysRef.current = keys + return useChat({ + connection: fetchServerSentEvents( + '/api/chat', + withByok(() => keysRef.current, { + onMissingKey: (provider) => { + if (status[provider]?.state === 'locked') void unlock() + else openKeyDialog(provider) + }, + }), + ), + }) +} +``` + +For a lower-level path, `byokHeaders(keys)` returns the raw header map and +`byokFetch(onMissingKey)` wraps a `fetch` to detect the `byokMissing` 401. + +### OpenRouter PKCE sign-in + +OpenRouter users can connect in one click via OAuth PKCE instead of pasting a +key. Import from `@tanstack/ai-byok/openrouter` (and +`@tanstack/ai-byok/openrouter/react` for `useOpenRouterPkce`) — the returned key +is stored as `openrouter` in the passkey-encrypted keyring like any other BYOK +key when using `defaultByokStorage()`. + +### Fetcher transport (`useChat`/`useGeneration` with a `fetcher`) + +`withByok` targets the `connection` transport. For the `fetcher` transport — +used by `useGeneration` (image/audio/video/speech/transcribe) and by +`useChat({ fetcher })` — use `byokFetcher`. It hands your fetcher body BYOK +`headers` and a missing-key-aware `fetch`, read fresh on every call, and works +for both fetcher styles: + +```ts +// A) fetch-based fetcher — spread headers, use the wrapped fetch for onMissingKey +useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers, fetch, signal }) => + fetch('/api/generate/audio', { + method: 'POST', + headers: { 'content-type': 'application/json', ...headers }, + body: JSON.stringify(input), + signal, + }), + { onMissingKey: (provider) => openKeyDialog(provider) }, + ), +}) + +// B) TanStack Start server function — forward headers at the call site +useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers }) => generateAudioFn({ data: input, headers }), + ), +}) +``` + +The key still travels in the `x-byok-` header, never the `data` +body. On the server, read it off the request with the same `getByokKey`: + +```ts +import { getByokKey, byokMissing } from '@tanstack/ai-byok/server' +import { getRequest } from '@tanstack/react-start/server' + +export const generateAudioFn = createServerFn({ method: 'POST' }) + .inputValidator((data: { prompt: string; provider: ProviderId }) => data) + .handler(({ data }) => { + const apiKey = getByokKey(getRequest(), data.provider) + if (!apiKey) return byokMissing(data.provider) + // ...hand apiKey to the adapter for one call + }) +``` + +The `onMissingKey` callback fires only on the fetch-based path (style A), where +`byokFetcher` owns the `fetch` and can see the `byokMissing` 401. A server +function (style B) surfaces the missing key as a thrown error instead — catch +it and inspect the message, or pre-check with `hasKey(provider)` before calling. + +## Server (stateless — no persist, no log) + +```ts +import { getByokKey, byokMissing } from '@tanstack/ai-byok/server' +import { chat, toServerSentEventsResponse } from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai/adapters' + +export async function POST(request: Request) { + const { messages } = await request.json() + + const apiKey = getByokKey(request, 'openai') + if (!apiKey) return byokMissing('openai') + + const stream = chat({ + adapter: createOpenaiChat('gpt-5.2', apiKey), + messages, + }) + return toServerSentEventsResponse(stream) +} +``` + +## Storage + +Pass one `storage` to ``. Built-ins: + +- **`defaultByokStorage()` (recommended)** — passkey-encrypted when supported; + session memory fallback. All keys (pasted and OpenRouter PKCE) share one tier. +- **`memoryStorage()` (default alone)** — session / in-memory. All keys vanish on + refresh. Zero at-rest liability. +- **`passkeyStorage()`** — encrypted at rest. The keyring is stored in + IndexedDB as AES-256-GCM ciphertext, with the key derived from a passkey's + WebAuthn **PRF** output (via HKDF) and unwrapped on demand with a + biometric/PIN tap. Fully client-side — no server, no custodian. + + Feature-detect with `isPasskeyStorageSupported()` and fall back to + `memoryStorage()` (PRF is solid on Android and Apple platform authenticators; + patchy on Firefox / roaming keys on Safari). + + **Honest scope:** this protects against at-rest theft (stolen device, + storage-dumping extension, backups). It does **not** defeat live in-page XSS — + an attacker running JS in the origin after you unlock can read the decrypted + keys from memory. `` states this while passkey storage is + active. + +`passkeyStorage` is `unlockable`, so `` does not decrypt on mount: +`useByok()` reports `locked: true` until you call `unlock()` (or save a key, +which registers a passkey on first use). It does, however, `peek()` on mount — +reading an unencrypted `provider → last-4` sidecar with **no** unlock ceremony — +so after a refresh the UI can immediately show saved keys as `locked` (with +their last-4) before the biometric tap. `KeyringStorage` is a small interface, +so you can supply your own strategy. + +Recovery is a non-issue: lose the device, re-paste the key from the provider +dashboard. + +## Validation + +`validateKey(provider, key)` pings the provider's cheapest authenticated +endpoint to confirm a key works before the user hits a wall mid-stream. It +returns `'valid' | 'invalid' | 'unsupported'`, and **throws** on a network/CORS +failure rather than guessing — some providers block browser origins entirely. diff --git a/packages/ai-byok/package.json b/packages/ai-byok/package.json new file mode 100644 index 000000000..a88b3f7a1 --- /dev/null +++ b/packages/ai-byok/package.json @@ -0,0 +1,88 @@ +{ + "name": "@tanstack/ai-byok", + "version": "0.1.0", + "description": "Bring-your-own-key toolkit for TanStack AI: client keyring, request headers, and stateless server helpers that never persist or log provider keys.", + "author": "Tanner Linsley", + "license": "MIT", + "homepage": "https://tanstack.com/ai", + "repository": { + "type": "git", + "url": "git+https://github.com/TanStack/ai.git", + "directory": "packages/ai-byok" + }, + "bugs": { + "url": "https://github.com/TanStack/ai/issues" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/tannerlinsley" + }, + "type": "module", + "module": "./dist/esm/index.js", + "types": "./dist/esm/index.d.ts", + "exports": { + ".": { + "types": "./dist/esm/index.d.ts", + "import": "./dist/esm/index.js" + }, + "./server": { + "types": "./dist/esm/server.d.ts", + "import": "./dist/esm/server.js" + }, + "./react": { + "types": "./dist/esm/react.d.ts", + "import": "./dist/esm/react.js" + }, + "./openrouter": { + "types": "./dist/esm/openrouter.d.ts", + "import": "./dist/esm/openrouter.js" + }, + "./openrouter/react": { + "types": "./dist/esm/openrouter-react.d.ts", + "import": "./dist/esm/openrouter-react.js" + } + }, + "files": [ + "dist", + "src" + ], + "scripts": { + "clean": "premove ./build ./dist", + "lint:fix": "eslint ./src --fix", + "test:eslint": "eslint ./src", + "test:lib": "vitest run", + "test:lib:dev": "pnpm test:lib --watch", + "test:types": "tsc", + "test:build": "publint --strict", + "build": "vite build" + }, + "keywords": [ + "ai", + "ai-sdk", + "typescript", + "tanstack", + "byok", + "api-key", + "react" + ], + "peerDependencies": { + "@types/react": ">=18.0.0", + "react": ">=18.0.0" + }, + "peerDependenciesMeta": { + "@types/react": { + "optional": true + }, + "react": { + "optional": true + } + }, + "devDependencies": { + "@testing-library/react": "^16.3.0", + "@types/react": "^19.2.7", + "@vitest/coverage-v8": "4.0.14", + "jsdom": "^27.2.0", + "react": "^19.2.3", + "vite": "^7.3.3" + } +} diff --git a/packages/ai-byok/src/client/keyring.ts b/packages/ai-byok/src/client/keyring.ts new file mode 100644 index 000000000..c9c27e9be --- /dev/null +++ b/packages/ai-byok/src/client/keyring.ts @@ -0,0 +1,43 @@ +import { PROVIDER_IDS, byokHeaderName, isProviderId } from '../shared/providers' +import type { ProviderId } from '../shared/providers' + +/** + * In-memory map of `provider → key`. The browser is the system of record for + * keys; this is the shape held in client state and persisted (when a + * persisting storage tier is selected). + */ +export type Keyring = Partial> + +/** Keep only known providers with non-empty string keys. */ +export function sanitizeKeyring(value: unknown): Keyring { + if (typeof value !== 'object' || value === null) return {} + const keys: Keyring = {} + for (const [provider, key] of Object.entries(value)) { + if (isProviderId(provider) && typeof key === 'string' && key.length > 0) { + keys[provider] = key + } + } + return keys +} + +/** + * Turns a keyring into request headers for the connection layer — one header + * per present provider. Absent/empty keys are skipped. + * + * @example + * ```ts + * useChat({ + * connection: fetchServerSentEvents('/api/chat', { + * headers: byokHeaders(keys), + * }), + * }) + * ``` + */ +export function byokHeaders(keys: Keyring): Record { + const headers: Record = {} + for (const provider of PROVIDER_IDS) { + const key = keys[provider] + if (key) headers[byokHeaderName(provider)] = key + } + return headers +} diff --git a/packages/ai-byok/src/client/openrouter-pkce.ts b/packages/ai-byok/src/client/openrouter-pkce.ts new file mode 100644 index 000000000..7b27083ec --- /dev/null +++ b/packages/ai-byok/src/client/openrouter-pkce.ts @@ -0,0 +1,290 @@ +/** + * OpenRouter OAuth PKCE — one-click login that yields a user-controlled API key. + * + * Flow (https://openrouter.ai/docs/guides/overview/auth/oauth): + * 1. Redirect the user to `openrouter.ai/auth` with `callback_url` and optional + * S256 `code_challenge`. + * 2. OpenRouter redirects back with a `code` query param. + * 3. POST the code (and `code_verifier` when a challenge was used) to + * `/api/v1/auth/keys` to receive `{ key }`. + * + * The returned key is stored via `setKey('openrouter', key)` in the configured + * keyring (passkey by default via {@link defaultByokStorage}). + */ + +const AUTH_ORIGIN = 'https://openrouter.ai' +const AUTH_PATH = '/auth' +const KEYS_URL = `${AUTH_ORIGIN}/api/v1/auth/keys` +const PENDING_STORAGE_KEY = 'byok:openrouter:pkce:v1' + +export type OpenRouterPkceChallengeMethod = 'S256' | 'plain' + +/** Pending PKCE state between the outbound redirect and the callback. */ +export interface OpenRouterPkcePending { + codeVerifier: string + codeChallengeMethod: OpenRouterPkceChallengeMethod + callbackUrl: string +} + +export interface OpenRouterAuthUrlOptions { + /** Where OpenRouter redirects after authorization (your app URL). */ + callbackUrl: string + /** S256 challenge (recommended). Omit for no PKCE challenge. */ + codeChallenge?: string + codeChallengeMethod?: OpenRouterPkceChallengeMethod +} + +export interface StartOpenRouterPkceOptions { + callbackUrl: string + /** + * Use S256 PKCE (recommended). When `false`, no challenge is sent — OpenRouter + * still returns a code, but the exchange omits verifier fields. + */ + useS256?: boolean + /** Override `window.location.assign` (for tests). */ + navigate?: (url: string) => void +} + +export interface ExchangeOpenRouterCodeOptions { + code: string + codeVerifier?: string + codeChallengeMethod?: OpenRouterPkceChallengeMethod + fetchImpl?: typeof fetch +} + +export interface CompleteOpenRouterPkceFromUrlOptions { + /** Defaults to `window.location.href`. */ + url?: string + fetchImpl?: typeof fetch + /** Clear pending session state after a successful exchange. Default `true`. */ + clearPending?: boolean + /** Strip `code` from the browser URL after success. Default `true` in browser. */ + cleanUrl?: boolean +} + +const VERIFIER_CHARS = + 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-._~' + +/** URL-safe random string for PKCE `code_verifier` (RFC 7636). */ +export function generateCodeVerifier(length = 64): string { + const bytes = crypto.getRandomValues(new Uint8Array(length)) + let out = '' + for (let i = 0; i < length; i++) { + out += VERIFIER_CHARS[(bytes[i] ?? 0) % VERIFIER_CHARS.length] + } + return out +} + +function base64UrlEncode(bytes: Uint8Array): string { + let binary = '' + for (const byte of bytes) binary += String.fromCharCode(byte) + return btoa(binary).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '') +} + +/** S256 `code_challenge` from a `code_verifier` (base64url-encoded SHA-256). */ +export async function createS256CodeChallenge( + codeVerifier: string, +): Promise { + const digest = await crypto.subtle.digest( + 'SHA-256', + new TextEncoder().encode(codeVerifier), + ) + return base64UrlEncode(new Uint8Array(digest)) +} + +/** Build the OpenRouter `/auth` URL that starts the PKCE redirect. */ +export function buildOpenRouterAuthUrl( + options: OpenRouterAuthUrlOptions, +): string { + const url = new URL(AUTH_PATH, AUTH_ORIGIN) + url.searchParams.set('callback_url', options.callbackUrl) + if (options.codeChallenge) { + url.searchParams.set('code_challenge', options.codeChallenge) + url.searchParams.set( + 'code_challenge_method', + options.codeChallengeMethod ?? 'S256', + ) + } + return url.toString() +} + +function sessionStorage(): Storage | null { + if (typeof globalThis.sessionStorage === 'undefined') return null + return globalThis.sessionStorage +} + +/** Persist pending PKCE state until the callback completes. */ +export function storeOpenRouterPkcePending( + pending: OpenRouterPkcePending, +): void { + sessionStorage()?.setItem(PENDING_STORAGE_KEY, JSON.stringify(pending)) +} + +/** Read pending PKCE state, or `null` when absent or corrupt. */ +export function loadOpenRouterPkcePending(): OpenRouterPkcePending | null { + const raw = sessionStorage()?.getItem(PENDING_STORAGE_KEY) + if (!raw) return null + try { + const parsed: unknown = JSON.parse(raw) + if (typeof parsed !== 'object' || parsed === null) return null + const { codeVerifier, codeChallengeMethod, callbackUrl } = parsed as { + codeVerifier?: unknown + codeChallengeMethod?: unknown + callbackUrl?: unknown + } + if ( + typeof codeVerifier !== 'string' || + (codeChallengeMethod !== 'S256' && codeChallengeMethod !== 'plain') || + typeof callbackUrl !== 'string' + ) { + return null + } + return { codeVerifier, codeChallengeMethod, callbackUrl } + } catch { + return null + } +} + +export function clearOpenRouterPkcePending(): void { + sessionStorage()?.removeItem(PENDING_STORAGE_KEY) +} + +/** + * Default callback URL: current origin + pathname (no query/hash). Suitable for + * SPAs where OpenRouter appends `?code=…` on return. + */ +export function defaultOpenRouterCallbackUrl(): string { + if (typeof globalThis.location === 'undefined') return '' + return `${globalThis.location.origin}${globalThis.location.pathname}` +} + +/** + * Start the OpenRouter PKCE login: generate verifier, store pending state, and + * redirect the browser to OpenRouter's `/auth` page. + */ +export async function startOpenRouterPkceLogin( + options: StartOpenRouterPkceOptions, +): Promise { + const useS256 = options.useS256 !== false + const codeVerifier = generateCodeVerifier() + const codeChallengeMethod: OpenRouterPkceChallengeMethod = useS256 + ? 'S256' + : 'plain' + + let codeChallenge: string | undefined + if (useS256) { + codeChallenge = await createS256CodeChallenge(codeVerifier) + } + + storeOpenRouterPkcePending({ + codeVerifier, + codeChallengeMethod, + callbackUrl: options.callbackUrl, + }) + + const authUrl = buildOpenRouterAuthUrl({ + callbackUrl: options.callbackUrl, + codeChallenge, + codeChallengeMethod: useS256 ? 'S256' : undefined, + }) + + const navigate = + options.navigate ?? + ((url: string) => { + globalThis.location.assign(url) + }) + navigate(authUrl) +} + +/** + * Exchange an authorization `code` for a user-controlled OpenRouter API key. + */ +export async function exchangeOpenRouterCode( + options: ExchangeOpenRouterCodeOptions, +): Promise { + const fetchImpl = options.fetchImpl ?? fetch + const body: Record = { code: options.code } + if (options.codeVerifier) { + body.code_verifier = options.codeVerifier + if (options.codeChallengeMethod) { + body.code_challenge_method = options.codeChallengeMethod + } + } + + const response = await fetchImpl(KEYS_URL, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body), + }) + + if (!response.ok) { + let detail = `${response.status} ${response.statusText}` + try { + const err: unknown = await response.json() + if (typeof err === 'object' && err !== null && 'error' in err) { + detail = String(err.error) + } + } catch { + // use status line + } + throw new Error(`OpenRouter PKCE exchange failed: ${detail}`) + } + + const data: unknown = await response.json() + if ( + typeof data !== 'object' || + data === null || + typeof (data as { key?: unknown }).key !== 'string' || + !(data as { key: string }).key + ) { + throw new Error('OpenRouter PKCE exchange returned no key') + } + return (data as { key: string }).key +} + +/** Remove `code` from the current URL without reloading. */ +export function stripOpenRouterCodeFromUrl(href?: string): void { + if (typeof globalThis.history === 'undefined') return + const base = href ?? globalThis.location.href + const url = new URL(base) + if (!url.searchParams.has('code')) return + url.searchParams.delete('code') + const next = `${url.pathname}${url.search}${url.hash}` + globalThis.history.replaceState({}, '', next) +} + +/** + * If the current URL carries an OpenRouter `code` and pending PKCE state exists, + * exchange the code for an API key. Returns `null` when there is nothing to + * complete (no code, or no pending verifier). + */ +export async function completeOpenRouterPkceFromUrl( + options: CompleteOpenRouterPkceFromUrlOptions = {}, +): Promise { + const href = + options.url ?? + (typeof globalThis.location !== 'undefined' ? globalThis.location.href : '') + if (!href) return null + + const code = new URL(href).searchParams.get('code') + if (!code) return null + + const pending = loadOpenRouterPkcePending() + if (!pending) { + throw new Error( + 'OpenRouter authorization code present but PKCE session expired — try signing in again', + ) + } + + const key = await exchangeOpenRouterCode({ + code, + codeVerifier: pending.codeVerifier, + codeChallengeMethod: pending.codeChallengeMethod, + fetchImpl: options.fetchImpl, + }) + + if (options.clearPending !== false) clearOpenRouterPkcePending() + if (options.cleanUrl !== false) stripOpenRouterCodeFromUrl(href) + + return key +} diff --git a/packages/ai-byok/src/client/passkey.ts b/packages/ai-byok/src/client/passkey.ts new file mode 100644 index 000000000..4a02acc17 --- /dev/null +++ b/packages/ai-byok/src/client/passkey.ts @@ -0,0 +1,379 @@ +import { sanitizeKeyring } from './keyring' +import { memoryStorage } from './storage' +import type { Keyring } from './keyring' +import type { KeyPreview, KeyringStorage } from './storage' +import type { ProviderId } from '../shared/providers' + +/** + * Passkey-encrypted keyring storage (WebAuthn PRF → HKDF → AES-256-GCM). + * + * The keyring is encrypted at rest in IndexedDB with an AES-256-GCM key derived + * from a passkey's PRF output, unwrapped on demand with a biometric/PIN tap. + * Decryption happens entirely client-side with the user present — no server, + * no custodian. + * + * **Honest scope:** this protects against at-rest theft (stolen device, + * storage-dumping extension, backups). It does NOT defeat live in-page XSS — an + * attacker running JS in the origin after the user unlocks can read the + * decrypted keys from memory. + */ + +const STORE_NAME = 'keyring' +const RECORD_ID = 'default' +const HKDF_INFO = 'byok:keyring:v1' +const DEFAULT_DB = 'byok' + +interface StoredRecord { + id: string + /** The passkey's raw credential id, replayed in the unlock ceremony. */ + credentialId: ArrayBuffer + /** Fixed per-install PRF evaluation input (not secret). */ + salt: ArrayBuffer + /** AES-GCM initialization vector for this ciphertext. */ + iv: ArrayBuffer + /** Encrypted keyring JSON. */ + ciphertext: ArrayBuffer + /** + * Unencrypted presence metadata (`provider → last 4`). Non-sensitive, so it + * can be read via {@link KeyringStorage.peek} without an unlock ceremony to + * show saved keys as "locked" after a refresh. + */ + preview: KeyPreview +} + +/** Build the non-sensitive `provider → last 4` preview from a keyring. */ +function previewOf(keys: Keyring): KeyPreview { + const preview: KeyPreview = {} + for (const [provider, key] of Object.entries(keys)) { + if (key) preview[provider as ProviderId] = key.slice(-4) + } + return preview +} + +/** + * Whether the current environment exposes WebAuthn. Note that actual PRF + * support can only be confirmed during registration; `passkeyStorage` throws a + * descriptive error if the chosen authenticator does not support PRF, which the + * caller should catch and fall back to {@link memoryStorage}. + */ +export function isPasskeyStorageSupported(): boolean { + return ( + typeof globalThis !== 'undefined' && + typeof globalThis.PublicKeyCredential !== 'undefined' && + typeof globalThis.navigator !== 'undefined' && + typeof globalThis.navigator.credentials.create === 'function' + ) +} + +// --------------------------------------------------------------------------- +// Crypto (exported for testing; the WebAuthn ceremony below feeds `deriveAesKey`) +// --------------------------------------------------------------------------- + +/** Derive a non-extractable AES-256-GCM key from a 32-byte PRF output. */ +export async function deriveAesKey( + prfOutput: BufferSource, +): Promise { + const base = await crypto.subtle.importKey('raw', prfOutput, 'HKDF', false, [ + 'deriveKey', + ]) + return crypto.subtle.deriveKey( + { + name: 'HKDF', + hash: 'SHA-256', + salt: new Uint8Array(0), + info: new TextEncoder().encode(HKDF_INFO), + }, + base, + { name: 'AES-GCM', length: 256 }, + false, + ['encrypt', 'decrypt'], + ) +} + +export async function encryptKeyring( + key: CryptoKey, + keys: Keyring, +): Promise<{ iv: ArrayBuffer; ciphertext: ArrayBuffer }> { + const iv = crypto.getRandomValues(new Uint8Array(12)) + const plaintext = new TextEncoder().encode(JSON.stringify(keys)) + const ciphertext = await crypto.subtle.encrypt( + { name: 'AES-GCM', iv }, + key, + plaintext, + ) + return { iv: iv.buffer, ciphertext } +} + +export async function decryptKeyring( + key: CryptoKey, + iv: BufferSource, + ciphertext: BufferSource, +): Promise { + const plaintext = await crypto.subtle.decrypt( + { name: 'AES-GCM', iv }, + key, + ciphertext, + ) + const parsed: unknown = JSON.parse(new TextDecoder().decode(plaintext)) + return sanitizeKeyring(parsed) +} + +// --------------------------------------------------------------------------- +// IndexedDB +// --------------------------------------------------------------------------- + +function openDb(dbName: string): Promise { + return new Promise((resolve, reject) => { + const request = indexedDB.open(dbName, 1) + request.onupgradeneeded = () => { + request.result.createObjectStore(STORE_NAME, { keyPath: 'id' }) + } + request.onsuccess = () => resolve(request.result) + request.onerror = () => reject(request.error) + }) +} + +function idbGet(dbName: string): Promise { + return openDb(dbName).then( + (db) => + new Promise((resolve, reject) => { + const request = db + .transaction(STORE_NAME, 'readonly') + .objectStore(STORE_NAME) + .get(RECORD_ID) + request.onsuccess = () => + resolve((request.result as StoredRecord | undefined) ?? null) + request.onerror = () => reject(request.error) + }), + ) +} + +function idbPut(dbName: string, record: StoredRecord): Promise { + return openDb(dbName).then( + (db) => + new Promise((resolve, reject) => { + const tx = db.transaction(STORE_NAME, 'readwrite') + tx.objectStore(STORE_NAME).put(record) + tx.oncomplete = () => resolve() + tx.onerror = () => reject(tx.error) + }), + ) +} + +function idbClear(dbName: string): Promise { + return openDb(dbName).then( + (db) => + new Promise((resolve, reject) => { + const tx = db.transaction(STORE_NAME, 'readwrite') + tx.objectStore(STORE_NAME).delete(RECORD_ID) + tx.oncomplete = () => resolve() + tx.onerror = () => reject(tx.error) + }), + ) +} + +// --------------------------------------------------------------------------- +// WebAuthn ceremonies +// --------------------------------------------------------------------------- + +function requirePublicKeyCredential( + credential: Credential | null, + action: string, +): PublicKeyCredential { + if (!credential) throw new Error(`Passkey ${action} was cancelled`) + if (!(credential instanceof PublicKeyCredential)) { + throw new Error(`Unexpected credential type during ${action}`) + } + return credential +} + +async function registerPasskey( + rpName: string, + userName: string, + rpId?: string, +): Promise<{ + credentialId: ArrayBuffer + salt: Uint8Array + prf?: BufferSource +}> { + const salt = crypto.getRandomValues(new Uint8Array(32)) + const credential = requirePublicKeyCredential( + await navigator.credentials.create({ + publicKey: { + challenge: crypto.getRandomValues(new Uint8Array(32)), + // Omit `id` to let the browser bind the passkey to the current origin's + // effective domain; set it to scope across subdomains of a self-host. + rp: rpId ? { name: rpName, id: rpId } : { name: rpName }, + user: { + id: crypto.getRandomValues(new Uint8Array(16)), + name: userName, + displayName: userName, + }, + pubKeyCredParams: [ + { type: 'public-key', alg: -7 }, + { type: 'public-key', alg: -257 }, + ], + authenticatorSelection: { + residentKey: 'required', + userVerification: 'required', + }, + extensions: { prf: { eval: { first: salt } } }, + }, + }), + 'registration', + ) + + const prf = credential.getClientExtensionResults().prf + if (!prf?.enabled) { + throw new Error( + 'This authenticator does not support the WebAuthn PRF extension', + ) + } + return { credentialId: credential.rawId, salt, prf: prf.results?.first } +} + +async function evaluatePrf( + credentialId: BufferSource, + salt: BufferSource, +): Promise { + const credential = requirePublicKeyCredential( + await navigator.credentials.get({ + publicKey: { + challenge: crypto.getRandomValues(new Uint8Array(32)), + allowCredentials: [{ type: 'public-key', id: credentialId }], + userVerification: 'required', + extensions: { prf: { eval: { first: salt } } }, + }, + }), + 'unlock', + ) + const result = credential.getClientExtensionResults().prf?.results?.first + if (!result) { + throw new Error('Authenticator did not return a PRF result') + } + return result +} + +// --------------------------------------------------------------------------- +// Storage strategy +// --------------------------------------------------------------------------- + +export interface PasskeyStorageOptions { + /** Relying-party name shown in the passkey prompt. */ + rpName?: string + /** Username label attached to the created passkey. */ + userName?: string + /** + * WebAuthn Relying Party ID. Omit to bind the passkey to the current origin's + * effective domain (the default — no central/hardcoded domain). Set it to a + * registrable parent domain to share the credential across subdomains of your + * own deployment. The encrypted keyring is never portable across unrelated + * domains. + */ + rpId?: string + /** IndexedDB database name. Defaults to `byok`. */ + dbName?: string +} + +/** + * Passkey-encrypted persistence. `` treats this as `unlockable`, + * so nothing is decrypted until the user calls `unlock()` (or saves a key, + * which registers a passkey on first use). The derived key is cached in memory + * for the session so repeated saves don't re-prompt. + */ +export function passkeyStorage( + options: PasskeyStorageOptions = {}, +): KeyringStorage { + const rpName = options.rpName ?? 'BYOK' + const userName = options.userName ?? 'byok-keyring' + const { rpId } = options + const dbName = options.dbName ?? DEFAULT_DB + + let cachedKey: CryptoKey | null = null + let cachedMeta: { + credentialId: ArrayBuffer + salt: Uint8Array + } | null = null + + // Obtain the AES key, running exactly one WebAuthn ceremony if it isn't + // already cached for this session (unlock if a passkey exists, else register). + async function ensureKey(): Promise<{ + key: CryptoKey + credentialId: ArrayBuffer + salt: Uint8Array + }> { + if (cachedKey && cachedMeta) { + return { key: cachedKey, ...cachedMeta } + } + const existing = await idbGet(dbName) + if (existing) { + const prf = await evaluatePrf(existing.credentialId, existing.salt) + cachedKey = await deriveAesKey(prf) + cachedMeta = { + credentialId: existing.credentialId, + salt: new Uint8Array(existing.salt), + } + } else { + const reg = await registerPasskey(rpName, userName, rpId) + const prf = reg.prf ?? (await evaluatePrf(reg.credentialId, reg.salt)) + cachedKey = await deriveAesKey(prf) + cachedMeta = { credentialId: reg.credentialId, salt: reg.salt } + } + return { key: cachedKey, ...cachedMeta } + } + + return { + id: 'passkey', + label: 'Passkey-encrypted (this device)', + persistent: true, + unlockable: true, + warning: + 'Keys are encrypted with your passkey and unlocked with biometrics. ' + + 'This protects saved keys if your device is stolen, but not against code ' + + 'running on this page after you unlock.', + peek: async () => { + // Read the unencrypted preview — no key material, no unlock ceremony. + const existing = await idbGet(dbName) + return existing?.preview ?? {} + }, + load: async () => { + const existing = await idbGet(dbName) + if (!existing) return {} // nothing stored yet — no ceremony needed + const { key } = await ensureKey() + return decryptKeyring(key, existing.iv, existing.ciphertext) + }, + save: async (keys) => { + const existing = await idbGet(dbName) + const hasKeys = Object.values(keys).some(Boolean) + // First save with an empty keyring is a no-op — avoids a passkey ceremony + // when another storage tier (e.g. OpenRouter PKCE in session memory) saves. + if (!hasKeys && !existing) return + + const { key, credentialId, salt } = await ensureKey() + const { iv, ciphertext } = await encryptKeyring(key, keys) + await idbPut(dbName, { + id: RECORD_ID, + credentialId, + salt: salt.buffer, + iv, + ciphertext, + preview: previewOf(keys), + }) + }, + clear: async () => { + cachedKey = null + cachedMeta = null + await idbClear(dbName) + }, + } +} + +/** + * Passkey-encrypted storage when supported, otherwise session memory. + * All keys — pasted or OpenRouter PKCE — use the same tier. + */ +export function defaultByokStorage( + options?: PasskeyStorageOptions, +): KeyringStorage { + return isPasskeyStorageSupported() ? passkeyStorage(options) : memoryStorage() +} diff --git a/packages/ai-byok/src/client/storage.ts b/packages/ai-byok/src/client/storage.ts new file mode 100644 index 000000000..de22d2827 --- /dev/null +++ b/packages/ai-byok/src/client/storage.ts @@ -0,0 +1,59 @@ +import type { Keyring } from './keyring' +import type { ProviderId } from '../shared/providers' + +/** + * Non-sensitive presence metadata: `provider → last 4 chars` of a stored key. + * Never contains the full key. + */ +export type KeyPreview = Partial> + +/** + * A client-side persistence strategy for the keyring. Methods may be sync or + * async so a strategy backed by async crypto (e.g. the passkey/PRF strategy) + * fits the same interface as the synchronous ones. + */ +export interface KeyringStorage { + /** A stable id for the strategy, surfaced in UI. */ + readonly id: string + /** Human-readable label. */ + readonly label: string + /** Whether keys written here survive a page refresh. `false` for memory. */ + readonly persistent: boolean + /** + * Whether this strategy requires an explicit unlock ceremony (e.g. a + * biometric tap) before stored keys can be read. When `true`, `` + * does not auto-load on mount — it waits for `unlock()`. + */ + readonly unlockable?: boolean + /** + * An honest, storage-specific caveat rendered by `` while this + * strategy is active (e.g. "protects at rest, not against in-page attacks"). + */ + readonly warning?: string + /** + * Report which providers have a stored key, and each key's last 4 chars, + * WITHOUT decrypting or running an unlock ceremony. Lets the UI surface saved + * keys as "locked" immediately after a refresh, before the user unlocks. Omit + * when not applicable (e.g. memory, which persists nothing). + */ + readonly peek?: () => Promise | KeyPreview + readonly load: () => Promise | Keyring + readonly save: (keys: Keyring) => Promise | void + readonly clear: () => Promise | void +} + +/** + * The default: session / in-memory. Keys live only in React state and vanish on + * refresh. Persists nothing, so it holds no state of its own — `load` always + * returns an empty keyring. Zero at-rest liability. + */ +export function memoryStorage(): KeyringStorage { + return { + id: 'memory', + label: 'Session only (not saved)', + persistent: false, + load: () => ({}), + save: () => {}, + clear: () => {}, + } +} diff --git a/packages/ai-byok/src/client/validate.ts b/packages/ai-byok/src/client/validate.ts new file mode 100644 index 000000000..9eee46468 --- /dev/null +++ b/packages/ai-byok/src/client/validate.ts @@ -0,0 +1,40 @@ +import { providerValidateConfig } from '../shared/providers' +import type { ProviderId } from '../shared/providers' + +/** + * Result of a key validation attempt. + * - `valid` — provider accepted the key. + * - `invalid` — provider rejected the key (401/403). + * - `unsupported` — provider has no browser-reachable validation endpoint. + */ +export type ValidationStatus = 'valid' | 'invalid' | 'unsupported' + +/** + * Pings the provider's cheapest authenticated endpoint (usually a models list) + * to confirm a key works before the user hits a wall mid-stream. + * + * Returns `'unsupported'` for providers with no browser-reachable endpoint. + * For any other failure (rate limit, 5xx, or a network/CORS error) this + * **throws** rather than guessing — the caller decides how to surface it. Note + * that some providers block browser origins entirely; a thrown `TypeError` + * ("Failed to fetch") is the honest signal there, not a silent `invalid`. + */ +export async function validateKey( + provider: ProviderId, + key: string, +): Promise { + const config = providerValidateConfig(provider) + if (!config) return 'unsupported' + + const response = await fetch(config.url, { + method: 'GET', + headers: config.headers(key), + }) + + if (response.ok) return 'valid' + if (response.status === 401 || response.status === 403) return 'invalid' + + throw new Error( + `Could not validate ${provider} key: ${response.status} ${response.statusText}`, + ) +} diff --git a/packages/ai-byok/src/client/with-byok.ts b/packages/ai-byok/src/client/with-byok.ts new file mode 100644 index 000000000..3fe06d6ac --- /dev/null +++ b/packages/ai-byok/src/client/with-byok.ts @@ -0,0 +1,150 @@ +import { isByokMissingBody } from '../shared/byok-missing' +import { byokHeaders } from './keyring' +import type { Keyring } from './keyring' +import type { ProviderId } from '../shared/providers' + +/** + * Wraps a `fetch` so a `byokMissing` 401 from the relay invokes `onMissingKey` + * with the provider that needs a key. The response is passed through unchanged — + * the connection still surfaces the error — so this only adds the side-channel + * callback the SSE error can't carry (it exposes only the status, not the body). + */ +export function byokFetch( + onMissingKey: (provider: ProviderId) => void, + fetchImpl: typeof fetch = fetch, +): typeof fetch { + return async (input, init) => { + const response = await fetchImpl(input, init) + if (response.status === 401) { + const body: unknown = await response + .clone() + .json() + .catch(() => null) + if (isByokMissingBody(body)) onMissingKey(body.error.provider) + } + return response + } +} + +export interface WithByokOptions { + /** Invoked when the relay reports a missing key (a `byokMissing` 401). */ + onMissingKey?: (provider: ProviderId) => void + /** Extra request headers; BYOK headers are merged on top. */ + headers?: Record + /** Underlying fetch (defaults to global `fetch`). */ + fetchClient?: typeof fetch +} + +/** The connection options `withByok` produces (a subset of a fetch adapter's). */ +export interface ByokConnectionOptions { + headers: Record + fetchClient?: typeof fetch +} + +/** Context a {@link byokFetcher} handler receives alongside the request input. */ +export interface ByokFetcherContext { + /** + * Per-provider BYOK headers read fresh for this request. Spread onto a + * `fetch` call's `headers`, or pass as a TanStack Start server function's + * call-site `headers` option — either way the key travels in the header, + * never the body. + */ + headers: Record + /** + * A `fetch` that also detects the relay's `byokMissing` 401 and invokes + * `onMissingKey` (identical to the global `fetch` when `onMissingKey` is + * unset). Only relevant to fetch-based fetchers; a server-function fetcher + * uses `headers` and surfaces a missing key as a thrown error instead. + */ + fetch: typeof fetch + /** The abort signal the transport forwards from `stop()`, when provided. */ + signal?: AbortSignal +} + +function buildByokFetch(options: WithByokOptions): typeof fetch { + return options.onMissingKey + ? byokFetch(options.onMissingKey, options.fetchClient) + : (options.fetchClient ?? fetch) +} + +/** Shared header + fetch wiring for {@link withByok} and {@link byokFetcher}. */ +export function buildByokRequestContext( + getKeys: () => Keyring, + options: WithByokOptions = {}, + signal?: AbortSignal, +): ByokFetcherContext { + return { + headers: { ...options.headers, ...byokHeaders(getKeys()) }, + fetch: buildByokFetch(options), + signal, + } +} + +/** + * Build BYOK connection options for a fetch-based connection adapter: attaches + * `byokHeaders(keys)` on every request and, when `onMissingKey` is set, detects + * the relay's `byokMissing` 401 so the UI can prompt for (or unlock) the key. + * + * Pass the result straight to the adapter. `getKeys` is read on every request, + * so back it with a ref/getter to stay current: + * + * ```ts + * useChat({ + * connection: fetchServerSentEvents( + * '/api/chat', + * withByok(() => keysRef.current, { + * onMissingKey: (provider) => openKeyDialog(provider), + * }), + * ), + * }) + * ``` + */ +export function withByok( + getKeys: () => Keyring, + options: WithByokOptions = {}, +): () => ByokConnectionOptions { + return () => { + const { headers, fetch: fetchClient } = buildByokRequestContext( + getKeys, + options, + ) + return { + headers, + ...(options.onMissingKey || options.fetchClient ? { fetchClient } : {}), + } + } +} + +/** + * The `useChat`/`useGeneration` **fetcher** counterpart to {@link withByok} + * (which targets the `connection` transport). Wraps a fetcher body so it + * receives BYOK `headers` and a missing-key-aware `fetch`, read fresh on every + * call. Works for both fetcher styles: + * + * ```ts + * // fetch-based: spread headers, use the wrapped fetch for onMissingKey + * fetcher: byokFetcher(() => keysRef.current, (input, { headers, fetch, signal }) => + * fetch('/api/generate/audio', { + * method: 'POST', + * headers: { 'content-type': 'application/json', ...headers }, + * body: JSON.stringify(input), + * signal, + * }), + * { onMissingKey: (provider) => openKeyDialog(provider) }, + * ) + * + * // TanStack Start server function: forward headers at the call site + * fetcher: byokFetcher(() => keysRef.current, (input, { headers }) => + * generateAudioFn({ data: input, headers }), + * ) + * // server handler: getByokKey(getRequest(), 'elevenlabs') + * ``` + */ +export function byokFetcher( + getKeys: () => Keyring, + handler: (input: TInput, context: ByokFetcherContext) => TReturn, + options: WithByokOptions = {}, +): (input: TInput, transport?: { signal?: AbortSignal }) => TReturn { + return (input, transport) => + handler(input, buildByokRequestContext(getKeys, options, transport?.signal)) +} diff --git a/packages/ai-byok/src/index.ts b/packages/ai-byok/src/index.ts new file mode 100644 index 000000000..69b51ff14 --- /dev/null +++ b/packages/ai-byok/src/index.ts @@ -0,0 +1,55 @@ +/** + * `@tanstack/ai-byok` — framework-agnostic BYOK client toolkit. + * + * Keys live client-side (the browser is the system of record) and travel to + * the relay in a per-provider header, never the request body. This entry has + * no framework or server dependencies; React bindings live in + * `@tanstack/ai-byok/react` and server helpers in `@tanstack/ai-byok/server`. + */ + +// Shared registry +export { + BYOK_PROVIDERS, + PROVIDER_IDS, + BYOK_HEADER_PREFIX, + byokHeaderName, + isProviderId, +} from './shared/providers' +export type { + ProviderId, + ProviderConfig, + ProviderValidateConfig, +} from './shared/providers' + +// Keyring + headers +export { byokHeaders, sanitizeKeyring } from './client/keyring' +export type { Keyring } from './client/keyring' + +// Connection helper: attach BYOK headers + detect a missing-key response +export { + withByok, + byokFetch, + byokFetcher, + buildByokRequestContext, +} from './client/with-byok' +export type { + WithByokOptions, + ByokConnectionOptions, + ByokFetcherContext, +} from './client/with-byok' +export { isByokMissingBody } from './shared/byok-missing' +export type { ByokMissingBody } from './shared/byok-missing' + +// Storage +export { memoryStorage } from './client/storage' +export type { KeyringStorage, KeyPreview } from './client/storage' +export { + passkeyStorage, + isPasskeyStorageSupported, + defaultByokStorage, +} from './client/passkey' +export type { PasskeyStorageOptions } from './client/passkey' + +// Validation +export { validateKey } from './client/validate' +export type { ValidationStatus } from './client/validate' diff --git a/packages/ai-byok/src/openrouter-react.ts b/packages/ai-byok/src/openrouter-react.ts new file mode 100644 index 000000000..fe9543790 --- /dev/null +++ b/packages/ai-byok/src/openrouter-react.ts @@ -0,0 +1,5 @@ +/** + * `@tanstack/ai-byok/openrouter/react` — React hook for OpenRouter PKCE login. + */ +export { useOpenRouterPkce } from './react/use-openrouter-pkce' +export type { UseOpenRouterPkceOptions } from './react/use-openrouter-pkce' diff --git a/packages/ai-byok/src/openrouter.ts b/packages/ai-byok/src/openrouter.ts new file mode 100644 index 000000000..d72096ed9 --- /dev/null +++ b/packages/ai-byok/src/openrouter.ts @@ -0,0 +1,29 @@ +/** + * `@tanstack/ai-byok/openrouter` — OpenRouter OAuth PKCE helpers. + * + * Optional vendor-specific login that yields a user-controlled API key stored + * via the core BYOK keyring. Import from this subpath only when you need + * OpenRouter sign-in; the core `@tanstack/ai-byok` package stays vendor-neutral. + */ + +export { + generateCodeVerifier, + createS256CodeChallenge, + buildOpenRouterAuthUrl, + storeOpenRouterPkcePending, + loadOpenRouterPkcePending, + clearOpenRouterPkcePending, + defaultOpenRouterCallbackUrl, + startOpenRouterPkceLogin, + exchangeOpenRouterCode, + stripOpenRouterCodeFromUrl, + completeOpenRouterPkceFromUrl, +} from './client/openrouter-pkce' +export type { + OpenRouterPkceChallengeMethod, + OpenRouterPkcePending, + OpenRouterAuthUrlOptions, + StartOpenRouterPkceOptions, + ExchangeOpenRouterCodeOptions, + CompleteOpenRouterPkceFromUrlOptions, +} from './client/openrouter-pkce' diff --git a/packages/ai-byok/src/react.ts b/packages/ai-byok/src/react.ts new file mode 100644 index 000000000..75bfd3806 --- /dev/null +++ b/packages/ai-byok/src/react.ts @@ -0,0 +1,48 @@ +/** + * `@tanstack/ai-byok/react` — React bindings for the BYOK keyring. + */ +export { ByokProvider, ByokContext } from './react/byok-context' +export type { + ByokProviderProps, + ByokContextValue, + KeyStatus, +} from './react/byok-context' +export { useByok } from './react/use-byok' +export { ByokKeyManager } from './react/byok-key-manager' +export type { ByokKeyManagerProps } from './react/byok-key-manager' +export { ByokKeyDialog } from './react/byok-key-dialog' +export type { ByokKeyDialogProps } from './react/byok-key-dialog' +export { ByokProviderRow } from './react/byok-provider-row' +export type { + ByokProviderRowProps, + ByokOpenRouterLoginProps, +} from './react/byok-provider-row' + +// Re-export the client toolkit so React consumers have one import path. +export { + byokHeaders, + withByok, + byokFetch, + byokFetcher, + buildByokRequestContext, + isByokMissingBody, + memoryStorage, + defaultByokStorage, + passkeyStorage, + isPasskeyStorageSupported, + validateKey, + BYOK_PROVIDERS, + PROVIDER_IDS, + byokHeaderName, + isProviderId, +} from './index' +export type { + Keyring, + KeyringStorage, + ProviderId, + ValidationStatus, + WithByokOptions, + ByokConnectionOptions, + ByokFetcherContext, + ByokMissingBody, +} from './index' diff --git a/packages/ai-byok/src/react/byok-context.tsx b/packages/ai-byok/src/react/byok-context.tsx new file mode 100644 index 000000000..bb2765a58 --- /dev/null +++ b/packages/ai-byok/src/react/byok-context.tsx @@ -0,0 +1,272 @@ +import { + createContext, + useCallback, + useEffect, + useMemo, + useRef, + useState, +} from 'react' +import { validateKey as pingProvider } from '../client/validate' +import { memoryStorage } from '../client/storage' +import { maskKey } from '../server/scrub' +import { PROVIDER_IDS } from '../shared/providers' +import type { ReactNode } from 'react' +import type { Keyring } from '../client/keyring' +import type { KeyringStorage } from '../client/storage' +import type { ProviderId } from '../shared/providers' +import type { ValidationStatus } from '../client/validate' + +/** + * Per-provider status surfaced to the UI. `masked` never contains more than the + * last 4 characters — the full key is never rendered back. + */ +export type KeyStatus = + | { state: 'empty' } + | { state: 'set'; masked: string } + // Stored but not yet decrypted this session (unlockable storage after a + // refresh). The key isn't usable until `unlock()`; only last-4 is known. + | { state: 'locked'; masked: string } + | { state: 'validating'; masked: string } + | { state: ValidationStatus; masked: string } + | { state: 'error'; masked: string; message: string } + +const EMPTY: KeyStatus = { state: 'empty' } + +export interface ByokContextValue { + /** + * The live keyring. Pass to `byokHeaders(keys)` when building the connection. + * The UI never renders these; treat them as write-only from the UI's view. + */ + keys: Keyring + /** Set (or overwrite) a provider's key and persist it to the configured storage. */ + setKey: (provider: ProviderId, key: string) => Promise + /** Remove a single provider's key. */ + clearKey: (provider: ProviderId) => Promise + /** Remove every key. */ + clearAll: () => Promise + /** + * Validate a key against the provider. Validates the given key, or the stored + * key when omitted. Records the outcome in {@link status} and returns it; + * never throws — a network/CORS failure is reported as an `error` status. + */ + validateKey: (provider: ProviderId, key?: string) => Promise + /** Per-provider status map. Providers with no key report `{ state: 'empty' }`. */ + status: Record + /** The configured persistence storage. */ + storage: KeyringStorage + /** + * `true` when an `unlockable` storage (e.g. passkey) may hold encrypted keys + * that haven't been decrypted this session. Always `false` for storage that + * needs no unlock ceremony (e.g. memory). + */ + locked: boolean + /** + * Decrypt and load keys from an `unlockable` storage, triggering its unlock + * ceremony (e.g. a biometric tap). No-op — and does not prompt — for storage + * that isn't unlockable or when nothing is stored yet. Rejects if the + * ceremony fails or is cancelled. + */ + unlock: () => Promise + hasKey: (provider: ProviderId) => boolean +} + +export const ByokContext = createContext(null) + +export interface ByokProviderProps { + children: ReactNode + /** + * Where keys are persisted. Defaults to the safest option + * ({@link memoryStorage}) — keys vanish on refresh and nothing is persisted. + * Fixed for the life of the provider. + */ + storage?: KeyringStorage +} + +export function ByokProvider({ + children, + storage: initialStorage, +}: ByokProviderProps) { + // Storage is chosen once and fixed for the life of the provider. + const [storage] = useState( + () => initialStorage ?? memoryStorage(), + ) + const [keys, setKeys] = useState({}) + const [statuses, setStatuses] = useState< + Partial> + >({}) + // Unlockable storage (passkey) starts locked; the user must unlock to decrypt. + const [locked, setLocked] = useState(() => Boolean(storage.unlockable)) + + // Keep a ref to the current keys so the persisting callbacks read the latest + // without re-creating on every keystroke. + const keysRef = useRef(keys) + keysRef.current = keys + + // Apply decrypted/loaded keys. Merge keys UNDER any edits the user made during + // the async load so an early setKey is never clobbered; promote a provider's + // status to `set` when it had no status or was a `locked` placeholder, while + // preserving a fresh user edit. + const applyLoaded = useCallback((loaded: Keyring) => { + setKeys((current) => ({ ...loaded, ...current })) + setStatuses((current) => { + const next = { ...current } + for (const [provider, key] of Object.entries(loaded)) { + if (!key) continue + const existing = current[provider as ProviderId] + if (!existing || existing.state === 'locked') { + next[provider as ProviderId] = { state: 'set', masked: maskKey(key) } + } + } + return next + }) + }, []) + + // On mount: unlockable storage peeks (no ceremony) to surface saved keys as + // `locked`; other storage auto-hydrates its keys. + useEffect(() => { + let cancelled = false + if (storage.unlockable) { + if (!storage.peek) return + void Promise.resolve(storage.peek()).then((preview) => { + if (cancelled) return + const present = Object.entries(preview).filter(([, last4]) => + Boolean(last4), + ) + setStatuses((current) => { + const next = { ...current } + for (const [provider, last4] of present) { + if (!next[provider as ProviderId]) { + next[provider as ProviderId] = { + state: 'locked', + masked: `…${last4}`, + } + } + } + return next + }) + // Nothing stored → nothing to unlock. + if (present.length === 0) setLocked(false) + }) + return () => { + cancelled = true + } + } + void Promise.resolve(storage.load()).then((loaded) => { + if (!cancelled) applyLoaded(loaded) + }) + return () => { + cancelled = true + } + }, [storage, applyLoaded]) + + const unlock = useCallback(async () => { + if (!storage.unlockable) return + applyLoaded(await Promise.resolve(storage.load())) + setLocked(false) + }, [storage, applyLoaded]) + + const persist = useCallback( + (next: Keyring) => Promise.resolve(storage.save(next)), + [storage], + ) + + const setKey = useCallback( + async (provider: ProviderId, key: string) => { + const next = { ...keysRef.current, [provider]: key } + setKeys(next) + setStatuses((prev) => ({ + ...prev, + [provider]: { state: 'set', masked: maskKey(key) }, + })) + await persist(next) + // A successful save ran any unlock/registration ceremony, so the session + // is now unlocked. + setLocked(false) + }, + [persist], + ) + + const clearKey = useCallback( + async (provider: ProviderId) => { + const next = { ...keysRef.current } + delete next[provider] + setKeys(next) + setStatuses((prev) => ({ ...prev, [provider]: EMPTY })) + await persist(next) + }, + [persist, storage], + ) + + const clearAll = useCallback(async () => { + setKeys({}) + setStatuses({}) + await Promise.resolve(storage.clear()) + setLocked(false) + }, [storage]) + + const validateKey = useCallback( + async (provider: ProviderId, key?: string): Promise => { + const target = key ?? keysRef.current[provider] + if (!target) { + setStatuses((prev) => ({ ...prev, [provider]: EMPTY })) + return EMPTY + } + const masked = maskKey(target) + setStatuses((prev) => ({ + ...prev, + [provider]: { state: 'validating', masked }, + })) + + let result: KeyStatus + try { + const status = await pingProvider(provider, target) + result = { state: status, masked } + } catch (error) { + result = { + state: 'error', + masked, + message: error instanceof Error ? error.message : String(error), + } + } + setStatuses((prev) => ({ ...prev, [provider]: result })) + return result + }, + [], + ) + + const status = useMemo(() => { + const full = {} as Record + for (const provider of PROVIDER_IDS) { + full[provider] = statuses[provider] ?? EMPTY + } + return full + }, [statuses]) + + const value = useMemo( + () => ({ + keys, + setKey, + clearKey, + clearAll, + validateKey, + status, + storage, + locked, + unlock, + hasKey: (provider) => Boolean(keys[provider]), + }), + [ + keys, + setKey, + clearKey, + clearAll, + validateKey, + status, + storage, + locked, + unlock, + ], + ) + + return {children} +} diff --git a/packages/ai-byok/src/react/byok-key-dialog.tsx b/packages/ai-byok/src/react/byok-key-dialog.tsx new file mode 100644 index 000000000..2c4ea0875 --- /dev/null +++ b/packages/ai-byok/src/react/byok-key-dialog.tsx @@ -0,0 +1,231 @@ +import { ByokKeyManager } from './byok-key-manager' +import { useByok } from './use-byok' +import type { ByokKeyManagerProps } from './byok-key-manager' +import type { CSSProperties, ReactNode } from 'react' +import type { ProviderId } from '../shared/providers' + +export interface ByokKeyDialogProps extends Omit< + ByokKeyManagerProps, + 'className' | 'style' +> { + /** Controlled open state. */ + open: boolean + onOpenChange: (open: boolean) => void + /** + * Provider for the currently selected model. When it needs a key, the trigger + * shows an attention indicator. + */ + activeProvider?: ProviderId | null + /** Custom trigger; defaults to a key-icon button with an attention dot. */ + trigger?: ReactNode + title?: string + description?: string + overlayClassName?: string + panelClassName?: string + panelStyle?: CSSProperties +} + +/** + * Modal wrapper around {@link ByokKeyManager} with a trigger button. Use when + * keys should live behind a dialog instead of inline on the page. + */ +export function ByokKeyDialog({ + open, + onOpenChange, + activeProvider = null, + envStatus, + trigger, + title = 'API keys', + description = 'Keys stay in your browser and are sent per-request in a header — never stored on the server.', + overlayClassName, + panelClassName, + panelStyle, + ...managerProps +}: ByokKeyDialogProps) { + const { keys } = useByok() + + const activeNeedsKey = + activeProvider != null && + !envStatus?.[activeProvider] && + !keys[activeProvider] + + return ( + <> + {trigger ?? ( + + )} + + {open ? ( +
onOpenChange(false)} + > +
event.stopPropagation()} + > +
+

+ {title} +

+ +
+ {description ? ( +

+ {description} +

+ ) : null} + +
+
+ ) : null} + + ) +} + +function KeyIcon() { + return ( + + + + + + ) +} + +const styles = { + trigger: { + position: 'relative', + display: 'inline-flex', + alignItems: 'center', + justifyContent: 'center', + padding: 8, + borderRadius: 8, + border: '1px solid #e5e7eb', + background: '#f9fafb', + color: '#374151', + cursor: 'pointer', + }, + attentionDot: { + position: 'absolute', + top: -2, + right: -2, + width: 10, + height: 10, + borderRadius: '50%', + background: '#fbbf24', + boxShadow: '0 0 0 2px #fff', + }, + overlay: { + position: 'fixed', + inset: 0, + zIndex: 50, + display: 'flex', + alignItems: 'center', + justifyContent: 'center', + padding: 16, + background: 'rgba(0, 0, 0, 0.6)', + }, + panel: { + width: '100%', + maxWidth: 480, + padding: 20, + borderRadius: 12, + border: '1px solid #e5e7eb', + background: '#fff', + boxShadow: '0 25px 50px -12px rgba(0, 0, 0, 0.25)', + }, + panelHeader: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + marginBottom: 4, + }, + title: { margin: 0, fontSize: 18, fontWeight: 600 }, + titleDark: { margin: 0, fontSize: 18, fontWeight: 600, color: '#fff' }, + closeButton: { + border: 'none', + background: 'transparent', + fontSize: 24, + lineHeight: 1, + cursor: 'pointer', + color: '#6b7280', + }, + closeButtonDark: { + border: 'none', + background: 'transparent', + fontSize: 24, + lineHeight: 1, + cursor: 'pointer', + color: '#9ca3af', + }, + description: { + margin: '0 0 16px', + fontSize: 14, + color: '#6b7280', + lineHeight: 1.4, + }, + descriptionDark: { + margin: '0 0 16px', + fontSize: 14, + color: '#9ca3af', + lineHeight: 1.4, + }, +} satisfies Record diff --git a/packages/ai-byok/src/react/byok-key-manager.tsx b/packages/ai-byok/src/react/byok-key-manager.tsx new file mode 100644 index 000000000..a491cdfe2 --- /dev/null +++ b/packages/ai-byok/src/react/byok-key-manager.tsx @@ -0,0 +1,142 @@ +import { useState } from 'react' +import { PROVIDER_IDS } from '../shared/providers' +import { useByok } from './use-byok' +import { ByokProviderRow } from './byok-provider-row' +import type { + ByokOpenRouterLoginProps, + ByokUiVariant, +} from './byok-provider-row' +import type { CSSProperties } from 'react' +import type { ProviderId } from '../shared/providers' + +export interface ByokKeyManagerProps { + /** Providers to show. Defaults to every registered provider. */ + providers?: Array + /** + * Which providers already have a server env key (booleans only — never the + * key values). Rows show "Server key" when true. + */ + envStatus?: Partial> + /** Highlight a provider row (e.g. after a missing-key prompt). */ + highlightProvider?: ProviderId | null + /** + * OpenRouter PKCE controls from `@tanstack/ai-byok/openrouter/react`. Omit + * when OpenRouter sign-in is not needed. + */ + openRouter?: ByokOpenRouterLoginProps & { + completing?: boolean + error?: string | null + } + variant?: ByokUiVariant + className?: string + style?: CSSProperties +} + +/** + * Drop-in settings UI for entering, validating, and clearing provider keys. + * + * Keys are write-only from this component's perspective: once saved, only the + * last 4 characters are ever shown. The full key is never rendered back. + */ +export function ByokKeyManager({ + providers = PROVIDER_IDS, + envStatus, + highlightProvider, + openRouter, + variant = 'light', + className, + style, +}: ByokKeyManagerProps) { + const { status, storage, locked, unlock } = useByok() + const [unlocking, setUnlocking] = useState(false) + const [unlockError, setUnlockError] = useState(null) + + return ( +
+ {locked ? ( +
+ Your saved keys are locked ({storage.label}). + +
+ ) : null} + {unlockError ?

{unlockError}

: null} + + {storage.warning ?

{storage.warning}

: null} + + {openRouter?.completing ? ( +

Completing OpenRouter sign-in…

+ ) : null} + {openRouter?.error ? ( +

{openRouter.error}

+ ) : null} + + {providers.map((provider) => ( + + ))} +
+ ) +} + +const styles = { + root: { + display: 'flex', + flexDirection: 'column', + gap: 16, + fontFamily: 'system-ui, sans-serif', + fontSize: 14, + maxWidth: 480, + }, + warning: { margin: 0, color: '#b45309', fontSize: 12, lineHeight: 1.4 }, + hint: { margin: 0, color: '#6b7280', fontSize: 12, lineHeight: 1.4 }, + unlockBanner: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + gap: 8, + padding: 12, + borderRadius: 8, + background: '#f3f4f6', + border: '1px solid #e5e7eb', + }, + primaryButton: { + padding: '6px 12px', + borderRadius: 6, + border: 'none', + background: '#111827', + color: '#fff', + cursor: 'pointer', + }, +} satisfies Record diff --git a/packages/ai-byok/src/react/byok-provider-row.tsx b/packages/ai-byok/src/react/byok-provider-row.tsx new file mode 100644 index 000000000..35eb79b4b --- /dev/null +++ b/packages/ai-byok/src/react/byok-provider-row.tsx @@ -0,0 +1,291 @@ +import { useState } from 'react' +import { BYOK_PROVIDERS } from '../shared/providers' +import { useByok } from './use-byok' +import type { CSSProperties } from 'react' +import type { KeyStatus } from './byok-context' +import type { ProviderId } from '../shared/providers' + +export interface ByokOpenRouterLoginProps { + onLogin: () => void + completing?: boolean + error?: string | null +} + +export type ByokUiVariant = 'light' | 'dark' + +export interface ByokProviderRowProps { + provider: ProviderId + status: KeyStatus + /** When true, the provider has a server env key and needs no BYOK key. */ + hasEnvKey?: boolean + highlight?: boolean + openRouter?: ByokOpenRouterLoginProps + variant?: ByokUiVariant + styles?: ByokRowStyles +} + +export type ByokRowStyles = Record + +export function ByokProviderRow({ + provider, + status, + hasEnvKey = false, + highlight = false, + openRouter, + variant = 'light', + styles, +}: ByokProviderRowProps) { + const resolvedStyles = + styles ?? (variant === 'dark' ? darkStyles : lightStyles) + const { keys, setKey, clearKey, validateKey } = useByok() + const [draft, setDraft] = useState('') + const yourKey = keys[provider] + const hasKey = status.state !== 'empty' + const lockedLast4 = + status.state === 'locked' && 'masked' in status + ? status.masked.slice(-4) + : null + + return ( +
+
+ + {BYOK_PROVIDERS[provider].label} + + +
+ + {hasKey && 'masked' in status ? ( +
+ + {status.masked} + +
+ + +
+
+ ) : null} + + {provider === 'openrouter' && + openRouter && + !yourKey && + status.state === 'empty' ? ( + + ) : null} + +
{ + event.preventDefault() + if (!draft) return + void setKey(provider, draft) + setDraft('') + }} + > + setDraft(event.target.value)} + style={resolvedStyles.input} + /> + +
+
+ ) +} + +function ProviderPresenceBadge({ + status, + yourKey, + lockedLast4, + hasEnvKey, + styles, +}: { + status: KeyStatus + yourKey?: string + lockedLast4: string | null + hasEnvKey: boolean + styles: ByokRowStyles +}) { + if (yourKey) { + return Your key + } + if (lockedLast4) { + return Locked + } + if (hasEnvKey) { + return Server key + } + return +} + +function StatusBadge({ + status, + styles, +}: { + status: KeyStatus + styles: ByokRowStyles +}) { + const config: Record = { + empty: { label: 'Not set', color: '#9ca3af' }, + set: { label: 'Saved', color: '#6b7280' }, + locked: { label: 'Locked', color: '#d97706' }, + validating: { label: 'Validating…', color: '#d97706' }, + valid: { label: 'Valid', color: '#059669' }, + invalid: { label: 'Invalid', color: '#dc2626' }, + unsupported: { label: 'Cannot verify', color: '#9ca3af' }, + error: { label: 'Check failed', color: '#dc2626' }, + } + const { label, color } = config[status.state] + const title = status.state === 'error' ? status.message : undefined + return ( + + {label} + + ) +} + +export const lightStyles = { + row: { + display: 'flex', + flexDirection: 'column', + gap: 8, + padding: 12, + border: '1px solid #e5e7eb', + borderRadius: 8, + }, + rowHighlight: { + border: '1px solid #fbbf24', + boxShadow: '0 0 0 1px rgba(251, 191, 36, 0.4)', + }, + rowHeader: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + }, + provider: { fontWeight: 600 }, + savedRow: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + gap: 8, + }, + masked: { + fontFamily: 'ui-monospace, monospace', + color: '#374151', + letterSpacing: 1, + }, + actions: { display: 'flex', gap: 6 }, + inputRow: { display: 'flex', gap: 6 }, + input: { + flex: 1, + padding: '6px 8px', + borderRadius: 6, + border: '1px solid #d1d5db', + }, + badge: { fontSize: 12, fontWeight: 600 }, + button: { + padding: '4px 10px', + borderRadius: 6, + border: '1px solid #d1d5db', + background: '#fff', + cursor: 'pointer', + }, + primaryButton: { + padding: '6px 12px', + borderRadius: 6, + border: 'none', + background: '#111827', + color: '#fff', + cursor: 'pointer', + }, + oauthButton: { + width: '100%', + padding: '8px 12px', + borderRadius: 6, + border: '1px solid #d1d5db', + background: '#fff', + cursor: 'pointer', + fontWeight: 600, + fontSize: 13, + }, +} satisfies ByokRowStyles + +export const darkStyles: ByokRowStyles = { + ...lightStyles, + row: { + ...lightStyles.row, + border: '1px solid #374151', + background: 'rgba(31, 41, 55, 0.5)', + }, + rowHighlight: { + border: '1px solid rgba(251, 191, 36, 0.6)', + boxShadow: '0 0 0 1px rgba(251, 191, 36, 0.4)', + }, + provider: { ...lightStyles.provider, color: '#fff' }, + masked: { ...lightStyles.masked, color: '#d1d5db' }, + input: { + ...lightStyles.input, + border: '1px solid #4b5563', + background: '#111827', + color: '#fff', + }, + button: { + ...lightStyles.button, + border: '1px solid #4b5563', + background: '#1f2937', + color: '#e5e7eb', + }, + primaryButton: { + ...lightStyles.primaryButton, + background: '#ea580c', + }, + oauthButton: { + ...lightStyles.oauthButton, + border: '1px solid #4b5563', + background: '#1f2937', + color: '#fff', + }, +} diff --git a/packages/ai-byok/src/react/use-byok.ts b/packages/ai-byok/src/react/use-byok.ts new file mode 100644 index 000000000..7a6a8cbf6 --- /dev/null +++ b/packages/ai-byok/src/react/use-byok.ts @@ -0,0 +1,25 @@ +import { useContext } from 'react' +import { ByokContext } from './byok-context' +import type { ByokContextValue } from './byok-context' + +/** + * Access the BYOK keyring and controls. Must be called under a + * {@link ByokProvider}. + * + * @example + * ```tsx + * const { keys } = useByok() + * useChat({ + * connection: fetchServerSentEvents('/api/chat', { + * headers: byokHeaders(keys), + * }), + * }) + * ``` + */ +export function useByok(): ByokContextValue { + const context = useContext(ByokContext) + if (!context) { + throw new Error('useByok must be used within a ') + } + return context +} diff --git a/packages/ai-byok/src/react/use-openrouter-pkce.ts b/packages/ai-byok/src/react/use-openrouter-pkce.ts new file mode 100644 index 000000000..6f6bdf611 --- /dev/null +++ b/packages/ai-byok/src/react/use-openrouter-pkce.ts @@ -0,0 +1,75 @@ +import { useCallback, useEffect, useState } from 'react' +import { + completeOpenRouterPkceFromUrl, + defaultOpenRouterCallbackUrl, + startOpenRouterPkceLogin, +} from '../client/openrouter-pkce' // openrouter subpath re-exports these +import { useByok } from './use-byok' + +export interface UseOpenRouterPkceOptions { + /** + * Where OpenRouter redirects after login. Defaults to + * `origin + pathname` of the current page. + */ + callbackUrl?: string + /** + * When `true` (default), completes the flow on mount if the URL contains + * `?code=…` from an OpenRouter redirect. + */ + autoComplete?: boolean + /** Use S256 PKCE (recommended). Default `true`. */ + useS256?: boolean +} + +/** + * OpenRouter PKCE login for BYOK. On return from OpenRouter, exchanges the + * authorization code and saves the key via `setKey('openrouter', key)`. + */ +export function useOpenRouterPkce(options: UseOpenRouterPkceOptions = {}) { + const { setKey } = useByok() + const [completing, setCompleting] = useState(false) + const [error, setError] = useState(null) + + const callbackUrl = options.callbackUrl ?? defaultOpenRouterCallbackUrl() + + useEffect(() => { + if (options.autoComplete === false) return + if (typeof globalThis.location !== 'undefined') { + const code = new URL(globalThis.location.href).searchParams.get('code') + if (!code) return + } + let cancelled = false + setCompleting(true) + setError(null) + void completeOpenRouterPkceFromUrl() + .then((key) => { + if (cancelled || !key) return + return setKey('openrouter', key) + }) + .catch((err: unknown) => { + if (!cancelled) { + setError(err instanceof Error ? err.message : String(err)) + } + }) + .finally(() => { + if (!cancelled) setCompleting(false) + }) + return () => { + cancelled = true + } + }, [options.autoComplete, setKey]) + + const login = useCallback(async () => { + setError(null) + try { + await startOpenRouterPkceLogin({ + callbackUrl, + useS256: options.useS256, + }) + } catch (err: unknown) { + setError(err instanceof Error ? err.message : String(err)) + } + }, [callbackUrl, options.useS256]) + + return { login, completing, error, callbackUrl } +} diff --git a/packages/ai-byok/src/server.ts b/packages/ai-byok/src/server.ts new file mode 100644 index 000000000..f79c914f6 --- /dev/null +++ b/packages/ai-byok/src/server.ts @@ -0,0 +1,23 @@ +/** + * `@tanstack/ai-byok/server` — stateless server helpers for BYOK relays. + * + * The server piece reads a provider key off the incoming request header, hands + * it to the TanStack AI adapter for that one call, and never persists or logs + * it. It is trivially self-hostable: there is no central endpoint baked in. + */ +export { getByokKey } from './server/get-byok-key' +export { byokMissing, isByokMissingBody } from './server/byok-missing' +export type { ByokMissingBody } from './server/byok-missing' +export { preferByokAdapter, requireByokOrEnv } from './server/prefer-byok' +export { lastFour, maskKey, scrubSecrets } from './server/scrub' + +// Re-export the shared registry so a server can enumerate/validate provider ids +// without a second import path. +export { + BYOK_PROVIDERS, + PROVIDER_IDS, + BYOK_HEADER_PREFIX, + byokHeaderName, + isProviderId, +} from './shared/providers' +export type { ProviderId, ProviderConfig } from './shared/providers' diff --git a/packages/ai-byok/src/server/byok-missing.ts b/packages/ai-byok/src/server/byok-missing.ts new file mode 100644 index 000000000..b906181a9 --- /dev/null +++ b/packages/ai-byok/src/server/byok-missing.ts @@ -0,0 +1,32 @@ +import { isByokMissingBody } from '../shared/byok-missing' +import type { ByokMissingBody } from '../shared/byok-missing' +import type { ProviderId } from '../shared/providers' + +export type { ByokMissingBody } +export { isByokMissingBody } + +/** + * Builds a typed JSON error response telling the client which provider key is + * missing, so it can render an "add your `` key" prompt. Defaults to + * HTTP 401. Carries no key material. + */ +export function byokMissing( + provider: ProviderId, + init?: ResponseInit, +): Response { + const body: ByokMissingBody = { + error: { + type: 'byok_missing', + provider, + message: `Missing API key for "${provider}". Add your ${provider} key to continue.`, + }, + } + return new Response(JSON.stringify(body), { + status: 401, + ...init, + headers: { + 'content-type': 'application/json', + ...init?.headers, + }, + }) +} diff --git a/packages/ai-byok/src/server/get-byok-key.ts b/packages/ai-byok/src/server/get-byok-key.ts new file mode 100644 index 000000000..6057b5080 --- /dev/null +++ b/packages/ai-byok/src/server/get-byok-key.ts @@ -0,0 +1,21 @@ +import { byokHeaderName } from '../shared/providers' +import type { ProviderId } from '../shared/providers' + +/** + * Reads a provider's BYOK key off the incoming request header. + * + * Returns the key or `null` if absent. The key is read from the header only — + * never the body — so it stays out of any persisted `messages` array and out + * of the event/observability stream. This function does not log the value and + * must not be wrapped in anything that attaches it to a logger context. + * + * Accepts either a `Request` or any object exposing a `Headers`-like `get`, so + * it works across Fetch-API runtimes (Workers, Deno, Bun, Node/undici). + */ +export function getByokKey( + request: { headers: Pick }, + provider: ProviderId, +): string | null { + const value = request.headers.get(byokHeaderName(provider)) + return value && value.length > 0 ? value : null +} diff --git a/packages/ai-byok/src/server/prefer-byok.ts b/packages/ai-byok/src/server/prefer-byok.ts new file mode 100644 index 000000000..b65e922ee --- /dev/null +++ b/packages/ai-byok/src/server/prefer-byok.ts @@ -0,0 +1,42 @@ +import { getByokKey } from './get-byok-key' +import { byokMissing } from './byok-missing' +import type { ProviderId } from '../shared/providers' + +/** + * Prefer a per-request BYOK header key over a server env-configured adapter. + * + * @example + * ```ts + * adapter: preferByokAdapter(request, 'openai', model, { + * byok: createOpenaiChat, + * env: openaiText, + * }) + * ``` + */ +export function preferByokAdapter( + request: { headers: Pick }, + provider: ProviderId, + model: TModel, + factories: { + byok: (model: TModel, apiKey: string) => TAdapter + env: (model: TModel) => TAdapter + }, +): TAdapter { + const key = getByokKey(request, provider) + return key ? factories.byok(model, key) : factories.env(model) +} + +/** + * Return a typed `byokMissing` response when neither a BYOK header nor any of + * the named env vars is present. Returns `null` when the request may proceed. + */ +export function requireByokOrEnv( + request: { headers: Pick }, + provider: ProviderId, + envVarNames: ReadonlyArray, +): Response | null { + const hasByokKey = Boolean(getByokKey(request, provider)) + const hasEnvKey = envVarNames.some((name) => Boolean(process.env[name])) + if (!hasByokKey && !hasEnvKey) return byokMissing(provider) + return null +} diff --git a/packages/ai-byok/src/server/scrub.ts b/packages/ai-byok/src/server/scrub.ts new file mode 100644 index 000000000..3443f71f1 --- /dev/null +++ b/packages/ai-byok/src/server/scrub.ts @@ -0,0 +1,33 @@ +/** + * Helpers for keeping key material out of logs and error responses. + * + * The server piece is a stateless pass-through: it must never write a key to a + * DB, cache, log, or observability stream, and must never echo a full key back + * to a client. These utilities make the last-4 the only representation that + * ever leaves the process. + */ + +/** Returns the last 4 characters of a key for display, e.g. `"...a1b2"`. */ +export function lastFour(key: string): string { + return key.slice(-4) +} + +/** Masks a key to `"…last4"`, hiding everything but the trailing 4 chars. */ +export function maskKey(key: string): string { + if (key.length <= 4) return '…' + return `…${lastFour(key)}` +} + +/** + * Replaces every occurrence of each secret in `input` with its masked form. + * Use before logging or returning any string that may have interpolated a key + * (URLs, provider SDK error messages, stack traces). + */ +export function scrubSecrets(input: string, secrets: Array): string { + let output = input + for (const secret of secrets) { + if (!secret) continue + output = output.split(secret).join(maskKey(secret)) + } + return output +} diff --git a/packages/ai-byok/src/shared/byok-missing.ts b/packages/ai-byok/src/shared/byok-missing.ts new file mode 100644 index 000000000..a220a4797 --- /dev/null +++ b/packages/ai-byok/src/shared/byok-missing.ts @@ -0,0 +1,24 @@ +import { isProviderId } from './providers' +import type { ProviderId } from './providers' + +/** Discriminated body returned by {@link byokMissing}. */ +export interface ByokMissingBody { + error: { + type: 'byok_missing' + provider: ProviderId + message: string + } +} + +/** Type guard for a {@link ByokMissingBody} parsed from a response. */ +export function isByokMissingBody(value: unknown): value is ByokMissingBody { + if (typeof value !== 'object' || value === null) return false + const { error } = value as { error?: unknown } + if (typeof error !== 'object' || error === null) return false + const typed = error as { type?: unknown; provider?: unknown } + return ( + typed.type === 'byok_missing' && + typeof typed.provider === 'string' && + isProviderId(typed.provider) + ) +} diff --git a/packages/ai-byok/src/shared/providers.ts b/packages/ai-byok/src/shared/providers.ts new file mode 100644 index 000000000..47ffcbf29 --- /dev/null +++ b/packages/ai-byok/src/shared/providers.ts @@ -0,0 +1,140 @@ +/** + * Shared provider registry and header-naming convention. + * + * This module is fully isomorphic (no browser or Node APIs) and is the single + * source of truth imported by both the client keyring and the server helpers. + */ + +/** + * Describes how to validate a key against a provider's cheapest authenticated + * endpoint (typically a models list). `headers` builds the request headers + * from the raw key. Only providers that expose a browser-reachable endpoint + * carry a `validate` entry; the rest report `'unsupported'` from + * {@link validateKey}. + */ +export interface ProviderValidateConfig { + /** Endpoint hit with a GET request to confirm the key works. */ + url: string + /** Builds the auth headers for the validation request from the raw key. */ + headers: (key: string) => Record +} + +/** Static metadata for a supported provider. */ +export interface ProviderConfig { + /** Stable id used in the keyring map and the per-provider header name. */ + id: string + /** Human-readable label for UI. */ + label: string + /** Optional validation endpoint metadata. */ + validate?: ProviderValidateConfig +} + +const ANTHROPIC_HEADERS = (key: string): Record => ({ + 'x-api-key': key, + 'anthropic-version': '2023-06-01', + // Required for Anthropic to serve the request from a browser origin. + 'anthropic-dangerous-direct-browser-access': 'true', +}) + +const bearer = (key: string): Record => ({ + Authorization: `Bearer ${key}`, +}) + +/** + * Registry of providers the BYOK toolkit understands. `ProviderId` is derived + * from these keys, so adding a provider here extends the typed union + * everywhere. + */ +export const BYOK_PROVIDERS = { + openai: { + id: 'openai', + label: 'OpenAI', + validate: { url: 'https://api.openai.com/v1/models', headers: bearer }, + }, + anthropic: { + id: 'anthropic', + label: 'Anthropic', + validate: { + url: 'https://api.anthropic.com/v1/models', + headers: ANTHROPIC_HEADERS, + }, + }, + gemini: { + id: 'gemini', + label: 'Google Gemini', + validate: { + // Gemini authenticates the models list via a query param, not a header. + url: 'https://generativelanguage.googleapis.com/v1beta/models', + headers: (key) => ({ 'x-goog-api-key': key }), + }, + }, + openrouter: { + id: 'openrouter', + label: 'OpenRouter', + validate: { url: 'https://openrouter.ai/api/v1/key', headers: bearer }, + }, + groq: { + id: 'groq', + label: 'Groq', + validate: { + url: 'https://api.groq.com/openai/v1/models', + headers: bearer, + }, + }, + grok: { + id: 'grok', + label: 'xAI Grok', + validate: { url: 'https://api.x.ai/v1/models', headers: bearer }, + }, + mistral: { + id: 'mistral', + label: 'Mistral', + validate: { url: 'https://api.mistral.ai/v1/models', headers: bearer }, + }, + elevenlabs: { + id: 'elevenlabs', + label: 'ElevenLabs', + validate: { + url: 'https://api.elevenlabs.io/v1/user', + headers: (key) => ({ 'xi-api-key': key }), + }, + }, + // No browser-reachable validation endpoint (key-scheme auth, no models list). + fal: { id: 'fal', label: 'fal.ai' }, + // Local runtime, no API key involved. + ollama: { id: 'ollama', label: 'Ollama' }, +} as const satisfies Record + +/** Union of every supported provider id. */ +export type ProviderId = keyof typeof BYOK_PROVIDERS + +/** All provider ids as a runtime array. */ +export const PROVIDER_IDS = Object.keys(BYOK_PROVIDERS) as Array + +/** Type guard narrowing an arbitrary string to a known {@link ProviderId}. */ +export function isProviderId(value: string): value is ProviderId { + return value in BYOK_PROVIDERS +} + +/** + * The validation config for a provider, or `undefined` when the provider has no + * browser-reachable validation endpoint. + */ +export function providerValidateConfig( + provider: ProviderId, +): ProviderValidateConfig | undefined { + const config = BYOK_PROVIDERS[provider] + return 'validate' in config ? config.validate : undefined +} + +/** Prefix for every per-provider BYOK header. */ +export const BYOK_HEADER_PREFIX = 'x-byok-' + +/** + * The HTTP header name that carries the key for a given provider. Keys always + * travel in this header, never in the request body or message history, so they + * stay out of persisted conversations and the event/observability stream. + */ +export function byokHeaderName(provider: ProviderId): string { + return `${BYOK_HEADER_PREFIX}${provider}` +} diff --git a/packages/ai-byok/tests/byok.test.ts b/packages/ai-byok/tests/byok.test.ts new file mode 100644 index 000000000..99d5efc06 --- /dev/null +++ b/packages/ai-byok/tests/byok.test.ts @@ -0,0 +1,361 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { + byokFetch, + byokFetcher, + byokHeaders, + byokHeaderName, + buildByokRequestContext, + sanitizeKeyring, + withByok, +} from '../src/index' +import type { Keyring } from '../src/index' +import { + byokMissing, + getByokKey, + isByokMissingBody, + maskKey, + preferByokAdapter, + requireByokOrEnv, + scrubSecrets, +} from '../src/server' +import { memoryStorage } from '../src/client/storage' +import { + decryptKeyring, + deriveAesKey, + encryptKeyring, + defaultByokStorage, +} from '../src/client/passkey' +import { + buildOpenRouterAuthUrl, + clearOpenRouterPkcePending, + completeOpenRouterPkceFromUrl, + createS256CodeChallenge, + exchangeOpenRouterCode, + generateCodeVerifier, + loadOpenRouterPkcePending, + storeOpenRouterPkcePending, +} from '../src/openrouter' + +describe('byokHeaders', () => { + it('emits one header per present provider and skips empty keys', () => { + const headers = byokHeaders({ + openai: 'sk-abc', + anthropic: '', + gemini: 'g-1', + }) + expect(headers).toEqual({ + 'x-byok-openai': 'sk-abc', + 'x-byok-gemini': 'g-1', + }) + }) +}) + +describe('sanitizeKeyring', () => { + it('keeps only known providers with non-empty string keys', () => { + expect( + sanitizeKeyring({ + openai: 'sk-1', + anthropic: '', + bogus: 'nope', + gemini: 42, + }), + ).toEqual({ openai: 'sk-1' }) + }) +}) + +describe('withByok / byokFetch / buildByokRequestContext', () => { + it('withByok attaches BYOK headers read fresh at request time', () => { + let keys: Keyring = { openai: 'sk-a' } + const build = withByok(() => keys) + expect(build().headers).toEqual({ 'x-byok-openai': 'sk-a' }) + keys = { openai: 'sk-b', gemini: 'g-1' } + expect(build().headers).toEqual({ + 'x-byok-openai': 'sk-b', + 'x-byok-gemini': 'g-1', + }) + }) + + it('buildByokRequestContext shares headers and fetch between helpers', async () => { + const onMissingKey = vi.fn() + const fetchImpl = vi.fn(async () => byokMissing('anthropic')) + const ctx = buildByokRequestContext(() => ({ openai: 'sk-a' }), { + onMissingKey, + fetchClient: fetchImpl, + }) + expect(ctx.headers).toEqual({ 'x-byok-openai': 'sk-a' }) + await ctx.fetch('https://x.test') + expect(onMissingKey).toHaveBeenCalledWith('anthropic') + }) + + it('byokFetch fires onMissingKey with the provider on a byokMissing 401', async () => { + const onMissingKey = vi.fn() + const fetchImpl = vi.fn(async () => byokMissing('anthropic')) + const wrapped = byokFetch(onMissingKey, fetchImpl) + await wrapped('https://x.test') + expect(onMissingKey).toHaveBeenCalledWith('anthropic') + }) + + it('byokFetch ignores non-401 and non-byok responses', async () => { + const onMissingKey = vi.fn() + const wrapped = byokFetch( + onMissingKey, + async () => new Response('ok', { status: 200 }), + ) + await wrapped('https://x.test') + expect(onMissingKey).not.toHaveBeenCalled() + }) + + it('isByokMissingBody rejects malformed provider ids', () => { + expect( + isByokMissingBody({ + error: { type: 'byok_missing', provider: 'not-a-provider' }, + }), + ).toBe(false) + }) +}) + +describe('byokFetcher', () => { + it('hands the handler fresh headers and forwards the transport signal', () => { + let keys: Keyring = { openai: 'sk-a' } + const handler = vi.fn((_input: { prompt: string }, ctx) => ctx) + const fetcher = byokFetcher(() => keys, handler) + + const first = fetcher({ prompt: 'hi' }) + expect(first.headers).toEqual({ 'x-byok-openai': 'sk-a' }) + expect(first.fetch).toBe(fetch) + expect(first.signal).toBeUndefined() + + keys = { openai: 'sk-b', gemini: 'g-1' } + const controller = new AbortController() + const second = fetcher({ prompt: 'yo' }, { signal: controller.signal }) + expect(second.headers).toEqual({ + 'x-byok-openai': 'sk-b', + 'x-byok-gemini': 'g-1', + }) + expect(second.signal).toBe(controller.signal) + }) + + it('supplies a missing-key-aware fetch when onMissingKey is set', async () => { + const onMissingKey = vi.fn() + const fetchImpl = vi.fn(async () => byokMissing('anthropic')) + const fetcher = byokFetcher( + () => ({ anthropic: 'sk-x' }), + (_input: null, ctx) => ctx.fetch('https://x.test'), + { onMissingKey, fetchClient: fetchImpl }, + ) + await fetcher(null) + expect(onMissingKey).toHaveBeenCalledWith('anthropic') + }) +}) + +describe('getByokKey', () => { + it('reads the key from the header, returns null when absent', () => { + const request = new Request('https://x.test', { + headers: { [byokHeaderName('openai')]: 'sk-live' }, + }) + expect(getByokKey(request, 'openai')).toBe('sk-live') + expect(getByokKey(request, 'anthropic')).toBeNull() + }) +}) + +describe('preferByokAdapter', () => { + it('uses the BYOK header when present, otherwise the env factory', () => { + const withKey = new Request('https://x.test', { + headers: { [byokHeaderName('openai')]: 'sk-byok' }, + }) + const withoutKey = new Request('https://x.test') + const byok = vi.fn((model: string, key: string) => ({ + kind: 'byok', + model, + key, + })) + const env = vi.fn((model: string) => ({ kind: 'env', model })) + + expect( + preferByokAdapter(withKey, 'openai', 'gpt-5.2', { byok, env }), + ).toEqual({ kind: 'byok', model: 'gpt-5.2', key: 'sk-byok' }) + expect( + preferByokAdapter(withoutKey, 'openai', 'gpt-5.2', { byok, env }), + ).toEqual({ kind: 'env', model: 'gpt-5.2' }) + }) +}) + +describe('requireByokOrEnv', () => { + const originalOpenAi = process.env.OPENAI_API_KEY + + afterEach(() => { + if (originalOpenAi === undefined) delete process.env.OPENAI_API_KEY + else process.env.OPENAI_API_KEY = originalOpenAi + }) + + it('returns byokMissing when neither header nor env is present', async () => { + delete process.env.OPENAI_API_KEY + const request = new Request('https://x.test') + const blocked = requireByokOrEnv(request, 'openai', ['OPENAI_API_KEY']) + expect(blocked?.status).toBe(401) + const body: unknown = await blocked!.json() + expect(isByokMissingBody(body)).toBe(true) + }) + + it('returns null when an env var is configured', () => { + process.env.OPENAI_API_KEY = 'sk-env' + const request = new Request('https://x.test') + expect(requireByokOrEnv(request, 'openai', ['OPENAI_API_KEY'])).toBeNull() + }) +}) + +describe('byokMissing', () => { + it('returns a typed 401 the client can detect', async () => { + const response = byokMissing('openai') + expect(response.status).toBe(401) + const body: unknown = await response.json() + expect(isByokMissingBody(body)).toBe(true) + if (isByokMissingBody(body)) { + expect(body.error.provider).toBe('openai') + } + }) +}) + +describe('scrub', () => { + it('masks a key down to the last 4 and redacts occurrences', () => { + expect(maskKey('sk-supersecret1234')).toBe('…1234') + expect( + scrubSecrets('failed with sk-supersecret1234!', ['sk-supersecret1234']), + ).toBe('failed with …1234!') + }) +}) + +describe('memory storage', () => { + it('persists nothing', () => { + const store = memoryStorage() + store.save({ openai: 'sk-1' }) + expect(store.load()).toEqual({}) + expect(store.persistent).toBe(false) + }) +}) + +describe('passkey crypto', () => { + const prf = new Uint8Array(32).fill(7) + + it('round-trips a keyring through AES-256-GCM', async () => { + const key = await deriveAesKey(prf) + const { iv, ciphertext } = await encryptKeyring(key, { + openai: 'sk-secret', + }) + expect(await decryptKeyring(key, iv, ciphertext)).toEqual({ + openai: 'sk-secret', + }) + }) + + it('drops unknown providers when decrypting', async () => { + const key = await deriveAesKey(prf) + const iv = crypto.getRandomValues(new Uint8Array(12)) + const plaintext = new TextEncoder().encode( + JSON.stringify({ openai: 'sk-secret', evil: 'bad' }), + ) + const ciphertext = await crypto.subtle.encrypt( + { name: 'AES-GCM', iv }, + key, + plaintext, + ) + expect(await decryptKeyring(key, iv.buffer, ciphertext)).toEqual({ + openai: 'sk-secret', + }) + }) + + it('fails to decrypt with a key derived from a different PRF output', async () => { + const key = await deriveAesKey(prf) + const { iv, ciphertext } = await encryptKeyring(key, { + openai: 'sk-secret', + }) + const wrongKey = await deriveAesKey(new Uint8Array(32).fill(9)) + await expect(decryptKeyring(wrongKey, iv, ciphertext)).rejects.toThrow() + }) +}) + +describe('defaultByokStorage', () => { + it('returns passkey or memory storage', () => { + const store = defaultByokStorage() + expect(['passkey', 'memory']).toContain(store.id) + }) +}) + +describe('OpenRouter PKCE', () => { + const storage = new Map() + + beforeEach(() => { + storage.clear() + vi.stubGlobal('sessionStorage', { + getItem: (key: string) => storage.get(key) ?? null, + setItem: (key: string, value: string) => { + storage.set(key, value) + }, + removeItem: (key: string) => { + storage.delete(key) + }, + }) + }) + + it('generateCodeVerifier produces URL-safe strings', () => { + const verifier = generateCodeVerifier(48) + expect(verifier).toHaveLength(48) + expect(verifier).toMatch(/^[A-Za-z0-9\-._~]+$/) + }) + + it('createS256CodeChallenge is deterministic for a fixed verifier', async () => { + const challenge = await createS256CodeChallenge('test-verifier-fixed') + expect(challenge).toMatch(/^[A-Za-z0-9\-_]+$/) + expect(await createS256CodeChallenge('test-verifier-fixed')).toBe(challenge) + }) + + it('buildOpenRouterAuthUrl encodes callback and S256 challenge', () => { + const url = buildOpenRouterAuthUrl({ + callbackUrl: 'https://app.test/chat', + codeChallenge: 'abc123', + codeChallengeMethod: 'S256', + }) + expect(url).toContain('https://openrouter.ai/auth?') + expect(url).toContain('callback_url=https%3A%2F%2Fapp.test%2Fchat') + expect(url).toContain('code_challenge=abc123') + expect(url).toContain('code_challenge_method=S256') + }) + + it('exchangeOpenRouterCode posts code and verifier', async () => { + let postedBody = '' + const fetchImpl = vi.fn(async (_url: string, init?: RequestInit) => { + postedBody = String(init?.body ?? '') + return Response.json({ key: 'sk-or-pkce-key' }) + }) as typeof fetch + const key = await exchangeOpenRouterCode({ + code: 'auth-code', + codeVerifier: 'verifier-1', + codeChallengeMethod: 'S256', + fetchImpl, + }) + expect(key).toBe('sk-or-pkce-key') + expect(JSON.parse(postedBody)).toEqual({ + code: 'auth-code', + code_verifier: 'verifier-1', + code_challenge_method: 'S256', + }) + }) + + it('completeOpenRouterPkceFromUrl exchanges when code and pending state exist', async () => { + storeOpenRouterPkcePending({ + codeVerifier: 'verifier-xyz', + codeChallengeMethod: 'S256', + callbackUrl: 'https://app.test/', + }) + const fetchImpl = vi.fn(async () => + Response.json({ key: 'sk-or-returned' }), + ) + const key = await completeOpenRouterPkceFromUrl({ + url: 'https://app.test/?code=returned-code', + fetchImpl, + cleanUrl: false, + }) + expect(key).toBe('sk-or-returned') + expect(loadOpenRouterPkcePending()).toBeNull() + clearOpenRouterPkcePending() + }) +}) diff --git a/packages/ai-byok/tests/react.test.tsx b/packages/ai-byok/tests/react.test.tsx new file mode 100644 index 000000000..fbd55830b --- /dev/null +++ b/packages/ai-byok/tests/react.test.tsx @@ -0,0 +1,90 @@ +import { describe, expect, it } from 'vitest' +import { act, renderHook, waitFor } from '@testing-library/react' +import { ByokProvider, useByok } from '../src/react' +import { byokHeaders } from '../src/index' +import type { ReactNode } from 'react' +import type { Keyring } from '../src/index' +import type { KeyringStorage } from '../src/index' + +/** In-memory stand-in for an unlockable (passkey-like) storage, no WebAuthn. */ +function fakeEncryptedStorage(initial: Keyring): KeyringStorage { + let data: Keyring = { ...initial } + return { + id: 'fake', + label: 'Fake encrypted', + persistent: true, + unlockable: true, + peek: () => + Object.fromEntries( + Object.entries(data) + .filter(([, key]) => Boolean(key)) + .map(([provider, key]) => [provider, key.slice(-4)]), + ), + load: () => ({ ...data }), + save: (keys) => { + data = { ...keys } + }, + clear: () => { + data = {} + }, + } +} + +describe('useByok', () => { + it('throws outside a provider', () => { + expect(() => renderHook(() => useByok())).toThrow(/ByokProvider/) + }) + + it('sets, exposes for headers, and clears keys', async () => { + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await act(async () => { + await result.current.setKey('openai', 'sk-live-1234') + }) + expect(byokHeaders(result.current.keys)).toEqual({ + 'x-byok-openai': 'sk-live-1234', + }) + // Only the last 4 are exposed as status. + expect(result.current.status.openai).toEqual({ + state: 'set', + masked: '…1234', + }) + + await act(async () => { + await result.current.clearKey('openai') + }) + expect(result.current.keys.openai).toBeUndefined() + }) + + it('surfaces saved keys as locked after refresh, then unlock reveals them', async () => { + const store = fakeEncryptedStorage({ anthropic: 'sk-ant-9999' }) + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + // Peek (no decryption) surfaces presence + last-4 as a locked status, but + // the key itself isn't available until unlock. + await waitFor(() => + expect(result.current.status.anthropic).toEqual({ + state: 'locked', + masked: '…9999', + }), + ) + expect(result.current.locked).toBe(true) + expect(result.current.keys.anthropic).toBeUndefined() + + await act(async () => { + await result.current.unlock() + }) + expect(result.current.locked).toBe(false) + expect(result.current.keys.anthropic).toBe('sk-ant-9999') + expect(result.current.status.anthropic).toEqual({ + state: 'set', + masked: '…9999', + }) + }) +}) diff --git a/packages/ai-byok/tsconfig.json b/packages/ai-byok/tsconfig.json new file mode 100644 index 000000000..b12da32b4 --- /dev/null +++ b/packages/ai-byok/tsconfig.json @@ -0,0 +1,10 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "jsx": "react-jsx", + "lib": ["ES2022", "DOM", "DOM.Iterable"], + "outDir": "dist" + }, + "include": ["src", "tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-byok/vite.config.ts b/packages/ai-byok/vite.config.ts new file mode 100644 index 000000000..4d794c30e --- /dev/null +++ b/packages/ai-byok/vite.config.ts @@ -0,0 +1,42 @@ +import { defineConfig, mergeConfig } from 'vitest/config' +import { tanstackViteConfig } from '@tanstack/vite-config' +import packageJson from './package.json' + +const config = defineConfig({ + test: { + name: packageJson.name, + dir: './', + watch: false, + globals: true, + environment: 'jsdom', + include: ['tests/**/*.test.ts', 'tests/**/*.test.tsx'], + coverage: { + provider: 'v8', + reporter: ['text', 'json', 'html', 'lcov'], + exclude: [ + 'node_modules/', + 'dist/', + 'tests/', + '**/*.test.ts', + '**/*.config.ts', + '**/types.ts', + ], + include: ['src/**/*.{ts,tsx}'], + }, + }, +}) + +export default mergeConfig( + config, + tanstackViteConfig({ + entry: [ + './src/index.ts', + './src/server.ts', + './src/react.ts', + './src/openrouter.ts', + './src/openrouter-react.ts', + ], + srcDir: './src', + cjs: false, + }), +) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e006a80cf..6f76ae500 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -622,6 +622,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-claude-code': specifier: workspace:* version: link:../../packages/ai-claude-code @@ -1470,6 +1473,27 @@ importers: specifier: ^7.3.3 version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-byok: + devDependencies: + '@testing-library/react': + specifier: ^16.3.0 + version: 16.3.0(@testing-library/dom@10.4.1)(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + '@types/react': + specifier: ^19.2.7 + version: 19.2.7 + '@vitest/coverage-v8': + specifier: 4.0.14 + version: 4.0.14(vitest@4.1.4) + jsdom: + specifier: ^27.2.0 + version: 27.3.0(postcss@8.5.15) + react: + specifier: ^19.2.3 + version: 19.2.3 + vite: + specifier: ^7.3.3 + version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-claude-code: devDependencies: '@tanstack/ai': @@ -2485,6 +2509,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-claude-code': specifier: workspace:* version: link:../../packages/ai-claude-code diff --git a/testing/e2e/README.md b/testing/e2e/README.md index cc06de6ca..0becdcb93 100644 --- a/testing/e2e/README.md +++ b/testing/e2e/README.md @@ -52,14 +52,15 @@ Deterministic scenarios covering tool execution flows: ### Advanced feature tests -| Spec file | What it covers | -| ------------------------------ | --------------------------------------------------------- | -| `tests/abort.spec.ts` | Stop button cancels in-flight generation | -| `tests/lazy-tools.spec.ts` | `__lazy__tool__discovery__` discovers and uses lazy tools | -| `tests/custom-events.spec.ts` | Server tool `emitCustomEvent` received by client | -| `tests/middleware.spec.ts` | `onChunk` transform, `onBeforeToolCall` skip | -| `tests/error-handling.spec.ts` | Server RUN_ERROR, aimock error fixture | -| `tests/tool-error.spec.ts` | Tool throws error, agentic loop continues | +| Spec file | What it covers | +| ------------------------------ | ---------------------------------------------------------------- | +| `tests/abort.spec.ts` | Stop button cancels in-flight generation | +| `tests/lazy-tools.spec.ts` | `__lazy__tool__discovery__` discovers and uses lazy tools | +| `tests/custom-events.spec.ts` | Server tool `emitCustomEvent` received by client | +| `tests/middleware.spec.ts` | `onChunk` transform, `onBeforeToolCall` skip | +| `tests/error-handling.spec.ts` | Server RUN_ERROR, aimock error fixture | +| `tests/tool-error.spec.ts` | Tool throws error, agentic loop continues | +| `tests/byok.spec.ts` | BYOK key rides in header (not body); missing key → `byokMissing` | ## 1. Quick Start diff --git a/testing/e2e/package.json b/testing/e2e/package.json index b761eb58a..e8e877b39 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -20,6 +20,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", + "@tanstack/ai-byok": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", diff --git a/testing/e2e/src/lib/providers.ts b/testing/e2e/src/lib/providers.ts index fa186a19e..6155a14b5 100644 --- a/testing/e2e/src/lib/providers.ts +++ b/testing/e2e/src/lib/providers.ts @@ -45,8 +45,13 @@ export function createTextAdapter( _aimockPort?: number, testId?: string, feature?: Feature, + // Per-request key override — the BYOK route passes the key read from the + // request header here, proving it flows client → header → adapter. aimock + // ignores auth, so the value only needs to be non-empty. + apiKeyOverride?: string, ): { adapter: AnyTextAdapter } { const model = modelOverride ?? defaultModels[provider] + const apiKey = apiKeyOverride ?? DUMMY_KEY // OpenAI, Grok SDKs need /v1 in baseURL. Groq SDK appends /openai/v1/ internally. // Anthropic, Gemini, Ollama SDKs include their path prefixes internally @@ -77,7 +82,7 @@ export function createTextAdapter( const factories: Record { adapter: AnyTextAdapter }> = { openai: () => createChatOptions({ - adapter: createOpenaiChat(model as 'gpt-4o', DUMMY_KEY, { + adapter: createOpenaiChat(model as 'gpt-4o', apiKey, { baseURL: openaiUrl, defaultHeaders: testHeaders, }), diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 14ffbf330..9422e01ec 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -19,6 +19,7 @@ import { Route as DevtoolsRouteARouteImport } from './routes/devtools-route-a' import { Route as DevtoolsGenerationHooksRouteImport } from './routes/devtools-generation-hooks' import { Route as DevtoolsChatRouteImport } from './routes/devtools-chat' import { Route as ChatClientDefaultBridgeRouteImport } from './routes/chat-client-default-bridge' +import { Route as ByokRouteImport } from './routes/byok' import { Route as IndexRouteImport } from './routes/index' import { Route as ProviderIndexRouteImport } from './routes/$provider/index' import { Route as ApiVideoRouteImport } from './routes/api.video' @@ -46,6 +47,7 @@ import { Route as ApiMcpAppsCallRouteImport } from './routes/api.mcp-apps-call' import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiChatRouteImport } from './routes/api.chat' +import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiAudioRouteImport } from './routes/api.audio' import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire' import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage' @@ -108,6 +110,11 @@ const ChatClientDefaultBridgeRoute = ChatClientDefaultBridgeRouteImport.update({ path: '/chat-client-default-bridge', getParentRoute: () => rootRouteImport, } as any) +const ByokRoute = ByokRouteImport.update({ + id: '/byok', + path: '/byok', + getParentRoute: () => rootRouteImport, +} as any) const IndexRoute = IndexRouteImport.update({ id: '/', path: '/', @@ -247,6 +254,11 @@ const ApiChatRoute = ApiChatRouteImport.update({ path: '/api/chat', getParentRoute: () => rootRouteImport, } as any) +const ApiByokChatRoute = ApiByokChatRouteImport.update({ + id: '/api/byok-chat', + path: '/api/byok-chat', + getParentRoute: () => rootRouteImport, +} as any) const ApiAudioRoute = ApiAudioRouteImport.update({ id: '/api/audio', path: '/api/audio', @@ -306,6 +318,7 @@ const ApiAudioStreamRoute = ApiAudioStreamRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -322,6 +335,7 @@ export interface FileRoutesByFullPath { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -356,6 +370,7 @@ export interface FileRoutesByFullPath { } export interface FileRoutesByTo { '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -372,6 +387,7 @@ export interface FileRoutesByTo { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -407,6 +423,7 @@ export interface FileRoutesByTo { export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -423,6 +440,7 @@ export interface FileRoutesById { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -459,6 +477,7 @@ export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -475,6 +494,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -509,6 +529,7 @@ export interface FileRouteTypes { fileRoutesByTo: FileRoutesByTo to: | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -525,6 +546,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -559,6 +581,7 @@ export interface FileRouteTypes { id: | '__root__' | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -575,6 +598,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -610,6 +634,7 @@ export interface FileRouteTypes { } export interface RootRouteChildren { IndexRoute: typeof IndexRoute + ByokRoute: typeof ByokRoute ChatClientDefaultBridgeRoute: typeof ChatClientDefaultBridgeRoute DevtoolsChatRoute: typeof DevtoolsChatRoute DevtoolsGenerationHooksRoute: typeof DevtoolsGenerationHooksRoute @@ -626,6 +651,7 @@ export interface RootRouteChildren { ApiAnthropicStructuredUsageRoute: typeof ApiAnthropicStructuredUsageRoute ApiArktypeToolWireRoute: typeof ApiArktypeToolWireRoute ApiAudioRoute: typeof ApiAudioRouteWithChildren + ApiByokChatRoute: typeof ApiByokChatRoute ApiChatRoute: typeof ApiChatRoute ApiImageRoute: typeof ApiImageRouteWithChildren ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -726,6 +752,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ChatClientDefaultBridgeRouteImport parentRoute: typeof rootRouteImport } + '/byok': { + id: '/byok' + path: '/byok' + fullPath: '/byok' + preLoaderRoute: typeof ByokRouteImport + parentRoute: typeof rootRouteImport + } '/': { id: '/' path: '/' @@ -915,6 +948,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatRouteImport parentRoute: typeof rootRouteImport } + '/api/byok-chat': { + id: '/api/byok-chat' + path: '/api/byok-chat' + fullPath: '/api/byok-chat' + preLoaderRoute: typeof ApiByokChatRouteImport + parentRoute: typeof rootRouteImport + } '/api/audio': { id: '/api/audio' path: '/api/audio' @@ -1055,6 +1095,7 @@ const ApiVideoRouteWithChildren = ApiVideoRoute._addFileChildren( const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, + ByokRoute: ByokRoute, ChatClientDefaultBridgeRoute: ChatClientDefaultBridgeRoute, DevtoolsChatRoute: DevtoolsChatRoute, DevtoolsGenerationHooksRoute: DevtoolsGenerationHooksRoute, @@ -1071,6 +1112,7 @@ const rootRouteChildren: RootRouteChildren = { ApiAnthropicStructuredUsageRoute: ApiAnthropicStructuredUsageRoute, ApiArktypeToolWireRoute: ApiArktypeToolWireRoute, ApiAudioRoute: ApiAudioRouteWithChildren, + ApiByokChatRoute: ApiByokChatRoute, ApiChatRoute: ApiChatRoute, ApiImageRoute: ApiImageRouteWithChildren, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, diff --git a/testing/e2e/src/routes/api.byok-chat.ts b/testing/e2e/src/routes/api.byok-chat.ts new file mode 100644 index 000000000..e726c5083 --- /dev/null +++ b/testing/e2e/src/routes/api.byok-chat.ts @@ -0,0 +1,58 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { byokMissing, getByokKey } from '@tanstack/ai-byok/server' +import { createTextAdapter } from '@/lib/providers' + +/** + * BYOK relay for E2E. Reads the OpenAI key from the per-provider request header + * (never the body), hands it to the adapter for this one call, and streams the + * aimock-backed response. Returns a typed `byokMissing` 401 when the header is + * absent. Never persists or logs the key. + */ +export const Route = createFileRoute('/api/byok-chat')({ + server: { + handlers: { + POST: async ({ request }) => { + await import('@/lib/llmock-server').then((m) => m.ensureLLMock()) + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + const fp = params.forwardedProps as Record + const testId = typeof fp.testId === 'string' ? fp.testId : undefined + + // Header-only read. No key → typed 401 the client renders. + const apiKey = getByokKey(request, 'openai') + if (!apiKey) return byokMissing('openai') + + const adapterOptions = createTextAdapter( + 'openai', + undefined, + undefined, + testId, + 'chat', + apiKey, + ) + + const stream = chat({ + ...adapterOptions, + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + }) + return toServerSentEventsResponse(stream) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/byok.tsx b/testing/e2e/src/routes/byok.tsx new file mode 100644 index 000000000..ba9c5a40a --- /dev/null +++ b/testing/e2e/src/routes/byok.tsx @@ -0,0 +1,95 @@ +import { createFileRoute } from '@tanstack/react-router' +import { useMemo, useRef } from 'react' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { + ByokKeyManager, + ByokProvider, + byokHeaders, + memoryStorage, + useByok, +} from '@tanstack/ai-byok/react' +import { ChatUI } from '@/components/ChatUI' +import type { Keyring, KeyringStorage } from '@tanstack/ai-byok/react' + +interface ByokSearch { + testId?: string + aimockPort?: number + key?: string +} + +export const Route = createFileRoute('/byok')({ + component: ByokRoute, + validateSearch: (search: Record): ByokSearch => ({ + testId: typeof search.testId === 'string' ? search.testId : undefined, + aimockPort: + search.aimockPort != null ? Number(search.aimockPort) : undefined, + key: typeof search.key === 'string' ? search.key : undefined, + }), +}) + +// A storage that hydrates the keyring with a preloaded key on mount — the same +// path `passkeyStorage` uses to restore keys, driven deterministically for the +// test. No key param → session-only memory (nothing stored). +function preloadedStorage(keys: Keyring): KeyringStorage { + return { + id: 'preload', + label: 'Preloaded (test)', + persistent: false, + load: () => keys, + save: () => {}, + clear: () => {}, + } +} + +function ByokRoute() { + const { key } = Route.useSearch() + const storage = useMemo( + () => (key ? preloadedStorage({ openai: key }) : memoryStorage()), + [key], + ) + return ( + + + + ) +} + +function ByokChat() { + const { testId, aimockPort } = Route.useSearch() + const { keys } = useByok() + const keysRef = useRef(keys) + keysRef.current = keys + + const connection = useMemo( + () => + fetchServerSentEvents('/api/byok-chat', () => ({ + headers: byokHeaders(keysRef.current), + })), + [], + ) + + const chat = useChat({ + id: 'byok-chat', + connection, + body: { testId, aimockPort }, + }) + + return ( +
+ + {chat.error ? ( +

+ {chat.error.message} +

+ ) : null} + { + void chat.sendMessage(text) + }} + onStop={chat.stop} + /> +
+ ) +} diff --git a/testing/e2e/tests/byok.spec.ts b/testing/e2e/tests/byok.spec.ts new file mode 100644 index 000000000..9991e97f0 --- /dev/null +++ b/testing/e2e/tests/byok.spec.ts @@ -0,0 +1,61 @@ +import { test, expect } from './fixtures' +import { + getLastAssistantMessage, + sendMessage, + waitForResponse, +} from './helpers' + +const KEY = 'sk-e2e-byok-secret-1234' + +function byokUrl(testId: string, aimockPort: number, key?: string): string { + let url = `/byok?testId=${encodeURIComponent(testId)}&aimockPort=${aimockPort}` + if (key) url += `&key=${encodeURIComponent(key)}` + return url +} + +test.describe('byok', () => { + test('key rides in the header (not the body), streams a response, shows only last-4', async ({ + page, + testId, + aimockPort, + }) => { + // The keyring is hydrated with the key (as passkey storage would restore it). + await page.goto(byokUrl(testId, aimockPort, KEY)) + + // Write-only UI: the manager shows only the last 4, never the full key. + await expect(page.getByTestId('byok-masked-openai')).toHaveText('…1234') + await expect(page.getByTestId('byok-masked-openai')).not.toContainText(KEY) + + // Capture the outgoing relay request to prove the key is in the header and + // absent from the persisted body. + const requestPromise = page.waitForRequest( + (req) => req.url().includes('/api/byok-chat') && req.method() === 'POST', + ) + + await sendMessage(page, '[chat] recommend a guitar') + + const request = await requestPromise + expect(request.headers()['x-byok-openai']).toBe(KEY) + expect(request.postData() ?? '').not.toContain(KEY) + + await waitForResponse(page) + expect(await getLastAssistantMessage(page)).toContain('Fender Stratocaster') + }) + + test('missing key surfaces a byokMissing error and does not answer', async ({ + page, + testId, + aimockPort, + }) => { + await page.goto(byokUrl(testId, aimockPort)) + + // No key entered — the relay returns a typed 401. + await sendMessage(page, '[chat] recommend a guitar') + + await expect(page.getByTestId('byok-error')).toBeVisible() + // The relay refused the call, so no real answer is produced. + expect(await getLastAssistantMessage(page)).not.toContain( + 'Fender Stratocaster', + ) + }) +})