diff --git a/.changeset/session-scoped-mock-tools.md b/.changeset/session-scoped-mock-tools.md new file mode 100644 index 000000000..b853974c8 --- /dev/null +++ b/.changeset/session-scoped-mock-tools.md @@ -0,0 +1,5 @@ +--- +'@livekit/agents': patch +--- + +Add session-scoped `mockTools(agent, mocks, session)` to `voice.testing`: assigns a mock set for an Agent type on a specific session, effective for the session's lifetime. Context-scoped `withMockTools` mocks take precedence when both are active. diff --git a/agents/src/voice/generation.ts b/agents/src/voice/generation.ts index 3d6a4c831..cc60052d7 100644 --- a/agents/src/voice/generation.ts +++ b/agents/src/voice/generation.ts @@ -1202,7 +1202,7 @@ export function performToolExecutions({ { functionCall: toolCall, speechHandle }, async () => { const runCtx = new RunContext(session, speechHandle, toolCall); - const mock = getMockTool(session.currentAgent, toolCall.name); + const mock = getMockTool(session.currentAgent, toolCall.name, session); const toolToExecute = mock ? { ...tool, diff --git a/agents/src/voice/testing/index.ts b/agents/src/voice/testing/index.ts index 435b04c3d..d185d0a35 100644 --- a/agents/src/voice/testing/index.ts +++ b/agents/src/voice/testing/index.ts @@ -30,6 +30,7 @@ export { MessageAssert, RunAssert, RunResult, + mockTools, withMockTools, type MockToolFn, type MockToolsMap, diff --git a/agents/src/voice/testing/run_result.test.ts b/agents/src/voice/testing/run_result.test.ts index 3aa01aba0..aaf17da53 100644 --- a/agents/src/voice/testing/run_result.test.ts +++ b/agents/src/voice/testing/run_result.test.ts @@ -9,7 +9,7 @@ import { ToolContext, tool } from '../../llm/tool_context.js'; import { Agent } from '../agent.js'; import { performToolExecutions } from '../generation.js'; import { SpeechHandle } from '../speech_handle.js'; -import { activeMockTools, withMockTools } from './run_result.js'; +import { activeMockTools, getMockTool, mockTools, withMockTools } from './run_result.js'; class AgentA extends Agent { constructor() { @@ -180,3 +180,122 @@ describe('withMockTools', () => { expect(output.output[0]?.toolCallOutput?.isError).toBe(true); }); }); + +describe('mockTools (session-scoped)', () => { + // mockTools only uses the session as a WeakMap key; a stub is sufficient, + // matching the session stubs used by the performToolExecutions tests above. + const makeSession = () => ({ currentAgent: undefined }) as never; + + it('registers mocks for a session and resolves them via getMockTool', () => { + const session = makeSession(); + const agent = new AgentA(); + const mock = () => 'session-mocked'; + + mockTools(AgentA, { tool1: mock }, session); + + expect(getMockTool(agent, 'tool1', session)).toBe(mock); + // without the session, session-scoped mocks are invisible + expect(getMockTool(agent, 'tool1')).toBeUndefined(); + }); + + it('isolates mock sets per session', () => { + const sessionA = makeSession(); + const sessionB = makeSession(); + const agent = new AgentA(); + + mockTools(AgentA, { tool1: () => 'a' }, sessionA); + + expect(getMockTool(agent, 'tool1', sessionA)).toBeDefined(); + expect(getMockTool(agent, 'tool1', sessionB)).toBeUndefined(); + }); + + it('replaces the mock set on re-registration and removes it with an empty record', () => { + const session = makeSession(); + const agent = new AgentA(); + const second = () => 'second'; + + mockTools(AgentA, { tool1: () => 'first' }, session); + mockTools(AgentA, { tool2: second }, session); + + // full replacement: tool1 is gone, tool2 is present + expect(getMockTool(agent, 'tool1', session)).toBeUndefined(); + expect(getMockTool(agent, 'tool2', session)).toBe(second); + + mockTools(AgentA, {}, session); + expect(getMockTool(agent, 'tool2', session)).toBeUndefined(); + }); + + it('context-scoped mocks take precedence over session-scoped ones', () => { + const session = makeSession(); + const agent = new AgentA(); + const sessionMock = () => 'session'; + const contextMock = () => 'context'; + + mockTools(AgentA, { tool1: sessionMock }, session); + + { + using _mock = withMockTools(AgentA, { tool1: contextMock }); + expect(getMockTool(agent, 'tool1', session)).toBe(contextMock); + } + + // context scope ended: session mocks apply again + expect(getMockTool(agent, 'tool1', session)).toBe(sessionMock); + }); + + it('falls through to session mocks for tools the context set does not cover', () => { + const session = makeSession(); + const agent = new AgentA(); + const sessionMock = () => 'session'; + + mockTools(AgentA, { getWeather: sessionMock }, session); + + { + // context mocks a *different* tool of the same agent type — the + // session-scoped mock for getWeather must remain active + using _mock = withMockTools(AgentA, { orderItem: () => 'ordered' }); + expect(getMockTool(agent, 'getWeather', session)).toBe(sessionMock); + expect(getMockTool(agent, 'orderItem', session)).toBeDefined(); + } + }); + + it('routes performToolExecutions to a session-scoped mock', async () => { + let realCalled = false; + const realTool = tool({ + name: 'greet', + description: 'real', + parameters: z.object({ name: z.string() }), + execute: async () => { + realCalled = true; + return 'real'; + }, + }); + const toolCtx = new ToolContext([realTool]); + const speechHandle = SpeechHandle.create({ allowInterruptions: false }); + const session = { currentAgent: new AgentA() } as never; + + mockTools(AgentA, { greet: () => 'session-mocked' }, session); + + const controller = new AbortController(); + const call = FunctionCall.create({ + callId: 'call_session_mock', + name: 'greet', + args: JSON.stringify({ name: 'world' }), + }); + const stream = new ReadableStream({ + start(c) { + c.enqueue(call); + c.close(); + }, + }); + const [task, output] = performToolExecutions({ + session, + speechHandle, + toolCtx, + toolCallStream: stream, + controller, + }); + await task.result; + expect(realCalled).toBe(false); + expect(output.output[0]?.rawOutput).toBe('session-mocked'); + }); +}); diff --git a/agents/src/voice/testing/run_result.ts b/agents/src/voice/testing/run_result.ts index 4033c9436..1621e02f0 100644 --- a/agents/src/voice/testing/run_result.ts +++ b/agents/src/voice/testing/run_result.ts @@ -9,6 +9,7 @@ import { tool } from '../../llm/tool_context.js'; import type { Task } from '../../utils.js'; import { Future } from '../../utils.js'; import type { Agent } from '../agent.js'; +import type { AgentSession } from '../agent_session.js'; import { type SpeechHandle, isSpeechHandle } from '../speech_handle.js'; import { type AgentHandoffAssertOptions, @@ -966,13 +967,25 @@ export type MockToolsMap = Map>; /** @internal */ export let activeMockTools: MockToolsMap | undefined; -/** @internal */ -export function getMockTool(agent: Agent, toolName: string): MockToolFn | undefined { - if (!activeMockTools) return undefined; +/** @internal – session-scoped mock sets registered via {@link mockTools}. */ +const sessionMockTools = new WeakMap(); - for (const [agentConstructor, mocks] of activeMockTools) { - if (agent.constructor === agentConstructor) { - return mocks[toolName]; +/** @internal */ +export function getMockTool( + agent: Agent, + toolName: string, + session?: AgentSession, +): MockToolFn | undefined { + // Context-scoped mocks (withMockTools) take precedence over session-scoped + // ones (mockTools), per tool: a scope whose mock set lacks the tool falls + // through to the next scope instead of unmocking it. + const scopes = [activeMockTools, session ? sessionMockTools.get(session) : undefined]; + for (const scope of scopes) { + if (!scope) continue; + for (const [agentConstructor, mocks] of scope) { + if (agent.constructor === agentConstructor && toolName in mocks) { + return mocks[toolName]; + } } } return undefined; @@ -1020,6 +1033,43 @@ export function withMockTools( }; } +/** + * Assign a set of mock tool callables to a specific Agent type on a session, + * effective immediately and for the session's lifetime. + * + * Mocks intercept tool *execution* only; the LLM keeps seeing the real tool + * schemas. Call again to replace the mock set for that Agent type, or pass an + * empty record to remove all mocks for it. When both this and + * {@link withMockTools} are active, the context-scoped mocks take precedence + * per tool; session mocks for tools the context set does not cover still apply. + * + * Mirrors the session form of Python's `mock_tools` (livekit/agents#6080). + * + * @param agent - The Agent constructor whose tools should be mocked. + * @param mocks - A record mapping tool name to a mock implementation. + * @param session - The session the mock set is bound to. + * + * @example + * ```typescript + * mockTools(DriveThruAgent, { getWeather: () => 'sunny' }, session); + * + * const result = await session.run({ userInput: "What's the weather?" }); + * result.expect.containsFunctionCallOutput({ output: 'sunny' }); + * ``` + */ +export function mockTools( + agent: AgentConstructor, + mocks: Record, + session: AgentSession, +): void { + let map = sessionMockTools.get(session); + if (!map) { + map = new Map(); + sessionMockTools.set(session, map); + } + map.set(agent, { ...mocks }); +} + /** * Format events for debug output, optionally marking a selected index. */