diff --git a/backend/src/auth/session-cookie-httponly.test.ts b/backend/src/auth/session-cookie-httponly.test.ts index 243d6ddde..f3bf436f2 100644 --- a/backend/src/auth/session-cookie-httponly.test.ts +++ b/backend/src/auth/session-cookie-httponly.test.ts @@ -34,8 +34,13 @@ const buildEmailDeps = (): AuthEmailDeps => ({ const parseSetCookie = (raw: string) => { const [nameValue, ...attrParts] = raw.split(';').map((p) => p.trim()) const name = nameValue.split('=')[0] - const flags = new Set(attrParts.map((p) => p.split('=')[0].toLowerCase())) - return { name, flags } + const attributes = new Map( + attrParts.map((part) => { + const [key, value = ''] = part.split('=', 2) + return [key.toLowerCase(), value] as const + }), + ) + return { name, attributes, flags: new Set(attributes.keys()) } } /** Find the session-token cookie among all Set-Cookie lines (tolerates the `__Secure-` prefix). */ @@ -86,7 +91,7 @@ describe('session cookie HttpOnly (CWE-1004 regression)', () => { } }) - it('email-OTP sign-in issues an HttpOnly, SameSite session cookie (and no JS-readable __session cookie)', async () => { + it('email-OTP sign-in issues an HttpOnly session cookie', async () => { const email = `httponly-otp-${crypto.randomUUID()}@example.com` await db.insert(waitlist).values({ id: crypto.randomUUID(), email, status: 'approved' }) diff --git a/backend/src/config/cors.test.ts b/backend/src/config/cors.test.ts index b36481c2b..6e7dadaaa 100644 --- a/backend/src/config/cors.test.ts +++ b/backend/src/config/cors.test.ts @@ -3,8 +3,8 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ import { describe, expect, it } from 'bun:test' +import { createCorsMiddleware } from '@/config/cors' import { getCorsOriginsList } from '@/config/settings' -import cors from '@elysiajs/cors' import { Elysia } from 'elysia' /** @@ -12,14 +12,14 @@ import { Elysia } from 'elysia' * Verifies that the actual HTTP headers are set correctly for various origins. */ describe('CORS integration', () => { - const createTestApp = (corsOrigins: string[]) => + const createTestApp = (corsOrigins: string[], corsAllowCredentials = true) => new Elysia() .use( - cors({ - origin: corsOrigins, - credentials: true, - methods: 'GET,POST,PUT,DELETE,PATCH,OPTIONS', - allowedHeaders: 'Content-Type,Authorization', + createCorsMiddleware({ + corsOrigins: corsOrigins.join(','), + corsAllowCredentials, + corsAllowMethods: 'GET,POST,PUT,DELETE,PATCH,OPTIONS', + corsExposeHeaders: '', }), ) .get('/test', () => ({ ok: true })) @@ -40,6 +40,7 @@ describe('CORS integration', () => { expect(res.headers.get('access-control-allow-origin')).toBe('https://app.example.com') expect(res.headers.get('access-control-allow-credentials')).toBe('true') + expect(res.headers.get('timing-allow-origin')).toBe('https://app.example.com') }) it('should allow tauri://localhost', async () => { @@ -75,6 +76,7 @@ describe('CORS integration', () => { ) expect(res.headers.get('access-control-allow-origin')).toBeNull() + expect(res.headers.get('timing-allow-origin')).toBeNull() }) it('should reject unknown origins', async () => { @@ -101,6 +103,24 @@ describe('CORS integration', () => { ) expect(res.headers.get('access-control-allow-origin')).toBeNull() + expect(res.headers.get('timing-allow-origin')).toBeNull() + }) + + it('should expose timing for preflight from an allowed origin', async () => { + const app = createTestApp(origins) + const res = await app.handle( + new Request('http://localhost/test', { + method: 'OPTIONS', + headers: { + Origin: 'https://app.example.com', + 'Access-Control-Request-Method': 'DELETE', + }, + }), + ) + + expect(res.headers.get('access-control-allow-origin')).toBe('https://app.example.com') + expect(res.headers.get('access-control-max-age')).toBe('600') + expect(res.headers.get('timing-allow-origin')).toBe('https://app.example.com') }) }) @@ -131,4 +151,31 @@ describe('CORS integration', () => { expect(res.headers.get('access-control-allow-origin')).toBeNull() }) }) + + it('echoes the request origin for credentialed wildcard access', async () => { + const app = createTestApp(['*']) + const res = await app.handle( + new Request('http://localhost/test', { + headers: { Origin: 'https://app.example.com' }, + }), + ) + + expect(res.headers.get('access-control-allow-origin')).toBe('https://app.example.com') + expect(res.headers.get('access-control-allow-credentials')).toBe('true') + expect(res.headers.get('timing-allow-origin')).toBe('https://app.example.com') + expect(res.headers.get('vary')).toContain('Origin') + }) + + it('uses wildcard access when credentials are disabled', async () => { + const app = createTestApp(['*'], false) + const res = await app.handle( + new Request('http://localhost/test', { + headers: { Origin: 'https://app.example.com' }, + }), + ) + + expect(res.headers.get('access-control-allow-origin')).toBe('*') + expect(res.headers.get('access-control-allow-credentials')).toBeNull() + expect(res.headers.get('timing-allow-origin')).toBe('*') + }) }) diff --git a/backend/src/config/cors.ts b/backend/src/config/cors.ts new file mode 100644 index 000000000..941703f71 --- /dev/null +++ b/backend/src/config/cors.ts @@ -0,0 +1,65 @@ +/* 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 { getCorsOriginsList, isOriginAllowed, type Settings } from '@/config/settings' +import cors from '@elysiajs/cors' +import { Elysia } from 'elysia' + +type CorsSettings = Pick + +/** Resolve a request origin using the configured exact-origin CORS policy. */ +const resolveCorsOrigin = ( + request: Request, + settings: Pick, + allowsAnyOrigin: boolean, + allowCredentials: boolean, +): string | null => { + if (allowsAnyOrigin) { + // Browsers reject wildcard ACAO with credentials, so credentialed wildcard policies must echo Origin. + return allowCredentials ? request.headers.get('Origin') : '*' + } + + const origin = request.headers.get('Origin') + if (origin === null || !isOriginAllowed(origin, settings)) { + return null + } + return origin +} + +/** Configure CORS and grant Resource Timing access to the origin CORS resolved for the request. */ +export const createCorsMiddleware = (settings: CorsSettings) => { + const corsOrigins = getCorsOriginsList(settings) + const allowsAnyOrigin = corsOrigins.includes('*') + const resolveRequestOrigin = (request: Request) => + resolveCorsOrigin(request, settings, allowsAnyOrigin, settings.corsAllowCredentials) + const corsOrigin = + allowsAnyOrigin && !settings.corsAllowCredentials + ? '*' + : (request: Request) => resolveRequestOrigin(request) !== null + + return new Elysia({ name: 'cors-with-resource-timing' }) + .onRequest(({ request, set }) => { + const allowedOrigin = resolveRequestOrigin(request) + if (allowedOrigin === null) { + return + } + + // Resource Timing consumes TAO without CORS exposure; timing is revealed only to the CORS-allowed origin. + set.headers['Timing-Allow-Origin'] = allowedOrigin + }) + .use( + cors({ + origin: corsOrigin, + credentials: settings.corsAllowCredentials, + methods: settings.corsAllowMethods, + // Echo back the client's Access-Control-Request-Headers. The universal + // proxy forwards arbitrary upstream headers as X-Proxy-Passthrough-*. + allowedHeaders: true, + exposeHeaders: settings.corsExposeHeaders, + // Preflights cost ~195ms measured. 10 minutes covers back-to-back chat sends while keeping a + // CORS policy change from lingering in browser caches (Safari caps around this value anyway). + maxAge: 600, + }), + ) +} diff --git a/backend/src/config/settings.test.ts b/backend/src/config/settings.test.ts index cb9ba83f4..45a5d6c17 100644 --- a/backend/src/config/settings.test.ts +++ b/backend/src/config/settings.test.ts @@ -52,7 +52,7 @@ describe('Config Settings', () => { }) describe('CORS default security', () => { - const corsEnvKeys = ['CORS_ORIGINS'] as const + const corsEnvKeys = ['CORS_ORIGINS', 'CORS_EXPOSE_HEADERS'] as const let savedEnv: Partial> @@ -99,6 +99,15 @@ describe('Config Settings', () => { expect(isOriginAllowed('http://localhost:1420', settings)).toBe(true) }) + it('should expose proxy and server timing to cross-origin clients by default', () => { + delete process.env.CORS_EXPOSE_HEADERS + const settings = getSettings() + + expect(settings.corsExposeHeaders.split(',')).toContain('X-Proxy-Timing') + expect(settings.corsExposeHeaders.split(',')).toContain('Server-Timing') + expect(settings.corsExposeHeaders.split(',')).not.toContain('Timing-Allow-Origin') + }) + it('should not match non-Tauri origins by default', () => { delete process.env.CORS_ORIGINS const settings = getSettings() diff --git a/backend/src/config/settings.ts b/backend/src/config/settings.ts index 3dc4b99ea..877963cf0 100644 --- a/backend/src/config/settings.ts +++ b/backend/src/config/settings.ts @@ -7,6 +7,8 @@ import { z } from 'zod' const betterAuthTimeString = z.string().regex(/^\d+[smhd]$/, { message: 'must be a Better Auth time string (digits followed by s, m, h, or d)', }) +const defaultCorsExposeHeaders = + '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,X-Proxy-Timing,Server-Timing' /** * Settings schema for environment variables validation @@ -102,11 +104,7 @@ const settingsSchema = z corsAllowMethods: z.string().default('GET,POST,PUT,DELETE,PATCH,OPTIONS'), corsAllowHeaders: z.string().default(''), // Protocol-required: frontend proxy-fetch.ts unwrap needs these visible cross-origin (cors does not echo expose-headers). - 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,Ehbp-Response-Nonce', - ), + corsExposeHeaders: z.string().default(defaultCorsExposeHeaders), // E2E encryption — when true, devices must complete the trust flow before syncing e2eeEnabled: z.boolean().default(false), @@ -218,9 +216,7 @@ const parseSettings = (): Settings => { corsAllowCredentials: process.env.CORS_ALLOW_CREDENTIALS !== 'false', corsAllowMethods: process.env.CORS_ALLOW_METHODS || 'GET,POST,PUT,DELETE,PATCH,OPTIONS', 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,Ehbp-Response-Nonce', + corsExposeHeaders: process.env.CORS_EXPOSE_HEADERS || defaultCorsExposeHeaders, e2eeEnabled: process.env.E2EE_ENABLED === 'true', minAppVersion: process.env.MIN_APP_VERSION || '', swaggerEnabled: process.env.SWAGGER_ENABLED === 'true', diff --git a/backend/src/index.ts b/backend/src/index.ts index 8ea408e1b..9c6a14675 100644 --- a/backend/src/index.ts +++ b/backend/src/index.ts @@ -9,6 +9,7 @@ import { createMicrosoftAuthRoutes } from '@/auth/microsoft' import { createOidcConfigRoutes } from '@/auth/oidc' import { createSsoDesktopCallbackRoutes } from '@/auth/sso-desktop-callback' import { createLoggerMiddleware, createStandaloneLogger } from '@/config/logger' +import { createCorsMiddleware } from '@/config/cors' import { getCorsOriginsList, getSettings } from '@/config/settings' import { runMigrations } from '@/db/client' import { createInferenceRoutes } from '@/inference/routes' @@ -22,6 +23,7 @@ import { createSearchRoutes } from '@/api/search' import { createPreviewRoutes } from '@/api/preview' import { createPostHogRoutes } from '@/posthog/routes' import { createProToolsRoutes } from '@/pro/routes' +import { createTinfoilKeepWarm } from '@/tinfoil/keep-warm' import { createTinfoilRoutes } from '@/tinfoil/routes' import { createWaitlistRoutes } from '@/waitlist/routes' import { createAccountRoutes } from '@/api/account' @@ -31,7 +33,6 @@ import { createConfigRoutes } from '@/api/config' import { createEncryptionRoutes } from '@/api/encryption' import { createPowerSyncRoutes } from '@/api/powersync' import type { AppDeps } from '@/types' -import { cors } from '@elysiajs/cors' import { Elysia } from 'elysia' /** @@ -83,31 +84,20 @@ export const createApp = async (deps?: AppDeps) => { ) const auth = deps?.auth ?? createdAuth + const appLogger = createStandaloneLogger(settings) + // Build the production observability recorder unless tests injected their own. // Proxy events go to Pino + OTel only — not PostHog (proxy traffic is infra // plumbing, not product analytics). const proxyObservability = deps?.proxyObservability ?? createObservabilityRecorder({ - logger: createStandaloneLogger(settings), + logger: appLogger, }) return ( configuredApp - .use( - cors({ - origin: getCorsOriginsList(settings), - credentials: settings.corsAllowCredentials, - methods: settings.corsAllowMethods, - // Echo back the client's Access-Control-Request-Headers. The universal - // proxy at /v1/proxy forwards arbitrary upstream headers as - // X-Proxy-Passthrough-* (provider SDKs add x-api-key, x-stainless-*, - // openai-organization, anthropic-beta, …). A static allowlist can't - // enumerate every upstream's header set without breaking preflight. - allowedHeaders: true, - exposeHeaders: settings.corsExposeHeaders, - }), - ) + .use(createCorsMiddleware(settings)) .use(createLoggerMiddleware(settings)) .use(createHttpLoggingMiddleware(settings.trustedProxy)) .use(createErrorHandlingMiddleware()) @@ -129,7 +119,7 @@ export const createApp = async (deps?: AppDeps) => { dnsLookup: deps?.dnsLookup, }), ) - .use(createTinfoilRoutes({ auth, fetchFn, rateLimit: proRateLimit })) + .use(createTinfoilRoutes({ auth, fetchFn, logger: appLogger, rateLimit: proRateLimit })) .use( createUniversalProxyWsRoutes({ auth, @@ -140,7 +130,14 @@ export const createApp = async (deps?: AppDeps) => { ) .use(createSearchRoutes(auth, proRateLimit, { exaClient: deps?.searchExaClient })) .use(createPreviewRoutes({ auth, fetchFn, rateLimit: proRateLimit, dnsLookup: deps?.dnsLookup })) - .use(createInferenceRoutes(auth, createInferenceRateLimit(database, rateLimitSettings))) + .use( + createInferenceRoutes({ + auth, + fetchFn: deps?.fetchFn, + logger: appLogger, + rateLimit: createInferenceRateLimit(database, rateLimitSettings), + }), + ) .use(createConfigRoutes(settings)) .use(createPostHogRoutes(fetchFn)) .use( @@ -166,6 +163,7 @@ export const createApp = async (deps?: AppDeps) => { const startServer = async () => { const settings = getSettings() const log = createStandaloneLogger(settings) + const tinfoilKeepWarm = createTinfoilKeepWarm(settings, { logger: log }) // Set up logging log.info('Starting Thunderbolt Server...') @@ -196,6 +194,9 @@ const startServer = async () => { hostname, port: settings.port, reusePort: process.env.NODE_ENV === 'production', + // Must stay strictly above the ALB idle timeout (60s) — see deploy/pulumi; Bun cap is 255. + // If the server closes first, the ALB reuses a dead connection and returns 502. + idleTimeout: 120, }, () => { log.info( @@ -206,6 +207,7 @@ const startServer = async () => { }, '🦊 Elysia server started', ) + tinfoilKeepWarm.start() if (settings.swaggerEnabled) { log.info( @@ -220,11 +222,13 @@ const startServer = async () => { // Graceful shutdown process.on('SIGINT', async () => { + tinfoilKeepWarm.stop() log.info('Received SIGINT, shutting down gracefully...') process.exit(0) }) process.on('SIGTERM', async () => { + tinfoilKeepWarm.stop() log.info('Received SIGTERM, shutting down gracefully...') process.exit(0) }) diff --git a/backend/src/inference/client.ts b/backend/src/inference/client.ts index 006bce913..9b974dc1d 100644 --- a/backend/src/inference/client.ts +++ b/backend/src/inference/client.ts @@ -4,16 +4,131 @@ import { getSettings } from '@/config/settings' import { getPostHogClient, isPostHogConfigured } from '@/posthog/client' +import { elapsedMs } from '@/utils/timing' import { OpenAI as PostHogOpenAI } from '@posthog/ai' +import { AsyncLocalStorage } from 'node:async_hooks' import OpenAI from 'openai' export type InferenceProvider = 'fireworks' | 'mistral' | 'anthropic' -type InferenceClient = { +export type InferenceClient = { client: OpenAI | PostHogOpenAI provider: InferenceProvider } +export type InferenceUpstreamAttemptLog = { + event: 'inference_upstream_attempt' + provider: InferenceProvider + attempt: number + method: string + host: string + status: number | null + duration_ms: number + retry_after?: string + rate_limit_headers?: Record +} + +export type InferenceLogger = { + info: (context: InferenceUpstreamAttemptLog | object, message: string) => void +} + +export type InferenceClientOptions = { + fetchFn?: typeof fetch + logger?: InferenceLogger + /** Monotonic clock used for upstream-attempt instrumentation. */ + nowFn?: () => number +} + +type InferenceFetchOptions = InferenceClientOptions & { + provider: InferenceProvider +} + +export type InferenceAttemptTracker = { + attempts: number +} + +const inferenceAttemptStorage = new AsyncLocalStorage() + +/** Create request-local state used to count OpenAI SDK fetch attempts. */ +export const createInferenceAttemptTracker = (): InferenceAttemptTracker => ({ attempts: 0 }) + +/** Run an inference SDK call with request-local attempt counting enabled. */ +export const runWithInferenceAttemptTracking = (tracker: InferenceAttemptTracker, callback: () => T): T => + inferenceAttemptStorage.run(tracker, callback) + +/** Read the one-based attempt index emitted by the OpenAI SDK. */ +const getAttemptIndex = (input: RequestInfo | URL, init?: RequestInit): number => { + const headers = new Headers(init?.headers ?? (input instanceof Request ? input.headers : undefined)) + const retryCount = Number(headers.get('X-Stainless-Retry-Count') ?? 0) + return Number.isFinite(retryCount) ? retryCount + 1 : 1 +} + +/** Resolve the upstream HTTP method without inspecting request content. */ +const getRequestMethod = (input: RequestInfo | URL, init?: RequestInit): string => + (init?.method ?? (input instanceof Request ? input.method : 'GET')).toUpperCase() + +/** Resolve only the upstream hostname, excluding path, query, and credentials. */ +const getRequestHost = (input: RequestInfo | URL): string => + new URL(input instanceof Request ? input.url : input.toString()).hostname + +/** Collect rate-limit diagnostics while preserving upstream header names. */ +const getRateLimitHeaders = (headers: Headers): Record => + Object.fromEntries([...headers.entries()].filter(([name]) => name.toLowerCase().startsWith('x-ratelimit-'))) + +/** Emit one structured, body-free log entry for an upstream attempt. */ +const logUpstreamAttempt = ( + logger: InferenceLogger | undefined, + context: Omit, +) => { + logger?.info({ event: 'inference_upstream_attempt', ...context }, 'Inference upstream attempt') +} + +/** Wrap fetch with safe, per-attempt upstream telemetry for OpenAI-compatible clients. */ +export const createInferenceFetch = ({ + provider, + fetchFn = globalThis.fetch, + logger, + nowFn = () => performance.now(), +}: InferenceFetchOptions): typeof fetch => { + const instrumentedFetch = async (input: RequestInfo | URL, init?: RequestInit) => { + const attempt = getAttemptIndex(input, init) + const tracker = inferenceAttemptStorage.getStore() + if (tracker) { + tracker.attempts = Math.max(tracker.attempts, attempt) + } + + const startedAt = nowFn() + const requestContext = { + provider, + attempt, + method: getRequestMethod(input, init), + host: getRequestHost(input), + } + + try { + const response = await fetchFn(input, init) + const rateLimitHeaders = getRateLimitHeaders(response.headers) + const retryAfter = response.headers.get('retry-after') + logUpstreamAttempt(logger, { + ...requestContext, + status: response.status, + duration_ms: elapsedMs(startedAt, nowFn()), + ...(retryAfter === null ? {} : { retry_after: retryAfter }), + ...(Object.keys(rateLimitHeaders).length === 0 ? {} : { rate_limit_headers: rateLimitHeaders }), + }) + return response + } catch (error) { + logUpstreamAttempt(logger, { + ...requestContext, + status: null, + duration_ms: elapsedMs(startedAt, nowFn()), + }) + throw error + } + } + return instrumentedFetch as unknown as typeof fetch +} + /** * Lazily initialized Fireworks client */ @@ -32,7 +147,8 @@ let anthropicClient: OpenAI | PostHogOpenAI | null = null /** * Get the Fireworks AI client */ -const getFireworksClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { +const getFireworksClient = (options: InferenceClientOptions = {}): OpenAI | PostHogOpenAI => { + const { fetchFn, logger, nowFn } = options // Don't use cache when fetchFn is provided (primarily for testing) if (fireworksClient && !fetchFn) { return fireworksClient @@ -47,7 +163,8 @@ const getFireworksClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { const params = { apiKey: settings.fireworksApiKey, baseURL: 'https://api.fireworks.ai/inference/v1', - ...(fetchFn && { fetch: fetchFn }), + fetch: createInferenceFetch({ provider: 'fireworks', fetchFn, logger, nowFn }), + // OpenAI SDK defaults to 2 retries; changing maxRetries is a follow-up decision after collecting attempt data. } const client = isPostHogConfigured() @@ -68,7 +185,8 @@ const getFireworksClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { /** * Get the Mistral AI client using OpenAI-compatible API */ -const getMistralClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { +const getMistralClient = (options: InferenceClientOptions = {}): OpenAI | PostHogOpenAI => { + const { fetchFn, logger, nowFn } = options if (mistralClient && !fetchFn) { return mistralClient } @@ -82,7 +200,7 @@ const getMistralClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { const params = { apiKey: settings.mistralApiKey, baseURL: 'https://api.mistral.ai/v1', - ...(fetchFn && { fetch: fetchFn }), + fetch: createInferenceFetch({ provider: 'mistral', fetchFn, logger, nowFn }), } const client = isPostHogConfigured() @@ -102,7 +220,8 @@ const getMistralClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { /** * Get the Anthropic AI client using OpenAI-compatible API */ -const getAnthropicClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { +const getAnthropicClient = (options: InferenceClientOptions = {}): OpenAI | PostHogOpenAI => { + const { fetchFn, logger, nowFn } = options if (anthropicClient && !fetchFn) { return anthropicClient } @@ -116,7 +235,7 @@ const getAnthropicClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { const params = { apiKey: settings.anthropicApiKey, baseURL: 'https://api.anthropic.com/v1/', - ...(fetchFn && { fetch: fetchFn }), + fetch: createInferenceFetch({ provider: 'anthropic', fetchFn, logger, nowFn }), } const client = isPostHogConfigured() @@ -137,11 +256,14 @@ const getAnthropicClient = (fetchFn?: typeof fetch): OpenAI | PostHogOpenAI => { * Get the appropriate inference client based on provider * Clients are lazily initialized and reused across requests */ -export const getInferenceClient = (provider: InferenceProvider, fetchFn?: typeof fetch): InferenceClient => { +export const getInferenceClient = ( + provider: InferenceProvider, + options: InferenceClientOptions = {}, +): InferenceClient => { const clientMap: Record OpenAI | PostHogOpenAI> = { - mistral: () => getMistralClient(fetchFn), - anthropic: () => getAnthropicClient(fetchFn), - fireworks: () => getFireworksClient(fetchFn), + mistral: () => getMistralClient(options), + anthropic: () => getAnthropicClient(options), + fireworks: () => getFireworksClient(options), } const client = clientMap[provider]() diff --git a/backend/src/inference/instrumentation.test.ts b/backend/src/inference/instrumentation.test.ts new file mode 100644 index 000000000..f8e6f6d92 --- /dev/null +++ b/backend/src/inference/instrumentation.test.ts @@ -0,0 +1,155 @@ +/* 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 { clearSettingsCache } from '@/config/settings' +import { clearPostHogClient } from '@/posthog/client' +import { mockAuth } from '@/test-utils/mock-auth' +import { afterEach, beforeEach, describe, expect, it } from 'bun:test' +import { Elysia } from 'elysia' +import { clearInferenceClientCache, type InferenceUpstreamAttemptLog } from './client' +import { createInferenceRoutes, type InferenceProxyLatencyLog } from './routes' + +type InferenceLog = InferenceUpstreamAttemptLog | InferenceProxyLatencyLog + +const successfulStream = + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":0,"model":"test","choices":[{"index":0,"delta":{"content":"ok"},"finish_reason":null}]}\n\n' + + 'data: [DONE]\n\n' + +describe('inference attempt instrumentation', () => { + let originalFireworksApiKey: string | undefined + let originalAnthropicApiKey: string | undefined + let originalPostHogApiKey: string | undefined + + beforeEach(() => { + originalFireworksApiKey = process.env.FIREWORKS_API_KEY + originalAnthropicApiKey = process.env.ANTHROPIC_API_KEY + originalPostHogApiKey = process.env.POSTHOG_API_KEY + process.env.FIREWORKS_API_KEY = 'test-fireworks-key' + process.env.ANTHROPIC_API_KEY = 'test-anthropic-key' + delete process.env.POSTHOG_API_KEY + clearSettingsCache() + clearInferenceClientCache() + clearPostHogClient() + }) + + afterEach(() => { + if (originalFireworksApiKey === undefined) { + delete process.env.FIREWORKS_API_KEY + } else { + process.env.FIREWORKS_API_KEY = originalFireworksApiKey + } + if (originalAnthropicApiKey === undefined) { + delete process.env.ANTHROPIC_API_KEY + } else { + process.env.ANTHROPIC_API_KEY = originalAnthropicApiKey + } + if (originalPostHogApiKey === undefined) { + delete process.env.POSTHOG_API_KEY + } else { + process.env.POSTHOG_API_KEY = originalPostHogApiKey + } + clearSettingsCache() + clearInferenceClientCache() + clearPostHogClient() + }) + + it.each([ + { + model: 'deepseek-v4-flash', + provider: 'fireworks' as const, + host: 'api.fireworks.ai', + apiKey: 'test-fireworks-key', + }, + { + model: 'opus-4.8', + provider: 'anthropic' as const, + host: 'api.anthropic.com', + apiKey: 'test-anthropic-key', + }, + ])('logs $provider 429 retry and surfaces attempts=2 after success', async ({ model, provider, host, apiKey }) => { + const logs: Array<{ context: InferenceLog; message: string }> = [] + let callCount = 0 + const fetchFn = (async () => { + callCount += 1 + if (callCount === 1) { + return new Response(JSON.stringify({ error: { message: 'Rate limited' } }), { + status: 429, + headers: { + 'Content-Type': 'application/json', + 'Retry-After': '0', + 'X-RateLimit-Limit-Requests': '60', + 'X-RateLimit-Remaining-Requests': '0', + }, + }) + } + return new Response(successfulStream, { + status: 200, + headers: { 'Content-Type': 'text/event-stream' }, + }) + }) as unknown as typeof fetch + const app = new Elysia().use( + createInferenceRoutes({ + auth: mockAuth, + fetchFn, + logger: { + info: (context, message) => logs.push({ context: context as InferenceLog, message }), + }, + }), + ) + + const response = await app.handle( + new Request('http://localhost/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model, + messages: [{ role: 'user', content: 'Hello' }], + stream: true, + }), + }), + ) + await response.text() + + const attemptLogs = logs + .map(({ context }) => context) + .filter((context): context is InferenceUpstreamAttemptLog => context.event === 'inference_upstream_attempt') + const latencyLogs = logs + .map(({ context }) => context) + .filter((context): context is InferenceProxyLatencyLog => context.event === 'inference_proxy_latency') + + expect(response.status).toBe(200) + expect(callCount).toBe(2) + expect(response.headers.get('x-proxy-timing')).toContain('attempts=2') + expect(attemptLogs).toHaveLength(2) + expect(attemptLogs[0]).toMatchObject({ + provider, + attempt: 1, + method: 'POST', + host, + status: 429, + retry_after: '0', + rate_limit_headers: { + 'x-ratelimit-limit-requests': '60', + 'x-ratelimit-remaining-requests': '0', + }, + }) + expect(attemptLogs[1]).toMatchObject({ + provider, + attempt: 2, + method: 'POST', + host, + status: 200, + }) + expect(attemptLogs.every(({ duration_ms }) => duration_ms >= 0)).toBeTrue() + expect(latencyLogs).toHaveLength(1) + expect(latencyLogs[0]).toMatchObject({ + provider, + status: 200, + attempts: 2, + }) + expect(JSON.stringify(logs)).not.toContain(apiKey) + expect(JSON.stringify(logs)).not.toContain('Hello') + expect(JSON.stringify(logs)).not.toContain('/inference/v1/chat/completions') + }) +}) diff --git a/backend/src/inference/posthog-privacy.test.ts b/backend/src/inference/posthog-privacy.test.ts index 379744a24..bbc3773b2 100644 --- a/backend/src/inference/posthog-privacy.test.ts +++ b/backend/src/inference/posthog-privacy.test.ts @@ -159,7 +159,7 @@ describe('Inference Routes - PostHog Privacy Integration', () => { clearInferenceClientCache() clearPostHogClient() - const { client } = getInferenceClient('fireworks', mockFetch) + const { client } = getInferenceClient('fireworks', { fetchFn: mockFetch }) // Verify it's a PostHog-wrapped client expect(client.constructor.name).toBe('PostHogOpenAI') @@ -176,7 +176,7 @@ describe('Inference Routes - PostHog Privacy Integration', () => { clearInferenceClientCache() clearPostHogClient() - const { client } = getInferenceClient('fireworks', mockFetch) + const { client } = getInferenceClient('fireworks', { fetchFn: mockFetch }) // Verify client exists and is functional expect(client).toBeDefined() @@ -197,7 +197,7 @@ describe('Inference Routes - PostHog Privacy Integration', () => { clearPostHogClient() // Get the wrapped client with injected mock fetch - const { client } = getInferenceClient('fireworks', mockFetch) + const { client } = getInferenceClient('fireworks', { fetchFn: mockFetch }) // Make a completion with sensitive data const completion = await (client as PostHogOpenAI).chat.completions.create({ @@ -258,7 +258,7 @@ describe('Inference Routes - PostHog Privacy Integration', () => { clearInferenceClientCache() clearPostHogClient() - const { client } = getInferenceClient('fireworks', mockFetch) + const { client } = getInferenceClient('fireworks', { fetchFn: mockFetch }) // Make multiple completions const conversations = [ diff --git a/backend/src/inference/routes.test.ts b/backend/src/inference/routes.test.ts index 60173e4a3..7b2aa592b 100644 --- a/backend/src/inference/routes.test.ts +++ b/backend/src/inference/routes.test.ts @@ -2,16 +2,13 @@ * 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 * as posthogClient from '@/posthog/client' import type { ConsoleSpies } from '@/test-utils/console-spies' import { setupConsoleSpy } from '@/test-utils/console-spies' import { mockAuth, mockAuthUnauthenticated } from '@/test-utils/mock-auth' -import * as streamingUtils from '@/utils/streaming' -import { afterAll, beforeAll, beforeEach, describe, expect, it, mock, spyOn } from 'bun:test' +import { afterAll, beforeAll, beforeEach, describe, expect, it, mock } from 'bun:test' import { Elysia } from 'elysia' import type OpenAI from 'openai' -import * as inferenceClient from './client' -import { createInferenceRoutes, supportedModels } from './routes' +import { createInferenceRoutes, supportedModels, type InferenceProxyLatencyLog } from './routes' import { defaultModels } from '@shared/defaults/models' describe('Thunderbolt model catalog parity', () => { @@ -27,9 +24,6 @@ describe('Thunderbolt model catalog parity', () => { describe('Inference Routes', () => { let app: { handle: Elysia['handle'] } - let getInferenceClientSpy: ReturnType - let isPostHogConfiguredSpy: ReturnType - let createSSEStreamSpy: ReturnType let consoleSpies: ConsoleSpies // Mock OpenAI client @@ -43,7 +37,13 @@ describe('Inference Routes', () => { }, } - const createMockStream = (chunks: any[] = []) => ({ + const getInferenceClientMock = mock(() => ({ + client: mockOpenAIClient as unknown as OpenAI, + provider: 'mistral' as const, + })) + const isPostHogConfiguredMock = mock(() => false) + + const createMockStream = (chunks: unknown[] = []) => ({ [Symbol.asyncIterator]: async function* () { for (const chunk of chunks) { yield chunk @@ -51,33 +51,18 @@ describe('Inference Routes', () => { }, }) - const createMockSSEStream = () => - new ReadableStream({ - start(controller) { - controller.enqueue(new TextEncoder().encode('data: {"test": "chunk"}\n\n')) - controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n')) - controller.close() - }, - }) - beforeAll(async () => { consoleSpies = setupConsoleSpy() - - // Mock dependencies - getInferenceClientSpy = spyOn(inferenceClient, 'getInferenceClient').mockReturnValue({ - client: mockOpenAIClient as unknown as OpenAI, - provider: 'mistral', - }) - isPostHogConfiguredSpy = spyOn(posthogClient, 'isPostHogConfigured').mockReturnValue(false) - createSSEStreamSpy = spyOn(streamingUtils, 'createSSEStreamFromCompletion').mockReturnValue(createMockSSEStream()) - - app = new Elysia().use(createInferenceRoutes(mockAuth)) + app = new Elysia().use( + createInferenceRoutes({ + auth: mockAuth, + getClient: getInferenceClientMock, + isPostHogConfiguredFn: isPostHogConfiguredMock, + }), + ) }) afterAll(() => { - getInferenceClientSpy?.mockRestore() - isPostHogConfiguredSpy?.mockRestore() - createSSEStreamSpy?.mockRestore() consoleSpies.restore() }) @@ -92,12 +77,13 @@ describe('Inference Routes', () => { beforeEach(() => { // Reset all mocks before each test mockCreateCompletion.mockClear() - createSSEStreamSpy.mockClear() - getInferenceClientSpy.mockClear() - getInferenceClientSpy.mockReturnValue({ + getInferenceClientMock.mockClear() + isPostHogConfiguredMock.mockClear() + isPostHogConfiguredMock.mockImplementation(() => false) + getInferenceClientMock.mockImplementation(() => ({ client: mockOpenAIClient as unknown as OpenAI, - provider: 'mistral', - }) + provider: 'mistral' as const, + })) }) it('should handle valid streaming request successfully', async () => { @@ -129,8 +115,6 @@ describe('Inference Routes', () => { tool_choice: undefined, stream: true, }) - - expect(createSSEStreamSpy).toHaveBeenCalledWith(mockCompletion) }) it('should route mistral models to mistral provider', async () => { @@ -146,7 +130,7 @@ describe('Inference Routes', () => { ) expect(response.status).toBe(200) - expect(getInferenceClientSpy).toHaveBeenCalledWith('mistral') + expect(getInferenceClientMock).toHaveBeenCalledWith('mistral') expect(mockCreateCompletion).toHaveBeenCalledWith( expect.objectContaining({ model: 'mistral-large-2512', @@ -182,7 +166,7 @@ describe('Inference Routes', () => { }) it('should include PostHog properties when configured', async () => { - isPostHogConfiguredSpy.mockReturnValue(true) + isPostHogConfiguredMock.mockImplementation(() => true) const mockCompletion = createMockStream() mockCreateCompletion.mockImplementation(() => Promise.resolve(mockCompletion)) @@ -207,7 +191,7 @@ describe('Inference Routes', () => { ) // Reset for other tests - isPostHogConfiguredSpy.mockReturnValue(false) + isPostHogConfiguredMock.mockImplementation(() => false) }) it('should reject non-streaming requests', async () => { @@ -261,6 +245,92 @@ describe('Inference Routes', () => { expect(response.status).toBe(500) }) + it('emits phase timing headers and a structured latency log on success', async () => { + const entries: Array<{ context: InferenceProxyLatencyLog; message: string }> = [] + const timestamps = [100, 120, 170] + const timingApp = new Elysia().use( + createInferenceRoutes({ + auth: mockAuth, + getClient: getInferenceClientMock, + logger: { + info: (context, message) => entries.push({ context: context as InferenceProxyLatencyLog, message }), + }, + nowFn: () => timestamps.shift() ?? 0, + }), + ) + mockCreateCompletion.mockImplementation(() => Promise.resolve(createMockStream())) + + const response = await timingApp.handle( + new Request('http://localhost/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(validRequestBody), + }), + ) + + expect(response.status).toBe(200) + expect(response.headers.get('x-proxy-timing')).toBe('pre=20;upstream=50;total=70;attempts=0') + expect(response.headers.get('server-timing')).toBe('pre;dur=20, upstream;dur=50, total;dur=70') + expect(entries).toEqual([ + { + context: { + event: 'inference_proxy_latency', + route: '/chat/completions', + provider: 'mistral', + status: 200, + preMs: 20, + upstreamMs: 50, + totalMs: 70, + attempts: 0, + }, + message: 'Inference proxy latency', + }, + ]) + }) + + it('emits phase timing headers and a structured latency log on upstream error', async () => { + const entries: Array<{ context: InferenceProxyLatencyLog; message: string }> = [] + const timestamps = [200, 230, 310] + const timingApp = new Elysia().use( + createInferenceRoutes({ + auth: mockAuth, + getClient: getInferenceClientMock, + logger: { + info: (context, message) => entries.push({ context: context as InferenceProxyLatencyLog, message }), + }, + nowFn: () => timestamps.shift() ?? 0, + }), + ) + mockCreateCompletion.mockImplementation(() => Promise.reject(new Error('Upstream failed'))) + + const response = await timingApp.handle( + new Request('http://localhost/chat/completions', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(validRequestBody), + }), + ) + + expect(response.status).toBe(500) + expect(response.headers.get('x-proxy-timing')).toBe('pre=30;upstream=80;total=110;attempts=0') + expect(response.headers.get('server-timing')).toBe('pre;dur=30, upstream;dur=80, total;dur=110') + expect(entries).toEqual([ + { + context: { + event: 'inference_proxy_latency', + route: '/chat/completions', + provider: 'mistral', + status: 500, + preMs: 30, + upstreamMs: 80, + totalMs: 110, + attempts: 0, + }, + message: 'Inference proxy latency', + }, + ]) + }) + it('should handle malformed JSON requests', async () => { const response = await app.handle( new Request('http://localhost/chat/completions', { @@ -280,7 +350,7 @@ describe('Inference Routes', () => { }) it('should handle requests with has_tools flag correctly', async () => { - isPostHogConfiguredSpy.mockReturnValue(true) + isPostHogConfiguredMock.mockImplementation(() => true) const mockCompletion = createMockStream() mockCreateCompletion.mockImplementation(() => Promise.resolve(mockCompletion)) @@ -306,14 +376,14 @@ describe('Inference Routes', () => { ) // Reset for other tests - isPostHogConfiguredSpy.mockReturnValue(false) + isPostHogConfiguredMock.mockImplementation(() => false) }) }) describe('authentication', () => { it('should return 401 when session is null', async () => { mockCreateCompletion.mockClear() - const unauthenticatedApp = new Elysia().use(createInferenceRoutes(mockAuthUnauthenticated)) + const unauthenticatedApp = new Elysia().use(createInferenceRoutes({ auth: mockAuthUnauthenticated })) const response = await unauthenticatedApp.handle( new Request('http://localhost/chat/completions', { @@ -335,12 +405,13 @@ describe('Inference Routes', () => { describe('message role sanitization', () => { beforeEach(() => { mockCreateCompletion.mockClear() - createSSEStreamSpy.mockClear() - getInferenceClientSpy.mockClear() - getInferenceClientSpy.mockReturnValue({ + getInferenceClientMock.mockClear() + isPostHogConfiguredMock.mockClear() + isPostHogConfiguredMock.mockImplementation(() => false) + getInferenceClientMock.mockImplementation(() => ({ client: mockOpenAIClient as unknown as OpenAI, - provider: 'mistral', - }) + provider: 'mistral' as const, + })) mockCreateCompletion.mockImplementation(() => Promise.resolve(createMockStream())) }) diff --git a/backend/src/inference/routes.ts b/backend/src/inference/routes.ts index f0799b128..ec9aceb91 100644 --- a/backend/src/inference/routes.ts +++ b/backend/src/inference/routes.ts @@ -4,18 +4,28 @@ import type { Auth } from '@/auth/elysia-plugin' import { createAuthMacro } from '@/auth/elysia-plugin' -import { safeErrorHandler } from '@/middleware/error-handling' +import { getErrorStatus, safeErrorHandler } from '@/middleware/error-handling' import { isPostHogConfigured } from '@/posthog/client' import { createSSEStreamFromCompletion } from '@/utils/streaming' +import { elapsedMs } from '@/utils/timing' import type { OpenAI as PostHogOpenAI } from '@posthog/ai' import { Elysia, type AnyElysia } from 'elysia' import { APIConnectionError, APIConnectionTimeoutError } from 'openai' import type { ChatCompletionMessageParam } from 'openai/resources/chat/completions' -import { getInferenceClient, type InferenceProvider } from './client' +import { + createInferenceAttemptTracker, + getInferenceClient, + runWithInferenceAttemptTracking, + type InferenceClient, + type InferenceLogger, + type InferenceProvider, +} from './client' type Message = { role: string; content: unknown } const privilegedRoles = new Set(['developer', 'system']) +const inferenceProxyTimingHeader = 'X-Proxy-Timing' +const serverTimingHeader = 'Server-Timing' /** Downgrade developer/system roles to user for all messages except the first (the legitimate system prompt). */ const sanitizeMessageRoles = (messages: Message[]): Message[] => @@ -53,13 +63,53 @@ export const supportedModels: Record = { }, } +export type InferenceProxyLatencyLog = { + event: 'inference_proxy_latency' + route: string + provider: InferenceProvider + status: number + preMs: number + upstreamMs: number + totalMs: number + attempts: number +} + +export type CreateInferenceRoutesOptions = { + auth: Auth + fetchFn?: typeof fetch + getClient?: (provider: InferenceProvider) => InferenceClient + isPostHogConfiguredFn?: () => boolean + logger?: InferenceLogger + /** Monotonic clock used for route latency and upstream-attempt instrumentation. */ + nowFn?: () => number + rateLimit?: AnyElysia +} + +/** Format inference phases using Server-Timing header syntax. */ +const formatServerTiming = (preMs: number, upstreamMs: number, totalMs: number): string => + `pre;dur=${preMs}, upstream;dur=${upstreamMs}, total;dur=${totalMs}` + /** * Inference API routes */ -export const createInferenceRoutes = (auth: Auth, rateLimit?: AnyElysia) => { +export const createInferenceRoutes = (options: CreateInferenceRoutesOptions) => { + const { auth, fetchFn, logger, rateLimit } = options + const nowFn = options.nowFn ?? (() => performance.now()) + const isPostHogConfiguredFn = options.isPostHogConfiguredFn ?? isPostHogConfigured + const getClient = + options.getClient ?? ((provider: InferenceProvider) => getInferenceClient(provider, { fetchFn, logger, nowFn })) const app = new Elysia({ prefix: '/chat', - }).onError(safeErrorHandler) + }) + .onError(safeErrorHandler) + .decorate('inferenceRequestStartedAt', 0) + .onRequest((ctx) => { + // onRequest hooks become app-wide when plugins merge, so avoid timing unrelated routes. + if (!ctx.request.url.includes('/chat/completions')) { + return + } + ctx.inferenceRequestStartedAt = nowFn() + }) return app.use(createAuthMacro(auth)).guard({ auth: true }, (guardedApp) => { if (rateLimit) { @@ -67,6 +117,8 @@ export const createInferenceRoutes = (auth: Auth, rateLimit?: AnyElysia) => { } return guardedApp.post('/completions', async (ctx) => { + const handlerStartedAt = nowFn() + const preMs = elapsedMs(ctx.inferenceRequestStartedAt, handlerStartedAt) const body = await ctx.request.json() if (!body.stream) { @@ -80,28 +132,52 @@ export const createInferenceRoutes = (auth: Auth, rateLimit?: AnyElysia) => { const { provider, internalName, omitTemperature } = modelConfig - const { client } = getInferenceClient(provider) + const { client } = getClient(provider) + const attemptTracker = createInferenceAttemptTracker() + const route = new URL(ctx.request.url).pathname + /** Emit route phase telemetry in structured logs and response headers. */ + const recordLatency = (status: number, completedAt: number) => { + const upstreamMs = elapsedMs(handlerStartedAt, completedAt) + const totalMs = elapsedMs(ctx.inferenceRequestStartedAt, completedAt) + const latency: InferenceProxyLatencyLog = { + event: 'inference_proxy_latency', + route, + provider, + status, + preMs, + upstreamMs, + totalMs, + attempts: attemptTracker.attempts, + } - console.info(`Routing model "${body.model}" to ${provider} provider`) + ctx.set.headers[inferenceProxyTimingHeader] = + `pre=${preMs};upstream=${upstreamMs};total=${totalMs};attempts=${attemptTracker.attempts}` + ctx.set.headers[serverTimingHeader] = formatServerTiming(preMs, upstreamMs, totalMs) + logger?.info(latency, 'Inference proxy latency') + } try { - const completion = await (client as PostHogOpenAI).chat.completions.create({ - model: internalName, - messages: sanitizeMessageRoles(body.messages) as ChatCompletionMessageParam[], - ...(omitTemperature ? {} : { temperature: body.temperature }), - tools: body.tools, - tool_choice: body.tool_choice, - stream: true, - ...(isPostHogConfigured() && { - posthogProperties: { - model_provider: provider, - endpoint: '/chat/completions', - has_tools: !!body.tools, - temperature: body.temperature, - // @todo add distinct id and trace id - }, + const completion = await runWithInferenceAttemptTracking(attemptTracker, () => + (client as PostHogOpenAI).chat.completions.create({ + model: internalName, + messages: sanitizeMessageRoles(body.messages) as ChatCompletionMessageParam[], + ...(omitTemperature ? {} : { temperature: body.temperature }), + tools: body.tools, + tool_choice: body.tool_choice, + stream: true, + ...(isPostHogConfiguredFn() && { + posthogProperties: { + model_provider: provider, + endpoint: '/chat/completions', + has_tools: !!body.tools, + temperature: body.temperature, + // @todo add distinct id and trace id + }, + }), }), - }) + ) + const upstreamResolvedAt = nowFn() + recordLatency(200, upstreamResolvedAt) const stream = createSSEStreamFromCompletion(completion) @@ -121,6 +197,7 @@ export const createInferenceRoutes = (auth: Auth, rateLimit?: AnyElysia) => { return new Response(stream, { headers: responseHeaders }) } catch (error) { + recordLatency(getErrorStatus(error), nowFn()) if (error instanceof APIConnectionError) { console.error('Failed to connect to inference provider', error.cause) throw new Error('Failed to connect to inference provider', { cause: error }) diff --git a/backend/src/tinfoil/keep-warm.test.ts b/backend/src/tinfoil/keep-warm.test.ts new file mode 100644 index 000000000..9d1324083 --- /dev/null +++ b/backend/src/tinfoil/keep-warm.test.ts @@ -0,0 +1,163 @@ +/* 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 } from 'bun:test' +import { createTinfoilKeepWarm, type TinfoilKeepWarmLogger } from './keep-warm' +import { createTinfoilUpstreamOriginStore } from './upstream-origin' + +const tinfoilSettings = { + tinfoilApiKey: 'test-tinfoil-key', + tinfoilEnclaveUrl: 'https://inference.tinfoil.sh/v1', +} + +const createLogger = () => { + const entries: Array<{ context: Record; message: string }> = [] + const logger: TinfoilKeepWarmLogger = { + debug: (context, message) => entries.push({ context, message }), + } + + return { entries, logger } +} + +describe('createTinfoilKeepWarm', () => { + it('probes the configured default before traffic, repeats, and stops cleanly', async () => { + const requests: Array<{ input: RequestInfo | URL; init?: RequestInit }> = [] + const fetchFn = (async (input: RequestInfo | URL, init?: RequestInit) => { + requests.push({ input, init }) + return new Response('{}', { status: 200 }) + }) as unknown as typeof fetch + const { logger } = createLogger() + const keepWarm = createTinfoilKeepWarm(tinfoilSettings, { + fetchFn, + intervalMs: 5, + logger, + upstreamOriginStore: createTinfoilUpstreamOriginStore(), + }) + + keepWarm.start() + try { + expect(requests).toHaveLength(1) + await new Promise((resolve) => setTimeout(resolve, 25)) + expect(requests.length).toBeGreaterThan(1) + } finally { + keepWarm.stop() + } + + const requestCountAfterStop = requests.length + await new Promise((resolve) => setTimeout(resolve, 15)) + + expect(requests).toHaveLength(requestCountAfterStop) + const firstRequest = requests[0] + expect(firstRequest?.input.toString()).toBe('https://inference.tinfoil.sh/v1/models') + expect(firstRequest?.init?.method).toBe('GET') + expect(new Headers(firstRequest?.init?.headers).get('authorization')).toBe('Bearer test-tinfoil-key') + }) + + it('probes the latest upstream origin while preserving the configured API prefix', () => { + const requests: Array = [] + const fetchFn = (async (input: RequestInfo | URL) => { + requests.push(input) + return new Response('{}', { status: 200 }) + }) as unknown as typeof fetch + const upstreamOriginStore = createTinfoilUpstreamOriginStore() + upstreamOriginStore.record('https://router.inf6.tinfoil.sh/v1/chat/completions?stream=true') + const { logger } = createLogger() + const keepWarm = createTinfoilKeepWarm(tinfoilSettings, { + fetchFn, + intervalMs: 100, + logger, + upstreamOriginStore, + }) + + keepWarm.start() + keepWarm.stop() + + expect(upstreamOriginStore.get()).toBe('https://router.inf6.tinfoil.sh') + expect(requests[0]?.toString()).toBe('https://router.inf6.tinfoil.sh/v1/models') + }) + + it('does not start without an API key', async () => { + const requests: Array = [] + const fetchFn = (async (input: RequestInfo | URL) => { + requests.push(input) + return new Response('{}', { status: 200 }) + }) as unknown as typeof fetch + const { logger } = createLogger() + const keepWarm = createTinfoilKeepWarm( + { ...tinfoilSettings, tinfoilApiKey: '' }, + { fetchFn, intervalMs: 1, logger }, + ) + + keepWarm.start() + await new Promise((resolve) => setTimeout(resolve, 10)) + keepWarm.stop() + + expect(requests).toHaveLength(0) + }) + + it('resumes probing after start, stop, and start', () => { + const requests: Array = [] + const fetchFn = (async (input: RequestInfo | URL) => { + requests.push(input) + return new Response('{}', { status: 200 }) + }) as unknown as typeof fetch + const { logger } = createLogger() + const keepWarm = createTinfoilKeepWarm(tinfoilSettings, { + fetchFn, + intervalMs: 100, + logger, + }) + + keepWarm.start() + expect(requests).toHaveLength(1) + keepWarm.stop() + keepWarm.start() + + try { + expect(requests).toHaveLength(2) + } finally { + keepWarm.stop() + } + }) + + it('times out failed probes, logs only safe metadata, and ignores the failure', async () => { + const aborted = Promise.withResolvers() + const fetchFn = ((_input: RequestInfo | URL, init?: RequestInit) => { + const signal = init?.signal as AbortSignal + + return new Promise((_resolve, reject) => { + const rejectOnAbort = () => { + aborted.resolve(signal.reason) + reject(signal.reason) + } + signal.addEventListener('abort', rejectOnAbort, { once: true }) + if (signal.aborted) { + rejectOnAbort() + } + }) + }) as unknown as typeof fetch + const { entries, logger } = createLogger() + const keepWarm = createTinfoilKeepWarm(tinfoilSettings, { + fetchFn, + intervalMs: 100, + timeoutMs: 1, + logger, + }) + + keepWarm.start() + try { + expect(await aborted.promise).toBeInstanceOf(DOMException) + await new Promise((resolve) => setImmediate(resolve)) + } finally { + keepWarm.stop() + } + + expect(entries).toHaveLength(1) + expect(entries[0]).toEqual({ + context: { errorName: 'TimeoutError' }, + message: 'Tinfoil enclave keep-warm request failed', + }) + expect(JSON.stringify(entries)).not.toContain(tinfoilSettings.tinfoilApiKey) + }) +}) diff --git a/backend/src/tinfoil/keep-warm.ts b/backend/src/tinfoil/keep-warm.ts new file mode 100644 index 000000000..be1284a86 --- /dev/null +++ b/backend/src/tinfoil/keep-warm.ts @@ -0,0 +1,88 @@ +/* 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 type { Settings } from '@/config/settings' +import { tinfoilUpstreamOriginStore, type TinfoilUpstreamOriginStore } from './upstream-origin' + +const defaultIntervalMs = 60_000 +const defaultTimeoutMs = 5_000 + +export type TinfoilKeepWarmLogger = { + debug: (context: Record, message: string) => void +} + +type TinfoilKeepWarmOptions = { + fetchFn?: typeof fetch + intervalMs?: number + timeoutMs?: number + logger: TinfoilKeepWarmLogger + upstreamOriginStore?: Pick +} + +type TinfoilKeepWarmController = { + start: () => void + stop: () => void +} + +/** + * Create a lifecycle controller that keeps Bun's Tinfoil connection pool warm. + */ +export const createTinfoilKeepWarm = ( + settings: Pick, + options: TinfoilKeepWarmOptions, +): TinfoilKeepWarmController => { + const fetchFn = options.fetchFn ?? globalThis.fetch + const intervalMs = options.intervalMs ?? defaultIntervalMs + const timeoutMs = options.timeoutMs ?? defaultTimeoutMs + const upstreamOriginStore = options.upstreamOriginStore ?? tinfoilUpstreamOriginStore + const defaultEnclaveUrl = settings.tinfoilEnclaveUrl.replace(/\/$/, '') + const apiPathPrefix = new URL(defaultEnclaveUrl).pathname.replace(/\/$/, '') + const state: { intervalId?: ReturnType } = {} + + const keepWarm = async () => { + try { + const latestOrigin = upstreamOriginStore.get() + const modelsUrl = `${latestOrigin ? `${latestOrigin}${apiPathPrefix}` : defaultEnclaveUrl}/models` + const response = await fetchFn(modelsUrl, { + method: 'GET', + headers: { Authorization: `Bearer ${settings.tinfoilApiKey}` }, + signal: AbortSignal.timeout(timeoutMs), + }) + await response.arrayBuffer() + + if (!response.ok) { + options.logger.debug( + { status: response.status }, + 'Tinfoil enclave keep-warm request returned a non-success status', + ) + } + } catch (error) { + // Background probe failures must never affect server availability. + options.logger.debug( + { errorName: error instanceof Error ? error.name : 'UnknownError' }, + 'Tinfoil enclave keep-warm request failed', + ) + } + } + + const start = () => { + if (state.intervalId) { + return + } + + if (!settings.tinfoilApiKey) { + return + } + + void keepWarm() + state.intervalId = setInterval(() => void keepWarm(), intervalMs) + } + + const stop = () => { + clearInterval(state.intervalId) + state.intervalId = undefined + } + + return { start, stop } +} diff --git a/backend/src/tinfoil/routes.test.ts b/backend/src/tinfoil/routes.test.ts index 1e8be6c9f..f70ff25bf 100644 --- a/backend/src/tinfoil/routes.test.ts +++ b/backend/src/tinfoil/routes.test.ts @@ -7,7 +7,9 @@ import { setupConsoleSpy } from '@/test-utils/console-spies' import { mockAuth, mockAuthUnauthenticated } from '@/test-utils/mock-auth' import { afterAll, beforeAll, beforeEach, describe, expect, it, mock } from 'bun:test' import { Elysia } from 'elysia' -import { createTinfoilRoutes } from './routes' +import { createTinfoilKeepWarm } from './keep-warm' +import { createTinfoilRoutes, type TinfoilProxyLatencyLog, type TinfoilProxyLogger } from './routes' +import { createTinfoilUpstreamOriginStore, type TinfoilUpstreamOriginStore } from './upstream-origin' const enclaveUrl = 'https://inference.tinfoil.sh' const testApiKey = 'test-tinfoil-key' @@ -95,18 +97,24 @@ describe('createTinfoilRoutes', () => { enclaveUrl?: string auth?: typeof mockAuth fetchFn?: typeof fetch + logger?: TinfoilProxyLogger + nowFn?: () => number upstreamHeadersTimeoutMs?: number upstreamIdleTimeoutMs?: number + upstreamOriginStore?: TinfoilUpstreamOriginStore } = {}, ) => new Elysia().use( createTinfoilRoutes({ auth: overrides.auth ?? mockAuth, fetchFn: overrides.fetchFn ?? (mockFetch as unknown as typeof fetch), + logger: overrides.logger, + nowFn: overrides.nowFn, apiKey: overrides.apiKey ?? testApiKey, enclaveUrl: overrides.enclaveUrl ?? enclaveUrl, upstreamHeadersTimeoutMs: overrides.upstreamHeadersTimeoutMs, upstreamIdleTimeoutMs: overrides.upstreamIdleTimeoutMs, + upstreamOriginStore: overrides.upstreamOriginStore, }), ) @@ -181,6 +189,23 @@ describe('createTinfoilRoutes', () => { expect(sent.get('connection')).toBeNull() }) + it('does not forward X-Tinfoil-Enclave-Url upstream', async () => { + const app = buildApp() + await drain( + await app.handle( + new Request('http://localhost/tinfoil/v1/chat/completions', { + method: 'POST', + headers: { 'X-Tinfoil-Enclave-Url': 'https://router.inf6.tinfoil.sh' }, + body: 'opaque-bytes', + }), + ), + ) + + const [, init] = mockFetch.mock.calls[0] as [string, RequestInit] + const sent = init.headers as Headers + expect(sent.get('x-tinfoil-enclave-url')).toBeNull() + }) + it('strips response hop-by-hop headers while preserving content encoding', async () => { mockFetch.mockResolvedValueOnce( makeOkResponse('opaque', { @@ -457,7 +482,316 @@ describe('createTinfoilRoutes', () => { }) }) + describe('latency instrumentation', () => { + it('logs and exposes pre-handler, fetch, and upstream-header phases', async () => { + const entries: Array<{ context: TinfoilProxyLatencyLog; message: string }> = [] + const timestamps = [90, 100, 112, 152] + const logger: TinfoilProxyLogger = { + info: (context, message) => entries.push({ context, message }), + } + const nowFn = () => { + const timestamp = timestamps.shift() + if (timestamp === undefined) { + throw new Error('Unexpected timing read') + } + return timestamp + } + const app = buildApp({ logger, nowFn }) + + const response = await drain( + await app.handle(new Request('http://localhost/tinfoil/v1/models?request-secret=do-not-log')), + ) + + expect(entries).toEqual([ + { + context: { + event: 'tinfoil_proxy_latency', + route: '/tinfoil/v1/models', + status: 200, + preHandlerMs: 10, + handlerToUpstreamFetchMs: 12, + upstreamFetchToHeadersMs: 40, + handlerToOutcomeMs: 52, + }, + message: 'Tinfoil proxy latency', + }, + ]) + expect(response.headers.get('x-proxy-timing')).toBe('pre=10;fetch=12;headers=40') + expect(response.headers.get('server-timing')).toBe('pre;dur=10, fetch;dur=12, headers;dur=40') + expect(JSON.stringify(entries)).not.toContain('request-secret') + expect(JSON.stringify(entries)).not.toContain(testApiKey) + }) + + it('logs one structured line for a request rejected before upstream fetch', async () => { + const entries: Array<{ context: TinfoilProxyLatencyLog; message: string }> = [] + const timestamps = [5, 10, 13] + const logger: TinfoilProxyLogger = { + info: (context, message) => entries.push({ context, message }), + } + const app = buildApp({ + apiKey: '', + logger, + nowFn: () => timestamps.shift() ?? 0, + }) + + const response = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + + expect(response.status).toBe(503) + expect(entries).toEqual([ + { + context: { + event: 'tinfoil_proxy_latency', + route: '/tinfoil/v1/models', + status: 503, + preHandlerMs: 5, + handlerToUpstreamFetchMs: null, + upstreamFetchToHeadersMs: null, + handlerToOutcomeMs: 3, + }, + message: 'Tinfoil proxy latency', + }, + ]) + }) + + it('logs exactly one 504 line when upstream headers time out', async () => { + const entries: Array<{ context: TinfoilProxyLatencyLog; message: string }> = [] + const timestamps = [90, 100, 105, 135] + const upstream = createAbortableFetch() + const logger: TinfoilProxyLogger = { + info: (context, message) => entries.push({ context, message }), + } + const app = buildApp({ + fetchFn: upstream.fetchFn, + logger, + nowFn: () => timestamps.shift() ?? 0, + upstreamHeadersTimeoutMs: 0, + }) + + const response = await app.handle( + new Request('http://localhost/tinfoil/v1/chat/completions', { + method: 'POST', + body: 'opaque-bytes', + }), + ) + + expect(response.status).toBe(504) + expect(response.headers.get('x-proxy-timing')).toBe('pre=10;fetch=5;headers=na') + expect(response.headers.get('server-timing')).toBe('pre;dur=10, fetch;dur=5') + expect(entries).toEqual([ + { + context: { + event: 'tinfoil_proxy_latency', + route: '/tinfoil/v1/chat/completions', + status: 504, + preHandlerMs: 10, + handlerToUpstreamFetchMs: 5, + upstreamFetchToHeadersMs: null, + handlerToOutcomeMs: 35, + }, + message: 'Tinfoil proxy latency', + }, + ]) + }) + + it('logs exactly one 499 line when the client aborts before upstream headers', async () => { + const entries: Array<{ context: TinfoilProxyLatencyLog; message: string }> = [] + const timestamps = [5, 10, 12, 20] + const upstream = createAbortableFetch() + const clientController = new AbortController() + const logger: TinfoilProxyLogger = { + info: (context, message) => entries.push({ context, message }), + } + const app = buildApp({ + fetchFn: upstream.fetchFn, + logger, + nowFn: () => timestamps.shift() ?? 0, + }) + 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() + await response + + expect(entries).toEqual([ + { + context: { + event: 'tinfoil_proxy_latency', + route: '/tinfoil/v1/chat/completions', + status: 499, + preHandlerMs: 5, + handlerToUpstreamFetchMs: 2, + upstreamFetchToHeadersMs: null, + handlerToOutcomeMs: 10, + }, + message: 'Tinfoil proxy latency', + }, + ]) + expect(entries.some(({ context }) => context.status === 500)).toBeFalse() + }) + + it('logs exactly one 500 line when upstream fetch rejects', async () => { + const entries: Array<{ context: TinfoilProxyLatencyLog; message: string }> = [] + const timestamps = [5, 10, 12, 20] + const logger: TinfoilProxyLogger = { + info: (context, message) => entries.push({ context, message }), + } + const fetchFn = (async () => { + throw new Error('Upstream connection failed') + }) as unknown as typeof fetch + const app = buildApp({ + fetchFn, + logger, + nowFn: () => timestamps.shift() ?? 0, + }) + + const response = await app.handle(new Request('http://localhost/tinfoil/v1/models')) + + expect(response.status).toBe(500) + expect(entries).toEqual([ + { + context: { + event: 'tinfoil_proxy_latency', + route: '/tinfoil/v1/models', + status: 500, + preHandlerMs: 5, + handlerToUpstreamFetchMs: 2, + upstreamFetchToHeadersMs: null, + handlerToOutcomeMs: 10, + }, + message: 'Tinfoil proxy latency', + }, + ]) + }) + }) + describe('upstream URL derivation', () => { + it('applies the configured API path prefix to the ATC-assigned enclave origin', async () => { + const assignedEnclaveUrl = 'https://router.inf6.tinfoil.sh' + const upstreamOriginStore = createTinfoilUpstreamOriginStore() + const app = buildApp({ enclaveUrl: 'https://inference.tinfoil.sh/v1', upstreamOriginStore }) + + await drain( + await app.handle( + new Request('http://localhost/tinfoil/chat/completions?stream=true', { + method: 'POST', + headers: { 'X-Tinfoil-Enclave-Url': assignedEnclaveUrl }, + body: 'opaque-bytes', + }), + ), + ) + + const [calledUrl] = mockFetch.mock.calls[0] as [string, RequestInit] + expect(calledUrl).toBe(`${assignedEnclaveUrl}/v1/chat/completions?stream=true`) + expect(upstreamOriginStore.get()).toBe(assignedEnclaveUrl) + + const keepWarmFetch = mock(() => Promise.resolve(makeOkResponse())) + const keepWarm = createTinfoilKeepWarm( + { tinfoilApiKey: testApiKey, tinfoilEnclaveUrl: 'https://inference.tinfoil.sh/v1' }, + { + fetchFn: keepWarmFetch as unknown as typeof fetch, + intervalMs: 100, + logger: { debug: () => undefined }, + upstreamOriginStore, + }, + ) + keepWarm.start() + keepWarm.stop() + + expect(keepWarmFetch).toHaveBeenCalledWith( + `${assignedEnclaveUrl}/v1/models`, + expect.objectContaining({ method: 'GET' }), + ) + }) + + it('avoids a double slash when the assigned enclave origin has a trailing slash', async () => { + const assignedEnclaveUrl = 'https://router.inf6.tinfoil.sh' + const app = buildApp({ enclaveUrl: 'https://inference.tinfoil.sh/v1' }) + + await drain( + await app.handle( + new Request('http://localhost/tinfoil/chat/completions', { + method: 'POST', + headers: { 'X-Tinfoil-Enclave-Url': `${assignedEnclaveUrl}/` }, + body: 'opaque-bytes', + }), + ), + ) + + const [calledUrl] = mockFetch.mock.calls[0] as [string, RequestInit] + expect(calledUrl).toBe(`${assignedEnclaveUrl}/v1/chat/completions`) + }) + + it('does not add a path prefix when the configured enclave URL has none', async () => { + const assignedEnclaveUrl = 'https://router.inf6.tinfoil.sh' + const app = buildApp({ enclaveUrl: 'https://inference.tinfoil.sh' }) + + await drain( + await app.handle( + new Request('http://localhost/tinfoil/chat/completions', { + method: 'POST', + headers: { 'X-Tinfoil-Enclave-Url': assignedEnclaveUrl }, + body: 'opaque-bytes', + }), + ), + ) + + const [calledUrl] = mockFetch.mock.calls[0] as [string, RequestInit] + expect(calledUrl).toBe(`${assignedEnclaveUrl}/chat/completions`) + }) + + it.each(['https://sub.tinfoil.sh', 'https://tinfoil.sh'])( + 'accepts allowlisted assigned enclave URL %s', + async (assignedEnclaveUrl) => { + const app = buildApp() + + const response = await app.handle( + new Request('http://localhost/tinfoil/v1/models', { + headers: { 'X-Tinfoil-Enclave-Url': assignedEnclaveUrl }, + }), + ) + + expect(response.status).toBe(200) + await drain(response) + const [calledUrl] = mockFetch.mock.calls[0] as [string, RequestInit] + expect(calledUrl).toBe(`${assignedEnclaveUrl}/v1/models`) + }, + ) + + it.each(['http://inference.tinfoil.sh', 'https://evil-tinfoil.sh', 'https://tinfoil.sh.evil.com'])( + 'rejects non-allowlisted assigned enclave URL %s', + async (assignedEnclaveUrl) => { + const app = buildApp() + + const response = await app.handle( + new Request('http://localhost/tinfoil/v1/models', { + headers: { 'X-Tinfoil-Enclave-Url': assignedEnclaveUrl }, + }), + ) + + expect(response.status).toBe(400) + expect(response.headers.get('content-type')).toBe('text/plain') + expect(response.headers.get('x-proxy-timing')).not.toBeNull() + expect(await response.text()).toContain('Invalid X-Tinfoil-Enclave-Url') + expect(mockFetch).not.toHaveBeenCalled() + }, + ) + + it('uses the configured enclave URL when no assignment header is provided', async () => { + const fallbackEnclaveUrl = 'https://fallback.tinfoil.sh/v1' + const app = buildApp({ enclaveUrl: fallbackEnclaveUrl }) + + await drain(await app.handle(new Request('http://localhost/tinfoil/models'))) + + const [calledUrl] = mockFetch.mock.calls[0] as [string, RequestInit] + expect(calledUrl).toBe(`${fallbackEnclaveUrl}/models`) + }) + 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. const app = new Elysia({ prefix: '/v2/alt' }).use( diff --git a/backend/src/tinfoil/routes.ts b/backend/src/tinfoil/routes.ts index e318a009c..a19d418ef 100644 --- a/backend/src/tinfoil/routes.ts +++ b/backend/src/tinfoil/routes.ts @@ -8,8 +8,10 @@ import { getSettings } from '@/config/settings' import { safeErrorHandler } from '@/middleware/error-handling' import { capStream } from '@/proxy/streaming' import { filterHeaders } from '@/utils/request' +import { elapsedMs } from '@/utils/timing' import { tinfoilUpstreamIdleTimeoutMessage, tinfoilUpstreamTimeoutMessage } from '@shared/tinfoil-proxy' import { Elysia, type AnyElysia } from 'elysia' +import { tinfoilUpstreamOriginStore, type TinfoilUpstreamOriginStore } from './upstream-origin' const allowedMethods = new Set(['GET', 'POST', 'OPTIONS']) const bodylessMethods = new Set(['GET', 'OPTIONS']) @@ -18,14 +20,81 @@ const defaultUpstreamIdleTimeoutMs = 60_000 const abruptResponseCloseTimeoutSeconds = 1 const upstreamHeadersTimeoutError = new DOMException(tinfoilUpstreamTimeoutMessage, 'TimeoutError') const upstreamIdleTimeoutError = new DOMException(tinfoilUpstreamIdleTimeoutMessage, 'TimeoutError') +const tinfoilEnclaveUrlHeader = 'x-tinfoil-enclave-url' +const tinfoilProxyTimingHeader = 'X-Proxy-Timing' +const serverTimingHeader = 'Server-Timing' + +export type TinfoilProxyLatencyLog = { + event: 'tinfoil_proxy_latency' + route: string + status: number + preHandlerMs: number + handlerToUpstreamFetchMs: number | null + upstreamFetchToHeadersMs: number | null + handlerToOutcomeMs: number +} + +export type TinfoilProxyLogger = { + info: (context: TinfoilProxyLatencyLog, message: string) => void +} const textResponse = (status: number, body: string): Response => new Response(body, { status, headers: { 'Content-Type': 'text/plain' } }) +/** Format available Tinfoil proxy phases using the Server-Timing header syntax. */ +const formatServerTiming = ( + preHandlerMs: number, + handlerToUpstreamFetchMs: number | null, + upstreamFetchToHeadersMs: number | null, +): string => + [ + `pre;dur=${preHandlerMs}`, + handlerToUpstreamFetchMs === null ? null : `fetch;dur=${handlerToUpstreamFetchMs}`, + upstreamFetchToHeadersMs === null ? null : `headers;dur=${upstreamFetchToHeadersMs}`, + ] + .filter((metric): metric is string => metric !== null) + .join(', ') + +/** Check whether an error's cause chain contains the target value. */ +const errorCauseIncludes = (error: unknown, target: unknown): boolean => + error instanceof Error && (error.cause === target || errorCauseIncludes(error.cause, target)) + +/** Identify fetch rejection caused by the downstream client closing its request. */ +const isClientAbortError = (error: unknown, requestSignal: AbortSignal): boolean => + requestSignal.aborted || + (error instanceof Error && error.name === 'AbortError' && errorCauseIncludes(error, requestSignal)) + +/** Resolve an optional ATC-assigned enclave origin, applying the configured API path prefix. */ +const resolveEnclaveUrl = ( + assignedEnclaveUrl: string | null, + fallbackEnclaveUrl: string, + apiPathPrefix: string, +): string | null => { + if (assignedEnclaveUrl === null) { + return fallbackEnclaveUrl + } + + try { + const parsedUrl = new URL(assignedEnclaveUrl) + const isAllowedHostname = parsedUrl.hostname === 'tinfoil.sh' || parsedUrl.hostname.endsWith('.tinfoil.sh') + + if (parsedUrl.protocol !== 'https:' || !isAllowedHostname) { + return null + } + + return `${parsedUrl.origin}${apiPathPrefix}` + } catch { + return null + } +} + /** Forwards HPKE-encrypted bodies to the Tinfoil enclave; injects the bearer key from env. */ export type CreateTinfoilRoutesOptions = { auth: Auth fetchFn?: typeof fetch + logger?: TinfoilProxyLogger + /** Monotonic clock used for latency instrumentation. */ + nowFn?: () => number rateLimit?: AnyElysia /** Override the enclave bearer key. Defaults to `TINFOIL_API_KEY`. */ apiKey?: string @@ -35,40 +104,108 @@ export type CreateTinfoilRoutesOptions = { upstreamHeadersTimeoutMs?: number /** Maximum idle time between upstream response chunks. Defaults to 60 seconds. */ upstreamIdleTimeoutMs?: number + upstreamOriginStore?: Pick } export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { const { auth, rateLimit } = options const fetchFn = options.fetchFn ?? globalThis.fetch + const logger = options.logger + const nowFn = options.nowFn ?? (() => performance.now()) const settings = getSettings() const apiKey = options.apiKey ?? settings.tinfoilApiKey const enclaveUrl = (options.enclaveUrl ?? settings.tinfoilEnclaveUrl).replace(/\/$/, '') + const enclaveApiPathPrefix = new URL(enclaveUrl).pathname.replace(/\/$/, '') const upstreamHeadersTimeoutMs = options.upstreamHeadersTimeoutMs ?? defaultUpstreamHeadersTimeoutMs const upstreamIdleTimeoutMs = options.upstreamIdleTimeoutMs ?? defaultUpstreamIdleTimeoutMs + const upstreamOriginStore = options.upstreamOriginStore ?? tinfoilUpstreamOriginStore const proxyToEnclave = async ( + requestStartedAt: number, request: Request, wildcard: string, server: Bun.Server | null, + setHeaders: Record, ): Promise => { + const handlerStartedAt = nowFn() + const preHandlerMs = elapsedMs(requestStartedAt, handlerStartedAt) + const requestUrl = new URL(request.url) + const route = requestUrl.pathname const method = request.method.toUpperCase() + const recordLatency = ({ + status, + completedAt, + upstreamFetchStartedAt = null, + upstreamHeadersReceivedAt = null, + }: { + status: number + completedAt: number + upstreamFetchStartedAt?: number | null + upstreamHeadersReceivedAt?: number | null + }) => { + const handlerToUpstreamFetchMs = + upstreamFetchStartedAt === null ? null : elapsedMs(handlerStartedAt, upstreamFetchStartedAt) + const upstreamFetchToHeadersMs = + upstreamFetchStartedAt === null || upstreamHeadersReceivedAt === null + ? null + : elapsedMs(upstreamFetchStartedAt, upstreamHeadersReceivedAt) + const latency: TinfoilProxyLatencyLog = { + event: 'tinfoil_proxy_latency', + route, + status, + preHandlerMs, + handlerToUpstreamFetchMs, + upstreamFetchToHeadersMs, + handlerToOutcomeMs: elapsedMs(handlerStartedAt, completedAt), + } + + setHeaders[tinfoilProxyTimingHeader] = + `pre=${preHandlerMs};fetch=${handlerToUpstreamFetchMs ?? 'na'};headers=${upstreamFetchToHeadersMs ?? 'na'}` + setHeaders[serverTimingHeader] = formatServerTiming( + preHandlerMs, + handlerToUpstreamFetchMs, + upstreamFetchToHeadersMs, + ) + logger?.info(latency, 'Tinfoil proxy latency') + } if (!allowedMethods.has(method)) { + recordLatency({ status: 405, completedAt: nowFn() }) return textResponse(405, 'Method not allowed') } if (!apiKey) { + recordLatency({ status: 503, completedAt: nowFn() }) return textResponse(503, 'Tinfoil provider not configured') } + // ATC dynamically assigns the enclave whose HPKE key sealed this body, so forwarding elsewhere guarantees a + // key-config 422; only HTTPS tinfoil.sh hosts are accepted here as the SSRF guard. + const upstreamBaseUrl = resolveEnclaveUrl( + request.headers.get(tinfoilEnclaveUrlHeader), + enclaveUrl, + enclaveApiPathPrefix, + ) + if (upstreamBaseUrl === null) { + recordLatency({ status: 400, completedAt: nowFn() }) + return textResponse(400, 'Invalid X-Tinfoil-Enclave-Url: expected an HTTPS tinfoil.sh host') + } + const subpath = wildcard.startsWith('/') ? wildcard : `/${wildcard}` - const search = new URL(request.url).search - const upstreamUrl = `${enclaveUrl}${subpath}${search}` + const search = requestUrl.search + const upstreamUrl = `${upstreamBaseUrl}${subpath}${search}` + upstreamOriginStore.record(upstreamUrl) const headers = new Headers() request.headers.forEach((value, key) => { const lower = key.toLowerCase() - if (lower === 'authorization' || lower === 'host' || lower === 'cookie' || lower === 'connection') { + if ( + lower === 'authorization' || + lower === 'host' || + lower === 'cookie' || + lower === 'connection' || + lower === tinfoilEnclaveUrlHeader + ) { return } headers.set(key, value) @@ -82,6 +219,7 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { () => upstreamController.abort(upstreamHeadersTimeoutError), upstreamHeadersTimeoutMs, ) + const upstreamFetchStartedAt = nowFn() // Bun-specific fetch options: `duplex: 'half'` enables streaming request // bodies; `decompress: false` keeps the HPKE-encrypted bytes opaque on @@ -96,6 +234,7 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { decompress: false, duplex: 'half', } as RequestInit & { decompress: boolean; duplex: 'half' }) + const upstreamHeadersReceivedAt = nowFn() // Strip upstream CORS headers: the enclave emits a duplicated // `Access-Control-Allow-Credentials: true, true`, which browsers reject @@ -115,12 +254,26 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { }).stream : null - return new Response(responseBody, { + const response = new Response(responseBody, { status: upstream.status, statusText: upstream.statusText, headers: responseHeaders, }) + recordLatency({ + status: upstream.status, + completedAt: upstreamHeadersReceivedAt, + upstreamFetchStartedAt, + upstreamHeadersReceivedAt, + }) + return response } catch (error) { + const completedAt = nowFn() + recordLatency({ + status: isClientAbortError(error, request.signal) ? 499 : error === upstreamHeadersTimeoutError ? 504 : 500, + completedAt, + upstreamFetchStartedAt, + }) + if (error === upstreamHeadersTimeoutError) { return textResponse(504, tinfoilUpstreamTimeoutMessage) } @@ -138,13 +291,37 @@ export const createTinfoilRoutes = (options: CreateTinfoilRoutesOptions) => { // and make `.all()` uncallable / fall back to `any`). return new Elysia({ prefix: '/tinfoil' }) .onError(safeErrorHandler) + .decorate('tinfoilRequestStartedAt', 0) + .onRequest((ctx) => { + // onRequest hooks become app-wide when plugins merge, so avoid timing unrelated routes. + if (!ctx.request.url.includes('/tinfoil/')) { + return + } + ctx.tinfoilRequestStartedAt = nowFn() + }) .use(createAuthMacro(auth)) .guard({ auth: true }, (g) => { if (rateLimit) { return g .use(rateLimit) - .all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? '', ctx.server), { parse: 'none' }) + .all( + '/*', + (ctx) => + proxyToEnclave( + ctx.tinfoilRequestStartedAt, + ctx.request, + ctx.params['*'] ?? '', + ctx.server, + ctx.set.headers, + ), + { parse: 'none' }, + ) } - return g.all('/*', (ctx) => proxyToEnclave(ctx.request, ctx.params['*'] ?? '', ctx.server), { parse: 'none' }) + return g.all( + '/*', + (ctx) => + proxyToEnclave(ctx.tinfoilRequestStartedAt, ctx.request, ctx.params['*'] ?? '', ctx.server, ctx.set.headers), + { parse: 'none' }, + ) }) } diff --git a/backend/src/tinfoil/upstream-origin.ts b/backend/src/tinfoil/upstream-origin.ts new file mode 100644 index 000000000..66f08aea7 --- /dev/null +++ b/backend/src/tinfoil/upstream-origin.ts @@ -0,0 +1,22 @@ +/* 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 type TinfoilUpstreamOriginStore = { + get: () => string | null + record: (upstreamUrl: string) => void +} + +/** Create a store that retains only the latest Tinfoil upstream origin. */ +export const createTinfoilUpstreamOriginStore = (): TinfoilUpstreamOriginStore => { + const state: { latestOrigin: string | null } = { latestOrigin: null } + + return { + get: () => state.latestOrigin, + record: (upstreamUrl) => { + state.latestOrigin = new URL(upstreamUrl).origin + }, + } +} + +export const tinfoilUpstreamOriginStore = createTinfoilUpstreamOriginStore() diff --git a/backend/src/utils/timing.ts b/backend/src/utils/timing.ts new file mode 100644 index 000000000..fa88b71e0 --- /dev/null +++ b/backend/src/utils/timing.ts @@ -0,0 +1,7 @@ +/* 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/. */ + +/** Round a monotonic duration to hundredths of a millisecond. */ +export const elapsedMs = (startedAt: number, completedAt: number): number => + Math.round((completedAt - startedAt) * 100) / 100 diff --git a/deploy/pulumi/src/alb.ts b/deploy/pulumi/src/alb.ts index c84732ade..e663b1176 100644 --- a/deploy/pulumi/src/alb.ts +++ b/deploy/pulumi/src/alb.ts @@ -31,6 +31,8 @@ export const createAlb = ({ name, vpcId, publicSubnetIds, albSgId, hostnames }: loadBalancerType: 'application', securityGroups: [albSgId], subnets: publicSubnetIds, + // Must stay strictly below the backend keep-alive timeout — see backend/src/index.ts (120s). + idleTimeout: 60, tags: { Name: `${name}-alb` }, }) diff --git a/deploy/pulumi/src/shared.ts b/deploy/pulumi/src/shared.ts index c83c6aa51..d932ba16c 100644 --- a/deploy/pulumi/src/shared.ts +++ b/deploy/pulumi/src/shared.ts @@ -189,6 +189,8 @@ export const createSharedStack = (args: SharedStackArgs): SharedStackOutputs => loadBalancerType: 'application', securityGroups: [albSg.id], subnets: publicSubnets.map((s) => s.id), + // Must stay strictly below the backend keep-alive timeout — see backend/src/index.ts (120s). + idleTimeout: 60, tags: { Name: `${name}-alb` }, }) diff --git a/src/acp/built-in-adapter.test.ts b/src/acp/built-in-adapter.test.ts index 351c73442..d9da379df 100644 --- a/src/acp/built-in-adapter.test.ts +++ b/src/acp/built-in-adapter.test.ts @@ -136,7 +136,6 @@ describe('createBuiltInAdapter persistent harness', () => { mcpToolsMetadata: undefined, stableSystemPrompt: 'stable prompt', volatileSystemPrompt: `timestamp ${index + 1}`, - systemPrompt: `stable prompt\n\ntimestamp ${index + 1}`, }), ) const prepareConfig = mock(async () => configs.shift()!) diff --git a/src/acp/built-in-adapter.ts b/src/acp/built-in-adapter.ts index 348b2bbdc..0628d412a 100644 --- a/src/acp/built-in-adapter.ts +++ b/src/acp/built-in-adapter.ts @@ -303,6 +303,7 @@ export const harnessSignature = ( } /** Compose Pi's cacheable prompt prefix while keeping the per-send timestamp last. */ +// Unlike assembleBuiltInModelInput, ACP harness deliberately keeps Pi's single-string system prompt shape. const composeAppHarnessSystemPrompt = (config: AppHarnessSystemPromptConfig): string => `${config.stableSystemPrompt}\n\n${appHarnessEnvironmentPrompt}\n\n${config.volatileSystemPrompt}` diff --git a/src/ai/fetch.test.ts b/src/ai/fetch.test.ts index f125bf35b..7d0f51c5f 100644 --- a/src/ai/fetch.test.ts +++ b/src/ai/fetch.test.ts @@ -3,7 +3,7 @@ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ import { describe, expect, it, mock } from 'bun:test' -import { createPrompt } from '@/ai/prompt' +import { assembleBuiltInModelInput, createPrompt } from '@/ai/prompt' import { defaultSkillResearch, defaultSkillWeather } from '@/defaults/skills' import { fetch as baseFetch } from '@/lib/fetch' import type { MCPClient, NamedMCPClient } from '@/lib/mcp-provider' @@ -14,6 +14,7 @@ import type { Model, Skill } from '@/types' import type { Tool } from 'ai' import { addSkillTool, + buildVolatileSystemNotes, mergeMcpTools, resolveOpenAiCompatConnection, sanitizeToolPrefix, @@ -142,6 +143,43 @@ describe('selectPromptSkillDefinitions', () => { }) }) +describe('buildVolatileSystemNotes', () => { + it('wires volatile notes into the front system block', () => { + const volatileSystemPrompt = 'Current date/time: Friday, July 10, 2026 at 9:00 AM GMT-3' + const currentUserMessage = { role: 'user' as const, content: 'What changed?' } + + const notes = buildVolatileSystemNotes({ + volatileSystemPrompt, + voiceNotes: ['Voice mode is active.'], + skillSystemMessages: ['Follow project style.'], + askResponsesNote: 'Ask responses: concise', + }) + + expect(notes[0]).toBe(volatileSystemPrompt) + expect(notes).toEqual([ + volatileSystemPrompt, + 'Voice mode is active.', + 'Follow project style.', + 'Ask responses: concise', + ]) + const input = assembleBuiltInModelInput( + 'stable prompt', + [ + { role: 'user', content: 'Earlier question' }, + { role: 'assistant', content: 'Earlier answer' }, + currentUserMessage, + ], + notes, + ) + + const messages = [{ role: 'system' as const, content: input.system }, ...input.messages] + + expect(input.system).toBe(`stable prompt\n\n${notes.join('\n\n')}`) + expect(messages.slice(1).every(({ role }) => role !== 'system')).toBeTrue() + expect(input.messages.at(-1)).toEqual(currentUserMessage) + }) +}) + describe('mergeMcpTools', () => { it('prefixes each tool with its sanitized server name', async () => { const render = named('render', async () => ({ list_services: tool('ls'), get_service: tool('gs') })) diff --git a/src/ai/fetch.ts b/src/ai/fetch.ts index e3638a3ca..aa87fea1a 100644 --- a/src/ai/fetch.ts +++ b/src/ai/fetch.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 { createPromptParts } from '@/ai/prompt' +import { assembleBuiltInModelInput, createPromptParts, type BuiltInModelInput } from '@/ai/prompt' import { createTurnBudget, createTurnBudgetExhaustedError, type TurnBudgetConsumer } from '@/ai/retry-budget' import { buildStepOverrides, @@ -475,7 +475,6 @@ export type PreparedAiRequestConfig = { readonly mcpToolsMetadata: UIMessageMetadata['mcpTools'] readonly stableSystemPrompt: string readonly volatileSystemPrompt: string - readonly systemPrompt: string } export type PrepareAiRequestConfigOptions = { @@ -586,10 +585,27 @@ export const prepareAiRequestConfig = async ({ mcpToolsMetadata: merged.mcpTools, stableSystemPrompt: prompt.stablePrompt, volatileSystemPrompt: prompt.volatilePrompt, - systemPrompt: prompt.fullPrompt, } } +/** Order per-send system notes for the trailing half of the system prompt, with date/time first. */ +export const buildVolatileSystemNotes = ({ + volatileSystemPrompt, + voiceNotes, + skillSystemMessages, + askResponsesNote, +}: { + volatileSystemPrompt: string + voiceNotes: readonly string[] + skillSystemMessages: readonly string[] + askResponsesNote: string | null +}): string[] => [ + volatileSystemPrompt, + ...voiceNotes, + ...skillSystemMessages, + ...(askResponsesNote ? [askResponsesNote] : []), +] + /** * Stream one response through the legacy built-in pipeline. * @@ -617,13 +633,22 @@ export const aiFetchStreamingResponse = async ({ // reach this function the user turn is already persisted. const db = getDb() - const { model, profile, supportsTools, sourceCollector, toolset, skills, mcpToolsMetadata, systemPrompt } = - await prepareAiRequestConfig({ - modelId, - mcpClients, - reconnectClient, - httpClient, - }) + const { + model, + profile, + supportsTools, + sourceCollector, + toolset, + skills, + mcpToolsMetadata, + stableSystemPrompt, + volatileSystemPrompt, + } = await prepareAiRequestConfig({ + modelId, + mcpClients, + reconnectClient, + httpClient, + }) if (!supportsTools) { console.log('Model does not support tools, skipping tool setup') } @@ -667,15 +692,15 @@ export const aiFetchStreamingResponse = async ({ /** * Run a single streamText attempt and return the result along with metadata */ - const runStreamText = (inputMessages: Awaited>) => { + const runStreamText = (input: BuiltInModelInput) => { 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, - messages: inputMessages, + system: input.system, + messages: input.messages, tools: supportsTools ? (toolset as ToolSet) : undefined, stopWhen: stepCountIs(maxSteps), providerOptions, @@ -696,7 +721,7 @@ export const aiFetchStreamingResponse = async ({ return buildStepOverrides({ steps, messages: stepMessages, - systemPrompt, + currentSystemPrompt: input.system, profile, maxSteps, nudgeThreshold, @@ -762,7 +787,7 @@ export const aiFetchStreamingResponse = async ({ // a snapshot from when the message was originally typed. // // The composer (`chat-prompt-input.tsx`) uses the same helpers to size - // the context-overflow estimate so the budget and the actual prepend + // the context-overflow estimate so the budget and the actual injection // stay in lockstep. const lastUserText = extractLastUserText(messages) const instructionBySlug = new Map(skills.map(({ name, instruction }) => [name, instruction])) @@ -827,11 +852,16 @@ export const aiFetchStreamingResponse = async ({ ).flat() : [] const askResponsesNote = formatAskResponsesNote(askEntries) - // Voice turns reuse this same send path; when voice is active, prepend the + // Voice turns reuse this same send path; when voice is active, include the // voice self-context so the model knows it's speaking aloud, keeps replies // brief, and answers about itself instead of web-searching its own identity. const voiceNotes = isVoiceModeActive() ? [voiceModeSystemNote] : [] - const systemNotes = [...voiceNotes, ...skillSystemMessages, ...(askResponsesNote ? [askResponsesNote] : [])] + const volatileSystemNotes = buildVolatileSystemNotes({ + volatileSystemPrompt, + voiceNotes, + skillSystemMessages, + askResponsesNote, + }) const stream = createUIMessageStream({ generateId: uuidv7, @@ -844,10 +874,7 @@ export const aiFetchStreamingResponse = async ({ const baseMessages = await convertToModelMessages( hydrateQuotesAsText(await hydrateAttachmentsAsFileParts(messages)), ) - let currentMessages: typeof baseMessages = [ - ...systemNotes.map((content) => ({ role: 'system' as const, content })), - ...baseMessages, - ] + let currentInput = assembleBuiltInModelInput(stableSystemPrompt, baseMessages, volatileSystemNotes) let attemptNumber = 1 let isRetry = false // Track tool calls across ALL attempts — a retry may produce no tool calls @@ -861,7 +888,7 @@ export const aiFetchStreamingResponse = async ({ throw createTurnBudgetExhaustedError() } - const result = runStreamText(currentMessages) + const result = runStreamText(currentInput) const messageMetadata = createMessageMetadata(modelId, sourceCollector, mcpToolsMetadata) // If this is not the last possible attempt, we need to check for empty response @@ -897,11 +924,14 @@ export const aiFetchStreamingResponse = async ({ : activeNudges.retry console.info(`Empty response detected, retrying (attempt ${attemptNumber + 1}/${maxAttempts})...`) - currentMessages = [ - ...currentMessages, - ...response.messages, - { role: 'user' as const, content: retryNudge }, - ] + currentInput = { + ...currentInput, + messages: [ + ...currentInput.messages, + ...response.messages, + { role: 'user' as const, content: retryNudge }, + ], + } isRetry = true attemptNumber++ diff --git a/src/ai/prompt.test.ts b/src/ai/prompt.test.ts index b835419c7..3efc431b0 100644 --- a/src/ai/prompt.test.ts +++ b/src/ai/prompt.test.ts @@ -6,7 +6,7 @@ import { describe, expect, test } from 'bun:test' import type { ModelProfile } from '@/types' import { widgetRegistry } from '@/widgets' import { appHarnessEnvironmentPrompt } from '@shared/agent-core/environment-prompt' -import { createPrompt, createPromptParts, type PromptParams } from './prompt' +import { assembleBuiltInModelInput, createPrompt, createPromptParts, type PromptParams } from './prompt' const createStubProfile = (overrides: Partial = {}): ModelProfile => ({ modelId: 'test-model', @@ -51,6 +51,69 @@ const baseParams: PromptParams = { hasWebTools: false, } +describe('assembleBuiltInModelInput', () => { + const sharedHistory = [ + { role: 'user' as const, content: 'first question' }, + { role: 'assistant' as const, content: 'first answer' }, + ] + + test('keeps all system content at the front and ends on the current user turn', () => { + const firstPrompt = createPromptParts(baseParams, new Date('2026-07-10T12:00:00Z')) + const secondPrompt = createPromptParts(baseParams, new Date('2026-07-10T12:01:00Z')) + const fixedVolatileNotes = ['Voice mode is active.', 'Follow project style.', 'Ask responses: concise'] + const firstUserMessage = { role: 'user' as const, content: 'first follow-up' } + const secondUserMessage = { role: 'user' as const, content: 'second follow-up' } + const first = assembleBuiltInModelInput( + firstPrompt.stablePrompt, + [...sharedHistory, firstUserMessage], + [firstPrompt.volatilePrompt, ...fixedVolatileNotes], + ) + const second = assembleBuiltInModelInput( + secondPrompt.stablePrompt, + [...sharedHistory, secondUserMessage], + [secondPrompt.volatilePrompt, ...fixedVolatileNotes], + ) + const firstMessages = [{ role: 'system' as const, content: first.system }, ...first.messages] + + expect(firstMessages.map(({ role }) => role)).toEqual(['system', 'user', 'assistant', 'user']) + expect(first.system).toContain(firstPrompt.stablePrompt) + expect(first.system).toContain(firstPrompt.volatilePrompt) + expect(first.system).toContain(fixedVolatileNotes.join('\n\n')) + expect(first.system).not.toBe(second.system) + expect(first.messages).toEqual([...sharedHistory, firstUserMessage]) + expect(first.messages.at(-1)).toEqual(firstUserMessage) + }) + + test('keeps volatile notes in the system prompt for sole-message and empty inputs', () => { + const userMessage = { role: 'user' as const, content: 'hello' } + const volatileNote = 'Current date/time: now' + + expect(assembleBuiltInModelInput('stable', [userMessage], [volatileNote])).toEqual({ + system: `stable\n\n${volatileNote}`, + messages: [userMessage], + }) + expect(assembleBuiltInModelInput('stable', [], [volatileNote])).toEqual({ + system: `stable\n\n${volatileNote}`, + messages: [], + }) + }) + + test('produces one front system block for a two-turn conversation', () => { + const currentUserMessage = { role: 'user' as const, content: 'second question' } + const input = assembleBuiltInModelInput( + 'stable', + [...sharedHistory, currentUserMessage], + ['Current date/time: now', 'Voice mode is active.'], + ) + const messages = [{ role: 'system' as const, content: input.system }, ...input.messages] + const firstNonSystemIndex = messages.findIndex(({ role }) => role !== 'system') + + expect(firstNonSystemIndex).toBe(1) + expect(messages.slice(firstNonSystemIndex).every(({ role }) => role !== 'system')).toBeTrue() + expect(messages.at(-1)).toEqual(currentUserMessage) + }) +}) + describe('createPrompt', () => { test('includes model name', () => { const result = createPrompt(baseParams) @@ -190,10 +253,8 @@ describe('createPrompt', () => { expect(result).toContain('Do not emit tags') }) - test('keeps the per-turn timestamp in the suffix (prefix-cache friendly)', () => { + test('keeps the per-turn timestamp at the end of the complete stateless prompt', () => { const result = createPrompt(baseParams) - // The timestamp is the only per-turn-volatile field, so it comes after the - // whole static instruction block to leave a stable cacheable prefix. expect(result.indexOf('Current date/time')).toBeGreaterThan(result.indexOf('# Output Format')) expect(result.indexOf('Current date/time')).toBeGreaterThan(result.indexOf('# Tools')) }) @@ -215,13 +276,15 @@ describe('createPrompt', () => { expect(result.indexOf('# Conversation Style')).toBeLessThan(result.indexOf('Current date/time')) }) - test('separates stable instructions from the volatile timestamp', () => { - const first = createPromptParts(baseParams, new Date('2026-07-10T12:00:00Z')) - const second = createPromptParts(baseParams, new Date('2026-07-10T12:01:00Z')) + test('separates stable instructions from the minute-precision timestamp', () => { + const first = createPromptParts(baseParams, new Date('2026-07-10T12:00:01Z')) + const sameMinute = createPromptParts(baseParams, new Date('2026-07-10T12:00:59Z')) + const nextMinute = createPromptParts(baseParams, new Date('2026-07-10T12:01:00Z')) - expect(first.stablePrompt).toBe(second.stablePrompt) + expect(first.stablePrompt).toBe(nextMinute.stablePrompt) expect(first.stablePrompt).not.toContain('Current date/time') - expect(first.volatilePrompt).not.toBe(second.volatilePrompt) + expect(first.volatilePrompt).toBe(sameMinute.volatilePrompt) + expect(first.volatilePrompt).not.toBe(nextMinute.volatilePrompt) expect(first.fullPrompt).toBe(`${first.stablePrompt}\n\n${first.volatilePrompt}`) }) diff --git a/src/ai/prompt.ts b/src/ai/prompt.ts index e4a70370b..d10c6cb74 100644 --- a/src/ai/prompt.ts +++ b/src/ai/prompt.ts @@ -6,6 +6,7 @@ import { chatPrompt } from '@/ai/prompts/chat' import { webToolsPrompt } from '@/ai/prompts/web-tools' import type { ModelProfile } from '@/types' import { buildFallbackSkillDisclosure, buildSkillListing, type SkillDefinition } from '@shared/agent-core/skills' +import type { ModelMessage } from 'ai' /** Parameters to build the system prompt */ export type PromptParams = { @@ -38,6 +39,21 @@ export type PromptParts = { readonly fullPrompt: string } +export type BuiltInModelInput = { + readonly system: string + readonly messages: ModelMessage[] +} + +/** Combine stable and volatile system content before conversation messages. */ +export const assembleBuiltInModelInput = ( + stableSystemPrompt: string, + baseMessages: readonly ModelMessage[], + volatileSystemNotes: readonly string[], +): BuiltInModelInput => ({ + system: [stableSystemPrompt, ...volatileSystemNotes].join('\n\n'), + messages: [...baseMessages], +}) + /** Build stable assistant instructions separately from per-send date/time. */ export const createPromptParts = ( { @@ -59,8 +75,7 @@ export const createPromptParts = ( // Chat is the only conversation style now (Search/Research ship as default // skills), so its per-model addendum is the only one applied. const chatAddendum = profile?.chatModeAddendum ?? undefined - // The date/time changes every send; it goes in the suffix (see the ordering - // note on the returned template), while the stable context stays in the prefix. + // The date/time changes every send, while the remaining context stays stable. const currentDateTime = `Current date/time: ${currentDate.toLocaleString('en-US', { weekday: 'long', year: 'numeric', @@ -68,7 +83,6 @@ export const createPromptParts = ( day: 'numeric', hour: 'numeric', minute: '2-digit', - second: '2-digit', timeZoneName: 'short', })}` const contextSection = [ @@ -87,15 +101,7 @@ export const createPromptParts = ( // `\(…\)` / `\[…\]`). The chat renderer (src/components/chat/memoized-markdown.tsx) // still normalizes `\(…\)` / `\[…\]` defensively because models drift — the two // are complementary, not redundant; don't drop either side. - // - // Ordering is prefix-cache-friendly: the per-turn-volatile timestamp is the - // ONLY part that changes every send, so it alone is appended LAST (the - // suffix). Everything before it — the static instruction block plus the - // stable `# Context` (user profile + integration status) — forms a prefix - // that prefix-caching backends (vLLM/Tinfoil, OpenAI) reuse across turns. - // Keeping the timestamp at the front would invalidate the cache on every - // send. User-controlled fields stay in `# Context` (not the trailing suffix) - // so settings text can't read as the most-recent instruction. + // Keep user-controlled settings under # Context, never trailing, so they cannot read as the most-recent instruction. const stablePrompt = `You are an executive assistant using the **${modelName}** model. You ALWAYS cite sources with [N] — place each [N] once after the final sentence using that source, with a space before the bracket. Reasoning: low diff --git a/src/ai/step-logic.test.ts b/src/ai/step-logic.test.ts index 5fba0aa45..dfd6dacb3 100644 --- a/src/ai/step-logic.test.ts +++ b/src/ai/step-logic.test.ts @@ -4,6 +4,8 @@ import { describe, expect, test } from 'bun:test' import type { ModelProfile } from '@/types' +import { buildVolatileSystemNotes } from './fetch' +import { assembleBuiltInModelInput } from './prompt' import { buildStepOverrides, extractTextFromMessages, @@ -293,7 +295,7 @@ describe('getNudgeMessagesFromProfile', () => { describe('buildStepOverrides', () => { const baseParams = { - systemPrompt: 'You are an assistant.', + currentSystemPrompt: 'You are an assistant.', profile: null as ModelProfile | null, maxSteps: 20, nudgeThreshold: 6, @@ -360,11 +362,52 @@ describe('buildStepOverrides', () => { }) const result = buildStepOverrides({ ...baseParams, + currentSystemPrompt: 'You are an assistant.\n\nCurrent date/time: Friday, July 31, 2026', profile, steps: toolCallSteps(2), messages: [{ role: 'user', content: 'hello' }], }) - expect(result?.system).toBe('You are an assistant.\nsources') + expect(result?.system).toBe( + 'You are an assistant.\n\nCurrent date/time: Friday, July 31, 2026\nsources', + ) + expect(result?.system).toContain('Current date/time') + }) + + test('preserves every volatile note in citation-reinforcement continuation steps', () => { + const volatileSystemPrompt = 'Current date/time: Friday, July 31, 2026' + const voiceNote = 'Voice mode is active.' + const skillInstructions = 'Follow project style.' + const askResponsesNote = 'User chose concise.' + const volatileSystemNotes = buildVolatileSystemNotes({ + volatileSystemPrompt, + voiceNotes: [voiceNote], + skillSystemMessages: [skillInstructions], + askResponsesNote, + }) + const input = assembleBuiltInModelInput( + 'You are an assistant.', + [{ role: 'user', content: 'hello' }], + volatileSystemNotes, + ) + const citationReinforcementPrompt = '\nsources' + + const result = buildStepOverrides({ + ...baseParams, + currentSystemPrompt: input.system, + profile: createStubProfile({ + citationReinforcementEnabled: 1, + citationReinforcementPrompt, + }), + steps: toolCallSteps(1), + messages: [{ role: 'user', content: 'hello' }], + }) + + expect(result?.system).toBe(input.system + citationReinforcementPrompt) + expect(result?.system).toContain(volatileSystemPrompt) + expect(result?.system).toContain(voiceNote) + expect(result?.system).toContain(skillInstructions) + expect(result?.system).toContain(askResponsesNote) + expect(result?.system).toContain('sources') }) test('no citation reinforcement when disabled', () => { diff --git a/src/ai/step-logic.ts b/src/ai/step-logic.ts index 87b4f1207..f64ee8a49 100644 --- a/src/ai/step-logic.ts +++ b/src/ai/step-logic.ts @@ -87,7 +87,7 @@ export const nudgeMessages: NudgeMessages = { export const buildStepOverrides = ({ steps, messages, - systemPrompt, + currentSystemPrompt, profile, maxSteps, nudgeThreshold, @@ -95,7 +95,7 @@ export const buildStepOverrides = ({ }: { steps: Step[] messages: TMessage[] - systemPrompt: string + currentSystemPrompt: string profile: ModelProfile | null maxSteps: number nudgeThreshold: number @@ -104,7 +104,7 @@ export const buildStepOverrides = ({ const hadToolCallSteps = steps.some((s) => s.finishReason === 'tool-calls') const citationSystem = profile?.citationReinforcementEnabled === 1 && hadToolCallSteps - ? systemPrompt + (profile.citationReinforcementPrompt ?? '') + ? currentSystemPrompt + (profile.citationReinforcementPrompt ?? '') : undefined if (isFinalStep(steps.length, maxSteps)) { diff --git a/src/chats/chat-store.test.ts b/src/chats/chat-store.test.ts index d5139a5d7..fd0833b37 100644 --- a/src/chats/chat-store.test.ts +++ b/src/chats/chat-store.test.ts @@ -20,7 +20,7 @@ import { resetStore, } from '@/test-utils/chat-store-mocks' import type { Model, ThunderboltUIMessage } from '@/types' -import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it } from 'bun:test' +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, mock } from 'bun:test' import { deriveToolKey, findAllowOption, useChatStore } from './chat-store' describe('chat-store', () => { @@ -326,6 +326,49 @@ describe('chat-store', () => { const session = getCurrentSession() expect(session?.selectedModel?.id).toBe('tracked-model') }) + + it('invokes the prewarm wrapper with the selected model (wrapper owns the tinfoil/system guard)', async () => { + const model = createMockModel({ id: 'tinfoil-system-model', provider: 'tinfoil', isSystem: 1 }) + const prewarmSystemModel = mock(async () => {}) + + hydrateStore({ + chatInstance: createMockChatInstanceWithValidation(), + chatThread: null, + id: 'test-id', + mcpClients: [], + models: [model], + selectedModel: model, + triggerData: null, + }) + + await useChatStore.getState().setSelectedModel('test-id', model.id, { prewarmSystemModel }) + + expect(prewarmSystemModel).toHaveBeenCalledTimes(1) + expect(prewarmSystemModel).toHaveBeenCalledWith(model) + }) + + it('does not propagate prewarm errors into model selection', async () => { + const model = createMockModel({ id: 'tinfoil-system-model', provider: 'tinfoil', isSystem: 1 }) + const prewarmFailure = Promise.reject(new Error('attestation failed')) + void prewarmFailure.catch(() => undefined) + const prewarmSystemModel = mock(() => prewarmFailure) + + hydrateStore({ + chatInstance: createMockChatInstanceWithValidation(), + chatThread: null, + id: 'test-id', + mcpClients: [], + models: [model], + selectedModel: model, + triggerData: null, + }) + + await expect( + useChatStore.getState().setSelectedModel('test-id', model.id, { prewarmSystemModel }), + ).resolves.toBeUndefined() + expect(prewarmSystemModel).toHaveBeenCalledTimes(1) + expect(getCurrentSession()?.selectedModel).toBe(model) + }) }) describe('setSelectedAgent', () => { diff --git a/src/chats/chat-store.ts b/src/chats/chat-store.ts index 592c24c34..eb3fe8dc9 100644 --- a/src/chats/chat-store.ts +++ b/src/chats/chat-store.ts @@ -2,6 +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 { prewarmSystemModel } from '@/ai/prewarm-system-model' import { updateSettings } from '@/dal' import { updateChatThread } from '@/dal/chat-threads' import { getDb } from '@/db/database' @@ -76,10 +77,14 @@ type ChatStoreActions = { setPendingPermission(id: string, permission: PendingPermission | null): void resolvePendingPermission(id: string, response: RequestPermissionResponse): void setSelectedAgent(id: string, agent: Agent): Promise - setSelectedModel(id: string, modelId: string | null): Promise + setSelectedModel(id: string, modelId: string | null, deps?: SetSelectedModelDeps): Promise updateSession(id: string, session: Partial>): void } +type SetSelectedModelDeps = { + prewarmSystemModel?: typeof prewarmSystemModel +} + type ChatStore = ChatStoreState & ChatStoreActions const initialState: ChatStoreState = { @@ -233,7 +238,7 @@ export const useChatStore = create()((set, get) => ({ trackEvent('agent_select', { agent: agent.id }) }, - setSelectedModel: async (id, modelId) => { + setSelectedModel: async (id, modelId, deps = {}) => { const { models, sessions } = get() const model = models.find((m) => m.id === modelId) @@ -253,6 +258,10 @@ export const useChatStore = create()((set, get) => ({ set({ sessions: nextSessions }) + // Fire-and-forget: the wrapper no-ops (before any dynamic import) unless + // this is a Tinfoil system model, so the first send finds a warm client. + void (deps.prewarmSystemModel ?? prewarmSystemModel)(model) + const db = getDb() await updateSettings(db, { selected_model: model.id })