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).

Author: dgokeeffe

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

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,16 +112,37 @@ 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,
     pi_base_urls: dict[str, str],
     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)
 

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