924 lines
40 KiB
TypeScript
924 lines
40 KiB
TypeScript
import { getDefaultProviderRetryConfig, type ProviderRetryConfig } from '../config.js';
|
||
import { logger } from '../logger.js';
|
||
import { recordLlmUsage } from './usage-recorder.js';
|
||
import {
|
||
IMAGE_CONTENT_TOKENS,
|
||
estimateMessageTokens,
|
||
estimateRequestTokens,
|
||
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
|
||
};
|
||
}
|
||
|
||
/**
|
||
* Machine-readable classification of an LLM request failure. Travels with
|
||
* error/retry events so downstream layers (agent-loop abort messages, the
|
||
* per-task LLM call log, the UI) never have to string-parse error text.
|
||
* - preflight_block: client-side prompt-size guard refused to send
|
||
* - cancelled: external AbortSignal (user cancel / job deadline)
|
||
* - idle_timeout: no chunk received for timeoutMs
|
||
* - hard_cap: total stream duration exceeded maxStreamMs
|
||
* - connection: fetch/network failure
|
||
* - http: non-2xx response (see httpStatus)
|
||
* - stream: SSE read error mid-stream
|
||
* - gateway_*: AAO Gateway sentinel errors (see gatewayErrorType docs)
|
||
*/
|
||
export type LlmErrorClass =
|
||
| 'preflight_block'
|
||
| 'cancelled'
|
||
| 'idle_timeout'
|
||
| 'hard_cap'
|
||
| 'connection'
|
||
| 'http'
|
||
| 'stream'
|
||
| 'gateway_shutdown'
|
||
| 'gateway_timeout'
|
||
| 'budget_exhausted'
|
||
| 'rate_limited'
|
||
| 'unknown';
|
||
|
||
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; errorClass?: LlmErrorClass; httpStatus?: number; gatewayErrorType?: 'gateway_shutdown' | 'gateway_timeout' | 'budget_exhausted' | 'rate_limited' }
|
||
/**
|
||
* Emitted just before a client-internal retry sleep (transient fetch
|
||
* error, retryable HTTP status, or mid-stream read error). Lets the
|
||
* caller surface "retrying 2/3: HTTP 500" instead of silence during
|
||
* the backoff wait. `attempt` is the attempt that just FAILED.
|
||
*/
|
||
| { type: 'retry'; attempt: number; maxAttempts: number; reason: string; errorClass: LlmErrorClass; httpStatus?: number; delayMs: number }
|
||
/**
|
||
* Per-chunk `reasoning_content` size (thinking models). The content
|
||
* itself is intentionally NOT forwarded — only char counts, so the UI
|
||
* can show "thinking…" liveness without leaking chain-of-thought into
|
||
* transcripts. Consumers accumulate.
|
||
*/
|
||
| { type: 'thinking'; chars: number }
|
||
/**
|
||
* 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;
|
||
|
||
// The estimate MUST stay byte-identical to what the agent-loop prompt guard
|
||
// computes (estimateRequestTokens in token-estimate.ts). A drift between the
|
||
// two creates a band of prompt sizes the guard passes but this preflight
|
||
// blocks — the error-recovery path then finds nothing to shrink and the loop
|
||
// resends the identical request until maxIterations.
|
||
|
||
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;
|
||
};
|
||
}
|
||
|
||
/**
|
||
* Emit accumulated tool calls as `tool_use` events (sorted by index) and
|
||
* clear the accumulator. Called both on `finish_reason === 'tool_calls'` and
|
||
* at stream end — some OpenAI-compat backends finish a forced/named tool call
|
||
* with finish_reason 'stop', so draining at the done boundary keeps the call
|
||
* from being silently dropped. Returns an empty array when nothing is pending,
|
||
* so the done-site flush is a no-op for the normal 'tool_calls' path (the map
|
||
* is already cleared).
|
||
*/
|
||
function drainToolCalls(accumulators: Map<number, ToolCallAccumulator>): LLMEvent[] {
|
||
if (accumulators.size === 0) return [];
|
||
const events: LLMEvent[] = [];
|
||
const sortedIndices = Array.from(accumulators.keys()).sort((a, b) => a - b);
|
||
for (const idx of sortedIndices) {
|
||
const acc = accumulators.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}`);
|
||
}
|
||
events.push({ type: 'tool_use', id: acc.id, name: acc.function.name, input });
|
||
}
|
||
accumulators.clear();
|
||
return events;
|
||
}
|
||
|
||
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;
|
||
/**
|
||
* Hard wall-clock ceiling (ms) for a single chat() call, INCLUDING retries.
|
||
* Unlike the idle timeout (which resets on every chunk), this timer never
|
||
* resets — so a degenerate generation that keeps emitting tokens without
|
||
* ever stopping (runaway repetition, no stop token) is still aborted.
|
||
* `0` disables it. When omitted, the constructor defaults to 2× the idle
|
||
* timeout so every client is bounded even if the caller forgets to set it.
|
||
*/
|
||
maxStreamMs?: number;
|
||
/**
|
||
* When true, add `return_progress: true` to the request body so
|
||
* llama.cpp's llama-server streams `prompt_progress` chunks during
|
||
* prompt evaluation (surfaced as 'prompt_progress' events). Opt-in
|
||
* per worker: non-llama.cpp backends (vLLM, some gateways) may reject
|
||
* unknown body fields, so this must never default to on.
|
||
*/
|
||
requestPromptProgress?: boolean;
|
||
}
|
||
|
||
/**
|
||
* Per-call attribution context. Threaded from each call site so the
|
||
* usage recorder can attribute the completion to a MAESTRO user. Absent
|
||
* userId falls back to the 'system' sentinel (never NULL).
|
||
*/
|
||
export interface LlmCallContext {
|
||
userId?: string;
|
||
}
|
||
|
||
/**
|
||
* Per-call request-shaping overrides. Used by callers that need to force a
|
||
* tool (reflection's forced submit_reflection) or pin sampling temperature.
|
||
* Kept off the hot agent path (which leaves these unset).
|
||
*/
|
||
export interface LlmRequestOptions {
|
||
temperature?: number;
|
||
/** OpenAI tool_choice (e.g. `{ type: 'function', function: { name } }`). */
|
||
toolChoice?: unknown;
|
||
}
|
||
|
||
export class OpenAICompatClient {
|
||
private retryConfig: ProviderRetryConfig;
|
||
readonly timeoutMs: number;
|
||
/** Hard wall-clock ceiling per chat() call (ms); 0 = disabled. See OpenAICompatClientOptions.maxStreamMs. */
|
||
readonly maxStreamMs: 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
|
||
// Hard total-duration cap. Default to 2× the idle timeout so a runaway
|
||
// stream that keeps emitting tokens (never tripping the idle timer) is
|
||
// still bounded. `?? ` not `||` so an explicit 0 stays 0 (disabled).
|
||
this.maxStreamMs = options?.maxStreamMs ?? this.timeoutMs * 2;
|
||
this.proxy = options?.proxy === true;
|
||
this.requestPromptProgress = options?.requestPromptProgress === true;
|
||
}
|
||
|
||
private readonly requestPromptProgress: boolean;
|
||
|
||
private buildAbortErrorMessage(externalSignal?: AbortSignal, hardCapHit = false): string {
|
||
if (externalSignal?.aborted) {
|
||
return 'Request cancelled by caller';
|
||
}
|
||
if (hardCapHit) {
|
||
const mins = Math.round(this.maxStreamMs / 60000);
|
||
return `Request exceeded maximum stream duration (${mins} minutes)`;
|
||
}
|
||
const mins = Math.round(this.timeoutMs / 60000);
|
||
return `Request timed out (${mins} minutes)`;
|
||
}
|
||
|
||
/** Classify an AbortError raised inside chat() into its actual trigger. */
|
||
private classifyAbort(externalSignal?: AbortSignal, hardCapHit = false): LlmErrorClass {
|
||
if (externalSignal?.aborted) return 'cancelled';
|
||
if (hardCapHit) return 'hard_cap';
|
||
return 'idle_timeout';
|
||
}
|
||
|
||
/**
|
||
* Backend the next request should prefer (gateway sticky routing for
|
||
* KV-cache reuse). Updated by the worker whenever the resolved backend
|
||
* changes; per-client so concurrent jobs never share affinity.
|
||
*/
|
||
private preferredBackendId: string | null = null;
|
||
|
||
setPreferredBackendId(backendId: string | null): void {
|
||
this.preferredBackendId = backendId;
|
||
}
|
||
|
||
/**
|
||
* Record one successful completion to the per-user daily usage ledger.
|
||
* Called from the single done funnel (both the `[DONE]` and EOF exits)
|
||
* so the two terminal paths can never double-count. `source` is the
|
||
* client's proxy flag, `model` is the first observed chunk.model (routing
|
||
* key fallback), `route` is the gateway backendId (proxy) or endpoint host
|
||
* (direct). Never records on abort / timeout / error (those don't `done`).
|
||
*/
|
||
private finalizeDone(
|
||
usage: { prompt_tokens: number; completion_tokens: number } | undefined,
|
||
observedModel: string,
|
||
observedBackendId: string,
|
||
context?: LlmCallContext,
|
||
): void {
|
||
const source: 'gateway' | 'direct' = this.proxy ? 'gateway' : 'direct';
|
||
const model = observedModel || this.model || 'unknown';
|
||
let route = 'unknown';
|
||
if (this.proxy) {
|
||
route = observedBackendId || 'unknown';
|
||
} else {
|
||
try {
|
||
route = new URL(this.baseUrl).host || 'unknown';
|
||
} catch {
|
||
route = 'unknown';
|
||
}
|
||
}
|
||
recordLlmUsage({
|
||
userId: context?.userId || 'system',
|
||
source,
|
||
model,
|
||
route,
|
||
tokensIn: usage?.prompt_tokens ?? 0,
|
||
tokensOut: usage?.completion_tokens ?? 0,
|
||
});
|
||
}
|
||
|
||
async *chat(messages: Message[], tools?: ToolDef[], externalSignal?: AbortSignal, context?: LlmCallContext, requestOptions?: LlmRequestOptions): AsyncGenerator<LLMEvent> {
|
||
const controller = new AbortController();
|
||
// アイドルタイムアウト: チャンク受信のたびにリセットされる
|
||
let timeoutId = setTimeout(() => controller.abort(), this.timeoutMs);
|
||
const resetIdleTimeout = () => {
|
||
clearTimeout(timeoutId);
|
||
timeoutId = setTimeout(() => controller.abort(), this.timeoutMs);
|
||
};
|
||
|
||
// 総時間ハードキャップ: チャンクが届き続けてもリセットしない。
|
||
// 停止トークンを出さずにトークンを吐き続ける暴走（反復生成など）は
|
||
// アイドルタイムアウトでは止まらないため、この上限で必ず打ち切る。
|
||
let hardCapHit = false;
|
||
let hardCapId: ReturnType<typeof setTimeout> | undefined;
|
||
if (this.maxStreamMs > 0) {
|
||
hardCapId = setTimeout(() => {
|
||
hardCapHit = true;
|
||
controller.abort();
|
||
}, this.maxStreamMs);
|
||
// Don't let this safety timer keep the process alive on its own. If a
|
||
// caller abandons the generator (e.g. a Promise.race timeout that never
|
||
// resumes the for-await), its `finally` never runs to clear the timer;
|
||
// unref'ing means an otherwise-idle process can still exit instead of
|
||
// waiting out maxStreamMs. During a live stream the socket keeps the
|
||
// loop alive, so the timer still fires normally.
|
||
(hardCapId as { unref?: () => void }).unref?.();
|
||
}
|
||
|
||
let onExternalAbort: (() => void) | undefined;
|
||
if (externalSignal) {
|
||
if (externalSignal.aborted) {
|
||
clearTimeout(timeoutId);
|
||
if (hardCapId) clearTimeout(hardCapId);
|
||
yield { type: 'error', error: 'Request cancelled by caller', errorClass: 'cancelled' };
|
||
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}`;
|
||
}
|
||
// Sticky routing hint: ask the gateway to keep serving this job from
|
||
// the backend that already holds its KV cache. The gateway only honors
|
||
// it while that backend is online with free capacity; otherwise it
|
||
// re-routes normally. Direct (non-proxy) backends ignore the header.
|
||
if (this.preferredBackendId) {
|
||
headers['x-aao-preferred-backend'] = this.preferredBackendId;
|
||
}
|
||
|
||
const body: Record<string, unknown> = {
|
||
messages,
|
||
stream: true,
|
||
stream_options: { include_usage: true },
|
||
};
|
||
if (this.requestPromptProgress) {
|
||
// llama.cpp llama-server extension: stream prompt-eval progress
|
||
// chunks. Opt-in per worker (see OpenAICompatClientOptions).
|
||
body['return_progress'] = true;
|
||
}
|
||
if (this.model) {
|
||
body['model'] = this.model;
|
||
}
|
||
if (tools && tools.length > 0) {
|
||
body['tools'] = tools;
|
||
}
|
||
if (requestOptions?.temperature != null) {
|
||
body['temperature'] = requestOptions.temperature;
|
||
}
|
||
if (requestOptions?.toolChoice != null) {
|
||
body['tool_choice'] = requestOptions.toolChoice;
|
||
}
|
||
// 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, errorClass: 'preflight_block' };
|
||
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;
|
||
// Usage attribution captured during this attempt's stream. Reset
|
||
// per attempt; only the attempt that reaches `done` records.
|
||
let observedModel = '';
|
||
let observedBackendId = '';
|
||
// Whether we saw at least one real (non-error) SSE chunk. An EOF
|
||
// that closes the stream without any chunk and without [DONE] is an
|
||
// abnormal completion we must NOT record as a request (issue #498).
|
||
let sawChunk = false;
|
||
|
||
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, hardCapHit), errorClass: this.classifyAbort(externalSignal, hardCapHit) };
|
||
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}`, errorClass: 'connection' };
|
||
return;
|
||
}
|
||
|
||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||
logger.warn(`OpenAICompatClient: transient fetch error on attempt ${attempt}/${maxAttempts}: ${lastErrorMessage}; retrying in ${delayMs}ms`);
|
||
yield { type: 'retry', attempt, maxAttempts, reason: lastErrorMessage, errorClass: 'connection', delayMs };
|
||
if (!(await waitForRetry(delayMs, controller.signal))) {
|
||
logger.error('OpenAICompatClient: request timed out');
|
||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal, hardCapHit), errorClass: this.classifyAbort(externalSignal, hardCapHit) };
|
||
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, errorClass: 'http', httpStatus: response.status };
|
||
return;
|
||
}
|
||
|
||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||
logger.warn(`OpenAICompatClient: retryable HTTP ${response.status} on attempt ${attempt}/${maxAttempts}; retrying in ${delayMs}ms`);
|
||
yield { type: 'retry', attempt, maxAttempts, reason: `HTTP ${response.status}`, errorClass: 'http', httpStatus: response.status, delayMs };
|
||
if (!(await waitForRetry(delayMs, controller.signal))) {
|
||
logger.error('OpenAICompatClient: request timed out');
|
||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal, hardCapHit), errorClass: this.classifyAbort(externalSignal, hardCapHit) };
|
||
return;
|
||
}
|
||
continue;
|
||
}
|
||
|
||
if (!response.body) {
|
||
yield { type: 'error', error: 'Response body is null', errorClass: 'connection' };
|
||
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) {
|
||
observedBackendId = backendId;
|
||
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]') {
|
||
// Flush any tool calls the backend left un-finished (some
|
||
// OpenAI-compat servers end a forced/named tool call with
|
||
// finish_reason 'stop' instead of 'tool_calls'). Without this
|
||
// the accumulated call would be silently dropped.
|
||
yield* drainToolCalls(toolCallAccumulators);
|
||
// usage 付きで done を emit
|
||
this.finalizeDone(usage, observedModel, observedBackendId, context);
|
||
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;
|
||
}
|
||
|
||
// Real model name for usage attribution. The gateway passes
|
||
// chunks through byte-for-byte, so chunk.model is the actual
|
||
// backend model for both direct and gateway paths. First
|
||
// non-empty value wins.
|
||
if (!observedModel) {
|
||
const m = chunk['model'];
|
||
if (typeof m === 'string' && m.length > 0) observedModel = m;
|
||
}
|
||
|
||
// 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}`,
|
||
errorClass: errObj.type as LlmErrorClass,
|
||
gatewayErrorType: errObj.type as 'gateway_shutdown' | 'gateway_timeout' | 'budget_exhausted' | 'rate_limited',
|
||
};
|
||
return;
|
||
}
|
||
// Unknown error.type: keep parsing (it may be followed by real
|
||
// content). It is NOT a real chunk, so it doesn't flip
|
||
// `sawChunk` — an error-only stream that then EOFs without
|
||
// [DONE] stays unrecorded (issue #498).
|
||
}
|
||
|
||
// usage (stream_options で末尾チャンクに付く)
|
||
if (chunk['usage'] != null) {
|
||
sawChunk = true; // a real completion payload
|
||
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;
|
||
sawChunk = true; // a real content/finish chunk
|
||
|
||
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) — 中身は流さず文字数のみ
|
||
// 'thinking' イベントで通知（UI の「思考中」表示用）。
|
||
const reasoning = delta['reasoning_content'];
|
||
if (typeof reasoning === 'string' && reasoning.length > 0) {
|
||
logger.debug(`OpenAICompatClient: reasoning_content (${reasoning.length} chars)`);
|
||
yield { type: 'thinking', chars: reasoning.length };
|
||
}
|
||
|
||
// テキストチャンク
|
||
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') {
|
||
yield* drainToolCalls(toolCallAccumulators);
|
||
}
|
||
}
|
||
}
|
||
} 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, hardCapHit), errorClass: this.classifyAbort(externalSignal, hardCapHit) };
|
||
return;
|
||
}
|
||
|
||
// 一時的なストリームエラー — 試行回数が残っていればリトライ
|
||
if (attempt >= maxAttempts) {
|
||
logger.error(`OpenAICompatClient: stream read error: ${message}`);
|
||
yield { type: 'error', error: `Stream error: ${message}`, errorClass: 'stream' };
|
||
return;
|
||
}
|
||
|
||
const delayMs = getRetryDelayMs(this.retryConfig, attempt);
|
||
logger.warn(`OpenAICompatClient: stream read error on attempt ${attempt}/${maxAttempts}: ${message}; retrying in ${delayMs}ms`);
|
||
yield { type: 'retry', attempt, maxAttempts, reason: message, errorClass: 'stream', delayMs };
|
||
if (delayMs > 0 && !(await waitForRetry(delayMs, controller.signal))) {
|
||
logger.error('OpenAICompatClient: request timed out during retry wait');
|
||
yield { type: 'error', error: this.buildAbortErrorMessage(externalSignal, hardCapHit), errorClass: this.classifyAbort(externalSignal, hardCapHit) };
|
||
return;
|
||
}
|
||
continue;
|
||
} finally {
|
||
reader.releaseLock();
|
||
}
|
||
|
||
// [DONE] なしにストリームが終了した場合。チャンクを1つも受け取らずに
|
||
// EOF した「不明完了」は requests に数えない (issue #498)。明示的な
|
||
// [DONE] 経路は従来どおり常に記録する。
|
||
yield* drainToolCalls(toolCallAccumulators);
|
||
if (sawChunk) this.finalizeDone(usage, observedModel, observedBackendId, context);
|
||
yield { type: 'done', usage };
|
||
return;
|
||
}
|
||
|
||
// 全試行が失敗した場合
|
||
yield { type: 'error', error: lastErrorMessage || 'Unknown request error', errorClass: 'unknown' };
|
||
} finally {
|
||
clearTimeout(timeoutId);
|
||
if (hardCapId) clearTimeout(hardCapId);
|
||
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 };
|
||
}
|