From 91d91caec1de8254ef62ec1fc7d6e034ecef7650 Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <gmassadas@cloudflare.com>
Date: Wed, 24 Apr 2024 18:07:25 +0100
Subject: [PATCH 1/4] Add workers ai staging flag

---
 .changeset/thin-wasps-bathe.md                      |  5 +++++
 packages/wrangler/src/config/index.ts               | 13 ++++++++++---
 .../deployment-bundle/create-worker-upload-form.ts  |  3 ++-
 packages/wrangler/src/deployment-bundle/worker.ts   |  1 +
 4 files changed, 18 insertions(+), 4 deletions(-)
 create mode 100644 .changeset/thin-wasps-bathe.md

diff --git a/.changeset/thin-wasps-bathe.md b/.changeset/thin-wasps-bathe.md
new file mode 100644
index 00000000000..78b158df840
--- /dev/null
+++ b/.changeset/thin-wasps-bathe.md
@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+Add `staging` flag to AI binding
diff --git a/packages/wrangler/src/config/index.ts b/packages/wrangler/src/config/index.ts
index 28f75b7de8e..9c36a2a938b 100644
--- a/packages/wrangler/src/config/index.ts
+++ b/packages/wrangler/src/config/index.ts
@@ -182,8 +182,10 @@ export function printBindings(bindings: CfWorkerInit["bindings"]) {
 		return `${s.substring(0, maxLength - 3)}...`;
 	};
 
-	const output: { type: string; entries: { key: string; value: string }[] }[] =
-		[];
+	const output: {
+		type: string;
+		entries: { key: string; value: string | boolean }[];
+	}[] = [];
 
 	const {
 		data_blobs,
@@ -418,9 +420,14 @@ export function printBindings(bindings: CfWorkerInit["bindings"]) {
 	}
 
 	if (ai !== undefined) {
+		const entries: [{ key: string; value: string | boolean }] = [
+			{ key: "Name", value: ai.binding },
+		];
+		if (ai.staging) entries.push({ key: "Staging", value: ai.staging });
+
 		output.push({
 			type: "AI",
-			entries: [{ key: "Name", value: ai.binding }],
+			entries: entries,
 		});
 	}
 
diff --git a/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts b/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts
index 03d5580a399..c4c17c87216 100644
--- a/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts
+++ b/packages/wrangler/src/deployment-bundle/create-worker-upload-form.ts
@@ -42,7 +42,7 @@ export type WorkerMetadataBinding =
 	| { type: "wasm_module"; name: string; part: string }
 	| { type: "text_blob"; name: string; part: string }
 	| { type: "browser"; name: string }
-	| { type: "ai"; name: string }
+	| { type: "ai"; name: string; staging?: boolean }
 	| { type: "version_metadata"; name: string }
 	| { type: "data_blob"; name: string; part: string }
 	| { type: "kv_namespace"; name: string; namespace_id: string }
@@ -329,6 +329,7 @@ export function createWorkerUploadForm(worker: CfWorkerInit): FormData {
 	if (bindings.ai !== undefined) {
 		metadataBindings.push({
 			name: bindings.ai.binding,
+			staging: bindings.ai.staging,
 			type: "ai",
 		});
 	}
diff --git a/packages/wrangler/src/deployment-bundle/worker.ts b/packages/wrangler/src/deployment-bundle/worker.ts
index 73a852f1e2d..e05a5020c48 100644
--- a/packages/wrangler/src/deployment-bundle/worker.ts
+++ b/packages/wrangler/src/deployment-bundle/worker.ts
@@ -111,6 +111,7 @@ export interface CfBrowserBinding {
 
 export interface CfAIBinding {
 	binding: string;
+	staging?: boolean;
 }
 
 /**

From 3cda70a3e3a8cbc4dcce14166ed28f016c8ff072 Mon Sep 17 00:00:00 2001
From: Rahul Sethi <5822355+RamIdeas@users.noreply.github.com>
Date: Fri, 26 Apr 2024 16:11:32 +0100
Subject: [PATCH 2/4] add `staging` option to ai binding config

---
 packages/wrangler/src/config/environment.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/packages/wrangler/src/config/environment.ts b/packages/wrangler/src/config/environment.ts
index 420dd252b7e..42e7238eedb 100644
--- a/packages/wrangler/src/config/environment.ts
+++ b/packages/wrangler/src/config/environment.ts
@@ -632,6 +632,7 @@ export interface EnvironmentNonInheritable {
 	ai:
 		| {
 				binding: string;
+				staging?: boolean;
 		  }
 		| undefined;
 

From b4193708af689e77d4e70c94e1964950f9af648b Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <gmassadas@cloudflare.com>
Date: Fri, 26 Apr 2024 20:12:56 +0100
Subject: [PATCH 3/4] Trigger build


From 0e5133e37774cf66c7beb1f6f8e9d70e048736fd Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <gmassadas@cloudflare.com>
Date: Sat, 27 Apr 2024 17:05:00 +0100
Subject: [PATCH 4/4] Trigger build