waggle-os/packages/server/tests/local/dynamic-provider-completion.test.ts

import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import type { FastifyInstance } from 'fastify';

vi.mock('../../src/local/lifecycle.js', () => ({
  getLiteLLMStatus: vi.fn(async (port = 4000) => ({ status: 'running', port })),
  startLiteLLM: vi.fn(async (port = 4000) => ({ status: 'started', port })),
  stopLiteLLM: vi.fn(async () => undefined),
}));

import { runAgentLoop } from '../../../agent/src/agent-loop.js';
import { buildLocalServer } from '../../src/local/index.js';
import { PROVIDER_ENV_NAMES } from '../../src/local/provider-env.js';
import { startLiteLLM } from '../../src/local/lifecycle.js';
import { injectWithAuth } from '../test-utils.js';

describe('dynamic provider model completion path', () => {
  let server: FastifyInstance;
  let dataDir: string;
  const originalProviderEnv = new Map<string, string | undefined>();

  beforeEach(async () => {
    vi.clearAllMocks();
    for (const envName of new Set(Object.values(PROVIDER_ENV_NAMES).flat())) {
      originalProviderEnv.set(envName, process.env[envName]);
      delete process.env[envName];
    }
    dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-dynamic-completion-'));
    server = await buildLocalServer({
      dataDir,
      port: 0,
      manageLiteLLM: true,
      managedLiteLLMPort: 4567,
    });
  }, 30_000);

  afterEach(async () => {
    await server.close();
    fs.rmSync(dataDir, { recursive: true, force: true });
    for (const [envName, value] of originalProviderEnv) {
      if (value === undefined) delete process.env[envName];
      else process.env[envName] = value;
    }
    originalProviderEnv.clear();
    vi.restoreAllMocks();
  }, 30_000);

  it('saves a key, discovers an unseen model, configures it, and completes with that exact id', async () => {
    const newModel = 'openai/model-released-after-this-build';
    vi.spyOn(globalThis, 'fetch').mockImplementation(async (input) => {
      const url = String(input);
      if (url === 'https://api.openai.com/v1/models') {
        return new Response(JSON.stringify({
          data: [{ id: 'model-released-after-this-build' }],
        }), { status: 200 });
      }
      throw new Error(`Unexpected discovery request: ${url}`);
    });

    const save = await injectWithAuth(server, {
      method: 'PUT',
      url: '/api/settings',
      payload: { providers: { openai: { apiKey: 'new-provider-key' } } },
    });

    expect(save.statusCode).toBe(200);
    expect(save.json().router).toMatchObject({
      managed: true,
      ready: true,
      models: [newModel],
    });
    expect(startLiteLLM).toHaveBeenCalledWith(4567, path.join(dataDir, 'litellm.runtime.json'));

    const routerConfig = JSON.parse(
      fs.readFileSync(path.join(dataDir, 'litellm.runtime.json'), 'utf8'),
    ) as { model_list: Array<{ model_name: string }> };
    expect(routerConfig.model_list.map((entry) => entry.model_name)).toContain(newModel);

    const completionFetch = vi.fn(async (_url: string, init?: RequestInit) => {
      const body = JSON.parse(String(init?.body)) as { model: string };
      const configured = routerConfig.model_list.some((entry) => entry.model_name === body.model);
      return new Response(JSON.stringify(configured
        ? {
            choices: [{ message: { role: 'assistant', content: 'Dynamic model completed.' }, finish_reason: 'stop' }],
            usage: { prompt_tokens: 4, completion_tokens: 3 },
          }
        : { error: { message: 'Model is not configured' } }), {
        status: configured ? 200 : 404,
        headers: { 'content-type': 'application/json' },
      });
    });

    const completion = await runAgentLoop({
      litellmUrl: server.localConfig.litellmUrl,
      litellmApiKey: server.agentState.litellmApiKey,
      model: newModel,
      systemPrompt: 'Be concise.',
      tools: [],
      messages: [{ role: 'user', content: 'Confirm routing.' }],
      fetch: completionFetch,
    });

    expect(completion.content).toBe('Dynamic model completed.');
    expect(JSON.parse(String(completionFetch.mock.calls[0]?.[1]?.body)).model).toBe(newModel);
  });

  it('makes a model released during the running session executable when selected', async () => {
    const existingModel = 'openai/existing-runtime-model';
    const newModel = 'openai/model-released-during-this-session';
    let released = false;
    vi.spyOn(globalThis, 'fetch').mockImplementation(async (input) => {
      const url = String(input);
      if (url === 'https://api.openai.com/v1/models') {
        return new Response(JSON.stringify({
          data: [
            { id: 'existing-runtime-model' },
            ...(released ? [{ id: 'model-released-during-this-session' }] : []),
          ],
        }), { status: 200 });
      }
      throw new Error(`Unexpected discovery request: ${url}`);
    });

    const save = await injectWithAuth(server, {
      method: 'PUT',
      url: '/api/settings',
      payload: { providers: { openai: { apiKey: 'hot-refresh-provider-key' } } },
    });
    expect(save.statusCode).toBe(200);
    expect(save.json().router.models).toEqual([existingModel]);

    released = true;
    const providers = await injectWithAuth(server, { method: 'GET', url: '/api/providers' });
    expect(providers.statusCode).toBe(200);
    expect(providers.json().providers.find((provider: { id: string }) => provider.id === 'openai').models)
      .toEqual(expect.arrayContaining([expect.objectContaining({ id: newModel })]));

    const selected = await injectWithAuth(server, {
      method: 'PUT',
      url: '/api/settings',
      payload: { defaultModel: newModel },
    });
    expect(selected.statusCode).toBe(200);
    expect(selected.json()).toMatchObject({ defaultModel: newModel });
    expect(server.agentState.currentModel).toBe(newModel);
    expect(startLiteLLM).toHaveBeenCalledTimes(2);

    const routerConfig = JSON.parse(
      fs.readFileSync(path.join(dataDir, 'litellm.runtime.json'), 'utf8'),
    ) as { model_list: Array<{ model_name: string }> };
    expect(routerConfig.model_list.map((entry) => entry.model_name)).toContain(newModel);

    const completionFetch = vi.fn(async (_url: string, init?: RequestInit) => {
      const body = JSON.parse(String(init?.body)) as { model: string };
      return new Response(JSON.stringify({
        choices: [{ message: { role: 'assistant', content: 'Hot model completed.' }, finish_reason: 'stop' }],
        usage: { prompt_tokens: 4, completion_tokens: 3 },
      }), { status: 200, headers: { 'content-type': 'application/json' } });
    });
    const completion = await runAgentLoop({
      litellmUrl: server.localConfig.litellmUrl,
      litellmApiKey: server.agentState.litellmApiKey,
      model: server.agentState.currentModel,
      systemPrompt: 'Be concise.',
      tools: [],
      messages: [{ role: 'user', content: 'Confirm hot routing.' }],
      fetch: completionFetch,
    });
    expect(completion.content).toBe('Hot model completed.');
    expect(JSON.parse(String(completionFetch.mock.calls[0]?.[1]?.body)).model).toBe(newModel);
  });
});