🪲 BUG-#68: Request stream_options.include_usage in SSE streaming

pycodeloop/providers/generic.py

@@ -121,6 +121,7 @@ class _ConnectionSnapshot:
     request_builder: RequestBuilder
     response_parser: ResponseParser
     supports_openai_sse: bool
+    include_usage_in_stream: bool
 
 
 class GenericProvider(Provider):
@@ -196,6 +197,7 @@ def __init__(
         repetition_repeats: int = _REPETITION_REPEATS,
         context_window: int | None = None,
         supports_openai_sse: bool = True,
+        include_usage_in_stream: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(model=model, api_key=api_key, **kwargs)
@@ -211,6 +213,7 @@ def __init__(
         self.repetition_repeats = repetition_repeats
         self.context_window = context_window
         self._supports_openai_sse = supports_openai_sse
+        self._include_usage_in_stream = include_usage_in_stream
         self._config_path: Path | None = None
         self._lock = threading.Lock()
 
@@ -258,6 +261,7 @@ def _build_from_json(cls, path: str | Path) -> GenericProvider:
             timeout=data.get("timeout", 60.0),
             context_window=data.get("context_window"),
             supports_openai_sse=response_shape != "anthropic",
+            include_usage_in_stream=data.get("include_usage_in_stream", True),
         )
 
     def reload(self) -> None:
@@ -281,6 +285,7 @@ def reload(self) -> None:
             self.timeout = fresh.timeout
             self.context_window = fresh.context_window
             self._supports_openai_sse = fresh._supports_openai_sse
+            self._include_usage_in_stream = fresh._include_usage_in_stream
 
     @staticmethod
     def _default_request(
@@ -308,6 +313,7 @@ def _snapshot_locked(self) -> _ConnectionSnapshot:
             request_builder=self.request_builder,
             response_parser=self.response_parser,
             supports_openai_sse=self._supports_openai_sse,
+            include_usage_in_stream=self._include_usage_in_stream,
         )
 
     def _headers(self, config: _ConnectionSnapshot) -> dict[str, str]:
@@ -381,6 +387,12 @@ def _stream(
         config: _ConnectionSnapshot,
     ) -> ProviderResponse:
         body = {**body, "stream": True}
+        if config.include_usage_in_stream:
+            existing_stream_options = body.get("stream_options") or {}
+            body["stream_options"] = {
+                "include_usage": True,
+                **existing_stream_options,
+            }
         text = ""
         pending: dict[int, dict] = {}
         stop_reason: str | None = None

tests/providers/test_generic.py

@@ -338,6 +338,123 @@ def test_streaming_flags_a_connection_dropped_mid_response(self):
         self.assertEqual(result.text, "cut off mid")
         self.assertEqual(result.stop_reason, "connection_lost")
 
+    def test_streaming_requests_usage_and_captures_it_from_final_chunk(self):
+        """Regression: streaming previously sent `stream: True` with no
+        `stream_options.include_usage`, so OpenAI-compatible servers that
+        only report usage when asked (e.g. Ollama's /v1/chat/completions)
+        never sent a usage chunk and every streamed response reported
+        0/0 tokens."""
+        path = self._write_config(
+            {"url": "http://fake/v1/chat/completions", "model": "my-model"}
+        )
+        provider = GenericProvider.from_json(path)
+
+        chunks = [
+            {"choices": [{"delta": {"content": "hi"}}]},
+            {"choices": [{"delta": {}, "finish_reason": "stop"}]},
+            {
+                "choices": [],
+                "usage": {"prompt_tokens": 12, "completion_tokens": 3},
+            },
+        ]
+        sse_body = (
+            "".join(f"data: {json.dumps(c)}\n" for c in chunks)
+            + "data: [DONE]\n"
+        ).encode()
+
+        captured_requests = []
+
+        def fake_urlopen(request, timeout=None):
+            captured_requests.append(json.loads(request.data))
+            return _FakeResponse(sse_body)
+
+        with mock.patch(
+            "pycodeloop.providers.generic.urllib.request.urlopen",
+            side_effect=fake_urlopen,
+        ):
+            result = provider.complete("sys", [], [], on_delta=lambda _: None)
+
+        self.assertEqual(
+            captured_requests[0]["stream_options"], {"include_usage": True}
+        )
+        self.assertEqual(result.usage.input_tokens, 12)
+        self.assertEqual(result.usage.output_tokens, 3)
+
+    def test_streaming_merges_include_usage_into_callers_stream_options(
+        self,
+    ):
+        """A caller opting out via params.stream_options.include_usage
+        (e.g. a provider that rejects the field) must not be silently
+        overwritten, and sibling flags must survive the merge."""
+        path = self._write_config(
+            {
+                "url": "http://fake/v1/chat/completions",
+                "model": "my-model",
+                "request": {
+                    "params": {
+                        "stream_options": {
+                            "include_usage": False,
+                            "include_intermediary_tokens": True,
+                        }
+                    }
+                },
+            }
+        )
+        provider = GenericProvider.from_json(path)
+
+        sse_body = (
+            b'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}\n'
+            b"data: [DONE]\n"
+        )
+
+        captured_requests = []
+
+        def fake_urlopen(request, timeout=None):
+            captured_requests.append(json.loads(request.data))
+            return _FakeResponse(sse_body)
+
+        with mock.patch(
+            "pycodeloop.providers.generic.urllib.request.urlopen",
+            side_effect=fake_urlopen,
+        ):
+            provider.complete("sys", [], [], on_delta=lambda _: None)
+
+        self.assertEqual(
+            captured_requests[0]["stream_options"],
+            {"include_usage": False, "include_intermediary_tokens": True},
+        )
+
+    def test_include_usage_in_stream_false_omits_stream_options(self):
+        """Strict OpenAI-compatible endpoints that 400 on unknown fields
+        can opt out entirely via config."""
+        path = self._write_config(
+            {
+                "url": "http://fake/v1/chat/completions",
+                "model": "my-model",
+                "include_usage_in_stream": False,
+            }
+        )
+        provider = GenericProvider.from_json(path)
+
+        sse_body = (
+            b'data: {"choices": [{"delta": {}, "finish_reason": "stop"}]}\n'
+            b"data: [DONE]\n"
+        )
+
+        captured_requests = []
+
+        def fake_urlopen(request, timeout=None):
+            captured_requests.append(json.loads(request.data))
+            return _FakeResponse(sse_body)
+
+        with mock.patch(
+            "pycodeloop.providers.generic.urllib.request.urlopen",
+            side_effect=fake_urlopen,
+        ):
+            provider.complete("sys", [], [], on_delta=lambda _: None)
+
+        self.assertNotIn("stream_options", captured_requests[0])
+
     def test_streaming_cuts_a_looping_response_short(self):
         path = self._write_config(
             {"url": "http://fake/v1/chat/completions", "model": "my-model"}