broker: stream chat responses token-by-token instead of buffering

The front's broker path buffered the entire SSE response and sent one envelope, so
broker clients (lisa, and any OpenAI client via AISBF) got the whole reply at once.
Now the front streams it:

- streaming.py: stream_chunk_envelope uses event="chunk" with payload.chunk (the
  shape the AISBF relay consumes), finalize_stream uses event="done".
- dispatcher.py: extracted resolve_broker_request() (op routing + body decode +
  validation) shared by the buffered dispatch and the new streaming path, so they
  can't diverge; added BrokerDispatchError.
- app.py: broker_execute_stream() proxies to the engine with stream=True and yields
  each SSE chunk (sharing the per-model queue + in-flight tracking); start_broker
  wires client.stream_dispatcher.
- client.py: handle_message, for stream=true requests, relays each yielded chunk as
  a `chunk` envelope and ends with a `done` (instead of the single buffered reply).

Requires the matching AISBF change (send_request returns early for streaming and
the relay drains the chunk queue) — committed in the aisbf repo.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011DDv7BchtZQWsnPG6Jm49m

codai/broker/client.py

@@ -18,7 +18,8 @@ from codai.broker.capabilities import (
     build_register_message,
 )
 from codai.broker.dispatcher import OP_ROUTE_MAP
-from codai.broker.models import BrokerRequestEnvelope, success_envelope
+from codai.broker.models import BrokerRequestEnvelope, success_envelope, error_envelope
+from codai.broker.streaming import stream_chunk_envelope, finalize_stream
 
 Dispatcher = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
 
@@ -29,6 +30,9 @@ class BrokerClient:
     def __init__(self, runtime, dispatcher: Dispatcher | None = None):
         self.runtime = runtime
         self.dispatcher = dispatcher
+        # Optional async-generator dispatcher for streaming requests: yields SSE
+        # chunks which we relay as ``chunk`` envelopes + a terminal ``done``.
+        self.stream_dispatcher = None
         self.websocket = None
         self.session_id = None
         self.session_metadata: dict[str, Any] = {}
@@ -394,6 +398,41 @@ class BrokerClient:
                     self._send_keepalives(request_id, interval=30.0, estimated_timeout=300.0)
                 )
 
+            # Streaming path: relay engine SSE chunks as `chunk` envelopes + a
+            # terminal `done`, so the broker side streams tokens to the client
+            # instead of buffering the whole reply.
+            if stream and self.stream_dispatcher is not None and request_id:
+                seq = 0
+                _t_start = time.monotonic()
+                try:
+                    async for chunk in self.stream_dispatcher(message):
+                        if not chunk:
+                            continue
+                        seq += 1
+                        await self.websocket.send(json.dumps(
+                            stream_chunk_envelope(request_id, seq, chunk)))
+                    await self.websocket.send(json.dumps(finalize_stream(
+                        request_id, seq,
+                        round((time.monotonic() - _t_start) * 1000, 3))))
+                    logger.info(
+                        "CoderAI broker streamed request_id=%s chunks=%d", request_id, seq)
+                except Exception as exc:
+                    logger.error("CoderAI broker stream error request_id=%s: %s",
+                                 request_id, exc, exc_info=True)
+                    try:
+                        await self.websocket.send(json.dumps(error_envelope(
+                            request_id, code="stream_error", message=str(exc))))
+                    except Exception:
+                        pass
+                finally:
+                    if keepalive_task is not None:
+                        keepalive_task.cancel()
+                        try:
+                            await keepalive_task
+                        except asyncio.CancelledError:
+                            pass
+                return None
+
             try:
                 response = await self.dispatcher(message)
             except Exception as exc:

codai/broker/dispatcher.py

@@ -60,19 +60,20 @@ def _is_text_response(content_type: str | None) -> bool:
     )
 
 
-async def execute_broker_request(app, envelope, executor=None):
-    """Validate and execute a broker request envelope.
+class BrokerDispatchError(Exception):
+    """Carries a ready error envelope for an unsupported op/path."""
 
-    ``executor`` is an ``async (method, path, headers, query, body) -> {status_code,
-    headers, body}`` callable. When omitted the request is run in-process against
-    ``app`` via the ASGI bridge (engine / single-process mode). The front passes its
-    own executor that proxies to the right engine over HTTP."""
+    def __init__(self, envelope: dict):
+        super().__init__(envelope.get("error") or "broker dispatch error")
+        self.envelope = envelope
 
-    logger.debug(
-        "broker dispatch → op=%s request_id=%s path=%r method=%r stream=%s",
-        envelope.op, envelope.request_id, envelope.path, envelope.method, envelope.stream,
-    )
 
+def resolve_broker_request(envelope):
+    """Resolve an envelope to (headers, body), applying op routing, body decoding
+    and validation (mutating envelope.method/path/query in place). Shared by the
+    buffered dispatch and the streaming path so they never diverge. Raises
+    :class:`BrokerDispatchError` (with a ready error envelope) on unsupported
+    op/path."""
     if envelope.op == "proxy":
         proxy_payload = envelope.payload or {}
         endpoint_path = str(proxy_payload.get("endpoint_path") or envelope.path or "").strip()
@@ -89,36 +90,24 @@ async def execute_broker_request(app, envelope, executor=None):
             envelope.payload = b64decode(proxy_payload.get("body_base64") or "")
         elif proxy_payload.get("multipart") is not None:
             envelope.payload = {"_broker_multipart": proxy_payload.get("multipart")}
-        logger.debug("broker dispatch proxy resolved → %s %s", envelope.method, envelope.path)
     elif envelope.op in OP_ROUTE_MAP:
         envelope.method, envelope.path = OP_ROUTE_MAP[envelope.op]
-        logger.debug("broker dispatch op mapped → %s %s", envelope.method, envelope.path)
     elif not envelope.path:
-        logger.warning("broker dispatch unsupported op=%s request_id=%s", envelope.op, envelope.request_id)
-        return error_envelope(
-            envelope.request_id,
-            code="unsupported_operation",
-            message=f"Unsupported broker op: {envelope.op}",
-        )
+        raise BrokerDispatchError(error_envelope(
+            envelope.request_id, code="unsupported_operation",
+            message=f"Unsupported broker op: {envelope.op}"))
 
     envelope.validate()
 
     if not is_supported_path(envelope.path):
-        logger.warning(
-            "broker dispatch unsupported path=%r op=%s request_id=%s",
-            envelope.path, envelope.op, envelope.request_id,
-        )
-        return error_envelope(
-            envelope.request_id,
-            code="unsupported_endpoint",
-            message=f"Unsupported endpoint: {envelope.path}",
-        )
+        raise BrokerDispatchError(error_envelope(
+            envelope.request_id, code="unsupported_endpoint",
+            message=f"Unsupported endpoint: {envelope.path}"))
 
-    body: bytes
     if isinstance(envelope.payload, dict) and "_broker_multipart" in envelope.payload:
         from codai.broker.asgi_bridge import _build_multipart_body
-
-        body, multipart_content_type = _build_multipart_body(envelope.payload["_broker_multipart"] or {})
+        body, multipart_content_type = _build_multipart_body(
+            envelope.payload["_broker_multipart"] or {})
         headers = dict(envelope.headers)
         headers["content-type"] = multipart_content_type
     elif isinstance(envelope.payload, (dict, list)):
@@ -139,6 +128,28 @@ async def execute_broker_request(app, envelope, executor=None):
 
     if body and "content-type" not in {key.lower() for key in headers}:
         headers["content-type"] = envelope.content_type
+    return headers, body
+
+
+async def execute_broker_request(app, envelope, executor=None):
+    """Validate and execute a broker request envelope.
+
+    ``executor`` is an ``async (method, path, headers, query, body) -> {status_code,
+    headers, body}`` callable. When omitted the request is run in-process against
+    ``app`` via the ASGI bridge (engine / single-process mode). The front passes its
+    own executor that proxies to the right engine over HTTP."""
+
+    logger.debug(
+        "broker dispatch → op=%s request_id=%s path=%r method=%r stream=%s",
+        envelope.op, envelope.request_id, envelope.path, envelope.method, envelope.stream,
+    )
+
+    try:
+        headers, body = resolve_broker_request(envelope)
+    except BrokerDispatchError as _bde:
+        logger.warning("broker dispatch unsupported op=%s path=%r request_id=%s",
+                       envelope.op, envelope.path, envelope.request_id)
+        return _bde.envelope
 
     started_at = perf_counter()
     if executor is not None:

codai/broker/streaming.py

@@ -8,21 +8,23 @@ from codai.broker.models import success_envelope
 
 
 def stream_chunk_envelope(request_id: str, sequence: int, data: Any) -> dict[str, Any]:
-    """Build a normalized streaming chunk envelope."""
+    """Build a streaming chunk envelope.
 
+    Uses ``event="chunk"`` with the SSE text under ``payload.chunk`` — the shape the
+    AISBF broker relay consumes (``_iter_broker_stream_chunks``: it yields
+    ``payload.chunk`` for ``event in {chunk,progress,output,log,data}``)."""
     return success_envelope(
         request_id,
-        event="stream",
-        payload={"sequence": sequence, "data": data},
+        event="chunk",
+        payload={"sequence": sequence, "chunk": data},
     )
 
 
 def finalize_stream(request_id: str, total_chunks: int, elapsed_ms: float) -> dict[str, Any]:
-    """Build the terminal streaming metadata envelope."""
-
+    """Terminal streaming envelope. ``event="done"`` ends the relay's stream loop."""
     return success_envelope(
         request_id,
-        event="stream_end",
+        event="done",
         payload={"total_chunks": total_chunks},
         metrics={"elapsed_ms": elapsed_ms},
     )

codai/frontproxy/app.py

@@ -160,6 +160,25 @@ class FrontProxy:
             return await execute_broker_request(None, envelope,
                                                 executor=self.broker_execute)
         client.dispatcher = _dispatch
+
+        # Streaming dispatcher: for stream=true inference, yield engine SSE chunks so
+        # the broker client relays them token-by-token (chunk envelopes) instead of
+        # buffering the whole reply.
+        from codai.broker.dispatcher import (resolve_broker_request,
+                                             BrokerDispatchError)
+
+        async def _stream_dispatch(message):
+            envelope = client.message_to_envelope(message)
+            try:
+                headers, body = resolve_broker_request(envelope)
+            except BrokerDispatchError:
+                yield 'data: {"error":"unsupported broker request"}\n\n'
+                return
+            async for chunk in self.broker_execute_stream(
+                    method=envelope.method, path=envelope.path, headers=headers,
+                    query=envelope.query, body=body):
+                yield chunk
+        client.stream_dispatcher = _stream_dispatch
         self._broker = BrokerService(client)   # app=None → keep our dispatcher
         self._broker.start()
         print("[front] AISBF broker started (front-managed, routes to engines)",
@@ -333,6 +352,80 @@ class FrontProxy:
         return {"status_code": r.status_code, "headers": dict(r.headers),
                 "body": r.content}
 
+    async def broker_execute_stream(self, *, method, path, headers, query, body):
+        """Streaming executor for brokered inference: route to the engine, open a
+        streamed SSE response, and YIELD each chunk as it arrives (as text). The
+        broker client wraps each yielded chunk in a ``chunk`` envelope and sends a
+        terminal ``done`` — so the AISBF relay streams tokens to the client instead
+        of buffering the whole reply. Shares the per-model queue + in-flight tracking
+        with the buffered path."""
+        import json as _json
+        model = None
+        if method.upper() == "POST" and _router.is_inference_path(path):
+            try:
+                model = (_json.loads(body or b"{}") or {}).get("model")
+            except Exception:
+                model = None
+        engine = _router.pick_engine(
+            self.registry, path, method, model,
+            required_cap=self._required_cap(path, model),
+            default_engine=self.default_engine, pinned=self._pin_for(model),
+            pin_fallback=bool(self._model_info(model).get("engine_fallback")))
+        if engine is None:
+            yield 'data: {"error":"No engine is ready yet."}\n\n'
+            return
+        send_headers = {k: v for k, v in (headers or {}).items()
+                        if k.lower() not in _DROP_REQ}
+        if self.internal_token:
+            send_headers["x-coderai-broker-authed"] = self.internal_token
+        _qkey = None
+        if (method.upper() == "POST" and _router.is_inference_path(path)
+                and self._task_kind(path) == "text"):
+            _qkey = self._queue_key(model)
+            try:
+                await self.reqqueue.acquire(
+                    _qkey, self._model_capacity(model), self._queue_max_waiting(),
+                    rid=engine.name + ":" + (model or ""), model=model or "",
+                    engine=engine.name)
+            except QueueFull:
+                yield ('data: {"error":"Server busy: the generation queue is full, '
+                       'please retry shortly."}\n\n')
+                return
+        _rid = engine.enter_request(
+            {"model": model or "", "kind": self._task_kind(path), "path": path}
+            if _router.is_inference_path(path) else None)
+        import time as _t
+        _started = _t.time()
+        _status = 502
+        rp_req = self._long.build_request(method, engine.url + path,
+                                          headers=send_headers, params=query or {},
+                                          content=body or b"")
+        try:
+            rp_resp = await self._long.send(rp_req, stream=True)
+            _status = rp_resp.status_code
+            _meas = (rp_resp.status_code == 200 and "text/event-stream"
+                     in (rp_resp.headers.get("content-type") or ""))
+            ntok = 0
+            async for raw in rp_resp.aiter_raw():
+                if not raw:
+                    continue
+                if _meas:
+                    ntok += raw.count(b"data:")
+                    m = (engine.active or {}).get(_rid)
+                    if m is not None:
+                        m["step"] = ntok
+                yield raw.decode("utf-8", "replace")
+            await rp_resp.aclose()
+        except Exception as exc:
+            yield ('data: {"error":"engine#%s unreachable: %s"}\n\n'
+                   % (engine.id, exc))
+        finally:
+            engine.exit_request(_rid)
+            if _qkey is not None:
+                await self.reqqueue.release(_qkey)
+            if _router.is_inference_path(path):
+                self._record_activity(model, self._task_kind(path), _status, _started)
+
     # ------------------------------------------------------------------ helpers
     @staticmethod
     def _filter_headers(headers, drop) -> list: