diff --git a/backend/src/config/settings.ts b/backend/src/config/settings.ts index 3b61e3b61..3dc4b99ea 100644 --- a/backend/src/config/settings.ts +++ b/backend/src/config/settings.ts @@ -105,7 +105,7 @@ const settingsSchema = z corsExposeHeaders: z .string() .default( - 'set-auth-token,X-Proxy-Final-Url,X-Proxy-Passthrough-Content-Type,X-Proxy-Passthrough-Mcp-Session-Id,X-Proxy-Passthrough-Mcp-Protocol-Version,X-Proxy-Passthrough-Location,X-Proxy-Passthrough-Anthropic-Version,WWW-Authenticate', + 'set-auth-token,X-Proxy-Final-Url,X-Proxy-Passthrough-Content-Type,X-Proxy-Passthrough-Mcp-Session-Id,X-Proxy-Passthrough-Mcp-Protocol-Version,X-Proxy-Passthrough-Location,X-Proxy-Passthrough-Anthropic-Version,WWW-Authenticate,Ehbp-Response-Nonce', ), // E2E encryption — when true, devices must complete the trust flow before syncing @@ -220,7 +220,7 @@ const parseSettings = (): Settings => { corsAllowHeaders: process.env.CORS_ALLOW_HEADERS || '', corsExposeHeaders: process.env.CORS_EXPOSE_HEADERS || - 'set-auth-token,X-Proxy-Final-Url,X-Proxy-Passthrough-Content-Type,X-Proxy-Passthrough-Mcp-Session-Id,X-Proxy-Passthrough-Mcp-Protocol-Version,X-Proxy-Passthrough-Location,X-Proxy-Passthrough-Anthropic-Version,WWW-Authenticate', + 'set-auth-token,X-Proxy-Final-Url,X-Proxy-Passthrough-Content-Type,X-Proxy-Passthrough-Mcp-Session-Id,X-Proxy-Passthrough-Mcp-Protocol-Version,X-Proxy-Passthrough-Location,X-Proxy-Passthrough-Anthropic-Version,WWW-Authenticate,Ehbp-Response-Nonce', e2eeEnabled: process.env.E2EE_ENABLED === 'true', minAppVersion: process.env.MIN_APP_VERSION || '', swaggerEnabled: process.env.SWAGGER_ENABLED === 'true', diff --git a/backend/src/proxy/streaming.test.ts b/backend/src/proxy/streaming.test.ts index ac8424ff8..5b1b531e8 100644 --- a/backend/src/proxy/streaming.test.ts +++ b/backend/src/proxy/streaming.test.ts @@ -50,6 +50,21 @@ describe('capStream', () => { expect(capped.bytesRead()).toBe(data.byteLength) }) + it('forwards all bytes when the byte cap is omitted', async () => { + const chunks = [new Uint8Array(6), new Uint8Array(5)] + const aborts: string[] = [] + const capped = capStream(makeStream(chunks), { + idleTimeoutMs: 5000, + onAbort: (reason) => aborts.push(reason), + }) + + const result = await collectStream(capped.stream) + + expect(result.byteLength).toBe(11) + expect(aborts).toEqual([]) + expect(capped.bytesRead()).toBe(11) + }) + it('calls onAbort("cap") and terminates when bytes exceed cap', async () => { const chunk1 = new Uint8Array(6) const chunk2 = new Uint8Array(5) @@ -79,6 +94,74 @@ describe('capStream', () => { expect(aborts).toEqual(['idle']) }) + it('errors with the configured reason when error mode expires', async () => { + const idleError = new DOMException('opaque stream idle timeout', 'TimeoutError') + const aborts: string[] = [] + const completions: number[] = [] + const stalled = new ReadableStream({ start() {} }) + const capped = capStream(stalled, { + idleTimeoutMs: 0, + onIdle: 'error', + idleError, + onAbort: (reason) => aborts.push(reason), + onComplete: (bytesRead) => completions.push(bytesRead), + }) + + await expect(capped.stream.getReader().read()).rejects.toBe(idleError) + expect(aborts).toEqual(['idle']) + expect(completions).toEqual([0]) + }) + + it('keeps an externally handled idle error pending after the source aborts', async () => { + const idleError = new DOMException('opaque stream idle timeout', 'TimeoutError') + const upstreamController = new AbortController() + const handledError = Promise.withResolvers() + const stalled = new ReadableStream({ + start(controller) { + upstreamController.signal.addEventListener('abort', () => controller.error(idleError), { once: true }) + }, + }) + const capped = capStream(stalled, { + idleTimeoutMs: 0, + onIdle: 'error', + idleError, + onAbort: () => upstreamController.abort(idleError), + onIdleError: (error) => handledError.resolve(error), + }) + const reader = capped.stream.getReader() + const pendingRead = reader.read() + + expect(await handledError.promise).toBe(idleError) + const readState = await Promise.race([ + pendingRead.then( + () => 'settled', + () => 'settled', + ), + new Promise<'pending'>((resolve) => setTimeout(() => resolve('pending'), 10)), + ]) + expect(readState).toBe('pending') + + await reader.cancel() + await pendingRead + }) + + it('still relays source errors before an external idle handler fires', async () => { + const sourceError = new Error('upstream failed') + const failed = new ReadableStream({ + start(controller) { + controller.error(sourceError) + }, + }) + const capped = capStream(failed, { + idleTimeoutMs: 1000, + onIdle: 'error', + onAbort: () => {}, + onIdleError: () => {}, + }) + + await expect(capped.stream.getReader().read()).rejects.toBe(sourceError) + }) + it('resets idle timer on each chunk so a slow-but-steady stream completes', async () => { const aborts: string[] = [] const chunkDelay = 10 @@ -158,6 +241,27 @@ describe('capStream', () => { expect(aborts).toEqual([]) }) + it('clears idle timer and completes once when the source errors', async () => { + const sourceError = new Error('upstream failed') + const aborts: string[] = [] + const completions: number[] = [] + const failed = new ReadableStream({ + start(controller) { + controller.error(sourceError) + }, + }) + const capped = capStream(failed, { + idleTimeoutMs: 0, + onAbort: (reason) => aborts.push(reason), + onComplete: (bytesRead) => completions.push(bytesRead), + }) + + await expect(capped.stream.getReader().read()).rejects.toBe(sourceError) + await new Promise((resolve) => setImmediate(resolve)) + expect(aborts).toEqual([]) + expect(completions).toEqual([0]) + }) + // --------------------------------------------------------------------------- // Byte counter + onComplete contract — the observability layer reads these // --------------------------------------------------------------------------- diff --git a/backend/src/proxy/streaming.ts b/backend/src/proxy/streaming.ts index 11f00f057..5d6f9ab14 100644 --- a/backend/src/proxy/streaming.ts +++ b/backend/src/proxy/streaming.ts @@ -2,39 +2,51 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ +const defaultIdleTimeoutError = new DOMException('stream idle timeout', 'TimeoutError') + +export type CappedStream = { + stream: ReadableStream + /** Total bytes that flowed through the cap. Read after stream completion. */ + bytesRead: () => number +} + +/** Controls byte limiting and idle-expiry behavior for a relayed stream. */ +export type CapStreamOptions = { + /** Maximum bytes accepted before termination. Omit to disable byte limiting. */ + maxBytes?: number + idleTimeoutMs: number + /** Whether idle expiry terminates or errors the relayed stream. Defaults to termination. */ + onIdle?: 'terminate' | 'error' + /** Error exposed to readers when `onIdle` is `error`. */ + idleError?: Error + /** Handles error-mode idle expiry at the transport layer. Caller must tear + * down the transport because the relayed stream stays pending. */ + onIdleError?: (error: Error) => void + onAbort: (reason: 'cap' | 'idle') => void + /** Fired exactly once after the stream finishes (graceful close, cap-hit, + * idle, source error, or downstream cancel). Receives the total bytes + * that flowed through. Use for post-stream observability emission. */ + onComplete?: (bytesRead: number) => void +} + /** - * Wraps a ReadableStream with a byte-cap and idle-timeout TransformStream. - * When either limit is exceeded the stream is terminated (not errored) so the - * client receives a truncated but valid chunked response — the proxy cannot - * retroactively change the HTTP status because headers have already been sent. - * `onAbort` is called first so the caller can abort the upstream connection. + * Wraps a ReadableStream with an optional byte cap and idle watchdog. + * Limit expiry terminates by default because response headers have already + * been sent. Error mode rejects unsafe truncation by default or delegates it + * to a transport handler. `onAbort` fires first so callers can abort upstream. * * Returns `bytesRead()` so observability can record the actual transferred byte * count after the stream has been consumed. With `content-encoding` passthrough * the bytes counted are post-compression (what the wire saw), which is exactly * what we want to log. */ -export type CappedStream = { - stream: ReadableStream - /** Total bytes that flowed through the cap. Read after stream completion. */ - bytesRead: () => number -} - -export const capStream = ( - source: ReadableStream, - opts: { - maxBytes: number - idleTimeoutMs: number - onAbort: (reason: 'cap' | 'idle') => void - /** Fired exactly once after the stream finishes (graceful close, cap-hit, - * idle, source error, or downstream cancel). Receives the total bytes - * that flowed through. Use for post-stream observability emission. */ - onComplete?: (bytesRead: number) => void - }, -): CappedStream => { +export const capStream = (source: ReadableStream, opts: CapStreamOptions): CappedStream => { let bytesReceived = 0 let idleTimer: ReturnType | undefined let completed = false + let idleHandled = false + const idleError = opts.idleError ?? defaultIdleTimeoutError + const externalIdleHandler = opts.onIdle === 'error' ? opts.onIdleError : undefined const fireComplete = () => { if (completed) { @@ -44,44 +56,79 @@ export const capStream = ( opts.onComplete?.(bytesReceived) } - const resetIdleTimer = (controller: TransformStreamDefaultController) => { - clearTimeout(idleTimer) + const armIdleTimer = (controller: TransformStreamDefaultController) => { + if (idleTimer) { + idleTimer.refresh() + return + } + idleTimer = setTimeout(() => { + idleTimer = undefined + idleHandled = externalIdleHandler !== undefined opts.onAbort('idle') - controller.terminate() + if (opts.onIdle !== 'error') { + controller.terminate() + } else if (externalIdleHandler) { + externalIdleHandler(idleError) + } else { + controller.error(idleError) + } fireComplete() }, opts.idleTimeoutMs) } const { readable, writable } = new TransformStream({ start(controller) { - resetIdleTimer(controller) + armIdleTimer(controller) }, transform(chunk, controller) { bytesReceived += chunk.byteLength - if (bytesReceived > opts.maxBytes) { + if (opts.maxBytes !== undefined && bytesReceived > opts.maxBytes) { clearTimeout(idleTimer) + idleTimer = undefined opts.onAbort('cap') controller.terminate() fireComplete() return } controller.enqueue(chunk) - resetIdleTimer(controller) + armIdleTimer(controller) }, flush() { clearTimeout(idleTimer) + idleTimer = undefined fireComplete() }, }) - source.pipeTo(writable).catch(() => { - // pipeTo rejects when source errors OR when writable is aborted (e.g., downstream - // was cancelled). Clear the idle timer here so it doesn't fire after the stream - // has been torn down — running terminate() on an errored controller throws. + // pipeTo rejects when the source errors OR the writable is aborted (downstream + // cancel). Clear the idle timer so it can't fire after teardown — terminate() + // on an errored controller throws. + const finishRelay = () => { clearTimeout(idleTimer) + idleTimer = undefined fireComplete() - }) + } + + if (externalIdleHandler) { + /** Keeps the relayed body pending after idle expiry so Bun can reset the transport. */ + const relaySource = async () => { + try { + await source.pipeTo(writable, { preventAbort: true, preventClose: true }) + if (!idleHandled) { + await writable.close() + } + } catch (error) { + if (!idleHandled) { + await writable.abort(error).catch(() => {}) + } + } + } + + void relaySource().finally(finishRelay) + } else { + source.pipeTo(writable).catch(finishRelay) + } return { stream: readable, diff --git a/backend/src/tinfoil/routes.test.ts b/backend/src/tinfoil/routes.test.ts index 8a4ae255a..1e8be6c9f 100644 --- a/backend/src/tinfoil/routes.test.ts +++ b/backend/src/tinfoil/routes.test.ts @@ -11,6 +11,9 @@ import { createTinfoilRoutes } from './routes' const enclaveUrl = 'https://inference.tinfoil.sh' const testApiKey = 'test-tinfoil-key' +const upstreamTimeoutMessage = 'tinfoil upstream timeout' +const upstreamIdleTimeoutMessage = 'tinfoil upstream idle timeout' +const realFetch = (globalThis as Record).__originalFetch as typeof fetch const makeOkResponse = (body = 'ok', extraHeaders: Record = {}) => new Response(body, { @@ -26,6 +29,47 @@ const drain = async (res: Response): Promise => { return res } +/** Creates a fetch implementation whose pending request or response body follows its signal. */ +const createAbortableFetch = (response?: Response) => { + const started = Promise.withResolvers() + const aborted = Promise.withResolvers() + const fetchFn = ((_input: RequestInfo | URL, init?: RequestInit) => { + const signal = init?.signal as AbortSignal + started.resolve(signal) + + if (!response) { + return new Promise((_resolve, reject) => { + const rejectOnAbort = () => { + aborted.resolve(signal.reason) + reject(signal.reason) + } + signal.addEventListener('abort', rejectOnAbort, { once: true }) + if (signal.aborted) { + rejectOnAbort() + } + }) + } + + const trackAbort = () => aborted.resolve(signal.reason) + signal.addEventListener('abort', trackAbort, { once: true }) + if (signal.aborted) { + trackAbort() + return Promise.reject(signal.reason) + } + + const body = response.body?.pipeThrough(new TransformStream(), { signal }) ?? null + return Promise.resolve( + new Response(body, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }), + ) + }) as unknown as typeof fetch + + return { fetchFn, started: started.promise, aborted: aborted.promise } +} + describe('createTinfoilRoutes', () => { let mockFetch: ReturnType let consoleSpies: ConsoleSpies @@ -45,13 +89,24 @@ describe('createTinfoilRoutes', () => { consoleSpies.error.mockClear() }) - const buildApp = (overrides: { apiKey?: string; enclaveUrl?: string; auth?: typeof mockAuth } = {}) => + const buildApp = ( + overrides: { + apiKey?: string + enclaveUrl?: string + auth?: typeof mockAuth + fetchFn?: typeof fetch + upstreamHeadersTimeoutMs?: number + upstreamIdleTimeoutMs?: number + } = {}, + ) => new Elysia().use( createTinfoilRoutes({ auth: overrides.auth ?? mockAuth, - fetchFn: mockFetch as unknown as typeof fetch, + fetchFn: overrides.fetchFn ?? (mockFetch as unknown as typeof fetch), apiKey: overrides.apiKey ?? testApiKey, enclaveUrl: overrides.enclaveUrl ?? enclaveUrl, + upstreamHeadersTimeoutMs: overrides.upstreamHeadersTimeoutMs, + upstreamIdleTimeoutMs: overrides.upstreamIdleTimeoutMs, }), ) @@ -125,9 +180,89 @@ describe('createTinfoilRoutes', () => { expect(sent.get('cookie')).toBeNull() expect(sent.get('connection')).toBeNull() }) + + it('strips response hop-by-hop headers while preserving content encoding', async () => { + mockFetch.mockResolvedValueOnce( + makeOkResponse('opaque', { + connection: 'keep-alive', + 'transfer-encoding': 'chunked', + 'content-encoding': 'br', + }), + ) + const app = buildApp() + + const res = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + + expect(res.headers.get('connection')).toBeNull() + expect(res.headers.get('transfer-encoding')).toBeNull() + expect(res.headers.get('content-encoding')).toBe('br') + await res.arrayBuffer() + }) + + it('strips upstream CORS headers so only our own middleware sets them', async () => { + // The enclave emits a duplicated `Access-Control-Allow-Credentials: true, true` + // that browsers reject; relaying any upstream access-control-* would fight + // our cors() middleware. + mockFetch.mockResolvedValueOnce( + makeOkResponse('opaque', { + 'access-control-allow-credentials': 'true, true', + 'access-control-allow-origin': '*', + 'access-control-expose-headers': 'X-Upstream-Only', + 'ehbp-response-nonce': 'abc123', + }), + ) + const app = buildApp() + + const res = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + + expect(res.headers.get('access-control-allow-credentials')).toBeNull() + expect(res.headers.get('access-control-allow-origin')).toBeNull() + expect(res.headers.get('access-control-expose-headers')).toBeNull() + expect(res.headers.get('ehbp-response-nonce')).toBe('abc123') + await res.arrayBuffer() + }) }) describe('body forwarding', () => { + it('relays EHBP request and response bytes and headers unchanged', async () => { + const requestPayload = new Uint8Array([0x00, 0xff, 0x01, 0x80]) + const responseChunks = [new Uint8Array([0xfe, 0x02]), new Uint8Array([0x00, 0x7f])] + const requestBody = Promise.withResolvers() + const requestEncapsulatedKey = Promise.withResolvers() + const fetchFn = (async (_input: RequestInfo | URL, init?: RequestInit) => { + const headers = init?.headers as Headers + const body = init?.body as ReadableStream + requestEncapsulatedKey.resolve(headers.get('ehbp-encapsulated-key')) + requestBody.resolve(new Uint8Array(await new Response(body).arrayBuffer())) + + return new Response( + new ReadableStream({ + start(controller) { + for (const chunk of responseChunks) { + controller.enqueue(chunk) + } + controller.close() + }, + }), + { headers: { 'Ehbp-Response-Nonce': 'response-nonce' } }, + ) + }) as unknown as typeof fetch + const app = buildApp({ fetchFn }) + + const res = await app.handle( + new Request('http://localhost/tinfoil/v1/chat/completions', { + method: 'POST', + headers: { 'Ehbp-Encapsulated-Key': 'encapsulated-key' }, + body: requestPayload, + }), + ) + + expect(await requestBody.promise).toEqual(requestPayload) + expect(await requestEncapsulatedKey.promise).toBe('encapsulated-key') + expect(new Uint8Array(await res.arrayBuffer())).toEqual(new Uint8Array([0xfe, 0x02, 0x00, 0x7f])) + expect(res.headers.get('ehbp-response-nonce')).toBe('response-nonce') + }) + it('forwards the request body for POST requests', async () => { const app = buildApp() const payload = new Uint8Array([0x01, 0x02, 0x03, 0x04]) @@ -147,6 +282,9 @@ describe('createTinfoilRoutes', () => { expect(calledUrl).toBe(`${enclaveUrl}/v1/chat/completions`) expect(init.body).not.toBeNull() expect(init.method).toBe('POST') + expect(init.redirect).toBe('manual') + expect((init as RequestInit & { decompress: boolean }).decompress).toBeFalse() + expect((init as RequestInit & { duplex: string }).duplex).toBe('half') }) it('forwards JSON bodies untouched (parse: none keeps the stream intact)', async () => { @@ -177,6 +315,148 @@ describe('createTinfoilRoutes', () => { }) }) + describe('upstream timeouts and aborts', () => { + it('returns 504 and aborts upstream when response headers time out', async () => { + const upstream = createAbortableFetch() + const app = buildApp({ + fetchFn: upstream.fetchFn, + upstreamHeadersTimeoutMs: 0, + }) + + const res = await app.handle( + new Request('http://localhost/tinfoil/v1/chat/completions', { + method: 'POST', + body: 'opaque-bytes', + }), + ) + + expect(res.status).toBe(504) + expect(res.headers.get('content-type')).toBe('text/plain') + const body = await res.text() + expect(body).toBe(upstreamTimeoutMessage) + expect(await upstream.aborted).toBeInstanceOf(DOMException) + expect(body).not.toContain(testApiKey) + }) + + it('errors a stalled response stream and aborts upstream', async () => { + const firstChunk = new Uint8Array([0x01, 0xff]) + const upstream = createAbortableFetch( + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(firstChunk) + }, + }), + ), + ) + const app = buildApp({ + fetchFn: upstream.fetchFn, + upstreamIdleTimeoutMs: 0, + }) + + const res = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + const reader = res.body!.getReader() + + expect(await reader.read()).toEqual({ done: false, value: firstChunk }) + const stalledRead = reader.read() + await expect(stalledRead).rejects.toThrow(upstreamIdleTimeoutMessage) + await expect(stalledRead).rejects.not.toThrow(testApiKey) + expect(await upstream.aborted).toBeInstanceOf(DOMException) + }) + + it('abruptly closes the socket when a relayed response stalls', async () => { + const responseChunks = [new Uint8Array([0x01, 0x02]), new Uint8Array([0x03, 0x04])] + const expectedBytes = new Uint8Array(responseChunks.flatMap((chunk) => [...chunk])) + const upstream = createAbortableFetch( + new Response( + new ReadableStream({ + start(controller) { + for (const chunk of responseChunks) { + controller.enqueue(chunk) + } + }, + }), + ), + ) + const app = buildApp({ + fetchFn: upstream.fetchFn, + upstreamIdleTimeoutMs: 20, + }) + // Real socket via Elysia's own listen so the route receives the genuine + // `ctx.server` — hand-wiring `app.server` would bypass the seam under test. + await new Promise((resolve) => { + app.listen({ port: 0, hostname: '127.0.0.1' }, () => resolve()) + }) + + try { + const res = await realFetch(new URL('/tinfoil/v1/models', app.server!.url)) + const reader = res.body!.getReader() + const receivedBytes: number[] = [] + + while (receivedBytes.length < expectedBytes.byteLength) { + const result = await reader.read() + if (result.done) { + throw new Error('response ended before relaying both upstream chunks') + } + receivedBytes.push(...result.value) + } + + expect(new Uint8Array(receivedBytes)).toEqual(expectedBytes) + await expect(reader.read()).rejects.toThrow() + expect(await upstream.aborted).toBeInstanceOf(DOMException) + } finally { + await app.stop(true) + } + }) + + it('aborts upstream when the client request aborts', async () => { + const upstream = createAbortableFetch() + const clientController = new AbortController() + const app = buildApp({ fetchFn: upstream.fetchFn }) + const response = app.handle( + new Request('http://localhost/tinfoil/v1/chat/completions', { + method: 'POST', + body: 'opaque-bytes', + signal: clientController.signal, + }), + ) + + await upstream.started + clientController.abort() + + expect(await upstream.aborted).toBe(clientController.signal.reason) + expect(await (await response).text()).not.toContain(testApiKey) + }) + + it('keeps a steady response stream alive beyond the headers timeout', async () => { + const chunks = [new Uint8Array([0x01]), new Uint8Array([0x02]), new Uint8Array([0x03]), new Uint8Array([0x04])] + const upstream = createAbortableFetch( + new Response( + new ReadableStream({ + async start(controller) { + for (const chunk of chunks) { + controller.enqueue(chunk) + await new Promise((resolve) => setImmediate(resolve)) + } + controller.close() + }, + }), + ), + ) + const app = buildApp({ + fetchFn: upstream.fetchFn, + upstreamHeadersTimeoutMs: 0, + upstreamIdleTimeoutMs: 50, + }) + + const res = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + + expect(new Uint8Array(await res.arrayBuffer())).toEqual(new Uint8Array([0x01, 0x02, 0x03, 0x04])) + await new Promise((resolve) => setImmediate(resolve)) + expect((await upstream.started).aborted).toBeFalse() + }) + }) + describe('upstream URL derivation', () => { it('derives the upstream path from the wildcard, not the outer mount prefix', async () => { // Mount at a non-default outer prefix to prove the path comes from the wildcard. diff --git a/backend/src/tinfoil/routes.ts b/backend/src/tinfoil/routes.ts index 6c9fd2c9e..e318a009c 100644 --- a/backend/src/tinfoil/routes.ts +++ b/backend/src/tinfoil/routes.ts @@ -6,10 +6,18 @@ import type { Auth } from '@/auth/elysia-plugin' import { createAuthMacro } from '@/auth/elysia-plugin' import { getSettings } from '@/config/settings' import { safeErrorHandler } from '@/middleware/error-handling' +import { capStream } from '@/proxy/streaming' +import { filterHeaders } from '@/utils/request' +import { tinfoilUpstreamIdleTimeoutMessage, tinfoilUpstreamTimeoutMessage } from '@shared/tinfoil-proxy' import { Elysia, type AnyElysia } from 'elysia' const allowedMethods = new Set(['GET', 'POST', 'OPTIONS']) const bodylessMethods = new Set(['GET', 'OPTIONS']) +const defaultUpstreamHeadersTimeoutMs = 30_000 +const defaultUpstreamIdleTimeoutMs = 60_000 +const abruptResponseCloseTimeoutSeconds = 1 +const upstreamHeadersTimeoutError = new DOMException(tinfoilUpstreamTimeoutMessage, 'TimeoutError') +const upstreamIdleTimeoutError = new DOMException(tinfoilUpstreamIdleTimeoutMessage, 'TimeoutError') const textResponse = (status: number, body: string): Response => new Response(body, { status, headers: { 'Content-Type': 'text/plain' } }) @@ -23,6 +31,10 @@ export type CreateTinfoilRoutesOptions = { apiKey?: string /** Override the upstream enclave URL. Defaults to `TINFOIL_ENCLAVE_URL`. */ enclaveUrl?: string + /** Time allowed for upstream response headers. Defaults to 30 seconds. */ + upstreamHeadersTimeoutMs?: number + /** Maximum idle time between upstream response chunks. Defaults to 60 seconds. */ + upstreamIdleTimeoutMs?: number } export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { @@ -31,8 +43,14 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { const settings = getSettings() const apiKey = options.apiKey ?? settings.tinfoilApiKey const enclaveUrl = (options.enclaveUrl ?? settings.tinfoilEnclaveUrl).replace(/\/$/, '') + const upstreamHeadersTimeoutMs = options.upstreamHeadersTimeoutMs ?? defaultUpstreamHeadersTimeoutMs + const upstreamIdleTimeoutMs = options.upstreamIdleTimeoutMs ?? defaultUpstreamIdleTimeoutMs - const proxyToEnclave = async (request: Request, wildcard: string): Promise => { + const proxyToEnclave = async ( + request: Request, + wildcard: string, + server: Bun.Server | null, + ): Promise => { const method = request.method.toUpperCase() if (!allowedMethods.has(method)) { @@ -58,33 +76,58 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { headers.set('Authorization', `Bearer ${apiKey}`) const body = bodylessMethods.has(method) ? null : request.body + const upstreamController = new AbortController() + const signal = AbortSignal.any([request.signal, upstreamController.signal]) + const headersTimeoutId = setTimeout( + () => upstreamController.abort(upstreamHeadersTimeoutError), + upstreamHeadersTimeoutMs, + ) // Bun-specific fetch options: `duplex: 'half'` enables streaming request // bodies; `decompress: false` keeps the HPKE-encrypted bytes opaque on // the response path so the frontend SDK can decrypt them as-is. - const upstream = await fetchFn(upstreamUrl, { - method, - headers, - body, - redirect: 'manual', - decompress: false, - duplex: 'half', - } as RequestInit & { decompress: boolean; duplex: 'half' }) - - const responseHeaders = new Headers() - upstream.headers.forEach((value, key) => { - const lower = key.toLowerCase() - if (lower === 'transfer-encoding' || lower === 'connection') { - return - } - responseHeaders.set(key, value) - }) + try { + const upstream = await fetchFn(upstreamUrl, { + method, + headers, + body, + signal, + redirect: 'manual', + decompress: false, + duplex: 'half', + } as RequestInit & { decompress: boolean; duplex: 'half' }) - return new Response(upstream.body, { - status: upstream.status, - statusText: upstream.statusText, - headers: responseHeaders, - }) + // Strip upstream CORS headers: the enclave emits a duplicated + // `Access-Control-Allow-Credentials: true, true`, which browsers reject + // outright. Our own cors() middleware sets the correct CORS headers for + // our origin (including Ehbp-Response-Nonce in expose-headers). + const responseHeaders = filterHeaders(upstream.headers, ['transfer-encoding', 'connection', /^access-control-/i]) + + const responseBody = upstream.body + ? capStream(upstream.body, { + idleTimeoutMs: upstreamIdleTimeoutMs, + onIdle: 'error', + idleError: upstreamIdleTimeoutError, + onAbort: () => upstreamController.abort(upstreamIdleTimeoutError), + // Bun serializes controller.error() after headers as clean chunked EOF. + // Keep body pending and let native request timeout reset socket instead. + onIdleError: server ? () => server.timeout(request, abruptResponseCloseTimeoutSeconds) : undefined, + }).stream + : null + + return new Response(responseBody, { + status: upstream.status, + statusText: upstream.statusText, + headers: responseHeaders, + }) + } catch (error) { + if (error === upstreamHeadersTimeoutError) { + return textResponse(504, tinfoilUpstreamTimeoutMessage) + } + throw error + } finally { + clearTimeout(headersTimeoutId) + } } // `{ parse: 'none' }` keeps the request stream untouched so the HPKE-encrypted @@ -100,8 +143,8 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { if (rateLimit) { return g .use(rateLimit) - .all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? ''), { parse: 'none' }) + .all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? '', ctx.server), { parse: 'none' }) } - return g.all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? ''), { parse: 'none' }) + return g.all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? '', ctx.server), { parse: 'none' }) }) } diff --git a/shared/tinfoil-proxy.ts b/shared/tinfoil-proxy.ts new file mode 100644 index 000000000..aa0a69f4c --- /dev/null +++ b/shared/tinfoil-proxy.ts @@ -0,0 +1,13 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +/** + * Wire contract between the backend tinfoil proxy and the frontend error + * classifier: the proxy returns these exact strings as the 504 body (headers + * timeout) and the mid-stream idle error message, and the frontend matches + * them to classify the failure as a timeout. Shared so a reword cannot + * silently break classification on the other side of the boundary. + */ +export const tinfoilUpstreamTimeoutMessage = 'tinfoil upstream timeout' +export const tinfoilUpstreamIdleTimeoutMessage = 'tinfoil upstream idle timeout' diff --git a/src/acp/built-in-adapter.ts b/src/acp/built-in-adapter.ts index 931441f7e..348b2bbdc 100644 --- a/src/acp/built-in-adapter.ts +++ b/src/acp/built-in-adapter.ts @@ -507,6 +507,7 @@ export const createBuiltInAdapter = (agent: Agent, options: BuiltInAdapterOption reconnectClient: context.reconnectClient, httpClient: context.httpClient, getProxyFetch: context.getProxyFetch, + turnBudget: context.turnBudget, }) // Route tool-capable Pi-serviceable models (anthropic + the OpenAI-wire family) diff --git a/src/ai/fetch.ts b/src/ai/fetch.ts index b7ddacfb8..e3638a3ca 100644 --- a/src/ai/fetch.ts +++ b/src/ai/fetch.ts @@ -3,6 +3,7 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ import { createPromptParts } from '@/ai/prompt' +import { createTurnBudget, createTurnBudgetExhaustedError, type TurnBudgetConsumer } from '@/ai/retry-budget' import { buildStepOverrides, extractTextFromMessages, @@ -24,7 +25,8 @@ import { getLocalSetting } from '@/stores/local-settings-store' import { hydrateAttachmentsAsFileParts } from '@/lib/attachments' import { hydrateQuotesAsText } from '@/lib/quotes' import { isSsoMode } from '@/lib/auth-mode' -import { getAuthToken, getUserCacheSecret } from '@/lib/auth-token' +import { getAuthToken } from '@/lib/auth-token' +import { classifyErrorKind } from '@/lib/error-utils' import { fetch as baseFetch } from '@/lib/fetch' import { isLoopbackHost } from '@/lib/mcp-url-validation' import { normalizeOpenAiBaseUrl } from '@/lib/openai-base-url' @@ -36,7 +38,6 @@ import { createAnthropic } from '@ai-sdk/anthropic' import { createOpenAI } from '@ai-sdk/openai' import { createOpenAICompatible } from '@ai-sdk/openai-compatible' import type { HttpClient } from '@/lib/http' -import type { SecureClient } from 'tinfoil' import { v7 as uuidv7 } from 'uuid' // Currently @openrouter/ai-sdk-provider is NOT compatible with Vercel AI SDK v5. If you enable this, you will get the following error: @@ -66,6 +67,13 @@ import { smoothStreamWordDelayMs } from '@/chats/chat-throttle' import type { SkillDefinition } from '@shared/agent-core/skills' import { detectStreamChunk } from './smooth-chunking' import { createMessageMetadata } from './message-metadata' +import { + evictSystemTinfoilClient, + evictUserTinfoilClient, + getSystemTinfoilClient, + getTinfoilClient, + isTinfoilTransportWedgedError, +} from './tinfoil-client' /** * Sanitizes a server name into a valid tool prefix. @@ -92,96 +100,6 @@ export const ollama = createOpenAI({ fetch, }) -// Cached so attestation runs once per page load. `tinfoil` is dynamically -// imported to code-split its attestation/crypto deps. -// -// system: HPKE body POSTs to /tinfoil; backend injects our key. -// user: BYOK — direct to the enclave with the user's own key. -// -// System cache is keyed by cloudUrl so a dev-tools URL switch hits the new -// backend on the next call. -const systemTinfoilClients = new Map>() -let userTinfoilClient: SecureClient | null = null - -/** - * Build a fresh system `SecureClient` promise and cache it synchronously - * (before the dynamic `import('tinfoil')` resolves) so a prewarm and an - * immediate first send share one client and attest once instead of racing into - * two attestations. On construction failure we drop the entry so the next call - * retries the import rather than inheriting a sticky rejection. - */ -const createSystemTinfoilClient = (cloudUrl: string): Promise => { - const clientPromise = import('tinfoil').then( - ({ SecureClient }) => new SecureClient({ baseURL: `${cloudUrl}/tinfoil`, userCacheSecret: getUserCacheSecret() }), - ) - void clientPromise.catch(() => systemTinfoilClients.delete(cloudUrl)) - systemTinfoilClients.set(cloudUrl, clientPromise) - return clientPromise -} - -export const getSystemTinfoilClient = async (): Promise => { - // cloudUrl already ends in /v1 (shared with the OpenAI chat baseURL). - const cloudUrl = getLocalSetting('cloudUrl').replace(/\/$/, '') - // Reuse the cached construction promise across concurrent callers; `ready()` - // is awaited per call below (idempotent once attested). - const client = await (systemTinfoilClients.get(cloudUrl) ?? createSystemTinfoilClient(cloudUrl)) - await client.ready() - return client -} - -/** - * Best-effort warm-up of the Tinfoil system enclave so the first chat send - * doesn't pay the attestation handshake on the critical path. Fired (fire-and- - * forget) from the chat-ready path for the built-in agent only — see - * {@link useHydrateChatStore}; ACP agents route over the wire and never reach - * {@link createModel}. No-op unless `model` is a Tinfoil *system* model (the - * only path that attests via {@link getSystemTinfoilClient}); BYO/other - * providers never attest here. - * - * Idempotent: `getSystemTinfoilClient` memoizes per cloudUrl, so repeated warm- - * ups and a concurrent real send share the same in-flight client. Errors are - * swallowed ONLY here because this is a speculative cache fill — the real send - * still surfaces attestation failures loudly through {@link createModel}. - */ -export const runSystemModelPrewarm = async (model: Pick | null | undefined) => { - if (!model || model.provider !== 'tinfoil' || !model.isSystem) { - return - } - try { - await getSystemTinfoilClient() - } catch (error) { - console.warn('runSystemModelPrewarm: warm-up skipped', error) - } -} - -/** Drop the cached `SecureClient` so the next send constructs a fresh one with - * a new attestation context. Use when a key-config error keeps repeating - * inside the SDK's own reset+retry — the cached client's transport is wedged - * and only a brand-new instance breaks the cycle. */ -export const evictSystemTinfoilClient = (): void => { - const cloudUrl = getLocalSetting('cloudUrl').replace(/\/$/, '') - systemTinfoilClients.delete(cloudUrl) -} - -export const getTinfoilClient = async (): Promise => { - if (!userTinfoilClient) { - const { SecureClient } = await import('tinfoil') - userTinfoilClient = new SecureClient({ userCacheSecret: getUserCacheSecret() }) - } - await userTinfoilClient.ready() - return userTinfoilClient -} - -const evictUserTinfoilClient = (): void => { - userTinfoilClient = null -} - -/** A KeyConfigMismatchError that survives the SDK's internal reset+retry means - * our cached `SecureClient` has a wedged transport. Evict it so the next call - * builds a fresh instance with a brand-new attestation context. */ -export const isKeyConfigMismatchError = (err: unknown): boolean => - err instanceof Error && err.name === 'KeyConfigMismatchError' - /** Reconnect a dropped MCP client; returns a fresh client or null. Supplied by * the MCP provider via the chat store. See `src/lib/mcp-provider.tsx`. */ type ReconnectClient = (client: MCPClient) => Promise @@ -192,6 +110,7 @@ type AiFetchStreamingResponseOptions = { mcpClients?: NamedMCPClient[] reconnectClient?: ReconnectClient httpClient: HttpClient + turnBudget?: TurnBudgetConsumer /** Returns the current proxy fetch. Production callers pass the getter from * `ProxyFetchProvider` (`useProxyFetchGetter()`); non-React callers (eval * scripts) build a `proxyFetch` directly and wrap it in `() => fn`. */ @@ -500,7 +419,7 @@ export const createModel = async (modelConfig: Model, getProxyFetch: () => Fetch try { return await client.fetch(input, upstreamInit) } catch (err) { - if (isKeyConfigMismatchError(err)) { + if (isTinfoilTransportWedgedError(err)) { evictSystemTinfoilClient() } throw err @@ -525,7 +444,7 @@ export const createModel = async (modelConfig: Model, getProxyFetch: () => Fetch try { return await client.fetch(input, init) } catch (err) { - if (isKeyConfigMismatchError(err)) { + if (isTinfoilTransportWedgedError(err)) { evictUserTinfoilClient() } throw err @@ -671,6 +590,13 @@ export const prepareAiRequestConfig = async ({ } } +/** + * Stream one response through the legacy built-in pipeline. + * + * Adapter callers supply the active turn consumer; direct callers receive a + * local budget. Only empty-response attempts after the first consume here + * because the routing fetch already consumed the initial request. + */ export const aiFetchStreamingResponse = async ({ init, modelId, @@ -678,11 +604,13 @@ export const aiFetchStreamingResponse = async ({ reconnectClient, httpClient, getProxyFetch, + turnBudget, }: AiFetchStreamingResponseOptions) => { const options = init as RequestInit & { body: string } const body = JSON.parse(options.body) const abortSignal: AbortSignal | undefined = options.signal ?? undefined const { messages } = body as { messages: ThunderboltUIMessage[]; id: string } + const requestBudget = turnBudget ?? createTurnBudget().consumer // The chat instance saves the user message via `saveMessages` before // invoking the adapter — see `src/chats/chat-instance.ts`. By the time we @@ -741,6 +669,9 @@ export const aiFetchStreamingResponse = async ({ */ const runStreamText = (inputMessages: Awaited>) => { return streamText({ + // SDK-internal retries are invisible to the shared per-turn request budget. + // Keep retries in the app layers where every request is counted. + maxRetries: 0, temperature: modelTemperature, model: wrappedModel, system: systemPrompt, @@ -849,6 +780,7 @@ export const aiFetchStreamingResponse = async ({ error: error.responseBody ?? error.message, status: error.statusCode, isRetryable: error.isRetryable, + kind: classifyErrorKind(error), }) } // A provider that can't serialize a part throws this client-side, before @@ -859,9 +791,16 @@ export const aiFetchStreamingResponse = async ({ // must NOT be tagged 422, or they'd masquerade as a fixable attachment. if (UnsupportedFunctionalityError.isInstance(error)) { const isFilePart = /file part|media type/i.test(`${error.functionality} ${error.message}`) - return JSON.stringify({ error: error.message, status: isFilePart ? 422 : undefined, isRetryable: false }) + return JSON.stringify({ + error: error.message, + status: isFilePart ? 422 : undefined, + isRetryable: false, + kind: classifyErrorKind(error), + }) } - return error instanceof Error ? error.message : String(error) + const message = error instanceof Error ? error.message : String(error) + const kind = classifyErrorKind(error) + return kind ? JSON.stringify({ error: message, kind }) : message } // Surface the user's persisted ask-widget responses (stored in each @@ -916,6 +855,12 @@ export const aiFetchStreamingResponse = async ({ let anyAttemptHadToolCalls = false while (attemptNumber <= maxAttempts) { + if (attemptNumber > 1 && !requestBudget.tryConsumeRequest()) { + // Mirror the routing choke point's denial so both layers surface the + // same named error instead of ending the turn as a silent finish. + throw createTurnBudgetExhaustedError() + } + const result = runStreamText(currentMessages) const messageMetadata = createMessageMetadata(modelId, sourceCollector, mcpToolsMetadata) @@ -987,7 +932,7 @@ export const aiFetchStreamingResponse = async ({ console.error('aiFetchStreamingResponse error', error) const status = (error as { status?: number }).status ?? (error as { response?: { status?: number } }).response?.status - return new Response(JSON.stringify({ error: (error as Error).message, status }), { + return new Response(JSON.stringify({ error: (error as Error).message, status, kind: classifyErrorKind(error) }), { status: status ?? 500, headers: { 'Content-Type': 'application/json' }, }) diff --git a/src/ai/prewarm-system-model.ts b/src/ai/prewarm-system-model.ts index e96c230f8..48fe92cb4 100644 --- a/src/ai/prewarm-system-model.ts +++ b/src/ai/prewarm-system-model.ts @@ -12,6 +12,6 @@ export const prewarmSystemModel = async (model: Pick { + it('consumes requests until request capacity is exhausted', () => { + const budget = createTurnBudget() + + for (let request = 0; request < maxRequestsPerTurn; request++) { + expect(budget.consumer.tryConsumeRequest()).toBe(true) + } + + expect(budget.consumer.tryConsumeRequest()).toBe(false) + expect(budget.probe.isExhausted).toBe(true) + }) + + it('exhausts when turn wall-clock limit is reached', () => { + const clock = { now: 1_000 } + const budget = createTurnBudget(() => clock.now) + + expect(budget.consumer.tryConsumeRequest()).toBe(true) + clock.now += maxTurnWallClockMs - 1 + expect(budget.probe.isExhausted).toBe(false) + + clock.now++ + expect(budget.consumer.tryConsumeRequest()).toBe(false) + expect(budget.probe.isExhausted).toBe(true) + }) + + it('separates consuming from exhaustion probing', () => { + const budget = createTurnBudget() + + expect(Object.keys(budget.consumer)).toEqual(['tryConsumeRequest']) + expect(Object.keys(budget.probe)).toEqual(['isExhausted']) + expect(budget.probe.isExhausted).toBe(false) + }) +}) diff --git a/src/ai/retry-budget.ts b/src/ai/retry-budget.ts new file mode 100644 index 000000000..137bdc05a --- /dev/null +++ b/src/ai/retry-budget.ts @@ -0,0 +1,52 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +export const maxRequestsPerTurn = 6 +export const maxTurnWallClockMs = 120_000 + +/** The one sentinel both consuming layers throw on denial — classified by name. */ +export const createTurnBudgetExhaustedError = (): Error => + Object.assign(new Error('Turn request budget exhausted'), { name: 'TurnBudgetExhaustedError' }) + +export type TurnBudgetConsumer = { + tryConsumeRequest: () => boolean +} + +export type TurnBudgetProbe = { + readonly isExhausted: boolean +} + +export type TurnBudget = { + consumer: TurnBudgetConsumer + probe: TurnBudgetProbe +} + +/** + * Create one turn budget. + * + * Invariant: every model request in a turn — first send, SDK-level retry, + * empty-response retry, or outer auto-retry — draws from this shared budget. + */ +export const createTurnBudget = (now: () => number = Date.now): TurnBudget => { + const startedAt = now() + let requestsConsumed = 0 + const isExhausted = () => requestsConsumed >= maxRequestsPerTurn || now() - startedAt >= maxTurnWallClockMs + + return { + consumer: { + tryConsumeRequest: () => { + if (isExhausted()) { + return false + } + requestsConsumed++ + return true + }, + }, + probe: { + get isExhausted() { + return isExhausted() + }, + }, + } +} diff --git a/src/ai/tinfoil-client.test.ts b/src/ai/tinfoil-client.test.ts new file mode 100644 index 000000000..53d970c4b --- /dev/null +++ b/src/ai/tinfoil-client.test.ts @@ -0,0 +1,242 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +import { describe, expect, it, mock } from 'bun:test' +import { getClock } from '@/testing-library' +import type { SecureClient } from 'tinfoil' +import { createTinfoilClientLifecycle, isTinfoilTransportWedgedError } from './tinfoil-client' + +/** Build the `SecureClient` slice consumed by the lifecycle. */ +const createClient = (ready: () => Promise = async () => {}): SecureClient => + ({ + ready, + fetch: async () => new Response(), + getBaseURL: () => 'https://enclave.example.com/v1/', + }) as unknown as SecureClient + +describe('Tinfoil client lifecycle', () => { + it('wraps malformed attestation responses, preserves telemetry fidelity, and evicts the client', async () => { + const readyError = new SyntaxError('Unexpected token in attestation document') + const failedClient = createClient(async () => { + throw readyError + }) + const healthyClient = createClient() + const createClientFactory = mock(async () => healthyClient) + createClientFactory.mockImplementationOnce(async () => failedClient) + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1', + trackAttestation, + }) + + await expect(lifecycle.getSystemTinfoilClient()).rejects.toMatchObject({ + name: 'TinfoilAttestationError', + message: 'Tinfoil attestation failed: SyntaxError: Unexpected token in attestation document', + cause: readyError, + }) + await expect(lifecycle.getSystemTinfoilClient()).resolves.toBe(healthyClient) + + expect(createClientFactory).toHaveBeenCalledTimes(2) + expect(trackAttestation).toHaveBeenCalledWith({ + outcome: 'error', + duration_ms: 0, + error_name: 'SyntaxError', + client: 'system', + }) + }) + + it('times out hung ready, evicts the client, and reports timeout telemetry', async () => { + const hungClient = createClient(() => new Promise(() => {})) + const healthyClient = createClient() + const createClientFactory = mock(async () => healthyClient) + createClientFactory.mockImplementationOnce(async () => hungClient) + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1', + timeoutMs: 1_000, + trackAttestation, + }) + + const clientPromise = lifecycle.getSystemTinfoilClient() + const resultPromise = clientPromise.then( + (client) => ({ client }), + (error: unknown) => ({ error }), + ) + await getClock().tickAsync(1_000) + await expect(resultPromise).resolves.toMatchObject({ + error: { + name: 'TinfoilAttestationTimeoutError', + message: 'Tinfoil attestation timed out after 1000ms', + }, + }) + + await expect(lifecycle.getSystemTinfoilClient()).resolves.toBe(healthyClient) + expect(createClientFactory).toHaveBeenCalledTimes(2) + expect(trackAttestation).toHaveBeenCalledWith({ + outcome: 'timeout', + duration_ms: 1_000, + error_name: 'TinfoilAttestationTimeoutError', + client: 'system', + }) + }) + + it('evicts a rejected construction promise', async () => { + const constructionError = new Error('module load failed') + const healthyClient = createClient() + const createClientFactory = mock(async () => healthyClient) + createClientFactory.mockImplementationOnce(async () => { + throw constructionError + }) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1', + }) + + await expect(lifecycle.getSystemTinfoilClient()).rejects.toBe(constructionError) + await expect(lifecycle.getSystemTinfoilClient()).resolves.toBe(healthyClient) + expect(createClientFactory).toHaveBeenCalledTimes(2) + }) + + it('shares one system-client construction across concurrent callers', async () => { + const client = createClient() + const createClientFactory = mock(async () => client) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1/', + }) + + const [first, second] = await Promise.all([lifecycle.getSystemTinfoilClient(), lifecycle.getSystemTinfoilClient()]) + + expect(first).toBe(client) + expect(second).toBe(client) + expect(createClientFactory).toHaveBeenCalledTimes(1) + expect(createClientFactory).toHaveBeenCalledWith('system', 'https://cloud.example.com/v1') + }) + + it('tracks successful system attestation', async () => { + const client = createClient() + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: async () => client, + getCloudUrl: () => 'https://cloud.example.com/v1', + trackAttestation, + }) + + await lifecycle.getSystemTinfoilClient() + + expect(trackAttestation).toHaveBeenCalledWith({ + outcome: 'success', + duration_ms: 0, + client: 'system', + }) + }) + + it('reports success once per client instance while awaiting ready on every call', async () => { + const ready = mock(async () => {}) + const client = createClient(ready) + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: async () => client, + getCloudUrl: () => 'https://cloud.example.com/v1', + trackAttestation, + }) + + await lifecycle.getSystemTinfoilClient() + await lifecycle.getSystemTinfoilClient() + await lifecycle.getSystemTinfoilClient() + + expect(ready).toHaveBeenCalledTimes(3) + expect(trackAttestation).toHaveBeenCalledTimes(1) + }) + + it('times out and evicts when a later ready call hangs during re-attestation', async () => { + const ready = mock(async (): Promise => new Promise(() => {})) + ready.mockImplementationOnce(async () => {}) + const reattestingClient = createClient(ready) + const healthyClient = createClient() + const createClientFactory = mock(async () => healthyClient) + createClientFactory.mockImplementationOnce(async () => reattestingClient) + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1', + timeoutMs: 1_000, + trackAttestation, + }) + + await expect(lifecycle.getSystemTinfoilClient()).resolves.toBe(reattestingClient) + const resultPromise = lifecycle.getSystemTinfoilClient().then( + (client) => ({ client }), + (error: unknown) => ({ error }), + ) + await getClock().tickAsync(1_000) + + await expect(resultPromise).resolves.toMatchObject({ + error: { name: 'TinfoilAttestationTimeoutError' }, + }) + await expect(lifecycle.getSystemTinfoilClient()).resolves.toBe(healthyClient) + expect(createClientFactory).toHaveBeenCalledTimes(2) + expect(trackAttestation).toHaveBeenCalledWith({ + outcome: 'timeout', + duration_ms: 1_000, + error_name: 'TinfoilAttestationTimeoutError', + client: 'system', + }) + }) + + it('wraps unreachable attestation endpoints, tracks the underlying error, and evicts the client', async () => { + const readyError = new TypeError('Failed to fetch') + const failedClient = createClient(async () => { + throw readyError + }) + const healthyClient = createClient() + const createClientFactory = mock(async () => healthyClient) + createClientFactory.mockImplementationOnce(async () => failedClient) + const trackAttestation = mock(() => {}) + const lifecycle = createTinfoilClientLifecycle({ + createClient: createClientFactory, + getCloudUrl: () => 'https://cloud.example.com/v1', + trackAttestation, + }) + + await expect(lifecycle.getTinfoilClient()).rejects.toMatchObject({ + name: 'TinfoilAttestationError', + message: 'Tinfoil attestation failed: TypeError: Failed to fetch', + cause: readyError, + }) + await expect(lifecycle.getTinfoilClient()).resolves.toBe(healthyClient) + + expect(createClientFactory).toHaveBeenCalledTimes(2) + expect(createClientFactory).toHaveBeenCalledWith('user', '') + expect(trackAttestation).toHaveBeenCalledWith({ + outcome: 'error', + duration_ms: 0, + error_name: 'TypeError', + client: 'user', + }) + }) +}) + +describe('isTinfoilTransportWedgedError', () => { + it('matches SDK key-config mismatch errors by name', () => { + const error = { name: 'KeyConfigMismatchError', message: 'key changed' } + + expect(isTinfoilTransportWedgedError(error)).toBe(true) + }) + + it('matches the SDK concurrent-reset null-transport TypeError', () => { + const error = { name: 'TypeError', message: "Cannot read properties of null (reading 'fetch')" } + + expect(isTinfoilTransportWedgedError(error)).toBe(true) + }) + + it('does not classify unrelated errors for eviction', () => { + expect(isTinfoilTransportWedgedError(new Error('network unavailable'))).toBe(false) + expect(isTinfoilTransportWedgedError(new TypeError("Cannot read properties of null (reading 'baseURL')"))).toBe( + false, + ) + }) +}) diff --git a/src/ai/tinfoil-client.ts b/src/ai/tinfoil-client.ts new file mode 100644 index 000000000..f186d0985 --- /dev/null +++ b/src/ai/tinfoil-client.ts @@ -0,0 +1,203 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +import { getUserCacheSecret } from '@/lib/auth-token' +import { getErrorName } from '@/lib/error-utils' +import { trackEvent } from '@/lib/posthog' +import { withDeadline } from '@/lib/timeout' +import { getLocalSetting } from '@/stores/local-settings-store' +import type { Model } from '@/types' +import type { SecureClient } from 'tinfoil' + +export const tinfoilAttestationTimeoutMs = 15_000 + +type TinfoilClientType = 'system' | 'user' +type TinfoilAttestationProperties = { + outcome: 'success' | 'error' | 'timeout' + duration_ms: number + error_name?: string + client: TinfoilClientType +} + +type TinfoilClientLifecycleOptions = { + createClient: (type: TinfoilClientType, cloudUrl: string) => Promise + getCloudUrl: () => string + timeoutMs?: number + trackAttestation?: (properties: TinfoilAttestationProperties) => void +} + +/** Create the timeout error surfaced when Tinfoil attestation does not settle. */ +const createAttestationTimeoutError = (timeoutMs: number): Error => + Object.assign(new Error(`Tinfoil attestation timed out after ${timeoutMs}ms`), { + name: 'TinfoilAttestationTimeoutError', + }) + +/** + * Stabilize attestation classification across the SDK's inconsistent failure + * names; live malformed and unreachable ATC responses use SyntaxError/TypeError. + */ +const createAttestationError = (cause: unknown): Error => { + const causeName = getErrorName(cause) ?? 'Error' + const causeMessage = cause instanceof Error ? cause.message : String(cause) + return Object.assign(new Error(`Tinfoil attestation failed: ${causeName}: ${causeMessage}`, { cause }), { + name: 'TinfoilAttestationError', + }) +} + +/** Bound the SDK's non-abortable attestation wait and emit structured telemetry. */ +const waitForAttestation = async ( + client: SecureClient, + clientType: TinfoilClientType, + timeoutMs: number, + trackAttestation: (properties: TinfoilAttestationProperties) => void, +): Promise => { + const startedAt = Date.now() + + try { + // Escaped FetchError resets SDK state, so next ready() re-attests against + // atc.tinfoil.sh; bound every call, including previously attested clients. + await withDeadline(client.ready(), timeoutMs, () => createAttestationTimeoutError(timeoutMs)) + return { + outcome: 'success', + duration_ms: Date.now() - startedAt, + client: clientType, + } + } catch (error) { + const isTimeout = getErrorName(error) === 'TinfoilAttestationTimeoutError' + trackAttestation({ + outcome: isTimeout ? 'timeout' : 'error', + duration_ms: Date.now() - startedAt, + error_name: getErrorName(error), + client: clientType, + }) + // Wrap only at the throw so telemetry reports the underlying SDK failure + // name while classification keys on our stable wrapper name. + throw isTimeout ? error : createAttestationError(error) + } +} + +/** + * Create an isolated Tinfoil lifecycle. Production uses one instance; tests + * inject client factories, timeout, and telemetry without shared-module mocks. + */ +export const createTinfoilClientLifecycle = ({ + createClient, + getCloudUrl, + timeoutMs = tinfoilAttestationTimeoutMs, + trackAttestation = (properties) => trackEvent('tinfoil_attestation', properties), +}: TinfoilClientLifecycleOptions) => { + const clients = new Map>() + const reportedClients = new WeakSet() + + /** Resolve the cache key for a client type. */ + const getClientKey = (type: TinfoilClientType): string => + type === 'user' ? 'user' : `system:${getCloudUrl().replace(/\/$/, '')}` + + /** Cache construction before it settles so concurrent callers share one client. */ + const construct = (key: string, type: TinfoilClientType): Promise => { + const cloudUrl = type === 'system' ? key.slice('system:'.length) : '' + const clientPromise = createClient(type, cloudUrl) + clients.set(key, clientPromise) + return clientPromise + } + + /** Return an attested client, evicting failed construction or attestation. */ + const acquire = async (type: TinfoilClientType): Promise => { + const key = getClientKey(type) + const clientPromise = clients.get(key) ?? construct(key, type) + try { + const client = await clientPromise + const properties = await waitForAttestation(client, type, timeoutMs, trackAttestation) + if (!reportedClients.has(client)) { + reportedClients.add(client) + trackAttestation(properties) + } + return client + } catch (error) { + // Delete only the failed generation; a replacement may already be cached. + if (clients.get(key) === clientPromise) { + clients.delete(key) + } + throw error + } + } + + /** Return the attested system client for the current cloud URL. */ + const getSystemTinfoilClient = (): Promise => acquire('system') + /** Return the attested user client. */ + const getTinfoilClient = (): Promise => acquire('user') + /** Drop the current cloud URL's system client. */ + const evictSystemTinfoilClient = (): void => { + clients.delete(getClientKey('system')) + } + /** Drop the user client. */ + const evictUserTinfoilClient = (): void => { + clients.delete(getClientKey('user')) + } + + return { + getSystemTinfoilClient, + evictSystemTinfoilClient, + getTinfoilClient, + evictUserTinfoilClient, + } +} + +// Cached so callers reuse one client and its SDK-managed attestation state. +// `tinfoil` is dynamically imported to code-split its attestation/crypto deps. +// +// system: HPKE body POSTs to /tinfoil; backend injects our key. +// user: BYOK — direct to the enclave with the user's own key. +// +// System cache is keyed by cloudUrl so a dev-tools URL switch hits the new +// backend on the next call. +const lifecycle = createTinfoilClientLifecycle({ + // Cache construction synchronously before this dynamic import resolves so + // prewarm and immediate first send share one client and attestation. + createClient: async (type, cloudUrl) => { + const { SecureClient } = await import('tinfoil') + if (type === 'user') { + return new SecureClient({ userCacheSecret: getUserCacheSecret() }) + } + return new SecureClient({ baseURL: `${cloudUrl}/tinfoil`, userCacheSecret: getUserCacheSecret() }) + }, + getCloudUrl: () => getLocalSetting('cloudUrl'), +}) + +export const { getSystemTinfoilClient, evictSystemTinfoilClient, getTinfoilClient, evictUserTinfoilClient } = lifecycle + +/** + * Best-effort warm-up of the Tinfoil system enclave so first chat send avoids + * attestation on its critical path. Errors surface on real sends. + */ +export const runSystemModelPrewarm = async (model: Pick | null | undefined) => { + if (!model || model.provider !== 'tinfoil' || !model.isSystem) { + return + } + try { + await getSystemTinfoilClient() + } catch (error) { + console.warn('runSystemModelPrewarm: warm-up skipped', error) + } +} + +/** + * Match transport failures that require replacing the cached SecureClient. + * Message matching is brittle but unavoidable: SDK concurrent reset race + * exposes only a generic TypeError, with no exported class or discriminator. + */ +export const isTinfoilTransportWedgedError = (error: unknown): boolean => { + const errorName = getErrorName(error) + if (errorName === 'KeyConfigMismatchError') { + return true + } + return ( + errorName === 'TypeError' && + typeof error === 'object' && + error !== null && + 'message' in error && + typeof error.message === 'string' && + error.message.includes("reading 'fetch'") + ) +} diff --git a/src/chats/chat-instance.test.ts b/src/chats/chat-instance.test.ts index 473b95652..4f1c8feae 100644 --- a/src/chats/chat-instance.test.ts +++ b/src/chats/chat-instance.test.ts @@ -8,16 +8,19 @@ * each call to the injected `connectToAgent`. */ -import '@/testing-library' - +import { createTurnBudget, maxRequestsPerTurn, type TurnBudget } from '@/ai/retry-budget' import { builtInAgent } from '@/defaults/agents' import type { HttpClient } from '@/lib/http' import type { FetchFn } from '@/lib/proxy-fetch' import { createMockChatInstance, hydrateStore, resetStore } from '@/test-utils/chat-store-mocks' +import { getClock } from '@/testing-library' +import type { ThunderboltUIMessage } from '@/types' import type { Agent, AgentAdapter } from '@/types/acp' -import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test' +import type { Chat } from '@ai-sdk/react' +import type { ChatInit, ChatOnFinishCallback } from 'ai' +import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test' import { useChatStore } from './chat-store' -import { createAgentRoutingFetch } from './chat-instance' +import { createAgentRoutingFetch, createChatInstance } from './chat-instance' const sessionId = 'sess-1' const httpClient: HttpClient = {} as HttpClient @@ -41,6 +44,76 @@ const hydrate = () => { }) } +/** Build a chat instance whose retry callbacks and original methods are observable. */ +const createRetryHarness = () => { + const regenerate = mock(async () => {}) + const sendMessage = mock(async () => {}) + const budgets: TurnBudget[] = [] + let onFinish: ChatOnFinishCallback | undefined + + const createTrackedTurnBudget = () => { + const budget = createTurnBudget() + budgets.push(budget) + return budget + } + + const createChat = (init: ChatInit) => { + onFinish = init.onFinish + return { + id: init.id ?? sessionId, + messages: init.messages ?? [], + regenerate, + sendMessage, + } as unknown as Chat + } + + const instance = createChatInstance(sessionId, [], async () => {}, httpClient, getProxyFetch, { + createChat, + createTurnBudget: createTrackedTurnBudget, + }) + hydrateStore({ + chatInstance: instance, + chatThread: null, + id: sessionId, + selectedModel: { id: 'm1', isConfidential: 0 } as never, + triggerData: null, + }) + + const finishWithError = () => + onFinish!({ + message: { id: 'failed-assistant', role: 'assistant', parts: [] }, + messages: [], + isAbort: false, + isDisconnect: false, + isError: true, + }) + + const finishSuccessfully = () => + onFinish!({ + message: { + id: 'successful-assistant', + role: 'assistant', + parts: [{ type: 'text', text: 'Done' }], + }, + messages: [], + isAbort: false, + isDisconnect: false, + isError: false, + }) + + const getTurnBudget = () => budgets.at(-1)! + + return { finishSuccessfully, finishWithError, getTurnBudget, instance, regenerate, sendMessage } +} + +/** Fully consume one injected chat-instance turn budget. */ +const exhaustTurnBudget = (budget: TurnBudget) => { + for (let request = 0; request < maxRequestsPerTurn; request++) { + budget.consumer.tryConsumeRequest() + } + return budget +} + describe('createAgentRoutingFetch — connection status', () => { beforeEach(() => { resetStore() @@ -64,7 +137,7 @@ describe('createAgentRoutingFetch — connection status', () => { getDb: (() => ({})) as never, }) - await fetch('http://x', { body: '{}' } as RequestInit) + await fetch('https://x', { body: '{}' } as RequestInit) expect(observed).toEqual(['connecting']) expect(useChatStore.getState().sessions.get(sessionId)!.connectionStatus).toBe('ready') @@ -82,7 +155,7 @@ describe('createAgentRoutingFetch — connection status', () => { getDb: (() => ({})) as never, }) - await expect(fetch('http://x', { body: '{}' } as RequestInit)).rejects.toThrow('boom') + await expect(fetch('https://x', { body: '{}' } as RequestInit)).rejects.toThrow('boom') const session = useChatStore.getState().sessions.get(sessionId)! expect(session.connectionStatus).toBe('error') @@ -98,8 +171,8 @@ describe('createAgentRoutingFetch — connection status', () => { getDb: (() => ({})) as never, }) - await fetch('http://x', { body: '{}' } as RequestInit) - await fetch('http://x', { body: '{}' } as RequestInit) + await fetch('https://x', { body: '{}' } as RequestInit) + await fetch('https://x', { body: '{}' } as RequestInit) expect(connectToAgent).toHaveBeenCalledTimes(1) expect(useChatStore.getState().sessions.get(sessionId)!.connectionStatus).toBe('ready') @@ -115,11 +188,134 @@ describe('createAgentRoutingFetch — connection status', () => { getDb: (() => ({})) as never, }) - await fetch('http://x', { body: '{}' } as RequestInit) + await fetch('https://x', { body: '{}' } as RequestInit) useChatStore.getState().updateSession(sessionId, { selectedAgent: altAgent }) - await fetch('http://x', { body: '{}' } as RequestInit) + await fetch('https://x', { body: '{}' } as RequestInit) expect(connectToAgent).toHaveBeenCalledTimes(2) expect(useChatStore.getState().sessions.get(sessionId)!.connectionStatus).toBe('ready') }) + + it('throws TurnBudgetExhaustedError before invoking the adapter when the consumer is drained', async () => { + const budget = exhaustTurnBudget(createTurnBudget()) + const adapterFetch = mock(async () => new Response('ok')) + const connectToAgent = mock(async (agent: Agent) => ({ + ...makeAdapter(agent), + fetch: adapterFetch, + })) + const fetch = createAgentRoutingFetch( + sessionId, + async () => {}, + httpClient, + getProxyFetch, + { + connectToAgent: connectToAgent as never, + updateChatThread: (async () => {}) as never, + getDb: (() => ({})) as never, + }, + { getTurnBudget: () => budget }, + ) + + await expect(fetch('https://x', { body: '{}' } as RequestInit)).rejects.toMatchObject({ + name: 'TurnBudgetExhaustedError', + }) + expect(adapterFetch).not.toHaveBeenCalled() + }) +}) + +describe('createChatInstance — retry policy', () => { + beforeEach(() => { + resetStore() + }) + + afterEach(() => { + resetStore() + }) + + it('uses 2s, 4s, and 8s exponential retry delays', async () => { + const random = spyOn(Math, 'random').mockReturnValue(0.5) + const { finishWithError, regenerate } = createRetryHarness() + + try { + for (const [completedRetries, delay] of [2_000, 4_000, 8_000].entries()) { + await finishWithError() + + await getClock().tickAsync(delay - 1) + expect(regenerate).toHaveBeenCalledTimes(completedRetries) + + await getClock().tickAsync(1) + expect(regenerate).toHaveBeenCalledTimes(completedRetries + 1) + } + } finally { + random.mockRestore() + } + }) + + it('marks retries exhausted without scheduling when turn budget is exhausted', async () => { + const { finishWithError, getTurnBudget, regenerate } = createRetryHarness() + exhaustTurnBudget(getTurnBudget()) + + await finishWithError() + await getClock().runAllAsync() + + const session = useChatStore.getState().sessions.get(sessionId)! + expect(session.retryCount).toBe(0) + expect(session.retriesExhausted).toBe(true) + expect(regenerate).not.toHaveBeenCalled() + }) + + it('manual regenerate resets the turn budget', async () => { + const { getTurnBudget, instance, regenerate } = createRetryHarness() + const exhaustedBudget = exhaustTurnBudget(getTurnBudget()) + + await instance.regenerate() + + expect(getTurnBudget()).not.toBe(exhaustedBudget) + expect(getTurnBudget().probe.isExhausted).toBe(false) + expect(getTurnBudget().consumer.tryConsumeRequest()).toBe(true) + expect(regenerate).toHaveBeenCalledTimes(1) + }) + + it('auto-retry does not reset the turn budget', async () => { + const random = spyOn(Math, 'random').mockReturnValue(0.5) + const { finishWithError, getTurnBudget, regenerate } = createRetryHarness() + const budget = getTurnBudget() + budget.consumer.tryConsumeRequest() + + try { + await finishWithError() + await getClock().tickAsync(2_000) + + expect(regenerate).toHaveBeenCalledTimes(1) + expect(getTurnBudget()).toBe(budget) + for (let request = 1; request < maxRequestsPerTurn; request++) { + expect(budget.consumer.tryConsumeRequest()).toBe(true) + } + expect(budget.consumer.tryConsumeRequest()).toBe(false) + } finally { + random.mockRestore() + } + }) + + it('resets the turn budget after success so the next automatic turn starts fresh', async () => { + const { finishSuccessfully, getTurnBudget } = createRetryHarness() + const exhaustedBudget = exhaustTurnBudget(getTurnBudget()) + + await finishSuccessfully() + + expect(getTurnBudget()).not.toBe(exhaustedBudget) + expect(getTurnBudget().probe.isExhausted).toBe(false) + expect(getTurnBudget().consumer.tryConsumeRequest()).toBe(true) + }) + + it('user send resets the turn budget', async () => { + const { getTurnBudget, instance, sendMessage } = createRetryHarness() + const exhaustedBudget = exhaustTurnBudget(getTurnBudget()) + + await instance.sendMessage({ text: 'new turn' }) + + expect(getTurnBudget()).not.toBe(exhaustedBudget) + expect(getTurnBudget().probe.isExhausted).toBe(false) + expect(sendMessage).toHaveBeenCalledTimes(1) + }) }) diff --git a/src/chats/chat-instance.ts b/src/chats/chat-instance.ts index cb6044c68..612cd55e4 100644 --- a/src/chats/chat-instance.ts +++ b/src/chats/chat-instance.ts @@ -6,29 +6,41 @@ import type { connectToAgent as defaultConnectToAgent } from '@/acp' import { getOrConnectAdapter as defaultGetOrConnectAdapter } from '@/acp/adapter-cache' import type { AcpCommand, SessionSideEffect } from '@/acp/translators/acp-to-ai-sdk' import { useAgentCommandsStore } from '@/acp/agent-commands-store' +import { + createTurnBudget as defaultCreateTurnBudget, + createTurnBudgetExhaustedError, + type TurnBudget, +} from '@/ai/retry-budget' import { updateChatThread as defaultUpdateChatThread } from '@/dal/chat-threads' import { getAllSkills as defaultGetAllSkills } from '@/dal' import { isBuiltInAgent } from '@/defaults/agents' import { extractLastUserText, resolveSkillTokenInstructions } from '@/skills/resolve-skill-system-messages' import { getDb as defaultGetDb } from '@/db/database' -import { getErrorRetryable, isContentRejectionError, isContextOverflowError, isRateLimitError } from '@/lib/error-utils' +import { + getChatErrorKind, + getErrorRetryable, + isContentRejectionError, + isContextOverflowError, + isRateLimitError, +} from '@/lib/error-utils' import type { HttpClient } from '@/lib/http' import { trackEvent } from '@/lib/posthog' import type { FetchFn } from '@/lib/proxy-fetch' import type { SaveMessagesFunction, ThunderboltUIMessage } from '@/types' import { Chat } from '@ai-sdk/react' import type { RequestPermissionRequest, RequestPermissionResponse } from '@agentclientprotocol/sdk' -import { DefaultChatTransport } from 'ai' +import { DefaultChatTransport, type ChatInit } from 'ai' import { v7 as uuidv7 } from 'uuid' import { deriveToolKey, findAllowOption, useChatStore } from './chat-store' export const maxRetries = 3 +const baseRetryDelayMs = 2000 /** * Calculate retry delay with exponential backoff and jitter. * Jitter prevents synchronized retries from overwhelming servers. */ -const getRetryDelay = (attempt: number) => 2000 * attempt * (0.5 + Math.random()) +const getRetryDelay = (attempt: number) => baseRetryDelayMs * 2 ** (attempt - 1) * (0.5 + Math.random()) /** Bridge an ACP `requestPermission` call to the chat-store dialog flow. * Auto-approves remembered allowances; otherwise stashes the request until @@ -89,10 +101,13 @@ export type CreateChatInstanceDeps = { updateChatThread?: typeof defaultUpdateChatThread getDb?: typeof defaultGetDb getAllSkills?: typeof defaultGetAllSkills + createChat?: (init: ChatInit) => Chat + createTurnBudget?: typeof defaultCreateTurnBudget } export type AgentRoutingState = { regenerationRevision?: number + getTurnBudget?: () => TurnBudget } /** @@ -127,6 +142,12 @@ export const createAgentRoutingFetch = ( const updateChatThread = deps.updateChatThread ?? defaultUpdateChatThread const getDb = deps.getDb ?? defaultGetDb const getAllSkills = deps.getAllSkills ?? defaultGetAllSkills + const getTurnBudget = + routingState.getTurnBudget ?? + (() => { + const turnBudget = (deps.createTurnBudget ?? defaultCreateTurnBudget)() + return () => turnBudget + })() let routedAgentId: string | null = null @@ -226,6 +247,11 @@ export const createAgentRoutingFetch = ( ? undefined : (request: RequestPermissionRequest) => requestPermissionViaStore(id, selectedAgent.id, request) + const turnBudget = getTurnBudget().consumer + if (!turnBudget.tryConsumeRequest()) { + throw createTurnBudgetExhaustedError() + } + return adapter.fetch(init, { threadId: id, chatThread, @@ -236,6 +262,7 @@ export const createAgentRoutingFetch = ( reconnectClient, httpClient, getProxyFetch, + turnBudget, regenerationRevision: routingState.regenerationRevision ?? 0, skillInstructions, onAcpSessionId: persistAcpSessionId, @@ -249,6 +276,11 @@ export const createAgentRoutingFetch = ( ) } +/** + * Create one chat instance with retry state and request budget scoped to its + * closure. New, successful, and aborted turns replace the budget while + * automatic retries preserve it. + */ export const createChatInstance = ( id: string, messages: ThunderboltUIMessage[], @@ -257,14 +289,41 @@ export const createChatInstance = ( getProxyFetch: () => FetchFn, deps: CreateChatInstanceDeps = {}, ) => { - const routingState: AgentRoutingState = { regenerationRevision: 0 } + const createTurnBudget = deps.createTurnBudget ?? defaultCreateTurnBudget + let turnBudget = createTurnBudget() + const routingState: AgentRoutingState = { + regenerationRevision: 0, + getTurnBudget: () => turnBudget, + } const customFetch = createAgentRoutingFetch(id, saveMessages, httpClient, getProxyFetch, deps, routingState) + const createChat = deps.createChat ?? ((init: ChatInit) => new Chat(init)) let retryCount = 0 let retryTimeout: ReturnType | null = null let lastError: Error | null = null - const instance = new Chat({ + /** Clear retry state and replace the completed turn's request budget. */ + const resetRetryStateForNewTurn = () => { + if (retryTimeout) { + clearTimeout(retryTimeout) + retryTimeout = null + } + turnBudget = createTurnBudget() + retryCount = 0 + lastError = null + useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + } + + /** Stop retrying this turn and record why it stopped. */ + const markRetriesExhausted = () => { + trackEvent('chat_retries_exhausted', { + reason: getChatErrorKind(lastError) ?? 'unknown', + attempts: retryCount, + }) + useChatStore.getState().updateSession(id, { retriesExhausted: true }) + } + + const instance = createChat({ id, messages, transport: new DefaultChatTransport({ fetch: customFetch }), @@ -273,14 +332,7 @@ export const createChatInstance = ( sendAutomaticallyWhen: ({ messages }) => messages.length > 0 && messages[messages.length - 1].role === 'user', onFinish: async ({ message, isError, isAbort }) => { if (isAbort) { - // Clear any pending retry timer and reset retry state when aborted - if (retryTimeout) { - clearTimeout(retryTimeout) - retryTimeout = null - } - retryCount = 0 - lastError = null - useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + resetRetryStateForNewTurn() // Persist whatever streamed before the user hit Stop. Streaming partial // saves are throttled and their pending trailing write is cancelled the @@ -296,9 +348,10 @@ export const createChatInstance = ( // Handle successful responses: message exists, no error, and has parts if (!isError && message && message.parts?.length) { - retryCount = 0 - lastError = null - useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + if (retryCount > 0) { + trackEvent('chat_retry_success', { attempts: retryCount }) + } + resetRetryStateForNewTurn() const { sessions } = useChatStore.getState() @@ -321,8 +374,8 @@ export const createChatInstance = ( // Don't auto-retry rate limit errors — retrying immediately makes it worse if (isRateLimitError(lastError)) { + markRetriesExhausted() lastError = null - useChatStore.getState().updateSession(id, { retriesExhausted: true }) return } @@ -343,16 +396,25 @@ export const createChatInstance = ( isContentRejectionError(lastError) || getErrorRetryable(lastError) === false ) { - useChatStore.getState().updateSession(id, { retriesExhausted: true }) + markRetriesExhausted() return } if (retryCount < maxRetries) { + if (turnBudget.probe.isExhausted) { + markRetriesExhausted() + return + } + retryCount++ useChatStore.getState().updateSession(id, { retryCount }) console.info(`Auto-retrying (${retryCount}/${maxRetries})...`) - trackEvent('chat_auto_retry', { attempt: retryCount, max_retries: maxRetries }) + trackEvent('chat_auto_retry', { + attempt: retryCount, + max_retries: maxRetries, + reason: getChatErrorKind(lastError) ?? 'unknown', + }) retryTimeout = setTimeout(() => { retryTimeout = null @@ -362,8 +424,7 @@ export const createChatInstance = ( if (!sessions.has(id) || currentSessionId !== id) { // Reset retry state when bailing out due to session switch, so the UI // doesn't show "Retrying..." when the user switches back to this session. - retryCount = 0 - useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + resetRetryStateForNewTurn() return } regenerateResponse().catch((err) => { @@ -375,7 +436,7 @@ export const createChatInstance = ( }) }, getRetryDelay(retryCount)) } else { - useChatStore.getState().updateSession(id, { retriesExhausted: true }) + markRetriesExhausted() } }, // Retry logic lives in onFinish (the SDK's finally block), not here. @@ -400,13 +461,7 @@ export const createChatInstance = ( // Reset retry count on manual regenerate (Retry button) so auto-retries work again instance.regenerate = async function () { - if (retryTimeout) { - clearTimeout(retryTimeout) - retryTimeout = null - } - retryCount = 0 - lastError = null - useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + resetRetryStateForNewTurn() return regenerateResponse() } @@ -415,13 +470,7 @@ export const createChatInstance = ( // Override the sendMessage method to check if the model is available for the chat thread instance.sendMessage = async function (message, options) { // Cancel any pending auto-retry and reset error state for the new message - if (retryTimeout) { - clearTimeout(retryTimeout) - retryTimeout = null - } - retryCount = 0 - lastError = null - useChatStore.getState().updateSession(id, { retryCount: 0, retriesExhausted: false }) + resetRetryStateForNewTurn() const { sessions } = useChatStore.getState() diff --git a/src/components/chat/error-message.test.tsx b/src/components/chat/error-message.test.tsx index b03ab36b6..20dc5b39a 100644 --- a/src/components/chat/error-message.test.tsx +++ b/src/components/chat/error-message.test.tsx @@ -10,6 +10,39 @@ import { ErrorMessage } from './error-message' afterEach(cleanup) describe('ErrorMessage', () => { + describe('cause-specific errors', () => { + const cases = [ + { + kind: 'attestation', + message: "Couldn't verify the secure AI connection. This is usually temporary — try again in a moment.", + }, + { + kind: 'timeout', + message: 'The AI provider took too long to respond. Try again.', + }, + { + kind: 'provider', + message: 'The AI provider is having trouble right now. Try again in a moment.', + }, + { + kind: 'network', + message: 'Connection problem. Check your internet and try again.', + }, + ] as const + + for (const { kind, message } of cases) { + it(`shows ${kind} guidance with Retry`, () => { + const onRetry = mock(() => {}) + const error = new Error(JSON.stringify({ error: 'Request failed', kind })) + + render() + + expect(screen.getByText(message)).toBeTruthy() + expect(screen.getByText('Retry')).toBeTruthy() + }) + } + }) + describe('rate limit errors', () => { it('should show rate limit message for JSON 429 status', () => { const error = new Error(JSON.stringify({ error: 'Rate limited', status: 429 })) diff --git a/src/components/chat/error-message.tsx b/src/components/chat/error-message.tsx index 95df9248c..f4addf6bf 100644 --- a/src/components/chat/error-message.tsx +++ b/src/components/chat/error-message.tsx @@ -3,10 +3,24 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ import { maxRetries } from '@/chats/chat-instance' -import { isContextOverflowError, isRateLimitError } from '@/lib/error-utils' +import { type ChatErrorKind, getChatErrorKind, isContextOverflowError, isRateLimitError } from '@/lib/error-utils' import { Loader2 } from 'lucide-react' import { memo } from 'react' +const defaultChatErrorMessage = 'Something went wrong. Please try again.' +const causeSpecificErrorMessages: Partial> = { + attestation: "Couldn't verify the secure AI connection. This is usually temporary — try again in a moment.", + timeout: 'The AI provider took too long to respond. Try again.', + provider: 'The AI provider is having trouble right now. Try again in a moment.', + network: 'Connection problem. Check your internet and try again.', +} + +/** Resolve final chat error copy after automatic retries stop. */ +const getFinalChatErrorMessage = (error?: Error | null): string => { + const kind = getChatErrorKind(error) + return (kind && causeSpecificErrorMessages[kind]) || defaultChatErrorMessage +} + type ErrorMessageProps = { retryCount: number retriesExhausted: boolean @@ -67,7 +81,7 @@ export const ErrorMessage = memo(

{deliveryExhausted ? "This model couldn't read the attached file. Try a different model." - : 'Something went wrong. Please try again.'} + : getFinalChatErrorMessage(error)}

{/* No Retry when delivery is exhausted — re-running identical input fails diff --git a/src/lib/error-utils.test.ts b/src/lib/error-utils.test.ts index 62c021ecd..88d12cab5 100644 --- a/src/lib/error-utils.test.ts +++ b/src/lib/error-utils.test.ts @@ -4,7 +4,11 @@ import { describe, expect, it } from 'bun:test' import { + type ChatErrorKind, + classifyErrorKind, createHandleError, + getChatErrorKind, + getErrorName, getErrorRetryable, getErrorStatusCode, isContentRejectionError, @@ -13,6 +17,132 @@ import { } from './error-utils' import type { HandleErrorCode } from '@/types/handle-errors' +describe('classifyErrorKind', () => { + const cases: { label: string; error: unknown; expected: ChatErrorKind }[] = [ + { + label: 'wrapped Tinfoil attestation fetch failures', + error: Object.assign(new Error('Tinfoil attestation failed: TypeError: Failed to fetch'), { + name: 'TinfoilAttestationError', + cause: new TypeError('Failed to fetch'), + }), + expected: 'attestation', + }, + { + label: 'Tinfoil attestation deadline errors', + error: Object.assign(new Error('Attestation timed out'), { name: 'TinfoilAttestationTimeoutError' }), + expected: 'timeout', + }, + { + label: 'Tinfoil key configuration transport errors', + error: Object.assign(new Error('Key configuration mismatch'), { name: 'KeyConfigMismatchError' }), + expected: 'provider', + }, + { + label: 'Tinfoil protocol transport errors', + error: Object.assign(new Error('Protocol failed'), { name: 'ProtocolError' }), + expected: 'provider', + }, + { + label: 'Tinfoil decryption transport errors', + error: Object.assign(new Error('Decryption failed'), { name: 'DecryptionError' }), + expected: 'provider', + }, + { + label: 'HTTP 429 responses', + error: { statusCode: 429, message: 'Too many requests' }, + expected: 'rate-limit', + }, + { + label: 'HTTP 500 responses', + error: { status: 500, message: 'Internal Server Error' }, + expected: 'provider', + }, + { + label: 'HTTP 502 responses', + error: { response: { status: 502 }, message: 'Bad Gateway' }, + expected: 'provider', + }, + { + label: 'HTTP 408 responses', + error: { status: 408, message: 'Request Timeout' }, + expected: 'timeout', + }, + { + label: 'Tinfoil 504 upstream timeout responses', + error: { statusCode: 504, responseBody: 'tinfoil upstream timeout' }, + expected: 'timeout', + }, + { + label: 'Tinfoil upstream idle timeout stream errors', + error: new Error('tinfoil upstream idle timeout'), + expected: 'timeout', + }, + { + label: 'bare non-attestation Chrome fetch failures', + error: new TypeError('Failed to fetch'), + expected: 'network', + }, + { + label: 'Safari fetch failures', + error: new TypeError('Load failed'), + expected: 'network', + }, + { + label: 'Firefox fetch failures', + error: new TypeError('NetworkError when attempting to fetch resource.'), + expected: 'network', + }, + ] + + for (const { label, error, expected } of cases) { + it(`classifies ${label}`, () => { + expect(classifyErrorKind(error)).toBe(expected) + }) + } + + it('returns undefined for unknown errors', () => { + expect(classifyErrorKind(new Error('Unknown failure'))).toBeUndefined() + expect(classifyErrorKind(null)).toBeUndefined() + }) +}) + +describe('getChatErrorKind', () => { + it('uses the kind carried by a serialized stream error', () => { + const error = new Error(JSON.stringify({ error: 'Bad Gateway', status: 502, kind: 'timeout' })) + + expect(getChatErrorKind(error)).toBe('timeout') + }) + + it('classifies kind-less legacy payloads from their status and message', () => { + expect(getChatErrorKind(new Error(JSON.stringify({ error: 'Bad Gateway', status: 502 })))).toBe('provider') + expect(getChatErrorKind(new Error(JSON.stringify({ error: 'tinfoil upstream timeout', statusCode: 504 })))).toBe( + 'timeout', + ) + }) + + it('classifies raw browser network errors', () => { + expect(getChatErrorKind(new TypeError('Failed to fetch'))).toBe('network') + }) + + it('returns undefined for unknown, null, and missing errors', () => { + expect(getChatErrorKind(new Error('Unknown failure'))).toBeUndefined() + expect(getChatErrorKind(null)).toBeUndefined() + expect(getChatErrorKind()).toBeUndefined() + }) +}) + +describe('getErrorName', () => { + it('returns string names from error-like values', () => { + expect(getErrorName(new TypeError('failed'))).toBe('TypeError') + expect(getErrorName({ name: 'SdkError' })).toBe('SdkError') + }) + + it('returns undefined when no string name exists', () => { + expect(getErrorName({ name: 42 })).toBeUndefined() + expect(getErrorName(null)).toBeUndefined() + }) +}) + describe('isRateLimitError', () => { it('detects 429 from JSON response body (DefaultChatTransport path)', () => { const error = new Error(JSON.stringify({ error: 'API call failed', statusCode: 429 })) diff --git a/src/lib/error-utils.ts b/src/lib/error-utils.ts index ebd27e070..6035e585d 100644 --- a/src/lib/error-utils.ts +++ b/src/lib/error-utils.ts @@ -3,6 +3,12 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ import type { HandleError, HandleErrorCode } from '@/types/handle-errors' +import { tinfoilUpstreamIdleTimeoutMessage, tinfoilUpstreamTimeoutMessage } from '@shared/tinfoil-proxy' + +const chatErrorKinds = ['attestation', 'timeout', 'rate-limit', 'provider', 'network'] as const +export type ChatErrorKind = (typeof chatErrorKinds)[number] + +const isChatErrorKind = (value: unknown): value is ChatErrorKind => chatErrorKinds.includes(value as ChatErrorKind) const parseJson = (str: string): Record | undefined => { try { @@ -11,6 +17,103 @@ const parseJson = (str: string): Record | undefined => { return undefined } } +const providerErrorNames = new Set(['KeyConfigMismatchError', 'ProtocolError', 'DecryptionError']) +const timeoutMarkers = [tinfoilUpstreamTimeoutMessage, tinfoilUpstreamIdleTimeoutMessage] +const networkErrorMarkers = ['failed to fetch', 'load failed', 'networkerror'] + +type ErrorClassificationFields = { + name?: string + status?: number + message?: string +} + +const firstNumber = (...values: unknown[]): number | undefined => + values.find((value): value is number => typeof value === 'number') + +const firstString = (...values: unknown[]): string | undefined => + values.find((value): value is string => typeof value === 'string') + +/** Normalize raw error fields used by chat error classification. */ +const getErrorClassificationFields = (error: unknown): ErrorClassificationFields => { + if (typeof error === 'string') { + return { message: error } + } + if (typeof error !== 'object' || error === null) { + return {} + } + + const value = error as Record + const response = (value.response ?? {}) as Record + return { + name: getErrorName(value), + status: firstNumber(value.status, value.statusCode, response.status), + message: firstString(value.responseBody, value.message), + } +} + +/** + * Classify a raw pipeline error into a stable user-facing chat error kind. + * Classification uses only normalized error name, HTTP status, and message. + */ +export const classifyErrorKind = (error: unknown): ChatErrorKind | undefined => { + const { name, status, message } = getErrorClassificationFields(error) + const normalizedMessage = message?.toLowerCase() + + if (name === 'TinfoilAttestationTimeoutError') { + return 'timeout' + } + if (name === 'TinfoilAttestationError') { + return 'attestation' + } + if (status === 429) { + return 'rate-limit' + } + if (status === 408 || (normalizedMessage && timeoutMarkers.some((marker) => normalizedMessage.includes(marker)))) { + return 'timeout' + } + if ((name && providerErrorNames.has(name)) || (status !== undefined && status >= 500)) { + return 'provider' + } + if ( + name === 'TypeError' && + normalizedMessage && + networkErrorMarkers.some((marker) => normalizedMessage.includes(marker)) + ) { + return 'network' + } + return undefined +} + +/** + * Read a serialized chat error kind, falling back to legacy status/message + * classification when older payloads do not carry one. + */ +export const getChatErrorKind = (error?: Error | null): ChatErrorKind | undefined => { + if (!error?.message) { + return undefined + } + + const parsed = parseJson(error.message) + if (!parsed) { + return classifyErrorKind(error) + } + if (isChatErrorKind(parsed.kind)) { + return parsed.kind + } + + return classifyErrorKind({ + status: firstNumber(parsed.status, parsed.statusCode), + message: firstString(parsed.error), + }) +} + +/** Return an error-like value's string name, when present. */ +export const getErrorName = (error: unknown): string | undefined => { + if (typeof error !== 'object' || error === null || !('name' in error)) { + return undefined + } + return typeof error.name === 'string' ? error.name : undefined +} /** Check whether an error represents a rate-limit (HTTP 429) response. */ export const isRateLimitError = (error?: Error | null): boolean => { @@ -39,13 +142,7 @@ export const getErrorStatusCode = (error?: Error | null): number | undefined => return undefined } const parsed = parseJson(error.message) - if (typeof parsed?.status === 'number') { - return parsed.status - } - if (typeof parsed?.statusCode === 'number') { - return parsed.statusCode - } - return undefined + return firstNumber(parsed?.status, parsed?.statusCode) } /** diff --git a/src/lib/posthog.tsx b/src/lib/posthog.tsx index 654ae9ed2..3392522a2 100644 --- a/src/lib/posthog.tsx +++ b/src/lib/posthog.tsx @@ -159,10 +159,13 @@ export type EventType = | 'chat_send_prompt_overflow' | 'chat_receive_reply' | 'chat_auto_retry' + | 'chat_retry_success' + | 'chat_retries_exhausted' | 'chat_select' | 'chat_new_clicked' | 'chat_delete' | 'chat_clear_all' + | 'tinfoil_attestation' // Model & Settings | 'model_select' | 'mode_select' diff --git a/src/lib/timeout.test.ts b/src/lib/timeout.test.ts index 5913e924f..82b7a0459 100644 --- a/src/lib/timeout.test.ts +++ b/src/lib/timeout.test.ts @@ -4,8 +4,8 @@ import { getClock } from '@/testing-library' import { act } from '@testing-library/react' -import { describe, expect, it } from 'bun:test' -import { withTimeout } from './timeout' +import { describe, expect, it, mock } from 'bun:test' +import { withDeadline, withTimeout } from './timeout' describe('withTimeout', () => { it('should resolve with promise value when promise resolves before timeout', async () => { @@ -47,3 +47,33 @@ describe('withTimeout', () => { expect(result).toBeUndefined() }) }) + +describe('withDeadline', () => { + it('rejects with an error created when the deadline fires', async () => { + const createError = mock(() => new Error('deadline exceeded')) + const resultPromise = withDeadline(new Promise(() => {}), 100, createError) + const errorPromise = resultPromise.then( + () => undefined, + (error: unknown) => error, + ) + + expect(createError).not.toHaveBeenCalled() + await act(async () => { + await getClock().tickAsync(100) + }) + + await expect(errorPromise).resolves.toMatchObject({ message: 'deadline exceeded' }) + expect(createError).toHaveBeenCalledTimes(1) + }) + + it('does not create the deadline error when the source promise settles first', async () => { + const createError = mock(() => new Error('deadline exceeded')) + + await expect(withDeadline(Promise.resolve('ok'), 100, createError)).resolves.toBe('ok') + await act(async () => { + await getClock().tickAsync(100) + }) + + expect(createError).not.toHaveBeenCalled() + }) +}) diff --git a/src/lib/timeout.ts b/src/lib/timeout.ts index 97b2e83e7..ee081ea60 100644 --- a/src/lib/timeout.ts +++ b/src/lib/timeout.ts @@ -22,3 +22,21 @@ export const withTimeout = (promise: Promise, ms: number, label: string): } }) } + +/** + * Races a promise against a rejecting deadline and clears the deadline timer + * when either promise settles. Creates the timeout error only when needed. + */ +export const withDeadline = (promise: Promise, ms: number, createError: () => Error): Promise => { + let timeoutId: ReturnType | undefined + + const deadlinePromise = new Promise((_, reject) => { + timeoutId = setTimeout(() => reject(createError()), ms) + }) + + return Promise.race([promise, deadlinePromise]).finally(() => { + if (timeoutId !== undefined) { + clearTimeout(timeoutId) + } + }) +} diff --git a/src/settings/models/model-catalog.ts b/src/settings/models/model-catalog.ts index 538aba78c..2c4fdc793 100644 --- a/src/settings/models/model-catalog.ts +++ b/src/settings/models/model-catalog.ts @@ -2,7 +2,7 @@ * License, v. 2.0. If a copy of the MPL was not distributed with this * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ -import { getTinfoilClient } from '@/ai/fetch' +import { getTinfoilClient } from '@/ai/tinfoil-client' import { fetch } from '@/lib/fetch' import { http } from '@/lib/http' import { normalizeOpenAiBaseUrl } from '@/lib/openai-base-url' diff --git a/src/types/acp.ts b/src/types/acp.ts index 2a6b9d942..576e3a5c1 100644 --- a/src/types/acp.ts +++ b/src/types/acp.ts @@ -15,6 +15,7 @@ import type { RequestPermissionRequest, RequestPermissionResponse } from '@agent import type { HttpClient } from '@/lib/http' import type { FetchFn } from '@/lib/proxy-fetch' import type { SessionSideEffectSink } from '@/acp/translators/acp-to-ai-sdk' +import type { TurnBudgetConsumer } from '@/ai/retry-budget' import type { ChatThread, Model, SaveMessagesFunction } from '@/types' /** Capabilities advertised by an ACP agent on `initialize`. Stored on the @@ -73,6 +74,7 @@ export type AgentAdapterContext = { reconnectClient: (client: MCPClient) => Promise httpClient: HttpClient getProxyFetch: () => FetchFn + turnBudget?: TurnBudgetConsumer /** Increments only when the current assistant response is regenerated. Built-in * persistent harnesses use it to rebuild from the edited transcript without * rebuilding during ordinary transcript growth. */ diff --git a/src/voice/engine/thunderbolt-engine.ts b/src/voice/engine/thunderbolt-engine.ts index d87e650c7..fba06a31b 100644 --- a/src/voice/engine/thunderbolt-engine.ts +++ b/src/voice/engine/thunderbolt-engine.ts @@ -11,7 +11,7 @@ * is processed enclave-private. Only the *transport* is Tinfoil-specific; the * request/response orchestration is shared via `createAudioEngine`. */ -import { evictSystemTinfoilClient, getSystemTinfoilClient, isKeyConfigMismatchError } from '@/ai/fetch' +import { evictSystemTinfoilClient, getSystemTinfoilClient, isTinfoilTransportWedgedError } from '@/ai/tinfoil-client' import { isSsoMode } from '@/lib/auth-mode' import { getAuthToken } from '@/lib/auth-token' import { type AudioTransport, createAudioEngine } from './audio-engine' @@ -41,9 +41,9 @@ const ttsInstructions = /** * POST to a Tinfoil `/v1/audio/*` endpoint via the attested `SecureClient`, * attaching the app's session auth (the `/tinfoil` route's guard needs the real - * token/SSO cookies, not the SDK placeholder). On a `KeyConfigMismatchError` - * that survives the SDK's own retry the cached client is wedged — evict it and - * retry once with a fresh attestation context (mirrors the chat path). + * token/SSO cookies, not the SDK placeholder). On a wedged transport failure + * that survives the SDK's own reset/retry, evict the cached client and retry + * once with a fresh attestation context (mirrors the chat path). */ const tinfoilTransport: AudioTransport = async (path, body, headers, signal) => { const sso = isSsoMode() @@ -64,8 +64,7 @@ const tinfoilTransport: AudioTransport = async (path, body, headers, signal) => } return client.fetch(`${baseUrl}${path}`, init) } - // Retry once on a wedged attestation key — whether the SDK throws - // KeyConfigMismatchError or the enclave returns it as a 422. + // Retry once on a wedged transport or a 422 key-mismatch response. try { const res = await attempt() if (res.status === 422) { @@ -74,7 +73,7 @@ const tinfoilTransport: AudioTransport = async (path, body, headers, signal) => } return res } catch (err) { - if (!isKeyConfigMismatchError(err)) { + if (!isTinfoilTransportWedgedError(err)) { throw err } evictSystemTinfoilClient()