From f5770713bfda5e9ae6ad90e6f3f686b9d184d576 Mon Sep 17 00:00:00 2001 From: "rosetta-livekit-bot[bot]" <282703043+rosetta-livekit-bot[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 05:56:06 +0000 Subject: [PATCH 1/3] feat(llm): prewarm provider connections --- .changeset/prewarm-llm-connections.md | 9 +++++++ agents/src/inference/llm.ts | 5 ++++ agents/src/llm/llm.ts | 34 +++++++++++++++++++++++---- agents/src/voice/agent_session.ts | 6 +++++ plugins/anthropic/src/llm.ts | 4 ++++ plugins/google/src/llm.ts | 5 ++++ plugins/mistralai/src/llm.ts | 4 ++++ plugins/openai/src/llm.ts | 4 ++++ 8 files changed, 67 insertions(+), 4 deletions(-) create mode 100644 .changeset/prewarm-llm-connections.md diff --git a/.changeset/prewarm-llm-connections.md b/.changeset/prewarm-llm-connections.md new file mode 100644 index 000000000..cd2c7234a --- /dev/null +++ b/.changeset/prewarm-llm-connections.md @@ -0,0 +1,9 @@ +--- +'@livekit/agents': patch +'@livekit/agents-plugin-anthropic': patch +'@livekit/agents-plugin-google': patch +'@livekit/agents-plugin-mistralai': patch +'@livekit/agents-plugin-openai': patch +--- + +Prewarm LLM provider connections before the first inference request. diff --git a/agents/src/inference/llm.ts b/agents/src/inference/llm.ts index 8925bab81..e3b56e427 100644 --- a/agents/src/inference/llm.ts +++ b/agents/src/inference/llm.ts @@ -261,6 +261,11 @@ export class LLM extends llm.LLM { return 'livekit'; } + protected override async _prewarmImpl(signal: AbortSignal): Promise { + this.client.apiKey = await createAccessToken(this.opts.apiKey, this.opts.apiSecret); + await this.client.models.list({ signal }); + } + static fromModelString(modelString: string): LLM { return new LLM({ model: modelString }); } diff --git a/agents/src/llm/llm.ts b/agents/src/llm/llm.ts index ae9bde31c..2476b0a5e 100644 --- a/agents/src/llm/llm.ts +++ b/agents/src/llm/llm.ts @@ -9,7 +9,7 @@ import { log } from '../log.js'; import type { LLMMetrics } from '../metrics/base.js'; import { recordException, traceTypes, tracer } from '../telemetry/index.js'; import { type APIConnectOptions, intervalForRetry } from '../types.js'; -import { AsyncIterableQueue, delay, startSoon, toError } from '../utils.js'; +import { AsyncIterableQueue, Task, delay, startSoon, toError } from '../utils.js'; import { type ChatContext, type ChatRole, type FunctionCall } from './chat_context.js'; import { type ToolChoice, @@ -65,6 +65,8 @@ export type LLMCallbacks = { }; export abstract class LLM extends (EventEmitter as new () => TypedEmitter) { + #prewarmTask?: Task; + constructor() { super(); } @@ -120,14 +122,38 @@ export abstract class LLM extends (EventEmitter as new () => TypedEmitter { + try { + await this._prewarmImpl(controller.signal); + } catch { + // Prewarm is best-effort and must not affect session startup. + } + }); + } + + protected async _prewarmImpl(_signal: AbortSignal): Promise { + // Providers can override with a cheap request that initializes DNS/TLS/keep-alive state. } async aclose(): Promise { - // Default implementation - subclasses can override + if (this.#prewarmTask) { + await this.#prewarmTask.cancelAndWait(); + } } } diff --git a/agents/src/voice/agent_session.ts b/agents/src/voice/agent_session.ts index a808deebb..bac0635b6 100644 --- a/agents/src/voice/agent_session.ts +++ b/agents/src/voice/agent_session.ts @@ -42,6 +42,7 @@ import type { ToolContextLike, } from '../llm/index.js'; import { ToolContext, toToolContext } from '../llm/index.js'; +import { LLM as BaseLLM } from '../llm/llm.js'; import type { LLMError } from '../llm/llm.js'; import { log } from '../log.js'; import { type ModelUsage, ModelUsageCollector, filterZeroValues } from '../metrics/model_usage.js'; @@ -546,6 +547,11 @@ export class AgentSession< this.llm = llm; } + // Eagerly establish DNS/TLS to the LLM provider so the first inference request is faster. + if (this.llm instanceof BaseLLM) { + this.llm.prewarm(); + } + if (typeof tts === 'string') { this.tts = InferenceTTS.fromModelString(tts); } else { diff --git a/plugins/anthropic/src/llm.ts b/plugins/anthropic/src/llm.ts index 0c85bbf90..782a43c32 100644 --- a/plugins/anthropic/src/llm.ts +++ b/plugins/anthropic/src/llm.ts @@ -89,6 +89,10 @@ export class LLM extends llm.LLM { } } + protected override async _prewarmImpl(signal: AbortSignal): Promise { + await this.#client.models.list({ limit: 1 }, { signal }); + } + /** * Converts a framework ChatContext into Anthropic's message format. * diff --git a/plugins/google/src/llm.ts b/plugins/google/src/llm.ts index 51874b26f..a28f9b3c5 100644 --- a/plugins/google/src/llm.ts +++ b/plugins/google/src/llm.ts @@ -108,6 +108,11 @@ export class LLM extends llm.LLM { return 'Gemini'; } + protected override async _prewarmImpl(signal: AbortSignal): Promise { + // Also fetches auth tokens ahead of time on Vertex AI. + await this.#client.models.list({ config: { pageSize: 1, abortSignal: signal } }); + } + /** * Create a new instance of Google GenAI LLM. * diff --git a/plugins/mistralai/src/llm.ts b/plugins/mistralai/src/llm.ts index d857c65f8..1edaa820e 100644 --- a/plugins/mistralai/src/llm.ts +++ b/plugins/mistralai/src/llm.ts @@ -94,6 +94,10 @@ export class LLM extends llm.LLM { return 'api.mistral.ai'; } + protected override async _prewarmImpl(signal: AbortSignal): Promise { + await this.#client.models.list(undefined, { signal }); + } + updateOptions(opts: { model?: MistralChatModels | string; maxCompletionTokens?: number; diff --git a/plugins/openai/src/llm.ts b/plugins/openai/src/llm.ts index 81d72fea1..88ba87b0f 100644 --- a/plugins/openai/src/llm.ts +++ b/plugins/openai/src/llm.ts @@ -101,6 +101,10 @@ export class LLM extends llm.LLM { } } + protected override async _prewarmImpl(signal: AbortSignal): Promise { + await this.#client.models.list({ signal }); + } + /** * Create a new instance of OpenAI LLM with Azure. * From faa1acb45444f451dd7ec2227bda7f87cc7ba9f7 Mon Sep 17 00:00:00 2001 From: Toubat Date: Fri, 24 Jul 2026 14:28:08 -0700 Subject: [PATCH 2/3] fix(llm): complete prewarm lifecycle parity Prewarm resolved agent-level models and guard close races, with deterministic lifecycle and provider-boundary coverage. Co-authored-by: Cursor --- agents/src/inference/llm.test.ts | 40 ++++++- agents/src/llm/llm.test.ts | 108 +++++++++++++++++- agents/src/llm/llm.ts | 18 ++- agents/src/voice/agent_activity.ts | 4 + .../src/voice/agent_activity_prewarm.test.ts | 60 ++++++++++ plugins/anthropic/src/llm.test.ts | 27 ++++- plugins/google/src/llm_prewarm.test.ts | 81 +++++++++++++ plugins/mistralai/src/llm.test.ts | 27 ++++- plugins/openai/src/llm.test.ts | 27 ++++- 9 files changed, 381 insertions(+), 11 deletions(-) create mode 100644 agents/src/voice/agent_activity_prewarm.test.ts create mode 100644 plugins/google/src/llm_prewarm.test.ts diff --git a/agents/src/inference/llm.test.ts b/agents/src/inference/llm.test.ts index 451f7980d..52d9d116a 100644 --- a/agents/src/inference/llm.test.ts +++ b/agents/src/inference/llm.test.ts @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: 2026 LiveKit, Inc. // // SPDX-License-Identifier: Apache-2.0 -import { beforeAll, describe, expect, it } from 'vitest'; +import { beforeAll, describe, expect, it, vi } from 'vitest'; import * as agents from '../index.js'; import { ChatContext } from '../llm/index.js'; import { initializeLogger } from '../log.js'; @@ -15,6 +15,44 @@ beforeAll(() => { type CapturedHeaders = Record; type CompletionChunk = Record; +describe('inference.LLM prewarm', () => { + it('refreshes the access token before listing models and forwards cancellation', async () => { + const llm = new LLM({ + model: 'openai/gpt-4o-mini', + apiKey: 'test-key', + apiSecret: 'test-secret', + baseURL: 'https://example.livekit.cloud', + }); + const internal = llm as unknown as { + client: { + apiKey: string; + models: { + list: (options: { signal?: AbortSignal }) => Promise; + }; + }; + }; + let apiKeyAtList = ''; + let prewarmSignal: AbortSignal | undefined; + const modelsList = vi.fn(async (options: { signal?: AbortSignal }) => { + apiKeyAtList = internal.client.apiKey; + prewarmSignal = options.signal; + }); + internal.client.models.list = modelsList; + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(modelsList).toHaveBeenCalledTimes(1); + expect(apiKeyAtList).not.toBe('placeholder'); + expect(apiKeyAtList.split('.')).toHaveLength(3); + expect(prewarmSignal).toBeInstanceOf(AbortSignal); + expect(prewarmSignal?.aborted).toBe(false); + + await llm.aclose(); + expect(prewarmSignal?.aborted).toBe(true); + }); +}); + /** * Build an LLM, stub its OpenAI client's chat.completions.create, start a chat * stream with the given per-call value, drain the stream, and return the headers diff --git a/agents/src/llm/llm.test.ts b/agents/src/llm/llm.test.ts index b921e62c2..887515331 100644 --- a/agents/src/llm/llm.test.ts +++ b/agents/src/llm/llm.test.ts @@ -1,10 +1,10 @@ // SPDX-FileCopyrightText: 2026 LiveKit, Inc. // // SPDX-License-Identifier: Apache-2.0 -import { beforeAll, describe, expect, it } from 'vitest'; +import { beforeAll, describe, expect, it, vi } from 'vitest'; import { initializeLogger } from '../log.js'; import { type APIConnectOptions, DEFAULT_API_CONNECT_OPTIONS } from '../types.js'; -import { delay } from '../utils.js'; +import { Future, Task, delay } from '../utils.js'; import { ChatContext, FunctionCall } from './chat_context.js'; import { type ChatChunk, LLM, LLMStream } from './llm.js'; import type { ToolChoice, ToolCtxInput } from './tool_context.js'; @@ -59,6 +59,110 @@ class MockLLM extends LLM { } } +class PrewarmLLM extends MockLLM { + constructor(private readonly prewarmImpl: (signal: AbortSignal) => Promise) { + super([]); + } + + protected override _prewarmImpl(signal: AbortSignal): Promise { + return this.prewarmImpl(signal); + } +} + +const waitForTasks = () => new Promise((resolve) => setImmediate(resolve)); + +describe('LLM prewarm lifecycle', () => { + it('is a no-op when the provider does not override _prewarmImpl', async () => { + const taskFrom = vi.spyOn(Task, 'from'); + const llm = new MockLLM([]); + + try { + llm.prewarm(); + await llm.aclose(); + + expect(taskFrom).not.toHaveBeenCalled(); + } finally { + taskFrom.mockRestore(); + } + }); + + it('schedules fire-and-forget work without surfacing provider rejection', async () => { + const unhandledRejection = vi.fn(); + process.on('unhandledRejection', unhandledRejection); + const llm = new PrewarmLLM(async () => { + throw new Error('provider unavailable'); + }); + + try { + expect(() => llm.prewarm()).not.toThrow(); + await waitForTasks(); + + expect(unhandledRejection).not.toHaveBeenCalled(); + } finally { + process.off('unhandledRejection', unhandledRejection); + } + }); + + it('invokes the provider hook once across repeated calls, including after failure', async () => { + const prewarmImpl = vi.fn(async () => { + throw new Error('provider unavailable'); + }); + const llm = new PrewarmLLM(prewarmImpl); + + llm.prewarm(); + llm.prewarm(); + await waitForTasks(); + llm.prewarm(); + await waitForTasks(); + + expect(prewarmImpl).toHaveBeenCalledTimes(1); + }); + + it('aborts an in-flight prewarm and waits for provider cleanup during close', async () => { + const started = new Future(); + let signal: AbortSignal | undefined; + let cleanedUp = false; + const llm = new PrewarmLLM( + (prewarmSignal) => + new Promise((resolve) => { + signal = prewarmSignal; + started.resolve(); + prewarmSignal.addEventListener( + 'abort', + () => { + setImmediate(() => { + cleanedUp = true; + resolve(); + }); + }, + { once: true }, + ); + }), + ); + + llm.prewarm(); + await started.await; + const closing = llm.aclose(); + + expect(signal?.aborted).toBe(true); + expect(cleanedUp).toBe(false); + + await closing; + expect(cleanedUp).toBe(true); + }); + + it('does not start prewarm work after close wins the lifecycle race', async () => { + const prewarmImpl = vi.fn(async () => {}); + const llm = new PrewarmLLM(prewarmImpl); + + await llm.aclose(); + llm.prewarm(); + await waitForTasks(); + + expect(prewarmImpl).not.toHaveBeenCalled(); + }); +}); + describe('LLMStream.collect', () => { beforeAll(() => { initializeLogger({ pretty: false }); diff --git a/agents/src/llm/llm.ts b/agents/src/llm/llm.ts index 2476b0a5e..3c1ac7bac 100644 --- a/agents/src/llm/llm.ts +++ b/agents/src/llm/llm.ts @@ -66,6 +66,7 @@ export type LLMCallbacks = { export abstract class LLM extends (EventEmitter as new () => TypedEmitter) { #prewarmTask?: Task; + #closed = false; constructor() { super(); @@ -126,10 +127,11 @@ export abstract class LLM extends (EventEmitter as new () => TypedEmitter TypedEmitter { - // Providers can override with a cheap request that initializes DNS/TLS/keep-alive state. - } + /** + * Performs a provider-specific, token-free request that initializes DNS, TLS, authentication, + * and keep-alive state. + * + * @remarks Exceptions are swallowed by {@link prewarm}. Implementations must honor `signal` so + * {@link aclose} can cancel and await in-flight work. + */ + protected async _prewarmImpl(_signal: AbortSignal): Promise {} async aclose(): Promise { + this.#closed = true; if (this.#prewarmTask) { await this.#prewarmTask.cancelAndWait(); } diff --git a/agents/src/voice/agent_activity.ts b/agents/src/voice/agent_activity.ts index 23d8b4ea9..54ce29b9e 100644 --- a/agents/src/voice/agent_activity.ts +++ b/agents/src/voice/agent_activity.ts @@ -478,6 +478,10 @@ export class AgentActivity implements RecognitionHooks { async start(options?: { reuseResources?: ReusableResources }): Promise { const unlock = await this.lock.lock(); try { + if (this.llm instanceof LLM) { + this.llm.prewarm(); + } + await this._startSession({ spanName: 'start_agent_activity', runOnEnter: true, diff --git a/agents/src/voice/agent_activity_prewarm.test.ts b/agents/src/voice/agent_activity_prewarm.test.ts new file mode 100644 index 000000000..26757631f --- /dev/null +++ b/agents/src/voice/agent_activity_prewarm.test.ts @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: 2026 LiveKit, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +import { describe, expect, it } from 'vitest'; +import { Agent } from './agent.js'; +import { AgentSession } from './agent_session.js'; +import { FakeLLM, type FakeLLMResponse } from './testing/fake_llm.js'; + +class TrackingLLM extends FakeLLM { + prewarmCalls = 0; + + constructor( + responses: FakeLLMResponse[] = [], + private readonly events: string[] = [], + ) { + super(responses); + } + + override chat(options: Parameters[0]) { + this.events.push('chat'); + return super.chat(options); + } + + protected override async _prewarmImpl(_signal: AbortSignal): Promise { + this.prewarmCalls += 1; + this.events.push('prewarm'); + } +} + +describe('AgentActivity LLM prewarm', () => { + it('prewarms an agent-level LLM exactly once before its first inference', async () => { + const sessionLLM = new TrackingLLM(); + const agentEvents: string[] = []; + const agentLLM = new TrackingLLM([{ input: 'hello', content: 'hi' }], agentEvents); + const session = new AgentSession({ + llm: sessionLLM, + vad: null, + turnHandling: { turnDetection: null }, + }); + + expect(sessionLLM.prewarmCalls).toBe(1); + expect(agentLLM.prewarmCalls).toBe(0); + + try { + await session.start({ + agent: new Agent({ instructions: '', llm: agentLLM }), + }); + + expect(agentLLM.prewarmCalls).toBe(1); + expect(sessionLLM.prewarmCalls).toBe(1); + + await session.run({ userInput: 'hello' }).wait(); + + expect(agentEvents).toEqual(['prewarm', 'chat']); + expect(agentLLM.prewarmCalls).toBe(1); + } finally { + await session.close(); + } + }); +}); diff --git a/plugins/anthropic/src/llm.test.ts b/plugins/anthropic/src/llm.test.ts index 3aa1eaece..3e125f4e6 100644 --- a/plugins/anthropic/src/llm.test.ts +++ b/plugins/anthropic/src/llm.test.ts @@ -3,10 +3,35 @@ // SPDX-License-Identifier: Apache-2.0 import type Anthropic from '@anthropic-ai/sdk'; import { llm } from '@livekit/agents'; -import { describe, expect, it } from 'vitest'; +import { describe, expect, it, vi } from 'vitest'; import { z } from 'zod'; import { LLM } from './llm.js'; +describe('Anthropic LLM prewarm', () => { + it('lists one model with the prewarm cancellation signal', async () => { + let prewarmSignal: AbortSignal | undefined; + const modelsList = vi.fn( + async (_params: { limit: number }, options: { signal?: AbortSignal }) => { + prewarmSignal = options.signal; + }, + ); + const client = { + baseURL: 'https://api.anthropic.test', + models: { list: modelsList }, + } as unknown as Anthropic; + const llm = new LLM({ model: 'claude-sonnet-4-6', client }); + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(modelsList).toHaveBeenCalledWith({ limit: 1 }, { signal: expect.any(AbortSignal) }); + expect(prewarmSignal?.aborted).toBe(false); + + await llm.aclose(); + expect(prewarmSignal?.aborted).toBe(true); + }); +}); + function messageStartEvent(): Anthropic.MessageStreamEvent { return { type: 'message_start', diff --git a/plugins/google/src/llm_prewarm.test.ts b/plugins/google/src/llm_prewarm.test.ts new file mode 100644 index 000000000..9dd640678 --- /dev/null +++ b/plugins/google/src/llm_prewarm.test.ts @@ -0,0 +1,81 @@ +// SPDX-FileCopyrightText: 2026 LiveKit, Inc. +// +// SPDX-License-Identifier: Apache-2.0 +import type * as googleGenai from '@google/genai'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { LLM } from './llm.js'; + +const googleMocks = vi.hoisted(() => ({ + constructorOptions: [] as unknown[], + modelLists: [] as Array<{ mock: { calls: unknown[][] } }>, +})); + +vi.mock('@google/genai', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + GoogleGenAI: class { + models: { list: ReturnType }; + + constructor(options: unknown) { + const list = vi.fn(async () => {}); + this.models = { list }; + googleMocks.constructorOptions.push(options); + googleMocks.modelLists.push(list); + } + }, + }; +}); + +describe('Google LLM prewarm', () => { + beforeEach(() => { + googleMocks.constructorOptions.length = 0; + googleMocks.modelLists.length = 0; + }); + + it('lists one model with the prewarm cancellation signal', async () => { + const llm = new LLM({ model: 'gemini-2.5-flash', apiKey: 'test-key' }); + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + const modelsList = googleMocks.modelLists[0]!; + expect(modelsList.mock.calls).toEqual([ + [{ config: { pageSize: 1, abortSignal: expect.any(AbortSignal) } }], + ]); + const signal = ( + modelsList.mock.calls[0]![0] as { + config: { abortSignal: AbortSignal }; + } + ).config.abortSignal; + expect(signal.aborted).toBe(false); + + await llm.aclose(); + expect(signal.aborted).toBe(true); + }); + + it('uses the same models request on the Vertex auth warmup path', async () => { + const llm = new LLM({ + model: 'gemini-2.5-flash', + vertexai: true, + project: 'test-project', + location: 'test-location', + }); + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(googleMocks.constructorOptions).toEqual([ + { + vertexai: true, + project: 'test-project', + location: 'test-location', + }, + ]); + expect(googleMocks.modelLists[0]!.mock.calls).toEqual([ + [{ config: { pageSize: 1, abortSignal: expect.any(AbortSignal) } }], + ]); + + await llm.aclose(); + }); +}); diff --git a/plugins/mistralai/src/llm.test.ts b/plugins/mistralai/src/llm.test.ts index dd7e1f016..d0162b2d6 100644 --- a/plugins/mistralai/src/llm.test.ts +++ b/plugins/mistralai/src/llm.test.ts @@ -2,9 +2,34 @@ // // SPDX-License-Identifier: Apache-2.0 import { llm as llmTest } from '@livekit/agents-plugins-test'; -import { describe, it } from 'vitest'; +import type { Mistral } from '@mistralai/mistralai'; +import { describe, expect, it, vi } from 'vitest'; import { LLM } from './llm.js'; +describe('Mistral LLM prewarm', () => { + it('lists models with the prewarm cancellation signal', async () => { + let prewarmSignal: AbortSignal | undefined; + const modelsList = vi.fn(async (_request: undefined, options: { signal?: AbortSignal }) => { + prewarmSignal = options.signal; + }); + const client = { + models: { list: modelsList }, + } as unknown as Mistral; + const llm = new LLM({ client }); + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(modelsList).toHaveBeenCalledWith(undefined, { + signal: expect.any(AbortSignal), + }); + expect(prewarmSignal?.aborted).toBe(false); + + await llm.aclose(); + expect(prewarmSignal?.aborted).toBe(true); + }); +}); + const hasMistralApiKey = Boolean(process.env.MISTRAL_API_KEY); if (hasMistralApiKey) { diff --git a/plugins/openai/src/llm.test.ts b/plugins/openai/src/llm.test.ts index 63e899bdd..18afc2dbc 100644 --- a/plugins/openai/src/llm.test.ts +++ b/plugins/openai/src/llm.test.ts @@ -2,9 +2,34 @@ // // SPDX-License-Identifier: Apache-2.0 import { llm, llmStrict } from '@livekit/agents-plugins-test'; -import { describe, it } from 'vitest'; +import type OpenAI from 'openai'; +import { describe, expect, it, vi } from 'vitest'; import { LLM } from './llm.js'; +describe('OpenAI LLM prewarm', () => { + it('lists models with the prewarm cancellation signal', async () => { + let prewarmSignal: AbortSignal | undefined; + const modelsList = vi.fn(async (options: { signal?: AbortSignal }) => { + prewarmSignal = options.signal; + }); + const client = { + baseURL: 'https://api.openai.test/v1', + models: { list: modelsList }, + } as unknown as OpenAI; + const llm = new LLM({ model: 'gpt-4.1', client }); + + llm.prewarm(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(modelsList).toHaveBeenCalledTimes(1); + expect(prewarmSignal).toBeInstanceOf(AbortSignal); + expect(prewarmSignal?.aborted).toBe(false); + + await llm.aclose(); + expect(prewarmSignal?.aborted).toBe(true); + }); +}); + const hasOpenAIApiKey = Boolean(process.env.OPENAI_API_KEY); if (hasOpenAIApiKey) { From 795981ec09298f80618a0b4e035ade1228c9d97d Mon Sep 17 00:00:00 2001 From: Toubat Date: Fri, 24 Jul 2026 14:36:48 -0700 Subject: [PATCH 3/3] test(google): cover real Vertex prewarm auth Exercise the installed SDK through token and network boundaries so auth ordering and cancellation cannot be hidden by a full-client mock. Co-authored-by: Cursor --- plugins/google/src/llm_prewarm.test.ts | 136 ++++++++++++++----------- 1 file changed, 78 insertions(+), 58 deletions(-) diff --git a/plugins/google/src/llm_prewarm.test.ts b/plugins/google/src/llm_prewarm.test.ts index 9dd640678..00cc63ec5 100644 --- a/plugins/google/src/llm_prewarm.test.ts +++ b/plugins/google/src/llm_prewarm.test.ts @@ -1,60 +1,75 @@ // SPDX-FileCopyrightText: 2026 LiveKit, Inc. // // SPDX-License-Identifier: Apache-2.0 -import type * as googleGenai from '@google/genai'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { GoogleAuth } from 'google-auth-library'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import { LLM } from './llm.js'; -const googleMocks = vi.hoisted(() => ({ - constructorOptions: [] as unknown[], - modelLists: [] as Array<{ mock: { calls: unknown[][] } }>, -})); - -vi.mock('@google/genai', async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - GoogleGenAI: class { - models: { list: ReturnType }; - - constructor(options: unknown) { - const list = vi.fn(async () => {}); - this.models = { list }; - googleMocks.constructorOptions.push(options); - googleMocks.modelLists.push(list); - } - }, - }; -}); +function withTimeout(promise: Promise, timeoutMs = 500): Promise { + let timer: ReturnType; + return Promise.race([ + promise, + new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error('timed out waiting for SDK request')), timeoutMs); + }), + ]).finally(() => clearTimeout(timer)); +} describe('Google LLM prewarm', () => { - beforeEach(() => { - googleMocks.constructorOptions.length = 0; - googleMocks.modelLists.length = 0; + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + vi.unstubAllEnvs(); }); - it('lists one model with the prewarm cancellation signal', async () => { - const llm = new LLM({ model: 'gemini-2.5-flash', apiKey: 'test-key' }); + it('acquires a Vertex token before the cancellable models-list request', async () => { + vi.stubEnv('GOOGLE_API_KEY', ''); + vi.stubEnv('GOOGLE_APPLICATION_CREDENTIALS', ''); + vi.stubEnv('GOOGLE_CLOUD_PROJECT', ''); + vi.stubEnv('GOOGLE_CLOUD_LOCATION', ''); + vi.stubEnv('GOOGLE_GENAI_USE_VERTEXAI', ''); - llm.prewarm(); - await new Promise((resolve) => setImmediate(resolve)); - - const modelsList = googleMocks.modelLists[0]!; - expect(modelsList.mock.calls).toEqual([ - [{ config: { pageSize: 1, abortSignal: expect.any(AbortSignal) } }], - ]); - const signal = ( - modelsList.mock.calls[0]![0] as { - config: { abortSignal: AbortSignal }; - } - ).config.abortSignal; - expect(signal.aborted).toBe(false); + const events: string[] = []; + const getAccessToken = vi.fn(async () => { + events.push('token'); + return { token: 'test-access-token' }; + }); + const authClient = { + getAccessToken, + getRequestHeaders: vi.fn(async () => { + const { token } = await getAccessToken(); + return new Headers({ authorization: `Bearer ${token}` }); + }), + }; + const getClient = vi + .spyOn(GoogleAuth.prototype, 'getClient') + .mockResolvedValue(authClient as never); - await llm.aclose(); - expect(signal.aborted).toBe(true); - }); + let fetchSignal: AbortSignal | undefined; + let resolveFetchStarted!: () => void; + const fetchStarted = new Promise((resolve) => { + resolveFetchStarted = resolve; + }); + const fetchMock = vi.fn( + async (_input: string | URL | Request, init?: RequestInit): Promise => { + events.push('models-list'); + fetchSignal = init?.signal instanceof AbortSignal ? init.signal : undefined; + resolveFetchStarted(); + return await new Promise((_resolve, reject) => { + if (fetchSignal?.aborted) { + reject(new DOMException('aborted', 'AbortError')); + return; + } + fetchSignal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + }, + ); + vi.stubGlobal('fetch', fetchMock); - it('uses the same models request on the Vertex auth warmup path', async () => { const llm = new LLM({ model: 'gemini-2.5-flash', vertexai: true, @@ -62,20 +77,25 @@ describe('Google LLM prewarm', () => { location: 'test-location', }); - llm.prewarm(); - await new Promise((resolve) => setImmediate(resolve)); + try { + llm.prewarm(); + await withTimeout(fetchStarted); - expect(googleMocks.constructorOptions).toEqual([ - { - vertexai: true, - project: 'test-project', - location: 'test-location', - }, - ]); - expect(googleMocks.modelLists[0]!.mock.calls).toEqual([ - [{ config: { pageSize: 1, abortSignal: expect.any(AbortSignal) } }], - ]); + expect(getClient).toHaveBeenCalledTimes(1); + expect(getAccessToken).toHaveBeenCalledTimes(1); + expect(events).toEqual(['token', 'models-list']); + expect(fetchMock).toHaveBeenCalledTimes(1); + + const [input, init] = fetchMock.mock.calls[0]!; + const url = new URL(String(input)); + expect(url.pathname).toContain('/publishers/google/models'); + expect(url.searchParams.get('pageSize')).toBe('1'); + expect(new Headers(init?.headers).get('authorization')).toBe('Bearer test-access-token'); + expect(fetchSignal?.aborted).toBe(false); + } finally { + await llm.aclose(); + } - await llm.aclose(); + expect(fetchSignal?.aborted).toBe(true); }); });