makelore/tests/unit/managed-model-capabilities.test.ts

// @vitest-environment node
import { createServer } from 'node:http';
import { mkdtemp, rm, writeFile } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { pathToFileURL } from 'node:url';
import { describe, expect, it, vi } from 'vitest';
import { streamSimple } from '@earendil-works/pi-ai/api/openai-completions';
import type { Model } from '@earendil-works/pi-ai';
import {
  buildManagedModelRequest, managedPiThinkingLevel, managedReasoningOptions,
  unknownManagedModelCapability,
  type ManagedModelCapability, type ManagedReasoningChoice,
} from '../../shared/managed-model-capabilities';
import { materializeMakelorePiExtension } from '../../electron/coding-runtime/pi/extensions/makelore-runtime';

function capability(format: 'qwen' | 'deepseek'): ManagedModelCapability {
  return {
    ...unknownManagedModelCapability(),
    inputModalities: ['text', 'image'], outputModalities: ['text'], resolutionStatus: 'ready',
    reasoning: {
      supported: true, canDisable: true, defaultEnabled: true,
      effortValues: ['low', 'medium', 'xhigh', 'max', 'future-native'],
      defaultEffort: 'low', controlFormat: format, budget: null,
    },
  };
}

describe('managed model reasoning', () => {
  it('distinguishes unknown, switch-only and always-on capabilities', () => {
    expect(managedReasoningOptions(unknownManagedModelCapability()).map((x) => x.value)).toEqual(['default']);
    const profile = capability('qwen');
    profile.reasoning.effortValues = [];
    expect(managedReasoningOptions(profile).map((x) => x.value)).toEqual(['default', 'disabled', 'enabled']);
    profile.reasoning.canDisable = false;
    profile.reasoning.effortValues = ['max'];
    expect(managedReasoningOptions(profile).map((x) => x.value)).toEqual(['default', 'effort:max']);
    expect(buildManagedModelRequest('model', { mode: 'enabled', effort: 'max' }, profile).reasoningFields)
      .toEqual({ reasoning_effort: 'max' });
    expect(() => buildManagedModelRequest('model', { mode: 'disabled' }, profile)).toThrow('不可用');
    expect(() => buildManagedModelRequest('model', { mode: 'enabled', effort: 'high' }, profile)).toThrow('不可用');
  });

  it.each(['qwen', 'deepseek'] as const)('sends exact %s controls through Pi and the bundled extension', async (format) => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-managed-wire-'));
    const bodies: Record<string, unknown>[] = [];
    const server = createServer(async (request, response) => {
      let raw = '';
      for await (const chunk of request) raw += String(chunk);
      bodies.push(JSON.parse(raw));
      response.writeHead(200, { 'content-type': 'text/event-stream' });
      response.end('data: ' + JSON.stringify({
        id: 'test-response', object: 'chat.completion.chunk', created: 0, model: 'catalog-model',
        choices: [{ index: 0, delta: { role: 'assistant', content: 'ok' }, finish_reason: 'stop' }],
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
      }) + '\n\ndata: [DONE]\n\n');
    });
    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
    try {
      const address = server.address();
      if (!address || typeof address === 'string') throw new Error('No test server address');
      const contextFile = path.join(root, 'worker.json');
      const profile = capability(format);
      await writeFile(contextFile, JSON.stringify({ runId: 'test-run', tools: [], allowedToolNames: [] }));
      const extensionPath = await materializeMakelorePiExtension(root);
      const extension = await import(/* @vite-ignore */ pathToFileURL(extensionPath).href);
      const hooks = new Map<string, (event: { payload: unknown }) => Promise<unknown>>();
      await extension.createMakeloreRuntime({ contextFile })({
        registerFlag: vi.fn(), registerTool: vi.fn(),
        on: (name: string, handler: (event: { payload: unknown }) => Promise<unknown>) => hooks.set(name, handler),
      });
      const beforeRequest = hooks.get('before_provider_request');
      expect(beforeRequest).toBeTypeOf('function');
      const model: Model<'openai-completions'> = {
        id: 'catalog-model', name: 'Catalog model', api: 'openai-completions',
        provider: 'managed-test', baseUrl: 'http://127.0.0.1:' + address.port + '/v1',
        reasoning: true, input: ['text', 'image'], contextWindow: 4096, maxTokens: 512,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        thinkingLevelMap: { off: 'off', low: 'low', medium: 'medium', xhigh: 'xhigh', max: 'max' },
        compat: { thinkingFormat: format, supportsReasoningEffort: true },
      };
      const choices: ManagedReasoningChoice[] = [
        { mode: 'default' }, { mode: 'disabled' }, { mode: 'enabled' },
        ...['low', 'medium', 'xhigh', 'max', 'future-native'].map((effort) => ({ mode: 'enabled' as const, effort })),
      ];
      for (const choice of choices) {
        const prepared = buildManagedModelRequest(model.id, choice, profile);
        await writeFile(contextFile, JSON.stringify({
          runId: 'test-run', tools: [], allowedToolNames: [], managedModelRequest: prepared,
        }));
        const result = await streamSimple(model, {
          messages: [{ role: 'user', content: 'test', timestamp: 0 }],
        }, {
          apiKey: 'synthetic-key', reasoning: managedPiThinkingLevel(choice),
          onPayload: (payload) => beforeRequest!({ payload }),
        }).result();
        expect(result.stopReason, result.errorMessage).not.toBe('error');
        const body = bodies.at(-1)!;
        for (const field of ['thinking', 'enable_thinking', 'reasoning_effort', 'thinking_budget']) {
          expect(body[field], field + ' for ' + JSON.stringify(choice)).toEqual(prepared.reasoningFields[field]);
        }
        expect(body.model).toBe(model.id);
        expect(body.messages).toEqual([{ role: 'user', content: 'test' }]);
      }
      expect(bodies).toHaveLength(choices.length);
    } finally {
      server.closeAllConnections();
      await new Promise<void>((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
      await rm(root, { recursive: true, force: true });
    }
  });
});