feat(ai): add native Bedrock Mantle support (#40119)

Co-authored-by: Aiden Cline <rekram1-node@users.noreply.github.com>

packages/ai/src/protocols/utils/bedrock-auth.ts

@@ -22,6 +22,8 @@ const signRequest = (input: {
   readonly body: string
   readonly headers: Headers.Headers
   readonly credentials: Credentials
+  readonly service: string
+  readonly name: string
 }) =>
   Effect.tryPromise({
     try: async () => {
@@ -34,23 +36,26 @@ const signRequest = (input: {
         accessKeyId: input.credentials.accessKeyId,
         secretAccessKey: input.credentials.secretAccessKey,
         sessionToken: input.credentials.sessionToken,
-        service: "bedrock",
+        service: input.service,
       }).sign()
       return Object.fromEntries(signed.headers.entries())
     },
     catch: (error) =>
       ProviderShared.invalidRequest(
-        `Bedrock Converse SigV4 signing failed: ${error instanceof Error ? error.message : String(error)}`,
+        `${input.name} SigV4 signing failed: ${error instanceof Error ? error.message : String(error)}`,
       ),
   })
 
 /** Sign the exact JSON bytes with SigV4 using credentials configured on the route. */
-export const sigV4 = (credentials: Credentials | undefined) =>
+export const sigV4 = (
+  credentials: Credentials | undefined,
+  options: { readonly service?: string; readonly name?: string } = {},
+) =>
   Auth.custom((input: AuthInput) => {
     return Effect.gen(function* () {
       if (!credentials) {
         return yield* ProviderShared.invalidRequest(
-          "Bedrock Converse requires either route bearer auth or AWS credentials configured on the route",
+          `${options.name ?? "Bedrock Converse"} requires either route bearer auth or AWS credentials configured on the route`,
         )
       }
       const headersForSigning = Headers.set(input.headers, "content-type", "application/json")
@@ -59,6 +64,8 @@ export const sigV4 = (credentials: Credentials | undefined) =>
         body: input.body,
         headers: headersForSigning,
         credentials,
+        service: options.service ?? "bedrock",
+        name: options.name ?? "Bedrock Converse",
       })
       return Headers.setAll(headersForSigning, signed)
     })

packages/ai/src/providers/amazon-bedrock-mantle.ts

@@ -0,0 +1,107 @@
+import { Auth } from "../route/auth"
+import type { Route as RouteDef, RouteDefaultsInput } from "../route/client"
+import type { ProviderPackage } from "../provider-package"
+import { OpenAIChat } from "../protocols/openai-chat"
+import { OpenAIResponses } from "../protocols/openai-responses"
+import { BedrockAuth, type Credentials } from "../protocols/utils/bedrock-auth"
+import { ProviderID, type ModelID } from "../schema"
+import { withOpenAIOptions, type OpenAIProviderOptionsInput } from "./openai-options"
+
+export const id = ProviderID.make("amazon-bedrock")
+
+export type Config = RouteDefaultsInput & {
+  readonly apiKey?: string
+  readonly baseURL?: string
+  readonly credentials?: Credentials
+  readonly region?: string
+  readonly providerOptions?: OpenAIProviderOptionsInput
+}
+
+export interface Settings extends ProviderPackage.Settings {
+  readonly apiKey?: string
+  readonly auth?: "bearer" | "sigv4"
+  readonly baseURL?: string
+  readonly credentials?: Credentials
+  readonly region?: string
+  readonly providerOptions?: OpenAIProviderOptionsInput
+}
+
+const responsesRoute = OpenAIResponses.route.with({
+  id: "bedrock-mantle-responses",
+  provider: id,
+})
+
+const chatRoute = OpenAIChat.route.with({
+  id: "bedrock-mantle-chat",
+  provider: id,
+})
+
+export const routes = [responsesRoute, chatRoute]
+
+const configuredRoute = <Body, Prepared>(route: RouteDef<Body, Prepared>, input: Config) => {
+  const region = input.region ?? input.credentials?.region ?? "us-east-1"
+  const credentials = input.credentials === undefined ? undefined : { ...input.credentials, region }
+  return route.with({
+    endpoint: { baseURL: input.baseURL ?? `https://bedrock-mantle.${region}.api.aws/v1` },
+    auth:
+      input.apiKey === undefined
+        ? BedrockAuth.sigV4(credentials, { service: "bedrock-mantle", name: "Bedrock Mantle" })
+        : Auth.bearer(input.apiKey),
+  })
+}
+
+const defaults = (input: Config) => {
+  const { apiKey: _, baseURL: _baseURL, credentials: _credentials, region: _region, ...rest } = input
+  return rest
+}
+
+export const configure = (input: Config = {}) => {
+  const configuredResponsesRoute = configuredRoute(responsesRoute, input)
+  const configuredChatRoute = configuredRoute(chatRoute, input)
+  const modelDefaults = defaults(input)
+  const responses = (modelID: string | ModelID) =>
+    configuredResponsesRoute
+      .with(withOpenAIOptions(modelID, modelDefaults))
+      .model<OpenAIProviderOptionsInput>({ id: modelID })
+  const chat = (modelID: string | ModelID) =>
+    configuredChatRoute
+      .with(withOpenAIOptions(modelID, modelDefaults))
+      .model<OpenAIProviderOptionsInput>({ id: modelID })
+
+  return {
+    id,
+    model: chat,
+    chat,
+    responses,
+    configure,
+  }
+}
+
+export const provider = configure()
+
+const config = (settings: Settings): Config => {
+  if (settings.auth === "bearer" && settings.apiKey === undefined)
+    throw new Error("Amazon Bedrock Mantle bearer auth requires apiKey")
+  if (settings.auth === "sigv4" && settings.apiKey !== undefined)
+    throw new Error("Amazon Bedrock Mantle SigV4 auth does not accept apiKey")
+  return {
+    apiKey: settings.auth === "sigv4" ? undefined : settings.apiKey,
+    baseURL: settings.baseURL,
+    credentials: settings.credentials,
+    headers: settings.headers === undefined ? undefined : { ...settings.headers },
+    http: settings.body === undefined ? undefined : { body: { ...settings.body } },
+    limits: settings.limits,
+    providerOptions: settings.providerOptions,
+    region: settings.region,
+  }
+}
+
+export const chatModel: ProviderPackage.Definition<Settings, OpenAIProviderOptionsInput>["model"] = (
+  modelID,
+  settings,
+) => configure(config(settings)).chat(modelID)
+export const responsesModel: ProviderPackage.Definition<Settings, OpenAIProviderOptionsInput>["model"] = (
+  modelID,
+  settings,
+) => configure(config(settings)).responses(modelID)
+export const model = chatModel

packages/ai/src/providers/amazon-bedrock/mantle.ts

@@ -0,0 +1,2 @@
+export { chatModel as model } from "../amazon-bedrock-mantle"
+export type { Settings } from "../amazon-bedrock-mantle"

packages/ai/src/providers/amazon-bedrock/mantle/chat.ts

@@ -0,0 +1,2 @@
+export { chatModel as model } from "../../amazon-bedrock-mantle"
+export type { Settings } from "../../amazon-bedrock-mantle"

packages/ai/src/providers/amazon-bedrock/mantle/responses.ts

@@ -0,0 +1,2 @@
+export { responsesModel as model } from "../../amazon-bedrock-mantle"
+export type { Settings } from "../../amazon-bedrock-mantle"

packages/ai/src/providers/index.ts

@@ -1,6 +1,7 @@
 export * as Anthropic from "./anthropic"
 export * as AnthropicCompatible from "./anthropic-compatible"
 export * as AmazonBedrock from "./amazon-bedrock"
+export * as AmazonBedrockMantle from "./amazon-bedrock-mantle"
 export * as Azure from "./azure"
 export * as Cloudflare from "./cloudflare"
 export { CloudflareAIGateway, CloudflareWorkersAI } from "./cloudflare"