From 289ca28ce281eb6861eb729ab69b63a66b904adc Mon Sep 17 00:00:00 2001
From: Harsh Kashyap <harsh.kashyap2001@gmail.com>
Date: Fri, 10 Jul 2026 08:29:32 +0530
Subject: [PATCH] Fix OpenAI agent stream chunk shape (#16402)

### What problem does this PR solve?

Closes #8175.

The Agent OpenAI-compatible streaming path uses `get_data_openai(...,
stream=True)`, but that helper currently returns a minimal chunk shape.
The main OpenAI-compatible chat endpoint already includes chunk metadata
such as `created`, `system_fingerprint`, `usage`, `logprobs`, and
assistant role/tool placeholders.

This PR aligns the Agent stream helper with that existing
OpenAI-compatible chunk shape while keeping the current `delta.content`
behavior and existing reference injection path intact.

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [ ] New Feature (non-breaking change which adds functionality)
- [ ] Documentation Update
- [ ] Refactoring
- [ ] Performance Improvement
- [ ] Other (please describe):

### Verification

- `./.venv/bin/python -m pytest
test/unit_test/api/utils/test_api_utils.py -q`
- `python3 -m py_compile api/utils/api_utils.py
test/unit_test/api/utils/test_api_utils.py`
- `uvx ruff check api/utils/api_utils.py
test/unit_test/api/utils/test_api_utils.py`

---------

Co-authored-by: Harsh Kashyap <harshkashyap@Harshs-MacBook-Pro.local>
---
 api/utils/api_utils.py                     | 23 ++++++++-
 test/unit_test/api/utils/test_api_utils.py | 57 ++++++++++++++++------
 2 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py
index acb3849726..0cde1eac2f 100644
--- a/api/utils/api_utils.py
+++ b/api/utils/api_utils.py
@@ -428,15 +428,36 @@ def get_data_openai(id=None, created=None, model=None, prompt_tokens=0, completi
     total_tokens = prompt_tokens + completion_tokens
 
     if stream:
+        usage = None
+        if finish_reason is not None:
+            usage = {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "completion_tokens_details": {
+                    "reasoning_tokens": 0,
+                    "accepted_prediction_tokens": 0,
+                    "rejected_prediction_tokens": 0,
+                },
+            }
         return {
             "id": f"{id}",
             "object": "chat.completion.chunk",
+            "created": created if created is not None else int(time.time()),
             "model": model,
+            "system_fingerprint": "",
+            "usage": usage,
             "choices": [
                 {
-                    "delta": {"content": content},
+                    "delta": {
+                        "content": content,
+                        "role": "assistant",
+                        "function_call": None,
+                        "tool_calls": None,
+                    },
                     "finish_reason": finish_reason,
                     "index": 0,
+                    "logprobs": None,
                 }
             ],
         }
diff --git a/test/unit_test/api/utils/test_api_utils.py b/test/unit_test/api/utils/test_api_utils.py
index 08d4d5ade1..ff8e895f64 100644
--- a/test/unit_test/api/utils/test_api_utils.py
+++ b/test/unit_test/api/utils/test_api_utils.py
@@ -17,32 +17,59 @@
 from api.utils import api_utils
 
 
-def test_get_data_openai_defaults_created_to_current_timestamp(monkeypatch):
+def test_get_data_openai_stream_chunk_matches_openai_shape(monkeypatch):
     monkeypatch.setattr(api_utils.time, "time", lambda: 1234567890.9)
 
-    data = api_utils.get_data_openai(content="answer")
-
-    assert data["created"] == 1234567890
-
-
-def test_get_data_openai_preserves_explicit_created_value():
-    data = api_utils.get_data_openai(created=0, content="answer")
-
-    assert data["created"] == 0
-
-
-def test_get_data_openai_stream_response_shape_is_unchanged():
-    data = api_utils.get_data_openai(id="chatcmpl-test", model="test-model", content="chunk", finish_reason=None, stream=True)
+    data = api_utils.get_data_openai(id="chatcmpl-test", model="test-model", content="chunk", stream=True)
 
     assert data == {
         "id": "chatcmpl-test",
         "object": "chat.completion.chunk",
+        "created": 1234567890,
         "model": "test-model",
+        "system_fingerprint": "",
+        "usage": None,
         "choices": [
             {
-                "delta": {"content": "chunk"},
+                "delta": {
+                    "content": "chunk",
+                    "role": "assistant",
+                    "function_call": None,
+                    "tool_calls": None,
+                },
                 "finish_reason": None,
                 "index": 0,
+                "logprobs": None,
             }
         ],
     }
+
+
+def test_get_data_openai_stream_preserves_explicit_created_value():
+    data = api_utils.get_data_openai(created=0, content="chunk", stream=True)
+
+    assert data["created"] == 0
+
+
+def test_get_data_openai_terminal_stream_chunk_includes_usage():
+    data = api_utils.get_data_openai(prompt_tokens=3, completion_tokens=5, content=None, finish_reason="stop", stream=True)
+
+    assert data["usage"] == {
+        "prompt_tokens": 3,
+        "completion_tokens": 5,
+        "total_tokens": 8,
+        "completion_tokens_details": {
+            "reasoning_tokens": 0,
+            "accepted_prediction_tokens": 0,
+            "rejected_prediction_tokens": 0,
+        },
+    }
+    assert data["choices"][0]["finish_reason"] == "stop"
+
+
+def test_get_data_openai_stream_delta_allows_reference_payload():
+    data = api_utils.get_data_openai(content="chunk", stream=True)
+
+    data["choices"][0]["delta"]["reference"] = {"chunks": []}
+
+    assert data["choices"][0]["delta"]["reference"] == {"chunks": []}