From 23d71a5f53c3b15f0ba05c7f56e0a6030393ec68 Mon Sep 17 00:00:00 2001
From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com>
Date: Thu, 16 Jul 2026 23:29:41 +1000
Subject: [PATCH 1/4] feat(pi): add databricks-mlflow OSS provider + broaden
 OSS cohort with reasoning/limits

Builds on the merged OpenCode OSS support (kimi/glm) to close the remaining gaps:

- Pi had NO OSS provider. Add databricks-mlflow (api openai-completions ->
  /ai-gateway/mlflow/v1) so Pi can select OSS chat models, with a local
  SSE-repair proxy (_mlflow_proxy) for models like inkling whose gateway stream
  omits finish_reason (Pi's parser rejects that; OpenCode's @ai-sdk/openai
  tolerates it, so it needs no proxy).
- Broaden _OSS_MODEL_FAMILIES from kimi/glm to the full chat-completions cohort
  (inkling, llama, qwen, gpt-oss, gemma), excluding non-chat services
  (embeddings/rerankers) that share a family substring (e.g. qwen3-embedding).
- Correct + extend _MODEL_TOKEN_LIMITS from probed gateway caps: glm is
  1M context / 65536 output (was 200k/25k); per-family output ceilings for the
  cohort. Add model_is_reasoning() from probed openai_reasoning capability;
  Pi sets reasoning:true so it renders streamed reasoning_content as thinking.

Both agents read the shared model_token_limits/model_is_reasoning tables.
Verified live: glm-5-2 answers through Pi (6*7 -> 42).
---
 src/ucode/agents/_mlflow_proxy.py | 150 ++++++++++++++++++++++++++++++
 src/ucode/agents/pi.py            |  93 ++++++++++++++++--
 src/ucode/cli.py                  |   2 +-
 src/ucode/databricks.py           |  58 ++++++++++--
 tests/test_agent_opencode.py      |  15 ++-
 tests/test_agent_pi.py            |  60 +++++++++++-
 tests/test_databricks.py          |  84 ++++++++++++++---
 tests/test_mlflow_proxy.py        | 122 ++++++++++++++++++++++++
 8 files changed, 553 insertions(+), 31 deletions(-)
 create mode 100644 src/ucode/agents/_mlflow_proxy.py
 create mode 100644 tests/test_mlflow_proxy.py

diff --git a/src/ucode/agents/_mlflow_proxy.py b/src/ucode/agents/_mlflow_proxy.py
new file mode 100644
index 0000000..70e26ed
--- /dev/null
+++ b/src/ucode/agents/_mlflow_proxy.py
@@ -0,0 +1,150 @@
+"""Local SSE-repair proxy for the MLflow chat-completions gateway.
+
+Some MLflow-served models (notably ``databricks-inkling``) stream a
+spec-incomplete response: on natural completion the terminal chunk omits the
+OpenAI-required ``finish_reason`` and the stream ends with only ``data: [DONE]``.
+Strict clients — Pi's ``openai-completions`` parser among them — reject this
+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.
+
+Tracked upstream as the gateway bug that should make this unnecessary; once the
+gateway emits ``finish_reason`` this module can be deleted.
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+import urllib.error
+import urllib.request
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from socketserver import ThreadingMixIn
+
+# Request headers we must not forward: hop-by-hop or ones urllib recomputes.
+_SKIP_REQUEST_HEADERS = frozenset({"host", "content-length", "accept-encoding", "connection"})
+
+
+def _make_handler(upstream: str) -> type[BaseHTTPRequestHandler]:
+    class _Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def do_POST(self) -> None:  # noqa: N802 (BaseHTTPRequestHandler API)
+            length = int(self.headers.get("Content-Length", 0))
+            body = self.rfile.read(length)
+            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)
+                return
+
+            self.send_response(upstream_resp.status)
+            self.send_header("Content-Type", "text/event-stream")
+            self.send_header("Cache-Control", "no-cache")
+            self.send_header("Connection", "close")
+            self.end_headers()
+            self._relay_stream(upstream_resp)
+
+        def _relay_stream(self, upstream_resp) -> None:
+            saw_finish = False
+            saw_done = False
+            saw_data = False
+            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]":
+                        saw_data = True
+                        try:
+                            obj = json.loads(stripped[6:])
+                            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]":
+                        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)
+            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)
+            self.wfile.flush()
+
+        def log_message(self, format: str, *args) -> None:  # silence default stderr logging
+            pass
+
+    return _Handler
+
+
+def _finish_chunk(chunk_id: str | None) -> bytes:
+    return json.dumps(
+        {
+            "id": chunk_id,
+            "object": "chat.completion.chunk",
+            "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}],
+        }
+    ).encode()
+
+
+class _Server(ThreadingMixIn, HTTPServer):
+    daemon_threads = True
+    allow_reuse_address = True
+
+
+def start(upstream: str) -> tuple[str, threading.Event]:
+    """Start the repair proxy on a free localhost port.
+
+    :param upstream: gateway origin to forward to, e.g.
+        ``https://workspace.databricks.com``.
+    :returns: ``(base_url, stop_event)`` where ``base_url`` is the local origin
+        Pi should target (``http://127.0.0.1:<port>``) and setting
+        ``stop_event`` shuts the server down.
+    """
+    server = _Server(("127.0.0.1", 0), _make_handler(upstream.rstrip("/")))
+    port = server.server_address[1]
+    stop_event = threading.Event()
+
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+
+    def _watch_stop() -> None:
+        stop_event.wait()
+        server.shutdown()
+
+    threading.Thread(target=_watch_stop, daemon=True).start()
+    return f"http://127.0.0.1:{port}", stop_event
diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py
index e7c1760..78d1ba6 100644
--- a/src/ucode/agents/pi.py
+++ b/src/ucode/agents/pi.py
@@ -1,12 +1,13 @@
 """Pi coding agent: writes ~/.pi/agent/models.json with Databricks-backed providers.
 
-Pi (https://pi.dev) is a multi-provider coding agent. We register three
+Pi (https://pi.dev) is a multi-provider coding agent. We register four
 providers in its `models.json`, each speaking the API dialect best suited to
 that family's gateway path:
 
 - `databricks-claude`  (api: anthropic-messages)       → /ai-gateway/anthropic
 - `databricks-openai`  (api: openai-responses)         → /ai-gateway/codex/v1
 - `databricks-gemini`  (api: google-generative-ai)     → /ai-gateway/gemini/v1beta
+- `databricks-mlflow`  (api: openai-completions)       → /ai-gateway/mlflow/v1
 
 Per-provider `compat` flags work around fields the gateway translators reject:
 
@@ -15,11 +16,20 @@
   pi uses for every request. With this flag pi omits the per-tool field and
   sends the legacy `anthropic-beta: fine-grained-tool-streaming-...` header
   instead, which the gateway accepts.
+- mlflow: `supportsStore: false` and `supportsStrictMode: false` — the MLflow
+  chat-completions gateway rejects OpenAI's `store` field and
+  `tools[].function.strict`.
 
-OSS / Databricks-foundation models (Llama, Qwen, etc.) are not exposed via
-pi today — they live behind /ai-gateway/mlflow/v1 with per-model
-`max_tokens` caps that pi has no global way to honor without per-model
-config we don't currently maintain.
+The `databricks-mlflow` provider carries the OSS chat-completions-only
+foundation models (Llama, Qwen, GLM, inkling, ...) discovered upstream. Per
+model it sets `contextWindow`/`maxTokens` from `databricks.model_token_limits`
+and `reasoning` from `databricks.model_is_reasoning` (so Pi renders the
+gateway's streamed reasoning_content as thinking).
+
+At launch the mlflow provider is routed through a local SSE-repair proxy (see
+`_mlflow_proxy`): some OSS models (inkling) omit the `finish_reason` on the
+streaming terminal chunk, which Pi's parser rejects. The proxy injects it and
+is a no-op for well-behaved models. Removable once the gateway is fixed.
 
 The bearer token is baked into the file and refreshed by a background thread
 while the session runs (same pattern as OpenCode/Copilot).
@@ -33,6 +43,7 @@
 import threading
 
 from ucode.agent_updates import available_npm_package_update
+from ucode.agents import _mlflow_proxy
 from ucode.config_io import (
     APP_DIR,
     ToolSpec,
@@ -45,6 +56,8 @@
     TOKEN_REFRESH_INTERVAL_SECONDS,
     build_pi_base_urls,
     get_databricks_token,
+    model_is_reasoning,
+    model_token_limits,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -68,6 +81,7 @@
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    "databricks-mlflow",
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
@@ -86,6 +100,7 @@ def _resolve_model_selector(
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    oss_models: list[str],
 ) -> str:
     """Return a Pi model selector in `<provider>/<model>` form when possible."""
     for name in PROVIDER_NAMES:
@@ -97,9 +112,29 @@ def _resolve_model_selector(
         return f"databricks-openai/{model}"
     if model in gemini_models:
         return f"databricks-gemini/{model}"
+    if model in oss_models:
+        return f"databricks-mlflow/{model}"
     return model
 
 
+def _pi_oss_model_entry(model_id: str) -> dict:
+    """Build a Pi mlflow model entry enriched from the shared limits/reasoning
+    tables: `reasoning:true` for reasoning models (Pi renders their streamed
+    reasoning_content as thinking), and `contextWindow`/`maxTokens` from
+    `model_token_limits`. Fields are omitted when unknown so Pi keeps its
+    default."""
+    entry: dict = {"id": model_id}
+    if model_is_reasoning(model_id):
+        entry["reasoning"] = True
+    limits = model_token_limits(model_id)
+    if limits:
+        if limits.get("context"):
+            entry["contextWindow"] = limits["context"]
+        if limits.get("output"):
+            entry["maxTokens"] = limits["output"]
+    return entry
+
+
 def render_overlay(
     model: str,
     token: str,
@@ -107,6 +142,7 @@ def render_overlay(
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    oss_models: list[str],
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for ~/.pi/agent/models.json."""
     providers: dict = {}
@@ -150,8 +186,23 @@ def render_overlay(
             "models": [{"id": m} for m in gemini_models],
         }
         keys.append(["providers", "databricks-gemini"])
+    if oss_models:
+        providers["databricks-mlflow"] = {
+            "baseUrl": pi_base_urls["oss"],
+            "api": "openai-completions",
+            "apiKey": token,
+            "authHeader": True,
+            # MLflow chat-completions gateway rejects OpenAI's `store` field
+            # and per-tool `strict`. Pi omits both when these are false.
+            "compat": {"supportsStore": False, "supportsStrictMode": False},
+            "headers": ua_headers,
+            "models": [_pi_oss_model_entry(m) for m in oss_models],
+        }
+        keys.append(["providers", "databricks-mlflow"])
     overlay: dict = {
-        "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
+        "model": _resolve_model_selector(
+            model, claude_models, codex_models, gemini_models, oss_models
+        ),
     }
     if providers:
         overlay["providers"] = providers
@@ -178,6 +229,7 @@ def write_tool_config(
         state.get("claude_models") or {},
         state.get("codex_models") or [],
         state.get("gemini_models") or [],
+        state.get("oss_models") or [],
     )
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
@@ -215,7 +267,10 @@ def default_model(state: dict) -> str | None:
     if codex_models:
         return codex_models[0]
     gemini_models = state.get("gemini_models") or []
-    return gemini_models[0] if gemini_models else None
+    if gemini_models:
+        return gemini_models[0]
+    oss_models = state.get("oss_models") or []
+    return oss_models[0] if oss_models else None
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
@@ -241,7 +296,29 @@ def build_runtime_env(token: str) -> dict[str, str]:
     return env
 
 
+def _start_oss_proxy(state: dict) -> threading.Event | None:
+    """Route the mlflow provider through the local SSE-repair proxy.
+
+    Rewrites ``state["base_urls"]["pi"]["oss"]`` in-memory to the proxy origin
+    so both the initial config write and every token-refresh rewrite point Pi
+    at the proxy. The gateway origin is always re-derived from the workspace
+    (never the persisted `oss` URL, which may hold a dead proxy port from a
+    prior session). Returns the proxy's stop event, or None when no OSS models
+    are configured. See `_mlflow_proxy`.
+    """
+    if not (state.get("oss_models") or []):
+        return None
+    pi_urls = state.setdefault("base_urls", {}).setdefault(
+        "pi", build_pi_base_urls(state["workspace"])
+    )
+    origin = build_pi_base_urls(state["workspace"])["oss"].split("/ai-gateway/", 1)[0]
+    proxy_base, stop_event = _mlflow_proxy.start(origin)
+    pi_urls["oss"] = f"{proxy_base}/ai-gateway/mlflow/v1"
+    return stop_event
+
+
 def launch(state: dict, tool_args: list[str]) -> None:
+    proxy_stop = _start_oss_proxy(state)
     token = _refresh_token_once(state)
     env = build_runtime_env(token)
 
@@ -262,6 +339,8 @@ def launch(state: dict, tool_args: list[str]) -> None:
     finally:
         stop_event.set()
         refresher.join(timeout=1)
+        if proxy_stop is not None:
+            proxy_stop.set()
 
     raise SystemExit(returncode)
 
diff --git a/src/ucode/cli.py b/src/ucode/cli.py
index 9ebdae9..9800d9d 100644
--- a/src/ucode/cli.py
+++ b/src/ucode/cli.py
@@ -324,7 +324,7 @@ def configure_shared_state(
     )
     want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools
     want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools
-    want_oss = fetch_all or "opencode" in tools
+    want_oss = fetch_all or "opencode" in tools or "pi" in tools
 
     claude_reason: str | None = None
     gemini_reason: str | None = None
diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index 9f89c6b..df3e27f 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -1097,8 +1097,31 @@ def build_auth_shell_command(
 _MODEL_SERVICE_REQUIRED_PREFIX = "system.ai."
 
 # Supported OSS chat families, matched by name substring. Add an entry to
-# support a new family.
-_OSS_MODEL_FAMILIES = ("kimi-", "glm-")
+# support a new family. Covers the full chat-completions-only cohort on the
+# `/ai-gateway/mlflow/v1` route (verified 2026-07-16): kimi, glm, inkling,
+# llama, qwen, gpt-oss, gemma.
+_OSS_MODEL_FAMILIES = (
+    "kimi-",
+    "glm-",
+    "inkling",
+    "llama-",
+    "qwen",
+    "gpt-oss-",
+    "gemma-",
+)
+
+# Non-chat services that can share a family substring with a chat model (e.g.
+# `qwen3-embedding-*` matches the `qwen` family) but can't back a chat agent.
+# The model-services API doesn't expose api_types, so exclude by name.
+_OSS_NON_CHAT_SUBSTRINGS = ("embedding", "embed", "rerank")
+
+
+def _is_oss_chat_model(model_id: str) -> bool:
+    """True if the id matches an OSS chat family and isn't a non-chat service."""
+    if any(bad in model_id for bad in _OSS_NON_CHAT_SUBSTRINGS):
+        return False
+    return any(family in model_id for family in _OSS_MODEL_FAMILIES)
+
 
 # Per-family token limits (context window + max output tokens). These are a
 # property of the model + its `/ai-gateway/mlflow/v1` route (the gateway rejects
@@ -1107,12 +1130,34 @@ def build_auth_shell_command(
 # config dialect. Both fields are provided because agents like OpenCode require
 # context and output together. Keyed by family substring; add an entry to bound
 # a new model.
+#
+# Output caps probed from the gateway 2026-07-16 (it 400s with "max_tokens (N)
+# cannot exceed <cap>"); context windows from each model's docs/description
+# (conservative when unstated). If the gateway raises a cap or ships a new
+# model, update this table.
 _MODEL_TOKEN_LIMITS: dict[str, dict[str, int]] = {
-    # GLM-4.6: 200k context, but the gateway caps output well below the model's
-    # native 128k — pin 25k so requests aren't rejected.
-    "glm": {"context": 200_000, "output": 25_000},
+    "glm": {"context": 1_000_000, "output": 65_536},
+    "inkling": {"context": 128_000, "output": 65_536},
+    "kimi": {"context": 128_000, "output": 65_536},
+    "gpt-oss": {"context": 128_000, "output": 25_000},
+    "qwen35": {"context": 128_000, "output": 25_000},
+    "qwen3-next": {"context": 128_000, "output": 10_000},
+    "llama-4": {"context": 128_000, "output": 8_192},
+    "llama-3": {"context": 128_000, "output": 8_192},
+    "gemma": {"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
+# (OpenCode's @ai-sdk/openai) need no flag.
+_OSS_REASONING_FAMILIES = ("glm", "inkling", "kimi", "gpt-oss", "qwen35")
+
+
+def model_is_reasoning(model_id: str) -> bool:
+    """True if the OSS model reports reasoning output (family-matched)."""
+    return any(family in model_id for family in _OSS_REASONING_FAMILIES)
+
 
 def model_token_limits(model_id: str) -> dict[str, int] | None:
     """Return ``{"context": ..., "output": ...}`` limits for ``model_id``, or None.
@@ -1252,7 +1297,7 @@ def discover_model_services(
     codex_models = [m for m in ids if "gpt-" in m]
     gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key)
 
-    oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)]
+    oss_models = [m for m in ids if _is_oss_chat_model(m)]
 
     if not (claude_models or codex_models or gemini_models or oss_models):
         sample = ", ".join(ids[:5])
@@ -2133,6 +2178,7 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]:
         "claude": build_tool_base_url("claude", workspace),
         "openai": build_tool_base_url("codex", workspace),
         "gemini": build_tool_base_url("gemini", workspace) + "/v1beta",
+        "oss": f"{workspace}/ai-gateway/mlflow/v1",
     }
 
 
diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py
index 0960c64..3cc7348 100644
--- a/tests/test_agent_opencode.py
+++ b/tests/test_agent_opencode.py
@@ -98,15 +98,24 @@ def test_glm_gets_token_limits(self):
         overlay, _ = opencode.render_overlay("system.ai.glm-5-2", "tok", _base_urls(), models)
         glm = overlay["provider"]["databricks-oss"]["models"]["system.ai.glm-5-2"]
         # OpenCode's schema requires both context and output on `limit`.
-        assert glm["limit"] == {"context": 200000, "output": 25000}
+        # Probed 2026-07-16: glm-5-2 is 1M context / 65536 output.
+        assert glm["limit"] == {"context": 1_000_000, "output": 65_536}
 
-    def test_non_glm_oss_model_has_no_output_cap(self):
+    def test_kimi_gets_token_limits(self):
+        # kimi is now a capped OSS family (128k context / 65536 output).
         models = {"oss": ["system.ai.kimi-k2-7-code"]}
         overlay, _ = opencode.render_overlay(
             "system.ai.kimi-k2-7-code", "tok", _base_urls(), models
         )
         kimi = overlay["provider"]["databricks-oss"]["models"]["system.ai.kimi-k2-7-code"]
-        assert "limit" not in kimi
+        assert kimi["limit"] == {"context": 128_000, "output": 65_536}
+
+    def test_uncapped_oss_model_has_no_limit(self):
+        # A model outside the limits table gets no `limit` (client default).
+        models = {"oss": ["system.ai.mystery-7b"]}
+        overlay, _ = opencode.render_overlay("system.ai.mystery-7b", "tok", _base_urls(), models)
+        entry = overlay["provider"]["databricks-oss"]["models"]["system.ai.mystery-7b"]
+        assert "limit" not in entry
 
     def test_token_in_api_key(self):
         models = {"anthropic": ["claude-sonnet"]}
diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py
index 0afc5fb..de002ba 100644
--- a/tests/test_agent_pi.py
+++ b/tests/test_agent_pi.py
@@ -17,6 +17,7 @@ def _base_urls() -> dict[str, str]:
         "claude": f"{WS}/ai-gateway/anthropic",
         "openai": f"{WS}/ai-gateway/codex/v1",
         "gemini": f"{WS}/ai-gateway/gemini/v1beta",
+        "oss": f"{WS}/ai-gateway/mlflow/v1",
     }
 
 
@@ -26,6 +27,7 @@ def _empty() -> dict:
         "claude_models": {},
         "codex_models": [],
         "gemini_models": [],
+        "oss_models": [],
     }
 
 
@@ -39,6 +41,7 @@ def _overlay(model: str, token: str = "tok", **kwargs):
         bundle["claude_models"],
         bundle["codex_models"],
         bundle["gemini_models"],
+        bundle["oss_models"],
     )
 
 
@@ -81,17 +84,59 @@ def test_gemini_provider_uses_google_generative_ai(self):
         assert provider["api"] == "google-generative-ai"
         assert provider["baseUrl"] == f"{WS}/ai-gateway/gemini/v1beta"
 
-    def test_all_three_providers_when_all_present(self):
+    def test_mlflow_provider_uses_openai_completions(self):
+        overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"])
+        provider = overlay["providers"]["databricks-mlflow"]
+        assert provider["api"] == "openai-completions"
+        assert provider["baseUrl"] == f"{WS}/ai-gateway/mlflow/v1"
+        assert provider["compat"] == {"supportsStore": False, "supportsStrictMode": False}
+
+    def test_no_mlflow_provider_when_no_oss_models(self):
+        overlay, _ = _overlay("gpt-5", codex_models=["gpt-5"])
+        assert "databricks-mlflow" not in overlay.get("providers", {})
+
+    def test_all_four_providers_when_all_present(self):
         overlay, _ = _overlay(
             "claude-sonnet",
             claude_models={"sonnet": "claude-sonnet"},
             codex_models=["gpt-5"],
             gemini_models=["gemini-2"],
+            oss_models=["system.ai.glm-5-2"],
         )
         assert set(overlay["providers"].keys()) == {
             "databricks-claude",
             "databricks-openai",
             "databricks-gemini",
+            "databricks-mlflow",
+        }
+
+
+class TestRenderOverlayOssEnrichment:
+    """OSS mlflow model entries carry reasoning + contextWindow + maxTokens
+    from the shared databricks.model_token_limits / model_is_reasoning tables."""
+
+    def test_reasoning_model_enriched(self):
+        overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"])
+        entry = overlay["providers"]["databricks-mlflow"]["models"][0]
+        assert entry["id"] == "system.ai.glm-5-2"
+        assert entry["reasoning"] is True
+        assert entry["contextWindow"] == 1_000_000
+        assert entry["maxTokens"] == 65_536
+
+    def test_non_reasoning_model_omits_reasoning(self):
+        overlay, _ = _overlay(
+            "system.ai.llama-4-maverick", oss_models=["system.ai.llama-4-maverick"]
+        )
+        entry = overlay["providers"]["databricks-mlflow"]["models"][0]
+        assert "reasoning" not in entry
+        assert entry["contextWindow"] == 128_000
+        assert entry["maxTokens"] == 8_192
+
+    def test_unknown_oss_model_bare(self):
+        # No limits/reasoning table entry -> only id, client keeps defaults.
+        overlay, _ = _overlay("system.ai.mystery-7b", oss_models=["system.ai.mystery-7b"])
+        assert overlay["providers"]["databricks-mlflow"]["models"][0] == {
+            "id": "system.ai.mystery-7b"
         }
 
 
@@ -193,6 +238,10 @@ def test_prefixes_gemini_model(self):
         overlay, _ = _overlay("gemini-2", gemini_models=["gemini-2"])
         assert overlay["model"] == "databricks-gemini/gemini-2"
 
+    def test_prefixes_oss_model(self):
+        overlay, _ = _overlay("system.ai.glm-5-2", oss_models=["system.ai.glm-5-2"])
+        assert overlay["model"] == "databricks-mlflow/system.ai.glm-5-2"
+
     def test_preserves_already_prefixed_model(self):
         overlay, _ = _overlay(
             "databricks-claude/claude-sonnet",
@@ -228,6 +277,15 @@ def test_falls_back_to_gemini(self):
         state = {"claude_models": {}, "codex_models": [], "gemini_models": ["gemini-2"]}
         assert pi.default_model(state) == "gemini-2"
 
+    def test_falls_back_to_oss_last(self):
+        state = {
+            "claude_models": {},
+            "codex_models": [],
+            "gemini_models": [],
+            "oss_models": ["system.ai.glm-5-2"],
+        }
+        assert pi.default_model(state) == "system.ai.glm-5-2"
+
     def test_returns_none_when_empty(self):
         assert pi.default_model({}) is None
         assert (
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index f59aa1d..ef5e13c 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -142,19 +142,48 @@ def _model_service(model_id: str) -> dict:
 
 class TestModelTokenLimits:
     def test_glm_is_capped(self):
+        # Probed 2026-07-16: glm-5-2 accepts 1M context / 65536 output.
         assert db_mod.model_token_limits("system.ai.glm-5-2") == {
-            "context": 200_000,
-            "output": 25_000,
+            "context": 1_000_000,
+            "output": 65_536,
         }
 
     def test_glm_matches_any_version(self):
         assert db_mod.model_token_limits("system.ai.glm-4-6-flash") == {
-            "context": 200_000,
-            "output": 25_000,
+            "context": 1_000_000,
+            "output": 65_536,
         }
 
+    def test_kimi_is_capped(self):
+        assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") == {
+            "context": 128_000,
+            "output": 65_536,
+        }
+
+    def test_llama_lower_output_cap(self):
+        assert db_mod.model_token_limits("system.ai.llama-4-maverick")["output"] == 8_192
+
+    def test_qwen_versions_distinct_caps(self):
+        # qwen35 and qwen3-next have different output ceilings; keyed precisely.
+        assert db_mod.model_token_limits("system.ai.qwen35-122b-a10b")["output"] == 25_000
+        assert (
+            db_mod.model_token_limits("system.ai.qwen3-next-80b-a3b-instruct")["output"] == 10_000
+        )
+
     def test_uncapped_model_returns_none(self):
-        assert db_mod.model_token_limits("system.ai.kimi-k2-7-code") is None
+        assert db_mod.model_token_limits("system.ai.deepseek-v3") is None
+
+
+class TestModelIsReasoning:
+    def test_reasoning_families(self):
+        assert db_mod.model_is_reasoning("system.ai.glm-5-2") is True
+        assert db_mod.model_is_reasoning("system.ai.inkling") is True
+        assert db_mod.model_is_reasoning("system.ai.qwen35-122b-a10b") is True
+
+    def test_non_reasoning_families(self):
+        assert db_mod.model_is_reasoning("system.ai.llama-4-maverick") is False
+        assert db_mod.model_is_reasoning("system.ai.qwen3-next-80b-a3b-instruct") is False
+        assert db_mod.model_is_reasoning("system.ai.gemma-3-12b") is False
 
 
 class TestDiscoverModelServices:
@@ -187,16 +216,21 @@ def test_buckets_families_by_name(self, monkeypatch):
         assert codex == ["system.ai.gpt-5"]
         # Gemini ordered newest-first via the shared sort key.
         assert gemini[0] == "system.ai.gemini-3-5-flash"
-        # kimi and glm are the allowlisted OSS families; llama is not.
-        assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"]
+        # Full OSS cohort now allowlisted: kimi, glm, AND llama.
+        assert oss == [
+            "system.ai.glm-5-2",
+            "system.ai.kimi-k2-7-code",
+            "system.ai.llama-4-maverick",
+        ]
 
     def test_oss_allowlist_drops_unsupported_families(self, monkeypatch):
-        # Only kimi/glm are allowlisted; other families are dropped.
+        # The cohort covers kimi/glm/qwen/llama/gpt-oss/gemma/inkling; families
+        # outside it (deepseek) and non-chat services (embeddings/rerankers) drop.
         payload = {
             "model_services": [
                 _model_service("system.ai.glm-5-2"),
-                _model_service("system.ai.kimi-k2-7-code"),
-                _model_service("system.ai.qwen-3-coder"),
+                _model_service("system.ai.qwen35-122b-a10b"),
+                _model_service("system.ai.inkling"),
                 _model_service("system.ai.deepseek-v3"),
                 _model_service("system.ai.gte-large-embed"),
                 _model_service("system.ai.bge-reranker-v2"),
@@ -210,7 +244,30 @@ def test_oss_allowlist_drops_unsupported_families(self, monkeypatch):
 
         assert reason is None
         assert (claude, codex, gemini) == ({}, [], [])
-        assert oss == ["system.ai.glm-5-2", "system.ai.kimi-k2-7-code"]
+        # deepseek/embeddings/rerankers dropped; cohort members kept.
+        assert "system.ai.glm-5-2" in oss
+        assert "system.ai.qwen35-122b-a10b" in oss
+        assert "system.ai.inkling" in oss
+        assert "system.ai.deepseek-v3" not in oss
+        assert "system.ai.gte-large-embed" not in oss
+
+    def test_embedding_model_sharing_family_substring_excluded(self, monkeypatch):
+        # qwen3-embedding matches the `qwen` family but is an embeddings model,
+        # not a chat model — it must not be bucketed as OSS chat.
+        payload = {
+            "model_services": [
+                _model_service("system.ai.qwen35-122b-a10b"),
+                _model_service("system.ai.qwen3-embedding-0-6b"),
+            ]
+        }
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None)
+        )
+
+        _, _, _, oss, _ = db_mod.discover_model_services(WS, "token")
+
+        assert "system.ai.qwen35-122b-a10b" in oss
+        assert "system.ai.qwen3-embedding-0-6b" not in oss
 
     def test_paginates_via_next_page_token(self, monkeypatch):
         pages = {
@@ -248,7 +305,8 @@ def test_http_failure_returns_reason(self, monkeypatch):
         assert reason == "HTTP 500 Server Error"
 
     def test_no_matching_families_reports_sample(self, monkeypatch):
-        payload = {"model_services": [_model_service("system.ai.llama-4-maverick")]}
+        # deepseek is outside every claude/gpt/gemini/oss family bucket.
+        payload = {"model_services": [_model_service("system.ai.deepseek-v3")]}
         monkeypatch.setattr(
             db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None)
         )
@@ -256,7 +314,7 @@ def test_no_matching_families_reports_sample(self, monkeypatch):
         claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token")
 
         assert (claude, codex, gemini, oss) == ({}, [], [], [])
-        assert reason is not None and "llama-4-maverick" in reason
+        assert reason is not None and "deepseek-v3" in reason
 
     def test_ignores_non_system_ai_schemas(self, monkeypatch):
         # The metastore listing returns services from every schema; only
diff --git a/tests/test_mlflow_proxy.py b/tests/test_mlflow_proxy.py
new file mode 100644
index 0000000..f71c0da
--- /dev/null
+++ b/tests/test_mlflow_proxy.py
@@ -0,0 +1,122 @@
+"""Tests for agents/_mlflow_proxy.py — the SSE finish_reason repair proxy."""
+
+from __future__ import annotations
+
+import threading
+import urllib.request
+from http.server import BaseHTTPRequestHandler, HTTPServer
+
+from ucode.agents import _mlflow_proxy
+
+# An inkling-style stream: content deltas then [DONE], NO finish_reason anywhere.
+_STREAM_NO_FINISH = (
+    b'data: {"id":"c1","choices":[{"delta":{"content":"ok"},"index":0}]}\n\ndata: [DONE]\n\n'
+)
+# A well-behaved stream: terminal chunk carries finish_reason.
+_STREAM_WITH_FINISH = (
+    b'data: {"id":"c2","choices":[{"delta":{"content":"ok"},"index":0}]}\n\n'
+    b'data: {"id":"c2","choices":[{"delta":{},"finish_reason":"stop","index":0}]}\n\n'
+    b"data: [DONE]\n\n"
+)
+# 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'
+
+
+def _make_fake_gateway(payload: bytes) -> 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_header("Content-Length", str(len(payload)))
+            self.end_headers()
+            self.wfile.write(payload)
+
+        def log_message(self, format: str, *args):
+            pass
+
+    server = HTTPServer(("127.0.0.1", 0), _Fake)
+    threading.Thread(target=server.serve_forever, daemon=True).start()
+    return f"http://127.0.0.1:{server.server_address[1]}", server
+
+
+def _post_through_proxy(upstream: str) -> str:
+    base, stop = _mlflow_proxy.start(upstream)
+    try:
+        req = urllib.request.Request(
+            f"{base}/ai-gateway/mlflow/v1/chat/completions",
+            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")
+    finally:
+        stop.set()
+
+
+class TestFinishReasonInjection:
+    def test_injects_finish_reason_when_absent(self):
+        upstream, server = _make_fake_gateway(_STREAM_NO_FINISH)
+        try:
+            out = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        assert '"finish_reason": "stop"' in out or '"finish_reason":"stop"' in out
+        # The injected chunk must come before the terminal [DONE].
+        assert out.index("finish_reason") < out.index("[DONE]")
+        assert out.rstrip().endswith("[DONE]")
+
+    def test_passes_through_when_finish_reason_present(self):
+        upstream, server = _make_fake_gateway(_STREAM_WITH_FINISH)
+        try:
+            out = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        # Exactly one finish_reason — the proxy did not add a second.
+        assert out.count("finish_reason") == 1
+
+    def test_forwards_content_unchanged(self):
+        upstream, server = _make_fake_gateway(_STREAM_NO_FINISH)
+        try:
+            out = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        assert '"content":"ok"' in out
+
+    def test_repairs_abrupt_end_with_finish_and_done(self):
+        # Upstream stopped after content with no finish_reason and no [DONE]
+        # (mid-stream 429/drop). Proxy must synthesize BOTH so a strict client
+        # doesn't error on a truncated stream.
+        upstream, server = _make_fake_gateway(_STREAM_ABRUPT_END)
+        try:
+            out = _post_through_proxy(upstream)
+        finally:
+            server.shutdown()
+        assert "finish_reason" in out
+        assert out.rstrip().endswith("[DONE]")
+        assert '"content":"ok"' in out
+
+
+class TestFinishChunk:
+    def test_shape(self):
+        import json
+
+        chunk = json.loads(_mlflow_proxy._finish_chunk("abc"))
+        assert chunk["id"] == "abc"
+        assert chunk["choices"][0]["finish_reason"] == "stop"
+        assert chunk["object"] == "chat.completion.chunk"
+
+
+class TestStart:
+    def test_returns_local_base_url_and_stop_event(self):
+        upstream, server = _make_fake_gateway(_STREAM_NO_FINISH)
+        try:
+            base, stop = _mlflow_proxy.start(upstream)
+            assert base.startswith("http://127.0.0.1:")
+            assert isinstance(stop, threading.Event)
+            stop.set()
+        finally:
+            server.shutdown()

From 3c9503e618333896bacf51312f8107020ebc9708 Mon Sep 17 00:00:00 2001
From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com>
Date: Thu, 16 Jul 2026 23:38:42 +1000
Subject: [PATCH 2/4] chore: bump version to 0.2.0

Pi databricks-mlflow OSS provider + broadened OSS cohort with reasoning/limits.
---
 pyproject.toml | 2 +-
 uv.lock        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 86bda87..b7db686 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ucode"
-version = "0.1.0"
+version = "0.2.0"
 description = "Bootstrap Codex against a Databricks workspace with minimal setup."
 readme = "README.md"
 requires-python = ">=3.12"
diff --git a/uv.lock b/uv.lock
index 1f599a0..512d256 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2927,7 +2927,7 @@ wheels = [
 
 [[package]]
 name = "ucode"
-version = "0.1.0"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "databricks-sql-connector" },

From acd2d06b11dcbfaf23399bb3770be33ff48b7e01 Mon Sep 17 00:00:00 2001
From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com>
Date: Fri, 17 Jul 2026 00:10:21 +1000
Subject: [PATCH 3/4] fix: exclude gpt-oss from codex bucket in model-services
 discovery
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

gpt-oss-* contains 'gpt-' so it matched the codex bucket, but it's a
chat-completions-only OSS model (served via /ai-gateway/mlflow/v1), not an
openai-responses model. Broadening _OSS_MODEL_FAMILIES to include gpt-oss put
it in BOTH codex_models and oss_models, so it was offered under Pi's
databricks-openai provider too — where the openai-responses route 400s for it.
Exclude gpt-oss from the codex bucket so it's OSS-only.
---
 src/ucode/databricks.py  |  5 ++++-
 tests/test_databricks.py | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index df3e27f..038321e 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -1294,7 +1294,10 @@ def discover_model_services(
         if candidates:
             claude_models[family] = candidates[0]
 
-    codex_models = [m for m in ids if "gpt-" in m]
+    # `gpt-oss-*` also contains "gpt-" but is a chat-completions-only OSS model
+    # (served via /ai-gateway/mlflow/v1), NOT an openai-responses codex model —
+    # exclude it here so it isn't offered under the codex provider (which 400s).
+    codex_models = [m for m in ids if "gpt-" in m and "gpt-oss" not in m]
     gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key)
 
     oss_models = [m for m in ids if _is_oss_chat_model(m)]
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index ef5e13c..6417faa 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -269,6 +269,26 @@ def test_embedding_model_sharing_family_substring_excluded(self, monkeypatch):
         assert "system.ai.qwen35-122b-a10b" in oss
         assert "system.ai.qwen3-embedding-0-6b" not in oss
 
+    def test_gpt_oss_is_oss_not_codex(self, monkeypatch):
+        # gpt-oss-* contains "gpt-" but is chat-completions-only, not an
+        # openai-responses codex model. It must bucket into oss ONLY, never
+        # codex (whose openai-responses route would 400 for it).
+        payload = {
+            "model_services": [
+                _model_service("system.ai.gpt-5"),
+                _model_service("system.ai.gpt-oss-120b"),
+            ]
+        }
+        monkeypatch.setattr(
+            db_mod, "_http_get_json", lambda url, token, timeout=10: (payload, None)
+        )
+
+        _, codex, _, oss, _ = db_mod.discover_model_services(WS, "token")
+
+        assert codex == ["system.ai.gpt-5"]
+        assert "system.ai.gpt-oss-120b" in oss
+        assert "system.ai.gpt-oss-120b" not in codex
+
     def test_paginates_via_next_page_token(self, monkeypatch):
         pages = {
             None: {

From ededa9c0ebe1042c114b0b6bff6c5700ef09a355 Mon Sep 17 00:00:00 2001
From: dgokeeffe <17697537+dgokeeffe@users.noreply.github.com>
Date: Fri, 17 Jul 2026 00:46:23 +1000
Subject: [PATCH 4/4] 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.
---
 src/ucode/agents/_mlflow_proxy.py | 107 ++++++++++++++++++++++--------
 src/ucode/databricks.py           |  18 ++++-
 tests/test_databricks.py          |  15 ++++-
 tests/test_mlflow_proxy.py        |  96 +++++++++++++++++++++++----
 4 files changed, 193 insertions(+), 43 deletions(-)

diff --git a/src/ucode/agents/_mlflow_proxy.py b/src/ucode/agents/_mlflow_proxy.py
index 70e26ed..80e763a 100644
--- a/src/ucode/agents/_mlflow_proxy.py
+++ b/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,14 +31,26 @@
 # 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:
@@ -47,13 +58,29 @@ def do_POST(self) -> None:  # noqa: N802 (BaseHTTPRequestHandler API)
             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,11 +95,11 @@ 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"):
@@ -80,7 +107,7 @@ def _relay_stream(self, upstream_resp) -> None:
                         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
@@ -88,19 +115,31 @@ def _relay_stream(self, upstream_resp) -> None:
                         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(
         {
diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index 038321e..4d89201 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -1147,6 +1147,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
@@ -1162,11 +1172,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
 
 
diff --git a/tests/test_databricks.py b/tests/test_databricks.py
index 6417faa..a7c6e69 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -170,9 +170,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):
diff --git a/tests/test_mlflow_proxy.py b/tests/test_mlflow_proxy.py
index f71c0da..5aeb16a 100644
--- a/tests/test_mlflow_proxy.py
+++ b/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"