feat(opencode): Amazon Bedrock Model Provider Service support

src/ucode/agents/__init__.py

@@ -426,6 +426,7 @@ def configure_tool(
     route_root_model: str | None = None,
     custom_model: str | None = None,
     coding_agent_config_defaults: dict[str, str] | None = None,
+    bedrock_targets: list[str] | None = None,
 ) -> dict:
     result: dict | tuple[dict, str]
     if tool == "codex":
@@ -446,17 +447,27 @@ def configure_tool(
             coding_agent_config_defaults=coding_agent_config_defaults,
         )
     else:
-        # Every tool in this branch needs a model — including gemini under a provider,
-        # which still pins the service's target model in the URL.
-        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.
+        # gemini under a provider still pins the service's target model in the URL.
+        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, provider=provider)
         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

@@ -52,6 +52,7 @@
     ["provider", "databricks-anthropic"],
     ["provider", "databricks-google"],
     ["provider", "databricks-oss"],
+    ["provider", "databricks-bedrock"],
 ]
 
 _AUTH_PLUGIN_TEMPLATE = """// Generated by ucode. Keep Databricks auth fresh for model requests.
@@ -60,6 +61,7 @@
 
 const DATABRICKS_PROVIDERS = new Set([
   "databricks-anthropic",
+  "databricks-bedrock",
   "databricks-google",
   "databricks-oss",
 ])
@@ -187,7 +189,9 @@ def write_auth_plugin(state: dict) -> 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-oss/", "databricks-bedrock/")
+    ):
         return model
 
     anthropic_models = opencode_models.get("anthropic") or []
@@ -220,10 +224,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}"}
@@ -241,6 +248,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
@@ -284,20 +308,31 @@ 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)
     if token is None:
-        token = get_databricks_token(state["workspace"], state.get("profile"))
+        token = get_databricks_token(
+            state["workspace"], state.get("profile"), force_refresh=force_refresh
+        )
     opencode_base_urls = state.get("base_urls", {}).get("opencode") or build_opencode_base_urls(
         state["workspace"]
     )
@@ -306,13 +341,16 @@ 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)
     write_auth_plugin(state)
     providers = existing.get("provider")
     if isinstance(providers, dict):
         for stale in (
             "databricks-anthropic",
+            "databricks-bedrock",
             "databricks-google",
             "databricks-openai",
             "databricks-oss",
@@ -374,11 +412,36 @@ def default_model(state: dict) -> str | None:
     return oss[0] if oss else None
 
 
-def _configure_launch(state: dict) -> str:
+def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> str:
+    # Preserve an existing databricks-bedrock provider block written by a
+    # --provider launch, so a relaunch 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.")
-    _, token = write_tool_config(state, model)
+    _, token = write_tool_config(state, model, force_refresh=force_refresh)
     return token
 
 
@@ -391,7 +454,7 @@ def build_runtime_env(token: str, state: dict | None = None) -> dict[str, str]:
 
 def launch(state: dict, tool_args: list[str], *, options: LaunchOptions) -> None:
     """Launch OpenCode with on-demand token refresh from its local plugin."""
-    token = _configure_launch(state)
+    token = _refresh_token_once(state)
     env = build_runtime_env(token, state)
 
     proc = subprocess.Popen([SPEC["binary"], *tool_args], env=env)

src/ucode/agents/pi.py

@@ -46,6 +46,7 @@
     build_pi_base_urls,
     classify_model_family,
     get_databricks_token,
+    model_token_limits,
 )
 from ucode.state import mark_tool_managed, save_state
 from ucode.telemetry import agent_version, ucode_version
@@ -71,6 +72,7 @@
     "databricks-claude",
     "databricks-openai",
     "databricks-gemini",
+    "databricks-bedrock",
 )
 
 PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES]
@@ -99,13 +101,31 @@ def _resolve_model_selector(
     return model
 
 
+def _bedrock_model_entry(model_id: str) -> dict:
+    """A Pi model entry for a Bedrock target, pinning known token limits.
+
+    Some Bedrock models cap output well below what Pi requests by default (e.g.
+    Nova rejects a `maxTokens` of 10k or more), so pin `maxTokens`/`contextWindow`
+    when the model has a known limit. Models with no known limit are left unbounded.
+    """
+    entry: dict = {"id": model_id}
+    limits = model_token_limits(model_id)
+    if limits is not None:
+        entry["contextWindow"] = limits["context"]
+        entry["maxTokens"] = limits["output"]
+    return entry
+
+
 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 = {}
@@ -149,20 +169,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": [_bedrock_model_entry(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:
@@ -183,6 +226,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")
@@ -261,6 +306,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.")