diff --git a/src/mcp/types.ts b/src/mcp/types.ts index 29efc18..600c832 100644 --- a/src/mcp/types.ts +++ b/src/mcp/types.ts @@ -51,6 +51,8 @@ export interface ModelConfig { baseURL: string; model: string; apiKey: string; + /** Optional provider-specific header used to correlate one logical call across retries. */ + callIdHeader?: string; secretRef?: string; temperature?: number; topP?: number; diff --git a/src/provider/runtime.ts b/src/provider/runtime.ts index 536f59e..4db7586 100644 --- a/src/provider/runtime.ts +++ b/src/provider/runtime.ts @@ -1,4 +1,5 @@ import OpenAI from 'openai'; +import { randomUUID } from 'node:crypto'; import type { RequestOptions } from 'openai/core'; import type { ChatCompletion, @@ -92,6 +93,25 @@ function nonNegativeInt(value: unknown, fallback: number): number { : fallback; } +function resolveCallIdHeader(value: unknown): string | null { + if (value === undefined) return null; + if (typeof value !== 'string' || !/^[!#$%&'*+.^_`|~0-9A-Za-z-]{1,128}$/.test(value)) { + throw new Error('provider call id header is invalid'); + } + return value.toLowerCase(); +} + +function callRequestOptions( + signal: AbortSignal | undefined, + header: string | null, + callId: string | null +): RequestOptions { + return { + signal, + ...(header && callId ? { headers: { [header]: callId } } : {}), + } as RequestOptions; +} + export function resolveProviderPolicy(model: ModelConfig): ProviderPolicy { const requestTimeoutMs = positiveInt( model.requestTimeoutMs, @@ -155,6 +175,10 @@ function isRetryableProviderError(err: unknown): boolean { if (err instanceof ProviderRequestTimeoutError) return true; if (isAbortLike(err)) return false; + const explicitRetryable = (err as any)?.error?.retryable ?? (err as any)?.retryable; + if (explicitRetryable === false) return false; + if (explicitRetryable === true) return true; + const status = (err as any)?.status; if ( status === 408 || @@ -294,6 +318,7 @@ export function createProviderRuntime( } const client = overrideClient ?? createProviderClient(model, policy); const maxAttempts = policy.maxRetries + 1; + const callIdHeader = resolveCallIdHeader(model.callIdHeader); async function runWithRetry( stream: boolean, @@ -348,13 +373,14 @@ export function createProviderRuntime( client, policy, createChatCompletion(request, options) { + const callId = callIdHeader ? `ma_call_${randomUUID()}` : null; return runWithRetry( false, () => raceWithTimeout( client.chat.completions.create( { ...request, stream: false }, - { signal: options?.signal } as RequestOptions + callRequestOptions(options?.signal, callIdHeader, callId) ) as unknown as Promise, policy.requestTimeoutMs, options?.signal @@ -363,14 +389,15 @@ export function createProviderRuntime( ); }, createStreamingChatCompletion(request, options) { + const callId = callIdHeader ? `ma_call_${randomUUID()}` : null; return runWithRetry( true, async () => { const start = await openStreamAndReadFirstChunk( () => client.chat.completions.create( - { ...request, stream: true }, - { signal: options?.signal } as RequestOptions + { ...request, stream: true }, + callRequestOptions(options?.signal, callIdHeader, callId) ) as unknown as Promise>, policy, options?.signal diff --git a/test/provider-runtime.test.ts b/test/provider-runtime.test.ts index e197846..552ea60 100644 --- a/test/provider-runtime.test.ts +++ b/test/provider-runtime.test.ts @@ -127,8 +127,89 @@ test('provider runtime: request timeout is retried', async () => { assert.ok(events.some((event) => event.type === 'retry')); }); +test('provider runtime: reuses one logical call id across retries and rotates it for the next call', async () => { + const requestOptions: any[] = []; + let attempts = 0; + const runtime = createProviderRuntime( + { + baseURL: 'http://example.test/v1', + model: 'stub', + apiKey: 'key', + callIdHeader: 'X-MTEAM-MA-Call-ID', + requestTimeoutMs: 50, + maxRetries: 1, + }, + fakeClient((_request, options) => { + requestOptions.push(options); + attempts++; + if (attempts === 1) { + const error = new Error('temporary failure') as Error & { status: number }; + error.status = 503; + return Promise.reject(error); + } + return Promise.resolve({ choices: [{ message: { content: 'ok' } }] }); + }) + ); + + await runtime.createChatCompletion({ model: 'stub', messages: [], stream: false }); + await runtime.createChatCompletion({ model: 'stub', messages: [], stream: false }); + + const ids = requestOptions.map((options) => options.headers['x-mteam-ma-call-id']); + assert.match(ids[0], /^ma_call_[0-9a-f-]{36}$/); + assert.equal(ids[1], ids[0]); + assert.notEqual(ids[2], ids[0]); +}); + +test('provider runtime: does not emit a call id header without an explicit provider opt-in', async () => { + let captured: any; + const runtime = createProviderRuntime( + { + baseURL: 'http://example.test/v1', + model: 'stub', + apiKey: 'key', + maxRetries: 0, + }, + fakeClient((_request, options) => { + captured = options; + return Promise.resolve({ choices: [] }); + }) + ); + + await runtime.createChatCompletion({ model: 'stub', messages: [], stream: false }); + assert.equal(captured.headers, undefined); +}); + +test('provider runtime: honors an explicit non-retryable provider error', async () => { + let calls = 0; + const runtime = createProviderRuntime( + { + baseURL: 'http://example.test/v1', + model: 'stub', + apiKey: 'key', + maxRetries: 3, + }, + fakeClient(() => { + calls++; + const error = new Error('authority unavailable') as Error & { + status: number; + error: { retryable: boolean }; + }; + error.status = 503; + error.error = { retryable: false }; + return Promise.reject(error); + }) + ); + + await assert.rejects( + () => runtime.createChatCompletion({ model: 'stub', messages: [], stream: false }), + /authority unavailable/ + ); + assert.equal(calls, 1); +}); + test('provider runtime: stream idle before first chunk is retried', async () => { let calls = 0; + const requestOptions: any[] = []; const runtime = createProviderRuntime( { baseURL: 'http://example.test/v1', @@ -137,8 +218,10 @@ test('provider runtime: stream idle before first chunk is retried', async () => requestTimeoutMs: 20, streamIdleTimeoutMs: 5, maxRetries: 1, + callIdHeader: 'x-mteam-ma-call-id', }, - fakeClient(() => { + fakeClient((_request, options) => { + requestOptions.push(options); calls++; if (calls === 1) return Promise.resolve(stallBeforeFirst()); return Promise.resolve(chunks([ @@ -157,6 +240,10 @@ test('provider runtime: stream idle before first chunk is retried', async () => assert.equal(calls, 2); assert.equal(out.length, 1); assert.equal(out[0].choices[0].delta.content, 'ok'); + assert.equal( + requestOptions[0].headers['x-mteam-ma-call-id'], + requestOptions[1].headers['x-mteam-ma-call-id'], + ); assert.ok(events.some((event) => event.type === 'retry')); });