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/examples/ts-react-chat/package.json b/examples/ts-react-chat/package.json index 67b58ad43..7faef3832 100644 --- a/examples/ts-react-chat/package.json +++ b/examples/ts-react-chat/package.json @@ -18,6 +18,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", + "@tanstack/ai-byok": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", "@tanstack/ai-code-mode": "workspace:*", diff --git a/examples/ts-react-chat/src/components/ByokKeyDialog.tsx b/examples/ts-react-chat/src/components/ByokKeyDialog.tsx new file mode 100644 index 000000000..6f9f20d1a --- /dev/null +++ b/examples/ts-react-chat/src/components/ByokKeyDialog.tsx @@ -0,0 +1,201 @@ +import { useState } from 'react' +import { KeyRound, Lock, X } from 'lucide-react' +import { BYOK_PROVIDERS, useByok } from '@tanstack/ai-byok/react' +import type { KeyStatus, ProviderId } from '@tanstack/ai-byok/react' + +// BYOK-supported providers in this example (see byok-config.ts). +const PROVIDERS: Array = [ + 'openai', + 'anthropic', + 'gemini', + 'grok', + 'groq', + 'openrouter', +] + +interface ByokKeyDialogProps { + /** Which providers have a key in the server env (never the key itself). */ + envStatus: Partial> + /** The provider for the currently selected model, flagged on the icon when keyless. */ + activeProvider: ProviderId | null + /** Controlled open state (the modal can be opened reactively, not just by the icon). */ + open: boolean + onOpenChange: (open: boolean) => void + /** Provider row to highlight when opened reactively (e.g. a missing key). */ + highlightProvider?: ProviderId | null +} + +/** + * Key-icon button + modal for entering per-provider BYOK keys. Keys stay in the + * browser (passkey-encrypted where supported, else session-only) and are + * attached per-request by the caller. Only the last 4 characters are ever shown + * back; saved-but-locked keys show a lock until unlocked. + */ +export function ByokKeyDialog({ + envStatus, + activeProvider, + open, + onOpenChange, + highlightProvider, +}: ByokKeyDialogProps) { + const { keys, status, locked, unlock } = useByok() + + // Attention needed when the selected model's provider isn't usable right now + // (no server key and no decrypted key — whether missing or saved-but-locked). + const activeNeedsKey = + activeProvider != null && + !envStatus[activeProvider] && + !keys[activeProvider] + + return ( + <> + + + {open ? ( +
onOpenChange(false)} + > +
event.stopPropagation()} + > +
+

API keys

+ +
+

+ 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. +

+ + {locked ? ( + + ) : null} + +
+ {PROVIDERS.map((provider) => ( + + ))} +
+
+
+ ) : null} + + ) +} + +function ProviderKeyRow({ + provider, + hasEnvKey, + status, + highlight, +}: { + provider: ProviderId + hasEnvKey: boolean + status: KeyStatus + highlight?: boolean +}) { + const { keys, setKey, clearKey } = useByok() + const [draft, setDraft] = useState('') + const yourKey = keys[provider] + // Saved-but-not-decrypted this session (persistent storage after a refresh). + const lockedLast4 = status.state === 'locked' ? status.masked.slice(-4) : null + + return ( +
+
+ + {BYOK_PROVIDERS[provider].label} + + {yourKey ? ( + + Your key ··{yourKey.slice(-4)} + + ) : lockedLast4 ? ( + + + ··{lockedLast4} + + ) : hasEnvKey ? ( + Server key + ) : ( + No key + )} +
+
{ + event.preventDefault() + if (!draft) return + void setKey(provider, draft) + setDraft('') + }} + > + setDraft(event.target.value)} + className="min-w-0 flex-1 rounded-md border border-gray-600 bg-gray-900 px-2 py-1 text-sm text-white focus:outline-none focus:ring-2 focus:ring-orange-500/50" + /> + + {yourKey ? ( + + ) : null} +
+
+ ) +} diff --git a/examples/ts-react-chat/src/lib/byok-config.ts b/examples/ts-react-chat/src/lib/byok-config.ts new file mode 100644 index 000000000..53c1974af --- /dev/null +++ b/examples/ts-react-chat/src/lib/byok-config.ts @@ -0,0 +1,41 @@ +import { createServerFn } from '@tanstack/react-start' +import type { ProviderId } from '@tanstack/ai-byok' +import type { Provider } from '@/lib/model-selection' + +/** + * Maps the example's providers to a BYOK provider id and the env var(s) the + * adapter reads server-side. Providers not listed (ollama = local, bedrock = + * AWS credentials) don't use a single bring-your-own key. + */ +export const BYOK_PROVIDER_MAP: Partial< + Record }> +> = { + openai: { byokId: 'openai', envVars: ['OPENAI_API_KEY'] }, + anthropic: { byokId: 'anthropic', envVars: ['ANTHROPIC_API_KEY'] }, + gemini: { byokId: 'gemini', envVars: ['GEMINI_API_KEY', 'GOOGLE_API_KEY'] }, + 'gemini-interactions': { + byokId: 'gemini', + envVars: ['GEMINI_API_KEY', 'GOOGLE_API_KEY'], + }, + grok: { byokId: 'grok', envVars: ['XAI_API_KEY'] }, + groq: { byokId: 'groq', envVars: ['GROQ_API_KEY'] }, + openrouter: { byokId: 'openrouter', envVars: ['OPENROUTER_API_KEY'] }, +} + +/** The BYOK provider id backing an example provider, if any. */ +export function byokIdForProvider(provider: Provider): ProviderId | null { + return BYOK_PROVIDER_MAP[provider]?.byokId ?? null +} + +/** + * Reports which BYOK-supported providers have their key set in the server env, + * so the client can warn before a user picks a model it can't run. Returns only + * booleans — never the key values. + */ +export const getEnvKeyStatus = createServerFn({ method: 'GET' }).handler(() => { + const status: Partial> = {} + for (const { byokId, envVars } of Object.values(BYOK_PROVIDER_MAP)) { + status[byokId] = envVars.some((name) => Boolean(process.env[name])) + } + return status +}) diff --git a/examples/ts-react-chat/src/routes/api.tanchat.ts b/examples/ts-react-chat/src/routes/api.tanchat.ts index a815d1e13..56adc435f 100644 --- a/examples/ts-react-chat/src/routes/api.tanchat.ts +++ b/examples/ts-react-chat/src/routes/api.tanchat.ts @@ -7,15 +7,21 @@ import { mergeAgentTools, toServerSentEventsResponse, } from '@tanstack/ai' -import { openaiText } from '@tanstack/ai-openai' +import { createOpenaiChat, openaiText } from '@tanstack/ai-openai' import { ollamaText } from '@tanstack/ai-ollama' -import { anthropicText } from '@tanstack/ai-anthropic' -import { geminiText } from '@tanstack/ai-gemini' -import { geminiTextInteractions } from '@tanstack/ai-gemini/experimental' -import { openRouterText } from '@tanstack/ai-openrouter' -import { grokText } from '@tanstack/ai-grok' -import { groqText } from '@tanstack/ai-groq' +import { anthropicText, createAnthropicChat } from '@tanstack/ai-anthropic' +import { createGeminiChat, geminiText } from '@tanstack/ai-gemini' +import { + createGeminiTextInteractions, + geminiTextInteractions, +} from '@tanstack/ai-gemini/experimental' +import { createOpenRouterText, openRouterText } from '@tanstack/ai-openrouter' +import { createGrokText, grokText } from '@tanstack/ai-grok' +import { createGroqText, groqText } from '@tanstack/ai-groq' import { bedrockText } from '@tanstack/ai-bedrock' +import { byokMissing, getByokKey } from '@tanstack/ai-byok/server' +import { BYOK_PROVIDER_MAP, byokIdForProvider } from '@/lib/byok-config' +import type { ProviderId } from '@tanstack/ai-byok/server' import type { AnyTextAdapter, ChatMiddleware } from '@tanstack/ai' import { addToCartToolDef, @@ -248,6 +254,19 @@ export const Route = createFileRoute('/api/tanchat')({ ? params.forwardedProps.previousInteractionId : undefined + // BYOK: prefer a per-request key from the request header, falling back + // to the env-configured adapter. The key is read from the header only, + // never persisted or logged. + function byokAdapter( + provider: ProviderId, + m: M, + byok: (model: M, apiKey: string) => AnyTextAdapter, + env: (model: M) => AnyTextAdapter, + ): AnyTextAdapter { + const key = getByokKey(request, provider) + return key ? byok(m, key) : env(m) + } + // Pre-define typed adapter configurations with full type inference // Model is passed to the adapter factory function for type-safe autocomplete const adapterConfig: Record< @@ -256,14 +275,20 @@ export const Route = createFileRoute('/api/tanchat')({ > = { anthropic: () => createChatOptions({ - adapter: anthropicText( + adapter: byokAdapter( + 'anthropic', (model || 'claude-sonnet-4-6') as 'claude-sonnet-4-6', + createAnthropicChat, + anthropicText, ), }), openrouter: () => createChatOptions({ - adapter: openRouterText( + adapter: byokAdapter( + 'openrouter', (model || 'openai/gpt-5.1') as 'openai/gpt-5.1', + createOpenRouterText, + openRouterText, ), modelOptions: { reasoning: { @@ -273,8 +298,11 @@ export const Route = createFileRoute('/api/tanchat')({ }), gemini: () => createChatOptions({ - adapter: geminiText( + adapter: byokAdapter( + 'gemini', (model || 'gemini-3.1-pro-preview') as 'gemini-3.1-pro-preview', + createGeminiChat, + geminiText, ), modelOptions: { thinkingConfig: { @@ -285,8 +313,11 @@ export const Route = createFileRoute('/api/tanchat')({ }), 'gemini-interactions': () => createChatOptions({ - adapter: geminiTextInteractions( + adapter: byokAdapter( + 'gemini', (model || 'gemini-3.1-pro-preview') as 'gemini-3.1-pro-preview', + createGeminiTextInteractions, + geminiTextInteractions, ), modelOptions: { previous_interaction_id: previousInteractionId, @@ -295,15 +326,21 @@ export const Route = createFileRoute('/api/tanchat')({ }), grok: () => createChatOptions({ - adapter: grokText( + adapter: byokAdapter( + 'grok', (model || 'grok-build-0.1') as 'grok-build-0.1', + createGrokText, + grokText, ), modelOptions: {}, }), groq: () => createChatOptions({ - adapter: groqText( + adapter: byokAdapter( + 'groq', (model || 'openai/gpt-oss-120b') as 'openai/gpt-oss-120b', + createGroqText, + groqText, ), }), bedrock: () => @@ -330,7 +367,12 @@ export const Route = createFileRoute('/api/tanchat')({ }), openai: () => createChatOptions({ - adapter: openaiText((model || 'gpt-5.2') as 'gpt-5.2'), + adapter: byokAdapter( + 'openai', + (model || 'gpt-5.2') as 'gpt-5.2', + createOpenaiChat, + openaiText, + ), modelOptions: { prompt_cache_key: 'user-session-12345', prompt_cache_retention: '24h', @@ -344,6 +386,19 @@ export const Route = createFileRoute('/api/tanchat')({ requestedProvider in adapterConfig ? (requestedProvider as Provider) : 'openai' + + // If this provider has no server env key and the request carried no + // BYOK key, tell the client which key to add — a typed 401 the client + // detects (via withByok) instead of a generic 500 from the SDK. + const byokId = byokIdForProvider(provider) + if (byokId) { + const hasByokKey = Boolean(getByokKey(request, byokId)) + const hasEnvKey = BYOK_PROVIDER_MAP[provider]?.envVars.some( + (name) => Boolean(process.env[name]), + ) + if (!hasByokKey && !hasEnvKey) return byokMissing(byokId) + } + // Get typed adapter options using createChatOptions pattern const options = adapterConfig[provider]() diff --git a/examples/ts-react-chat/src/routes/index.tsx b/examples/ts-react-chat/src/routes/index.tsx index fb1ce78f0..8e8c1eb6d 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, @@ -25,7 +25,18 @@ import { useChat, useTranscription, } from '@tanstack/ai-react' +import { + ByokProvider, + isPasskeyStorageSupported, + memoryStorage, + passkeyStorage, + useByok, + withByok, +} from '@tanstack/ai-byok/react' import { clientTools } from '@tanstack/ai-client' +import type { ProviderId } from '@tanstack/ai-byok/react' +import { ByokKeyDialog } from '@/components/ByokKeyDialog' +import { byokIdForProvider, getEnvKeyStatus } from '@/lib/byok-config' import { ThinkingPart } from '@tanstack/ai-react-ui' import type { UIMessage } from '@tanstack/ai-react' import type { ContentPart, TranscriptionResult } from '@tanstack/ai' @@ -371,6 +382,49 @@ 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 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 }> >([]) @@ -414,7 +468,10 @@ function ChatPage() { addToolApprovalResponse, stop, } = useChat({ - connection: fetchServerSentEvents('/api/tanchat'), + connection: fetchServerSentEvents( + '/api/tanchat', + withByok(() => keysRef.current, { onMissingKey: handleMissingKey }), + ), tools, body, onCustomEvent: (eventType, data, context) => { @@ -615,6 +672,13 @@ function ChatPage() { ))} + setKeyDialog((s) => ({ ...s, open }))} + highlightProvider={keyDialog.provider} + /> + {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} | null>(null) + useEffect(() => { + let active = true + async function pick() { + const platformAuth = + isPasskeyStorageSupported() && + (await globalThis.PublicKeyCredential?.isUserVerifyingPlatformAuthenticatorAvailable().catch( + () => false, + )) + if (active) setStorage(platformAuth ? passkeyStorage() : memoryStorage()) + } + void pick() + return () => { + active = false + } + }, []) + + if (!storage) return null + return ( + + + + ) +} + export const Route = createFileRoute('/')({ - component: ChatPage, + component: ChatRoute, }) diff --git a/knip.json b/knip.json index c7995c1aa..19c41a25c 100644 --- a/knip.json +++ b/knip.json @@ -48,6 +48,9 @@ "packages/ai-react": { "ignoreDependencies": ["@mcp-ui/client"] }, + "packages/ai-byok": { + "ignoreDependencies": ["react", "@types/react"] + }, "packages/ai-preact": { "ignoreDependencies": ["@mcp-ui/client"] }, diff --git a/packages/ai-byok/README.md b/packages/ai-byok/README.md new file mode 100644 index 000000000..6effc2a15 --- /dev/null +++ b/packages/ai-byok/README.md @@ -0,0 +1,187 @@ +# @tanstack/ai-byok + +Bring-your-own-key toolkit for [TanStack AI](https://tanstack.com/ai). Users +supply their own provider API keys; the library collects them **client-side**, +attaches them **per-request in a header**, and uses them **server-side without +ever persisting or logging them**. + +## First principle: never be a custodian + +- **Keys live client-side.** The browser is the system of record. +- **The server piece is a stateless pass-through.** It reads the key off the + incoming request header, hands it to the TanStack AI adapter for one call, and + never writes it to a DB, cache, log, or observability stream. +- **No central endpoint.** The server helper is trivially self-hostable; there + is no hardcoded relay URL. + +Keys travel in the `x-byok-` header — never the request body +or message history — so they stay out of persisted conversations and the +event/observability stream. + +## Client (React) + +```tsx +import { useRef } from 'react' +import { ByokProvider, ByokKeyManager, useByok } from '@tanstack/ai-byok/react' +import { + withByok, + memoryStorage, + passkeyStorage, + isPasskeyStorageSupported, +} from '@tanstack/ai-byok/react' +import { fetchServerSentEvents } from '@tanstack/ai-client' +import { useChat } from '@tanstack/ai-react' + +// Wrap your app. `storage` is chosen once. Defaults to the safest option +// (session-only, nothing persisted). Opt into encrypted persistence with +// passkeyStorage() where supported, falling back to memory otherwise. +function App({ children }) { + const storage = isPasskeyStorageSupported() + ? passkeyStorage() + : memoryStorage() + return {children} +} + +// Drop-in settings UI — shows the last 4 chars of a saved key, never the whole key. +function Settings() { + return +} + +// Attach the keys to the connection. `withByok` adds the per-request BYOK +// headers and, via `onMissingKey`, detects the relay's `byokMissing` 401 so you +// can prompt for (or unlock) the key the server was missing. +function Chat() { + const { keys, status, unlock } = useByok() + const keysRef = useRef(keys) + keysRef.current = keys + return useChat({ + connection: fetchServerSentEvents( + '/api/chat', + withByok(() => keysRef.current, { + onMissingKey: (provider) => { + if (status[provider]?.state === 'locked') void unlock() + else openKeyDialog(provider) + }, + }), + ), + }) +} +``` + +For a lower-level path, `byokHeaders(keys)` returns the raw header map and +`byokFetch(onMissingKey)` wraps a `fetch` to detect the `byokMissing` 401. + +### Fetcher transport (`useChat`/`useGeneration` with a `fetcher`) + +`withByok` targets the `connection` transport. For the `fetcher` transport — +used by `useGeneration` (image/audio/video/speech/transcribe) and by +`useChat({ fetcher })` — use `byokFetcher`. It hands your fetcher body BYOK +`headers` and a missing-key-aware `fetch`, read fresh on every call, and works +for both fetcher styles: + +```ts +// A) fetch-based fetcher — spread headers, use the wrapped fetch for onMissingKey +useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers, fetch, signal }) => + fetch('/api/generate/audio', { + method: 'POST', + headers: { 'content-type': 'application/json', ...headers }, + body: JSON.stringify(input), + signal, + }), + { onMissingKey: (provider) => openKeyDialog(provider) }, + ), +}) + +// B) TanStack Start server function — forward headers at the call site +useGenerateAudio({ + fetcher: byokFetcher( + () => keysRef.current, + (input, { headers }) => generateAudioFn({ data: input, headers }), + ), +}) +``` + +The key still travels in the `x-byok-` header, never the `data` +body. On the server, read it off the request with the same `getByokKey`: + +```ts +import { getByokKey, byokMissing } from '@tanstack/ai-byok/server' +import { getRequest } from '@tanstack/react-start/server' + +export const generateAudioFn = createServerFn({ method: 'POST' }) + .inputValidator((data: { prompt: string; provider: ProviderId }) => data) + .handler(({ data }) => { + const apiKey = getByokKey(getRequest(), data.provider) + if (!apiKey) return byokMissing(data.provider) + // ...hand apiKey to the adapter for one call + }) +``` + +The `onMissingKey` callback fires only on the fetch-based path (style A), where +`byokFetcher` owns the `fetch` and can see the `byokMissing` 401. A server +function (style B) surfaces the missing key as a thrown error instead — catch +it and inspect the message, or pre-check with `hasKey(provider)` before calling. + +## Server (stateless — no persist, no log) + +```ts +import { getByokKey, byokMissing } from '@tanstack/ai-byok/server' +import { chat, toServerSentEventsResponse } from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai/adapters' + +export async function POST(request: Request) { + const { messages } = await request.json() + + const apiKey = getByokKey(request, 'openai') + if (!apiKey) return byokMissing('openai') + + const stream = chat({ + adapter: createOpenaiChat('gpt-5.2', apiKey), + messages, + }) + return toServerSentEventsResponse(stream) +} +``` + +## Storage + +Pass one `storage` to ``. Two are built in — **there is no +plaintext persistence**: + +- **`memoryStorage()` (default)** — session / in-memory. Keys vanish on refresh. + Never persisted. Zero at-rest liability. +- **`passkeyStorage()` (opt-in)** — 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..1f235449b --- /dev/null +++ b/packages/ai-byok/package.json @@ -0,0 +1,80 @@ +{ + "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" + } + }, + "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..cdc45076c --- /dev/null +++ b/packages/ai-byok/src/client/keyring.ts @@ -0,0 +1,31 @@ +import { PROVIDER_IDS, byokHeaderName } 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> + +/** + * 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/passkey.ts b/packages/ai-byok/src/client/passkey.ts new file mode 100644 index 000000000..b286f03fc --- /dev/null +++ b/packages/ai-byok/src/client/passkey.ts @@ -0,0 +1,362 @@ +import type { Keyring } from './keyring' +import type { KeyPreview, KeyringStorage } from './storage' +import type { ProviderId } from '../shared/providers' + +/** + * Passkey-encrypted keyring storage (WebAuthn PRF → HKDF → AES-256-GCM). + * + * The keyring is encrypted at rest in IndexedDB with an AES-256-GCM key derived + * from a passkey's PRF output, unwrapped on demand with a biometric/PIN tap. + * Decryption happens entirely client-side with the user present — no server, + * no custodian. + * + * **Honest scope:** this protects against at-rest theft (stolen device, + * storage-dumping extension, backups). It does NOT defeat live in-page XSS — an + * attacker running JS in the origin after the user unlocks can read the + * decrypted keys from memory. + */ + +const STORE_NAME = 'keyring' +const RECORD_ID = 'default' +const HKDF_INFO = 'byok:keyring:v1' +const DEFAULT_DB = 'byok' + +interface StoredRecord { + id: string + /** The passkey's raw credential id, replayed in the unlock ceremony. */ + credentialId: ArrayBuffer + /** Fixed per-install PRF evaluation input (not secret). */ + salt: ArrayBuffer + /** AES-GCM initialization vector for this ciphertext. */ + iv: ArrayBuffer + /** Encrypted keyring JSON. */ + ciphertext: ArrayBuffer + /** + * Unencrypted presence metadata (`provider → last 4`). Non-sensitive, so it + * can be read via {@link KeyringStorage.peek} without an unlock ceremony to + * show saved keys as "locked" after a refresh. + */ + preview: KeyPreview +} + +/** Build the non-sensitive `provider → last 4` preview from a keyring. */ +function previewOf(keys: Keyring): KeyPreview { + const preview: KeyPreview = {} + for (const [provider, key] of Object.entries(keys)) { + if (key) preview[provider as ProviderId] = key.slice(-4) + } + return preview +} + +/** + * Whether the current environment exposes WebAuthn. Note that actual PRF + * support can only be confirmed during registration; `passkeyStorage` throws a + * descriptive error if the chosen authenticator does not support PRF, which the + * caller should catch and fall back to {@link memoryStorage}. + */ +export function isPasskeyStorageSupported(): boolean { + return ( + typeof globalThis !== 'undefined' && + typeof globalThis.PublicKeyCredential !== 'undefined' && + typeof globalThis.navigator !== 'undefined' && + typeof globalThis.navigator.credentials.create === 'function' + ) +} + +// --------------------------------------------------------------------------- +// Crypto (exported for testing; the WebAuthn ceremony below feeds `deriveAesKey`) +// --------------------------------------------------------------------------- + +/** Derive a non-extractable AES-256-GCM key from a 32-byte PRF output. */ +export async function deriveAesKey( + prfOutput: BufferSource, +): Promise { + const base = await crypto.subtle.importKey('raw', prfOutput, 'HKDF', false, [ + 'deriveKey', + ]) + return crypto.subtle.deriveKey( + { + name: 'HKDF', + hash: 'SHA-256', + salt: new Uint8Array(0), + info: new TextEncoder().encode(HKDF_INFO), + }, + base, + { name: 'AES-GCM', length: 256 }, + false, + ['encrypt', 'decrypt'], + ) +} + +export async function encryptKeyring( + key: CryptoKey, + keys: Keyring, +): Promise<{ iv: ArrayBuffer; ciphertext: ArrayBuffer }> { + const iv = crypto.getRandomValues(new Uint8Array(12)) + const plaintext = new TextEncoder().encode(JSON.stringify(keys)) + const ciphertext = await crypto.subtle.encrypt( + { name: 'AES-GCM', iv }, + key, + plaintext, + ) + return { iv: iv.buffer, ciphertext } +} + +export async function decryptKeyring( + key: CryptoKey, + iv: BufferSource, + ciphertext: BufferSource, +): Promise { + const plaintext = await crypto.subtle.decrypt( + { name: 'AES-GCM', iv }, + key, + ciphertext, + ) + const parsed: unknown = JSON.parse(new TextDecoder().decode(plaintext)) + if (typeof parsed !== 'object' || parsed === null) return {} + return 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 { 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) + }, + } +} 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..f949fc202 --- /dev/null +++ b/packages/ai-byok/src/client/with-byok.ts @@ -0,0 +1,136 @@ +import { isByokMissingBody } from '../server/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 +} + +/** + * 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 fetchClient = options.onMissingKey + ? byokFetch(options.onMissingKey, options.fetchClient) + : options.fetchClient + return { + headers: { ...options.headers, ...byokHeaders(getKeys()) }, + ...(fetchClient ? { fetchClient } : {}), + } + } +} + +/** 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 +} + +/** + * 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, { + headers: { ...options.headers, ...byokHeaders(getKeys()) }, + fetch: options.onMissingKey + ? byokFetch(options.onMissingKey, options.fetchClient) + : (options.fetchClient ?? fetch), + signal: transport?.signal, + }) +} diff --git a/packages/ai-byok/src/index.ts b/packages/ai-byok/src/index.ts new file mode 100644 index 000000000..cbae5ee5f --- /dev/null +++ b/packages/ai-byok/src/index.ts @@ -0,0 +1,46 @@ +/** + * `@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 } from './client/keyring' +export type { Keyring } from './client/keyring' + +// Connection helper: attach BYOK headers + detect a missing-key response +export { withByok, byokFetch, byokFetcher } from './client/with-byok' +export type { + WithByokOptions, + ByokConnectionOptions, + ByokFetcherContext, +} from './client/with-byok' +export { isByokMissingBody } from './server/byok-missing' +export type { ByokMissingBody } from './server/byok-missing' + +// Storage +export { memoryStorage } from './client/storage' +export type { KeyringStorage, KeyPreview } from './client/storage' +export { passkeyStorage, isPasskeyStorageSupported } 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/react.ts b/packages/ai-byok/src/react.ts new file mode 100644 index 000000000..281bb940b --- /dev/null +++ b/packages/ai-byok/src/react.ts @@ -0,0 +1,39 @@ +/** + * `@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' + +// Re-export the client toolkit so React consumers have one import path. +export { + byokHeaders, + withByok, + byokFetch, + byokFetcher, + isByokMissingBody, + memoryStorage, + 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..11cfd792b --- /dev/null +++ b/packages/ai-byok/src/react/byok-context.tsx @@ -0,0 +1,272 @@ +import { + createContext, + useCallback, + useEffect, + useMemo, + useRef, + useState, +} from 'react' +import { validateKey as pingProvider } from '../client/validate' +import { memoryStorage } from '../client/storage' +import { maskKey } from '../server/scrub' +import { PROVIDER_IDS } from '../shared/providers' +import type { ReactNode } from 'react' +import type { Keyring } from '../client/keyring' +import type { KeyringStorage } from '../client/storage' +import type { ProviderId } from '../shared/providers' +import type { ValidationStatus } from '../client/validate' + +/** + * Per-provider status surfaced to the UI. `masked` never contains more than the + * last 4 characters — the full key is never rendered back. + */ +export type KeyStatus = + | { state: 'empty' } + | { state: 'set'; masked: string } + // Stored but not yet decrypted this session (unlockable storage after a + // refresh). The key isn't usable until `unlock()`; only last-4 is known. + | { state: 'locked'; masked: string } + | { state: 'validating'; masked: string } + | { state: ValidationStatus; masked: string } + | { state: 'error'; masked: string; message: string } + +const EMPTY: KeyStatus = { state: 'empty' } + +export interface ByokContextValue { + /** + * The live keyring. Pass to `byokHeaders(keys)` when building the connection. + * The UI never renders these; treat them as write-only from the UI's view. + */ + keys: Keyring + /** Set (or overwrite) a provider's key and persist it to the configured storage. */ + setKey: (provider: ProviderId, key: string) => Promise + /** Remove a single provider's key. */ + clearKey: (provider: ProviderId) => Promise + /** Remove every key. */ + clearAll: () => Promise + /** + * Validate a key against the provider. Validates the given key, or the stored + * key when omitted. Records the outcome in {@link status} and returns it; + * never throws — a network/CORS failure is reported as an `error` status. + */ + validateKey: (provider: ProviderId, key?: string) => Promise + /** Per-provider status map. Providers with no key report `{ state: 'empty' }`. */ + status: Record + /** The configured persistence storage. */ + storage: KeyringStorage + /** + * `true` when an `unlockable` storage (e.g. passkey) may hold encrypted keys + * that haven't been decrypted this session. Always `false` for storage that + * needs no unlock ceremony (e.g. memory). + */ + locked: boolean + /** + * Decrypt and load keys from an `unlockable` storage, triggering its unlock + * ceremony (e.g. a biometric tap). No-op — and does not prompt — for storage + * that isn't unlockable or when nothing is stored yet. Rejects if the + * ceremony fails or is cancelled. + */ + unlock: () => Promise + hasKey: (provider: ProviderId) => boolean +} + +export const ByokContext = createContext(null) + +export interface ByokProviderProps { + children: ReactNode + /** + * Where keys are persisted. Defaults to the safest option + * ({@link memoryStorage}) — keys vanish on refresh and nothing is persisted. + * Fixed for the life of the provider. + */ + storage?: KeyringStorage +} + +export function ByokProvider({ + children, + storage: initialStorage, +}: ByokProviderProps) { + // Storage is chosen once and fixed for the life of the provider. + const [storage] = useState( + () => initialStorage ?? memoryStorage(), + ) + const [keys, setKeys] = useState({}) + const [statuses, setStatuses] = useState< + Partial> + >({}) + // Unlockable storage (passkey) starts locked; the user must unlock to decrypt. + const [locked, setLocked] = useState(() => Boolean(storage.unlockable)) + + // Keep a ref to the current keys so the persisting callbacks read the latest + // without re-creating on every keystroke. + const keysRef = useRef(keys) + keysRef.current = keys + + // Apply decrypted/loaded keys. Merge keys UNDER any edits the user made during + // the async load so an early setKey is never clobbered; promote a provider's + // status to `set` when it had no status or was a `locked` placeholder, while + // preserving a fresh user edit. + const applyLoaded = useCallback((loaded: Keyring) => { + setKeys((current) => ({ ...loaded, ...current })) + setStatuses((current) => { + const next = { ...current } + for (const [provider, key] of Object.entries(loaded)) { + if (!key) continue + const existing = current[provider as ProviderId] + if (!existing || existing.state === 'locked') { + next[provider as ProviderId] = { state: 'set', masked: maskKey(key) } + } + } + return next + }) + }, []) + + // On mount: unlockable storage peeks (no ceremony) to surface saved keys as + // `locked`; other storage auto-hydrates its keys. + useEffect(() => { + let cancelled = false + if (storage.unlockable) { + if (!storage.peek) return + void Promise.resolve(storage.peek()).then((preview) => { + if (cancelled) return + const present = Object.entries(preview).filter(([, last4]) => + Boolean(last4), + ) + setStatuses((current) => { + const next = { ...current } + for (const [provider, last4] of present) { + if (!next[provider as ProviderId]) { + next[provider as ProviderId] = { + state: 'locked', + masked: `…${last4}`, + } + } + } + return next + }) + // Nothing stored → nothing to unlock. + if (present.length === 0) setLocked(false) + }) + return () => { + cancelled = true + } + } + void Promise.resolve(storage.load()).then((loaded) => { + if (!cancelled) applyLoaded(loaded) + }) + return () => { + cancelled = true + } + }, [storage, applyLoaded]) + + const unlock = useCallback(async () => { + if (!storage.unlockable) return + applyLoaded(await Promise.resolve(storage.load())) + setLocked(false) + }, [storage, applyLoaded]) + + const persist = useCallback( + (next: Keyring) => Promise.resolve(storage.save(next)), + [storage], + ) + + const setKey = useCallback( + async (provider: ProviderId, key: string) => { + const next = { ...keysRef.current, [provider]: key } + setKeys(next) + setStatuses((prev) => ({ + ...prev, + [provider]: { state: 'set', masked: maskKey(key) }, + })) + await persist(next) + // A successful save ran any unlock/registration ceremony, so the session + // is now unlocked. + setLocked(false) + }, + [persist], + ) + + const clearKey = useCallback( + async (provider: ProviderId) => { + const next = { ...keysRef.current } + delete next[provider] + setKeys(next) + setStatuses((prev) => ({ ...prev, [provider]: EMPTY })) + await persist(next) + }, + [persist], + ) + + const clearAll = useCallback(async () => { + setKeys({}) + setStatuses({}) + await Promise.resolve(storage.clear()) + setLocked(false) + }, [storage]) + + const validateKey = useCallback( + async (provider: ProviderId, key?: string): Promise => { + const target = key ?? keysRef.current[provider] + if (!target) { + setStatuses((prev) => ({ ...prev, [provider]: EMPTY })) + return EMPTY + } + const masked = maskKey(target) + setStatuses((prev) => ({ + ...prev, + [provider]: { state: 'validating', masked }, + })) + + let result: KeyStatus + try { + const status = await pingProvider(provider, target) + result = { state: status, masked } + } catch (error) { + result = { + state: 'error', + masked, + message: error instanceof Error ? error.message : String(error), + } + } + setStatuses((prev) => ({ ...prev, [provider]: result })) + return result + }, + [], + ) + + const status = useMemo(() => { + const full = {} as Record + for (const provider of PROVIDER_IDS) { + full[provider] = statuses[provider] ?? EMPTY + } + return full + }, [statuses]) + + const value = useMemo( + () => ({ + keys, + setKey, + clearKey, + clearAll, + validateKey, + status, + storage, + locked, + unlock, + hasKey: (provider) => Boolean(keys[provider]), + }), + [ + keys, + setKey, + clearKey, + clearAll, + validateKey, + status, + storage, + locked, + unlock, + ], + ) + + return {children} +} diff --git a/packages/ai-byok/src/react/byok-key-manager.tsx b/packages/ai-byok/src/react/byok-key-manager.tsx new file mode 100644 index 000000000..65e32f6f6 --- /dev/null +++ b/packages/ai-byok/src/react/byok-key-manager.tsx @@ -0,0 +1,227 @@ +import { useState } from 'react' +import { BYOK_PROVIDERS, PROVIDER_IDS } 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 ByokKeyManagerProps { + /** Providers to show. Defaults to every registered provider. */ + providers?: Array + 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, + className, + style, +}: ByokKeyManagerProps) { + const { status, storage, locked, unlock } = useByok() + const [unlocking, setUnlocking] = useState(false) + const [unlockError, setUnlockError] = useState(null) + + return ( +
+ {locked ? ( +
+ Your saved keys are locked ({storage.label}). + +
+ ) : null} + {unlockError ?

{unlockError}

: null} + + {storage.warning ?

{storage.warning}

: null} + + {providers.map((provider) => ( + + ))} +
+ ) +} + +function ProviderRow({ + provider, + status, +}: { + provider: ProviderId + status: KeyStatus +}) { + const { setKey, clearKey, validateKey } = useByok() + const [draft, setDraft] = useState('') + const hasKey = status.state !== 'empty' + + return ( +
+
+ {BYOK_PROVIDERS[provider].label} + +
+ + {hasKey && 'masked' in status ? ( +
+ + {status.masked} + +
+ + +
+
+ ) : null} + +
{ + event.preventDefault() + if (!draft) return + void setKey(provider, draft) + setDraft('') + }} + > + setDraft(event.target.value)} + style={styles.input} + /> + +
+
+ ) +} + +function StatusBadge({ status }: { status: KeyStatus }) { + 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} + + ) +} + +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 }, + unlockBanner: { + display: 'flex', + alignItems: 'center', + justifyContent: 'space-between', + gap: 8, + padding: 12, + borderRadius: 8, + background: '#f3f4f6', + border: '1px solid #e5e7eb', + }, + row: { + display: 'flex', + flexDirection: 'column', + gap: 8, + padding: 12, + border: '1px solid #e5e7eb', + borderRadius: 8, + }, + 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', + }, +} satisfies Record 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/server.ts b/packages/ai-byok/src/server.ts new file mode 100644 index 000000000..fc14070cf --- /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 } 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..6fd23091e --- /dev/null +++ b/packages/ai-byok/src/server/byok-missing.ts @@ -0,0 +1,44 @@ +import type { ProviderId } from '../shared/providers' + +/** Discriminated body returned by {@link byokMissing}. */ +export interface ByokMissingBody { + error: { + type: 'byok_missing' + provider: ProviderId + message: string + } +} + +/** + * 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, + }, + }) +} + +/** 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 + return (error as { type?: unknown }).type === 'byok_missing' +} diff --git a/packages/ai-byok/src/server/get-byok-key.ts b/packages/ai-byok/src/server/get-byok-key.ts new file mode 100644 index 000000000..6057b5080 --- /dev/null +++ b/packages/ai-byok/src/server/get-byok-key.ts @@ -0,0 +1,21 @@ +import { byokHeaderName } from '../shared/providers' +import type { ProviderId } from '../shared/providers' + +/** + * Reads a provider's BYOK key off the incoming request header. + * + * Returns the key or `null` if absent. The key is read from the header only — + * never the body — so it stays out of any persisted `messages` array and out + * of the event/observability stream. This function does not log the value and + * must not be wrapped in anything that attaches it to a logger context. + * + * Accepts either a `Request` or any object exposing a `Headers`-like `get`, so + * it works across Fetch-API runtimes (Workers, Deno, Bun, Node/undici). + */ +export function getByokKey( + request: { headers: Pick }, + provider: ProviderId, +): string | null { + const value = request.headers.get(byokHeaderName(provider)) + return value && value.length > 0 ? value : null +} diff --git a/packages/ai-byok/src/server/scrub.ts b/packages/ai-byok/src/server/scrub.ts new file mode 100644 index 000000000..3443f71f1 --- /dev/null +++ b/packages/ai-byok/src/server/scrub.ts @@ -0,0 +1,33 @@ +/** + * Helpers for keeping key material out of logs and error responses. + * + * The server piece is a stateless pass-through: it must never write a key to a + * DB, cache, log, or observability stream, and must never echo a full key back + * to a client. These utilities make the last-4 the only representation that + * ever leaves the process. + */ + +/** Returns the last 4 characters of a key for display, e.g. `"...a1b2"`. */ +export function lastFour(key: string): string { + return key.slice(-4) +} + +/** Masks a key to `"…last4"`, hiding everything but the trailing 4 chars. */ +export function maskKey(key: string): string { + if (key.length <= 4) return '…' + return `…${lastFour(key)}` +} + +/** + * Replaces every occurrence of each secret in `input` with its masked form. + * Use before logging or returning any string that may have interpolated a key + * (URLs, provider SDK error messages, stack traces). + */ +export function scrubSecrets(input: string, secrets: Array): string { + let output = input + for (const secret of secrets) { + if (!secret) continue + output = output.split(secret).join(maskKey(secret)) + } + return output +} diff --git a/packages/ai-byok/src/shared/providers.ts b/packages/ai-byok/src/shared/providers.ts new file mode 100644 index 000000000..47ffcbf29 --- /dev/null +++ b/packages/ai-byok/src/shared/providers.ts @@ -0,0 +1,140 @@ +/** + * Shared provider registry and header-naming convention. + * + * This module is fully isomorphic (no browser or Node APIs) and is the single + * source of truth imported by both the client keyring and the server helpers. + */ + +/** + * Describes how to validate a key against a provider's cheapest authenticated + * endpoint (typically a models list). `headers` builds the request headers + * from the raw key. Only providers that expose a browser-reachable endpoint + * carry a `validate` entry; the rest report `'unsupported'` from + * {@link validateKey}. + */ +export interface ProviderValidateConfig { + /** Endpoint hit with a GET request to confirm the key works. */ + url: string + /** Builds the auth headers for the validation request from the raw key. */ + headers: (key: string) => Record +} + +/** Static metadata for a supported provider. */ +export interface ProviderConfig { + /** Stable id used in the keyring map and the per-provider header name. */ + id: string + /** Human-readable label for UI. */ + label: string + /** Optional validation endpoint metadata. */ + validate?: ProviderValidateConfig +} + +const ANTHROPIC_HEADERS = (key: string): Record => ({ + 'x-api-key': key, + 'anthropic-version': '2023-06-01', + // Required for Anthropic to serve the request from a browser origin. + 'anthropic-dangerous-direct-browser-access': 'true', +}) + +const bearer = (key: string): Record => ({ + Authorization: `Bearer ${key}`, +}) + +/** + * Registry of providers the BYOK toolkit understands. `ProviderId` is derived + * from these keys, so adding a provider here extends the typed union + * everywhere. + */ +export const BYOK_PROVIDERS = { + openai: { + id: 'openai', + label: 'OpenAI', + validate: { url: 'https://api.openai.com/v1/models', headers: bearer }, + }, + anthropic: { + id: 'anthropic', + label: 'Anthropic', + validate: { + url: 'https://api.anthropic.com/v1/models', + headers: ANTHROPIC_HEADERS, + }, + }, + gemini: { + id: 'gemini', + label: 'Google Gemini', + validate: { + // Gemini authenticates the models list via a query param, not a header. + url: 'https://generativelanguage.googleapis.com/v1beta/models', + headers: (key) => ({ 'x-goog-api-key': key }), + }, + }, + openrouter: { + id: 'openrouter', + label: 'OpenRouter', + validate: { url: 'https://openrouter.ai/api/v1/key', headers: bearer }, + }, + groq: { + id: 'groq', + label: 'Groq', + validate: { + url: 'https://api.groq.com/openai/v1/models', + headers: bearer, + }, + }, + grok: { + id: 'grok', + label: 'xAI Grok', + validate: { url: 'https://api.x.ai/v1/models', headers: bearer }, + }, + mistral: { + id: 'mistral', + label: 'Mistral', + validate: { url: 'https://api.mistral.ai/v1/models', headers: bearer }, + }, + elevenlabs: { + id: 'elevenlabs', + label: 'ElevenLabs', + validate: { + url: 'https://api.elevenlabs.io/v1/user', + headers: (key) => ({ 'xi-api-key': key }), + }, + }, + // No browser-reachable validation endpoint (key-scheme auth, no models list). + fal: { id: 'fal', label: 'fal.ai' }, + // Local runtime, no API key involved. + ollama: { id: 'ollama', label: 'Ollama' }, +} as const satisfies Record + +/** Union of every supported provider id. */ +export type ProviderId = keyof typeof BYOK_PROVIDERS + +/** All provider ids as a runtime array. */ +export const PROVIDER_IDS = Object.keys(BYOK_PROVIDERS) as Array + +/** Type guard narrowing an arbitrary string to a known {@link ProviderId}. */ +export function isProviderId(value: string): value is ProviderId { + return value in BYOK_PROVIDERS +} + +/** + * The validation config for a provider, or `undefined` when the provider has no + * browser-reachable validation endpoint. + */ +export function providerValidateConfig( + provider: ProviderId, +): ProviderValidateConfig | undefined { + const config = BYOK_PROVIDERS[provider] + return 'validate' in config ? config.validate : undefined +} + +/** Prefix for every per-provider BYOK header. */ +export const BYOK_HEADER_PREFIX = 'x-byok-' + +/** + * The HTTP header name that carries the key for a given provider. Keys always + * travel in this header, never in the request body or message history, so they + * stay out of persisted conversations and the event/observability stream. + */ +export function byokHeaderName(provider: ProviderId): string { + return `${BYOK_HEADER_PREFIX}${provider}` +} diff --git a/packages/ai-byok/tests/byok.test.ts b/packages/ai-byok/tests/byok.test.ts new file mode 100644 index 000000000..c844a315a --- /dev/null +++ b/packages/ai-byok/tests/byok.test.ts @@ -0,0 +1,166 @@ +import { describe, expect, it, vi } from 'vitest' +import { + byokFetch, + byokFetcher, + byokHeaders, + byokHeaderName, + withByok, +} from '../src/index' +import type { Keyring } from '../src/index' +import { + byokMissing, + getByokKey, + isByokMissingBody, + maskKey, + scrubSecrets, +} from '../src/server' +import { memoryStorage } from '../src/client/storage' +import { + decryptKeyring, + deriveAesKey, + encryptKeyring, +} from '../src/client/passkey' + +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('withByok / byokFetch', () => { + 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('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() + }) +}) + +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('byokMissing', () => { + it('returns a typed 401 the client can detect', async () => { + const response = byokMissing('openai') + expect(response.status).toBe(401) + const body: unknown = await response.json() + expect(isByokMissingBody(body)).toBe(true) + if (isByokMissingBody(body)) { + expect(body.error.provider).toBe('openai') + } + }) +}) + +describe('scrub', () => { + it('masks a key down to the last 4 and redacts occurrences', () => { + expect(maskKey('sk-supersecret1234')).toBe('…1234') + expect( + scrubSecrets('failed with sk-supersecret1234!', ['sk-supersecret1234']), + ).toBe('failed with …1234!') + }) +}) + +describe('memory storage', () => { + it('persists nothing', () => { + const store = memoryStorage() + store.save({ openai: 'sk-1' }) + expect(store.load()).toEqual({}) + expect(store.persistent).toBe(false) + }) +}) + +describe('passkey crypto', () => { + // The WebAuthn ceremony can't run in jsdom, but the PRF → AES-GCM path can: + // derive a key from a fixed 32-byte "PRF output" and round-trip a keyring. + 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('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() + }) +}) diff --git a/packages/ai-byok/tests/react.test.tsx b/packages/ai-byok/tests/react.test.tsx new file mode 100644 index 000000000..fbd55830b --- /dev/null +++ b/packages/ai-byok/tests/react.test.tsx @@ -0,0 +1,90 @@ +import { describe, expect, it } from 'vitest' +import { act, renderHook, waitFor } from '@testing-library/react' +import { ByokProvider, useByok } from '../src/react' +import { byokHeaders } from '../src/index' +import type { ReactNode } from 'react' +import type { Keyring } from '../src/index' +import type { KeyringStorage } from '../src/index' + +/** In-memory stand-in for an unlockable (passkey-like) storage, no WebAuthn. */ +function fakeEncryptedStorage(initial: Keyring): KeyringStorage { + let data: Keyring = { ...initial } + return { + id: 'fake', + label: 'Fake encrypted', + persistent: true, + unlockable: true, + peek: () => + Object.fromEntries( + Object.entries(data) + .filter(([, key]) => Boolean(key)) + .map(([provider, key]) => [provider, key.slice(-4)]), + ), + load: () => ({ ...data }), + save: (keys) => { + data = { ...keys } + }, + clear: () => { + data = {} + }, + } +} + +describe('useByok', () => { + it('throws outside a provider', () => { + expect(() => renderHook(() => useByok())).toThrow(/ByokProvider/) + }) + + it('sets, exposes for headers, and clears keys', async () => { + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + await act(async () => { + await result.current.setKey('openai', 'sk-live-1234') + }) + expect(byokHeaders(result.current.keys)).toEqual({ + 'x-byok-openai': 'sk-live-1234', + }) + // Only the last 4 are exposed as status. + expect(result.current.status.openai).toEqual({ + state: 'set', + masked: '…1234', + }) + + await act(async () => { + await result.current.clearKey('openai') + }) + expect(result.current.keys.openai).toBeUndefined() + }) + + it('surfaces saved keys as locked after refresh, then unlock reveals them', async () => { + const store = fakeEncryptedStorage({ anthropic: 'sk-ant-9999' }) + const wrapper = ({ children }: { children: ReactNode }) => ( + {children} + ) + const { result } = renderHook(() => useByok(), { wrapper }) + + // Peek (no decryption) surfaces presence + last-4 as a locked status, but + // the key itself isn't available until unlock. + await waitFor(() => + expect(result.current.status.anthropic).toEqual({ + state: 'locked', + masked: '…9999', + }), + ) + expect(result.current.locked).toBe(true) + expect(result.current.keys.anthropic).toBeUndefined() + + await act(async () => { + await result.current.unlock() + }) + expect(result.current.locked).toBe(false) + expect(result.current.keys.anthropic).toBe('sk-ant-9999') + expect(result.current.status.anthropic).toEqual({ + state: 'set', + masked: '…9999', + }) + }) +}) diff --git a/packages/ai-byok/tsconfig.json b/packages/ai-byok/tsconfig.json new file mode 100644 index 000000000..b12da32b4 --- /dev/null +++ b/packages/ai-byok/tsconfig.json @@ -0,0 +1,10 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "jsx": "react-jsx", + "lib": ["ES2022", "DOM", "DOM.Iterable"], + "outDir": "dist" + }, + "include": ["src", "tests"], + "exclude": ["node_modules", "dist"] +} diff --git a/packages/ai-byok/vite.config.ts b/packages/ai-byok/vite.config.ts new file mode 100644 index 000000000..f81606c79 --- /dev/null +++ b/packages/ai-byok/vite.config.ts @@ -0,0 +1,36 @@ +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'], + srcDir: './src', + cjs: false, + }), +) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5eab53f64..13719bc61 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -622,6 +622,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-claude-code': specifier: workspace:* version: link:../../packages/ai-claude-code @@ -1470,6 +1473,27 @@ importers: specifier: ^7.3.3 version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-byok: + devDependencies: + '@testing-library/react': + specifier: ^16.3.0 + version: 16.3.0(@testing-library/dom@10.4.1)(@types/react-dom@19.2.3(@types/react@19.2.7))(@types/react@19.2.7)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + '@types/react': + specifier: ^19.2.7 + version: 19.2.7 + '@vitest/coverage-v8': + specifier: 4.0.14 + version: 4.0.14(vitest@4.1.4) + jsdom: + specifier: ^27.2.0 + version: 27.3.0(postcss@8.5.15) + react: + specifier: ^19.2.3 + version: 19.2.3 + vite: + specifier: ^7.3.3 + version: 7.3.3(@types/node@24.10.3)(jiti@2.6.1)(less@4.6.6)(lightningcss@1.30.2)(sass@1.101.0)(terser@5.44.1)(tsx@4.21.0)(yaml@2.8.2) + packages/ai-claude-code: devDependencies: '@tanstack/ai': @@ -2485,6 +2509,9 @@ importers: '@tanstack/ai-bedrock': specifier: workspace:* version: link:../../packages/ai-bedrock + '@tanstack/ai-byok': + specifier: workspace:* + version: link:../../packages/ai-byok '@tanstack/ai-claude-code': specifier: workspace:* version: link:../../packages/ai-claude-code @@ -5511,12 +5538,6 @@ packages: '@emnapi/core': ^1.7.1 '@emnapi/runtime': ^1.7.1 - '@napi-rs/wasm-runtime@1.1.6': - resolution: {integrity: sha512-ZLv/JdUfkvOy9eCnnBaGfiO+XimbjebAeO+MRQqD/B+FR1tnRN0tpKSJHRbE8sFfS6aqsXZ67TQjfwfsxULVbg==} - peerDependencies: - '@emnapi/core': ^1.7.1 - '@emnapi/runtime': ^1.7.1 - '@ngrok/ngrok-android-arm64@1.7.0': resolution: {integrity: sha512-8tco3ID6noSaNy+CMS7ewqPoIkIM6XO5COCzsUp3Wv3XEbMSyn65RN6cflX2JdqLfUCHcMyD0ahr9IEiHwqmbQ==} engines: {node: '>= 10'} @@ -8802,9 +8823,6 @@ packages: '@tybys/wasm-util@0.10.1': resolution: {integrity: sha512-9tTaPJLSiejZKx+Bmog4uSubteqTvFrVrURwkmHixBo0G4seD0zUxp98E1DzUBJxLQ3NPwXrGKDiVjwx/DpPsg==} - '@tybys/wasm-util@0.10.3': - resolution: {integrity: sha512-F3fo1MYrRJYL3zER0OUOmkutjr1Vp23m7OsSgp7nq4SP6OqX6C/56XFIPAl5bt3zaBRjmW7SGz3u/6LwFpYcOg==} - '@tybys/wasm-util@0.9.0': resolution: {integrity: sha512-6+7nlbMVX/PVDCwaIQ8nTOPveOcFLSt8GcXdx8hD0bt39uWxYT88uXzqTd4fTvqta7oeUJqudepapKNt2DYJFw==} @@ -19290,13 +19308,6 @@ snapshots: '@tybys/wasm-util': 0.10.1 optional: true - '@napi-rs/wasm-runtime@1.1.6(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)': - dependencies: - '@emnapi/core': 1.11.1 - '@emnapi/runtime': 1.11.1 - '@tybys/wasm-util': 0.10.3 - optional: true - '@ngrok/ngrok-android-arm64@1.7.0': optional: true @@ -19827,7 +19838,7 @@ snapshots: '@oxc-resolver/binding-wasm32-wasi@11.15.0(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1)': dependencies: - '@napi-rs/wasm-runtime': 1.1.6(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) + '@napi-rs/wasm-runtime': 1.1.4(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) transitivePeerDependencies: - '@emnapi/core' - '@emnapi/runtime' @@ -20997,7 +21008,6 @@ snapshots: dependencies: '@emnapi/core': 1.11.1 '@emnapi/runtime': 1.11.1 - '@napi-rs/wasm-runtime': 1.1.6(@emnapi/core@1.11.1)(@emnapi/runtime@1.11.1) optional: true '@rolldown/binding-win32-arm64-msvc@1.0.0-beta.53': @@ -23064,11 +23074,6 @@ snapshots: tslib: 2.8.1 optional: true - '@tybys/wasm-util@0.10.3': - dependencies: - tslib: 2.8.1 - optional: true - '@tybys/wasm-util@0.9.0': dependencies: tslib: 2.8.1 diff --git a/testing/e2e/README.md b/testing/e2e/README.md index cc06de6ca..0becdcb93 100644 --- a/testing/e2e/README.md +++ b/testing/e2e/README.md @@ -52,14 +52,15 @@ Deterministic scenarios covering tool execution flows: ### Advanced feature tests -| Spec file | What it covers | -| ------------------------------ | --------------------------------------------------------- | -| `tests/abort.spec.ts` | Stop button cancels in-flight generation | -| `tests/lazy-tools.spec.ts` | `__lazy__tool__discovery__` discovers and uses lazy tools | -| `tests/custom-events.spec.ts` | Server tool `emitCustomEvent` received by client | -| `tests/middleware.spec.ts` | `onChunk` transform, `onBeforeToolCall` skip | -| `tests/error-handling.spec.ts` | Server RUN_ERROR, aimock error fixture | -| `tests/tool-error.spec.ts` | Tool throws error, agentic loop continues | +| Spec file | What it covers | +| ------------------------------ | ---------------------------------------------------------------- | +| `tests/abort.spec.ts` | Stop button cancels in-flight generation | +| `tests/lazy-tools.spec.ts` | `__lazy__tool__discovery__` discovers and uses lazy tools | +| `tests/custom-events.spec.ts` | Server tool `emitCustomEvent` received by client | +| `tests/middleware.spec.ts` | `onChunk` transform, `onBeforeToolCall` skip | +| `tests/error-handling.spec.ts` | Server RUN_ERROR, aimock error fixture | +| `tests/tool-error.spec.ts` | Tool throws error, agentic loop continues | +| `tests/byok.spec.ts` | BYOK key rides in header (not body); missing key → `byokMissing` | ## 1. Quick Start diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 206aff7bf..5bdc0f331 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -19,6 +19,7 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-anthropic": "workspace:*", "@tanstack/ai-bedrock": "workspace:*", + "@tanstack/ai-byok": "workspace:*", "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", diff --git a/testing/e2e/src/lib/providers.ts b/testing/e2e/src/lib/providers.ts index fa186a19e..6155a14b5 100644 --- a/testing/e2e/src/lib/providers.ts +++ b/testing/e2e/src/lib/providers.ts @@ -45,8 +45,13 @@ export function createTextAdapter( _aimockPort?: number, testId?: string, feature?: Feature, + // Per-request key override — the BYOK route passes the key read from the + // request header here, proving it flows client → header → adapter. aimock + // ignores auth, so the value only needs to be non-empty. + apiKeyOverride?: string, ): { adapter: AnyTextAdapter } { const model = modelOverride ?? defaultModels[provider] + const apiKey = apiKeyOverride ?? DUMMY_KEY // OpenAI, Grok SDKs need /v1 in baseURL. Groq SDK appends /openai/v1/ internally. // Anthropic, Gemini, Ollama SDKs include their path prefixes internally @@ -77,7 +82,7 @@ export function createTextAdapter( const factories: Record { adapter: AnyTextAdapter }> = { openai: () => createChatOptions({ - adapter: createOpenaiChat(model as 'gpt-4o', DUMMY_KEY, { + adapter: createOpenaiChat(model as 'gpt-4o', apiKey, { baseURL: openaiUrl, defaultHeaders: testHeaders, }), diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 14ffbf330..9422e01ec 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -19,6 +19,7 @@ import { Route as DevtoolsRouteARouteImport } from './routes/devtools-route-a' import { Route as DevtoolsGenerationHooksRouteImport } from './routes/devtools-generation-hooks' import { Route as DevtoolsChatRouteImport } from './routes/devtools-chat' import { Route as ChatClientDefaultBridgeRouteImport } from './routes/chat-client-default-bridge' +import { Route as ByokRouteImport } from './routes/byok' import { Route as IndexRouteImport } from './routes/index' import { Route as ProviderIndexRouteImport } from './routes/$provider/index' import { Route as ApiVideoRouteImport } from './routes/api.video' @@ -46,6 +47,7 @@ import { Route as ApiMcpAppsCallRouteImport } from './routes/api.mcp-apps-call' import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiChatRouteImport } from './routes/api.chat' +import { Route as ApiByokChatRouteImport } from './routes/api.byok-chat' import { Route as ApiAudioRouteImport } from './routes/api.audio' import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire' import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage' @@ -108,6 +110,11 @@ const ChatClientDefaultBridgeRoute = ChatClientDefaultBridgeRouteImport.update({ path: '/chat-client-default-bridge', getParentRoute: () => rootRouteImport, } as any) +const ByokRoute = ByokRouteImport.update({ + id: '/byok', + path: '/byok', + getParentRoute: () => rootRouteImport, +} as any) const IndexRoute = IndexRouteImport.update({ id: '/', path: '/', @@ -247,6 +254,11 @@ const ApiChatRoute = ApiChatRouteImport.update({ path: '/api/chat', getParentRoute: () => rootRouteImport, } as any) +const ApiByokChatRoute = ApiByokChatRouteImport.update({ + id: '/api/byok-chat', + path: '/api/byok-chat', + getParentRoute: () => rootRouteImport, +} as any) const ApiAudioRoute = ApiAudioRouteImport.update({ id: '/api/audio', path: '/api/audio', @@ -306,6 +318,7 @@ const ApiAudioStreamRoute = ApiAudioStreamRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -322,6 +335,7 @@ export interface FileRoutesByFullPath { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -356,6 +370,7 @@ export interface FileRoutesByFullPath { } export interface FileRoutesByTo { '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -372,6 +387,7 @@ export interface FileRoutesByTo { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -407,6 +423,7 @@ export interface FileRoutesByTo { export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute + '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute @@ -423,6 +440,7 @@ export interface FileRoutesById { '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute '/api/audio': typeof ApiAudioRouteWithChildren + '/api/byok-chat': typeof ApiByokChatRoute '/api/chat': typeof ApiChatRoute '/api/image': typeof ApiImageRouteWithChildren '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -459,6 +477,7 @@ export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -475,6 +494,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -509,6 +529,7 @@ export interface FileRouteTypes { fileRoutesByTo: FileRoutesByTo to: | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -525,6 +546,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -559,6 +581,7 @@ export interface FileRouteTypes { id: | '__root__' | '/' + | '/byok' | '/chat-client-default-bridge' | '/devtools-chat' | '/devtools-generation-hooks' @@ -575,6 +598,7 @@ export interface FileRouteTypes { | '/api/anthropic-structured-usage' | '/api/arktype-tool-wire' | '/api/audio' + | '/api/byok-chat' | '/api/chat' | '/api/image' | '/api/lazy-tools-wire' @@ -610,6 +634,7 @@ export interface FileRouteTypes { } export interface RootRouteChildren { IndexRoute: typeof IndexRoute + ByokRoute: typeof ByokRoute ChatClientDefaultBridgeRoute: typeof ChatClientDefaultBridgeRoute DevtoolsChatRoute: typeof DevtoolsChatRoute DevtoolsGenerationHooksRoute: typeof DevtoolsGenerationHooksRoute @@ -626,6 +651,7 @@ export interface RootRouteChildren { ApiAnthropicStructuredUsageRoute: typeof ApiAnthropicStructuredUsageRoute ApiArktypeToolWireRoute: typeof ApiArktypeToolWireRoute ApiAudioRoute: typeof ApiAudioRouteWithChildren + ApiByokChatRoute: typeof ApiByokChatRoute ApiChatRoute: typeof ApiChatRoute ApiImageRoute: typeof ApiImageRouteWithChildren ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -726,6 +752,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ChatClientDefaultBridgeRouteImport parentRoute: typeof rootRouteImport } + '/byok': { + id: '/byok' + path: '/byok' + fullPath: '/byok' + preLoaderRoute: typeof ByokRouteImport + parentRoute: typeof rootRouteImport + } '/': { id: '/' path: '/' @@ -915,6 +948,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatRouteImport parentRoute: typeof rootRouteImport } + '/api/byok-chat': { + id: '/api/byok-chat' + path: '/api/byok-chat' + fullPath: '/api/byok-chat' + preLoaderRoute: typeof ApiByokChatRouteImport + parentRoute: typeof rootRouteImport + } '/api/audio': { id: '/api/audio' path: '/api/audio' @@ -1055,6 +1095,7 @@ const ApiVideoRouteWithChildren = ApiVideoRoute._addFileChildren( const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, + ByokRoute: ByokRoute, ChatClientDefaultBridgeRoute: ChatClientDefaultBridgeRoute, DevtoolsChatRoute: DevtoolsChatRoute, DevtoolsGenerationHooksRoute: DevtoolsGenerationHooksRoute, @@ -1071,6 +1112,7 @@ const rootRouteChildren: RootRouteChildren = { ApiAnthropicStructuredUsageRoute: ApiAnthropicStructuredUsageRoute, ApiArktypeToolWireRoute: ApiArktypeToolWireRoute, ApiAudioRoute: ApiAudioRouteWithChildren, + ApiByokChatRoute: ApiByokChatRoute, ApiChatRoute: ApiChatRoute, ApiImageRoute: ApiImageRouteWithChildren, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, diff --git a/testing/e2e/src/routes/api.byok-chat.ts b/testing/e2e/src/routes/api.byok-chat.ts new file mode 100644 index 000000000..e726c5083 --- /dev/null +++ b/testing/e2e/src/routes/api.byok-chat.ts @@ -0,0 +1,58 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + toServerSentEventsResponse, +} from '@tanstack/ai' +import { byokMissing, getByokKey } from '@tanstack/ai-byok/server' +import { createTextAdapter } from '@/lib/providers' + +/** + * BYOK relay for E2E. Reads the OpenAI key from the per-provider request header + * (never the body), hands it to the adapter for this one call, and streams the + * aimock-backed response. Returns a typed `byokMissing` 401 when the header is + * absent. Never persists or logs the key. + */ +export const Route = createFileRoute('/api/byok-chat')({ + server: { + handlers: { + POST: async ({ request }) => { + await import('@/lib/llmock-server').then((m) => m.ensureLLMock()) + + let params + try { + params = await chatParamsFromRequestBody(await request.json()) + } catch (error) { + return new Response( + error instanceof Error ? error.message : 'Bad request', + { status: 400 }, + ) + } + + const fp = params.forwardedProps as Record + const testId = typeof fp.testId === 'string' ? fp.testId : undefined + + // Header-only read. No key → typed 401 the client renders. + const apiKey = getByokKey(request, 'openai') + if (!apiKey) return byokMissing('openai') + + const adapterOptions = createTextAdapter( + 'openai', + undefined, + undefined, + testId, + 'chat', + apiKey, + ) + + const stream = chat({ + ...adapterOptions, + messages: params.messages, + threadId: params.threadId, + runId: params.runId, + }) + return toServerSentEventsResponse(stream) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/byok.tsx b/testing/e2e/src/routes/byok.tsx new file mode 100644 index 000000000..ba9c5a40a --- /dev/null +++ b/testing/e2e/src/routes/byok.tsx @@ -0,0 +1,95 @@ +import { createFileRoute } from '@tanstack/react-router' +import { useMemo, useRef } from 'react' +import { fetchServerSentEvents, useChat } from '@tanstack/ai-react' +import { + ByokKeyManager, + ByokProvider, + byokHeaders, + memoryStorage, + useByok, +} from '@tanstack/ai-byok/react' +import { ChatUI } from '@/components/ChatUI' +import type { Keyring, KeyringStorage } from '@tanstack/ai-byok/react' + +interface ByokSearch { + testId?: string + aimockPort?: number + key?: string +} + +export const Route = createFileRoute('/byok')({ + component: ByokRoute, + validateSearch: (search: Record): ByokSearch => ({ + testId: typeof search.testId === 'string' ? search.testId : undefined, + aimockPort: + search.aimockPort != null ? Number(search.aimockPort) : undefined, + key: typeof search.key === 'string' ? search.key : undefined, + }), +}) + +// A storage that hydrates the keyring with a preloaded key on mount — the same +// path `passkeyStorage` uses to restore keys, driven deterministically for the +// test. No key param → session-only memory (nothing stored). +function preloadedStorage(keys: Keyring): KeyringStorage { + return { + id: 'preload', + label: 'Preloaded (test)', + persistent: false, + load: () => keys, + save: () => {}, + clear: () => {}, + } +} + +function ByokRoute() { + const { key } = Route.useSearch() + const storage = useMemo( + () => (key ? preloadedStorage({ openai: key }) : memoryStorage()), + [key], + ) + return ( + + + + ) +} + +function ByokChat() { + const { testId, aimockPort } = Route.useSearch() + const { keys } = useByok() + const keysRef = useRef(keys) + keysRef.current = keys + + const connection = useMemo( + () => + fetchServerSentEvents('/api/byok-chat', () => ({ + headers: byokHeaders(keysRef.current), + })), + [], + ) + + const chat = useChat({ + id: 'byok-chat', + connection, + body: { testId, aimockPort }, + }) + + return ( +
+ + {chat.error ? ( +

+ {chat.error.message} +

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