fix(pi): address code-review findings in mlflow proxy + OSS limits

Proxy (_mlflow_proxy.py):
- Close the connection at stream end (self.close_connection = True) so HTTP/1.1
  clients detect the message boundary instead of hanging after [DONE] (the
  proxy sends no Content-Length/chunking). [critical]
- Only repair text/event-stream responses; relay non-streaming and error bodies
  verbatim with their real Content-Type instead of forcing SSE. [high]
- Catch URLError/OSError around urlopen and return a clean 502 instead of a dead
  socket on connection-level failures. [high]
- Restrict forwarded paths to /ai-gateway/mlflow/ (404 otherwise) so the
  localhost relay can't be used as an authenticated arbitrary-path workspace
  client. [security]
- Move terminator-repair writes inside the OSError guard so a client
  disconnect doesn't traceback on the daemon thread. [medium]
- Parse SSE data lines with or without the optional space after the colon. [low]

Discovery (databricks.py):
- OSS chat models matching a broad family (qwen, llama-) but lacking a specific
  _MODEL_TOKEN_LIMITS entry now fall back to a conservative floor
  (128k/8192) so they're never offered uncapped (which 400s). Non-OSS and
  embedding ids still return None.

Verified live: glm-5-2 streams content (42) + finish_reason + [DONE] through the
proxy; non-stream/error/404/502 paths covered by tests. 898 passed.

Author: dgokeeffe

src/ucode/agents/_mlflow_proxy.py

@@ -7,14 +7,13 @@
 with "Stream ended without finish_reason", so the model is unusable.
 
 This proxy sits between Pi and the gateway. It forwards every request verbatim
-and passes every response byte through unchanged, except that it repairs the
-stream terminator: when a streaming response ends — cleanly at ``[DONE]`` OR
-abnormally (an upstream drop / mid-stream 429, which otherwise reaches Pi as
-"Stream ended without finish_reason") — after some data but without a
-``finish_reason``, it synthesizes a ``finish_reason: "stop"`` chunk (and a
-``[DONE]`` if the upstream never sent one). Models that already terminate
-correctly (glm, llama, qwen, ...) never trigger the injection, so the proxy is
-a no-op for them.
+and, for streaming (``text/event-stream``) responses, repairs the terminator:
+when the stream ends — cleanly at ``[DONE]`` OR abnormally (an upstream drop /
+mid-stream 429) — after some data but without a ``finish_reason``, it
+synthesizes a ``finish_reason: "stop"`` chunk (and a ``[DONE]`` if the upstream
+never sent one). Non-streaming and error responses are relayed byte-for-byte
+with their original ``Content-Type``. Well-behaved streams never trigger the
+injection, so the proxy is a no-op for them.
 
 Tracked upstream as the gateway bug that should make this unnecessary; once the
 gateway emits ``finish_reason`` this module can be deleted.
@@ -32,28 +31,56 @@
 # Request headers we must not forward: hop-by-hop or ones urllib recomputes.
 _SKIP_REQUEST_HEADERS = frozenset({"host", "content-length", "accept-encoding", "connection"})
 
+# The proxy only exists to front the mlflow chat-completions route. Refuse any
+# other path so a co-located process can't turn the localhost relay into an
+# arbitrary authenticated workspace client (SSRF-to-workspace) using the token
+# we forward.
+_ALLOWED_PATH_PREFIX = "/ai-gateway/mlflow/"
+
 
 def _make_handler(upstream: str) -> type[BaseHTTPRequestHandler]:
     class _Handler(BaseHTTPRequestHandler):
         protocol_version = "HTTP/1.1"
 
         def do_POST(self) -> None:  # noqa: N802 (BaseHTTPRequestHandler API)
+            # Read the body regardless so the socket is drained before we close.
             length = int(self.headers.get("Content-Length", 0))
             body = self.rfile.read(length)
+
+            if not self.path.startswith(_ALLOWED_PATH_PREFIX):
+                self._send_bytes(404, b'{"error":"path not allowed"}', "application/json")
+                return
+
             req = urllib.request.Request(upstream + self.path, data=body, method="POST")
             for key, value in self.headers.items():
                 if key.lower() not in _SKIP_REQUEST_HEADERS:
                     req.add_header(key, value)
             try:
                 upstream_resp = urllib.request.urlopen(req)  # noqa: S310 (trusted upstream)
             except urllib.error.HTTPError as exc:
-                payload = exc.read()
-                self.send_response(exc.code)
-                self.send_header("Content-Length", str(len(payload)))
-                self.end_headers()
-                self.wfile.write(payload)
+                # Relay the gateway's own error verbatim (status + body).
+                payload = exc.read() if exc.fp else b""
+                content_type = exc.headers.get("Content-Type", "application/json")
+                self._send_bytes(exc.code, payload, content_type)
+                return
+            except (urllib.error.URLError, OSError) as exc:
+                # Connection-level failure (DNS/TLS/timeout/reset): give the
+                # client a clean 502 instead of a dead socket.
+                msg = json.dumps({"error": f"upstream unreachable: {exc}"}).encode()
+                self._send_bytes(502, msg, "application/json")
+                return
+
+            content_type = upstream_resp.headers.get("Content-Type", "")
+            if "text/event-stream" not in content_type:
+                # Non-streaming (or error-envelope) 200: relay unchanged.
+                self._send_bytes(
+                    upstream_resp.status, upstream_resp.read(), content_type or "application/json"
+                )
                 return
 
+            # Streaming: close the connection at end-of-stream so an HTTP/1.1
+            # client can detect the message boundary (we send no length/chunking).
+            self.close_connection = True
             self.send_response(upstream_resp.status)
             self.send_header("Content-Type", "text/event-stream")
             self.send_header("Cache-Control", "no-cache")
@@ -68,39 +95,51 @@ def _relay_stream(self, upstream_resp) -> None:
             last_id: str | None = None
             try:
                 for raw in upstream_resp:  # yields line-by-line as bytes arrive
-                    stripped = raw.strip()
-                    if stripped.startswith(b"data: ") and stripped[6:] != b"[DONE]":
+                    payload = _sse_data_payload(raw)
+                    if payload is not None and payload != b"[DONE]":
                         saw_data = True
                         try:
-                            obj = json.loads(stripped[6:])
+                            obj = json.loads(payload)
                             last_id = obj.get("id", last_id)
                             choice = (obj.get("choices") or [{}])[0]
                             if choice.get("finish_reason"):
                                 saw_finish = True
                         except (ValueError, AttributeError, IndexError):
                             pass
                         self._write(raw)
-                    elif stripped == b"data: [DONE]":
+                    elif payload == b"[DONE]":
                         if not saw_finish:
                             self._write(b"data: " + _finish_chunk(last_id) + b"\n\n")
                             saw_finish = True
                         self._write(raw)
                         saw_done = True
                     else:
                         self._write(raw)
+                # Clean or abnormal end after data but without a terminator:
+                # synthesize one so strict clients don't error on a truncated
+                # stream (a mid-stream 429 is the common trigger).
+                if saw_data and not saw_finish:
+                    self._write(b"data: " + _finish_chunk(last_id) + b"\n\n")
+                if saw_data and not saw_done:
+                    self._write(b"data: [DONE]\n\n")
+            except OSError:
+                # Upstream dropped, or the client disconnected mid-write. Either
+                # way the socket is gone; swallow so the daemon thread exits
+                # quietly rather than tracebacking.
+                pass
+
+        def _send_bytes(self, status: int, data: bytes, content_type: str) -> None:
+            self.close_connection = True
+            try:
+                self.send_response(status)
+                self.send_header("Content-Type", content_type)
+                self.send_header("Content-Length", str(len(data)))
+                self.send_header("Connection", "close")
+                self.end_headers()
+                self.wfile.write(data)
+                self.wfile.flush()
             except OSError:
-                # Upstream stream dropped mid-flight (throttle/reset). Fall
-                # through to the terminator repair below so the client still
-                # sees a well-formed end instead of a truncated stream.
                 pass
-            # If the stream ended (cleanly or abnormally) after some data but
-            # without a finish_reason and/or [DONE], synthesize the terminator
-            # so strict clients (Pi) don't error on an incomplete stream. A
-            # gateway 429 mid-stream is the common trigger.
-            if saw_data and not saw_finish:
-                self._write(b"data: " + _finish_chunk(last_id) + b"\n\n")
-            if saw_data and not saw_done:
-                self._write(b"data: [DONE]\n\n")
 
         def _write(self, data: bytes) -> None:
             self.wfile.write(data)
@@ -112,6 +151,18 @@ def log_message(self, format: str, *args) -> None:  # silence default stderr log
     return _Handler
 
 
+def _sse_data_payload(raw: bytes) -> bytes | None:
+    """Return the payload of an SSE ``data:`` line (the spec allows an optional
+    single space after the colon), or None if the line isn't a data line."""
+    stripped = raw.strip()
+    if not stripped.startswith(b"data:"):
+        return None
+    payload = stripped[len(b"data:") :]
+    if payload.startswith(b" "):
+        payload = payload[1:]
+    return payload
+
+
 def _finish_chunk(chunk_id: str | None) -> bytes:
     return json.dumps(
         {

src/ucode/databricks.py

@@ -1189,6 +1189,16 @@ def _is_oss_chat_model(model_id: str) -> bool:
     "gemma": {"context": 128_000, "output": 8_192},
 }
 
+# Conservative fallback for a model that IS an OSS chat model (bucketed by
+# `_is_oss_chat_model`) but has no specific `_MODEL_TOKEN_LIMITS` entry — e.g. a
+# `qwen3-coder` or future `llama-5`. The discovery families are intentionally
+# broader than the limits table, so without a floor such a model would be
+# offered with no output cap and 400 on a default-size request. 8192 is the
+# lowest cap observed across the cohort, so every mlflow chat model accepts it;
+# pinning it only risks truncating a longer response (the safe failure
+# direction), never a hard 400.
+_OSS_FALLBACK_LIMITS = {"context": 128_000, "output": 8_192}
+
 # OSS families that emit reasoning (openai_reasoning capability, verified
 # 2026-07-16). Pi renders their streamed reasoning_content as thinking when the
 # model entry sets reasoning:true; agents whose SDK surfaces reasoning natively
@@ -1204,11 +1214,15 @@ def model_is_reasoning(model_id: str) -> bool:
 def model_token_limits(model_id: str) -> dict[str, int] | None:
     """Return ``{"context": ..., "output": ...}`` limits for ``model_id``, or None.
 
-    Matches by family substring (e.g. any ``*glm*`` id). None means the model
-    has no known limits and the agent should not pin any."""
+    Prefers a specific `_MODEL_TOKEN_LIMITS` family entry (e.g. any ``*glm*``
+    id). Any other OSS chat model falls back to a conservative floor so it is
+    never offered uncapped (which would 400). None only for non-OSS ids, where
+    the agent should not pin any limit."""
     for family, limits in _MODEL_TOKEN_LIMITS.items():
         if family in model_id:
             return dict(limits)
+    if _is_oss_chat_model(model_id):
+        return dict(_OSS_FALLBACK_LIMITS)
     return None
 
 

tests/test_databricks.py

@@ -185,9 +185,22 @@ def test_qwen_versions_distinct_caps(self):
             db_mod.model_token_limits("system.ai.qwen3-next-80b-a3b-instruct")["output"] == 10_000
         )
 
-    def test_uncapped_model_returns_none(self):
+    def test_oss_chat_model_without_specific_entry_uses_fallback(self):
+        # qwen3-coder is an OSS chat model (matches the `qwen` family) but has
+        # no specific limits entry — it must get the conservative floor, never
+        # None (which would leave it uncapped and 400 on a default request).
+        limits = db_mod.model_token_limits("system.ai.qwen3-coder-480b")
+        assert limits == {"context": 128_000, "output": 8_192}
+
+    def test_non_oss_model_returns_none(self):
+        # deepseek is outside every family (not an OSS chat model) -> no cap.
         assert db_mod.model_token_limits("system.ai.deepseek-v3") is None
 
+    def test_embedding_model_returns_none_not_fallback(self):
+        # An embeddings model is excluded from OSS chat, so it must NOT get the
+        # chat fallback limits.
+        assert db_mod.model_token_limits("system.ai.qwen3-embedding-0-6b") is None
+
 
 class TestModelIsReasoning:
     def test_reasoning_families(self):

tests/test_mlflow_proxy.py

@@ -2,7 +2,9 @@
 
 from __future__ import annotations
 
+import json
 import threading
+import urllib.error
 import urllib.request
 from http.server import BaseHTTPRequestHandler, HTTPServer
 
@@ -21,15 +23,23 @@
 # An abnormally-terminated stream: content delta, then the upstream just stops —
 # no finish_reason AND no [DONE] (what a mid-stream 429/drop looks like).
 _STREAM_ABRUPT_END = b'data: {"id":"c3","choices":[{"delta":{"content":"ok"},"index":0}]}\n\n'
+# SSE data line with NO space after the colon (spec-legal) and no finish_reason.
+_STREAM_NO_SPACE = (
+    b'data:{"id":"c4","choices":[{"delta":{"content":"ok"},"index":0}]}\n\ndata:[DONE]\n\n'
+)
+# A non-streaming JSON response (some models / non-stream requests).
+_JSON_BODY = b'{"id":"c5","choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}'
 
 
-def _make_fake_gateway(payload: bytes) -> tuple[str, HTTPServer]:
+def _make_fake_gateway(
+    payload: bytes, content_type: str = "text/event-stream", status: int = 200
+) -> tuple[str, HTTPServer]:
     class _Fake(BaseHTTPRequestHandler):
         protocol_version = "HTTP/1.1"
 
         def do_POST(self):  # noqa: N802
-            self.send_response(200)
-            self.send_header("Content-Type", "text/event-stream")
+            self.send_response(status)
+            self.send_header("Content-Type", content_type)
             self.send_header("Content-Length", str(len(payload)))
             self.end_headers()
             self.wfile.write(payload)
@@ -42,26 +52,35 @@ def log_message(self, format: str, *args):
     return f"http://127.0.0.1:{server.server_address[1]}", server
 
 
-def _post_through_proxy(upstream: str) -> str:
+def _post_through_proxy(upstream: str, path: str = "/ai-gateway/mlflow/v1/chat/completions"):
+    """POST through the proxy; returns (status, body, content_type)."""
     base, stop = _mlflow_proxy.start(upstream)
     try:
         req = urllib.request.Request(
-            f"{base}/ai-gateway/mlflow/v1/chat/completions",
+            f"{base}{path}",
             data=b'{"model":"databricks-inkling","stream":true}',
             headers={"Content-Type": "application/json"},
             method="POST",
         )
-        with urllib.request.urlopen(req, timeout=5) as resp:
-            return resp.read().decode("utf-8")
+        try:
+            with urllib.request.urlopen(req, timeout=5) as resp:
+                return resp.status, resp.read().decode("utf-8"), resp.headers.get("Content-Type")
+        except urllib.error.HTTPError as exc:
+            return exc.code, exc.read().decode("utf-8"), exc.headers.get("Content-Type")
     finally:
         stop.set()
 
 
+def _body(upstream: str) -> str:
+    """Convenience: body text of a normal streaming POST through the proxy."""
+    return _post_through_proxy(upstream)[1]
+
+
 class TestFinishReasonInjection:
     def test_injects_finish_reason_when_absent(self):
         upstream, server = _make_fake_gateway(_STREAM_NO_FINISH)
         try:
-            out = _post_through_proxy(upstream)
+            out = _body(upstream)
         finally:
             server.shutdown()
         assert '"finish_reason": "stop"' in out or '"finish_reason":"stop"' in out
@@ -72,7 +91,7 @@ def test_injects_finish_reason_when_absent(self):
     def test_passes_through_when_finish_reason_present(self):
         upstream, server = _make_fake_gateway(_STREAM_WITH_FINISH)
         try:
-            out = _post_through_proxy(upstream)
+            out = _body(upstream)
         finally:
             server.shutdown()
         # Exactly one finish_reason — the proxy did not add a second.
@@ -81,7 +100,7 @@ def test_passes_through_when_finish_reason_present(self):
     def test_forwards_content_unchanged(self):
         upstream, server = _make_fake_gateway(_STREAM_NO_FINISH)
         try:
-            out = _post_through_proxy(upstream)
+            out = _body(upstream)
         finally:
             server.shutdown()
         assert '"content":"ok"' in out
@@ -92,17 +111,70 @@ def test_repairs_abrupt_end_with_finish_and_done(self):
         # doesn't error on a truncated stream.
         upstream, server = _make_fake_gateway(_STREAM_ABRUPT_END)
         try:
-            out = _post_through_proxy(upstream)
+            out = _body(upstream)
         finally:
             server.shutdown()
         assert "finish_reason" in out
         assert out.rstrip().endswith("[DONE]")
         assert '"content":"ok"' in out
 
+    def test_parses_data_line_without_space(self):
+        # SSE allows `data:{...}` (no space). The already-present finish_reason
+        # must be detected so no duplicate is injected.
+        upstream, server = _make_fake_gateway(_STREAM_NO_SPACE)
+        try:
+            out = _body(upstream)
+        finally:
+            server.shutdown()
+        # One synthetic finish (none was present) — not two.
+        assert out.count("finish_reason") == 1
+
+
+class TestNonStreamingAndErrors:
+    def test_non_streaming_json_relayed_unchanged(self):
+        # A non-SSE 200 must pass through verbatim with its content-type — no
+        # finish_reason injection, no SSE relabeling.
+        upstream, server = _make_fake_gateway(_JSON_BODY, content_type="application/json")
+        try:
+            status, body, ctype = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        assert status == 200
+        assert "application/json" in ctype
+        assert body == _JSON_BODY.decode()
+        assert "chat.completion.chunk" not in body  # no injected SSE chunk
+
+    def test_gateway_error_relayed_with_status(self):
+        # A gateway 400 (e.g. max_tokens exceeded) is relayed with its status.
+        err = b'{"error":"max_tokens cannot exceed 8192"}'
+        upstream, server = _make_fake_gateway(err, content_type="application/json", status=400)
+        try:
+            status, body, _ = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        assert status == 400
+        assert "max_tokens" in body
+
+    def test_connection_failure_returns_502(self):
+        # Point the proxy at a dead port: urlopen raises URLError, and the proxy
+        # must return a clean 502, not a dead socket.
+        status, body, _ = _post_through_proxy("http://127.0.0.1:1")
+        assert status == 502
+        assert "upstream unreachable" in body
+
+    def test_disallowed_path_returns_404(self):
+        # Only the mlflow route may be forwarded (SSRF guard).
+        upstream, server = _make_fake_gateway(_JSON_BODY, content_type="application/json")
+        try:
+            status, body, _ = _post_through_proxy(upstream, path="/api/2.0/secrets/get")
+        finally:
+            server.shutdown()
+        assert status == 404
+        assert "not allowed" in body
+
 
 class TestFinishChunk:
     def test_shape(self):
-        import json
 
         chunk = json.loads(_mlflow_proxy._finish_chunk("abc"))
         assert chunk["id"] == "abc"