From b15a53dfde2ed8e6f8700247c9271ac84f833b2f Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 10:54:48 -0500 Subject: [PATCH 1/7] feat(codex): support amazon_bedrock Model Provider Services Codex speaks the OpenAI-compatible API, which Bedrock also exposes. `_TOOL_PROVIDER_TYPES` previously restricted codex to `openai` only, so `ucode codex --provider ` always failed with "which codex can't route to (supported: openai)." Three changes in databricks.py: - Add `amazon_bedrock` to codex's allowed provider types in `_TOOL_PROVIDER_TYPES`. - Gate the "exposes no Claude models" check in `resolve_provider_service` on `tool == "claude"` so a Bedrock MPS with OpenAI-compatible (non-Claude) targets isn't rejected when codex selects it. - Apply the same `tool == "claude"` guard in `service_usable_for_tool` so Bedrock services without Claude targets appear in the list when codex is the active tool. Co-Authored-By: Claude Sonnet 4.6 Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/databricks.py | 13 ++++++++----- tests/test_databricks.py | 35 +++++++++++++++++++++++++++++++++-- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f875a5c7..daa633ed 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2096,7 +2096,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: # form produced by `_provider_type_tag` (e.g. `amazon_bedrock`). _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = { "claude": ("anthropic", "amazon_bedrock"), - "codex": ("openai",), + "codex": ("openai", "amazon_bedrock"), } # Provider types that expose Bedrock-style model ids (e.g. @@ -2324,12 +2324,13 @@ def service_usable_for_tool(tool: str, service: dict) -> bool: Beyond the provider-type match, a Bedrock service is only usable for claude if it exposes at least one Claude model in its targets — otherwise there's no routable model id to pin. (Anthropic services use canonical names, so any - match is usable.) + match is usable.) Codex uses the OpenAI-compatible Bedrock endpoint, so any + Bedrock service is usable for it regardless of declared targets. """ provider_type = service.get("provider_type", "") if not tool_supports_provider_type(tool, provider_type): return False - if provider_type in BEDROCK_PROVIDER_TYPES: + if tool == "claude" and provider_type in BEDROCK_PROVIDER_TYPES: return bool(map_claude_family_models(service.get("targets") or [])) return True @@ -2370,8 +2371,10 @@ def resolve_provider_service( f"Model provider service '{service_name}' is a '{provider_type}' provider, " f"which {tool} can't route to (supported: {supported})." ) - if provider_type in BEDROCK_PROVIDER_TYPES and not map_claude_family_models( - match.get("targets") or [] + if ( + tool == "claude" + and provider_type in BEDROCK_PROVIDER_TYPES + and not map_claude_family_models(match.get("targets") or []) ): return None, ( f"Model provider service '{service_name}' exposes no Claude models — " diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 2d9f61a5..a43e1eb1 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -605,12 +605,16 @@ def test_claude_includes_anthropic_and_usable_bedrock(self, monkeypatch): "main.schema2.bedrock-svc", ] - def test_codex_filters_to_openai(self, monkeypatch): + def test_codex_filters_to_openai_and_bedrock(self, monkeypatch): + # codex supports both openai and amazon_bedrock provider types. monkeypatch.setattr( db_mod, "_http_get_json", lambda url, token, timeout=30: (self._PAYLOAD, None) ) names, _ = db_mod.list_tool_provider_services("codex", WS, "token") - assert names == ["main.schema1.openai-svc"] + assert "main.schema1.openai-svc" in names + assert "main.schema2.bedrock-svc" in names + assert "main.schema2.bedrock-titan-svc" in names + assert "main.schema1.anthropic-svc" not in names class TestMapClaudeFamilyModels: @@ -872,6 +876,33 @@ def test_bedrock_without_claude_rejected(self, monkeypatch): assert service is None assert "no Claude models" in error + def test_codex_bedrock_openai_compat_ok(self, monkeypatch): + # Bedrock MPS exposing non-Claude (OpenAI-compatible) models must work for codex. + self._patch(monkeypatch) + service, error = db_mod.resolve_provider_service( + "codex", "main.schema2.bedrock-titan-svc", WS, "token" + ) + assert error is None + assert service["provider_type"] == "amazon_bedrock" + + def test_codex_bedrock_with_claude_targets_ok(self, monkeypatch): + # Bedrock MPS that happens to expose Claude targets is also valid for codex. + self._patch(monkeypatch) + service, error = db_mod.resolve_provider_service( + "codex", "main.schema2.bedrock-svc", WS, "token" + ) + assert error is None + assert service["provider_type"] == "amazon_bedrock" + + def test_codex_anthropic_rejected(self, monkeypatch): + # codex does not speak the Anthropic Messages API. + self._patch(monkeypatch) + service, error = db_mod.resolve_provider_service( + "codex", "main.schema1.anthropic-svc", WS, "token" + ) + assert service is None + assert "can't route to" in error + def test_not_found_lists_usable(self, monkeypatch): self._patch(monkeypatch) service, error = db_mod.resolve_provider_service("claude", "main.x.missing", WS, "token") From a093d6949fd6a9db7a71f81d022b76e6342acd3d Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 11:03:00 -0500 Subject: [PATCH 2/7] feat: add `ucode providers list/show` commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds two new subcommands under `ucode providers` to inspect Model Provider Services on the workspace: - `ucode providers list [--tool TOOL]` — lists all MPS services with name, provider type, and declared targets. `--tool claude|codex` filters to services the given tool can actually route through. - `ucode providers show ` — shows full detail for one service: provider type, relay flag, allow_all_targets, and the complete targets list. Motivation: after `ucode codex --provider eng_dev.ai_gateway.amazonbedrock` launched without showing expected Bedrock models, there was no CLI to inspect what targets an MPS exposes. Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/cli.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6a1662ae..94c699b1 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -51,15 +51,18 @@ find_profile_name_for_host, get_databricks_profiles, get_databricks_token, + get_model_provider_service, install_databricks_cli, is_model_provider_feature_unavailable, is_workspace_admin, + list_model_provider_services, list_profile_entries, list_tool_provider_services, normalize_workspace_url, resolve_pat_token, resolve_provider_launch_model, run_databricks_login, + service_usable_for_tool, ) from ucode.managed_budget import ( budget_usage_percent, @@ -125,6 +128,7 @@ from ucode.ui import ( console, heading, + muted, print_err, print_heading, print_kv, @@ -136,6 +140,7 @@ prompt_for_tools, prompt_for_workspace, prompt_yes_no, + render_box_table, set_verbosity, spinner, status_badge, @@ -1159,6 +1164,8 @@ def revert() -> int: app.add_typer(configure_app, name="configure", help="Configure workspace and tool settings.") mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") +providers_app = typer.Typer(add_completion=False, no_args_is_help=True) +app.add_typer(providers_app, name="providers", help="Inspect Model Provider Services on the workspace.") setup_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer( setup_app, @@ -3256,6 +3263,82 @@ def upgrade_cmd() -> None: print_success("ucode upgraded") +@providers_app.command("list") +def providers_list_cmd( + tool: Annotated[ + str | None, + typer.Option("--tool", help="Filter to services usable by a specific tool (claude, codex)."), + ] = None, +) -> None: + """List Model Provider Services on the workspace.""" + state = load_state() + workspace = state.get("workspace") + if not workspace: + print_err("No workspace configured. Run `ucode configure` first.") + raise typer.Exit(1) from None + token = get_databricks_token(workspace, state.get("profile")) + with spinner("Fetching model provider services..."): + services, reason = list_model_provider_services(workspace, token) + if reason is not None: + print_err(f"Could not list model provider services: {reason}") + raise typer.Exit(1) from None + if tool: + services = [s for s in services if service_usable_for_tool(tool, s)] + if not services: + msg = "No model provider services found" + (f" for {tool}" if tool else "") + "." + print_note(msg) + return + rows = [ + [ + s["name"], + s["provider_type"], + ", ".join(s["targets"]) if s["targets"] else ("(all)" if s["allow_all_targets"] else "—"), + ] + for s in services + ] + print_section("Model Provider Services") + console.print(render_box_table(["Service", "Provider", "Targets"], rows, max_widths=[60, 20, 60])) + if tool: + console.print(muted(f" Filtered to services usable by {tool}.")) + + +@providers_app.command("show") +def providers_show_cmd( + service_name: Annotated[ + str, + typer.Argument(help="Fully qualified service name (catalog.schema.service)."), + ], +) -> None: + """Show targets and configuration for a Model Provider Service.""" + state = load_state() + workspace = state.get("workspace") + if not workspace: + print_err("No workspace configured. Run `ucode configure` first.") + raise typer.Exit(1) from None + token = get_databricks_token(workspace, state.get("profile")) + with spinner(f"Fetching {service_name}..."): + service, reason = get_model_provider_service(service_name, workspace, token) + if reason is not None: + print_err(f"Could not fetch '{service_name}': {reason}") + raise typer.Exit(1) from None + if service is None: + print_err(f"Model provider service '{service_name}' not found.") + raise typer.Exit(1) from None + print_section(service["name"]) + print_kv("Provider type", service["provider_type"]) + if service["relayed"]: + print_kv("Relay", "yes (subscription-backed, no credential stored)") + if service["allow_all_targets"]: + print_kv("Allow all targets", "yes") + targets = service["targets"] + if targets: + print_kv("Targets", targets[0]) + for t in targets[1:]: + print_kv("", t) + else: + print_kv("Targets", "none declared") + + def main() -> None: app() From a6ae6ace7e5c7549bcc7812e2340c4fdf0c03d46 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 11:14:28 -0500 Subject: [PATCH 3/7] feat: pin Bedrock target model when launching codex with an MPS provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When `ucode codex --provider ` is used, Codex's built-in model picker queries OpenAI for its model list — showing gpt-5-codex and gpt-5 instead of the Bedrock targets declared on the MPS. Fix this by: - Fetching the MPS targets at launch time and offering a picker (or auto-pinning when there's only one target) - Honoring an explicit `--model` flag for codex in the provider path, which was previously a no-op - Making `codex.write_tool_config` actually use the `model` parameter when a provider is active (it was silently ignored before) Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/codex.py | 9 ++++---- src/ucode/cli.py | 43 +++++++++++++++++++++++++++++++++++---- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index b53ca113..96467404 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -300,11 +300,12 @@ def revert_legacy_shared_config() -> bool: def write_tool_config(state: dict, model: str | None = None, provider: str | None = None) -> dict: workspace = state["workspace"] - # Leave model selection to Codex. The gateway still receives the configured - # provider and authentication settings, while Codex uses its own default. - # A managed default is the sole exception. + # Leave model selection to Codex — except when a provider is set and a target + # model was resolved from its MPS targets, or an admin managed default exists. managed_model = state.get("codex_default_model") - chosen_model = managed_model if isinstance(managed_model, str) else None + chosen_model = (model if provider else None) or ( + managed_model if isinstance(managed_model, str) else None + ) databricks_profile = state.get("profile") if _use_legacy_layout(): diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 94c699b1..8a7e0ae8 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -1165,7 +1165,9 @@ def revert() -> int: mcp_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ucode.") providers_app = typer.Typer(add_completion=False, no_args_is_help=True) -app.add_typer(providers_app, name="providers", help="Inspect Model Provider Services on the workspace.") +app.add_typer( + providers_app, name="providers", help="Inspect Model Provider Services on the workspace." +) setup_app = typer.Typer(add_completion=False, no_args_is_help=False) app.add_typer( setup_app, @@ -2058,6 +2060,33 @@ def _launch_tool( # Relayed services forward --model to Claude Code's own flag at launch (below), not env. if tool == "claude" and not relayed and (model or provider_models): route_root_model = resolve_provider_launch_model(model, provider_models or {}) + elif tool == "codex": + # Codex's built-in model picker queries OpenAI, not the MPS, so it shows the + # wrong model list when routing through a Bedrock provider. Pin a target from + # the MPS so Codex never reaches its picker. + if model: + resolved_model = model + else: + _token = get_databricks_token(state["workspace"], state.get("profile")) + with spinner("Fetching provider model targets..."): + _svc, _ = get_model_provider_service(provider, state["workspace"], _token) + if _svc: + _targets: list[str] = _svc.get("targets") or [] + if len(_targets) == 1: + resolved_model = _targets[0] + elif len(_targets) > 1: + _picked = prompt_for_selection( + "Select a model from the provider service:", + [(_t, _t) for _t in _targets], + ) + if _picked is None: + raise KeyboardInterrupt + resolved_model = _picked + elif _svc.get("allow_all_targets"): + print_warning( + f"'{provider}' allows all targets but has none declared. " + "Pass --model with the Bedrock model ID you want to use." + ) else: # A managed default_model is the model the admin wants sessions to start on, so it goes # in as the explicit model rather than being applied afterwards: for codex the proto has @@ -3267,7 +3296,9 @@ def upgrade_cmd() -> None: def providers_list_cmd( tool: Annotated[ str | None, - typer.Option("--tool", help="Filter to services usable by a specific tool (claude, codex)."), + typer.Option( + "--tool", help="Filter to services usable by a specific tool (claude, codex)." + ), ] = None, ) -> None: """List Model Provider Services on the workspace.""" @@ -3292,12 +3323,16 @@ def providers_list_cmd( [ s["name"], s["provider_type"], - ", ".join(s["targets"]) if s["targets"] else ("(all)" if s["allow_all_targets"] else "—"), + ", ".join(s["targets"]) + if s["targets"] + else ("(all)" if s["allow_all_targets"] else "—"), ] for s in services ] print_section("Model Provider Services") - console.print(render_box_table(["Service", "Provider", "Targets"], rows, max_widths=[60, 20, 60])) + console.print( + render_box_table(["Service", "Provider", "Targets"], rows, max_widths=[60, 20, 60]) + ) if tool: console.print(muted(f" Filtered to services usable by {tool}.")) From 50c41e33d9a533d2b813449c17d15d7b7d25b90c Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:41:08 -0500 Subject: [PATCH 4/7] feat: add Pi Bedrock provider support via correct gateway base URL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire `ucode pi --provider ` end-to-end: - `build_pi_base_urls`: add "bedrock" key pointing at `{workspace}/ai-gateway` (NOT `/ai-gateway/amazonbedrock` — that path maps to the Bedrock control plane; the standard path routes to the runtime via the MPS header) - `pi.render_overlay`: add `databricks-bedrock` provider block when `bedrock_targets` is supplied; defaults the session to the first target - `pi.write_tool_config`: accept `provider` and `bedrock_targets` kwargs - `agents.__init__.configure_tool`: pass `bedrock_targets` to Pi; allow Pi to launch without a model when a Bedrock provider + targets cover it - `cli.py`: fetch MPS targets for Pi in the provider launch path; handle `allow_all_targets` with a text prompt; thread `bedrock_targets` through to `configure_tool` Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/__init__.py | 13 ++++++-- src/ucode/agents/pi.py | 32 +++++++++++++++--- src/ucode/cli.py | 65 +++++++++++++++++++++++++++++++++--- src/ucode/databricks.py | 42 +++++++++++++++++++++++ 4 files changed, 140 insertions(+), 12 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 578aa208..edbcd1dc 100644 --- a/src/ucode/agents/__init__.py +++ b/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,22 @@ 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 == "pi" 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 + ) 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): diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index a673a548..2c8785eb 100644 --- a/src/ucode/agents/pi.py +++ b/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,9 +151,23 @@ 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, + "headers": {**ua_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) + # When launching with a Bedrock provider, default to the first target. + if not resolved and "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 @@ -157,10 +175,12 @@ def render_overlay( 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 +201,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") diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 8a7e0ae8..54397540 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -37,7 +37,7 @@ ) from ucode.agents.codex import revert_legacy_shared_config from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH -from ucode.config_io import is_dry_run, restore_file, set_dry_run +from ucode.config_io import is_dry_run, read_toml_safe, restore_file, set_dry_run from ucode.databricks import ( apply_pat_environment, build_shared_base_urls, @@ -56,6 +56,7 @@ is_model_provider_feature_unavailable, is_workspace_admin, list_model_provider_services, + list_mps_codex_models, list_profile_entries, list_tool_provider_services, normalize_workspace_url, @@ -137,6 +138,7 @@ print_success, print_warning, prompt_for_selection, + prompt_for_text, prompt_for_tools, prompt_for_workspace, prompt_yes_no, @@ -2051,6 +2053,7 @@ def _launch_tool( # The router's per-launch pick for the root session. Codex pins it as the # resolved model; claude pins it via ANTHROPIC_MODEL (route_root_model). route_root_model = None + bedrock_targets: list[str] | None = None if provider: # Routing through a Model Provider Service pins no Databricks model; # the agent uses its own canonical model names (header selects the @@ -2083,10 +2086,63 @@ def _launch_tool( raise KeyboardInterrupt resolved_model = _picked elif _svc.get("allow_all_targets"): - print_warning( - f"'{provider}' allows all targets but has none declared. " - "Pass --model with the Bedrock model ID you want to use." + # No declared targets but the service allows any — query the + # provider's OpenAI-compatible /models endpoint to get the list. + # Reuse the previously-saved model only when the config was last + # written with this same provider; a workspace model from a + # non-MPS run must not bleed into a Bedrock session. + _prev_cfg = read_toml_safe(codex_agent.CODEX_CONFIG_PATH) + _stored_provider = ( + _prev_cfg.get("model_providers", {}) + .get(codex_agent.CODEX_MODEL_PROVIDER_NAME, {}) + .get("http_headers", {}) + .get("Databricks-Model-Provider-Service") ) + _prev_model: str | None = ( + _prev_cfg.get("model") if _stored_provider == provider else None + ) + with spinner("Querying available models from provider..."): + _mps_models, _mps_err = list_mps_codex_models( + provider, state["workspace"], _token + ) + if _mps_err is None and _mps_models: + if len(_mps_models) == 1: + resolved_model = _mps_models[0] + else: + _mpicked = prompt_for_selection( + "Select a model from the provider service:", + [(_t, _t) for _t in _mps_models], + ) + if _mpicked is None: + raise KeyboardInterrupt + resolved_model = _mpicked + else: + # Live query failed or returned nothing — fall back to a + # free-text prompt, defaulting to the previously-saved model + # so subsequent launches don't ask again. + resolved_model = prompt_for_text( + f"Enter the model ID to use with '{provider}'", + default=_prev_model, + required=not _prev_model, + ) + elif tool == "pi": + # Pi receives the MPS targets as its databricks-bedrock model list; + # a single model is also set as the default for the session. + _pi_token = get_databricks_token(state["workspace"], state.get("profile")) + with spinner("Fetching provider model targets..."): + _pi_svc, _ = get_model_provider_service(provider, state["workspace"], _pi_token) + if _pi_svc: + bedrock_targets = _pi_svc.get("targets") or [] + if bedrock_targets: + resolved_model = bedrock_targets[0] + elif _pi_svc.get("allow_all_targets"): + _pi_entered = prompt_for_text( + f"Enter a Bedrock model ID to use with '{provider}'", + required=True, + ) + if _pi_entered: + bedrock_targets = [_pi_entered] + resolved_model = _pi_entered else: # A managed default_model is the model the admin wants sessions to start on, so it goes # in as the explicit model rather than being applied afterwards: for codex the proto has @@ -2123,6 +2179,7 @@ def _launch_tool( # the latter pins a raw id into every family alias, which would clobber the service's # per-family target pins. custom_model=model if (tool == "claude" and not provider) else None, + bedrock_targets=bedrock_targets, ) # Relayed = a Claude subscription: forward --model to Claude Code's own flag, like `-- --model X`. if tool == "claude" and provider and relayed and model and not forwarded_model: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index daa633ed..7d545329 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3049,6 +3049,44 @@ def fetch_codex_models(workspace: str, token: str) -> list[str]: return models +def list_mps_codex_models( + service_name: str, workspace: str, token: str +) -> tuple[list[str], str | None]: + """List models available through a Bedrock MPS's OpenAI-compatible endpoint. + + Queries ``{workspace}/ai-gateway/codex/v1/models`` with the + ``Databricks-Model-Provider-Service`` header so the gateway asks the MPS + what models it exposes. Used when a service has ``allow_all_targets`` set + and no explicit targets are declared. + + Returns ``(model_ids, reason)`` where ``reason`` is non-None on failure. + """ + url = f"{build_tool_base_url('codex', workspace)}/models" + req = urllib_request.Request( + url, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/json", + "Databricks-Model-Provider-Service": service_name, + }, + ) + try: + with urllib_request.urlopen(req, timeout=15) as resp: + body = resp.read().decode("utf-8") + payload = json.loads(body) + except urllib_error.HTTPError as exc: + return [], f"HTTP {exc.code}" + except Exception as exc: + return [], str(exc) + if not isinstance(payload, dict): + return [], "unexpected response shape" + data = payload.get("data") or [] + models = sorted( + str(m["id"]) for m in data if isinstance(m, dict) and isinstance(m.get("id"), str) + ) + return models, None + + def _probe_ai_gateway_v2(workspace: str, token: str) -> tuple[bool, str | None]: hostname = workspace_hostname(workspace) url = f"https://{hostname}/api/ai-gateway/v2/endpoints?page_size=1" @@ -3367,6 +3405,10 @@ def build_pi_base_urls(workspace: str) -> dict[str, str]: "claude": build_tool_base_url("claude", workspace), "openai": build_tool_base_url("codex", workspace), "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", + # Bedrock routes through the standard gateway; MPS header selects the provider. + # Do NOT include the MPS name in the path — /ai-gateway/amazonbedrock/ maps to + # the control plane (bedrock.amazonaws.com), not the runtime. + "bedrock": f"{workspace}/ai-gateway", } From d3c9e875d18ee6a6b40c30cc980a050a754dce7d Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:45:32 -0500 Subject: [PATCH 5/7] fix: add --provider option to ucode pi command Without it, --provider fell into ctx.args and was forwarded to Pi itself rather than being parsed by ucode, so the Bedrock target-fetching branch never ran. Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/cli.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 54397540..ca669123 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2593,12 +2593,20 @@ def copilot_cmd( @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def pi_cmd( ctx: typer.Context, + provider: Annotated[ + str | None, + typer.Option( + "--provider", + help="Route through a Unity Catalog Model Provider Service " + "(..). Pass before any `--` separator.", + ), + ] = None, skip_preflight: SkipPreflightOption = False, skip_managed_config: SkipManagedConfigOption = False, ) -> None: """Launch Pi coding agent via Databricks.""" _disable_managed_config_if_requested(skip_managed_config) - _launch_tool("pi", ctx, skip_preflight=skip_preflight) + _launch_tool("pi", ctx, provider=provider, skip_preflight=skip_preflight) @app.command("cursor", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) From e6c045c30d55b8109f3a1348f7480d96f5785840 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:47:15 -0500 Subject: [PATCH 6/7] fix: add pi to _TOOL_PROVIDER_TYPES for amazon_bedrock support Without this entry, ucode pi --provider rejects any Bedrock MPS with "pi can't route to (supported: none)" before ever fetching targets. Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/databricks.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 7d545329..6b98e353 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2097,6 +2097,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: _TOOL_PROVIDER_TYPES: dict[str, tuple[str, ...]] = { "claude": ("anthropic", "amazon_bedrock"), "codex": ("openai", "amazon_bedrock"), + "pi": ("anthropic", "amazon_bedrock"), } # Provider types that expose Bedrock-style model ids (e.g. From 96c468cfb7104a99080a3a2266e140ed4e4fea2a Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 21:10:40 -0500 Subject: [PATCH 7/7] fix: always prefix Bedrock selector with databricks-bedrock/ in Pi config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_resolve_model_selector` returns Bedrock model IDs (e.g. `anthropic.claude-3-haiku-20240307-v1:0`) unprefixed because they contain no `/`. The old `if not resolved` guard never fired since the ID is truthy. `_write_settings` then gets an empty model half from `partition("/")` and exits early — defaultProvider stays on databricks-claude instead of databricks-bedrock. Fix: unconditionally set `resolved = f"databricks-bedrock/{targets[0]}"` when the Bedrock provider block is present. Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/pi.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 2c8785eb..c0765c36 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -164,8 +164,11 @@ def render_overlay( } keys.append(["providers", "databricks-bedrock"]) resolved = _resolve_model_selector(model or "", claude_models, codex_models, gemini_models) - # When launching with a Bedrock provider, default to the first target. - if not resolved and "databricks-bedrock" in providers and bedrock_targets: + # 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: