feat: expose MLflow chat-completions-only foundation models in Pi and OpenCode

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"

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/opencode.py

@@ -1,4 +1,11 @@
-"""OpenCode agent: writes opencode.json with two Databricks-backed providers."""
+"""OpenCode agent: writes opencode.json with Databricks-backed providers.
+
+Providers: `databricks-anthropic` (@ai-sdk/anthropic), `databricks-google`
+(@ai-sdk/google), and `databricks-oss` (@ai-sdk/openai, chat-completions mode)
+for mlflow-only foundation models (inkling, glm, llama, qwen, ...). Unlike Pi,
+opencode's @ai-sdk/openai tolerates the `finish_reason` the gateway omits for
+inkling, so no SSE-repair proxy is needed here.
+"""
 
 from __future__ import annotations
 
@@ -41,6 +48,7 @@
 PROVIDER_KEYS: list[list[str]] = [
     ["provider", "databricks-anthropic"],
     ["provider", "databricks-google"],
+    ["provider", "databricks-oss"],
 ]
 
 
@@ -50,7 +58,11 @@ def is_update_available() -> tuple[str, str] | None:
 
 def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -> str:
     """Return an OpenCode model selector in provider/model form when possible."""
-    if model.startswith("databricks-anthropic/") or model.startswith("databricks-google/"):
+    if (
+        model.startswith("databricks-anthropic/")
+        or model.startswith("databricks-google/")
+        or model.startswith("databricks-oss/")
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
@@ -61,14 +73,37 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -
     if model in gemini_models:
         return f"databricks-google/{model}"
 
+    oss_models = opencode_models.get("oss") or []
+    if model in oss_models:
+        return f"databricks-oss/{model}"
+
     return model
 
 
+def _oss_model_entry(ua_header: dict, spec: dict | None) -> dict:
+    """Build an opencode oss model entry. @ai-sdk/openai surfaces reasoning
+    automatically, so we only set `limit`: `context` from the gateway's stated
+    context window and `output` from its per-model max-output ceiling. Each is
+    included only when known; omit the whole `limit` if neither is."""
+    entry: dict = {"headers": ua_header}
+    if not spec:
+        return entry
+    limit = {}
+    if spec.get("context_window"):
+        limit["context"] = spec["context_window"]
+    if spec.get("max_tokens"):
+        limit["output"] = spec["max_tokens"]
+    if limit:
+        entry["limit"] = limit
+    return entry
+
+
 def render_overlay(
     model: str,
     token: str,
     opencode_base_urls: dict[str, str],
     opencode_models: dict[str, list[str]],
+    oss_specs: list[dict] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for opencode.json."""
     auth_headers = {"Authorization": f"Bearer {token}"}
@@ -82,6 +117,7 @@ def render_overlay(
 
     anthropic_models = opencode_models.get("anthropic") or []
     gemini_models = opencode_models.get("gemini") or []
+    oss_models = opencode_models.get("oss") or []
 
     providers: dict = {}
     keys: list[list[str]] = [["model"]]
@@ -116,6 +152,23 @@ def render_overlay(
             "models": {m: {"headers": ua_header} for m in gemini_models},
         }
         keys.append(["provider", "databricks-google"])
+    if oss_models:
+        # mlflow chat-completions-only foundation models (inkling, glm, llama,
+        # qwen, ...) via @ai-sdk/openai in legacy chat-completions mode — NOT
+        # the Responses API (no useResponsesApi), unlike the codex provider.
+        # Unlike Pi, opencode's @ai-sdk/openai tolerates the finish_reason the
+        # gateway omits for reasoning models (inkling), so no proxy is needed.
+        oss_specs_by_id = {s["id"]: s for s in (oss_specs or [])}
+        providers["databricks-oss"] = {
+            "npm": "@ai-sdk/openai",
+            "options": {
+                "baseURL": opencode_base_urls["oss"],
+                "apiKey": token,
+                "headers": auth_headers,
+            },
+            "models": {m: _oss_model_entry(ua_header, oss_specs_by_id.get(m)) for m in oss_models},
+        }
+        keys.append(["provider", "databricks-oss"])
 
     overlay: dict = {"model": _resolve_model_selector(model, opencode_models)}
     if providers:
@@ -141,11 +194,17 @@ def write_tool_config(
         token,
         opencode_base_urls,
         state.get("opencode_models") or {},
+        state.get("oss_model_specs") or [],
     )
     existing = read_json_safe(OPENCODE_CONFIG_PATH)
     providers = existing.get("provider")
     if isinstance(providers, dict):
-        for stale in ("databricks-anthropic", "databricks-google", "databricks-openai"):
+        for stale in (
+            "databricks-anthropic",
+            "databricks-google",
+            "databricks-openai",
+            "databricks-oss",
+        ):
             providers.pop(stale, None)
     merged = deep_merge_dict(existing, overlay)
     write_json_file(OPENCODE_CONFIG_PATH, merged)
@@ -195,7 +254,10 @@ def default_model(state: dict) -> str | None:
     if anthropic:
         return anthropic[0]
     gemini = opencode_models.get("gemini") or []
-    return gemini[0] if gemini else None
+    if gemini:
+        return gemini[0]
+    oss = opencode_models.get("oss") or []
+    return oss[0] if oss else None
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str: