import { describe, expect, it, vi } from 'vitest'; import { ModelToolRegistry, } from '../../electron/coding-runtime/pi/model-tools/model-tool-registry'; import type { PiProviderDescriptor, PiProviderSelection, PiWorkerCredentialProjection, } from '../../electron/coding-runtime/pi/provider-config'; import type { ProviderAccount } from '../../electron/shared/providers/types'; const capability = { schemaVersion: 1, adapter: 'bailian-chat-completions', supportsForcedSearch: true, sourceMode: 'inline-or-structured', billingAuthority: 'model-request', } as const; const account: ProviderAccount = { id: 'niancode-user-models', vendorId: 'custom', label: 'Makelore Models', authMode: 'api_key', apiProtocol: 'openai-completions', baseUrl: 'http://127.0.0.1:13210/api/ai-proxy/v1', model: 'deepseek-v4-pro', enabled: true, isDefault: true, createdAt: '2026-09-02T00:00:00.000Z', updatedAt: '2026-09-02T00:00:00.000Z', metadata: { worksSquareModelCapabilities: { 'deepseek-v4-pro': { reasoningEfforts: ['low', 'high', 'max'], reasoningCanDisable: true, webSearch: capability, }, }, }, }; const selection: PiProviderSelection = { accountId: account.id, runtimeProviderId: 'makelore-account-1', modelId: 'deepseek-v4-pro', thinkingLevel: 'high', input: ['text'], }; const descriptor: PiProviderDescriptor = { accountId: account.id, runtimeProviderId: selection.runtimeProviderId, api: 'openai-completions', baseUrl: account.baseUrl, apiKeyEnv: 'MAKELORE_PI_API_KEY', headers: { Authorization: '$MAKELORE_PI_AUTHORIZATION', }, models: [{ id: selection.modelId, name: selection.modelId, input: ['text'], reasoning: true, }], }; const credential: PiWorkerCredentialProjection = { env: { MAKELORE_PI_API_KEY: 'frozen-token', MAKELORE_PI_AUTHORIZATION: 'Bearer frozen-token', }, sensitiveValues: ['frozen-token', 'Bearer frozen-token'], }; describe('ModelToolRegistry', () => { it('projects and invokes Web Search only for the frozen supported model generation', async () => { const search = vi.fn().mockResolvedValue({ schema: 'makelore-model-tool.v1', tool: 'web_search', status: 'succeeded', modelId: 'deepseek-v4-pro', answer: 'Grounded answer', sources: [{ title: 'Source', url: 'https://example.test/' }], sourceMode: 'inline-or-structured', }); const registry = new ModelToolRegistry({ adapter: { search } }); const registration = registry.registerWorker({ conversationId: 'conversation-1', generation: 3, account, descriptor, selection, credential, }); expect(registration.tools.map(({ name }) => name)).toEqual(['web_search']); expect(registration.tools[0]?.description).toContain('Do not use agent_browser as a fallback'); await expect(registry.invoke('web_search', { conversationId: 'conversation-1', workerGeneration: 3, runId: 'run-1', resourceId: 'tool-call-1', }, { query: ' latest fact ' })).resolves.toEqual({ content: [{ type: 'text', text: 'Grounded answer\n\nSources:\n- Source: https://example.test/', }], details: expect.objectContaining({ schema: 'makelore-model-tool.v1', tool: 'web_search', status: 'succeeded', }), }); expect(search).toHaveBeenCalledWith({ query: 'latest fact', selectedModel: expect.objectContaining({ accountId: account.id, modelId: selection.modelId, generation: 3, headers: { Authorization: 'Bearer frozen-token' }, }), parentTurnId: 'run-1', toolCallId: 'tool-call-1', }, expect.any(AbortSignal)); registration.dispose(); const stale = await registry.invoke('web_search', { conversationId: 'conversation-1', workerGeneration: 3, runId: 'run-2', resourceId: 'tool-call-2', }, { query: 'new fact' }); expect(stale.details).toMatchObject({ status: 'failed', error: { code: 'model_context_changed', httpStatus: 409, retryable: false }, }); expect(search).toHaveBeenCalledTimes(1); }); it('does not project a tool when the selected model has no capability', () => { const registry = new ModelToolRegistry({ adapter: { search: vi.fn() } }); const registration = registry.registerWorker({ conversationId: 'conversation-2', generation: 1, account: { ...account, metadata: undefined }, descriptor, selection, credential, }); expect(registration.tools).toEqual([]); registration.dispose(); }); });