Merge pull request #126 from secondly-com/feat/broker-streaming-and-tests

Stream broker responses through and add pytest coverage

Author: adamcohenhillel

.github/workflows/ci.yml

@@ -45,5 +45,10 @@ jobs:
       - name: Run OpenPhone checks
         run: ./scripts/check.sh
 
+      - name: Run model broker unit tests
+        run: |
+          pip install pytest
+          python3 -m pytest services/model-broker/tests -q
+
       - name: Check whitespace
         run: git diff --check

scripts/smoke-test-model-broker.sh

@@ -65,14 +65,31 @@ import http.server
 import json
 import pathlib
 import sys
+import time
 
 state_path = pathlib.Path(sys.argv[1])
 
 
 class Handler(http.server.BaseHTTPRequestHandler):
     def do_POST(self):
         length = int(self.headers.get("Content-Length", "0"))
-        self.rfile.read(length)
+        body = self.rfile.read(length)
+        try:
+            payload = json.loads(body.decode("utf-8"))
+        except (UnicodeDecodeError, json.JSONDecodeError):
+            payload = {}
+        if isinstance(payload, dict) and payload.get("stream") is True:
+            self.send_response(200)
+            self.send_header("Content-Type", "text/event-stream")
+            self.send_header("Connection", "close")
+            self.end_headers()
+            for index in range(3):
+                self.wfile.write(f"data: chunk-{index}\n\n".encode("utf-8"))
+                self.wfile.flush()
+                time.sleep(0.4)
+            self.wfile.write(b"data: [DONE]\n\n")
+            self.wfile.flush()
+            return
         attempts = int(state_path.read_text(encoding="utf-8")) if state_path.exists() else 0
         attempts += 1
         state_path.write_text(str(attempts), encoding="utf-8")
@@ -318,6 +335,59 @@ expect_status 200 \
   -H 'Content-Type: application/json' \
   --data '{"model":"gpt-4.1-mini","metadata":{"openphone_task":"true","risk_flags":""},"input":[{"role":"user","content":[{"type":"input_text","text":"hello"}]}]}'
 
+# Streaming: the broker must relay SSE chunks incrementally instead of
+# buffering the full upstream response before replying.
+python3 - <<'PY' "$port" "$static_token"
+import http.client
+import json
+import sys
+import time
+
+port, token = sys.argv[1:]
+body = json.dumps({
+    "model": "gpt-4.1-mini",
+    "stream": True,
+    "metadata": {"openphone_task": "true", "risk_flags": ""},
+    "input": [{"role": "user", "content": [{"type": "input_text", "text": "hi"}]}],
+})
+connection = http.client.HTTPConnection("127.0.0.1", int(port), timeout=30)
+connection.request(
+    "POST",
+    "/v1/responses",
+    body=body,
+    headers={
+        "Authorization": f"Bearer {token}",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream",
+    },
+)
+response = connection.getresponse()
+if response.status != 200:
+    raise SystemExit(f"streaming request failed: HTTP {response.status}")
+content_type = response.headers.get("Content-Type", "")
+if not content_type.startswith("text/event-stream"):
+    raise SystemExit(f"streaming response has wrong content-type: {content_type}")
+
+chunks = []
+arrival_times = []
+while True:
+    chunk = response.read1(65536)
+    if not chunk:
+        break
+    chunks.append(chunk)
+    arrival_times.append(time.monotonic())
+payload = b"".join(chunks).decode("utf-8")
+for expected in ("data: chunk-0", "data: chunk-1", "data: chunk-2", "data: [DONE]"):
+    if expected not in payload:
+        raise SystemExit(f"streaming response missing {expected!r}: {payload!r}")
+if len(arrival_times) < 2:
+    raise SystemExit("streaming response arrived as a single buffered chunk")
+spread = arrival_times[-1] - arrival_times[0]
+if spread < 0.2:
+    raise SystemExit(f"streaming chunks were not delivered incrementally (spread={spread:.3f}s)")
+print(f"streaming smoke case passed ({len(chunks)} chunks over {spread:.2f}s)")
+PY
+
 expect_status 429 \
   -X POST "http://127.0.0.1:$port/v1/responses" \
   -H "Authorization: Bearer $static_token" \
@@ -374,7 +444,7 @@ for event in events:
     if "body" in event or "request" in event or "response" in event:
         raise SystemExit(f"audit event contains body-like key: {event}")
 proxied = [event for event in events if event.get("outcome") == "proxied"]
-if not proxied or proxied[-1].get("provider_attempts") != 2:
+if not any(event.get("provider_attempts") == 2 for event in proxied):
     raise SystemExit("broker did not record retried provider forwarding")
 PY
 

services/model-broker/openphone_model_broker.py

@@ -211,6 +211,7 @@ def do_POST(self) -> None:  # noqa: N802
                 body_size,
                 started_at,
                 model,
+                stream=self._wants_stream(payload),
             )
             return
 
@@ -326,6 +327,7 @@ def _proxy_to_openai(
         body_size: int,
         started_at: float,
         model: str | None,
+        stream: bool = False,
     ) -> None:
         headers = {
             "Authorization": f"Bearer {self.server.config.api_key}",
@@ -338,6 +340,17 @@ def _proxy_to_openai(
             attempts += 1
             try:
                 with urllib.request.urlopen(request, timeout=120) as response:
+                    if stream:
+                        self._stream_provider_response(
+                            response,
+                            token_subject,
+                            body_size,
+                            started_at,
+                            endpoint,
+                            model,
+                            attempts,
+                        )
+                        return
                     response_body = response.read()
                     self._audit_request(
                         "proxied",
@@ -404,6 +417,53 @@ def _proxy_to_openai(
                 )
                 return
 
+    def _wants_stream(self, payload: dict[str, object]) -> bool:
+        accept = self.headers.get("Accept", "")
+        if "text/event-stream" in accept.lower():
+            return True
+        return payload.get("stream") is True
+
+    def _stream_provider_response(
+        self,
+        response: object,
+        token_subject: str,
+        body_size: int,
+        started_at: float,
+        endpoint: str,
+        model: str | None,
+        attempts: int,
+    ) -> None:
+        """Relay the upstream body chunk-by-chunk instead of buffering it.
+
+        The response is delimited by connection close (HTTP/1.0 semantics),
+        which every SSE-capable client handles; no Content-Length is sent.
+        """
+        self._audit_request(
+            "proxied",
+            token_subject,
+            response.status,
+            body_size,
+            started_at,
+            endpoint,
+            model,
+            provider_attempts=attempts,
+        )
+        self.send_response(response.status)
+        self.send_header(
+            "Content-Type",
+            response.headers.get("Content-Type", "text/event-stream"),
+        )
+        self.send_header("Cache-Control", "no-store")
+        self.send_header("Connection", "close")
+        self.end_headers()
+        self.close_connection = True
+        while True:
+            chunk = response.read1(65536)
+            if not chunk:
+                break
+            self.wfile.write(chunk)
+            self.wfile.flush()
+
     def _should_retry_provider(self, status: int, attempts: int) -> bool:
         if attempts > self.server.config.provider_max_retries:
             return False

services/model-broker/tests/conftest.py

@@ -0,0 +1,8 @@
+"""Make the stdlib-only broker module importable from the tests directory."""
+
+import pathlib
+import sys
+
+BROKER_DIR = pathlib.Path(__file__).resolve().parent.parent
+if str(BROKER_DIR) not in sys.path:
+    sys.path.insert(0, str(BROKER_DIR))