From 0199abfd110bfc225512985ee62bab64ca332011 Mon Sep 17 00:00:00 2001
From: Mario Zechner <badlogicgames@gmail.com>
Date: Wed, 4 Mar 2026 20:47:41 +0100
Subject: [PATCH] fix(binary): preload Bedrock provider module for Bun bundle

refs #1814
---
 packages/ai/bedrock-provider.d.ts              |  1 +
 packages/ai/bedrock-provider.js                |  1 +
 packages/ai/package.json                       |  6 ++++++
 packages/ai/src/bedrock-provider.ts            |  6 ++++++
 packages/ai/src/providers/register-builtins.ts | 17 +++++++++++++----
 packages/coding-agent/src/cli.ts               |  3 +++
 6 files changed, 30 insertions(+), 4 deletions(-)
 create mode 100644 packages/ai/bedrock-provider.d.ts
 create mode 100644 packages/ai/bedrock-provider.js
 create mode 100644 packages/ai/src/bedrock-provider.ts

diff --git a/packages/ai/bedrock-provider.d.ts b/packages/ai/bedrock-provider.d.ts
new file mode 100644
index 00000000..a66eabee
--- /dev/null
+++ b/packages/ai/bedrock-provider.d.ts
@@ -0,0 +1 @@
+export * from "./dist/bedrock-provider.js";
diff --git a/packages/ai/bedrock-provider.js b/packages/ai/bedrock-provider.js
new file mode 100644
index 00000000..a66eabee
--- /dev/null
+++ b/packages/ai/bedrock-provider.js
@@ -0,0 +1 @@
+export * from "./dist/bedrock-provider.js";
diff --git a/packages/ai/package.json b/packages/ai/package.json
index 9bd2721a..b4ea5268 100644
--- a/packages/ai/package.json
+++ b/packages/ai/package.json
@@ -13,6 +13,10 @@
 		"./oauth": {
 			"types": "./oauth.d.ts",
 			"import": "./oauth.js"
+		},
+		"./bedrock-provider": {
+			"types": "./bedrock-provider.d.ts",
+			"import": "./bedrock-provider.js"
 		}
 	},
 	"bin": {
@@ -22,6 +26,8 @@
 		"dist",
 		"oauth.js",
 		"oauth.d.ts",
+		"bedrock-provider.js",
+		"bedrock-provider.d.ts",
 		"README.md"
 	],
 	"scripts": {
diff --git a/packages/ai/src/bedrock-provider.ts b/packages/ai/src/bedrock-provider.ts
new file mode 100644
index 00000000..07305418
--- /dev/null
+++ b/packages/ai/src/bedrock-provider.ts
@@ -0,0 +1,6 @@
+import { streamBedrock, streamSimpleBedrock } from "./providers/amazon-bedrock.js";
+
+export const bedrockProviderModule = {
+	streamBedrock,
+	streamSimpleBedrock,
+};
diff --git a/packages/ai/src/providers/register-builtins.ts b/packages/ai/src/providers/register-builtins.ts
index 6c13095c..b497d133 100644
--- a/packages/ai/src/providers/register-builtins.ts
+++ b/packages/ai/src/providers/register-builtins.ts
@@ -22,26 +22,35 @@ interface BedrockProviderModule {
 		model: Model<"bedrock-converse-stream">,
 		context: Context,
 		options?: StreamOptions,
-	) => AssistantMessageEventStream;
+	) => AsyncIterable<AssistantMessageEvent>;
 	streamSimpleBedrock: (
 		model: Model<"bedrock-converse-stream">,
 		context: Context,
 		options?: SimpleStreamOptions,
-	) => AssistantMessageEventStream;
+	) => AsyncIterable<AssistantMessageEvent>;
 }
 
 type DynamicImport = (specifier: string) => Promise<unknown>;
 
 const dynamicImport = new Function("specifier", "return import(specifier);") as DynamicImport;
 
+let bedrockProviderModuleOverride: BedrockProviderModule | undefined;
+
+export function setBedrockProviderModule(module: BedrockProviderModule): void {
+	bedrockProviderModuleOverride = module;
+}
+
 async function loadBedrockProviderModule(): Promise<BedrockProviderModule> {
+	if (bedrockProviderModuleOverride) {
+		return bedrockProviderModuleOverride;
+	}
 	const module = await dynamicImport("./amazon-bedrock.js");
 	return module as BedrockProviderModule;
 }
 
-function forwardStream(target: AssistantMessageEventStream, source: AssistantMessageEventStream): void {
+function forwardStream(target: AssistantMessageEventStream, source: AsyncIterable<AssistantMessageEvent>): void {
 	(async () => {
-		for await (const event of source as AsyncIterable<AssistantMessageEvent>) {
+		for await (const event of source) {
 			target.push(event);
 		}
 		target.end();
diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts
index 131fb6ad..d05106c8 100644
--- a/packages/coding-agent/src/cli.ts
+++ b/packages/coding-agent/src/cli.ts
@@ -7,9 +7,12 @@
  */
 process.title = "pi";
 
+import { setBedrockProviderModule } from "@mariozechner/pi-ai";
+import { bedrockProviderModule } from "@mariozechner/pi-ai/bedrock-provider";
 import { EnvHttpProxyAgent, setGlobalDispatcher } from "undici";
 import { main } from "./main.js";
 
 setGlobalDispatcher(new EnvHttpProxyAgent());
+setBedrockProviderModule(bedrockProviderModule);
 
 main(process.argv.slice(2));