From 892efd2077be40a903d788189aac50a5684a08c1 Mon Sep 17 00:00:00 2001
From: Nikhil Sinha <nsinha@cloudflare.com>
Date: Fri, 24 Jul 2026 15:21:22 -0700
Subject: [PATCH] Update bindings code to match REST API changes

---
 src/cloudflare/internal/ai-api.ts             | 39 ++++++++++++++++++-
 .../internal/test/ai/ai-api-test.js           | 24 +++++++++++-
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/src/cloudflare/internal/ai-api.ts b/src/cloudflare/internal/ai-api.ts
index 38acedaeae2..d28b2b85172 100644
--- a/src/cloudflare/internal/ai-api.ts
+++ b/src/cloudflare/internal/ai-api.ts
@@ -109,6 +109,27 @@ function isFormData(obj: unknown): obj is FormData {
   return obj instanceof FormData;
 }
 
+/**
+ * Check if a model ID should be handled by AI Gateway
+ * AI Gateway models follow the pattern "provider/model-name" (e.g., "xai/grok-voice")
+ * Workers AI models start with "@cf/" or "@hf/" (e.g., "@cf/meta/llama-3.1-8b-instruct")
+ */
+function isHandledByAiGateway(modelId: string): boolean {
+  // Workers AI models
+  if (modelId.startsWith('@cf/') || modelId.startsWith('@hf/')) {
+    return false;
+  }
+
+  // Skip processing models (e.g., "skipProcessing:...")
+  if (modelId.includes(':')) {
+    return false;
+  }
+
+  // AI Gateway models follow "provider/model-name" pattern
+  const HANDLED_BY_AIG = /^[a-zA-Z0-9._-]+\/[a-zA-Z0-9._-]+$/;
+  return HANDLED_BY_AIG.test(modelId);
+}
+
 /**
  * Find keys in inputs that have a ReadableStream
  * */
@@ -284,8 +305,22 @@ export class Ai {
     }
 
     const aiEndpoint = new URL(`${this.#endpointURL}/run`);
-    aiEndpoint.searchParams.set('version', '3');
-    aiEndpoint.searchParams.set('body', body);
+    
+    // AI Gateway models expect inputs as query params, not version/body
+    if (isHandledByAiGateway(model)) {
+      aiEndpoint.searchParams.set('model', model);
+      
+      // Add all inputs as query parameters
+      for (const [key, value] of Object.entries(inputs)) {
+        if (value !== null && value !== undefined) {
+          aiEndpoint.searchParams.set(key, String(value));
+        }
+      }
+    } else {
+      // Workers AI models use version and body query params
+      aiEndpoint.searchParams.set('version', '3');
+      aiEndpoint.searchParams.set('body', body);
+    }
 
     return await this.#fetcher.fetch(aiEndpoint, fetchOptions);
   }
diff --git a/src/cloudflare/internal/test/ai/ai-api-test.js b/src/cloudflare/internal/test/ai/ai-api-test.js
index e9b31898c43..78e16f2cdc7 100644
--- a/src/cloudflare/internal/test/ai/ai-api-test.js
+++ b/src/cloudflare/internal/test/ai/ai-api-test.js
@@ -345,7 +345,7 @@ export const tests = {
     }
 
     {
-      // Test websocket option with basic inputs
+      // Test websocket option with basic inputs (Workers AI model)
       const resp = await env.ai.run(
         '@cf/test/websocket',
         { encoding: 'utf8' },
@@ -366,6 +366,28 @@ export const tests = {
       });
     }
 
+    {
+      // Test websocket option with AI Gateway model (inputs as query params)
+      const resp = await env.ai.run(
+        'xai/grok-voice',
+        { text: 'hello', voice: 'alloy' },
+        { websocket: true }
+      );
+      assert.deepStrictEqual(resp instanceof Response, true);
+      const respData = await resp.json();
+      assert.deepStrictEqual(respData, {
+        inputs: { text: 'hello', voice: 'alloy' },
+        options: { websocket: true },
+        requestUrl:
+          'https://workers-binding.ai/run?model=xai%2Fgrok-voice&text=hello&voice=alloy',
+        headers: {
+          'cf-consn-sdk-version': '2.0.0',
+          'cf-consn-model-id': 'xai/grok-voice',
+          upgrade: 'websocket',
+        },
+      });
+    }
+
     {
       // Test signal option is not included in request body
       const controller = new AbortController();