makelore/tests/unit/model-tool-registry.test.ts

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();
  });
});