fix: stream Work model tool responses

agent_platform/runtime/provider.py

@@ -30,6 +30,96 @@ def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float:
     return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1)))
 
 
+async def decode_complete_response(response: httpx.Response) -> dict[str, Any]:
+    content_type = response.headers.get("Content-Type", "").lower()
+    if "text/event-stream" not in content_type:
+        await response.aread()
+        data = response.json()
+        if not data.get("choices"):
+            raise RuntimeError("Model provider returned no choices")
+        return data
+
+    envelope: dict[str, Any] = {"object": "chat.completion"}
+    message: dict[str, Any] = {"role": "assistant", "content": None}
+    content_parts: list[str] = []
+    reasoning_parts: list[str] = []
+    tool_calls: dict[int, dict[str, Any]] = {}
+    finish_reason: str | None = None
+    usage: dict[str, Any] | None = None
+    saw_choice = False
+
+    async for line in response.aiter_lines():
+        if not line.startswith("data:"):
+            continue
+        raw = line[5:].strip()
+        if not raw or raw == "[DONE]":
+            continue
+        chunk = httpx.Response(200, content=raw).json()
+        if chunk.get("error"):
+            raise RuntimeError(f"Model provider stream error: {chunk['error']}")
+        for key in ("id", "created", "model", "system_fingerprint"):
+            if chunk.get(key) is not None:
+                envelope[key] = chunk[key]
+        if chunk.get("usage"):
+            usage = chunk["usage"]
+        choices = chunk.get("choices") or []
+        if not choices:
+            continue
+        saw_choice = True
+        choice = choices[0]
+        delta = choice.get("delta") or {}
+        if delta.get("role"):
+            message["role"] = delta["role"]
+        if isinstance(delta.get("content"), str):
+            content_parts.append(delta["content"])
+        reasoning = delta.get("reasoning_content")
+        if not isinstance(reasoning, str):
+            reasoning = delta.get("reasoning")
+        if isinstance(reasoning, str):
+            reasoning_parts.append(reasoning)
+        for position, fragment in enumerate(delta.get("tool_calls") or []):
+            index = int(fragment.get("index", position))
+            merged = tool_calls.setdefault(
+                index,
+                {
+                    "id": "",
+                    "type": "function",
+                    "function": {"name": "", "arguments": ""},
+                },
+            )
+            if fragment.get("id"):
+                merged["id"] = fragment["id"]
+            if fragment.get("type"):
+                merged["type"] = fragment["type"]
+            function = fragment.get("function") or {}
+            if function.get("name"):
+                merged["function"]["name"] += function["name"]
+            if function.get("arguments"):
+                merged["function"]["arguments"] += function["arguments"]
+        if choice.get("finish_reason") is not None:
+            finish_reason = choice["finish_reason"]
+
+    if not saw_choice:
+        raise RuntimeError("Model provider stream returned no choices")
+    if content_parts:
+        message["content"] = "".join(content_parts)
+    if reasoning_parts:
+        message["reasoning_content"] = "".join(reasoning_parts)
+    if tool_calls:
+        message["tool_calls"] = [tool_calls[index] for index in sorted(tool_calls)]
+
+    envelope["choices"] = [
+        {
+            "index": 0,
+            "message": message,
+            "finish_reason": finish_reason,
+        }
+    ]
+    if usage is not None:
+        envelope["usage"] = usage
+    return envelope
+
+
 class ModelProvider:
     def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
         self.settings = settings
@@ -61,7 +151,7 @@ class ModelProvider:
         payload: dict[str, Any] = {
             "model": model,
             "messages": messages,
-            "stream": False,
+            "stream": True,
         }
         if tools:
             payload["tools"] = tools
@@ -69,29 +159,27 @@ class ModelProvider:
         if temperature is not None:
             payload["temperature"] = temperature
         for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1):
+            response: httpx.Response | None = None
             try:
-                response = await self.client.post(
+                async with self.client.stream(
+                    "POST",
                     completions_url(self.settings.model_api_base_url),
                     headers=self.headers,
                     json=payload,
-                )
+                ) as response:
+                    if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
+                        await response.aread()
+                        await asyncio.sleep(retry_delay_seconds(response, attempt))
+                        continue
+
+                    response.raise_for_status()
+                    return await decode_complete_response(response)
             except httpx.TransportError:
                 if attempt >= COMPLETE_MAX_ATTEMPTS:
                     raise
-                await asyncio.sleep(retry_delay_seconds(None, attempt))
-                continue
-
-            if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
-                await response.aread()
                 await asyncio.sleep(retry_delay_seconds(response, attempt))
                 continue
 
-            response.raise_for_status()
-            data = response.json()
-            if not data.get("choices"):
-                raise RuntimeError("Model provider returned no choices")
-            return data
-
         raise RuntimeError("Model provider retry loop exited unexpectedly")
 
     async def forward(self, payload: dict[str, Any]) -> httpx.Response:

tests/test_provider.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from collections.abc import Awaitable, Callable
 
 import httpx
@@ -78,3 +79,80 @@ async def test_complete_does_not_retry_client_error(settings: Settings) -> None:
         await _complete_with_handler(settings, handler)
 
     assert calls == 1
+
+
+async def test_complete_aggregates_streamed_tool_call(settings: Settings) -> None:
+    chunks = [
+        {
+            "id": "completion-1",
+            "object": "chat.completion.chunk",
+            "created": 123,
+            "model": "test-model",
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "id": "call-1",
+                                "type": "function",
+                                "function": {"name": "ping", "arguments": '{"val'},
+                            }
+                        ],
+                    },
+                    "finish_reason": None,
+                }
+            ],
+        },
+        {
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "tool_calls": [
+                            {
+                                "index": 0,
+                                "function": {"arguments": 'ue":"ok"}'},
+                            }
+                        ]
+                    },
+                    "finish_reason": None,
+                }
+            ]
+        },
+        {
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "tool_calls",
+                }
+            ],
+            "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7},
+        },
+    ]
+    body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        assert json.loads(request.content)["stream"] is True
+        return httpx.Response(
+            200,
+            request=request,
+            headers={"Content-Type": "text/event-stream"},
+            content=body,
+        )
+
+    result = await _complete_with_handler(settings, handler)
+
+    choice = result["choices"][0]
+    assert choice["finish_reason"] == "tool_calls"
+    assert choice["message"]["tool_calls"] == [
+        {
+            "id": "call-1",
+            "type": "function",
+            "function": {"name": "ping", "arguments": '{"value":"ok"}'},
+        }
+    ]
+    assert result["usage"]["total_tokens"] == 7