From bb8b234d85d2568d4458dc0f9a3b3d0adcf846af Mon Sep 17 00:00:00 2001
From: Willem Jiang <willem.jiang@gmail.com>
Date: Sat, 2 May 2026 15:04:11 +0800
Subject: [PATCH] chroe(2585): keep polishing the code of codex token usage
 (#2689)

---
 .../deerflow/models/openai_codex_provider.py  | 31 +++++++++++++++++--
 backend/tests/test_codex_provider.py          | 14 ++++-----
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/backend/packages/harness/deerflow/models/openai_codex_provider.py b/backend/packages/harness/deerflow/models/openai_codex_provider.py
index d8e46c2ae..95cf12b09 100644
--- a/backend/packages/harness/deerflow/models/openai_codex_provider.py
+++ b/backend/packages/harness/deerflow/models/openai_codex_provider.py
@@ -21,13 +21,40 @@ from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage
 from langchain_core.outputs import ChatGeneration, ChatResult
-from langchain_openai.chat_models.base import _create_usage_metadata_responses
 
 from deerflow.models.credential_loader import CodexCliCredential, load_codex_cli_credential
 
 logger = logging.getLogger(__name__)
 
 CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
+
+
+def _build_usage_metadata(oai_usage: dict) -> dict:
+    """Convert Codex/Responses API usage dict to LangChain usage_metadata format.
+
+    Maps OpenAI Responses API token usage fields to the dict structure that
+    LangChain AIMessage.usage_metadata expects. This avoids depending on
+    langchain_openai private helpers like ``_create_usage_metadata_responses``.
+    """
+    input_tokens = oai_usage.get("input_tokens", 0)
+    output_tokens = oai_usage.get("output_tokens", 0)
+    total_tokens = oai_usage.get("total_tokens", input_tokens + output_tokens)
+    metadata: dict = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+    }
+    input_details = oai_usage.get("input_tokens_details") or {}
+    output_details = oai_usage.get("output_tokens_details") or {}
+    cache_read = input_details.get("cached_tokens")
+    if cache_read is not None:
+        metadata["input_token_details"] = {"cache_read": cache_read}
+    reasoning = output_details.get("reasoning_tokens")
+    if reasoning is not None:
+        metadata["output_token_details"] = {"reasoning": reasoning}
+    return metadata
+
+
 MAX_RETRIES = 3
 
 
@@ -347,7 +374,7 @@ class CodexChatModel(BaseChatModel):
                 )
 
         usage = response.get("usage", {})
-        usage_metadata = _create_usage_metadata_responses(usage) if usage else None
+        usage_metadata = _build_usage_metadata(usage) if usage else None
         additional_kwargs = {}
         if reasoning_content:
             additional_kwargs["reasoning_content"] = reasoning_content
diff --git a/backend/tests/test_codex_provider.py b/backend/tests/test_codex_provider.py
index 512154564..1b9136b85 100644
--- a/backend/tests/test_codex_provider.py
+++ b/backend/tests/test_codex_provider.py
@@ -103,13 +103,13 @@ def test_parse_response_populates_usage_metadata():
 
     result = model._parse_response(response)
 
-    assert result.generations[0].message.usage_metadata == {
-        "input_tokens": 10,
-        "output_tokens": 5,
-        "total_tokens": 15,
-        "input_token_details": {"cache_read": 3},
-        "output_token_details": {"reasoning": 2},
-    }
+    meta = result.generations[0].message.usage_metadata
+    assert meta is not None
+    assert meta["input_tokens"] == 10
+    assert meta["output_tokens"] == 5
+    assert meta["total_tokens"] == 15
+    assert meta["input_token_details"]["cache_read"] == 3
+    assert meta["output_token_details"]["reasoning"] == 2
 
 
 def test_parse_response_reasoning_content():