@lota-sdk/core 0.4.10 → 0.4.12

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lota-sdk/core",
-  "version": "0.4.10",
+  "version": "0.4.12",
   "type": "module",
   "main": "./src/index.ts",
   "types": "./src/index.ts",
@@ -31,10 +31,10 @@
     "@ai-sdk/openai": "^3.0.53",
     "@chat-adapter/slack": "^4.26.0",
     "@chat-adapter/state-ioredis": "^4.26.0",
-    "@lota-sdk/shared": "0.4.10",
+    "@lota-sdk/shared": "0.4.12",
     "@mendable/firecrawl-js": "^4.18.3",
     "@surrealdb/node": "^3.0.3",
-    "ai": "^6.0.167",
+    "ai": "^6.0.168",
     "bullmq": "^5.74.1",
     "chat": "^4.26.0",
     "effect": "^4.0.0-beta.50",

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

@@ -6,7 +6,6 @@ import { Cause, Clock, Context, Duration, Effect, ExecutionPlan, Fiber, Layer, S
 
 import { DEFAULT_AI_GATEWAY_URL } from '../config/constants'
 import { AiGenerationError, ConfigurationError } from '../effect/errors'
-import { resolveLotaService } from '../effect/runtime'
 import { RuntimeConfigServiceTag } from '../effect/services'
 import { getDirectOpenRouterProvider, normalizeDirectOpenRouterModelId } from '../openrouter/direct-provider'
 import { isRecord, readString } from '../utils/string'
@@ -24,6 +23,8 @@ type AiGatewayGeneratedContent = AiGatewayGenerateResult['content'][number]
 type AiGatewayStreamPart = AiGatewayStreamResult['stream'] extends ReadableStream<infer T> ? T : never
 type AiGatewayProviderOptions = NonNullable<AiGatewayCallOptions['providerOptions']>
 type AiGatewayAttemptResult<A> = { source: string; result: A }
+// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
+type AiGatewayRunFork = <A, E>(effect: Effect.Effect<A, E, never>) => Fiber.Fiber<A, E | unknown>
 
 class AiGatewayGenerateAttempt extends Context.Service<
   AiGatewayGenerateAttempt,
@@ -271,12 +272,13 @@ function withAiGatewayResilience<A>(source: string, effect: Effect.Effect<A, AiG
 function withAiGatewayStreamIdleTimeout(
   stream: ReadableStream<AiGatewayStreamPart>,
   source: string,
+  runFork: AiGatewayRunFork,
   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 idleTimeoutFiber: Fiber.Fiber<unknown, unknown> | null = null
+  let bodyPumpFiber: Fiber.Fiber<unknown, unknown> | null = null
   let finalized = false
 
   const finalize = () => {
@@ -285,9 +287,9 @@ function withAiGatewayStreamIdleTimeout(
     onFinalize?.()
   }
 
-  const interruptFiber = (fiber: ReturnType<typeof Effect.runFork> | null) => {
+  const interruptFiber = (fiber: Fiber.Fiber<unknown, unknown> | null) => {
     if (!fiber) return
-    void Effect.runFork(Fiber.interrupt(fiber))
+    void runFork(Fiber.interrupt(fiber))
   }
 
   const stopIdleTimeout = () => {
@@ -351,7 +353,7 @@ function withAiGatewayStreamIdleTimeout(
 
   const resetIdleTimeout = (controller: ReadableStreamDefaultController<AiGatewayStreamPart>) => {
     stopIdleTimeout()
-    idleTimeoutFiber = Effect.runFork(
+    idleTimeoutFiber = runFork(
       Effect.sleep(Duration.millis(AI_GATEWAY_STREAM_IDLE_TIMEOUT_MS)).pipe(
         Effect.flatMap(() =>
           Effect.gen(function* () {
@@ -417,7 +419,7 @@ function withAiGatewayStreamIdleTimeout(
     start(controller) {
       const streamReader = stream.getReader()
       reader = streamReader
-      bodyPumpFiber = Effect.runFork(pumpStreamEffect(streamReader, controller))
+      bodyPumpFiber = runFork(pumpStreamEffect(streamReader, controller))
     },
     cancel(reason) {
       closed = true
@@ -493,40 +495,22 @@ export const AiGatewayLive = Layer.effect(
 
 type AiGatewayRuntimeConfig = Context.Service.Shape<typeof RuntimeConfigServiceTag>
 
-let currentAiGateway: AiGatewayTag['Service'] | null = null
-let currentAiGatewayRuntimeConfig: AiGatewayRuntimeConfig | null = null
-
-export function configureAiGatewayRuntimeAccessors(params: {
-  aiGateway: AiGatewayTag['Service']
-  runtimeConfig: AiGatewayRuntimeConfig
-}): void {
-  currentAiGateway = params.aiGateway
-  currentAiGatewayRuntimeConfig = params.runtimeConfig
-}
-
-export function clearAiGatewayRuntimeAccessors(): void {
-  currentAiGateway = null
-  currentAiGatewayRuntimeConfig = null
-}
-
-function getAiGateway(): AiGatewayTag['Service'] {
-  return currentAiGateway ?? resolveLotaService(AiGatewayTag)
-}
-
-function getAiGatewayRuntimeConfig(): AiGatewayRuntimeConfig {
-  return currentAiGatewayRuntimeConfig ?? resolveLotaService(RuntimeConfigServiceTag)
-}
-
-function withAiGatewayConcurrency<A>(effect: Effect.Effect<A, AiGenerationError>): Effect.Effect<A, AiGenerationError> {
-  return getAiGateway().semaphore.withPermit(effect)
+function withAiGatewayConcurrency<A>(
+  effect: Effect.Effect<A, AiGenerationError>,
+): Effect.Effect<A, AiGenerationError, AiGatewayTag> {
+  return Effect.gen(function* () {
+    const gateway = yield* AiGatewayTag
+    return yield* gateway.semaphore.withPermit(effect)
+  })
 }
 
 function withAiGatewayStreamConcurrency(
   effect: Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError>,
-): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
+  runFork: AiGatewayRunFork,
+): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError, AiGatewayTag> {
   return Effect.uninterruptibleMask((restore) =>
     Effect.gen(function* () {
-      const { semaphore } = getAiGateway()
+      const { semaphore } = yield* AiGatewayTag
       const currentContext = yield* Effect.context<never>()
       yield* semaphore.take(1)
 
@@ -551,7 +535,7 @@ function withAiGatewayStreamConcurrency(
         ...attempt,
         result: {
           ...attempt.result,
-          stream: withAiGatewayStreamIdleTimeout(attempt.result.stream, attempt.source, release),
+          stream: withAiGatewayStreamIdleTimeout(attempt.result.stream, attempt.source, runFork, release),
         },
       }
     }),
@@ -611,6 +595,52 @@ export function extractAiGatewayChatReasoningDeltaText(rawChunk: unknown): strin
   return null
 }
 
+function findAiGatewayChatReasoningOverlap(previousReasoningText: string, nextReasoningText: string): number {
+  const maxOverlap = Math.min(previousReasoningText.length, nextReasoningText.length)
+
+  for (let overlapLength = maxOverlap; overlapLength > 0; overlapLength -= 1) {
+    if (previousReasoningText.slice(-overlapLength) === nextReasoningText.slice(0, overlapLength)) {
+      return overlapLength
+    }
+  }
+
+  return 0
+}
+
+function deriveAiGatewayChatReasoningDeltaText(params: { previousReasoningText: string; rawChunk: unknown }): {
+  delta: string | null
+  nextReasoningText: string
+} {
+  const extractedText = extractAiGatewayChatReasoningDeltaText(params.rawChunk)
+  if (!extractedText) {
+    return { delta: null, nextReasoningText: params.previousReasoningText }
+  }
+
+  if (params.previousReasoningText.length === 0) {
+    return { delta: extractedText, nextReasoningText: extractedText }
+  }
+
+  if (extractedText === params.previousReasoningText) {
+    return { delta: null, nextReasoningText: params.previousReasoningText }
+  }
+
+  if (extractedText.startsWith(params.previousReasoningText)) {
+    const delta = extractedText.slice(params.previousReasoningText.length)
+    return { delta: delta.length > 0 ? delta : null, nextReasoningText: extractedText }
+  }
+
+  const overlapLength = findAiGatewayChatReasoningOverlap(params.previousReasoningText, extractedText)
+  if (overlapLength > 0) {
+    const delta = extractedText.slice(overlapLength)
+    return { delta: delta.length > 0 ? delta : null, nextReasoningText: `${params.previousReasoningText}${delta}` }
+  }
+
+  // Some providers emit true deltas, others resend the full reasoning-so-far.
+  // If the chunk is not a prefix extension, treat it as a standalone delta and
+  // append it to the accumulated reasoning text.
+  return { delta: extractedText, nextReasoningText: `${params.previousReasoningText}${extractedText}` }
+}
+
 export function injectAiGatewayChatReasoningContent(
   content: readonly AiGatewayGeneratedContent[],
   response?: AiGatewayChatResponse,
@@ -638,18 +668,20 @@ function isOpenRouterModel(modelId: string): boolean {
   return modelId.trim().toLowerCase().startsWith('openrouter/')
 }
 
-function hasDirectOpenRouterFallback(modelId: string): boolean {
-  const config = getAiGatewayRuntimeConfig()
+function hasDirectOpenRouterFallback(config: AiGatewayRuntimeConfig, modelId: string): boolean {
   return isOpenRouterModel(modelId) && Boolean(config.aiGateway.openRouterApiKey?.trim())
 }
 
-function getDirectOpenRouterChatModel(modelId: string): AiGatewayLanguageModel {
-  const config = getAiGatewayRuntimeConfig()
+function getDirectOpenRouterChatModel(config: AiGatewayRuntimeConfig, modelId: string): AiGatewayLanguageModel {
   return getDirectOpenRouterProvider(config.aiGateway.openRouterApiKey).chat(normalizeDirectOpenRouterModelId(modelId))
 }
 
-function shouldFallbackToDirectOpenRouter(modelId: string, error: AiGenerationError): boolean {
-  return hasDirectOpenRouterFallback(modelId) && isRetryableAiGatewayError(error)
+function shouldFallbackToDirectOpenRouter(
+  config: AiGatewayRuntimeConfig,
+  modelId: string,
+  error: AiGenerationError,
+): boolean {
+  return hasDirectOpenRouterFallback(config, modelId) && isRetryableAiGatewayError(error)
 }
 
 function attemptAiGatewayGenerate(
@@ -681,22 +713,25 @@ function attemptAiGatewayStream(
 }
 
 function attemptDirectOpenRouterGenerate(
+  config: AiGatewayRuntimeConfig,
   modelId: string,
   params: AiGatewayCallOptions,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayGenerateResult>, AiGenerationError> {
-  const model = getDirectOpenRouterChatModel(modelId)
+  const model = getDirectOpenRouterChatModel(config, modelId)
   return attemptAiGatewayGenerate('openrouter.generate', () => model.doGenerate(params))
 }
 
 function attemptDirectOpenRouterStream(
+  config: AiGatewayRuntimeConfig,
   modelId: string,
   params: AiGatewayCallOptions,
 ): Effect.Effect<AiGatewayAttemptResult<AiGatewayStreamResult>, AiGenerationError> {
-  const model = getDirectOpenRouterChatModel(modelId)
+  const model = getDirectOpenRouterChatModel(config, modelId)
   return attemptAiGatewayStream('openrouter.stream', () => model.doStream(params))
 }
 
 function executeGenerateAttemptPlan(
+  config: AiGatewayRuntimeConfig,
   modelId: string,
   params: AiGatewayCallOptions,
   doGenerate: () => PromiseLike<AiGatewayGenerateResult>,
@@ -709,7 +744,7 @@ function executeGenerateAttemptPlan(
     return yield* attempt.execute
   })
 
-  if (!hasDirectOpenRouterFallback(modelId)) {
+  if (!hasDirectOpenRouterFallback(config, modelId)) {
     return effect.pipe(
       Effect.provide(primary),
       Effect.withSpan('AiGateway.executeGeneratePlan'),
@@ -723,9 +758,9 @@ function executeGenerateAttemptPlan(
         { provide: primary },
         {
           provide: Layer.succeed(AiGatewayGenerateAttempt, {
-            execute: attemptDirectOpenRouterGenerate(modelId, params),
+            execute: attemptDirectOpenRouterGenerate(config, modelId, params),
           }),
-          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(modelId, error),
+          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(config, modelId, error),
         },
       ),
     ),
@@ -735,6 +770,7 @@ function executeGenerateAttemptPlan(
 }
 
 function executeStreamAttemptPlan(
+  config: AiGatewayRuntimeConfig,
   modelId: string,
   params: AiGatewayCallOptions,
   doStream: () => PromiseLike<AiGatewayStreamResult>,
@@ -747,7 +783,7 @@ function executeStreamAttemptPlan(
     return yield* attempt.execute
   })
 
-  if (!hasDirectOpenRouterFallback(modelId)) {
+  if (!hasDirectOpenRouterFallback(config, modelId)) {
     return effect.pipe(
       Effect.provide(primary),
       Effect.withSpan('AiGateway.executeStreamPlan'),
@@ -760,8 +796,10 @@ function executeStreamAttemptPlan(
       ExecutionPlan.make(
         { provide: primary },
         {
-          provide: Layer.succeed(AiGatewayStreamAttempt, { execute: attemptDirectOpenRouterStream(modelId, params) }),
-          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(modelId, error),
+          provide: Layer.succeed(AiGatewayStreamAttempt, {
+            execute: attemptDirectOpenRouterStream(config, modelId, params),
+          }),
+          while: (error: AiGenerationError) => shouldFallbackToDirectOpenRouter(config, modelId, error),
         },
       ),
     ),
@@ -794,6 +832,8 @@ export function injectAiGatewayChatReasoningStream(
   const reasoningId = 'ai-gateway-reasoning-0'
   let reasoningOpen = false
   let reasoningClosed = false
+  let reasoningText = ''
+  let nativeReasoningSeen = false
 
   return stream.pipeThrough(
     new TransformStream<AiGatewayStreamPart, AiGatewayStreamPart>({
@@ -805,11 +845,25 @@ export function injectAiGatewayChatReasoningStream(
           reasoningClosed = true
         }
 
+        if (chunk.type === 'reasoning-start' || chunk.type === 'reasoning-delta' || chunk.type === 'reasoning-end') {
+          nativeReasoningSeen = true
+          closeReasoning()
+          controller.enqueue(chunk)
+          return
+        }
+
         if (chunk.type === 'raw') {
-          const reasoningDelta = reasoningClosed ? null : extractAiGatewayChatReasoningDeltaText(chunk.rawValue)
+          const reasoningDeltaState =
+            reasoningClosed || nativeReasoningSeen
+              ? null
+              : deriveAiGatewayChatReasoningDeltaText({
+                  previousReasoningText: reasoningText,
+                  rawChunk: chunk.rawValue,
+                })
           controller.enqueue(chunk)
 
-          if (reasoningDelta) {
+          if (reasoningDeltaState?.delta) {
+            reasoningText = reasoningDeltaState.nextReasoningText
             if (!reasoningOpen) {
               controller.enqueue({ type: 'reasoning-start', id: reasoningId } satisfies AiGatewayStreamPart)
               reasoningOpen = true
@@ -818,7 +872,7 @@ export function injectAiGatewayChatReasoningStream(
             controller.enqueue({
               type: 'reasoning-delta',
               id: reasoningId,
-              delta: reasoningDelta,
+              delta: reasoningDeltaState.delta,
             } satisfies AiGatewayStreamPart)
           }
           return
@@ -849,7 +903,56 @@ function addAiGatewayReasoningRawChunks(
   return { ...params, includeRawChunks: true }
 }
 
-function createAiGatewayLanguageModelMiddleware(modelId: string): LanguageModelMiddleware {
+function resolveProviderModel(
+  provider: ReturnType<typeof createOpenAI>,
+  modelId: string,
+  providerId: string,
+): AiGatewayLanguageModel {
+  return providerId === OPENAI_CHAT_PROVIDER_ID ? provider.chat(modelId) : provider(modelId)
+}
+
+// Module-level Promise slot that `createLotaRuntime` populates during boot.
+// This is a legitimate per-process singleton (mirrors the worker bootstrap
+// pattern in `workers/bootstrap.ts`): the AI gateway middleware is dispatched
+// by AI SDK callers that live outside Effect context, so the middleware needs
+// a way to run gateway Effects without capturing a `ManagedRuntime` through
+// every `aiGatewayModel(modelId)` call site.
+//
+// Only `createLotaRuntime` writes to the slot; resetting on disconnect is a
+// Phase 3b concern — for now it stays alive for the process lifetime.
+let aiGatewayRuntimeReady: Promise<{
+  gateway: Context.Service.Shape<typeof AiGatewayTag>
+  runtimeConfig: Context.Service.Shape<typeof RuntimeConfigServiceTag>
+  runPromise: <A, E>(effect: Effect.Effect<A, E, never>) => Promise<A>
+  runFork: AiGatewayRunFork
+}> | null = null
+
+export function bindAiGatewayRuntime(params: {
+  gateway: Context.Service.Shape<typeof AiGatewayTag>
+  runtimeConfig: Context.Service.Shape<typeof RuntimeConfigServiceTag>
+  runPromise: <A, E>(effect: Effect.Effect<A, E, never>) => Promise<A>
+  runFork: AiGatewayRunFork
+}): void {
+  aiGatewayRuntimeReady = Promise.resolve(params)
+}
+
+export function clearAiGatewayRuntime(): void {
+  aiGatewayRuntimeReady = null
+}
+
+async function getAiGatewayRuntime(): Promise<{
+  gateway: Context.Service.Shape<typeof AiGatewayTag>
+  runtimeConfig: Context.Service.Shape<typeof RuntimeConfigServiceTag>
+  runPromise: <A, E>(effect: Effect.Effect<A, E, never>) => Promise<A>
+  runFork: AiGatewayRunFork
+}> {
+  if (!aiGatewayRuntimeReady) {
+    throw new Error('AI gateway runtime has not been initialized. Call createLotaRuntime() first.')
+  }
+  return aiGatewayRuntimeReady
+}
+
+function createAiGatewayLanguageModelMiddleware(modelId: string, providerId: string): LanguageModelMiddleware {
   return {
     specificationVersion: 'v3',
     transformParams: ({ params, type }) =>
@@ -858,10 +961,12 @@ function createAiGatewayLanguageModelMiddleware(modelId: string): LanguageModelM
           addAiGatewayReasoningRawChunks(normalizeAiGatewayChatProviderOptions(params, modelId), type),
         ),
       ),
-    wrapGenerate: ({ doGenerate, params }) =>
-      Effect.runPromise(
+    wrapGenerate: async ({ params }) => {
+      const { gateway, runtimeConfig, runPromise } = await getAiGatewayRuntime()
+      const model = resolveProviderModel(gateway.provider, modelId, providerId)
+      return runPromise(
         withAiGatewayConcurrency(
-          executeGenerateAttemptPlan(modelId, params, doGenerate).pipe(
+          executeGenerateAttemptPlan(runtimeConfig, modelId, params, () => model.doGenerate(params)).pipe(
             Effect.map(({ result }) => ({
               ...result,
               content: injectAiGatewayChatReasoningContent(
@@ -870,12 +975,15 @@ function createAiGatewayLanguageModelMiddleware(modelId: string): LanguageModelM
               ),
             })),
           ),
-        ),
-      ),
-    wrapStream: ({ doStream, params }) =>
-      Effect.runPromise(
+        ).pipe(Effect.provideService(AiGatewayTag, gateway)),
+      )
+    },
+    wrapStream: async ({ params }) => {
+      const { gateway, runtimeConfig, runPromise, runFork } = await getAiGatewayRuntime()
+      const model = resolveProviderModel(gateway.provider, modelId, providerId)
+      return runPromise(
         withAiGatewayStreamConcurrency(
-          executeStreamAttemptPlan(modelId, params, doStream).pipe(
+          executeStreamAttemptPlan(runtimeConfig, modelId, params, () => model.doStream(params)).pipe(
             Effect.map((attempt) => ({
               ...attempt,
               result: isReasoningEnabled(params)
@@ -883,8 +991,12 @@ function createAiGatewayLanguageModelMiddleware(modelId: string): LanguageModelM
                 : attempt.result,
             })),
           ),
-        ).pipe(Effect.map(({ result }) => result)),
-      ),
+          runFork,
+        )
+          .pipe(Effect.map(({ result }) => result))
+          .pipe(Effect.provideService(AiGatewayTag, gateway)),
+      )
+    },
   }
 }
 
@@ -921,36 +1033,42 @@ function withAiGatewayDevTools<TModel extends AiGatewayLanguageModel>(model: TMo
   return wrapLanguageModel({ model, middleware: devToolsMiddleware() }) as TModel
 }
 
-function createLazyAiGatewayLanguageModel(params: {
-  modelId: string
-  providerId: string
-  resolve: () => AiGatewayLanguageModel
-}): AiGatewayLanguageModel {
+function createAiGatewayLanguageModelPlaceholder(modelId: string, providerId: string): AiGatewayLanguageModel {
+  const unreachable = (method: string) =>
+    Promise.reject(
+      new Error(
+        `[ai-gateway] AiGateway language model ${modelId}.${method} was invoked without the gateway middleware; ` +
+          'this call path should be fully handled by createAiGatewayLanguageModelMiddleware.',
+      ),
+    )
+
   return {
     specificationVersion: 'v3',
-    provider: params.providerId,
-    modelId: params.modelId,
+    provider: providerId,
+    modelId,
     supportedUrls: {},
-    doGenerate: (options) => params.resolve().doGenerate(options),
-    doStream: (options) => params.resolve().doStream(options),
+    doGenerate: () => unreachable('doGenerate'),
+    doStream: () => unreachable('doStream'),
   }
 }
 
-function createLazyAiGatewayEmbeddingModel(modelId: string): AiGatewayEmbeddingModel {
+function createAiGatewayEmbeddingModelPlaceholder(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),
+    doEmbed: () =>
+      Promise.reject(
+        new Error(
+          `[ai-gateway] AiGateway embedding model ${modelId}.doEmbed was invoked without the gateway middleware; ` +
+            'this call path should be fully handled by aiGatewayEmbeddingModel middleware.',
+        ),
+      ),
   }
 }
 
-export function getAiGatewayProvider() {
-  return getAiGateway().provider
-}
-
 export function aiGatewayModel(modelId: string) {
   if (isOpenRouterModel(modelId)) {
     return aiGatewayChatModel(modelId)
@@ -958,12 +1076,8 @@ export function aiGatewayModel(modelId: string) {
 
   return withAiGatewayDevTools(
     wrapLanguageModel({
-      model: createLazyAiGatewayLanguageModel({
-        modelId,
-        providerId: OPENAI_RESPONSES_PROVIDER_ID,
-        resolve: () => getAiGatewayProvider()(modelId),
-      }),
-      middleware: createAiGatewayLanguageModelMiddleware(modelId),
+      model: createAiGatewayLanguageModelPlaceholder(modelId, OPENAI_RESPONSES_PROVIDER_ID),
+      middleware: createAiGatewayLanguageModelMiddleware(modelId, OPENAI_RESPONSES_PROVIDER_ID),
     }),
   )
 }
@@ -975,30 +1089,32 @@ export function aiGatewayOpenRouterResponseHealingModel(modelId: string) {
 export function aiGatewayChatModel(modelId: string) {
   return withAiGatewayDevTools(
     wrapLanguageModel({
-      model: createLazyAiGatewayLanguageModel({
-        modelId,
-        providerId: OPENAI_CHAT_PROVIDER_ID,
-        resolve: () => getAiGatewayProvider().chat(modelId),
-      }),
-      middleware: createAiGatewayLanguageModelMiddleware(modelId),
+      model: createAiGatewayLanguageModelPlaceholder(modelId, OPENAI_CHAT_PROVIDER_ID),
+      middleware: createAiGatewayLanguageModelMiddleware(modelId, OPENAI_CHAT_PROVIDER_ID),
     }),
   )
 }
 
 export function aiGatewayEmbeddingModel(modelId: string) {
   return wrapEmbeddingModel({
-    model: createLazyAiGatewayEmbeddingModel(modelId),
+    model: createAiGatewayEmbeddingModelPlaceholder(modelId),
     middleware: {
       specificationVersion: 'v3',
-      wrapEmbed: ({ doEmbed }) =>
-        Effect.runPromise(
+      wrapEmbed: async ({ params }) => {
+        const { gateway, runPromise } = await getAiGatewayRuntime()
+        const embeddingModel = gateway.provider.embeddingModel(modelId)
+        return 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 })),
-        ),
+              Effect.tryPromise({
+                try: () => embeddingModel.doEmbed(params),
+                catch: (cause) => classifyAiGatewayError('ai-gateway.embed', cause),
+              }),
+            ).pipe(Effect.withSpan('AiGateway.embed'), Effect.annotateSpans({ modelId })),
+          ).pipe(Effect.provideService(AiGatewayTag, gateway)),
+        )
+      },
     },
   })
 }

package/src/ai-gateway/index.ts

@@ -1,2 +1,17 @@
-export * from './ai-gateway'
+export {
+  AiGatewayLive,
+  AiGatewayTag,
+  DEFAULT_AI_GATEWAY_URL,
+  aiGatewayChatModel,
+  aiGatewayEmbeddingModel,
+  aiGatewayModel,
+  aiGatewayOpenRouterResponseHealingModel,
+  bindAiGatewayRuntime,
+  extractAiGatewayChatReasoningDeltaText,
+  extractAiGatewayChatReasoningText,
+  injectAiGatewayChatReasoningContent,
+  injectAiGatewayChatReasoningStream,
+  normalizeAiGatewayChatProviderOptions,
+  normalizeAiGatewayUrl,
+} from './ai-gateway'
 export * from './cache-headers'