From 4097759b6fbef4cd9aa81d3a6f01fc868ff50dd8 Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <5445926+G4brym@users.noreply.github.com>
Date: Mon, 29 Apr 2024 17:23:35 +0100
Subject: [PATCH] Add workers ai staging flag (#5705)

* Add workers ai staging flag

* add `staging` option to ai binding config

---------

Co-authored-by: Rahul Sethi <5822355+RamIdeas@users.noreply.github.com>
---
 .changeset/thin-wasps-bathe.md                      |  5 +++++
 packages/wrangler/src/config/environment.ts         |  1 +
 packages/wrangler/src/config/index.ts               | 13 ++++++++++---
 .../deployment-bundle/create-worker-upload-form.ts  |  3 ++-
 packages/wrangler/src/deployment-bundle/worker.ts   |  1 +
 5 files changed, 19 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/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;
 
diff --git a/packages/wrangler/src/config/index.ts b/packages/wrangler/src/config/index.ts
index e0660107d8b..32dd152820f 100644
--- a/packages/wrangler/src/config/index.ts
+++ b/packages/wrangler/src/config/index.ts
@@ -183,8 +183,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,
@@ -419,9 +421,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;
 }
 
 /**