From dbc080c23b8de5245a4bbcf42711a8143b84e602 Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <gmassadas@cloudflare.com>
Date: Wed, 24 Apr 2024 18:07:25 +0100
Subject: [PATCH] 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 000000000000..78b158df840f
--- /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 28f75b7de8e8..9c36a2a938b3 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 03d5580a3994..c4c17c87216e 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 73a852f1e2da..e05a5020c48b 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;
 }
 
 /**