From f61c2b2b0a758a079c38575e044e4fc5915a4e51 Mon Sep 17 00:00:00 2001
From: CrystalVibe28 <82704512+CrystalVibe28@users.noreply.github.com>
Date: Sun, 17 May 2026 01:01:08 +0800
Subject: [PATCH] fix(chatgpt): avoid recursive image usage parsing

---
 .../image_generation/response_parsing.py      | 23 +++++++--
 .../chatgpt/image_generation/test_usage.py    | 48 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/chatgpt/image_generation/response_parsing.py b/litellm/llms/chatgpt/image_generation/response_parsing.py
index 2b3b8545351..df848d696ca 100644
--- a/litellm/llms/chatgpt/image_generation/response_parsing.py
+++ b/litellm/llms/chatgpt/image_generation/response_parsing.py
@@ -86,12 +86,29 @@ def get_image_generation_usage(response_payload: Any) -> Optional[dict]:
     if not isinstance(response_payload, dict):
         return None
 
-    response = response_payload.get("response")
-    if isinstance(response, dict):
-        image_gen_usage = get_image_generation_usage(response)
+    payloads_to_check: List[dict] = []
+    current_payload = response_payload
+    visited_container_ids = set()
+    while isinstance(current_payload, dict):
+        container_id = id(current_payload)
+        if container_id in visited_container_ids:
+            break
+        visited_container_ids.add(container_id)
+        payloads_to_check.append(current_payload)
+
+        next_payload = current_payload.get("response")
+        if not isinstance(next_payload, dict):
+            break
+        current_payload = next_payload
+
+    for payload in reversed(payloads_to_check):
+        image_gen_usage = _get_image_generation_usage_from_payload(payload)
         if image_gen_usage is not None:
             return image_gen_usage
+    return None
 
+
+def _get_image_generation_usage_from_payload(response_payload: dict) -> Optional[dict]:
     tool_usage = response_payload.get("tool_usage")
     if not isinstance(tool_usage, dict):
         return None
diff --git a/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py b/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py
index 703b10b3dd7..cfe64b2f069 100644
--- a/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py
+++ b/tests/test_litellm/llms/chatgpt/image_generation/test_usage.py
@@ -28,6 +28,54 @@ def test_chatgpt_image_generation_usage_helpers_ignore_invalid_payloads():
     )
 
 
+def test_chatgpt_image_generation_usage_helper_prefers_nested_response_usage():
+    config = ChatGPTImageGenerationConfig()
+
+    usage = config._get_image_generation_usage(
+        {
+            "tool_usage": {
+                "image_gen": {
+                    "input_tokens": 1,
+                    "output_tokens": 2,
+                }
+            },
+            "response": {
+                "tool_usage": {
+                    "image_gen": {
+                        "input_tokens": 10,
+                        "output_tokens": 20,
+                    }
+                }
+            },
+        }
+    )
+
+    assert usage is not None
+    assert usage["input_tokens"] == 10
+    assert usage["output_tokens"] == 20
+    assert usage["total_tokens"] == 30
+
+
+def test_chatgpt_image_generation_usage_helper_handles_cyclic_response_payload():
+    config = ChatGPTImageGenerationConfig()
+    payload = {
+        "tool_usage": {
+            "image_gen": {
+                "input_tokens": 3,
+                "output_tokens": 4,
+            }
+        }
+    }
+    payload["response"] = payload
+
+    usage = config._get_image_generation_usage(payload)
+
+    assert usage is not None
+    assert usage["input_tokens"] == 3
+    assert usage["output_tokens"] == 4
+    assert usage["total_tokens"] == 7
+
+
 def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response():
     config = ChatGPTImageGenerationConfig()
     raw_response = httpx.Response(