@lota-sdk/core 0.4.7 → 0.4.9

package/src/ai-gateway/ai-gateway.ts

@@ -1,33 +1,438 @@
 import { devToolsMiddleware } from '@ai-sdk/devtools'
 import { createOpenAI } from '@ai-sdk/openai'
-import { wrapLanguageModel } from 'ai'
+import { wrapEmbeddingModel, wrapLanguageModel } from 'ai'
 import type { LanguageModelMiddleware } from 'ai'
+import { Cause, Clock, Context, Duration, Effect, ExecutionPlan, Fiber, Layer, Semaphore } from 'effect'
 
-import { getRuntimeConfig } from '../runtime/runtime-config'
+import { DEFAULT_AI_GATEWAY_URL } from '../config/constants'
+import { AiGenerationError, ConfigurationError } from '../effect/errors'
+import { getLotaSdkRuntime } from '../effect/runtime'
+import { RuntimeConfigServiceTag } from '../effect/services'
+import { getDirectOpenRouterProvider, normalizeDirectOpenRouterModelId } from '../openrouter/direct-provider'
 import { isRecord, readString } from '../utils/string'
 import { buildAiGatewayCacheHeaders } from './cache-headers'
 
-type AiGatewayLanguageModel = Parameters<typeof wrapLanguageModel>[0]['model']
-type AiGatewayExtraParams = Record<string, unknown>
 type AiGatewayChatResponse = { body?: unknown }
 type AiGatewayTransformParamsOptions = Parameters<NonNullable<LanguageModelMiddleware['transformParams']>>[0]
 type WrapStreamOptions = Parameters<NonNullable<LanguageModelMiddleware['wrapStream']>>[0]
+type AiGatewayLanguageModel = Parameters<typeof wrapLanguageModel>[0]['model']
+type AiGatewayEmbeddingModel = Parameters<typeof wrapEmbeddingModel>[0]['model']
 type AiGatewayCallOptions = WrapStreamOptions['params']
 type AiGatewayGenerateResult = Awaited<ReturnType<WrapStreamOptions['doGenerate']>>
 type AiGatewayStreamResult = Awaited<ReturnType<WrapStreamOptions['doStream']>>
 type AiGatewayGeneratedContent = AiGatewayGenerateResult['content'][number]
 type AiGatewayStreamPart = AiGatewayStreamResult['stream'] extends ReadableStream<infer T> ? T : never
-type AiGatewayConfig = { apiKey: string; baseURL: string }
 type AiGatewayProviderOptions = NonNullable<AiGatewayCallOptions['providerOptions']>
+type AiGatewayAttemptResult<A> = { source: string; result: A }
+
+class AiGatewayGenerateAttempt extends Context.Service<
+  AiGatewayGenerateAttempt,
+  { readonly execute: Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> }
+>()('AiGatewayGenerateAttempt') {}
+
+class AiGatewayStreamAttempt extends Context.Service<
+  AiGatewayStreamAttempt,
+  { readonly execute: Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> }
+>()('AiGatewayStreamAttempt') {}
 
 const EXPECTED_GATEWAY_KEY_PREFIX = 'sk-bf-'
 const AI_GATEWAY_VIRTUAL_KEY_HEADER = 'x-bf-vk'
-const AI_GATEWAY_EXTRA_PARAMS_HEADER = 'x-bf-passthrough-extra-params'
-const DEFAULT_AI_GATEWAY_URL = 'https://ai-gateway.gobrainy.ai' as const
-const OPENAI_PROMPT_CACHE_RETENTION = '24h' as const
-const OPENROUTER_RESPONSE_HEALING_EXTRA_PARAMS = {
-  plugins: [{ id: 'response-healing' }],
-} as const satisfies AiGatewayExtraParams
+const AI_GATEWAY_TIMEOUT_MS = 30_000
+const AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS = 30_000
+const AI_GATEWAY_MAX_RETRIES = 4
+const AI_GATEWAY_MAX_RETRY_DELAY_MS = 15_000
+const OPENAI_RESPONSES_PROVIDER_ID = 'openai.responses'
+const OPENAI_CHAT_PROVIDER_ID = 'openai.chat'
+const OPENAI_EMBEDDING_PROVIDER_ID = 'openai.embedding'
+const OPENAI_EMBEDDING_MAX_PER_CALL = 2_048
+const RETRYABLE_NETWORK_ERROR_CODES = new Set([
+  'ECONNABORTED',
+  'ECONNREFUSED',
+  'ECONNRESET',
+  'EAI_AGAIN',
+  'EHOSTUNREACH',
+  'ENETDOWN',
+  'ENETUNREACH',
+  'ENOTFOUND',
+  'ETIMEDOUT',
+  'UND_ERR_CONNECT_TIMEOUT',
+  'UND_ERR_SOCKET',
+])
+const RETRYABLE_NETWORK_ERROR_PATTERNS = [
+  /fetch failed/i,
+  /network error/i,
+  /socket hang up/i,
+  /socket closed/i,
+  /connection (?:reset|refused|closed|timed out|timeout)/i,
+  /dns/i,
+  /temporary failure in name resolution/i,
+  /timed out/i,
+]
+
+function getNumericField(value: Record<string, unknown>, key: string): number | null {
+  const field = value[key]
+  if (typeof field === 'number' && Number.isFinite(field)) return field
+  if (typeof field === 'string') {
+    const parsed = Number(field)
+    if (Number.isFinite(parsed)) return parsed
+  }
+  return null
+}
+
+function getErrorStatus(error: unknown): number | null {
+  if (!isRecord(error)) return null
+  return getNumericField(error, 'status') ?? getNumericField(error, 'statusCode')
+}
+
+function isHeaderEntry(value: unknown): value is readonly [string, string | readonly string[]] {
+  return Array.isArray(value) && value.length >= 2 && typeof value[0] === 'string'
+}
+
+function readHeaderValue(headers: unknown, name: string): string | null {
+  const normalizedName = name.toLowerCase()
+
+  if (headers instanceof Headers) {
+    return headers.get(name)
+  }
+
+  if (Array.isArray(headers)) {
+    for (const entry of headers) {
+      if (!isHeaderEntry(entry)) continue
+      const [headerName, headerValue] = entry
+      if (headerName.toLowerCase() !== normalizedName) continue
+      if (typeof headerValue === 'string') return headerValue
+      if (Array.isArray(headerValue)) {
+        const firstValue = headerValue.find((value): value is string => typeof value === 'string')
+        if (firstValue) return firstValue
+      }
+    }
+    return null
+  }
+
+  if (!isRecord(headers)) {
+    return null
+  }
+
+  for (const [headerName, headerValue] of Object.entries(headers)) {
+    if (headerName.toLowerCase() !== normalizedName) continue
+    if (typeof headerValue === 'string') return headerValue
+    if (Array.isArray(headerValue)) {
+      const firstValue = headerValue.find((value): value is string => typeof value === 'string')
+      if (firstValue) return firstValue
+    }
+  }
+
+  return null
+}
+
+function extractRetryAfter(error: unknown): string | undefined {
+  if (!isRecord(error)) return undefined
+
+  const direct =
+    readHeaderValue(error.responseHeaders, 'retry-after') ??
+    readHeaderValue(error.headers, 'retry-after') ??
+    (isRecord(error.response) ? readHeaderValue(error.response.headers, 'retry-after') : null)
+  if (direct) {
+    return direct
+  }
+
+  if (!isRecord(error.cause)) return undefined
+  return (
+    readHeaderValue(error.cause.responseHeaders, 'retry-after') ??
+    readHeaderValue(error.cause.headers, 'retry-after') ??
+    (isRecord(error.cause.response) ? readHeaderValue(error.cause.response.headers, 'retry-after') : undefined) ??
+    undefined
+  )
+}
+
+function stringifyProviderField(value: unknown, maxLength: number): string | undefined {
+  if (value === undefined) return undefined
+  try {
+    const stringified = typeof value === 'string' ? value : JSON.stringify(value)
+    return stringified.length > maxLength ? `${stringified.slice(0, maxLength)}...` : stringified
+  } catch {
+    return undefined
+  }
+}
+
+function classifyAiGatewayError(source: string, error: unknown): AiGenerationError {
+  if (error instanceof AiGenerationError) {
+    return error
+  }
+
+  const status = getErrorStatus(error)
+  const rateLimited = status === 429
+  const retryAfter = extractRetryAfter(error)
+  const errorRecord = isRecord(error) ? error : null
+  const message = error instanceof Error ? error.message : String(error)
+  const providerData = errorRecord ? stringifyProviderField(errorRecord.data, 600) : undefined
+  const cause = errorRecord ? stringifyProviderField(errorRecord.cause, 600) : undefined
+  const responseBody = errorRecord ? stringifyProviderField(errorRecord.responseBody, 600) : undefined
+  const url = errorRecord ? stringifyProviderField(errorRecord.url, 200) : undefined
+
+  const parts = [`[${source}]`]
+  if (status !== null) parts.push(`status=${status}`)
+  if (rateLimited) parts.push('rate_limited')
+  parts.push(message)
+  if (providerData) parts.push(`provider_data=${providerData}`)
+  if (cause) parts.push(`cause=${cause}`)
+  if (responseBody) parts.push(`response_body=${responseBody}`)
+  if (url) parts.push(`url=${url}`)
+
+  return new AiGenerationError({
+    source,
+    message: parts.join(' '),
+    ...(status !== null ? { status } : {}),
+    ...(rateLimited ? { rateLimited: true } : {}),
+    ...(retryAfter ? { retryAfter } : {}),
+    ...(providerData ? { providerData } : {}),
+    ...(responseBody ? { responseBody } : {}),
+    ...(url ? { url } : {}),
+  })
+}
+
+function isRetryableAiGatewayError(error: AiGenerationError): boolean {
+  if (error.rateLimited) return true
+
+  if (typeof error.status === 'number' && error.status >= 500) {
+    return true
+  }
+
+  const message = error.message.toLowerCase()
+  if (RETRYABLE_NETWORK_ERROR_PATTERNS.some((pattern) => pattern.test(message))) {
+    return true
+  }
+
+  for (const code of RETRYABLE_NETWORK_ERROR_CODES) {
+    if (message.includes(code.toLowerCase())) {
+      return true
+    }
+  }
+
+  return false
+}
+
+function parseRetryAfterDelayMs(retryAfter: string | undefined, nowMillis: number): number | undefined {
+  if (!retryAfter) return undefined
+
+  const trimmed = retryAfter.trim()
+  if (!trimmed) return undefined
+
+  const seconds = Number(trimmed)
+  if (Number.isFinite(seconds) && seconds >= 0) {
+    return Math.min(AI_GATEWAY_MAX_RETRY_DELAY_MS, Math.round(seconds * 1_000))
+  }
+
+  const retryAt = Date.parse(trimmed)
+  if (Number.isNaN(retryAt)) return undefined
+
+  return Math.min(AI_GATEWAY_MAX_RETRY_DELAY_MS, Math.max(0, retryAt - nowMillis))
+}
+
+function computeRetryDelayMs(attempt: number, error: AiGenerationError, nowMillis: number): number {
+  const baseDelayMs = Math.min(AI_GATEWAY_MAX_RETRY_DELAY_MS, 500 * 2 ** attempt)
+  const retryAfterDelayMs = parseRetryAfterDelayMs(error.retryAfter, nowMillis)
+  return retryAfterDelayMs === undefined ? baseDelayMs : Math.max(baseDelayMs, retryAfterDelayMs)
+}
+
+function withAiGatewayResilience<A>(source: string, effect: Effect.Effect<A, AiGenerationError>) {
+  const retryEffect = Effect.gen(function* () {
+    const runAttempt = (attempt: number): Effect.Effect<A, AiGenerationError> =>
+      effect.pipe(
+        Effect.catchTag('AiGenerationError', (error) =>
+          Effect.gen(function* () {
+            if (!isRetryableAiGatewayError(error) || attempt >= AI_GATEWAY_MAX_RETRIES - 1) {
+              return yield* error
+            }
+
+            const nowMillis = yield* Clock.currentTimeMillis
+            const retryDelayMs = computeRetryDelayMs(attempt, error, nowMillis)
+            yield* Effect.sleep(Duration.millis(retryDelayMs))
+            return yield* runAttempt(attempt + 1)
+          }),
+        ),
+      )
+
+    return yield* runAttempt(0)
+  })
+
+  return retryEffect.pipe(
+    Effect.timeout(Duration.millis(AI_GATEWAY_TIMEOUT_MS)),
+    Effect.catchIf(Cause.isTimeoutError, () =>
+      Effect.fail(new AiGenerationError({ source, message: `[${source}] Timed out after ${AI_GATEWAY_TIMEOUT_MS}ms` })),
+    ),
+  )
+}
+
+function withAiGatewayStreamIdleTimeout(
+  stream: ReadableStream<AiGatewayStreamPart>,
+  source: string,
+  onFinalize?: () => void,
+): ReadableStream<AiGatewayStreamPart> {
+  let closed = false
+  let reader: ReadableStreamDefaultReader<AiGatewayStreamPart> | null = null
+  let idleTimeoutFiber: ReturnType<typeof Effect.runFork> | null = null
+  let bodyPumpFiber: ReturnType<typeof Effect.runFork> | null = null
+  let finalized = false
+
+  const finalize = () => {
+    if (finalized) return
+    finalized = true
+    onFinalize?.()
+  }
+
+  const interruptFiber = (fiber: ReturnType<typeof Effect.runFork> | null) => {
+    if (!fiber) return
+    void Effect.runFork(Fiber.interrupt(fiber))
+  }
+
+  const stopIdleTimeout = () => {
+    const fiber = idleTimeoutFiber
+    idleTimeoutFiber = null
+    interruptFiber(fiber)
+  }
+
+  const stopBodyPump = () => {
+    const fiber = bodyPumpFiber
+    bodyPumpFiber = null
+    interruptFiber(fiber)
+  }
+
+  const releaseReader = (streamReader: ReadableStreamDefaultReader<AiGatewayStreamPart>) => {
+    try {
+      streamReader.releaseLock()
+    } catch {
+      // Best-effort cleanup.
+    }
+  }
+
+  const enqueueChunk = (
+    controller: ReadableStreamDefaultController<AiGatewayStreamPart>,
+    chunk: AiGatewayStreamPart,
+  ): boolean => {
+    if (closed) return false
+
+    try {
+      controller.enqueue(chunk)
+      return true
+    } catch {
+      closed = true
+      return false
+    }
+  }
+
+  const closeStream = (controller: ReadableStreamDefaultController<AiGatewayStreamPart>) => {
+    if (closed) return
+    closed = true
+    finalize()
+
+    try {
+      controller.close()
+    } catch {
+      // Best-effort cleanup.
+    }
+  }
+
+  const errorStream = (controller: ReadableStreamDefaultController<AiGatewayStreamPart>, error: unknown) => {
+    if (closed) return
+    closed = true
+    finalize()
+
+    try {
+      controller.error(error)
+    } catch {
+      // Best-effort cleanup.
+    }
+  }
+
+  const resetIdleTimeout = (controller: ReadableStreamDefaultController<AiGatewayStreamPart>) => {
+    stopIdleTimeout()
+    idleTimeoutFiber = Effect.runFork(
+      Effect.sleep(Duration.millis(AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS)).pipe(
+        Effect.flatMap(() =>
+          Effect.gen(function* () {
+            if (closed) return
+
+            const timeoutError = new AiGenerationError({
+              source,
+              message: `[${source}] Stream stalled after ${AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS}ms`,
+            })
+
+            yield* Effect.sync(() => errorStream(controller, timeoutError))
+            stopBodyPump()
+
+            const streamReader = reader
+            reader = null
+            if (!streamReader) return
+
+            yield* Effect.tryPromise(() => streamReader.cancel(timeoutError)).pipe(Effect.catch(() => Effect.void))
+          }),
+        ),
+      ),
+    )
+  }
+
+  const pumpStreamEffect = (
+    streamReader: ReadableStreamDefaultReader<AiGatewayStreamPart>,
+    controller: ReadableStreamDefaultController<AiGatewayStreamPart>,
+  ): Effect.Effect<void> =>
+    Effect.gen(function* () {
+      resetIdleTimeout(controller)
+
+      for (;;) {
+        if (closed) return
+
+        const { done, value } = yield* Effect.tryPromise(() => streamReader.read())
+        if (done) {
+          stopIdleTimeout()
+          yield* Effect.sync(() => closeStream(controller))
+          return
+        }
+
+        if (!enqueueChunk(controller, value)) {
+          return
+        }
+
+        resetIdleTimeout(controller)
+      }
+    }).pipe(
+      Effect.catch((error: unknown) => Effect.sync(() => errorStream(controller, error))),
+      Effect.ensuring(
+        Effect.sync(() => {
+          closed = true
+          finalize()
+          stopIdleTimeout()
+          bodyPumpFiber = null
+          reader = null
+          releaseReader(streamReader)
+        }),
+      ),
+    )
+
+  return new ReadableStream<AiGatewayStreamPart>({
+    start(controller) {
+      const streamReader = stream.getReader()
+      reader = streamReader
+      bodyPumpFiber = Effect.runFork(pumpStreamEffect(streamReader, controller))
+    },
+    cancel(reason) {
+      closed = true
+      finalize()
+      stopIdleTimeout()
+      stopBodyPump()
+
+      const streamReader = reader
+      reader = null
+      if (!streamReader) {
+        return
+      }
+
+      return Effect.runPromise(
+        Effect.tryPromise(() => streamReader.cancel(reason)).pipe(Effect.catch(() => Effect.void)),
+      )
+    },
+  })
+}
 
 function mergeAiGatewayHeaders(
   existingHeaders: AiGatewayCallOptions['headers'] | undefined,
@@ -42,19 +447,6 @@ function mergeAiGatewayHeaders(
   return Object.fromEntries(merged.entries())
 }
 
-function parseAiGatewayJsonRequestBody(body: BodyInit | null | undefined): Record<string, unknown> | null {
-  if (typeof body !== 'string') return null
-
-  let parsed: unknown
-  try {
-    parsed = JSON.parse(body)
-  } catch {
-    return null
-  }
-
-  return isRecord(parsed) ? parsed : null
-}
-
 function withDefaultAiGatewayCacheHeaders(params: AiGatewayCallOptions): AiGatewayCallOptions {
   return { ...params, headers: mergeAiGatewayHeaders(params.headers, buildAiGatewayCacheHeaders('lota-sdk')) }
 }
@@ -62,31 +454,80 @@ function withDefaultAiGatewayCacheHeaders(params: AiGatewayCallOptions): AiGatew
 function normalizeAiGatewayUrl(value: string): string {
   const trimmed = value.trim()
   if (!trimmed) {
-    throw new Error('[ai-gateway] AI gateway URL is required.')
+    throw new ConfigurationError({ message: '[ai-gateway] AI gateway URL is required.', key: 'aiGateway.url' })
   }
 
   const normalized = trimmed.replace(/\/+$/, '')
   return normalized.endsWith('/v1') ? normalized : `${normalized}/v1`
 }
 
-function readDirectEnvAiGatewayConfig(): AiGatewayConfig {
-  const apiKey = (process.env.AI_GATEWAY_KEY ?? '').trim()
-  if (!apiKey) {
-    throw new Error(
-      '[ai-gateway] Missing AI gateway key. Set AI_GATEWAY_KEY, or configure createLotaRuntime({ aiGateway: { key } }).',
-    )
-  }
+// ── AiGateway Effect Layer ────────────────────────────────────────────
+export class AiGatewayTag extends Context.Service<
+  AiGatewayTag,
+  { readonly semaphore: Semaphore.Semaphore; readonly provider: ReturnType<typeof createOpenAI> }
+>()('AiGateway') {}
+
+export const AiGatewayLive = Layer.effect(
+  AiGatewayTag,
+  Effect.gen(function* () {
+    const config = yield* RuntimeConfigServiceTag
+    const semaphore = yield* Semaphore.make(config.aiGateway.maxConcurrency)
+
+    const apiKey = config.aiGateway.key.trim()
+    if (!apiKey.startsWith(EXPECTED_GATEWAY_KEY_PREFIX)) {
+      return yield* new ConfigurationError({
+        message: `[ai-gateway] Gateway keys must use the ${EXPECTED_GATEWAY_KEY_PREFIX}* format.`,
+        key: 'aiGateway.key',
+      })
+    }
+    const baseURL = normalizeAiGatewayUrl(config.aiGateway.url)
+    const provider = createOpenAI({ baseURL, apiKey, headers: { [AI_GATEWAY_VIRTUAL_KEY_HEADER]: apiKey } })
 
-  return { apiKey, baseURL: normalizeAiGatewayUrl(process.env.AI_GATEWAY_URL?.trim() || DEFAULT_AI_GATEWAY_URL) }
+    return AiGatewayTag.of({ semaphore, provider })
+  }),
+)
+
+function resolveFromRuntime<I, T>(tag: Context.Key<I, T>): T {
+  return getLotaSdkRuntime().runSync(Effect.service(tag))
 }
 
-function readAiGatewayConfig(): AiGatewayConfig {
-  try {
-    const { aiGateway } = getRuntimeConfig()
-    return { apiKey: aiGateway.key.trim(), baseURL: normalizeAiGatewayUrl(aiGateway.url) }
-  } catch {
-    return readDirectEnvAiGatewayConfig()
-  }
+function getAiGateway(): AiGatewayTag['Service'] {
+  return resolveFromRuntime(AiGatewayTag)
+}
+
+function withAiGatewayConcurrency<A>(effect: Effect.Effect<A, AiGenerationError>): Effect.Effect<A, AiGenerationError> {
+  return getAiGateway().semaphore.withPermit(effect)
+}
+
+function withAiGatewayStreamConcurrency(
+  effect: Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError>,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
+  return Effect.uninterruptibleMask((restore) =>
+    Effect.gen(function* () {
+      const { semaphore } = getAiGateway()
+      const currentContext = yield* Effect.context<never>()
+      yield* semaphore.take(1)
+
+      let released = false
+      const release = () => {
+        if (released) return
+        released = true
+        void Effect.runForkWith(currentContext)(semaphore.release(1))
+      }
+
+      const attempt = yield* restore(effect).pipe(
+        Effect.catchTag('AiGenerationError', (error) => Effect.sync(release).pipe(Effect.andThen(Effect.fail(error)))),
+      )
+
+      return {
+        ...attempt,
+        result: {
+          ...attempt.result,
+          stream: withAiGatewayStreamIdleTimeout(attempt.result.stream, attempt.source, release),
+        },
+      }
+    }),
+  )
 }
 
 function readReasoningDetailsText(value: unknown): string | null {
@@ -142,19 +583,6 @@ export function extractAiGatewayChatReasoningDeltaText(rawChunk: unknown): strin
   return null
 }
 
-type AiGatewayResponsesReasoningDelta = { id: string; delta: string; itemId: string }
-
-export function extractAiGatewayResponsesReasoningDelta(rawChunk: unknown): AiGatewayResponsesReasoningDelta | null {
-  if (!isRecord(rawChunk) || rawChunk.type !== 'response.reasoning_summary_text.delta') return null
-  if ('summary_index' in rawChunk) return null
-
-  const itemId = readString(rawChunk.item_id)
-  const delta = readReasoningDeltaText(rawChunk.delta)
-  if (!itemId || !delta) return null
-
-  return { id: `${itemId}:0`, delta, itemId }
-}
-
 export function injectAiGatewayChatReasoningContent(
   content: readonly AiGatewayGeneratedContent[],
   response?: AiGatewayChatResponse,
@@ -178,8 +606,158 @@ function isReasoningEnabled(params: AiGatewayCallOptions): boolean {
   return typeof openaiOptions.reasoningEffort === 'string' && openaiOptions.reasoningEffort !== 'none'
 }
 
+function isOpenRouterModel(modelId: string): boolean {
+  return modelId.trim().toLowerCase().startsWith('openrouter/')
+}
+
+function hasDirectOpenRouterFallback(modelId: string): boolean {
+  const config = resolveFromRuntime(RuntimeConfigServiceTag)
+  return isOpenRouterModel(modelId) && Boolean(config.aiGateway.openRouterApiKey?.trim())
+}
+
+function getDirectOpenRouterChatModel(modelId: string): AiGatewayLanguageModel {
+  const config = resolveFromRuntime(RuntimeConfigServiceTag)
+  return getDirectOpenRouterProvider(config.aiGateway.openRouterApiKey).chat(normalizeDirectOpenRouterModelId(modelId))
+}
+
+function shouldFallbackToDirectOpenRouter(modelId: string, error: AiGenerationError): boolean {
+  return hasDirectOpenRouterFallback(modelId) && isRetryableAiGatewayError(error)
+}
+
+function attemptAiGatewayGenerate(
+  source: string,
+  evaluate: () => PromiseLike<AiGatewayGenerateResult>,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
+  return withAiGatewayResilience(
+    source,
+    Effect.tryPromise({ try: evaluate, catch: (cause) => classifyAiGatewayError(source, cause) }),
+  ).pipe(
+    Effect.map((result) => ({ source, result })),
+    Effect.withSpan('AiGateway.generateAttempt'),
+    Effect.annotateSpans({ gatewaySource: source }),
+  )
+}
+
+function attemptAiGatewayStream(
+  source: string,
+  evaluate: () => PromiseLike<AiGatewayStreamResult>,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
+  return withAiGatewayResilience(
+    source,
+    Effect.tryPromise({ try: evaluate, catch: (cause) => classifyAiGatewayError(source, cause) }),
+  ).pipe(
+    Effect.map((result) => ({ source, result })),
+    Effect.withSpan('AiGateway.streamAttempt'),
+    Effect.annotateSpans({ gatewaySource: source }),
+  )
+}
+
+function attemptDirectOpenRouterGenerate(
+  modelId: string,
+  params: AiGatewayCallOptions,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
+  const model = getDirectOpenRouterChatModel(modelId)
+  return attemptAiGatewayGenerate('openrouter.generate', () => model.doGenerate(params))
+}
+
+function attemptDirectOpenRouterStream(
+  modelId: string,
+  params: AiGatewayCallOptions,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
+  const model = getDirectOpenRouterChatModel(modelId)
+  return attemptAiGatewayStream('openrouter.stream', () => model.doStream(params))
+}
+
+function executeGenerateAttemptPlan(
+  modelId: string,
+  params: AiGatewayCallOptions,
+  doGenerate: () => PromiseLike<AiGatewayGenerateResult>,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
+  const primary = Layer.succeed(AiGatewayGenerateAttempt, {
+    execute: attemptAiGatewayGenerate('ai-gateway.generate', doGenerate),
+  })
+  const effect = Effect.gen(function* () {
+    const attempt = yield* AiGatewayGenerateAttempt
+    return yield* attempt.execute
+  })
+
+  if (!hasDirectOpenRouterFallback(modelId)) {
+    return effect.pipe(
+      Effect.provide(primary),
+      Effect.withSpan('AiGateway.executeGeneratePlan'),
+      Effect.annotateSpans({ modelId, directOpenRouterFallbackEnabled: false }),
+    )
+  }
+
+  return effect.pipe(
+    Effect.withExecutionPlan(
+      ExecutionPlan.make(
+        { provide: primary },
+        {
+          provide: Layer.succeed(AiGatewayGenerateAttempt, {
+            execute: attemptDirectOpenRouterGenerate(modelId, params),
+          }),
+          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(modelId, error),
+        },
+      ),
+    ),
+    Effect.withSpan('AiGateway.executeGeneratePlan'),
+    Effect.annotateSpans({ modelId, directOpenRouterFallbackEnabled: true }),
+  )
+}
+
+function executeStreamAttemptPlan(
+  modelId: string,
+  params: AiGatewayCallOptions,
+  doStream: () => PromiseLike<AiGatewayStreamResult>,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
+  const primary = Layer.succeed(AiGatewayStreamAttempt, {
+    execute: attemptAiGatewayStream('ai-gateway.stream', doStream),
+  })
+  const effect = Effect.gen(function* () {
+    const attempt = yield* AiGatewayStreamAttempt
+    return yield* attempt.execute
+  })
+
+  if (!hasDirectOpenRouterFallback(modelId)) {
+    return effect.pipe(
+      Effect.provide(primary),
+      Effect.withSpan('AiGateway.executeStreamPlan'),
+      Effect.annotateSpans({ modelId, directOpenRouterFallbackEnabled: false }),
+    )
+  }
+
+  return effect.pipe(
+    Effect.withExecutionPlan(
+      ExecutionPlan.make(
+        { provide: primary },
+        {
+          provide: Layer.succeed(AiGatewayStreamAttempt, { execute: attemptDirectOpenRouterStream(modelId, params) }),
+          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(modelId, error),
+        },
+      ),
+    ),
+    Effect.withSpan('AiGateway.executeStreamPlan'),
+    Effect.annotateSpans({ modelId, directOpenRouterFallbackEnabled: true }),
+  )
+}
+
+function isOpenRouterOpenAIReasoningModel(modelId: string): boolean {
+  return modelId.trim().toLowerCase().startsWith('openrouter/openai/gpt-5')
+}
+
 function shouldCloseInjectedReasoning(chunk: AiGatewayStreamPart): boolean {
-  return chunk.type !== 'stream-start' && chunk.type !== 'response-metadata' && chunk.type !== 'raw'
+  switch (chunk.type) {
+    case 'stream-start':
+    case 'response-metadata':
+    case 'raw':
+    case 'text-start':
+      return false
+    case 'text-delta':
+      return chunk.delta.length > 0
+    default:
+      return true
+  }
 }
 
 export function injectAiGatewayChatReasoningStream(
@@ -232,30 +810,6 @@ export function injectAiGatewayChatReasoningStream(
   )
 }
 
-export function injectAiGatewayResponsesReasoningStream(
-  stream: ReadableStream<AiGatewayStreamPart>,
-): ReadableStream<AiGatewayStreamPart> {
-  return stream.pipeThrough(
-    new TransformStream<AiGatewayStreamPart, AiGatewayStreamPart>({
-      transform(chunk, controller) {
-        controller.enqueue(chunk)
-
-        if (chunk.type !== 'raw') return
-
-        const reasoningDelta = extractAiGatewayResponsesReasoningDelta(chunk.rawValue)
-        if (!reasoningDelta) return
-
-        controller.enqueue({
-          type: 'reasoning-delta',
-          id: reasoningDelta.id,
-          delta: reasoningDelta.delta,
-          providerMetadata: { openai: { itemId: reasoningDelta.itemId } },
-        } satisfies AiGatewayStreamPart)
-      },
-    }),
-  )
-}
-
 function addAiGatewayReasoningRawChunks(
   params: AiGatewayCallOptions,
   type: AiGatewayTransformParamsOptions['type'],
@@ -267,7 +821,49 @@ function addAiGatewayReasoningRawChunks(
   return { ...params, includeRawChunks: true }
 }
 
-export function normalizeAiGatewayChatProviderOptions(params: AiGatewayCallOptions): AiGatewayCallOptions {
+function createAiGatewayLanguageModelMiddleware(modelId: string): LanguageModelMiddleware {
+  return {
+    specificationVersion: 'v3',
+    transformParams: ({ params, type }) =>
+      Promise.resolve(
+        withDefaultAiGatewayCacheHeaders(
+          addAiGatewayReasoningRawChunks(normalizeAiGatewayChatProviderOptions(params, modelId), type),
+        ),
+      ),
+    wrapGenerate: ({ doGenerate, params }) =>
+      Effect.runPromise(
+        withAiGatewayConcurrency(
+          executeGenerateAttemptPlan(modelId, params, doGenerate).pipe(
+            Effect.map(({ result }) => ({
+              ...result,
+              content: injectAiGatewayChatReasoningContent(
+                result.content,
+                result.response as AiGatewayChatResponse | undefined,
+              ),
+            })),
+          ),
+        ),
+      ),
+    wrapStream: ({ doStream, params }) =>
+      Effect.runPromise(
+        withAiGatewayStreamConcurrency(
+          executeStreamAttemptPlan(modelId, params, doStream).pipe(
+            Effect.map((attempt) => ({
+              ...attempt,
+              result: isReasoningEnabled(params)
+                ? { ...attempt.result, stream: injectAiGatewayChatReasoningStream(attempt.result.stream) }
+                : attempt.result,
+            })),
+          ),
+        ).pipe(Effect.map(({ result }) => result)),
+      ),
+  }
+}
+
+export function normalizeAiGatewayChatProviderOptions(
+  params: AiGatewayCallOptions,
+  modelId?: string,
+): AiGatewayCallOptions {
   const providerOptions = isRecord(params.providerOptions)
     ? ({ ...params.providerOptions } as AiGatewayProviderOptions)
     : ({} as AiGatewayProviderOptions)
@@ -275,189 +871,108 @@ export function normalizeAiGatewayChatProviderOptions(params: AiGatewayCallOptio
     ? { ...providerOptions.openai }
     : ({} as Record<string, unknown>)
 
-  if (openaiOptions.systemMessageMode === 'system') {
+  if (modelId && isOpenRouterOpenAIReasoningModel(modelId) && openaiOptions.forceReasoning === undefined) {
+    openaiOptions.forceReasoning = true
+  }
+
+  if (providerOptions.openai === openaiOptions || Object.keys(openaiOptions).length === 0) {
     return params
   }
 
   return {
     ...params,
-    providerOptions: {
-      ...providerOptions,
-      openai: {
-        ...openaiOptions,
-        ...(openaiOptions.systemMessageMode === 'remove' ? {} : { systemMessageMode: 'system' }),
-      },
-    },
+    providerOptions: { ...providerOptions, openai: openaiOptions as AiGatewayProviderOptions['openai'] },
   }
 }
 
-export function injectAiGatewayExtraParamsRequestBody(
-  body: BodyInit | null | undefined,
-  extraParams: AiGatewayExtraParams,
-): BodyInit | null | undefined {
-  const parsed = parseAiGatewayJsonRequestBody(body)
-  if (!parsed) return body
-
-  const mergedExtraParams = isRecord(parsed.extra_params)
-    ? { ...parsed.extra_params, ...extraParams }
-    : { ...extraParams }
-
-  return JSON.stringify({ ...parsed, extra_params: mergedExtraParams })
-}
-
-function createAiGatewayFetch(extraParams?: AiGatewayExtraParams): typeof fetch {
-  const fetchWithMutations = (input: RequestInfo | URL, init?: RequestInit | BunFetchRequestInit) => {
-    const parsedBody = parseAiGatewayJsonRequestBody(init?.body)
-    let nextBody = init?.body
-    let nextParsedBody = parsedBody
-
-    if (
-      nextParsedBody &&
-      readString(nextParsedBody.model)?.startsWith('openai/') &&
-      !readString(nextParsedBody.prompt_cache_retention)
-    ) {
-      nextParsedBody = { ...nextParsedBody, prompt_cache_retention: OPENAI_PROMPT_CACHE_RETENTION }
-      nextBody = JSON.stringify(nextParsedBody)
-    }
-
-    if (nextParsedBody && extraParams !== undefined) {
-      nextParsedBody = {
-        ...nextParsedBody,
-        extra_params: isRecord(nextParsedBody.extra_params)
-          ? { ...nextParsedBody.extra_params, ...extraParams }
-          : { ...extraParams },
-      }
-      nextBody = JSON.stringify(nextParsedBody)
-    }
-
-    const headers = new Headers(init?.headers)
-    if (
-      extraParams !== undefined ||
-      (readString(nextParsedBody?.model)?.startsWith('openai/') &&
-        readString(nextParsedBody?.prompt_cache_retention) !== null)
-    ) {
-      // Bifrost only forwards provider-specific extra params when passthrough is enabled.
-      headers.set(AI_GATEWAY_EXTRA_PARAMS_HEADER, 'true')
-    }
-
-    return globalThis.fetch(input, { ...init, headers, body: nextBody })
-  }
-
-  const preconnect = globalThis.fetch.preconnect
-
-  if (typeof preconnect !== 'function') {
-    return fetchWithMutations as typeof fetch
+function withAiGatewayDevTools<TModel extends AiGatewayLanguageModel>(model: TModel): TModel {
+  if (Bun.env.NODE_ENV === 'production') {
+    return model
   }
 
-  return Object.assign(fetchWithMutations, { preconnect: preconnect.bind(globalThis.fetch) })
+  return wrapLanguageModel({ model, middleware: devToolsMiddleware() }) as TModel
 }
 
-function createAiGatewayProvider(extraParams?: AiGatewayExtraParams) {
-  const { apiKey, baseURL } = readAiGatewayConfig()
-  if (!apiKey.startsWith(EXPECTED_GATEWAY_KEY_PREFIX)) {
-    throw new Error(`[ai-gateway] Gateway keys must use the ${EXPECTED_GATEWAY_KEY_PREFIX}* format.`)
+function createLazyAiGatewayLanguageModel(params: {
+  modelId: string
+  providerId: string
+  resolve: () => AiGatewayLanguageModel
+}): AiGatewayLanguageModel {
+  return {
+    specificationVersion: 'v3',
+    provider: params.providerId,
+    modelId: params.modelId,
+    supportedUrls: {},
+    doGenerate: (options) => params.resolve().doGenerate(options),
+    doStream: (options) => params.resolve().doStream(options),
   }
-
-  return createOpenAI({
-    baseURL,
-    apiKey,
-    headers: { [AI_GATEWAY_VIRTUAL_KEY_HEADER]: apiKey },
-    fetch: createAiGatewayFetch(extraParams),
-  })
 }
 
-function withAiGatewayDevTools<TModel extends AiGatewayLanguageModel>(model: TModel): TModel {
-  if (process.env.NODE_ENV === 'production') {
-    return model
+function createLazyAiGatewayEmbeddingModel(modelId: string): AiGatewayEmbeddingModel {
+  return {
+    specificationVersion: 'v3',
+    provider: OPENAI_EMBEDDING_PROVIDER_ID,
+    modelId,
+    maxEmbeddingsPerCall: OPENAI_EMBEDDING_MAX_PER_CALL,
+    supportsParallelCalls: true,
+    doEmbed: (options) => getAiGatewayProvider().embeddingModel(modelId).doEmbed(options),
   }
-
-  return wrapLanguageModel({ model, middleware: devToolsMiddleware() }) as TModel
 }
 
-let provider: ReturnType<typeof createOpenAI> | null = null
-let openRouterResponseHealingProvider: ReturnType<typeof createOpenAI> | null = null
-
 export function getAiGatewayProvider() {
-  if (provider) return provider
-
-  provider = createAiGatewayProvider()
-
-  return provider
-}
-
-export function getAiGatewayOpenRouterResponseHealingProvider() {
-  if (openRouterResponseHealingProvider) return openRouterResponseHealingProvider
-
-  openRouterResponseHealingProvider = createAiGatewayProvider(OPENROUTER_RESPONSE_HEALING_EXTRA_PARAMS)
-
-  return openRouterResponseHealingProvider
+  return getAiGateway().provider
 }
 
 export function aiGatewayModel(modelId: string) {
+  if (isOpenRouterModel(modelId)) {
+    return aiGatewayChatModel(modelId)
+  }
+
   return withAiGatewayDevTools(
     wrapLanguageModel({
-      model: getAiGatewayProvider()(modelId),
-      middleware: {
-        specificationVersion: 'v3',
-        transformParams: async ({ params, type }) =>
-          withDefaultAiGatewayCacheHeaders(addAiGatewayReasoningRawChunks(params, type)),
-        wrapStream: async ({ doStream, params }) => {
-          const result = await doStream()
-          if (!isReasoningEnabled(params)) return result
-
-          return { ...result, stream: injectAiGatewayResponsesReasoningStream(result.stream) }
-        },
-      },
+      model: createLazyAiGatewayLanguageModel({
+        modelId,
+        providerId: OPENAI_RESPONSES_PROVIDER_ID,
+        resolve: () => getAiGatewayProvider()(modelId),
+      }),
+      middleware: createAiGatewayLanguageModelMiddleware(modelId),
     }),
   )
 }
 
 export function aiGatewayOpenRouterResponseHealingModel(modelId: string) {
-  return withAiGatewayDevTools(
-    wrapLanguageModel({
-      model: getAiGatewayOpenRouterResponseHealingProvider()(modelId),
-      middleware: {
-        specificationVersion: 'v3',
-        transformParams: async ({ params }) => withDefaultAiGatewayCacheHeaders(params),
-      },
-    }),
-  )
+  return aiGatewayChatModel(modelId)
 }
 
 export function aiGatewayChatModel(modelId: string) {
   return withAiGatewayDevTools(
     wrapLanguageModel({
-      model: getAiGatewayProvider().chat(modelId),
-      middleware: {
-        specificationVersion: 'v3',
-        transformParams: async ({ params, type }) =>
-          normalizeAiGatewayChatProviderOptions(
-            withDefaultAiGatewayCacheHeaders(addAiGatewayReasoningRawChunks(params, type)),
-          ),
-        wrapGenerate: async ({ doGenerate }) => {
-          const result = await doGenerate()
-
-          return {
-            ...result,
-            content: injectAiGatewayChatReasoningContent(
-              result.content,
-              result.response as AiGatewayChatResponse | undefined,
-            ),
-          }
-        },
-        wrapStream: async ({ doStream, params }) => {
-          const result = await doStream()
-          if (!isReasoningEnabled(params)) return result
-
-          return { ...result, stream: injectAiGatewayChatReasoningStream(result.stream) }
-        },
-      },
+      model: createLazyAiGatewayLanguageModel({
+        modelId,
+        providerId: OPENAI_CHAT_PROVIDER_ID,
+        resolve: () => getAiGatewayProvider().chat(modelId),
+      }),
+      middleware: createAiGatewayLanguageModelMiddleware(modelId),
     }),
   )
 }
 
 export function aiGatewayEmbeddingModel(modelId: string) {
-  return getAiGatewayProvider().embeddingModel(modelId)
+  return wrapEmbeddingModel({
+    model: createLazyAiGatewayEmbeddingModel(modelId),
+    middleware: {
+      specificationVersion: 'v3',
+      wrapEmbed: ({ doEmbed }) =>
+        Effect.runPromise(
+          withAiGatewayConcurrency(
+            withAiGatewayResilience(
+              'ai-gateway.embed',
+              Effect.tryPromise({ try: doEmbed, catch: (cause) => classifyAiGatewayError('ai-gateway.embed', cause) }),
+            ),
+          ).pipe(Effect.withSpan('AiGateway.embed'), Effect.annotateSpans({ modelId })),
+        ),
+    },
+  })
 }
 
 export { DEFAULT_AI_GATEWAY_URL, normalizeAiGatewayUrl }