feat: initial public release (MAESTRO)
This commit is contained in:
oss-sync
@@ -0,0 +1,760 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { OpenAICompatClient, type LLMEvent, type Message } from './openai-compat.js';
|
||||
|
||||
function createSSEResponse(chunks: string[], shouldError = false, errorAfterChunks = 0): Response {
|
||||
let chunkIndex = 0;
|
||||
const encoder = new TextEncoder();
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
pull(controller) {
|
||||
if (shouldError && chunkIndex >= errorAfterChunks) {
|
||||
controller.error(new Error('Connection reset'));
|
||||
return;
|
||||
}
|
||||
if (chunkIndex < chunks.length) {
|
||||
controller.enqueue(encoder.encode(chunks[chunkIndex] + '\n'));
|
||||
chunkIndex++;
|
||||
} else {
|
||||
controller.close();
|
||||
}
|
||||
},
|
||||
});
|
||||
return new Response(stream, { status: 200, headers: { 'Content-Type': 'text/event-stream' } });
|
||||
}
|
||||
|
||||
function makeSseResponse(chunks: Array<Record<string, unknown> | '[DONE]'>): Response {
|
||||
const encoder = new TextEncoder();
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
for (const chunk of chunks) {
|
||||
const data = chunk === '[DONE]' ? chunk : JSON.stringify(chunk);
|
||||
controller.enqueue(encoder.encode(`data: ${data}\n\n`));
|
||||
}
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(stream, {
|
||||
status: 200,
|
||||
headers: { 'Content-Type': 'text/event-stream' },
|
||||
});
|
||||
}
|
||||
|
||||
async function collectEvents(client: OpenAICompatClient, messages: Message[], tools?: undefined, externalSignal?: AbortSignal): Promise<LLMEvent[]> {
|
||||
const events: LLMEvent[] = [];
|
||||
for await (const event of client.chat(messages, tools, externalSignal)) {
|
||||
events.push(event);
|
||||
}
|
||||
return events;
|
||||
}
|
||||
|
||||
describe('OpenAICompatClient retry', () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
it('retries retryable HTTP responses and continues streaming on success', async () => {
|
||||
const fetchMock = vi.fn()
|
||||
.mockResolvedValueOnce(new Response('busy', { status: 503 }))
|
||||
.mockResolvedValueOnce(makeSseResponse([
|
||||
{ choices: [{ delta: { content: 'hello' } }] },
|
||||
'[DONE]',
|
||||
]));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 2, backoffMs: [0], retryableStatus: [503] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
expect(events).toEqual([
|
||||
{ type: 'text', text: 'hello' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
});
|
||||
|
||||
it('retries transient fetch errors', async () => {
|
||||
const fetchMock = vi.fn()
|
||||
.mockRejectedValueOnce(new TypeError('socket hang up'))
|
||||
.mockResolvedValueOnce(makeSseResponse([
|
||||
{ choices: [{ delta: { content: 'ok' } }] },
|
||||
'[DONE]',
|
||||
]));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 2, backoffMs: [0], retryableStatus: [503] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
expect(events).toEqual([
|
||||
{ type: 'text', text: 'ok' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
});
|
||||
|
||||
it('does not retry non-retryable HTTP responses', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(new Response('bad request', { status: 400 }));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 3, backoffMs: [0, 0], retryableStatus: [503] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
expect(events).toEqual([
|
||||
{ type: 'error', error: 'HTTP 400: bad request' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('blocks oversized requests before fetch', async () => {
|
||||
const fetchMock = vi.fn();
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 3, backoffMs: [0, 0], retryableStatus: [503] },
|
||||
10_000,
|
||||
1_000,
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'x'.repeat(4_000) }]);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(events).toEqual([
|
||||
expect.objectContaining({
|
||||
type: 'error',
|
||||
error: expect.stringContaining('LLM request blocked before send'),
|
||||
}),
|
||||
]);
|
||||
});
|
||||
|
||||
it('allows requests that are only slightly above the old conservative cap', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(makeSseResponse([
|
||||
{ choices: [{ delta: { content: 'ok' } }] },
|
||||
'[DONE]',
|
||||
]));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [503] },
|
||||
10_000,
|
||||
32_000,
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'x'.repeat(15_950) }]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
expect(events).toEqual([
|
||||
{ type: 'text', text: 'ok' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
});
|
||||
|
||||
it('uses the provided context limit directly when larger than the legacy 32k cap', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(makeSseResponse([
|
||||
{ choices: [{ delta: { content: 'ok' } }] },
|
||||
'[DONE]',
|
||||
]));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [503] },
|
||||
10_000,
|
||||
200_000,
|
||||
);
|
||||
|
||||
// 30k ASCII chars are well under 200_000 * 0.8 = 160_000.
|
||||
// Before the cap was removed, this fired the 32k preflight guard.
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'x'.repeat(30_000) }]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
expect(events).toEqual([
|
||||
{ type: 'text', text: 'ok' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
});
|
||||
|
||||
it('still blocks prompts that exceed the provided context limit', async () => {
|
||||
const fetchMock = vi.fn();
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [503] },
|
||||
10_000,
|
||||
200_000,
|
||||
);
|
||||
|
||||
// 200_000 * 0.8 = 160_000 max prompt.
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'x'.repeat(600_000) }]);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(events).toEqual([
|
||||
expect.objectContaining({
|
||||
type: 'error',
|
||||
error: expect.stringContaining('context 200,000'),
|
||||
}),
|
||||
]);
|
||||
});
|
||||
|
||||
it('counts image_url parts as fixed image cost instead of data URL text', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(makeSseResponse([
|
||||
{ choices: [{ delta: { content: 'ok' } }] },
|
||||
'[DONE]',
|
||||
]));
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
const preflightLines: string[] = [];
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [503] },
|
||||
10_000,
|
||||
32_000,
|
||||
0.8,
|
||||
(line) => preflightLines.push(line),
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{
|
||||
role: 'user',
|
||||
content: [
|
||||
{ type: 'text', text: 'inspect this' },
|
||||
{ type: 'image_url', image_url: { url: `data:image/png;base64,${'A'.repeat(200_000)}` } },
|
||||
],
|
||||
}]);
|
||||
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
expect(events).toEqual([
|
||||
{ type: 'text', text: 'ok' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
expect(preflightLines[0]).toContain('images=1,imageTokenCost=1024');
|
||||
expect(preflightLines[0]).toContain('requestJsonChars=200,');
|
||||
});
|
||||
});
|
||||
|
||||
describe('OpenAICompatClient stream retry', () => {
|
||||
let fetchMock: ReturnType<typeof vi.fn>;
|
||||
const originalFetch = globalThis.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
fetchMock = vi.fn();
|
||||
globalThis.fetch = fetchMock;
|
||||
});
|
||||
afterEach(() => {
|
||||
globalThis.fetch = originalFetch;
|
||||
});
|
||||
|
||||
const retryConfig = {
|
||||
maxAttempts: 3,
|
||||
backoffMs: [0, 0, 0],
|
||||
retryableStatus: [500],
|
||||
};
|
||||
|
||||
it('正常なストリームを処理する', async () => {
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"hello"},"finish_reason":null}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
expect(events).toContainEqual({ type: 'text', text: 'hello' });
|
||||
expect(events).toContainEqual({ type: 'done', usage: undefined });
|
||||
});
|
||||
|
||||
it('tool_calls を正しく蓄積して emit する', async () => {
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"Read","arguments":""}}]},"finish_reason":null}]}',
|
||||
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"path\\":\\"test\\"}"}}]},"finish_reason":null}]}',
|
||||
'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
const toolEvent = events.find(e => e.type === 'tool_use');
|
||||
expect(toolEvent).toBeDefined();
|
||||
expect(toolEvent).toMatchObject({
|
||||
type: 'tool_use',
|
||||
id: 'call_1',
|
||||
name: 'Read',
|
||||
input: { path: 'test' },
|
||||
});
|
||||
});
|
||||
|
||||
it('emits tool_use_delta snapshots (full accumulated args) while streaming', async () => {
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"Write","arguments":""}}]},"finish_reason":null}]}',
|
||||
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"content\\":\\"ab"}}]},"finish_reason":null}]}',
|
||||
'data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"cd\\"}"}}]},"finish_reason":null}]}',
|
||||
'data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
const deltas = events.filter(e => e.type === 'tool_use_delta');
|
||||
expect(deltas).toHaveLength(2);
|
||||
// Each delta carries the FULL accumulated args so far (snapshot), so a
|
||||
// late-attaching client always receives the opening JSON structure.
|
||||
expect(deltas[0]).toMatchObject({ type: 'tool_use_delta', index: 0, callId: 'call_1', name: 'Write', chunk: '{"content":"ab' });
|
||||
expect(deltas[1]).toMatchObject({ type: 'tool_use_delta', index: 0, callId: 'call_1', name: 'Write', chunk: '{"content":"abcd"}' });
|
||||
// Final aggregated tool_use must still be emitted unchanged
|
||||
const toolEvent = events.find(e => e.type === 'tool_use');
|
||||
expect(toolEvent).toMatchObject({ type: 'tool_use', id: 'call_1', name: 'Write', input: { content: 'abcd' } });
|
||||
});
|
||||
|
||||
it('ストリーム途中エラーでリトライし完了する', async () => {
|
||||
// 1回目: 1チャンク後にエラー
|
||||
fetchMock.mockResolvedValueOnce(
|
||||
createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"partial"},"finish_reason":null}]}',
|
||||
], true, 1)
|
||||
);
|
||||
// 2回目: 正常完了
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"hello"},"finish_reason":null}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
expect(events.some(e => e.type === 'done')).toBe(true);
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('maxAttempts 超過でエラーを返す', async () => {
|
||||
// 全試行でストリームエラー
|
||||
for (let i = 0; i < 3; i++) {
|
||||
fetchMock.mockResolvedValueOnce(
|
||||
createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"x"},"finish_reason":null}]}',
|
||||
], true, 1)
|
||||
);
|
||||
}
|
||||
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, {
|
||||
maxAttempts: 3,
|
||||
backoffMs: [0, 0, 0],
|
||||
retryableStatus: [500],
|
||||
});
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
expect(events.some(e => e.type === 'error')).toBe(true);
|
||||
expect(fetchMock).toHaveBeenCalledTimes(3);
|
||||
});
|
||||
|
||||
it('fetch 接続エラーでリトライする', async () => {
|
||||
fetchMock.mockRejectedValueOnce(new Error('ECONNREFUSED'));
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"ok"},"finish_reason":null}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
expect(events.some(e => e.type === 'done')).toBe(true);
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('HTTP 500 でリトライする', async () => {
|
||||
fetchMock.mockResolvedValueOnce(new Response('Internal Server Error', { status: 500 }));
|
||||
fetchMock.mockResolvedValueOnce(createSSEResponse([
|
||||
'data: {"choices":[{"delta":{"content":"ok"},"finish_reason":null}]}',
|
||||
'data: [DONE]',
|
||||
]));
|
||||
|
||||
const client = new OpenAICompatClient('http://test', 'model', undefined, retryConfig);
|
||||
const events = [];
|
||||
for await (const event of client.chat([{ role: 'user', content: 'hi' }])) {
|
||||
events.push(event);
|
||||
}
|
||||
expect(events.some(e => e.type === 'done')).toBe(true);
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
|
||||
describe('OpenAICompatClient external AbortSignal', () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
it('returns cancelled error immediately when external signal is already aborted', async () => {
|
||||
const fetchMock = vi.fn();
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient('http://llm.test/v1', 'test-model');
|
||||
const abortController = new AbortController();
|
||||
abortController.abort();
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }], undefined, abortController.signal);
|
||||
|
||||
expect(fetchMock).not.toHaveBeenCalled();
|
||||
expect(events).toEqual([
|
||||
{ type: 'error', error: 'Request cancelled by caller' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('includes dynamic timeout minutes in error message (default 10 minutes)', async () => {
|
||||
// Use a short timeout so the test completes quickly
|
||||
const timeoutMs = 100;
|
||||
const abortError = new DOMException('The operation was aborted', 'AbortError');
|
||||
const fetchMock = vi.fn().mockRejectedValue(abortError);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
timeoutMs,
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
|
||||
expect(events).toEqual([
|
||||
{ type: 'error', error: 'Request timed out (0 minutes)' },
|
||||
]);
|
||||
});
|
||||
|
||||
it('shows correct minutes for custom timeout', async () => {
|
||||
const timeoutMs = 5 * 60 * 1000; // 5 minutes
|
||||
const abortError = new DOMException('The operation was aborted', 'AbortError');
|
||||
const fetchMock = vi.fn().mockRejectedValue(abortError);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://llm.test/v1',
|
||||
'test-model',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
timeoutMs,
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
|
||||
expect(events).toEqual([
|
||||
{ type: 'error', error: 'Request timed out (5 minutes)' },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe('OpenAICompatClient proxy backend headers', () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
function makeSseResponseWithHeaders(
|
||||
chunks: Array<Record<string, unknown> | '[DONE]'>,
|
||||
headers: Record<string, string>,
|
||||
): Response {
|
||||
const encoder = new TextEncoder();
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
for (const chunk of chunks) {
|
||||
const data = chunk === '[DONE]' ? chunk : JSON.stringify(chunk);
|
||||
controller.enqueue(encoder.encode(`data: ${data}\n\n`));
|
||||
}
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
return new Response(stream, {
|
||||
status: 200,
|
||||
headers: { 'Content-Type': 'text/event-stream', ...headers },
|
||||
});
|
||||
}
|
||||
|
||||
it('emits a backend event when proxy=true and x-litellm-model-id header is present', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
{ 'x-litellm-model-id': 'gpu-rtx-a' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
'qwen3:8b',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events).toEqual([
|
||||
{ type: 'backend', backendId: 'gpu-rtx-a', cacheKey: null },
|
||||
{ type: 'text', text: 'ok' },
|
||||
{ type: 'done', usage: undefined },
|
||||
]);
|
||||
});
|
||||
|
||||
it('includes cacheKey when x-litellm-cache-key is present', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'cached' } }] }, '[DONE]'],
|
||||
{ 'x-litellm-model-id': 'gpu-h100-b', 'x-litellm-cache-key': 'sha:abc123' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'q' }]);
|
||||
expect(events[0]).toEqual({ type: 'backend', backendId: 'gpu-h100-b', cacheKey: 'sha:abc123' });
|
||||
});
|
||||
|
||||
it('does not emit a backend event when proxy=false (direct worker)', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
// Even if the upstream happens to set the header, direct mode must ignore it.
|
||||
{ 'x-litellm-model-id': 'should-be-ignored' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://gpu.test/v1',
|
||||
'qwen3:8b',
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events.find(e => e.type === 'backend')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('does not emit a backend event when proxy=true but header is missing', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
{},
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events.find(e => e.type === 'backend')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('trims whitespace from x-litellm-model-id and x-litellm-cache-key', async () => {
|
||||
// config-api.ts trims its side too; without symmetric trim here, the
|
||||
// backend id keyed by the worker would never match what the UI shows.
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
{ 'x-litellm-model-id': ' gpu-a ', 'x-litellm-cache-key': ' sha:xyz ' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events[0]).toEqual({ type: 'backend', backendId: 'gpu-a', cacheKey: 'sha:xyz' });
|
||||
});
|
||||
|
||||
it('does not emit a backend event when x-litellm-model-id is whitespace-only', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
{ 'x-litellm-model-id': ' ' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events.find(e => e.type === 'backend')).toBeUndefined();
|
||||
});
|
||||
|
||||
it('drops whitespace-only cacheKey to null while keeping a valid backendId', async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(
|
||||
makeSseResponseWithHeaders(
|
||||
[{ choices: [{ delta: { content: 'ok' } }] }, '[DONE]'],
|
||||
{ 'x-litellm-model-id': 'gpu-a', 'x-litellm-cache-key': ' ' },
|
||||
),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://litellm.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
undefined,
|
||||
{ proxy: true },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
expect(events[0]).toEqual({ type: 'backend', backendId: 'gpu-a', cacheKey: null });
|
||||
});
|
||||
});
|
||||
|
||||
describe('OpenAICompatClient gateway sentinel error events', () => {
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
function makeSseResponseWithChunks(chunks: Array<string>): Response {
|
||||
const encoder = new TextEncoder();
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
for (const chunk of chunks) {
|
||||
controller.enqueue(encoder.encode(`data: ${chunk}\n\n`));
|
||||
}
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
return new Response(stream, {
|
||||
status: 200,
|
||||
headers: { 'Content-Type': 'text/event-stream' },
|
||||
});
|
||||
}
|
||||
|
||||
for (const errType of ['gateway_shutdown', 'gateway_timeout', 'budget_exhausted', 'rate_limited'] as const) {
|
||||
it(`emits structured error event with gatewayErrorType=${errType}`, async () => {
|
||||
const fetchMock = vi.fn(async () =>
|
||||
makeSseResponseWithChunks([
|
||||
JSON.stringify({ choices: [{ delta: { content: 'partial' } }] }),
|
||||
JSON.stringify({ error: { type: errType, message: 'server says nope' } }),
|
||||
]),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://gateway.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
// partial text event then structured error, no [DONE]
|
||||
const errorEvent = events.find(e => e.type === 'error');
|
||||
expect(errorEvent).toBeDefined();
|
||||
expect((errorEvent as { gatewayErrorType?: string }).gatewayErrorType).toBe(errType);
|
||||
expect((errorEvent as { error: string }).error).toContain(errType);
|
||||
});
|
||||
}
|
||||
|
||||
it('unknown error.type falls through to generic stream parse (no early return)', async () => {
|
||||
const fetchMock = vi.fn(async () =>
|
||||
makeSseResponseWithChunks([
|
||||
JSON.stringify({ error: { type: 'mystery_error', message: 'x' } }),
|
||||
JSON.stringify({ choices: [{ delta: { content: 'ok' } }] }),
|
||||
'[DONE]',
|
||||
]),
|
||||
);
|
||||
vi.stubGlobal('fetch', fetchMock);
|
||||
|
||||
const client = new OpenAICompatClient(
|
||||
'http://gateway.test/v1',
|
||||
undefined,
|
||||
undefined,
|
||||
{ maxAttempts: 1, backoffMs: [0], retryableStatus: [] },
|
||||
);
|
||||
|
||||
const events = await collectEvents(client, [{ role: 'user', content: 'hi' }]);
|
||||
// Unknown gateway error.type should not short-circuit — we keep parsing
|
||||
// and the normal text/done flow continues.
|
||||
expect(events.some(e => e.type === 'text' && e.text === 'ok')).toBe(true);
|
||||
expect(events.at(-1)?.type).toBe('done');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,686 @@
|
||||
import { getDefaultProviderRetryConfig, type ProviderRetryConfig } from '../config.js';
|
||||
import { logger } from '../logger.js';
|
||||
import {
|
||||
IMAGE_CONTENT_TOKENS,
|
||||
estimateMessageTokens,
|
||||
estimateToolsTokens,
|
||||
} from '../engine/context/token-estimate.js';
|
||||
|
||||
export type ContentPart =
|
||||
| { type: 'text'; text: string }
|
||||
| { type: 'image_url'; image_url: { url: string } };
|
||||
|
||||
export interface Message {
|
||||
role: 'system' | 'user' | 'assistant' | 'tool';
|
||||
content?: string | ContentPart[];
|
||||
tool_calls?: ToolCall[];
|
||||
tool_call_id?: string; // role: 'tool' の時
|
||||
name?: string; // role: 'tool' の時
|
||||
}
|
||||
|
||||
export interface ToolCall {
|
||||
id: string;
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string; // JSON string
|
||||
};
|
||||
}
|
||||
|
||||
export interface ToolDef {
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: Record<string, unknown>; // JSON Schema
|
||||
};
|
||||
}
|
||||
|
||||
export type LLMEvent =
|
||||
| { type: 'text'; text: string }
|
||||
| { type: 'tool_use'; id: string; name: string; input: Record<string, unknown> }
|
||||
/**
|
||||
* Tool-call argument SNAPSHOT, emitted as `function.arguments` deltas
|
||||
* stream in (before the aggregated `tool_use`). `chunk` is the FULL
|
||||
* accumulated arguments so far, not just the latest piece — so a client
|
||||
* that attaches to the SSE stream mid-generation still receives the
|
||||
* opening `{"...":"..."` structure the UI's field extractor needs.
|
||||
* Consumers REPLACE their buffer with `chunk` (do not append).
|
||||
* `callId`/`name` come from the accumulator and are stable once the
|
||||
* first chunk has set them.
|
||||
*/
|
||||
| { type: 'tool_use_delta'; index: number; callId: string; name: string; chunk: string }
|
||||
| { type: 'done'; usage?: { prompt_tokens: number; completion_tokens: number } }
|
||||
/**
|
||||
* SSE / response error. `gatewayErrorType` is set when the error came
|
||||
* from an AAO Gateway sentinel SSE event (`data: {"error":{"type":...}}`):
|
||||
* - `gateway_shutdown`: upstream is draining; retrying soon will hit
|
||||
* another worker. Caller should treat as transient.
|
||||
* - `gateway_timeout`: upstream took too long; backend may be unhealthy.
|
||||
* - `budget_exhausted` / `rate_limited`: client-side over-quota, retry
|
||||
* won't help until the period resets.
|
||||
* Unset for generic transport / parse errors.
|
||||
*/
|
||||
| { type: 'error'; error: string; gatewayErrorType?: 'gateway_shutdown' | 'gateway_timeout' | 'budget_exhausted' | 'rate_limited' }
|
||||
/**
|
||||
* Emitted once per request, immediately after response headers arrive,
|
||||
* for proxy-backed clients (LiteLLM Proxy etc.). Carries the physical
|
||||
* backend identity so callers can attribute the call to a specific
|
||||
* GPU pool member, distinct from the worker the request was sent through.
|
||||
*
|
||||
* Only fired when `proxy: true` was passed to OpenAICompatClient and the
|
||||
* response actually surfaced one of the proxy headers (e.g.
|
||||
* `x-litellm-model-id`). For direct (non-proxy) workers, this event is
|
||||
* never emitted. Cache hits include cacheKey; cold calls leave it null.
|
||||
*/
|
||||
| { type: 'backend'; backendId: string; cacheKey: string | null }
|
||||
| { type: 'prompt_progress'; processed: number; total: number; timeMs: number; cache: number };
|
||||
|
||||
export type PromptPreflightLogger = (line: string) => void;
|
||||
|
||||
const DEFAULT_CONTEXT_LIMIT_TOKENS = 32_000;
|
||||
const DEFAULT_PROMPT_GUARD_RATIO = 0.8;
|
||||
|
||||
function estimateRequestTokens(messages: Message[], tools?: ToolDef[]): number {
|
||||
const messageTokens = messages.reduce((total, message) => total + estimateMessageTokens(message), 0);
|
||||
const toolTokens = tools && tools.length > 0 ? estimateToolsTokens(tools) : 0;
|
||||
return messageTokens + toolTokens + 128;
|
||||
}
|
||||
|
||||
function contentChars(message: Message): number {
|
||||
if (typeof message.content === 'string') return message.content.length;
|
||||
if (!Array.isArray(message.content)) return 0;
|
||||
return message.content.reduce((total, part) => {
|
||||
if (part.type === 'text') return total + part.text.length;
|
||||
return total;
|
||||
}, 0);
|
||||
}
|
||||
|
||||
function imageCount(message: Message): number {
|
||||
if (!Array.isArray(message.content)) return 0;
|
||||
return message.content.filter((part) => part.type === 'image_url').length;
|
||||
}
|
||||
|
||||
function toolCallChars(message: Message): number {
|
||||
return (message.tool_calls ?? []).reduce((total, toolCall) => {
|
||||
return total + toolCall.id.length + toolCall.function.name.length + toolCall.function.arguments.length;
|
||||
}, 0);
|
||||
}
|
||||
|
||||
function summarizeLargestMessages(messages: Message[]): string {
|
||||
return messages
|
||||
.map((message, index) => ({
|
||||
index,
|
||||
role: message.role,
|
||||
tokens: estimateMessageTokens(message),
|
||||
contentChars: contentChars(message),
|
||||
images: imageCount(message),
|
||||
toolCallChars: toolCallChars(message),
|
||||
toolCallNames: (message.tool_calls ?? []).map((toolCall) => toolCall.function.name),
|
||||
toolName: message.name,
|
||||
}))
|
||||
.sort((a, b) => b.tokens - a.tokens)
|
||||
.slice(0, 5)
|
||||
.map((item) => {
|
||||
const names = item.toolCallNames.length > 0
|
||||
? ` calls=${item.toolCallNames.join('|')}`
|
||||
: item.toolName
|
||||
? ` name=${item.toolName}`
|
||||
: '';
|
||||
return `#${item.index}:${item.role} tokens=${item.tokens.toLocaleString()} contentChars=${item.contentChars.toLocaleString()} images=${item.images} toolCallChars=${item.toolCallChars.toLocaleString()}${names}`;
|
||||
})
|
||||
.join('; ');
|
||||
}
|
||||
|
||||
function summarizeRoleTotals(messages: Message[]): string {
|
||||
const totals = new Map<Message['role'], { count: number; tokens: number; chars: number; images: number }>();
|
||||
for (const message of messages) {
|
||||
const current = totals.get(message.role) ?? { count: 0, tokens: 0, chars: 0, images: 0 };
|
||||
current.count++;
|
||||
current.tokens += estimateMessageTokens(message);
|
||||
current.chars += contentChars(message) + toolCallChars(message);
|
||||
current.images += imageCount(message);
|
||||
totals.set(message.role, current);
|
||||
}
|
||||
return [...totals.entries()]
|
||||
.map(([role, total]) => `${role}:count=${total.count},tokens=${total.tokens.toLocaleString()},chars=${total.chars.toLocaleString()},images=${total.images}`)
|
||||
.join(' ');
|
||||
}
|
||||
|
||||
function summarizeTools(tools: ToolDef[] | undefined): string {
|
||||
if (!tools || tools.length === 0) return 'count=0 tokens=0 jsonChars=0 largest=none';
|
||||
const toolJson = JSON.stringify(tools);
|
||||
const largest = tools
|
||||
.map((tool) => ({
|
||||
name: tool.function.name,
|
||||
jsonChars: JSON.stringify(tool).length,
|
||||
}))
|
||||
.sort((a, b) => b.jsonChars - a.jsonChars)
|
||||
.slice(0, 5)
|
||||
.map((tool) => `${tool.name}:${tool.jsonChars.toLocaleString()}chars`)
|
||||
.join('|');
|
||||
return `count=${tools.length} tokens=${estimateToolsTokens(tools).toLocaleString()} jsonChars=${toolJson.length.toLocaleString()} largest=${largest}`;
|
||||
}
|
||||
|
||||
function buildPromptBreakdownLine(
|
||||
label: 'ok' | 'blocked',
|
||||
requestBody: Record<string, unknown>,
|
||||
messages: Message[],
|
||||
tools: ToolDef[] | undefined,
|
||||
estimatedPromptTokens: number,
|
||||
maxPromptTokens: number,
|
||||
contextLimitTokens: number,
|
||||
): string {
|
||||
const requestJsonChars = JSON.stringify(requestBody).length;
|
||||
const messageTokens = messages.reduce((total, message) => total + estimateMessageTokens(message), 0);
|
||||
const messageChars = messages.reduce((total, message) => total + contentChars(message) + toolCallChars(message), 0);
|
||||
const images = messages.reduce((total, message) => total + imageCount(message), 0);
|
||||
const toolsTokens = tools && tools.length > 0 ? estimateToolsTokens(tools) : 0;
|
||||
const baseOverheadTokens = Math.max(0, estimatedPromptTokens - messageTokens - toolsTokens);
|
||||
return [
|
||||
`[llm-preflight:${label}]`,
|
||||
`model=${requestBody['model'] != null ? String(requestBody['model']) : '<none>'}`,
|
||||
`estimated=${estimatedPromptTokens.toLocaleString()}`,
|
||||
`safe=${maxPromptTokens.toLocaleString()}`,
|
||||
`context=${contextLimitTokens.toLocaleString()}`,
|
||||
`requestJsonChars=${requestJsonChars.toLocaleString()}`,
|
||||
`messages=count=${messages.length},tokens=${messageTokens.toLocaleString()},chars=${messageChars.toLocaleString()},images=${images},imageTokenCost=${IMAGE_CONTENT_TOKENS}`,
|
||||
`tools=${summarizeTools(tools)}`,
|
||||
`baseOverheadTokens=${baseOverheadTokens.toLocaleString()}`,
|
||||
`roles=[${summarizeRoleTotals(messages)}]`,
|
||||
`largestMessages=[${summarizeLargestMessages(messages)}]`,
|
||||
].join(' ');
|
||||
}
|
||||
|
||||
function logPromptBreakdown(
|
||||
label: 'ok' | 'blocked',
|
||||
requestBody: Record<string, unknown>,
|
||||
messages: Message[],
|
||||
tools: ToolDef[] | undefined,
|
||||
estimatedPromptTokens: number,
|
||||
maxPromptTokens: number,
|
||||
contextLimitTokens: number,
|
||||
onPromptPreflight?: PromptPreflightLogger,
|
||||
): void {
|
||||
const line = buildPromptBreakdownLine(
|
||||
label,
|
||||
requestBody,
|
||||
messages,
|
||||
tools,
|
||||
estimatedPromptTokens,
|
||||
maxPromptTokens,
|
||||
contextLimitTokens,
|
||||
);
|
||||
onPromptPreflight?.(line);
|
||||
if (label === 'blocked') {
|
||||
logger.warn(line);
|
||||
} else {
|
||||
logger.info(line);
|
||||
}
|
||||
}
|
||||
|
||||
function buildPromptTooLargeError(estimatedTokens: number, maxPromptTokens: number, contextLimitTokens: number, ratio: number): string {
|
||||
return [
|
||||
'LLM request blocked before send:',
|
||||
`estimated prompt size ${estimatedTokens.toLocaleString()} tokens exceeds safe limit ${maxPromptTokens.toLocaleString()} tokens`,
|
||||
`(${Math.round(ratio * 100)}% of context ${contextLimitTokens.toLocaleString()}).`,
|
||||
'Narrow the requested content with Read(offset/limit), Read(byte_offset/byte_length), Grep, or targeted Bash before continuing.',
|
||||
].join(' ');
|
||||
}
|
||||
|
||||
// SSE チャンク内の tool_call delta を蓄積するための内部型
|
||||
interface ToolCallAccumulator {
|
||||
id: string;
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string;
|
||||
};
|
||||
}
|
||||
|
||||
export interface OpenAICompatClientOptions {
|
||||
/**
|
||||
* When true, this client treats its endpoint as an LLM gateway / proxy
|
||||
* (e.g. LiteLLM Proxy). The chat() stream will emit a one-shot 'backend'
|
||||
* event after the response headers arrive, carrying the physical backend
|
||||
* identity derived from `x-litellm-model-id` (and cacheKey from
|
||||
* `x-litellm-cache-key` when present).
|
||||
*
|
||||
* Direct (non-proxy) workers leave this false; no 'backend' event is
|
||||
* ever emitted in that mode.
|
||||
*/
|
||||
proxy?: boolean;
|
||||
}
|
||||
|
||||
export class OpenAICompatClient {
|
||||
private retryConfig: ProviderRetryConfig;
|
||||
readonly timeoutMs: number;
|
||||
private readonly proxy: boolean;
|
||||
|
||||
constructor(
|
||||
private baseUrl: string,
|
||||
private model: string | undefined,
|
||||
private apiKey?: string,
|
||||
retryConfig?: ProviderRetryConfig,
|
||||
timeoutMs?: number,
|
||||
private contextLimitTokens: number = DEFAULT_CONTEXT_LIMIT_TOKENS,
|
||||
private promptGuardRatio: number = DEFAULT_PROMPT_GUARD_RATIO,
|
||||
private onPromptPreflight?: PromptPreflightLogger,
|
||||
options?: OpenAICompatClientOptions,
|
||||
) {
|
||||
this.retryConfig = retryConfig ?? getDefaultProviderRetryConfig();
|
||||
this.timeoutMs = timeoutMs ?? 10 * 60 * 1000; // default: 10 minutes
|
||||
this.proxy = options?.proxy === true;
|
||||
}
|
||||
|
||||
private buildAbortErrorMessage(externalSignal?: AbortSignal): string {
|
||||
if (externalSignal?.aborted) {
|
||||
return 'Request cancelled by caller';
|
||||
}
|
||||
const mins = Math.round(this.timeoutMs / 60000);
|
||||
return `Request timed out (${mins} minutes)`;
|
||||
}
|
||||
|
||||
async *chat(messages: Message[], tools?: ToolDef[], externalSignal?: AbortSignal): AsyncGenerator<LLMEvent> {
|
||||
const controller = new AbortController();
|
||||
// アイドルタイムアウト: チャンク受信のたびにリセットされる
|
||||
let timeoutId = setTimeout(() => controller.abort(), this.timeoutMs);
|
||||
const resetIdleTimeout = () => {
|
||||
clearTimeout(timeoutId);
|
||||
timeoutId = setTimeout(() => controller.abort(), this.timeoutMs);
|
||||
};
|
||||
|
||||
let onExternalAbort: (() => void) | undefined;
|
||||
if (externalSignal) {
|
||||
if (externalSignal.aborted) {
|
||||
clearTimeout(timeoutId);
|
||||
yield { type: 'error', error: 'Request cancelled by caller' };
|
||||
return;
|
||||
}
|
||||
onExternalAbort = () => controller.abort();
|
||||
externalSignal.addEventListener('abort', onExternalAbort, { once: true });
|
||||
}
|
||||
|
||||
try {
|
||||
const headers: Record<string, string> = {
|
||||
'Content-Type': 'application/json',
|
||||
};
|
||||
if (this.apiKey) {
|
||||
headers['Authorization'] = `Bearer ${this.apiKey}`;
|
||||
}
|
||||
|
||||
const body: Record<string, unknown> = {
|
||||
messages,
|
||||
stream: true,
|
||||
stream_options: { include_usage: true },
|
||||
};
|
||||
if (this.model) {
|
||||
body['model'] = this.model;
|
||||
}
|
||||
if (tools && tools.length > 0) {
|
||||
body['tools'] = tools;
|
||||
}
|
||||
// Block oversized prompts before the HTTP request so callers see a
|
||||
// structured error instead of an opaque HTTP 400. The runtime context
|
||||
// limit is fetched per-model (see fetchOllamaContextLimit) and passed
|
||||
// in via contextLimitTokens, so we trust it directly here.
|
||||
const maxPromptTokens = Math.floor(this.contextLimitTokens * this.promptGuardRatio);
|
||||
const estimatedPromptTokens = estimateRequestTokens(messages, tools);
|
||||
if (estimatedPromptTokens > maxPromptTokens) {
|
||||
logPromptBreakdown('blocked', body, messages, tools, estimatedPromptTokens, maxPromptTokens, this.contextLimitTokens, this.onPromptPreflight);
|
||||
const error = buildPromptTooLargeError(estimatedPromptTokens, maxPromptTokens, this.contextLimitTokens, this.promptGuardRatio);
|
||||
logger.warn(`OpenAICompatClient: ${error}`);
|
||||
yield { type: 'error', error };
|
||||
return;
|
||||
}
|
||||
logPromptBreakdown('ok', body, messages, tools, estimatedPromptTokens, maxPromptTokens, this.contextLimitTokens, this.onPromptPreflight);
|
||||
|
||||
const maxAttempts = Math.max(1, this.retryConfig.maxAttempts || 1);
|
||||
let lastErrorMessage = '';
|
||||
|
||||
for (let attempt = 1; attempt <= maxAttempts; attempt++) {
|
||||
let response: Response | null = null;
|
||||
|
||||
try {
|
||||
response = await fetch(`${this.baseUrl}/chat/completions`, {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: JSON.stringify(body),
|
||||
signal: controller.signal,
|
||||
});
|
||||
} catch (err) {
|
||||
if ((err as Error)?.name === 'AbortError') {
|
||||
logger.error('OpenAICompatClient: request timed out');
|
||||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal) };
|
||||
return;
|
||||
}
|
||||
|
||||
lastErrorMessage = err instanceof Error ? err.message : String(err);
|
||||
if (!isTransientFetchError(err) || attempt >= maxAttempts) {
|
||||
logger.error(`OpenAICompatClient: fetch failed: ${lastErrorMessage}`);
|
||||
yield { type: 'error', error: `Connection error: ${lastErrorMessage}` };
|
||||
return;
|
||||
}
|
||||
|
||||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||||
logger.warn(`OpenAICompatClient: transient fetch error on attempt ${attempt}/${maxAttempts}: ${lastErrorMessage}; retrying in ${delayMs}ms`);
|
||||
if (!(await waitForRetry(delayMs, controller.signal))) {
|
||||
logger.error('OpenAICompatClient: request timed out');
|
||||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal) };
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// レスポンスヘッダー受信 = サーバーが応答開始 → アイドルタイマーリセット
|
||||
resetIdleTimeout();
|
||||
|
||||
if (!response.ok) {
|
||||
let errorBody = '';
|
||||
try {
|
||||
errorBody = await response.text();
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
lastErrorMessage = `HTTP ${response.status}: ${errorBody}`;
|
||||
|
||||
if (!isRetryableHttpStatus(response.status, this.retryConfig) || attempt >= maxAttempts) {
|
||||
logger.error(`OpenAICompatClient: ${lastErrorMessage}`);
|
||||
yield { type: 'error', error: lastErrorMessage };
|
||||
return;
|
||||
}
|
||||
|
||||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||||
logger.warn(`OpenAICompatClient: retryable HTTP ${response.status} on attempt ${attempt}/${maxAttempts}; retrying in ${delayMs}ms`);
|
||||
if (!(await waitForRetry(delayMs, controller.signal))) {
|
||||
logger.error('OpenAICompatClient: request timed out');
|
||||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal) };
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!response.body) {
|
||||
yield { type: 'error', error: 'Response body is null' };
|
||||
return;
|
||||
}
|
||||
|
||||
// Proxy backend identification: surface the physical backend id
|
||||
// (and optional cache hit key) so the worker / agent-loop can
|
||||
// attribute each call to a specific GPU pool member.
|
||||
// We trust the very first headers we get for this request; once
|
||||
// a backend is selected for a streaming completion, LiteLLM
|
||||
// doesn't switch mid-stream. See:
|
||||
// docs/superpowers/specs/2026-05-18-multi-team-gpu-pool-and-node-status-design.md
|
||||
if (this.proxy) {
|
||||
// Trim whitespace so that whitespace-only header values are
|
||||
// treated as missing. Without trim, a header like
|
||||
// `x-litellm-model-id: " "` would emit " " as the backend
|
||||
// id, but config-api's /v1/models reader trims its side — the
|
||||
// two ids would never key-match and the UI Pet mapping would
|
||||
// mysteriously misbehave.
|
||||
const rawBackendId = response.headers.get('x-litellm-model-id');
|
||||
const backendId = rawBackendId ? rawBackendId.trim() : '';
|
||||
if (backendId.length > 0) {
|
||||
const rawCacheKey = response.headers.get('x-litellm-cache-key');
|
||||
const cacheKey = rawCacheKey ? rawCacheKey.trim() : '';
|
||||
yield { type: 'backend', backendId, cacheKey: cacheKey.length > 0 ? cacheKey : null };
|
||||
}
|
||||
}
|
||||
|
||||
// ストリーム読み取り（リトライループ内）
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
|
||||
// tool_calls を index ごとに蓄積
|
||||
const toolCallAccumulators = new Map<number, ToolCallAccumulator>();
|
||||
let usage: { prompt_tokens: number; completion_tokens: number } | undefined;
|
||||
let buffer = '';
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
resetIdleTimeout();
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
|
||||
// 行単位で処理
|
||||
const lines = buffer.split('\n');
|
||||
// 最後の要素は不完全な行の可能性があるのでバッファに残す
|
||||
buffer = lines.pop() ?? '';
|
||||
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
if (!trimmed || !trimmed.startsWith('data: ')) continue;
|
||||
|
||||
const data = trimmed.slice('data: '.length);
|
||||
|
||||
if (data === '[DONE]') {
|
||||
// usage 付きで done を emit
|
||||
yield { type: 'done', usage };
|
||||
return;
|
||||
}
|
||||
|
||||
let chunk: Record<string, unknown>;
|
||||
try {
|
||||
chunk = JSON.parse(data) as Record<string, unknown>;
|
||||
} catch (err) {
|
||||
logger.warn(`OpenAICompatClient: failed to parse SSE chunk: ${data}`);
|
||||
continue;
|
||||
}
|
||||
|
||||
// AAO Gateway / LiteLLM sentinel error event:
|
||||
// data: {"error":{"type":"gateway_shutdown","message":"..."}}
|
||||
// gateway_shutdown / gateway_timeout は他 worker に retry すれば
|
||||
// 通る可能性が高い transient エラーなので、generic stream error と
|
||||
// 区別して呼び出し元に伝える。
|
||||
if (chunk['error'] && typeof chunk['error'] === 'object') {
|
||||
const errObj = chunk['error'] as { type?: unknown; message?: unknown };
|
||||
const knownTypes = new Set(['gateway_shutdown', 'gateway_timeout', 'budget_exhausted', 'rate_limited']);
|
||||
if (typeof errObj.type === 'string' && knownTypes.has(errObj.type)) {
|
||||
const msg = typeof errObj.message === 'string' ? errObj.message : errObj.type;
|
||||
logger.warn(`OpenAICompatClient: gateway sentinel error mid-stream type=${errObj.type} msg=${msg}`);
|
||||
yield {
|
||||
type: 'error',
|
||||
error: `gateway ${errObj.type}: ${msg}`,
|
||||
gatewayErrorType: errObj.type as 'gateway_shutdown' | 'gateway_timeout' | 'budget_exhausted' | 'rate_limited',
|
||||
};
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// usage (stream_options で末尾チャンクに付く)
|
||||
if (chunk['usage'] != null) {
|
||||
const u = chunk['usage'] as Record<string, unknown>;
|
||||
usage = {
|
||||
prompt_tokens: (u['prompt_tokens'] as number) ?? 0,
|
||||
completion_tokens: (u['completion_tokens'] as number) ?? 0,
|
||||
};
|
||||
}
|
||||
|
||||
// llama-server prompt_progress (prompt eval 進捗)
|
||||
const pp = chunk['prompt_progress'] as Record<string, unknown> | undefined;
|
||||
if (pp && typeof pp['processed'] === 'number' && typeof pp['total'] === 'number') {
|
||||
yield {
|
||||
type: 'prompt_progress',
|
||||
processed: pp['processed'] as number,
|
||||
total: pp['total'] as number,
|
||||
timeMs: (pp['time_ms'] as number) ?? 0,
|
||||
cache: (pp['cache'] as number) ?? 0,
|
||||
};
|
||||
}
|
||||
|
||||
const choices = chunk['choices'] as Array<Record<string, unknown>> | undefined;
|
||||
if (!choices || choices.length === 0) continue;
|
||||
|
||||
const choice = choices[0] as Record<string, unknown>;
|
||||
const delta = choice['delta'] as Record<string, unknown> | undefined;
|
||||
const finishReason = choice['finish_reason'] as string | null | undefined;
|
||||
|
||||
if (delta) {
|
||||
// reasoning_content (thinking models) — スキップしてログのみ
|
||||
const reasoning = delta['reasoning_content'];
|
||||
if (typeof reasoning === 'string' && reasoning.length > 0) {
|
||||
logger.debug(`OpenAICompatClient: reasoning_content (${reasoning.length} chars), skipping`);
|
||||
}
|
||||
|
||||
// テキストチャンク
|
||||
const content = delta['content'];
|
||||
if (typeof content === 'string' && content.length > 0) {
|
||||
yield { type: 'text', text: content };
|
||||
}
|
||||
|
||||
// tool_calls delta の蓄積
|
||||
const deltaToolCalls = delta['tool_calls'] as Array<Record<string, unknown>> | undefined;
|
||||
if (deltaToolCalls) {
|
||||
for (const tc of deltaToolCalls) {
|
||||
const index = tc['index'] as number;
|
||||
const fn = tc['function'] as Record<string, unknown> | undefined;
|
||||
|
||||
if (!toolCallAccumulators.has(index)) {
|
||||
toolCallAccumulators.set(index, {
|
||||
id: (tc['id'] as string) ?? '',
|
||||
type: 'function',
|
||||
function: {
|
||||
name: (fn?.['name'] as string) ?? '',
|
||||
arguments: (fn?.['arguments'] as string) ?? '',
|
||||
},
|
||||
});
|
||||
} else {
|
||||
const acc = toolCallAccumulators.get(index)!;
|
||||
// id が来た場合は上書き（最初のチャンクのみ）
|
||||
if (tc['id']) acc.id = tc['id'] as string;
|
||||
if (fn?.['name']) acc.function.name += fn['name'] as string;
|
||||
if (fn?.['arguments']) acc.function.arguments += fn['arguments'] as string;
|
||||
}
|
||||
|
||||
// Live streaming: surface the FULL accumulated arguments
|
||||
// so far (a snapshot) whenever a new args chunk arrives.
|
||||
// Sending the whole prefix (not just the latest piece)
|
||||
// lets a client that attaches mid-generation still get the
|
||||
// opening JSON structure. The aggregated tool_use is still
|
||||
// emitted later on finish_reason.
|
||||
const argsChunk = (fn?.['arguments'] as string) ?? '';
|
||||
if (argsChunk.length > 0) {
|
||||
const acc = toolCallAccumulators.get(index)!;
|
||||
yield {
|
||||
type: 'tool_use_delta',
|
||||
index,
|
||||
callId: acc.id,
|
||||
name: acc.function.name,
|
||||
chunk: acc.function.arguments,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// tool_calls が完了したら emit
|
||||
if (finishReason === 'tool_calls') {
|
||||
const sortedIndices = Array.from(toolCallAccumulators.keys()).sort((a, b) => a - b);
|
||||
for (const idx of sortedIndices) {
|
||||
const acc = toolCallAccumulators.get(idx)!;
|
||||
let input: Record<string, unknown> = {};
|
||||
try {
|
||||
input = JSON.parse(acc.function.arguments) as Record<string, unknown>;
|
||||
} catch {
|
||||
logger.warn(`OpenAICompatClient: failed to parse tool arguments: ${acc.function.arguments}`);
|
||||
}
|
||||
yield {
|
||||
type: 'tool_use',
|
||||
id: acc.id,
|
||||
name: acc.function.name,
|
||||
input,
|
||||
};
|
||||
}
|
||||
toolCallAccumulators.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
if ((err as Error)?.name === 'AbortError') {
|
||||
logger.error('OpenAICompatClient: request timed out');
|
||||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal) };
|
||||
return;
|
||||
}
|
||||
|
||||
// 一時的なストリームエラー — 試行回数が残っていればリトライ
|
||||
if (attempt >= maxAttempts) {
|
||||
logger.error(`OpenAICompatClient: stream read error: ${message}`);
|
||||
yield { type: 'error', error: `Stream error: ${message}` };
|
||||
return;
|
||||
}
|
||||
|
||||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||||
logger.warn(`OpenAICompatClient: stream read error on attempt ${attempt}/${maxAttempts}: ${message}; retrying in ${delayMs}ms`);
|
||||
if (delayMs > 0 && !(await waitForRetry(delayMs, controller.signal))) {
|
||||
logger.error('OpenAICompatClient: request timed out during retry wait');
|
||||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal) };
|
||||
return;
|
||||
}
|
||||
continue;
|
||||
} finally {
|
||||
reader.releaseLock();
|
||||
}
|
||||
|
||||
// [DONE] なしにストリームが終了した場合
|
||||
yield { type: 'done', usage };
|
||||
return;
|
||||
}
|
||||
|
||||
// 全試行が失敗した場合
|
||||
yield { type: 'error', error: lastErrorMessage || 'Unknown request error' };
|
||||
} finally {
|
||||
clearTimeout(timeoutId);
|
||||
if (onExternalAbort && externalSignal) {
|
||||
externalSignal.removeEventListener('abort', onExternalAbort);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function isRetryableHttpStatus(status: number, retryConfig: ProviderRetryConfig): boolean {
|
||||
return retryConfig.retryableStatus.includes(status);
|
||||
}
|
||||
|
||||
function isTransientFetchError(err: unknown): boolean {
|
||||
return err instanceof Error && err.name !== 'AbortError';
|
||||
}
|
||||
|
||||
function getRetryDelayMs(retryConfig: ProviderRetryConfig, attempt: number): number {
|
||||
const delays = retryConfig.backoffMs;
|
||||
if (!Array.isArray(delays) || delays.length === 0) return 0;
|
||||
const index = Math.min(Math.max(attempt - 1, 0), delays.length - 1);
|
||||
return Math.max(0, delays[index] ?? 0);
|
||||
}
|
||||
|
||||
function waitForRetry(delayMs: number, signal: AbortSignal): Promise<boolean> {
|
||||
if (delayMs <= 0) return Promise.resolve(true);
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const timeout = setTimeout(() => {
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
resolve(true);
|
||||
}, delayMs);
|
||||
|
||||
const onAbort = () => {
|
||||
clearTimeout(timeout);
|
||||
signal.removeEventListener('abort', onAbort);
|
||||
resolve(false);
|
||||
};
|
||||
|
||||
signal.addEventListener('abort', onAbort, { once: true });
|
||||
});
|
||||
}
|
||||
|
||||
// ツール実行結果を Message に変換
|
||||
export function toolResultMessage(toolCallId: string, result: string): Message {
|
||||
return { role: 'tool', content: result, tool_call_id: toolCallId };
|
||||
}
|
||||
|
||||
// assistant の tool_calls を Message に変換
|
||||
export function assistantToolCallMessage(toolCalls: ToolCall[]): Message {
|
||||
return { role: 'assistant', tool_calls: toolCalls };
|
||||
}
|
||||
Reference in New Issue
Block a user