224 lines
8.2 KiB
TypeScript
224 lines
8.2 KiB
TypeScript
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 { VaultStore } from '@waggle/core';
|
|
import {
|
|
buildLiteLLMRuntimeConfig,
|
|
ensureManagedLiteLLMModel,
|
|
prepareLiteLLMRuntimeConfig,
|
|
refreshManagedLiteLLM,
|
|
} from '../../src/local/litellm-runtime-config.js';
|
|
import { clearProviderModelCache, type DiscoveredProviderModel } from '../../src/local/provider-model-catalog.js';
|
|
import { startLiteLLM, stopLiteLLM } from '../../src/local/lifecycle.js';
|
|
import { PROVIDER_ENV_NAMES } from '../../src/local/provider-env.js';
|
|
import type { FastifyInstance } from 'fastify';
|
|
|
|
vi.mock('../../src/local/lifecycle.js', () => ({
|
|
startLiteLLM: vi.fn(async (port: number) => ({ status: 'started', port })),
|
|
stopLiteLLM: vi.fn(async () => undefined),
|
|
}));
|
|
|
|
const tempDirs: string[] = [];
|
|
const originalProviderEnv = new Map<string, string | undefined>();
|
|
|
|
function model(id: string): DiscoveredProviderModel {
|
|
return { id, name: id, cost: '$$', speed: 'medium', source: 'provider-api' };
|
|
}
|
|
|
|
beforeEach(() => {
|
|
clearProviderModelCache();
|
|
vi.clearAllMocks();
|
|
for (const envName of new Set(Object.values(PROVIDER_ENV_NAMES).flat())) {
|
|
originalProviderEnv.set(envName, process.env[envName]);
|
|
delete process.env[envName];
|
|
}
|
|
});
|
|
afterEach(() => {
|
|
for (const dir of tempDirs.splice(0)) fs.rmSync(dir, { recursive: true, force: true });
|
|
for (const [envName, value] of originalProviderEnv) {
|
|
if (value === undefined) delete process.env[envName];
|
|
else process.env[envName] = value;
|
|
}
|
|
originalProviderEnv.clear();
|
|
});
|
|
|
|
describe('dynamic LiteLLM runtime config', () => {
|
|
it('routes provider-discovered ids without maintaining a model inventory', () => {
|
|
const config = buildLiteLLMRuntimeConfig(new Map([
|
|
['google', [model('google/gemini-model-released-tomorrow')]],
|
|
['alibaba', [model('alibaba/qwen-model-released-tomorrow')]],
|
|
]));
|
|
|
|
expect(config.model_list).toEqual([
|
|
{
|
|
model_name: 'google/gemini-model-released-tomorrow',
|
|
litellm_params: {
|
|
model: 'gemini/gemini-model-released-tomorrow',
|
|
api_key: 'os.environ/GEMINI_API_KEY',
|
|
},
|
|
},
|
|
{
|
|
model_name: 'alibaba/qwen-model-released-tomorrow',
|
|
litellm_params: {
|
|
model: 'openai/qwen-model-released-tomorrow',
|
|
api_key: 'os.environ/DASHSCOPE_API_KEY',
|
|
api_base: 'https://dashscope-intl.aliyuncs.com/compatible-mode/v1',
|
|
},
|
|
},
|
|
]);
|
|
});
|
|
|
|
it('writes every model returned by the provider and never writes the API secret', async () => {
|
|
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-litellm-config-'));
|
|
tempDirs.push(dataDir);
|
|
const vault = new VaultStore(dataDir);
|
|
vault.set('openai', 'super-secret-provider-key');
|
|
const fetchImpl = vi.fn<typeof fetch>(async () => new Response(JSON.stringify({
|
|
data: [
|
|
{ id: 'existing-model' },
|
|
{ id: 'brand-new-unseen-model' },
|
|
],
|
|
}), { status: 200 }));
|
|
|
|
const result = await prepareLiteLLMRuntimeConfig(dataDir, vault, { fetchImpl });
|
|
|
|
expect(result.modelIds).toEqual([
|
|
'openai/existing-model',
|
|
'openai/brand-new-unseen-model',
|
|
]);
|
|
expect(result.configPath).toBe(path.join(dataDir, 'litellm.runtime.json'));
|
|
const raw = fs.readFileSync(result.configPath!, 'utf8');
|
|
expect(raw).toContain('openai/brand-new-unseen-model');
|
|
expect(raw).toContain('os.environ/OPENAI_API_KEY');
|
|
expect(raw).not.toContain('super-secret-provider-key');
|
|
});
|
|
|
|
it('discovers with a working Google alias when the first Gemini alias is stale', async () => {
|
|
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-litellm-google-alias-'));
|
|
tempDirs.push(dataDir);
|
|
const vault = new VaultStore(dataDir);
|
|
process.env.GEMINI_API_KEY = 'stale-gemini-key';
|
|
process.env.GOOGLE_API_KEY = 'working-google-key';
|
|
const fetchImpl = vi.fn<typeof fetch>(async (_url, init) => {
|
|
const key = (init?.headers as Record<string, string>)['x-goog-api-key'];
|
|
if (key === 'stale-gemini-key') {
|
|
return new Response(JSON.stringify({ error: 'invalid key' }), { status: 400 });
|
|
}
|
|
return new Response(JSON.stringify({
|
|
models: [{ name: 'models/gemini-2.5-flash' }],
|
|
}), { status: 200 });
|
|
});
|
|
|
|
const result = await prepareLiteLLMRuntimeConfig(dataDir, vault, { fetchImpl });
|
|
|
|
expect(fetchImpl).toHaveBeenCalledTimes(2);
|
|
expect(result.modelIds).toEqual(['google/gemini-2.5-flash']);
|
|
expect(process.env.GEMINI_API_KEY).toBe('working-google-key');
|
|
expect(process.env.GOOGLE_API_KEY).toBe('working-google-key');
|
|
});
|
|
|
|
it('restarts the managed router with a newly discovered model and switches runtime routing', async () => {
|
|
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-litellm-refresh-'));
|
|
tempDirs.push(dataDir);
|
|
const vault = new VaultStore(dataDir);
|
|
vault.set('openai', 'runtime-refresh-key');
|
|
const originalFetch = globalThis.fetch;
|
|
globalThis.fetch = vi.fn<typeof fetch>(async () => new Response(JSON.stringify({
|
|
data: [{ id: 'model-that-did-not-exist-at-build-time' }],
|
|
}), { status: 200 }));
|
|
const server = {
|
|
localConfig: {
|
|
dataDir,
|
|
port: 3333,
|
|
host: '127.0.0.1',
|
|
litellmUrl: 'http://127.0.0.1:3333/v1',
|
|
manageLiteLLM: true,
|
|
managedLiteLLMPort: 4567,
|
|
},
|
|
vault,
|
|
agentState: {
|
|
litellmApiKey: 'fallback-key',
|
|
llmProvider: {
|
|
provider: 'anthropic-proxy',
|
|
health: 'degraded',
|
|
detail: 'fallback',
|
|
checkedAt: new Date(0).toISOString(),
|
|
},
|
|
},
|
|
} as unknown as FastifyInstance;
|
|
|
|
try {
|
|
const result = await refreshManagedLiteLLM(server);
|
|
|
|
expect(result).toMatchObject({
|
|
managed: true,
|
|
ready: true,
|
|
port: 4567,
|
|
models: ['openai/model-that-did-not-exist-at-build-time'],
|
|
});
|
|
expect(stopLiteLLM).toHaveBeenCalled();
|
|
expect(startLiteLLM).toHaveBeenCalledWith(4567, path.join(dataDir, 'litellm.runtime.json'));
|
|
expect(server.localConfig.litellmUrl).toBe('http://localhost:4567');
|
|
expect(server.agentState.llmProvider.provider).toBe('litellm');
|
|
expect(server.agentState.llmProvider.health).toBe('healthy');
|
|
} finally {
|
|
globalThis.fetch = originalFetch;
|
|
}
|
|
});
|
|
|
|
it('hot-loads a model released after startup and does not restart for an existing model', async () => {
|
|
const dataDir = fs.mkdtempSync(path.join(os.tmpdir(), 'waggle-litellm-hot-model-'));
|
|
tempDirs.push(dataDir);
|
|
const vault = new VaultStore(dataDir);
|
|
vault.set('openai', 'hot-model-key');
|
|
let released = false;
|
|
const originalFetch = globalThis.fetch;
|
|
globalThis.fetch = vi.fn<typeof fetch>(async () => new Response(JSON.stringify({
|
|
data: [
|
|
{ id: 'existing-model' },
|
|
...(released ? [{ id: 'model-released-while-waggle-is-open' }] : []),
|
|
],
|
|
}), { status: 200 }));
|
|
const server = {
|
|
localConfig: {
|
|
dataDir,
|
|
port: 3333,
|
|
host: '127.0.0.1',
|
|
litellmUrl: 'http://localhost:4567',
|
|
manageLiteLLM: true,
|
|
managedLiteLLMPort: 4567,
|
|
},
|
|
vault,
|
|
agentState: {
|
|
litellmApiKey: 'router-key',
|
|
llmProvider: {
|
|
provider: 'litellm',
|
|
health: 'healthy',
|
|
detail: 'ready',
|
|
checkedAt: new Date().toISOString(),
|
|
},
|
|
},
|
|
} as unknown as FastifyInstance;
|
|
|
|
try {
|
|
await prepareLiteLLMRuntimeConfig(dataDir, vault);
|
|
expect(await ensureManagedLiteLLMModel(server, 'openai/existing-model')).toBe(true);
|
|
expect(stopLiteLLM).not.toHaveBeenCalled();
|
|
expect(startLiteLLM).not.toHaveBeenCalled();
|
|
|
|
released = true;
|
|
expect(await ensureManagedLiteLLMModel(
|
|
server,
|
|
'openai/model-released-while-waggle-is-open',
|
|
)).toBe(true);
|
|
expect(stopLiteLLM).toHaveBeenCalledTimes(1);
|
|
expect(startLiteLLM).toHaveBeenCalledWith(4567, path.join(dataDir, 'litellm.runtime.json'));
|
|
expect(fs.readFileSync(path.join(dataDir, 'litellm.runtime.json'), 'utf8'))
|
|
.toContain('openai/model-released-while-waggle-is-open');
|
|
} finally {
|
|
globalThis.fetch = originalFetch;
|
|
}
|
|
});
|
|
});
|