From 910816d38612b893d8bf052fcff48965f50d8131 Mon Sep 17 00:00:00 2001
From: Joaquin Gimenez <joaquingimenez@cloudflare.com>
Date: Wed, 18 Dec 2024 15:30:37 -0300
Subject: [PATCH] Add http and internal codes to AI binding

---
 src/cloudflare/internal/ai-api.ts              | 4 ++++
 src/cloudflare/internal/test/ai/ai-api-test.js | 5 +++++
 2 files changed, 9 insertions(+)

diff --git a/src/cloudflare/internal/ai-api.ts b/src/cloudflare/internal/ai-api.ts
index 4eb2cdd5c17..7077e6765e8 100644
--- a/src/cloudflare/internal/ai-api.ts
+++ b/src/cloudflare/internal/ai-api.ts
@@ -49,6 +49,8 @@ export class Ai {
   private options: AiOptions = {};
   public lastRequestId: string | null = null;
   public aiGatewayLogId: string | null = null;
+  public lastRequestHttpStatusCode: number | null = null;
+  public lastRequestInternalStatusCode: number | null = null;
 
   public constructor(fetcher: Fetcher) {
     this.fetcher = fetcher;
@@ -101,6 +103,7 @@ export class Ai {
 
     this.lastRequestId = res.headers.get('cf-ai-req-id');
     this.aiGatewayLogId = res.headers.get('cf-aig-log-id');
+    this.lastRequestHttpStatusCode = res.status;
 
     if (inputs['stream']) {
       if (!res.ok) {
@@ -135,6 +138,7 @@ export class Ai {
 
     try {
       const parsedContent = JSON.parse(content) as AiError;
+      this.lastRequestInternalStatusCode = parsedContent.internalCode;
       return new InferenceUpstreamError(
         `${parsedContent.internalCode}: ${parsedContent.description}`,
         parsedContent.name
diff --git a/src/cloudflare/internal/test/ai/ai-api-test.js b/src/cloudflare/internal/test/ai/ai-api-test.js
index c362b46a707..49956a4521e 100644
--- a/src/cloudflare/internal/test/ai/ai-api-test.js
+++ b/src/cloudflare/internal/test/ai/ai-api-test.js
@@ -16,6 +16,8 @@ export const tests = {
         env.ai.lastRequestId,
         '3a1983d7-1ddd-453a-ab75-c4358c91b582'
       );
+      // Test request http status code is present
+      assert.deepStrictEqual(env.ai.lastRequestHttpStatusCode, 200);
     }
 
     {
@@ -55,6 +57,9 @@ export const tests = {
             message: '1001: prompt and messages are mutually exclusive',
           }
         );
+        // Test request internal status code is present
+        assert.deepEqual;
+        assert.deepStrictEqual(env.ai.lastRequestInternalStatusCode, 1001);
       }
     }