From ea14b3b17b7cb4111f033197518be7adfac59d77 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 10:54:48 -0500 Subject: [PATCH 01/13] 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 49dba37f..9dade7c7 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2104,7 +2104,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: # 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"), "gemini": ("gemini_enterprise",), } @@ -2333,12 +2333,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 @@ -2379,8 +2380,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 20301da3..73efd22d 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -638,12 +638,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: @@ -905,6 +909,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 553d9f3059b7f02db733fc3036b07a75f8ea822b Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 11:03:00 -0500 Subject: [PATCH 02/13] 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 9857d78a..b981c94b 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -57,9 +57,11 @@ 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, @@ -67,6 +69,7 @@ resolve_pat_token, resolve_provider_launch_model, run_databricks_login, + service_usable_for_tool, ) from ucode.managed_budget import ( budget_usage_percent, @@ -132,6 +135,7 @@ from ucode.ui import ( console, heading, + muted, print_err, print_heading, print_kv, @@ -143,6 +147,7 @@ prompt_for_tools, prompt_for_workspace, prompt_yes_no, + render_box_table, set_verbosity, spinner, status_badge, @@ -1051,6 +1056,8 @@ def revert() -> int: app.add_typer(mcp_app, name="mcp", help="MCP servers exposed by ug.") skill_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(skill_app, name="skill", help="Databricks Skills for your coding tools.") +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, @@ -3353,6 +3360,82 @@ def _verify_upgraded_commands() -> None: ) +@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 9f6269eb8e3274385de775bcf931be163abea886 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:41:08 -0500 Subject: [PATCH 03/13] 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 | 15 +++++++++---- src/ucode/agents/pi.py | 32 ++++++++++++++++++++++----- src/ucode/cli.py | 24 ++++++++++++++++++++- src/ucode/databricks.py | 42 ++++++++++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 10 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 8ea3e20d..352a5a26 100644 --- a/src/ucode/agents/__init__.py +++ b/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,23 @@ 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 == "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, 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 + ) 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 a193e8ff..fd733d90 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -71,6 +71,7 @@ "databricks-claude", "databricks-openai", "databricks-gemini", + "databricks-bedrock", ) PROVIDER_KEYS: list[list[str]] = [["providers", name] for name in PROVIDER_NAMES] @@ -100,12 +101,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 = {} @@ -149,9 +153,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 @@ -159,10 +177,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: @@ -183,6 +203,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 b981c94b..ba8dfb7d 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -44,7 +44,7 @@ from ucode.agents.args import has_explicit_model_arg 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, @@ -62,6 +62,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, @@ -144,6 +145,7 @@ print_success, print_warning, prompt_for_selection, + prompt_for_text, prompt_for_tools, prompt_for_workspace, prompt_yes_no, @@ -2023,6 +2025,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 @@ -2038,6 +2041,24 @@ def _launch_tool( resolved_model, gemini_error = resolve_gemini_provider_model(state, provider, model) if gemini_error: raise RuntimeError(gemini_error) + 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 @@ -2071,6 +2092,7 @@ def _launch_tool( # Claude's explicit model is launch-scoped and is passed through LaunchOptions below. custom_model=None, coding_agent_config_defaults=coding_agent_config_defaults, + 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 9dade7c7..4ec2317a 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -3079,6 +3079,44 @@ class GatewayProbe(NamedTuple): conclusive: bool = True +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 + + _MODEL_SERVICE_PROBE_PAGE_SIZE = 50 _MODEL_SERVICE_PROBE_MAX_PAGES = 20 _MODEL_SERVICE_EMPTY_DETAIL = ( @@ -3416,6 +3454,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 0626d541765c01f32dc38df5639513c4fba1d5d3 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:45:32 -0500 Subject: [PATCH 04/13] 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 ba8dfb7d..4518b51a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2531,10 +2531,18 @@ 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, ) -> None: """Launch Pi coding agent via Databricks.""" - _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 5144331a2352bfd7e42fc32b5d7757b753bda0fc Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 16:47:15 -0500 Subject: [PATCH 05/13] 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 4ec2317a..d509aa06 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2106,6 +2106,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: "claude": ("anthropic", "amazon_bedrock"), "codex": ("openai", "amazon_bedrock"), "gemini": ("gemini_enterprise",), + "pi": ("anthropic", "amazon_bedrock"), } # Provider types that expose Bedrock-style model ids (e.g. From 6c2c868ad0c78d21d3a03fa9d9b030c6b50cec56 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 21:10:40 -0500 Subject: [PATCH 06/13] 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 fd733d90..6382d2de 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -166,8 +166,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: From 685e605ffc8c82a1c4d5308b13d8d2d574927d2d Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Thu, 3 Sep 2026 11:28:10 -0500 Subject: [PATCH 07/13] fix(pi): keep Bedrock provider across token refresh and drop duplicate UA The launch-time and 30-minute token refresh re-rendered Pi's models.json without the Bedrock provider, dropping the databricks-bedrock block and falling back to a system-hosted model. _refresh_token_once now reads the existing config and preserves a databricks-bedrock block, re-applying it with a freshly refreshed token. Also stop sending ucode's User-Agent on the Bedrock block: Pi's bedrock-converse-stream client sets its own, and two values made the gateway reject the request ("Header field 'user-agent' must only have a single value"). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/pi.py | 33 ++++++++++- tests/test_agent_pi.py | 124 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index 6382d2de..c21ba569 100644 --- a/src/ucode/agents/pi.py +++ b/src/ucode/agents/pi.py @@ -161,7 +161,11 @@ def render_overlay( "api": "bedrock-converse-stream", "apiKey": token, "authHeader": True, - "headers": {**ua_headers, "Databricks-Model-Provider-Service": provider}, + # 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"]) @@ -286,6 +290,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.") diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index ff7f172d..30290972 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -475,3 +475,127 @@ def test_pi_default_model_wins_over_allowlist(self): def test_falls_back_to_pi_models_without_default(self): state = {"pi_models": ["system.ai.claude-opus-4-8"]} assert pi.default_model(state) == "system.ai.claude-opus-4-8" + + +class TestRefreshTokenOnceBedrockPreservation: + """_refresh_token_once must preserve an existing databricks-bedrock provider block.""" + + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.pi as pi_mod + import ucode.config_io as config_io_mod + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + config_file = tmp_path / "models.json" + settings_file = tmp_path / "settings.json" + monkeypatch.setattr(pi_mod, "PI_CONFIG_PATH", config_file) + monkeypatch.setattr(pi_mod, "PI_SETTINGS_PATH", settings_file) + monkeypatch.setattr(pi_mod, "PI_BACKUP_PATH", tmp_path / "pi-backup.json") + monkeypatch.setattr(pi_mod, "PI_SETTINGS_BACKUP_PATH", tmp_path / "pi-settings-backup.json") + return pi_mod, config_file, settings_file + + def _state(self) -> dict: + return { + "workspace": WS, + "base_urls": {"pi": _base_urls()}, + "claude_models": {"sonnet": "claude-sonnet"}, + "codex_models": [], + "gemini_models": [], + "managed_configs": {}, + } + + def test_bedrock_block_survives_token_refresh(self, tmp_path, monkeypatch): + """Regression: token refresh must not clobber the databricks-bedrock provider block.""" + pi_mod, config_file, settings_file = self._setup(tmp_path, monkeypatch) + + # Pre-write a models.json that already has a bedrock provider block, + # as written by write_tool_config(..., provider=..., bedrock_targets=[...]). + bedrock_config = { + "model": "databricks-bedrock/anthropic.claude-3-haiku-20240307-v1:0", + "providers": { + "databricks-bedrock": { + "baseUrl": f"{WS}/ai-gateway", + "api": "bedrock-converse-stream", + "apiKey": "old-token", + "authHeader": True, + "headers": { + "User-Agent": "ucode/0.1.0 pi/0.74.0", + "Databricks-Model-Provider-Service": "my-mps-provider", + }, + "models": [ + {"id": "anthropic.claude-3-haiku-20240307-v1:0"}, + {"id": "anthropic.claude-3-sonnet-20240229-v1:0"}, + ], + } + }, + } + config_file.parent.mkdir(parents=True, exist_ok=True) + config_file.write_text(json.dumps(bedrock_config), encoding="utf-8") + + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="new-token"), + patch("ucode.agents.pi.save_state"), + ): + token = pi_mod._refresh_token_once(self._state()) + + assert token == "new-token" + + written = json.loads(config_file.read_text()) + providers = written.get("providers", {}) + + # The bedrock provider block must still be present. + assert "databricks-bedrock" in providers + + bedrock = providers["databricks-bedrock"] + # MPS header preserved. + assert bedrock["headers"]["Databricks-Model-Provider-Service"] == "my-mps-provider" + # Model ids preserved. + model_ids = [m["id"] for m in bedrock.get("models", [])] + assert "anthropic.claude-3-haiku-20240307-v1:0" in model_ids + assert "anthropic.claude-3-sonnet-20240229-v1:0" in model_ids + # Token refreshed. + assert bedrock["apiKey"] == "new-token" + + # settings.json must pin defaultProvider to databricks-bedrock. + settings = json.loads(settings_file.read_text()) + assert settings["defaultProvider"] == "databricks-bedrock" + + def test_no_bedrock_block_uses_default_model(self, tmp_path, monkeypatch): + """Without a bedrock block, _refresh_token_once falls back to the normal path.""" + pi_mod, config_file, settings_file = self._setup(tmp_path, monkeypatch) + + # Config has only a Claude provider — no bedrock. + existing_config = { + "model": "databricks-claude/claude-sonnet", + "providers": { + "databricks-claude": { + "baseUrl": f"{WS}/ai-gateway/anthropic", + "api": "anthropic-messages", + "apiKey": "old-token", + "authHeader": True, + "headers": {}, + "models": [{"id": "claude-sonnet"}], + } + }, + } + config_file.parent.mkdir(parents=True, exist_ok=True) + config_file.write_text(json.dumps(existing_config), encoding="utf-8") + + with ( + patch("ucode.agents.pi.get_databricks_token", return_value="new-token"), + patch("ucode.agents.pi.save_state"), + ): + token = pi_mod._refresh_token_once(self._state()) + + assert token == "new-token" + + written = json.loads(config_file.read_text()) + providers = written.get("providers", {}) + + # No bedrock block should be written. + assert "databricks-bedrock" in providers is False or "databricks-bedrock" not in providers + # Claude provider still present. + assert "databricks-claude" in providers + + # settings.json must pin defaultProvider to databricks-claude (normal path). + settings = json.loads(settings_file.read_text()) + assert settings["defaultProvider"] == "databricks-claude" From a84e018186faa39d453133486e0b6542ba7dbf6d Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Thu, 3 Sep 2026 12:26:07 -0500 Subject: [PATCH 08/13] fix(pi): resolve provider-support test and lint after dropping codex path Reflect that Pi now supports anthropic/amazon_bedrock provider services, and clean up imports left unused once the codex Bedrock launch branch is excluded. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/cli.py | 19 +++++++++++++------ tests/test_managed_setup.py | 6 +++++- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 4518b51a..b7a25591 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -44,7 +44,7 @@ from ucode.agents.args import has_explicit_model_arg 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, read_toml_safe, restore_file, set_dry_run +from ucode.config_io import is_dry_run, restore_file, set_dry_run from ucode.databricks import ( apply_pat_environment, build_shared_base_urls, @@ -62,7 +62,6 @@ 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, @@ -1059,7 +1058,9 @@ def revert() -> int: skill_app = typer.Typer(add_completion=False, no_args_is_help=True) app.add_typer(skill_app, name="skill", help="Databricks Skills for your coding tools.") 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, @@ -3394,7 +3395,9 @@ def _verify_upgraded_commands() -> 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.""" @@ -3419,12 +3422,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}.")) diff --git a/tests/test_managed_setup.py b/tests/test_managed_setup.py index 140bc3d9..07ef6f14 100644 --- a/tests/test_managed_setup.py +++ b/tests/test_managed_setup.py @@ -362,8 +362,12 @@ def test_codex_supports_openai(self): def test_claude_does_not_support_openai(self): assert not supports_provider_service("claude", "openai") + def test_pi_supports_anthropic_and_bedrock(self): + assert supports_provider_service("pi", "anthropic") + assert supports_provider_service("pi", "amazon_bedrock") + def test_other_agents_have_no_provider_support(self): - for tool in ("gemini", "opencode", "pi", "copilot"): + for tool in ("gemini", "opencode", "copilot"): assert not supports_provider_service(tool, "anthropic"), tool From 96b077c02fb09230186670053b4e52ce6de700e6 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Thu, 3 Sep 2026 14:00:20 -0500 Subject: [PATCH 09/13] fix(pi): pin per-model output caps for Bedrock targets Some Bedrock models cap output well below Pi's default request (Nova rejects maxTokens >= 10000). Pin maxTokens/contextWindow on a Bedrock model entry when the model has a known limit (new `nova` entry in _MODEL_TOKEN_LIMITS); models with no known low cap, like Claude, stay unbounded. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/pi.py | 18 +++++++++++++++++- src/ucode/databricks.py | 4 ++++ tests/test_agent_pi.py | 27 +++++++++++++++++++++++++++ 3 files changed, 48 insertions(+), 1 deletion(-) diff --git a/src/ucode/agents/pi.py b/src/ucode/agents/pi.py index c21ba569..bbbded17 100644 --- a/src/ucode/agents/pi.py +++ b/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 @@ -100,6 +101,21 @@ 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 | None, token: str, @@ -166,7 +182,7 @@ def render_overlay( # 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], + "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) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index d509aa06..81155d64 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1541,6 +1541,10 @@ def classify_model_family(model_id: str) -> str | None: # GLM-4.6: 200k context, but the gateway caps output well below the model's # native 128k — pin 25k so requests aren't rejected. "glm": {"context": 200_000, "output": 25_000}, + # Amazon Bedrock Nova (Micro/Lite/Pro), served over Converse: the gateway + # rejects an output cap of 10k or more ("model limit of 10000"), so pin a + # value safely under it. Claude/others have no known low cap and stay unset. + "nova": {"context": 300_000, "output": 8_192}, } diff --git a/tests/test_agent_pi.py b/tests/test_agent_pi.py index 30290972..1604d179 100644 --- a/tests/test_agent_pi.py +++ b/tests/test_agent_pi.py @@ -477,6 +477,33 @@ def test_falls_back_to_pi_models_without_default(self): assert pi.default_model(state) == "system.ai.claude-opus-4-8" +class TestRenderOverlayBedrockLimits: + """Bedrock model entries pin known per-model token caps, and only those.""" + + def _bedrock_models(self, targets): + overlay, _ = pi.render_overlay( + targets[0], + "tok", + _base_urls(), + {}, + [], + [], + provider="cat.sch.mps", + bedrock_targets=targets, + ) + return overlay["providers"]["databricks-bedrock"]["models"] + + def test_nova_target_pins_max_tokens(self): + entry = self._bedrock_models(["us.amazon.nova-lite-v1:0"])[0] + assert entry["id"] == "us.amazon.nova-lite-v1:0" + assert entry["maxTokens"] == 8192 + assert entry["contextWindow"] == 300_000 + + def test_claude_target_has_no_cap(self): + entry = self._bedrock_models(["us.anthropic.claude-sonnet-4-20250514-v1:0"])[0] + assert entry == {"id": "us.anthropic.claude-sonnet-4-20250514-v1:0"} + + class TestRefreshTokenOnceBedrockPreservation: """_refresh_token_once must preserve an existing databricks-bedrock provider block.""" From 9ea89a864b68e6ccf35491e99732b21e7276303c Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Thu, 3 Sep 2026 13:29:15 -0500 Subject: [PATCH 10/13] feat(opencode): Amazon Bedrock Model Provider Service support Route OpenCode to an Amazon Bedrock MPS through the Databricks AI Gateway, mirroring the Pi support. OpenCode uses the @ai-sdk/amazon-bedrock provider with a bearer apiKey (no SigV4, no region) against {workspace}/ai-gateway, which the SDK turns into /model/{id}/converse-stream. The Databricks-Model-Provider-Service header rides per-model, since OpenCode clobbers provider-level headers. _refresh_token_once now preserves an existing databricks-bedrock block across the launch-time and 30-minute token refresh (reading the saved MPS name and target ids back out of the per-model headers), so the session keeps routing to Bedrock instead of dropping to a system-hosted model. Verified end to end against the live gateway: the generated config survives a refresh and a real `opencode run` returns Bedrock output. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_017G9kjrbhqvWucH26GwykSn --- src/ucode/agents/__init__.py | 6 +- src/ucode/agents/opencode.py | 79 +++++++++-- src/ucode/cli.py | 16 ++- src/ucode/databricks.py | 4 + tests/test_agent_opencode.py | 248 +++++++++++++++++++++++++++++++++++ tests/test_managed_wizard.py | 3 +- 6 files changed, 342 insertions(+), 14 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 352a5a26..7a1c3918 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -450,7 +450,7 @@ def configure_tool( # 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 == "pi" and provider and bedrock_targets): + 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 @@ -462,6 +462,10 @@ def configure_tool( 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) diff --git a/src/ucode/agents/opencode.py b/src/ucode/agents/opencode.py index 30d6dba8..0119e595 100644 --- a/src/ucode/agents/opencode.py +++ b/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,7 +308,12 @@ 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 @@ -292,12 +321,18 @@ def render_overlay( 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,6 +341,8 @@ 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) @@ -313,6 +350,7 @@ def write_tool_config( 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) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b7a25591..293dc5fa 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2042,9 +2042,9 @@ def _launch_tool( resolved_model, gemini_error = resolve_gemini_provider_model(state, provider, model) if gemini_error: raise RuntimeError(gemini_error) - 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. + elif tool in ("pi", "opencode"): + # Pi and OpenCode receive the MPS targets as their 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) @@ -2514,10 +2514,18 @@ def gemini_cmd( ) def opencode_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, ) -> None: """Launch OpenCode via Databricks.""" - _launch_tool("opencode", ctx, skip_preflight=skip_preflight) + _launch_tool("opencode", ctx, provider=provider, skip_preflight=skip_preflight) @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 81155d64..f48cc566 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2110,6 +2110,7 @@ def build_skills_mcp_url(workspace: str, locations: list[str]) -> str: "claude": ("anthropic", "amazon_bedrock"), "codex": ("openai", "amazon_bedrock"), "gemini": ("gemini_enterprise",), + "opencode": ("amazon_bedrock",), "pi": ("anthropic", "amazon_bedrock"), } @@ -3440,6 +3441,9 @@ def build_opencode_base_urls(workspace: str) -> dict[str, str]: "anthropic": build_tool_base_url("claude", workspace) + "/v1", "gemini": build_tool_base_url("gemini", workspace) + "/v1beta", "oss": f"{workspace}/ai-gateway/mlflow/v1", + # Bedrock routes through the standard gateway; MPS header selects the + # provider. Do NOT include the MPS name in the path. + "bedrock": f"{workspace}/ai-gateway", } diff --git a/tests/test_agent_opencode.py b/tests/test_agent_opencode.py index 45c04acd..ea7ce750 100644 --- a/tests/test_agent_opencode.py +++ b/tests/test_agent_opencode.py @@ -500,3 +500,251 @@ def test_config_written_with_correct_model(self, tmp_path, monkeypatch): written = json.loads(config_file.read_text()) assert written["model"] == "databricks-anthropic/claude-sonnet" + + +def _bedrock_base_urls() -> dict[str, str]: + return { + **_base_urls(), + "bedrock": f"{WS}/ai-gateway", + } + + +class TestRenderOverlayBedrock: + def test_bedrock_provider_added_when_provider_and_targets(self): + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert "databricks-bedrock" in overlay["provider"] + + def test_bedrock_uses_amazon_bedrock_npm_package(self): + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert overlay["provider"]["databricks-bedrock"]["npm"] == "@ai-sdk/amazon-bedrock" + + def test_bedrock_uses_gateway_base_url(self): + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + options = overlay["provider"]["databricks-bedrock"]["options"] + assert options["baseURL"] == f"{WS}/ai-gateway" + + def test_bedrock_uses_token_as_api_key(self): + overlay, _ = opencode.render_overlay( + None, + "mytoken", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert overlay["provider"]["databricks-bedrock"]["options"]["apiKey"] == "mytoken" + + def test_bedrock_no_region_in_options(self): + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert "region" not in overlay["provider"]["databricks-bedrock"]["options"] + + def test_bedrock_mps_header_is_per_model(self): + target = "anthropic.claude-3-haiku-20240307-v1:0" + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=[target], + ) + model_entry = overlay["provider"]["databricks-bedrock"]["models"][target] + assert model_entry["headers"]["Databricks-Model-Provider-Service"] == "main.ai.my-mps" + + def test_bedrock_ua_header_is_per_model(self, monkeypatch): + monkeypatch.setattr("ucode.agents.opencode.ucode_version", lambda: "1.0.0") + monkeypatch.setattr("ucode.agents.opencode.agent_version", lambda _: "2.0.0") + target = "anthropic.claude-3-haiku-20240307-v1:0" + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=[target], + ) + ua = overlay["provider"]["databricks-bedrock"]["models"][target]["headers"]["User-Agent"] + assert ua == "ucode/1.0.0 opencode/2.0.0" + + def test_bedrock_no_authorization_header_at_provider_level(self): + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert "headers" not in overlay["provider"]["databricks-bedrock"]["options"] + + def test_bedrock_model_selector_prefixed(self): + target = "anthropic.claude-3-haiku-20240307-v1:0" + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=[target], + ) + assert overlay["model"] == f"databricks-bedrock/{target}" + + def test_bedrock_all_targets_listed_as_models(self): + targets = [ + "anthropic.claude-3-haiku-20240307-v1:0", + "anthropic.claude-3-sonnet-20240229-v1:0", + ] + overlay, _ = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=targets, + ) + models = overlay["provider"]["databricks-bedrock"]["models"] + assert set(models.keys()) == set(targets) + + def test_bedrock_managed_key_tracked(self): + _, keys = opencode.render_overlay( + None, + "tok", + _bedrock_base_urls(), + {}, + provider="main.ai.my-mps", + bedrock_targets=["anthropic.claude-3-haiku-20240307-v1:0"], + ) + assert ["provider", "databricks-bedrock"] in keys + + +class TestRefreshTokenOnceBedrockPreservation: + """_refresh_token_once must preserve an existing databricks-bedrock provider block.""" + + def _setup(self, tmp_path, monkeypatch): + import ucode.agents.opencode as oc_mod + import ucode.config_io as config_io_mod + + monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path) + config_file = tmp_path / "opencode.json" + backup_file = tmp_path / "opencode-backup.json" + monkeypatch.setattr(oc_mod, "OPENCODE_CONFIG_PATH", config_file) + monkeypatch.setattr(oc_mod, "OPENCODE_BACKUP_PATH", backup_file) + return oc_mod, config_file + + def _state(self) -> dict: + return { + "workspace": WS, + "base_urls": {"opencode": _bedrock_base_urls()}, + "opencode_models": {"anthropic": ["claude-sonnet"]}, + "managed_configs": {}, + } + + def test_bedrock_block_survives_token_refresh(self, tmp_path, monkeypatch): + """Regression: token refresh must not clobber the databricks-bedrock provider block.""" + oc_mod, config_file = self._setup(tmp_path, monkeypatch) + + bedrock_config = { + "model": "databricks-bedrock/anthropic.claude-3-haiku-20240307-v1:0", + "provider": { + "databricks-bedrock": { + "npm": "@ai-sdk/amazon-bedrock", + "options": { + "baseURL": f"{WS}/ai-gateway", + "apiKey": "old-token", + }, + "models": { + "anthropic.claude-3-haiku-20240307-v1:0": { + "headers": { + "User-Agent": "ucode/0.1.0 opencode/0.74.0", + "Databricks-Model-Provider-Service": "my-mps-provider", + } + }, + "anthropic.claude-3-sonnet-20240229-v1:0": { + "headers": { + "User-Agent": "ucode/0.1.0 opencode/0.74.0", + "Databricks-Model-Provider-Service": "my-mps-provider", + } + }, + }, + } + }, + } + config_file.parent.mkdir(parents=True, exist_ok=True) + config_file.write_text(json.dumps(bedrock_config), encoding="utf-8") + + with ( + patch("ucode.agents.opencode.get_databricks_token", return_value="new-token"), + patch("ucode.agents.opencode.save_state"), + ): + token = oc_mod._refresh_token_once(self._state()) + + assert token == "new-token" + + written = json.loads(config_file.read_text()) + providers = written.get("provider", {}) + + assert "databricks-bedrock" in providers + bedrock = providers["databricks-bedrock"] + assert bedrock["options"]["apiKey"] == "new-token" + + model_ids = list(bedrock["models"].keys()) + assert "anthropic.claude-3-haiku-20240307-v1:0" in model_ids + assert "anthropic.claude-3-sonnet-20240229-v1:0" in model_ids + + for entry in bedrock["models"].values(): + assert entry["headers"]["Databricks-Model-Provider-Service"] == "my-mps-provider" + + def test_no_bedrock_block_uses_default_model(self, tmp_path, monkeypatch): + """Without a bedrock block, _refresh_token_once falls back to the normal path.""" + oc_mod, config_file = self._setup(tmp_path, monkeypatch) + + existing_config = { + "model": "databricks-anthropic/claude-sonnet", + "provider": { + "databricks-anthropic": { + "npm": "@ai-sdk/anthropic", + "options": {"baseURL": f"{WS}/ai-gateway/anthropic/v1", "apiKey": "old-token"}, + "models": {"claude-sonnet": {}}, + } + }, + } + config_file.parent.mkdir(parents=True, exist_ok=True) + config_file.write_text(json.dumps(existing_config), encoding="utf-8") + + with ( + patch("ucode.agents.opencode.get_databricks_token", return_value="new-token"), + patch("ucode.agents.opencode.save_state"), + ): + token = oc_mod._refresh_token_once(self._state()) + + assert token == "new-token" + written = json.loads(config_file.read_text()) + assert "databricks-anthropic" in written.get("provider", {}) diff --git a/tests/test_managed_wizard.py b/tests/test_managed_wizard.py index a7173414..667696ec 100644 --- a/tests/test_managed_wizard.py +++ b/tests/test_managed_wizard.py @@ -1297,8 +1297,9 @@ def fake_spinner(message): class TestProviderServiceSelection: def test_agents_without_provider_support_skip_the_prompt(self): + # gemini has no provider type support; the prompt must be skipped entirely. with patch.object(wizard, "list_model_provider_services") as listing: - assert wizard._select_provider_service("opencode", WORKSPACE, "token") is None + assert wizard._select_provider_service("gemini", WORKSPACE, "token") is None assert not listing.called def test_feature_disabled_is_silent(self): From 578a0b333cd4e6e28cd25304eb0b7a13667291cc Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 2 Sep 2026 11:14:28 -0500 Subject: [PATCH 11/13] 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 | 27 +++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1f25252f..39a2b61a 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -306,11 +306,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 293dc5fa..5c72f664 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -2060,6 +2060,33 @@ def _launch_tool( if _pi_entered: bedrock_targets = [_pi_entered] resolved_model = _pi_entered + 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 From 6ca522df679525e46fdc750262079ca2e22be714 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 9 Sep 2026 10:53:14 -0500 Subject: [PATCH 12/13] fix(codex): preserve Bedrock target model across launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit clear_model_preferences runs at the start of launch() and wipes any model key that is not backed by codex_default_model in state. When a Bedrock MPS provider is active the model written to the config is the Bedrock target id (e.g. us.openai.gpt-5.6-luna) that the gateway uses for routing — clearing it causes Codex to fall back to its own model picker, which queries OpenAI directly and returns an id the MPS does not recognise, producing a 403. Add _has_active_mps() and return early from clear_model_preferences when Databricks-Model-Provider-Service is set in the provider block. This covers both the direct call in launch() and the indirect path through default_model() inside _app_server_start_model(). Co-Authored-By: Claude Sonnet 4.6 (1M context) Claude-Session: https://claude.ai/code/session_01HaJ5LAQfUTfJPhTAismXmU --- src/ucode/agents/codex.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 39a2b61a..1879cab6 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -473,10 +473,31 @@ def default_model(state: dict) -> str | None: return None +def _has_active_mps() -> bool: + """True when the ucode Codex config has a Databricks-Model-Provider-Service header set. + + When an MPS is active the model written to the config is the Bedrock target id + (e.g. ``us.openai.gpt-5.6-luna``) that the gateway uses for routing. Clearing + that id here would cause Codex to fall back to its own model picker, which queries + OpenAI directly and returns an id the MPS doesn't recognise. + """ + doc = read_toml_safe(CODEX_CONFIG_PATH) + providers = doc.get("model_providers") + if not isinstance(providers, dict): + return False + gateway = providers.get(CODEX_MODEL_PROVIDER_NAME) + if not isinstance(gateway, dict): + return False + headers = gateway.get("http_headers") + return isinstance(headers, dict) and bool(headers.get("Databricks-Model-Provider-Service")) + + def clear_model_preferences(state: dict) -> bool: """Remove ucode profile model preferences so Codex selects its default.""" if isinstance(state.get("codex_default_model"), str): return False + if _has_active_mps(): + return False doc = read_toml_safe(CODEX_CONFIG_PATH) changed = False for key in ("model", "model_reasoning_effort"): From 5eb41e7000b74991cfeb0f11c2ab221483b62f55 Mon Sep 17 00:00:00 2001 From: Billy Janssen Date: Wed, 9 Sep 2026 13:40:49 -0500 Subject: [PATCH 13/13] feat(codex): inject model catalog for Bedrock targets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex refreshes its model catalog from the gateway's codex/v1/models endpoint, which is disabled on this workspace ("not enabled for this workspace"). Codex then has no metadata for the Bedrock provider-side target ids it resolves to (e.g. us.openai.gpt-5.6-luna) and falls back to generic metadata — printing a warning and, worse, losing web search and reasoning configuration for the session. The gateway can't serve the catalog either: Bedrock has no models-list operation, and enabling the codex dialect is an undocumented platform-team toggle. So generate the catalog client-side instead: - `codex debug models` (no --profile, so it reads the base config, not ucode's profile where model_catalog_json lives — avoids a compounding feedback loop) yields the pristine built-in catalog. - For each OpenAI-dialect MPS target, clone the matching built-in entry under the Bedrock slug so get_model_info resolves real metadata. Non- OpenAI targets (claude/grok/qwen) have no Codex metadata to borrow and are skipped. - Write it to ~/.codex/ucode-model-catalog.json and pin it via model_catalog_json in the ucode profile. Regeneration is idempotent; the pin is dropped when a run isn't routing through a provider. Verified end to end against the live gateway: the warning and the "web search is not supported" errors are gone and codex returns Bedrock output. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01HaJ5LAQfUTfJPhTAismXmU --- src/ucode/agents/codex.py | 121 ++++++++++++++++++++++++++++++++++++++ tests/test_agent_codex.py | 112 +++++++++++++++++++++++++++++++++++ 2 files changed, 233 insertions(+) diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index 1879cab6..3120d43f 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -3,8 +3,10 @@ from __future__ import annotations import copy +import json import os import re +import subprocess from collections.abc import Callable from pathlib import Path @@ -24,6 +26,7 @@ build_auth_token_argv, build_tool_base_url, get_databricks_token, + get_model_provider_service, ) from ucode.launcher import exec_or_spawn from ucode.managed_files import ( @@ -57,6 +60,11 @@ CODEX_BACKUP_PATH = APP_DIR / "codex-ucode-config.backup.toml" LEGACY_CODEX_CONFIG_PATH = CODEX_CONFIG_DIR / "config.toml" LEGACY_CODEX_BACKUP_PATH = APP_DIR / "codex-config.backup.toml" +# Augmented model catalog ucode writes when routing Codex at a Bedrock MPS. Codex's +# remote catalog endpoint (codex/v1/models) is disabled on the gateway, so without +# this Codex has no metadata for Bedrock target ids (e.g. us.openai.gpt-5.6-luna) and +# degrades to fallback metadata — losing web search and reasoning config. +CODEX_MODEL_CATALOG_PATH = CODEX_CONFIG_DIR / "ucode-model-catalog.json" CODEX_MODEL_PROVIDER_NAME = "ucode-databricks" MINIMUM_CODEX_VERSION = (0, 134, 0) MINIMUM_CODEX_VERSION_TEXT = "0.134.0" @@ -77,6 +85,7 @@ MANAGED_KEYS: list[list[str]] = [ ["model_provider"], ["model"], + ["model_catalog_json"], ["model_providers", CODEX_MODEL_PROVIDER_NAME], ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], ] @@ -304,6 +313,110 @@ def revert_legacy_shared_config() -> bool: return _strip_legacy_ucode_entries(_legacy_config_path()) +def _canonical_codex_slug(target: str) -> str: + """Return the built-in Codex slug a Bedrock target id derives from. + + Bedrock exposes OpenAI models under region/vendor-prefixed ids + (``us.openai.gpt-5.6-luna``, ``openai.gpt-oss-120b-1:0``). Codex's built-in + catalog keys off the bare model name (``gpt-5.6-luna``), so strip everything up + to and including the ``openai.`` segment. A target with no ``openai.`` segment + (e.g. ``anthropic.claude-opus-4-8``, ``us.xai.grok-4.6``) has no Codex metadata + to borrow and is returned unchanged, so the by-slug lookup simply misses it. + """ + marker = "openai." + idx = target.rfind(marker) + return target[idx + len(marker) :] if idx != -1 else target + + +def _codex_builtin_catalog() -> dict | None: + """Return Codex's pristine built-in model catalog, or None if unavailable. + + Runs ``codex debug models`` with no ``--profile``, so it reads the base config + rather than ucode's profile — the profile is where ucode writes + ``model_catalog_json``, and reading through it would fold prior clones back into + the source (a compounding feedback loop). + """ + try: + result = subprocess.run( + [SPEC["binary"], "debug", "models"], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return None + if result.returncode != 0: + return None + try: + catalog = json.loads(result.stdout) + except json.JSONDecodeError: + return None + return catalog if isinstance(catalog, dict) else None + + +def build_model_catalog(targets: list[str]) -> Path | None: + """Write an augmented Codex catalog cloning built-in entries under Bedrock slugs. + + For each Bedrock target whose canonical slug matches a built-in Codex model, add + a copy of that model's metadata keyed by the full Bedrock id, so Codex resolves + real metadata for ids like ``us.openai.gpt-5.6-luna`` instead of falling back. + Best-effort: returns the written path, or None when the catalog can't be built + (Codex missing, no matching targets) so callers degrade to plain fallback. + """ + catalog = _codex_builtin_catalog() + if catalog is None: + return None + models = catalog.get("models") + if not isinstance(models, list): + return None + target_set = set(targets) + # Drop any stale clones (a slug equal to a current target) so regeneration is + # idempotent, then re-clone from the pristine canonical entries below. + models = [m for m in models if not (isinstance(m, dict) and m.get("slug") in target_set)] + by_slug = {m["slug"]: m for m in models if isinstance(m, dict) and "slug" in m} + added = False + for target in targets: + canonical = by_slug.get(_canonical_codex_slug(target)) + if canonical is None: + continue + clone = copy.deepcopy(canonical) + clone["slug"] = target + models.append(clone) + added = True + if not added: + return None + catalog["models"] = models + try: + CODEX_MODEL_CATALOG_PATH.write_text(json.dumps(catalog)) + except OSError: + return None + return CODEX_MODEL_CATALOG_PATH + + +def _provider_catalog_path(state: dict, provider: str | None) -> Path | None: + """Build the augmented catalog for ``provider``'s MPS targets, or None. + + Best-effort and self-contained: fetches the service's targets and generates the + catalog. Any failure (no provider, listing error, no OpenAI targets) returns None + so config writing proceeds without the ``model_catalog_json`` pin. + """ + if not provider: + return None + workspace = state.get("workspace") + if not workspace: + return None + try: + token = get_databricks_token(workspace, state.get("profile")) + service, _ = get_model_provider_service(provider, workspace, token) + except Exception: + return None + if not service: + return None + targets = service.get("targets") or [] + return build_model_catalog(targets) if targets else None + + def write_tool_config(state: dict, model: str | None = None, provider: str | None = None) -> dict: workspace = state["workspace"] # Leave model selection to Codex — except when a provider is set and a target @@ -352,6 +465,11 @@ def write_tool_config(state: dict, model: str | None = None, provider: str | Non use_pat=bool(state.get("use_pat")), provider=provider, ) + # Under a Bedrock MPS, pin an augmented catalog so Codex has metadata for the + # provider-side target ids (the gateway's codex/v1/models endpoint is disabled). + catalog_path = _provider_catalog_path(state, provider) + if catalog_path is not None: + overlay["model_catalog_json"] = str(catalog_path) def compose(base: dict) -> dict: deep_merge_dict(base, copy.deepcopy(overlay)) @@ -359,6 +477,9 @@ def compose(base: dict) -> dict: if chosen_model is None: for key in ("model", "model_reasoning_effort"): base.pop(key, None) + # Drop a stale catalog pin when this run isn't routing through a provider. + if catalog_path is None: + base.pop("model_catalog_json", None) return base doc = read_toml_safe(CODEX_CONFIG_PATH) diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index cbd6c051..478878fc 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os from pathlib import Path @@ -188,6 +189,7 @@ def test_provider_writes_header_and_drops_stale_model(self, tmp_path, monkeypatc monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") monkeypatch.setattr(codex, "save_state", lambda state: None) + monkeypatch.setattr(codex, "_provider_catalog_path", lambda *a, **k: None) # An earlier run pinned a model. codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) @@ -376,6 +378,116 @@ def test_legacy_write_preserves_other_profiles_in_shared_config(self, tmp_path, assert doc["profiles"]["ucode"]["model_provider"] == "ucode-databricks" +_BUILTIN_CATALOG = { + "models": [ + {"slug": "gpt-5.6-luna", "display_name": "GPT-5.6-Luna", "context_window": 272000}, + {"slug": "gpt-5.6-sol", "display_name": "GPT-5.6-Sol", "context_window": 272000}, + ] +} + + +class TestCodexCanonicalSlug: + def test_strips_region_and_vendor(self): + assert codex._canonical_codex_slug("us.openai.gpt-5.6-luna") == "gpt-5.6-luna" + + def test_strips_vendor_without_region(self): + assert codex._canonical_codex_slug("openai.gpt-oss-120b-1:0") == "gpt-oss-120b-1:0" + + def test_leaves_non_openai_unchanged(self): + assert ( + codex._canonical_codex_slug("anthropic.claude-opus-4-8") == "anthropic.claude-opus-4-8" + ) + assert codex._canonical_codex_slug("us.xai.grok-4.6") == "us.xai.grok-4.6" + + +class TestCodexBuildModelCatalog: + def _patch(self, tmp_path, monkeypatch, catalog=_BUILTIN_CATALOG): + monkeypatch.setattr(codex, "CODEX_MODEL_CATALOG_PATH", tmp_path / "catalog.json") + monkeypatch.setattr( + codex, "_codex_builtin_catalog", lambda: json.loads(json.dumps(catalog)) + ) + + def test_clones_openai_targets_and_skips_others(self, tmp_path, monkeypatch): + self._patch(tmp_path, monkeypatch) + path = codex.build_model_catalog( + ["us.openai.gpt-5.6-luna", "anthropic.claude-opus-4-8", "us.xai.grok-4.6"] + ) + assert path is not None + slugs = [m["slug"] for m in json.loads(path.read_text())["models"]] + assert "us.openai.gpt-5.6-luna" in slugs # cloned + assert "anthropic.claude-opus-4-8" not in slugs # no codex metadata to borrow + assert "us.xai.grok-4.6" not in slugs + + def test_clone_copies_canonical_metadata(self, tmp_path, monkeypatch): + self._patch(tmp_path, monkeypatch) + path = codex.build_model_catalog(["us.openai.gpt-5.6-luna"]) + clone = next( + m + for m in json.loads(path.read_text())["models"] + if m["slug"] == "us.openai.gpt-5.6-luna" + ) + assert clone["context_window"] == 272000 + assert clone["display_name"] == "GPT-5.6-Luna" + + def test_idempotent_no_duplicate_clones(self, tmp_path, monkeypatch): + self._patch(tmp_path, monkeypatch) + codex.build_model_catalog(["us.openai.gpt-5.6-luna"]) + path = codex.build_model_catalog(["us.openai.gpt-5.6-luna"]) + slugs = [m["slug"] for m in json.loads(path.read_text())["models"]] + assert slugs.count("us.openai.gpt-5.6-luna") == 1 + + def test_no_matching_targets_returns_none(self, tmp_path, monkeypatch): + self._patch(tmp_path, monkeypatch) + assert codex.build_model_catalog(["anthropic.claude-opus-4-8"]) is None + + def test_returns_none_when_codex_catalog_unavailable(self, tmp_path, monkeypatch): + monkeypatch.setattr(codex, "CODEX_MODEL_CATALOG_PATH", tmp_path / "catalog.json") + monkeypatch.setattr(codex, "_codex_builtin_catalog", lambda: None) + assert codex.build_model_catalog(["us.openai.gpt-5.6-luna"]) is None + + +class TestCodexProviderCatalogIntegration: + def _patch(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + return config_path + + def test_provider_pins_model_catalog_json(self, tmp_path, monkeypatch): + config_path = self._patch(tmp_path, monkeypatch) + catalog = tmp_path / "ucode-model-catalog.json" + monkeypatch.setattr(codex, "_provider_catalog_path", lambda state, provider: catalog) + codex.write_tool_config({"workspace": WS}, provider="eng.ai.bedrock") + assert read_toml_safe(config_path)["model_catalog_json"] == str(catalog) + + def test_no_provider_drops_stale_catalog_pin(self, tmp_path, monkeypatch): + config_path = self._patch(tmp_path, monkeypatch) + config_path.parent.mkdir(parents=True, exist_ok=True) + config_path.write_text('model_catalog_json = "/old/catalog.json"\n', encoding="utf-8") + monkeypatch.setattr(codex, "_provider_catalog_path", lambda state, provider: None) + codex.write_tool_config({"workspace": WS}) + assert "model_catalog_json" not in read_toml_safe(config_path) + + +class TestClearModelPreferencesWithMps: + def test_skips_clear_when_mps_header_present(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + config_path.parent.mkdir() + config_path.write_text( + 'model = "us.openai.gpt-5.6-luna"\n' + "[model_providers.ucode-databricks.http_headers]\n" + 'Databricks-Model-Provider-Service = "eng.ai.bedrock"\n', + encoding="utf-8", + ) + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", tmp_path / "backup.toml") + + assert codex.clear_model_preferences({}) is False + assert read_toml_safe(config_path)["model"] == "us.openai.gpt-5.6-luna" + + class TestCodexLegacyLayoutDetection: def test_new_codex_uses_modern_layout(self, monkeypatch): monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0")