feat(opencode): Amazon Bedrock Model Provider Service support

src/ucode/agents/__init__.py

@@ -350,6 +350,7 @@ def configure_tool(
     relayed: bool = False,
     route_root_model: str | None = None,
     custom_model: str | None = None,
+    bedrock_targets: list[str] | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
@@ -370,16 +371,26 @@ def configure_tool(
             custom_model=custom_model,
         )
     else:
-        # provider routing is claude/codex-only; every other tool needs a model.
-        if not model:
+        # provider routing is claude/codex-only; every other tool needs a model —
+        # except pi with a Bedrock provider, where targets replace the model list.
+        if not model and not (tool in ("pi", "opencode") and provider and bedrock_targets):
             raise RuntimeError(f"A {tool} model must be selected before configuration.")
         if tool == "gemini":
+            assert model is not None
             result = gemini.write_tool_config(state, model)
         elif tool == "copilot":
+            assert model is not None
             result = copilot.write_tool_config(state, model)
         elif tool == "pi":
-            result = pi.write_tool_config(state, model)
+            result = pi.write_tool_config(
+                state, model, provider=provider, bedrock_targets=bedrock_targets
+            )
+        elif tool == "opencode":
+            result = opencode.write_tool_config(
+                state, model, provider=provider, bedrock_targets=bedrock_targets
+            )
         else:
+            assert model is not None
             result = opencode.write_tool_config(state, model)
     # gemini/opencode/copilot/pi return (state, token); codex/claude return state
     if isinstance(result, tuple):

src/ucode/agents/opencode.py

@@ -41,12 +41,15 @@
     ["provider", "databricks-anthropic"],
     ["provider", "databricks-google"],
     ["provider", "databricks-oss"],
+    ["provider", "databricks-bedrock"],
 ]
 
 
 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-oss/", "databricks-bedrock/")
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
@@ -79,10 +82,13 @@ def _oss_model_overlay(model: str, ua_header: dict[str, str]) -> dict:
 
 
 def render_overlay(
-    model: str,
+    model: str | None,
     token: str,
     opencode_base_urls: dict[str, str],
     opencode_models: dict[str, list[str]],
+    *,
+    provider: str | None = None,
+    bedrock_targets: list[str] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for opencode.json."""
     auth_headers = {"Authorization": f"Bearer {token}"}
@@ -100,6 +106,23 @@ def render_overlay(
 
     providers: dict = {}
     keys: list[list[str]] = [["model"]]
+    if provider and bedrock_targets:
+        # Bedrock routes through Databricks AI Gateway using bearer auth only
+        # (no AWS SigV4, no region). MPS and UA headers must be per-model because
+        # OpenCode clobbers provider-level headers in session/llm.ts.
+        bedrock_model_header = {
+            "User-Agent": ua_header["User-Agent"],
+            "Databricks-Model-Provider-Service": provider,
+        }
+        providers["databricks-bedrock"] = {
+            "npm": "@ai-sdk/amazon-bedrock",
+            "options": {
+                "baseURL": opencode_base_urls["bedrock"],
+                "apiKey": token,
+            },
+            "models": {t: {"headers": bedrock_model_header} for t in bedrock_targets},
+        }
+        keys.append(["provider", "databricks-bedrock"])
     if anthropic_models:
         # @ai-sdk/anthropic injects `eager_input_streaming: true` on tool defs;
         # the Databricks gateway's strict validator rejects it. opencode's
@@ -143,17 +166,24 @@ def render_overlay(
         }
         keys.append(["provider", "databricks-oss"])
 
-    overlay: dict = {"model": _resolve_model_selector(model, opencode_models)}
+    if provider and bedrock_targets:
+        model_selector = f"databricks-bedrock/{bedrock_targets[0]}"
+    else:
+        assert model is not None
+        model_selector = _resolve_model_selector(model, opencode_models)
+    overlay: dict = {"model": model_selector}
     if providers:
         overlay["provider"] = providers
     return overlay, keys
 
 
 def write_tool_config(
     state: dict,
-    model: str,
+    model: str | None,
     token: str | None = None,
     *,
+    provider: str | None = None,
+    bedrock_targets: list[str] | None = None,
     force_refresh: bool = False,
 ) -> tuple[dict, str]:
     backup_existing_file(OPENCODE_CONFIG_PATH, OPENCODE_BACKUP_PATH)
@@ -169,12 +199,15 @@ def write_tool_config(
         token,
         opencode_base_urls,
         state.get("opencode_models") or {},
+        provider=provider,
+        bedrock_targets=bedrock_targets,
     )
     existing = read_json_safe(OPENCODE_CONFIG_PATH)
     providers = existing.get("provider")
     if isinstance(providers, dict):
         for stale in (
             "databricks-anthropic",
+            "databricks-bedrock",
             "databricks-google",
             "databricks-openai",
             "databricks-oss",
@@ -237,6 +270,31 @@ def default_model(state: dict) -> str | None:
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
+    # Preserve an existing databricks-bedrock provider block written by a
+    # --provider launch, so token refresh does not silently drop it. The MPS
+    # name lives in each model entry's headers (per-model, not provider-level).
+    existing = read_json_safe(OPENCODE_CONFIG_PATH)
+    bedrock = (existing.get("provider") or {}).get("databricks-bedrock")
+    if isinstance(bedrock, dict):
+        models_dict = bedrock.get("models") or {}
+        saved_targets = list(models_dict.keys()) if models_dict else None
+        saved_provider: str | None = None
+        for entry in models_dict.values():
+            if isinstance(entry, dict):
+                saved_provider = (entry.get("headers") or {}).get(
+                    "Databricks-Model-Provider-Service"
+                )
+                if saved_provider:
+                    break
+        if saved_targets and saved_provider:
+            _, token = write_tool_config(
+                state,
+                None,
+                force_refresh=force_refresh,
+                provider=saved_provider,
+                bedrock_targets=saved_targets,
+            )
+            return token
     model = default_model(state)
     if not model:
         raise RuntimeError("No OpenCode model is configured.")

src/ucode/agents/pi.py

@@ -69,6 +69,7 @@
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    "databricks-bedrock",
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
@@ -98,12 +99,15 @@ def _resolve_model_selector(
 
 
 def render_overlay(
-    model: str,
+    model: str | None,
     token: str,
     pi_base_urls: dict[str, str],
     claude_models: dict[str, str],
     codex_models: list[str],
     gemini_models: list[str],
+    *,
+    provider: str | None = None,
+    bedrock_targets: list[str] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Pi's private agent config."""
     providers: dict = {}
@@ -147,20 +151,43 @@ def render_overlay(
             "models": [{"id": m} for m in gemini_models],
         }
         keys.append(["providers", "databricks-gemini"])
-    overlay: dict = {
-        "model": _resolve_model_selector(model, claude_models, codex_models, gemini_models),
-    }
+    if provider and bedrock_targets:
+        providers["databricks-bedrock"] = {
+            "baseUrl": pi_base_urls.get(
+                "bedrock", f"{pi_base_urls['claude'].rsplit('/ai-gateway', 1)[0]}/ai-gateway"
+            ),
+            "api": "bedrock-converse-stream",
+            "apiKey": token,
+            "authHeader": True,
+            # Pi's bedrock-converse-stream client (AWS SDK style) sets its own
+            # User-Agent; adding ours produces two `user-agent` values and the
+            # gateway rejects the request ("Header field ... must only have a
+            # single value"). Send only the MPS selector header here.
+            "headers": {"Databricks-Model-Provider-Service": provider},
+            "models": [{"id": t} for t in bedrock_targets],
+        }
+        keys.append(["providers", "databricks-bedrock"])
+    resolved = _resolve_model_selector(model or "", claude_models, codex_models, gemini_models)
+    # Bedrock model IDs contain no `/` (e.g. `anthropic.claude-3-haiku-20240307-v1:0`), so
+    # _resolve_model_selector returns them unprefixed. _write_settings splits on `/` to get
+    # provider/model — without the prefix it gets an empty model_id and skips defaultProvider.
+    # Always force the `databricks-bedrock/` prefix when the Bedrock provider is active.
+    if "databricks-bedrock" in providers and bedrock_targets:
+        resolved = f"databricks-bedrock/{bedrock_targets[0]}"
+    overlay: dict = {"model": resolved}
     if providers:
         overlay["providers"] = providers
     return overlay, keys
 
 
 def write_tool_config(
     state: dict,
-    model: str,
+    model: str | None,
     token: str | None = None,
     *,
     force_refresh: bool = False,
+    provider: str | None = None,
+    bedrock_targets: list[str] | None = None,
 ) -> tuple[dict, str]:
     backup_existing_file(PI_CONFIG_PATH, PI_BACKUP_PATH)
     if token is None:
@@ -181,6 +208,8 @@ def write_tool_config(
         claude_models,
         codex_models,
         gemini_models,
+        provider=provider,
+        bedrock_targets=bedrock_targets,
     )
     existing = read_json_safe(PI_CONFIG_PATH)
     providers = existing.get("providers")
@@ -259,6 +288,33 @@ def default_model(state: dict) -> str | None:
 
 
 def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
+    # Preserve a Bedrock provider block across token refreshes. The block is
+    # self-describing: its MPS header + model ids are enough to re-render it,
+    # so a refresh keeps routing through Bedrock instead of dropping to a
+    # system-hosted model. When the config has no Bedrock block (a non-Bedrock
+    # session, or after a non-Bedrock reconfigure overwrote it), fall through
+    # to the normal path.
+    existing = read_json_safe(PI_CONFIG_PATH)
+    bedrock = (existing.get("providers") or {}).get("databricks-bedrock")
+    provider: str | None = None
+    bedrock_targets: list[str] | None = None
+    if isinstance(bedrock, dict):
+        headers = bedrock.get("headers") or {}
+        provider = headers.get("Databricks-Model-Provider-Service")
+        bedrock_targets = [
+            m["id"]
+            for m in (bedrock.get("models") or [])
+            if isinstance(m, dict) and isinstance(m.get("id"), str)
+        ] or None
+    if provider and bedrock_targets:
+        _, token = write_tool_config(
+            state,
+            bedrock_targets[0],
+            force_refresh=force_refresh,
+            provider=provider,
+            bedrock_targets=bedrock_targets,
+        )
+        return token
     model = default_model(state)
     if not model:
         raise RuntimeError("No Pi model is available on this workspace.")