coderai broker: actually stream responses (don't wait for done before yielding)

Streaming was scaffolded (stream_queue, _publish_stream_response,
wait_for_stream_event, _iter_broker_stream_chunks) but send_request always awaited
the terminal future and then popped the pending entry in finally — so by the time
_broker_stream began consuming, the stream was over and its queue was gone. Result:
the whole reply arrived at once.

Fixes:
- send_request: for streaming requests, return a {event:"stream_start"} immediately
  and do NOT pop _pending in finally — the consumer drives it. Pop on error.
- _publish_stream_response: queue EVERY event (chunks AND the terminal done) so the
  consumer's loop receives the end marker; still resolves the future for safety.
- finish_stream(): new — removes the streaming request's _pending once drained.
- wait_for_stream_event: simplified to a queue.get (dropped the event_log path).
- providers/coderai._iter_broker_stream_chunks: wrap in try/finally that calls
  finish_stream(request_id) on completion/error; treat stream_start as a no-data
  continue.

Pairs with the coderai-side change that emits chunk/done envelopes.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011DDv7BchtZQWsnPG6Jm49m

aisbf/coderai_broker.py

@@ -596,6 +596,14 @@ class CoderAIBroker:
                     op, provider_id, target_node_id, request_id,
                 )
 
+            if stream_queue is not None:
+                # Streaming: hand control to the consumer immediately instead of
+                # waiting for the terminal event. _iter_broker_stream_chunks() drains
+                # chunks via wait_for_stream_event() and calls finish_stream() to
+                # remove _pending when the stream ends — so we must NOT pop it here.
+                return {"v": 1, "status": "ok", "event": "stream_start",
+                        "request_id": request_id, "payload": {}}
+
             async def wait_reply() -> Dict[str, Any]:
                 if use_local_fast_path:
                     # Poll with short intervals so keepalive messages can extend the deadline.
@@ -666,9 +674,16 @@ class CoderAIBroker:
             if not future.done():
                 future.set_result(result)
             return result
-        finally:
+        except Exception:
             async with self._lock:
                 self._pending.pop(request_id, None)
+            raise
+        finally:
+            # Non-streaming requests are done here; streaming requests keep _pending
+            # alive until the consumer calls finish_stream().
+            if stream_queue is None:
+                async with self._lock:
+                    self._pending.pop(request_id, None)
 
     async def consume_request(self, session_id: str, timeout: int = HEARTBEAT_POLL_SECONDS) -> Optional[Dict[str, Any]]:
         """Block-pop the next request for this session from the shared queue."""
@@ -716,27 +731,29 @@ class CoderAIBroker:
             queue = pending.stream_queue if pending else None
 
         if queue is not None:
-            # Fast path: local asyncio.Queue for this instance's pending request.
-            pending.event_log.append(message)
-            if message.get("event") in {"done", "completed"} and not pending.future.done():
+            # Fast path: local asyncio.Queue. Queue EVERY event (chunks AND the
+            # terminal done/completed) so the consumer's loop receives the end
+            # marker and stops; also resolve the future for any non-stream waiter.
+            if message.get("event") in {"done", "completed"}:
                 await self._record_request_metric(request_id, message)
-                pending.future.set_result(message)
-            else:
-                await queue.put(message)
+                if not pending.future.done():
+                    pending.future.set_result(message)
+            await queue.put(message)
             return
 
         # Slow path: push chunk to the shared cache list for the remote requester.
         self._cache.broker_push(self._reply_key(request_id), message, ttl=REPLY_TTL_SECONDS)
 
+    async def finish_stream(self, request_id: str) -> None:
+        """Remove a streaming request's pending state once the consumer has drained
+        it (called by the relay's stream iterator on completion/error)."""
+        async with self._lock:
+            self._pending.pop(request_id, None)
+
     async def wait_for_stream_event(self, request_id: str, timeout: float = 300.0) -> Dict[str, Any]:
         async with self._lock:
             pending = self._pending.get(request_id)
             queue = pending.stream_queue if pending else None
-            if pending and pending.event_log:
-                for idx, event in enumerate(pending.event_log):
-                    if event.get("event") not in {"done", "completed"}:
-                        pending.event_log.pop(idx)
-                        return event
 
         if queue is not None:
             return await asyncio.wait_for(queue.get(), timeout=timeout)

aisbf/providers/coderai.py

@@ -290,40 +290,47 @@ class CoderAIProviderHandler(BaseProviderHandler):
 
     async def _iter_broker_stream_chunks(self, initial_message: Dict[str, Any], timeout: float) -> AsyncIterator[bytes]:
         message = initial_message
-        while True:
-            status = message.get("status") or "ok"
-            if status == "error":
-                raise Exception(message.get("error") or "CoderAI broker bridge error")
-
-            event = message.get("event")
-            payload_data = message.get("payload") or {}
-
-            # Broker HTTP envelope (ASGI bridge buffered the full response into payload.body).
-            # Yield the body as raw bytes — SSE or JSON — then stop.
-            if "status_code" in payload_data and "body" in payload_data:
-                body = payload_data.get("body")
-                if isinstance(body, str):
-                    yield body.encode("utf-8")
-                elif isinstance(body, bytes):
-                    yield body
-                elif isinstance(body, (dict, list)):
-                    yield json.dumps(body, separators=(",", ":")).encode("utf-8")
-                return
-
-            if event in {"chunk", "progress", "output", "log", "data"}:
-                chunk = payload_data.get("chunk", payload_data)
-                yield self._decode_broker_chunk(chunk)
-            elif isinstance(payload_data.get("chunks"), list):
-                for item in payload_data.get("chunks", []):
-                    yield self._decode_broker_chunk(item)
-
-            if event in {None, "done", "completed"}:
-                return
-
-            next_request_id = message.get("request_id")
-            if not next_request_id:
-                return
-            message = await coderai_broker.wait_for_stream_event(next_request_id, timeout=timeout)
+        request_id = message.get("request_id")
+        try:
+            while True:
+                status = message.get("status") or "ok"
+                if status == "error":
+                    raise Exception(message.get("error") or "CoderAI broker bridge error")
+
+                event = message.get("event")
+                payload_data = message.get("payload") or {}
+
+                # Broker HTTP envelope (ASGI bridge buffered the full response into
+                # payload.body). Yield the body as raw bytes — SSE or JSON — then stop.
+                if "status_code" in payload_data and "body" in payload_data:
+                    body = payload_data.get("body")
+                    if isinstance(body, str):
+                        yield body.encode("utf-8")
+                    elif isinstance(body, bytes):
+                        yield body
+                    elif isinstance(body, (dict, list)):
+                        yield json.dumps(body, separators=(",", ":")).encode("utf-8")
+                    return
+
+                if event in {"chunk", "progress", "output", "log", "data"}:
+                    chunk = payload_data.get("chunk", payload_data)
+                    yield self._decode_broker_chunk(chunk)
+                elif isinstance(payload_data.get("chunks"), list):
+                    for item in payload_data.get("chunks", []):
+                        yield self._decode_broker_chunk(item)
+
+                # stream_start (event="stream_start") carries no data and isn't in
+                # this set, so we keep waiting for the first real chunk; done/None end it.
+                if event in {None, "done", "completed"}:
+                    return
+
+                next_request_id = message.get("request_id") or request_id
+                if not next_request_id:
+                    return
+                message = await coderai_broker.wait_for_stream_event(next_request_id, timeout=timeout)
+        finally:
+            if request_id:
+                await coderai_broker.finish_stream(request_id)
 
     def validate_credentials(self) -> bool:
         if self._broker_mode: