fix(oci): terminate stream on finishReason — OCI does not send [DONE] (#757)

* fix(oci): terminate stream on finishReason instead of waiting for [DONE]

OCI Generative AI does not send a `data: [DONE]` SSE marker to signal
end-of-stream. It sends a final event with `finishReason` and keeps
the connection open, causing chat_stream() to hang indefinitely.

Emit closing events (message-end / stream-end) and return from the
generator when `finishReason` is detected. The [DONE] path is kept
as a fallback for forward compatibility.

Fixes #756

* test(oci): assert stream termination and full event lifecycle

Strengthen V1 and V2 streaming integration tests to verify streams
terminate correctly and produce the expected event sequence:

- V1: stream-start → text-generation(s) → stream-end
- V2: message-start → content-start → content-delta(s) → content-end → message-end

Without these assertions the previous tests would have hung forever
on the streaming bug rather than failing with a clear error.

Author: fede-kamel

src/cohere/oci_client.py

@@ -1055,36 +1055,45 @@ def _transform_v1_event(oci_event: typing.Dict[str, typing.Any]) -> typing.Itera
                 final_v1_finish_reason = oci_event.get("finishReason", final_v1_finish_reason)
             yield _emit_v1_event(event)
 
+    stream_finished = False
+
+    def _emit_closing_events() -> typing.Iterator[bytes]:
+        """Emit the final closing events for the stream."""
+        if is_v2:
+            if emitted_start:
+                if not emitted_content_end:
+                    yield _emit_v2_event({"type": "content-end", "index": current_content_index})
+                message_end_event: typing.Dict[str, typing.Any] = {
+                    "type": "message-end",
+                    "id": generation_id,
+                    "delta": {"finish_reason": final_finish_reason},
+                }
+                if final_usage:
+                    message_end_event["delta"]["usage"] = final_usage
+                yield _emit_v2_event(message_end_event)
+        else:
+            yield _emit_v1_event(
+                {
+                    "event_type": "stream-end",
+                    "finish_reason": final_v1_finish_reason,
+                    "response": {
+                        "text": full_v1_text,
+                        "generation_id": generation_id,
+                        "finish_reason": final_v1_finish_reason,
+                    },
+                }
+            )
+
     def _process_line(line: str) -> typing.Iterator[bytes]:
+        nonlocal stream_finished
         if not line.startswith("data: "):
             return
 
         data_str = line[6:]
         if data_str.strip() == "[DONE]":
-            if is_v2:
-                if emitted_start:
-                    if not emitted_content_end:
-                        yield _emit_v2_event({"type": "content-end", "index": current_content_index})
-                    message_end_event: typing.Dict[str, typing.Any] = {
-                        "type": "message-end",
-                        "id": generation_id,
-                        "delta": {"finish_reason": final_finish_reason},
-                    }
-                    if final_usage:
-                        message_end_event["delta"]["usage"] = final_usage
-                    yield _emit_v2_event(message_end_event)
-            else:
-                yield _emit_v1_event(
-                    {
-                        "event_type": "stream-end",
-                        "finish_reason": final_v1_finish_reason,
-                        "response": {
-                            "text": full_v1_text,
-                            "generation_id": generation_id,
-                            "finish_reason": final_v1_finish_reason,
-                        },
-                    }
-                )
+            for event_bytes in _emit_closing_events():
+                yield event_bytes
+            stream_finished = True
             return
 
         try:
@@ -1102,15 +1111,23 @@ def _process_line(line: str) -> typing.Iterator[bytes]:
         except Exception as exc:
             raise RuntimeError(f"OCI stream event transformation failed for endpoint '{endpoint}': {exc}") from exc
 
+        # OCI may not send [DONE] — treat finishReason as stream termination
+        if "finishReason" in oci_event:
+            for event_bytes in _emit_closing_events():
+                yield event_bytes
+            stream_finished = True
+
     for chunk in stream:
         buffer += chunk
         while b"\n" in buffer:
             line_bytes, buffer = buffer.split(b"\n", 1)
             line = line_bytes.decode("utf-8").strip()
             for event_bytes in _process_line(line):
                 yield event_bytes
+            if stream_finished:
+                return
 
-    if buffer.strip():
+    if buffer.strip() and not stream_finished:
         line = buffer.decode("utf-8").strip()
         for event_bytes in _process_line(line):
             yield event_bytes

tests/test_oci_client.py

@@ -113,7 +113,7 @@ def test_chat(self):
         self.assertIn("4", response.text)
 
     def test_chat_stream(self):
-        """Test V1 streaming chat."""
+        """Test V1 streaming chat terminates and produces correct events."""
         events = []
         for event in self.client.chat_stream(
             model="command-r-08-2024",
@@ -125,6 +125,11 @@ def test_chat_stream(self):
         text_events = [e for e in events if hasattr(e, "text") and e.text]
         self.assertTrue(len(text_events) > 0)
 
+        # Verify stream terminates with correct event lifecycle
+        event_types = [getattr(e, "event_type", None) for e in events]
+        self.assertEqual(event_types[0], "stream-start")
+        self.assertEqual(event_types[-1], "stream-end")
+
 
 @unittest.skipIf(os.getenv("TEST_OCI") is None, "TEST_OCI not set")
 class TestOciClientV2(unittest.TestCase):
@@ -186,7 +191,7 @@ def test_chat_v2(self):
         self.assertIsNotNone(response.message)
 
     def test_chat_stream_v2(self):
-        """Test streaming chat with v2 client."""
+        """Test V2 streaming chat terminates and produces correct event lifecycle."""
         events = []
         for event in self.client.chat_stream(
             model="command-a-03-2025",
@@ -195,11 +200,16 @@ def test_chat_stream_v2(self):
             events.append(event)
 
         self.assertTrue(len(events) > 0)
-        # Verify we received content-delta events with text
-        content_delta_events = [e for e in events if hasattr(e, "type") and e.type == "content-delta"]
-        self.assertTrue(len(content_delta_events) > 0)
 
-        # Verify we can extract text from events
+        # Verify full event lifecycle: message-start → content-start → content-delta(s) → content-end → message-end
+        event_types = [e.type for e in events]
+        self.assertEqual(event_types[0], "message-start")
+        self.assertIn("content-start", event_types)
+        self.assertIn("content-delta", event_types)
+        self.assertIn("content-end", event_types)
+        self.assertEqual(event_types[-1], "message-end")
+
+        # Verify we can extract text from content-delta events
         full_text = ""
         for event in events:
             if (
@@ -214,7 +224,6 @@ def test_chat_stream_v2(self):
             ):
                 full_text += event.delta.message.content.text
 
-        # Should have received some text
         self.assertTrue(len(full_text) > 0)
 
 @unittest.skipIf(os.getenv("TEST_OCI") is None, "TEST_OCI not set")