fix(pi): extend forceAdaptiveThinking to Claude 5+ models

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "uv_build"
 
 [project]
 name = "ucode"
-version = "0.1.0"
+version = "0.3.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,201 @@
+"""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, 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.
+"""
+
+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"})
+
+# The proxy only exists to front the mlflow chat-completions route. Refuse any
+# other path so a co-located process can't turn the localhost relay into an
+# arbitrary authenticated workspace client (SSRF-to-workspace) using the token
+# we forward.
+_ALLOWED_PATH_PREFIX = "/ai-gateway/mlflow/"
+
+
+def _make_handler(upstream: str) -> type[BaseHTTPRequestHandler]:
+    class _Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def do_POST(self) -> None:  # noqa: N802 (BaseHTTPRequestHandler API)
+            # Read the body regardless so the socket is drained before we close.
+            length = int(self.headers.get("Content-Length", 0))
+            body = self.rfile.read(length)
+
+            if not self.path.startswith(_ALLOWED_PATH_PREFIX):
+                self._send_bytes(404, b'{"error":"path not allowed"}', "application/json")
+                return
+
+            req = urllib.request.Request(upstream + self.path, data=body, method="POST")
+            for key, value in self.headers.items():
+                if key.lower() not in _SKIP_REQUEST_HEADERS:
+                    req.add_header(key, value)
+            try:
+                upstream_resp = urllib.request.urlopen(req)  # noqa: S310 (trusted upstream)
+            except urllib.error.HTTPError as exc:
+                # 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")
+            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
+                    payload = _sse_data_payload(raw)
+                    if payload is not None and payload != b"[DONE]":
+                        saw_data = True
+                        try:
+                            obj = json.loads(payload)
+                            last_id = obj.get("id", last_id)
+                            choice = (obj.get("choices") or [{}])[0]
+                            if choice.get("finish_reason"):
+                                saw_finish = True
+                        except (ValueError, AttributeError, IndexError):
+                            pass
+                        self._write(raw)
+                    elif payload == b"[DONE]":
+                        if not saw_finish:
+                            self._write(b"data: " + _finish_chunk(last_id) + b"\n\n")
+                            saw_finish = True
+                        self._write(raw)
+                        saw_done = True
+                    else:
+                        self._write(raw)
+                # Clean or abnormal end after data but without a terminator:
+                # synthesize one so strict clients don't error on a truncated
+                # stream (a mid-stream 429 is the common trigger).
+                if saw_data and not saw_finish:
+                    self._write(b"data: " + _finish_chunk(last_id) + b"\n\n")
+                if saw_data and not saw_done:
+                    self._write(b"data: [DONE]\n\n")
+            except OSError:
+                # Upstream dropped, or the client disconnected mid-write. Either
+                # way the socket is gone; swallow so the daemon thread exits
+                # quietly rather than tracebacking.
+                pass
+
+        def _send_bytes(self, status: int, data: bytes, content_type: str) -> None:
+            self.close_connection = True
+            try:
+                self.send_response(status)
+                self.send_header("Content-Type", content_type)
+                self.send_header("Content-Length", str(len(data)))
+                self.send_header("Connection", "close")
+                self.end_headers()
+                self.wfile.write(data)
+                self.wfile.flush()
+            except OSError:
+                pass
+
+        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 _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(
+        {
+            "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/claude.py

@@ -251,8 +251,12 @@ def _maybe_add_1m_suffix(model: str) -> str:
     family, major_raw, minor_raw, _ = match.groups()
     major = int(major_raw)
     minor = int(minor_raw)
+    # 1M-token context window support: Opus from 4.6, Sonnet from 4.5 (Sonnet
+    # 4.5 gained the 1M beta window before Opus did). Without the `[1m]`
+    # suffix the gateway serves these models with the default 200k context, so
+    # gate each family at its own floor rather than a shared version.
     should_suffix = (family == "opus" and (major, minor) >= (4, 6)) or (
-        family == "sonnet" and (major, minor) >= (4, 6)
+        family == "sonnet" and (major, minor) >= (4, 5)
     )
     return f"{model}[1m]" if should_suffix else model
 

src/ucode/agents/opencode.py

@@ -42,6 +42,7 @@
 PROVIDER_KEYS: list[list[str]] = [
     ["provider", "databricks-anthropic"],
     ["provider", "databricks-google"],
+    ["provider", "databricks-openai"],
     ["provider", "databricks-oss"],
 ]
 
@@ -52,7 +53,14 @@ 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/", "databricks-google/", "databricks-oss/")):
+    if model.startswith(
+        (
+            "databricks-anthropic/",
+            "databricks-google/",
+            "databricks-openai/",
+            "databricks-oss/",
+        )
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
@@ -63,6 +71,10 @@ def _resolve_model_selector(model: str, opencode_models: dict[str, list[str]]) -
     if model in gemini_models:
         return f"databricks-google/{model}"
 
+    openai_models = opencode_models.get("openai") or []
+    if model in openai_models:
+        return f"databricks-openai/{model}"
+
     oss_models = opencode_models.get("oss") or []
     if model in oss_models:
         return f"databricks-oss/{model}"
@@ -102,6 +114,7 @@ def render_overlay(
 
     anthropic_models = opencode_models.get("anthropic") or []
     gemini_models = opencode_models.get("gemini") or []
+    openai_models = opencode_models.get("openai") or []
     oss_models = opencode_models.get("oss") or []
 
     providers: dict = {}
@@ -137,6 +150,28 @@ def render_overlay(
             "models": {m: {"headers": ua_header} for m in gemini_models},
         }
         keys.append(["provider", "databricks-google"])
+    if openai_models:
+        # @ai-sdk/openai supports both the Responses API and the legacy
+        # chat-completions API. Databricks GPT-5 / GPT-5.6 / Codex models are
+        # Responses-only on /ai-gateway/codex/v1, so the per-model flag
+        # `useResponsesApi: true` lives in models.<m>.options where opencode
+        # reads it (provider-level options is read by the SDK only).
+        providers["databricks-openai"] = {
+            "npm": "@ai-sdk/openai",
+            "options": {
+                "baseURL": opencode_base_urls["openai"],
+                "apiKey": token,
+                "headers": auth_headers,
+            },
+            "models": {
+                m: {
+                    "headers": ua_header,
+                    "options": {"useResponsesApi": True},
+                }
+                for m in openai_models
+            },
+        }
+        keys.append(["provider", "databricks-openai"])
     if oss_models:
         providers["databricks-oss"] = {
             "npm": "@ai-sdk/openai",
@@ -233,6 +268,9 @@ def default_model(state: dict) -> str | None:
     anthropic = opencode_models.get("anthropic") or []
     if anthropic:
         return anthropic[0]
+    openai = opencode_models.get("openai") or []
+    if openai:
+        return openai[0]
     gemini = opencode_models.get("gemini") or []
     if gemini:
         return gemini[0]