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/api/ai-byok.md b/docs/api/ai-byok.md new file mode 100644 index 000000000..ae210f52c --- /dev/null +++ b/docs/api/ai-byok.md @@ -0,0 +1,337 @@ +--- +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"); + // ... +} +``` + +### `getByokOrEnvKey(request, provider, envVarNames)` + +Header key if present, otherwise the first non-empty named env var. Returns `null` when neither is set — callers should `return byokMissing(provider)` and pass the key to `createOpenaiChat(model, apiKey)` (or the matching `create*` factory). + +### `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)` + +- `maskKey` — display-safe form (`…last4`, or `…` when the key is 4 characters or shorter). +- `lastFour` — raw trailing 4 characters. Not display-safe: a short key is returned in full. Use `maskKey` in UI and logs. + +## `@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, // true only for a decrypted key this session — not for locked keys + storageError, // mount/load failure from the configured storage, if any + } = 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 or persist failed; 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: + +```http +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 82ec31072..eb9bf7721 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 ca6ac5918..f0a00a17b 100644 --- a/docs/config.json +++ b/docs/config.json @@ -780,6 +780,11 @@ "to": "api/ai-angular", "addedAt": "2026-06-15", "updatedAt": "2026-08-18" + }, + { + "label": "@tanstack/ai-byok", + "to": "api/ai-byok", + "addedAt": "2026-08-13" } ] }, diff --git a/docs/getting-started/overview.md b/docs/getting-started/overview.md index 80d37a4ee..f7f71796b 100644 --- a/docs/getting-started/overview.md +++ b/docs/getting-started/overview.md @@ -103,6 +103,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 673d79dc7..b31b0aaee 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -22,6 +22,7 @@ "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", "@tanstack/ai-byteplus": "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/lib/model-selection.ts b/examples/ts-react-chat/src/lib/model-selection.ts index 6c26bf27a..a9e44f675 100644 --- a/examples/ts-react-chat/src/lib/model-selection.ts +++ b/examples/ts-react-chat/src/lib/model-selection.ts @@ -1,22 +1,6 @@ -export type Provider = - | 'openai' - | 'anthropic' - | 'gemini' - | 'gemini-interactions' - | 'ollama' - | 'grok' - | 'groq' - | 'openrouter' - | 'bedrock' - | 'byteplus' +import { z } from 'zod' -export interface ModelOption { - provider: Provider - model: string - label: string -} - -export const MODEL_OPTIONS: Array = [ +export const MODEL_OPTIONS = [ // OpenAI { provider: 'openai', model: 'gpt-5.6', label: 'OpenAI - GPT-5.6' }, { provider: 'openai', model: 'gpt-5.6-sol', label: 'OpenAI - GPT-5.6 Sol' }, @@ -186,8 +170,8 @@ export const MODEL_OPTIONS: Array = [ }, { provider: 'openrouter', - model: 'x-ai/grok-4', - label: 'OpenRouter - xAI Grok 4', + model: 'x-ai/grok-4.3', + label: 'OpenRouter - xAI Grok 4.3', }, { provider: 'openrouter', @@ -308,6 +292,28 @@ export const MODEL_OPTIONS: Array = [ // gpt-oss-120b-250805 is deliberately absent: this route always merges the // server tool set into the request, and that model's tool support is // undeclared in model-meta and unverified against the live API. -] +] as const + +export type ModelOption = (typeof MODEL_OPTIONS)[number] +export type Provider = ModelOption['provider'] + +export type ChatSelection = { + [TProvider in Provider]: { + provider: TProvider + model: Extract['model'] + } +}[Provider] + +export const chatSelectionSchema = z.custom( + (value) => { + if (typeof value !== 'object' || value === null) return false + if (!('provider' in value) || !('model' in value)) return false + return MODEL_OPTIONS.some( + (option) => + option.provider === value.provider && option.model === value.model, + ) + }, + { message: 'Unknown provider/model combination' }, +) export const DEFAULT_MODEL_OPTION = MODEL_OPTIONS[0] diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index 91f64eb35..9d5a70162 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, @@ -9,6 +9,7 @@ import { Github, Image, ImagePlus, + KeyRound, Mic, Music, PauseCircle, @@ -28,7 +29,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 { BoundInterrupts } from '@tanstack/ai-client' import type { UIMessage } from '@tanstack/ai-react' @@ -389,6 +400,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 }> >([]) @@ -434,7 +489,10 @@ function ChatPage() { interrupts, stop, } = useChat({ - connection: fetchServerSentEvents('/api/tanchat'), + connection: fetchServerSentEvents( + '/api/tanchat', + withByok(() => keysRef.current, { onMissingKey: handleMissingKey }), + ), tools, byok, byokProvider: () => toByokProvider(selectedProviderRef.current), @@ -633,6 +691,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} @@ -804,6 +919,15 @@ function ChatPage() { ) } +function ChatRoute() { + const [storage] = useState(() => defaultByokStorage()) + return ( + + + + ) +} + export const Route = createFileRoute('/')({ - component: ChatPage, + component: ChatRoute, }) diff --git a/knip.json b/knip.json index 98ba9561f..a70646653 100644 --- a/knip.json +++ b/knip.json @@ -59,6 +59,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..fce6d3e5a --- /dev/null +++ b/packages/ai-byok/README.md @@ -0,0 +1,190 @@ +# @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 + useChat({ + connection: fetchServerSentEvents( + '/api/chat', + withByok(() => keysRef.current, { + onMissingKey: (provider) => { + if (status[provider]?.state === 'locked') void unlock() + else openKeyDialog(provider) + }, + }), + ), + }) + return null +} +``` + +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) || status[provider].state === 'locked'` (then `unlock()`) +before calling. `hasKey` is only true for a **decrypted** key this session. + +## 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' + +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.5', 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..37c8723c6 --- /dev/null +++ b/packages/ai-byok/src/client/openrouter-pkce.ts @@ -0,0 +1,304 @@ +/** + * 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 = useS256 ? generateCodeVerifier() : '' + const codeChallengeMethod: OpenRouterPkceChallengeMethod = useS256 + ? 'S256' + : 'plain' + + storeOpenRouterPkcePending({ + codeVerifier, + codeChallengeMethod, + callbackUrl: options.callbackUrl, + }) + + const authUrl = buildOpenRouterAuthUrl({ + callbackUrl: options.callbackUrl, + codeChallenge: useS256 + ? await createS256CodeChallenge(codeVerifier) + : undefined, + 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 the URL has no `code`. + * Throws if a `code` is present but PKCE session state is missing (expired) or + * the token exchange fails. Concurrent callers with the same code share one + * exchange — authorization codes are single-use. + */ +const inflightByCode = new Map>() + +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 existing = inflightByCode.get(code) + if (existing) return existing + + const pending = loadOpenRouterPkcePending() + if (!pending) { + throw new Error( + 'OpenRouter authorization code present but PKCE session expired — try signing in again', + ) + } + + const exchange = exchangeOpenRouterCode({ + code, + ...(pending.codeChallengeMethod === 'S256' + ? { + codeVerifier: pending.codeVerifier, + codeChallengeMethod: pending.codeChallengeMethod, + } + : {}), + fetchImpl: options.fetchImpl, + }).then((key) => { + if (options.clearPending !== false) clearOpenRouterPkcePending() + if (options.cleanUrl !== false) stripOpenRouterCodeFromUrl(href) + return key + }) + + inflightByCode.set(code, exchange) + try { + return await exchange + } finally { + inflightByCode.delete(code) + } +} diff --git a/packages/ai-byok/src/client/passkey.ts b/packages/ai-byok/src/client/passkey.ts new file mode 100644 index 000000000..a7a29ade9 --- /dev/null +++ b/packages/ai-byok/src/client/passkey.ts @@ -0,0 +1,381 @@ +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) continue + // Keys of length ≤ 4 would make last-4 the whole secret — store presence only. + preview[provider as ProviderId] = key.length > 4 ? 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..235cf2288 --- /dev/null +++ b/packages/ai-byok/src/react/byok-context.tsx @@ -0,0 +1,379 @@ +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 + /** + * True when this session holds a decrypted key for the provider. False for + * empty and for `locked` keys — call `unlock()` before treating a locked + * provider as usable. + */ + hasKey: (provider: ProviderId) => boolean + /** Mount/load failure from the configured storage, if any. */ + storageError: string | null +} + +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)) + const [storageError, setStorageError] = useState(null) + + // 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 + const statusesRef = useRef(statuses) + statusesRef.current = statuses + const lockedRef = useRef(locked) + lockedRef.current = locked + + // 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 + const fail = (error: unknown) => { + if (!cancelled) { + setStorageError(error instanceof Error ? error.message : String(error)) + } + } + if (storage.unlockable) { + if (!storage.peek) return + void Promise.resolve(storage.peek()) + .then((preview) => { + if (cancelled) return + const present = Object.entries(preview) + setStatuses((current) => { + const next = { ...current } + for (const [provider, last4] of present) { + if (!next[provider as ProviderId]) { + next[provider as ProviderId] = { + state: 'locked', + masked: last4 ? `…${last4}` : '…', + } + } + } + return next + }) + // Nothing stored → nothing to unlock. + if (present.length === 0) setLocked(false) + }) + .catch(fail) + return () => { + cancelled = true + } + } + void Promise.resolve(storage.load()) + .then((loaded) => { + if (!cancelled) applyLoaded(loaded) + }) + .catch(fail) + return () => { + cancelled = true + } + }, [storage, applyLoaded]) + + const unlock = useCallback(async () => { + if (!storage.unlockable) return + applyLoaded(await Promise.resolve(storage.load())) + setLocked(false) + setStorageError(null) + }, [storage, applyLoaded]) + + const persist = useCallback( + (next: Keyring) => Promise.resolve(storage.save(next)), + [storage], + ) + + // Decrypt the stored ring before mutating it. Persist replaces the whole + // ciphertext, so a locked session (in-memory keys are `{}`) must load first + // or Save/Clear/PKCE would wipe every other provider. + const hydrateIfLocked = useCallback(async (): Promise => { + if (!storage.unlockable || !lockedRef.current) return keysRef.current + const loaded = await Promise.resolve(storage.load()) + applyLoaded(loaded) + setLocked(false) + lockedRef.current = false + setStorageError(null) + return { ...loaded, ...keysRef.current } + }, [storage, applyLoaded]) + + const failMutation = ( + provider: ProviderId, + previousKeys: Keyring, + previousStatuses: Partial>, + key: string | undefined, + error: unknown, + ) => { + setKeys(previousKeys) + keysRef.current = previousKeys + const previous = previousStatuses[provider] + const masked = key + ? maskKey(key) + : previous && 'masked' in previous + ? previous.masked + : '…' + setStatuses({ + ...previousStatuses, + [provider]: { + state: 'error', + masked, + message: error instanceof Error ? error.message : String(error), + }, + }) + } + + const setKey = useCallback( + async (provider: ProviderId, key: string) => { + const previousKeys = keysRef.current + const previousStatuses = statusesRef.current + const wasLocked = lockedRef.current + try { + const next = { ...(await hydrateIfLocked()), [provider]: key } + setKeys(next) + keysRef.current = next + setStatuses((prev) => ({ + ...prev, + [provider]: { state: 'set', masked: maskKey(key) }, + })) + await persist(next) + setLocked(false) + lockedRef.current = false + setStorageError(null) + } catch (error) { + failMutation(provider, previousKeys, previousStatuses, key, error) + if (wasLocked) { + setLocked(true) + lockedRef.current = true + } + throw error + } + }, + [hydrateIfLocked, persist], + ) + + const clearKey = useCallback( + async (provider: ProviderId) => { + const previousKeys = keysRef.current + const previousStatuses = statusesRef.current + const wasLocked = lockedRef.current + try { + const next = { ...(await hydrateIfLocked()) } + delete next[provider] + setKeys(next) + keysRef.current = next + setStatuses((prev) => ({ ...prev, [provider]: EMPTY })) + await persist(next) + } catch (error) { + failMutation( + provider, + previousKeys, + previousStatuses, + previousKeys[provider], + error, + ) + if (wasLocked) { + setLocked(true) + lockedRef.current = true + } + throw error + } + }, + [hydrateIfLocked, persist], + ) + + const clearAll = useCallback(async () => { + const previousKeys = keysRef.current + const previousStatuses = statusesRef.current + try { + setKeys({}) + keysRef.current = {} + setStatuses({}) + await Promise.resolve(storage.clear()) + setLocked(false) + lockedRef.current = false + setStorageError(null) + } catch (error) { + setKeys(previousKeys) + keysRef.current = previousKeys + setStatuses(previousStatuses) + setStorageError(error instanceof Error ? error.message : String(error)) + throw error + } + }, [storage]) + + const validateKey = useCallback( + async (provider: ProviderId, key?: string): Promise => { + const target = key ?? keysRef.current[provider] + if (!target) { + const existing = statusesRef.current[provider] + if (existing?.state === 'locked') return existing + 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]), + storageError, + }), + [ + keys, + setKey, + clearKey, + clearAll, + validateKey, + status, + storage, + locked, + unlock, + storageError, + ], + ) + + 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..7f0ebf2b1 --- /dev/null +++ b/packages/ai-byok/src/react/byok-key-manager.tsx @@ -0,0 +1,143 @@ +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, storageError } = useByok() + const [unlocking, setUnlocking] = useState(false) + const [unlockError, setUnlockError] = useState(null) + + return ( +
+ {locked ? ( +
+ Your saved keys are locked ({storage.label}). + +
+ ) : null} + {unlockError ?

{unlockError}

: null} + {storageError ?

{storageError}

: 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..befc030bf --- /dev/null +++ b/packages/ai-byok/src/react/byok-provider-row.tsx @@ -0,0 +1,309 @@ +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 isLocked = status.state === 'locked' + const hasKey = status.state !== 'empty' + const lockedLast4 = + isLocked && '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} + + {status.state === 'error' ? ( +

{status.message}

+ ) : null} + +
{ + event.preventDefault() + if (!draft || isLocked) return + void setKey(provider, draft) + .then(() => setDraft('')) + .catch(() => undefined) + }} + > + 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, + }, + error: { margin: 0, color: '#dc2626', fontSize: 12, lineHeight: 1.4 }, +} 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..91838823c --- /dev/null +++ b/packages/ai-byok/src/server.ts @@ -0,0 +1,22 @@ +/** + * `@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, getByokOrEnvKey } from './server/get-byok-key' +export { byokMissing, isByokMissingBody } from './server/byok-missing' +export type { ByokMissingBody } from './server/byok-missing' +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..248923ed5 --- /dev/null +++ b/packages/ai-byok/src/server/get-byok-key.ts @@ -0,0 +1,39 @@ +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 +} + +/** + * Header key if present, otherwise the first non-empty named env var. + * Returns `null` when neither is set — callers should `return byokMissing(provider)`. + */ +export function getByokOrEnvKey( + request: { headers: Pick }, + provider: ProviderId, + envVarNames: ReadonlyArray, +): string | null { + const header = getByokKey(request, provider) + if (header) return header + for (const name of envVarNames) { + const value = process.env[name] + if (value) return value + } + 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..cb950c946 --- /dev/null +++ b/packages/ai-byok/src/server/scrub.ts @@ -0,0 +1,36 @@ +/** + * 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. + */ + +/** + * Raw trailing 4 characters. Not display-safe: a key of length ≤ 4 is returned + * in full. Use {@link maskKey} in UI and logs. + */ +export function lastFour(key: string): string { + return key.slice(-4) +} + +/** Masks a key to `"…last4"`, or `"…"` when the key is 4 characters or shorter. */ +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..fbf2a4f9a --- /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 models list accepts `x-goog-api-key` (header) or `?key=` (query). + 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..4757963b9 --- /dev/null +++ b/packages/ai-byok/tests/byok.test.ts @@ -0,0 +1,433 @@ +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, + getByokOrEnvKey, + isByokMissingBody, + maskKey, + 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, + startOpenRouterPkceLogin, + 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('withByok({ onMissingKey }) attaches a fetchClient that fires on byokMissing', async () => { + const onMissingKey = vi.fn() + const fetchImpl = vi.fn(async () => byokMissing('anthropic')) + const build = withByok(() => ({ openai: 'sk-a' }), { + onMissingKey, + headers: { Authorization: 'Bearer x' }, + fetchClient: fetchImpl, + }) + const options = build() + expect(options.headers).toEqual({ + Authorization: 'Bearer x', + 'x-byok-openai': 'sk-a', + }) + expect(options.fetchClient).toEqual(expect.any(Function)) + await options.fetchClient?.('https://x.test') + expect(onMissingKey).toHaveBeenCalledWith('anthropic') + }) + + 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('getByokOrEnvKey', () => { + 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('prefers the BYOK header over env', () => { + process.env.OPENAI_API_KEY = 'sk-env' + const request = new Request('https://x.test', { + headers: { [byokHeaderName('openai')]: 'sk-byok' }, + }) + expect(getByokOrEnvKey(request, 'openai', ['OPENAI_API_KEY'])).toBe( + 'sk-byok', + ) + }) + + it('falls back to the first non-empty env var', () => { + process.env.OPENAI_API_KEY = 'sk-env' + const request = new Request('https://x.test') + expect(getByokOrEnvKey(request, 'openai', ['OPENAI_API_KEY'])).toBe( + 'sk-env', + ) + }) + + it('returns null when neither header nor env is present', () => { + delete process.env.OPENAI_API_KEY + const request = new Request('https://x.test') + expect(getByokOrEnvKey(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(maskKey('abcd')).toBe('…') + 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() + }) + + it('useS256:false omits the auth challenge and the exchange verifier', async () => { + await startOpenRouterPkceLogin({ + callbackUrl: 'https://app.test/chat', + useS256: false, + navigate: (url) => { + expect(url).toContain('callback_url=https%3A%2F%2Fapp.test%2Fchat') + expect(url).not.toContain('code_challenge') + }, + }) + expect(loadOpenRouterPkcePending()?.codeChallengeMethod).toBe('plain') + + let postedBody = '' + const fetchImpl = vi.fn(async (_url: string, init?: RequestInit) => { + postedBody = String(init?.body ?? '') + return Response.json({ key: 'sk-or-plain' }) + }) as typeof fetch + const key = await completeOpenRouterPkceFromUrl({ + url: 'https://app.test/chat?code=plain-code', + fetchImpl, + cleanUrl: false, + }) + expect(key).toBe('sk-or-plain') + expect(JSON.parse(postedBody)).toEqual({ code: 'plain-code' }) + }) + + it('completeOpenRouterPkceFromUrl throws when code is present without pending state', async () => { + await expect( + completeOpenRouterPkceFromUrl({ + url: 'https://app.test/?code=orphan', + cleanUrl: false, + }), + ).rejects.toThrow(/PKCE session expired/) + }) + + it('completeOpenRouterPkceFromUrl shares one exchange for the same code', async () => { + storeOpenRouterPkcePending({ + codeVerifier: 'verifier-xyz', + codeChallengeMethod: 'S256', + callbackUrl: 'https://app.test/', + }) + let resolveFetch!: (value: Response) => void + const fetchImpl = vi.fn( + () => + new Promise((resolve) => { + resolveFetch = resolve + }), + ) + const url = 'https://app.test/?code=same-code' + const first = completeOpenRouterPkceFromUrl({ + url, + fetchImpl, + cleanUrl: false, + }) + const second = completeOpenRouterPkceFromUrl({ + url, + fetchImpl, + cleanUrl: false, + }) + resolveFetch(Response.json({ key: 'sk-or-shared' })) + expect(await Promise.all([first, second])).toEqual([ + 'sk-or-shared', + 'sk-or-shared', + ]) + expect(fetchImpl).toHaveBeenCalledTimes(1) + }) +}) diff --git a/packages/ai-byok/tests/react.test.tsx b/packages/ai-byok/tests/react.test.tsx new file mode 100644 index 000000000..8a4046589 --- /dev/null +++ b/packages/ai-byok/tests/react.test.tsx @@ -0,0 +1,191 @@ +import { describe, expect, it, vi } 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', + }) + }) + + it('setKey while locked hydrates before persist so other keys are kept', async () => { + const store = fakeEncryptedStorage({ + openai: 'sk-openai-1111', + anthropic: 'sk-ant-2222', + }) + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await waitFor(() => expect(result.current.locked).toBe(true)) + + await act(async () => { + await result.current.setKey('openrouter', 'sk-or-3333') + }) + + expect(store.load()).toEqual({ + openai: 'sk-openai-1111', + anthropic: 'sk-ant-2222', + openrouter: 'sk-or-3333', + }) + expect(result.current.keys).toEqual({ + openai: 'sk-openai-1111', + anthropic: 'sk-ant-2222', + openrouter: 'sk-or-3333', + }) + expect(result.current.locked).toBe(false) + }) + + it('setKey writes the full keyring to storage.save', async () => { + const save = vi.fn() + const store: KeyringStorage = { + id: 'spy', + label: 'Spy', + persistent: false, + load: () => ({}), + save, + clear: () => {}, + } + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await act(async () => { + await result.current.setKey('openai', 'sk-live-1234') + }) + expect(save).toHaveBeenCalledWith({ openai: 'sk-live-1234' }) + }) + + it('rolls back keys when save throws', async () => { + const store: KeyringStorage = { + id: 'fail', + label: 'Fail', + persistent: false, + load: () => ({}), + save: () => { + throw new Error('PRF unsupported') + }, + clear: () => {}, + } + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await act(async () => { + await expect( + result.current.setKey('openai', 'sk-live-1234'), + ).rejects.toThrow('PRF unsupported') + }) + expect(result.current.keys.openai).toBeUndefined() + expect(result.current.status.openai).toEqual({ + state: 'error', + masked: '…1234', + message: 'PRF unsupported', + }) + }) + + it('validateKey does not clear a locked status', async () => { + const store = fakeEncryptedStorage({ anthropic: 'sk-ant-9999' }) + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await waitFor(() => + expect(result.current.status.anthropic.state).toBe('locked'), + ) + + let status + await act(async () => { + status = await result.current.validateKey('anthropic') + }) + expect(status).toEqual({ + state: 'locked', + masked: '…9999', + }) + expect(result.current.status.anthropic.state).toBe('locked') + }) +}) 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 1e77cd9c6..f9badcdf2 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -682,6 +682,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-byteplus': specifier: workspace:* version: link:../../packages/ai-byteplus @@ -1594,6 +1597,27 @@ importers: specifier: ^8.2.1 version: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + packages/ai-byok: + devDependencies: + '@testing-library/react': + specifier: ^16.3.0 + version: 16.3.2(@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(supports-color@7.2.0)(vitest@4.1.10) + jsdom: + specifier: ^27.2.0 + version: 27.4.0(postcss@8.5.26) + react: + specifier: ^19.2.3 + version: 19.2.3 + vite: + specifier: ^7.3.3 + version: 7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + packages/ai-byteplus: dependencies: '@tanstack/ai-utils': @@ -2787,6 +2811,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-byteplus': specifier: workspace:* version: link:../../packages/ai-byteplus @@ -9847,6 +9874,15 @@ packages: vite: ^5.0.0 || ^6.0.0 || ^7.0.0 || ^8.0.0 vue: ^3.2.25 + '@vitest/coverage-v8@4.0.14': + resolution: {integrity: sha512-EYHLqN/BY6b47qHH7gtMxAg++saoGmsjWmAq9MlXxAz4M0NcHh9iOyKhBZyU4yxZqOd8Xnqp80/5saeitz4Cng==} + peerDependencies: + '@vitest/browser': 4.0.14 + vitest: 4.0.14 + peerDependenciesMeta: + '@vitest/browser': + optional: true + '@vitest/coverage-v8@4.1.10': resolution: {integrity: sha512-IM49HmthevbgAO4anp1hwtoT9wYe59w0LR00gr+eagHE+ZJ5lK4sLPeO0ubgoJcwLk6dehU3R24N+FbEEKDc8g==} peerDependencies: @@ -9870,6 +9906,9 @@ packages: vite: optional: true + '@vitest/pretty-format@4.0.14': + resolution: {integrity: sha512-SOYPgujB6TITcJxgd3wmsLl+wZv+fy3av2PpiPpsWPZ6J1ySUYfScfpIt2Yv56ShJXR2MOA6q2KjKHN4EpdyRQ==} + '@vitest/pretty-format@4.1.10': resolution: {integrity: sha512-W1HsjSH4MXQ9YfmmhLAoIYf1HRfekQCGngeIgcei6MP5QQGWUe0gkopdZQaVCFO+JDJMrAJGwa5pRpNpvy4P8Q==} @@ -9882,6 +9921,9 @@ packages: '@vitest/spy@4.1.10': resolution: {integrity: sha512-PLf/Ugvoq5wO/b4rwYCR1h2PSIdXz7wnkQFMiUpLdtM7l6pqVFcQIBEHyT1+l+cj7mNwAfZHzqXqDyjvOuwbDw==} + '@vitest/utils@4.0.14': + resolution: {integrity: sha512-hLqXZKAWNg8pI+SQXyXxWCTOpA3MvsqcbVeNgSi8x/CSN2wi26dSzn1wrOhmCmFjEvN9p8/kLFRHa6PI8jHazw==} + '@vitest/utils@4.1.10': resolution: {integrity: sha512-fy9am/HWxbaGt/Sawrp90vt6Y6jQwf1RX77cz3uwoJwJVMli/e1IEwRPnMNJ7vKfPTwo0diXifkpPvwH9v7nGA==} @@ -9955,10 +9997,12 @@ packages: '@xmldom/xmldom@0.8.13': resolution: {integrity: sha512-KRYzxepc14G/CEpEGc3Yn+JKaAeT63smlDr+vjB8jRfgTBBI9wRj/nkQEO+ucV8p8I9bfKLWp37uHgFrbntPvw==} engines: {node: '>=10.0.0'} + deprecated: this version has critical issues, please update to the latest version '@xmldom/xmldom@0.9.10': resolution: {integrity: sha512-A9gOqLdi6cV4ibazAjcQufGj0B1y/vDqYrcuP6d/6x8P27gRS8643Dj9o1dEKtB6O7fwxb2FgBmJS2mX7gpvdw==} engines: {node: '>=14.6'} + deprecated: this version has critical issues, please update to the latest version '@yarnpkg/lockfile@1.1.0': resolution: {integrity: sha512-GpSwvyXOcOOlV70vbnzjj4fW5xW/FdUF6nQEt1ENy7m4ZCczi1+/buVUPAqmGfqznsORNFzUMjctTIp8a9tuCQ==} @@ -10164,6 +10208,9 @@ packages: resolution: {integrity: sha512-6t10qk83GOG8p0vKmaCr8eiilZwO171AvbROMtvvNiwrTly62t+7XkA8RdIIVbpMhCASAsxgAzdRSwh6nw/5Dg==} engines: {node: '>=4'} + ast-v8-to-istanbul@0.3.12: + resolution: {integrity: sha512-BRRC8VRZY2R4Z4lFIL35MwNXmwVqBityvOIwETtsCSwvjl0IdgFsy9NhdaA6j74nUdtJJlIypeRhpDam19Wq3g==} + ast-v8-to-istanbul@1.0.5: resolution: {integrity: sha512-UPAgKJFSEGMWSDr3LX4tqnAb4f7KGT8O40Tyx8wbYmmZ/yn58lNCm8h3svs3eXgiGd5AXxz8NDOvXWvicq+rJA==} @@ -12731,6 +12778,10 @@ packages: resolution: {integrity: sha512-GCfE1mtsHGOELCU8e/Z7YWzpmybrx/+dSTfLrvY8qRmaY6zXTKWn6WQIjaAFw069icm6GVMNkgu0NzI4iPZUNw==} engines: {node: '>=10'} + istanbul-lib-source-maps@5.0.6: + resolution: {integrity: sha512-yg2d+Em4KizZC5niWhQaIomgf5WlL4vOOjZ5xGCmF8SnPE/mDWWXgvRExdcpCgh9lLRRa1/fSYp2ymmbJ1pI+A==} + engines: {node: '>=10'} + istanbul-reports@3.2.0: resolution: {integrity: sha512-HGYWWS/ehqTV3xN10i23tkPkpH46MLCIMFNCaaKNavAXTF1RkqxawEPtnjnGZ6XKSInBKkiOA5BKS+aZiY3AvA==} engines: {node: '>=8'} @@ -17531,7 +17582,7 @@ snapshots: '@babel/code-frame': 7.29.7 '@babel/generator': 7.29.7 '@babel/helper-compilation-targets': 7.29.7 - '@babel/helper-module-transforms': 7.29.7(@babel/core@7.29.0) + '@babel/helper-module-transforms': 7.29.7(@babel/core@7.29.0(supports-color@7.2.0)) '@babel/helpers': 7.29.7 '@babel/parser': 7.29.7 '@babel/template': 7.29.7 @@ -17769,7 +17820,7 @@ snapshots: transitivePeerDependencies: - supports-color - '@babel/helper-module-transforms@7.29.7(@babel/core@7.29.0)': + '@babel/helper-module-transforms@7.29.7(@babel/core@7.29.0(supports-color@7.2.0))': dependencies: '@babel/core': 7.29.0(supports-color@7.2.0) '@babel/helper-module-imports': 7.29.7 @@ -24941,6 +24992,23 @@ snapshots: vite: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) vue: 3.5.25(typescript@5.9.3) + '@vitest/coverage-v8@4.0.14(supports-color@7.2.0)(vitest@4.1.10)': + dependencies: + '@bcoe/v8-coverage': 1.0.2 + '@vitest/utils': 4.0.14 + ast-v8-to-istanbul: 0.3.12 + istanbul-lib-coverage: 3.2.2 + istanbul-lib-report: 3.0.1 + istanbul-lib-source-maps: 5.0.6(supports-color@7.2.0) + istanbul-reports: 3.2.0 + magicast: 0.5.2 + obug: 2.1.1 + std-env: 3.10.0 + tinyrainbow: 3.1.0 + vitest: 4.1.10(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.0.14)(happy-dom@20.11.2)(jsdom@27.4.0(postcss@8.5.26))(vite@7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)) + transitivePeerDependencies: + - supports-color + '@vitest/coverage-v8@4.1.10(vitest@4.1.10)': dependencies: '@bcoe/v8-coverage': 1.0.2 @@ -24989,6 +25057,10 @@ snapshots: optionalDependencies: vite: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + '@vitest/pretty-format@4.0.14': + dependencies: + tinyrainbow: 3.1.0 + '@vitest/pretty-format@4.1.10': dependencies: tinyrainbow: 3.1.0 @@ -25007,6 +25079,11 @@ snapshots: '@vitest/spy@4.1.10': {} + '@vitest/utils@4.0.14': + dependencies: + '@vitest/pretty-format': 4.0.14 + tinyrainbow: 3.1.0 + '@vitest/utils@4.1.10': dependencies: '@vitest/pretty-format': 4.1.10 @@ -25352,6 +25429,12 @@ snapshots: dependencies: tslib: 2.8.1 + ast-v8-to-istanbul@0.3.12: + dependencies: + '@jridgewell/trace-mapping': 0.3.31 + estree-walker: 3.0.3 + js-tokens: 10.0.0 + ast-v8-to-istanbul@1.0.5: dependencies: '@jridgewell/trace-mapping': 0.3.31 @@ -28296,6 +28379,14 @@ snapshots: make-dir: 4.0.0 supports-color: 7.2.0 + istanbul-lib-source-maps@5.0.6(supports-color@7.2.0): + dependencies: + '@jridgewell/trace-mapping': 0.3.31 + debug: 4.4.3(supports-color@7.2.0) + istanbul-lib-coverage: 3.2.2 + transitivePeerDependencies: + - supports-color + istanbul-reports@3.2.0: dependencies: html-escaper: 2.0.2 @@ -29718,7 +29809,7 @@ snapshots: less: 4.6.6 ora: 9.4.0 piscina: 5.2.0 - postcss: 8.5.19 + postcss: 8.5.26 rollup-plugin-dts: 6.4.1(rollup@4.60.1)(typescript@5.9.3) rxjs: 7.8.2 sass: 1.101.0 @@ -33409,6 +33500,37 @@ snapshots: optionalDependencies: vite: 8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + vitest@4.1.10(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.0.14)(happy-dom@20.11.2)(jsdom@27.4.0(postcss@8.5.26))(vite@7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)): + dependencies: + '@vitest/expect': 4.1.10 + '@vitest/mocker': 4.1.10(vite@7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)) + '@vitest/pretty-format': 4.1.10 + '@vitest/runner': 4.1.10 + '@vitest/snapshot': 4.1.10 + '@vitest/spy': 4.1.10 + '@vitest/utils': 4.1.10 + es-module-lexer: 2.0.0 + expect-type: 1.3.0 + magic-string: 0.30.21 + obug: 2.1.1 + pathe: 2.0.3 + picomatch: 4.0.5 + std-env: 4.0.0 + tinybench: 2.9.0 + tinyexec: 1.1.1 + tinyglobby: 0.2.17 + tinyrainbow: 3.1.0 + vite: 7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0) + why-is-node-running: 2.3.0 + optionalDependencies: + '@opentelemetry/api': 1.9.1 + '@types/node': 24.10.3 + '@vitest/coverage-v8': 4.0.14(supports-color@7.2.0)(vitest@4.1.10) + happy-dom: 20.11.2 + jsdom: 27.4.0(postcss@8.5.26) + transitivePeerDependencies: + - msw + vitest@4.1.10(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.1.10)(happy-dom@20.11.2)(jsdom@27.4.0(postcss@8.5.19))(vite@7.3.6(@types/node@24.10.3)(jiti@2.7.0)(less@4.6.6)(lightningcss@1.33.0)(sass@1.97.3)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0)): dependencies: '@vitest/expect': 4.1.10 diff --git a/testing/e2e/README.md b/testing/e2e/README.md index 309860a84..f175f61b0 100644 --- a/testing/e2e/README.md +++ b/testing/e2e/README.md @@ -77,14 +77,15 @@ Notes: ### 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` | ### Durable / detachable run tests diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 5f07fce4c..e8136fcc5 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -21,6 +21,7 @@ "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", "@tanstack/ai-byteplus": "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 e617949df..486bba82f 100644 --- a/testing/e2e/src/lib/providers.ts +++ b/testing/e2e/src/lib/providers.ts @@ -54,8 +54,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 @@ -91,7 +96,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/routes/api.byok-chat.ts b/testing/e2e/src/routes/api.byok-chat.ts index 817961cb9..e726c5083 100644 --- a/testing/e2e/src/routes/api.byok-chat.ts +++ b/testing/e2e/src/routes/api.byok-chat.ts @@ -4,22 +4,20 @@ import { chatParamsFromRequestBody, toServerSentEventsResponse, } from '@tanstack/ai' -import { byokMissing, getByokKey } from '@tanstack/ai/byok' +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()) - if (request.signal.aborted) { - return new Response(null, { status: 499 }) - } - - const key = getByokKey(request, 'openai') - if (!key) return byokMissing('openai') - - const abortController = new AbortController() let params try { @@ -31,42 +29,29 @@ export const Route = createFileRoute('/api/byok-chat')({ ) } - const fp = params.forwardedProps + const fp = params.forwardedProps as Record const testId = typeof fp.testId === 'string' ? fp.testId : undefined - const aimockPort = - fp.aimockPort != null ? Number(fp.aimockPort) : undefined - try { - const adapterOptions = createTextAdapter( - 'openai', - undefined, - aimockPort, - testId, - 'chat', - ) - const stream = chat({ - ...adapterOptions, - systemPrompts: ['You are a helpful assistant for a guitar store.'], - messages: params.messages, - threadId: params.threadId, - runId: params.runId, - abortController, - }) - return toServerSentEventsResponse(stream, { abortController }) - } catch (error) { - if ( - (error instanceof Error && error.name === 'AbortError') || - abortController.signal.aborted - ) { - return new Response(null, { status: 499 }) - } - const message = - error instanceof Error ? error.message : 'An error occurred' - return new Response(JSON.stringify({ error: message }), { - status: 500, - headers: { 'Content-Type': 'application/json' }, - }) - } + // 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 index 573916fbd..ba9c5a40a 100644 --- a/testing/e2e/src/routes/byok.tsx +++ b/testing/e2e/src/routes/byok.tsx @@ -1,93 +1,94 @@ -import { useEffect, useState, type FormEvent } from 'react' import { createFileRoute } from '@tanstack/react-router' -import { defineByok, memoryStorage } from '@tanstack/ai-client/byok' -import { fetchServerSentEvents, useByok, useChat } from '@tanstack/ai-react' +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: ByokPage, - validateSearch: (search: Record) => { - const port = - typeof search.aimockPort === 'number' - ? search.aimockPort - : typeof search.aimockPort === 'string' - ? parseInt(search.aimockPort, 10) - : undefined - return { - testId: typeof search.testId === 'string' ? search.testId : undefined, - aimockPort: port != null && !Number.isNaN(port) ? port : undefined, - } - }, + 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, + }), }) -function ByokPage() { +// 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 [byok] = useState(() => defineByok({ storage: memoryStorage() })) - const snapshot = useByok(byok) - const [hydrated, setHydrated] = useState(false) - const openaiStatus = snapshot.status.openai - const last4 = openaiStatus.state === 'empty' ? '' : openaiStatus.masked + const { keys } = useByok() + const keysRef = useRef(keys) + keysRef.current = keys - useEffect(() => { - setHydrated(true) - }, []) + const connection = useMemo( + () => + fetchServerSentEvents('/api/byok-chat', () => ({ + headers: byokHeaders(keysRef.current), + })), + [], + ) - const { messages, sendMessage, isLoading, stop } = useChat({ - connection: fetchServerSentEvents('/api/byok-chat'), - byok, - forwardedProps: { - provider: 'openai', - model: 'gpt-5.5', - testId, - aimockPort, - }, + const chat = useChat({ + id: 'byok-chat', + connection, + body: { testId, aimockPort }, }) - const handleSave = (event: FormEvent) => { - event.preventDefault() - const form = event.currentTarget - const data = new FormData(form) - const raw = data.get('key') - const next = typeof raw === 'string' ? raw.trim() : '' - if (!next) return - void byok.update('openai', next).then(() => { - form.reset() - }) - } - return ( -
-
-
{snapshot.prompt?.provider ?? ''}
-
{last4}
-
- - -
-
+
+ + {chat.error ? ( +

+ {chat.error.message} +

+ ) : null} { - void sendMessage(text).catch(() => { - // ByokBlockedError / ByokMissingError set snapshot.prompt - }) + void chat.sendMessage(text) }} + onStop={chat.stop} />
) diff --git a/testing/e2e/tests/byok.spec.ts b/testing/e2e/tests/byok.spec.ts index 107faacf7..9991e97f0 100644 --- a/testing/e2e/tests/byok.spec.ts +++ b/testing/e2e/tests/byok.spec.ts @@ -1,44 +1,61 @@ import { test, expect } from './fixtures' -import { sendMessage } from './helpers' - -const RAW_KEY = 'sk-e2e-byok-test-1234' - -function byokUrl(testId: string, aimockPort: number): string { - return `/byok?testId=${encodeURIComponent(testId)}&aimockPort=${aimockPort}` +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('saves a key and sends it in the x-byok-openai header, not the body', async ({ + test('key rides in the header (not the body), streams a response, shows only last-4', async ({ page, testId, aimockPort, }) => { - await page.goto(byokUrl(testId, aimockPort)) + // The keyring is hydrated with the key (as passkey storage would restore it). + await page.goto(byokUrl(testId, aimockPort, KEY)) - const keyInput = page.getByTestId('byok-key-input') - await keyInput.fill(RAW_KEY) - await expect(keyInput).toHaveValue(RAW_KEY) - await page.getByTestId('byok-save-button').click() - await expect(page.getByTestId('byok-last4')).toHaveText('1234') + // 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) - expect(request.headers()['x-byok-openai']).toBe(RAW_KEY) - expect(request.postData() ?? '').not.toContain(RAW_KEY) + await waitForResponse(page) + expect(await getLastAssistantMessage(page)).toContain('Fender Stratocaster') }) - test('prompts for openai when sending without a key', async ({ + test('missing key surfaces a byokMissing error and does not answer', async ({ page, testId, aimockPort, }) => { await page.goto(byokUrl(testId, aimockPort)) - await sendMessage(page, 'hello') - await expect(page.getByTestId('byok-prompt')).toHaveText('openai') + // 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', + ) }) })