Support custom Claude model endpoints via --model-prefix

README.md

@@ -61,6 +61,25 @@ ucode configure --workspaces https://first.databricks.com,https://second.databri
 
 When multiple workspaces are provided, `ucode` logs into and saves state for each workspace. Launch commands such as `ucode codex` use the first workspace in the list.
 
+### Custom Claude model endpoints
+
+By default `ucode` discovers Claude models named `databricks-claude-<family>-<version>` on the
+AI Gateway anthropic route and picks the newest per family (`opus`, `sonnet`, `haiku`). If your
+workspace exposes Claude through custom serving endpoints — for example Bedrock-backed models
+named `acme-bedrock-claude-opus-4-8` — pass `--model-prefix` with the shared prefix before the
+family token:
+
+```bash
+ucode configure --agents claude --model-prefix acme-bedrock-claude-
+```
+
+This scopes discovery to those endpoints (the built-in `databricks-claude-*` hosted models no
+longer match the prefix, so they are excluded), picks the newest version per family as the
+default, and keeps every matching endpoint selectable. `ucode` writes the matching set to Claude
+Code's `availableModels` allowlist so the `/model` picker is restricted to them, and persists the
+prefix to state so later `ucode configure` runs reuse it without the flag. `ucode status` shows
+the resolved prefix, per-family defaults, and the number of selectable models.
+
 ### MCP servers (optional)
 
 ```bash

src/ucode/agents/claude.py

@@ -17,6 +17,7 @@
     write_json_file,
 )
 from ucode.databricks import (
+    CLAUDE_FAMILIES,
     build_auth_shell_command,
     build_tool_base_url,
     get_databricks_token,
@@ -81,12 +82,38 @@ def _web_search_mcp_entry(workspace: str, search_model: str, profile: str | None
     }
 
 
+def _build_available_models(claude_models: dict[str, str], claude_allowed: list[str]) -> list[str]:
+    """Allowlist written to Claude Code's ``availableModels`` setting.
+
+    Combines the discovered endpoint ids (every allowed version) with the
+    family aliases (``opus``/``sonnet``/``haiku``) that have a default, so the
+    friendly family picker entries stay visible and resolve through the
+    ``ANTHROPIC_DEFAULT_*_MODEL`` env vars.
+
+    Including aliases intentionally loosens strictness — Claude Code's matcher
+    treats an ``opus`` alias as "any opus version", so a same-family model would
+    pass the allowlist. We accept that to keep the default picker entries usable;
+    the Databricks gateway remains the hard backstop. With an empty allowlist we
+    fall back to the flat defaults' values (old state shape / back-compat).
+    """
+    ids = [m for m in claude_allowed if isinstance(m, str) and m]
+    if not ids:
+        ids = [v for v in claude_models.values() if isinstance(v, str) and v]
+    aliases = [family for family in CLAUDE_FAMILIES if claude_models.get(family)]
+    out: list[str] = []
+    for item in [*aliases, *sorted(set(ids))]:
+        if item not in out:
+            out.append(item)
+    return out
+
+
 def render_overlay(
     workspace: str,
     model: str,
     claude_models: dict[str, str] | None = None,
     disable_web_search: bool = False,
     profile: str | None = None,
+    claude_allowed: list[str] | None = None,
 ) -> tuple[dict, list[list[str]]]:
     """Return (overlay, managed_key_paths) for Claude settings.json.
 
@@ -121,6 +148,14 @@ def render_overlay(
     overlay: dict = {"apiKeyHelper": build_auth_shell_command(workspace, profile), "env": env}
     keys: list[list[str]] = [["apiKeyHelper"]] + [["env", k] for k in env]
 
+    # Restrict the `/model` picker to the discovered models via Claude Code's
+    # `availableModels` allowlist, so hosted models that don't match the
+    # configured prefix aren't selectable.
+    available_models = _build_available_models(claude_models or {}, claude_allowed or [])
+    if available_models:
+        overlay["availableModels"] = available_models
+        keys.append(["availableModels"])
+
     # Disable Claude Code's built-in WebSearch (it routes through Anthropic's
     # hosted infra and fails through the Databricks gateway). The replacement
     # `web_search` MCP server is registered separately via the claude CLI.
@@ -197,6 +232,7 @@ def write_tool_config(state: dict, model: str) -> dict:
         state.get("claude_models") or {},
         disable_web_search=web_search_model is not None,
         profile=state.get("profile"),
+        claude_allowed=state.get("claude_allowed_models") or [],
     )
     existing = read_json_safe(CLAUDE_SETTINGS_PATH)
     merged = deep_merge_dict(existing, overlay)

src/ucode/cli.py

@@ -28,6 +28,8 @@
 from ucode.agents.pi import PI_SETTINGS_BACKUP_PATH, PI_SETTINGS_PATH
 from ucode.config_io import restore_file, set_dry_run
 from ucode.databricks import (
+    CLAUDE_FAMILIES,
+    DEFAULT_CLAUDE_MODEL_PREFIX,
     build_shared_base_urls,
     discover_claude_models,
     discover_codex_models,
@@ -39,6 +41,7 @@
     get_databricks_token,
     install_databricks_cli,
     normalize_workspace_url,
+    resolve_claude_model_prefix,
     run_databricks_login,
 )
 from ucode.mcp import (
@@ -47,7 +50,7 @@
     purge_cross_workspace_mcp_residue,
     revert_mcp_configs,
 )
-from ucode.state import STATE_PATH, clear_state, load_state, save_state
+from ucode.state import STATE_PATH, clear_state, load_full_state, load_state, save_state
 from ucode.ui import (
     console,
     heading,
@@ -146,6 +149,7 @@ def configure_shared_state(
     profile: str | None = None,
     tools: list[str] | None = None,
     force_login: bool = False,
+    model_prefix: str | None = None,
 ) -> dict:
     """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state.
 
@@ -178,12 +182,22 @@ def configure_shared_state(
     want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools
     want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools
 
+    # Resolve the Claude endpoint name prefix (--model-prefix flag > this
+    # workspace's persisted value > built-in default) so discovery scopes to the
+    # right family of endpoints — e.g. Bedrock-backed `acme-bedrock-claude-*`
+    # instead of hosted `databricks-claude-*`.
+    existing_ws_state = load_full_state().get("workspaces", {}).get(workspace, {})
+    claude_model_prefix = resolve_claude_model_prefix(existing_ws_state, override=model_prefix)
+
     claude_reason: str | None = None
     gemini_reason: str | None = None
     codex_reason: str | None = None
+    claude_allowed: list[str] = []
     with spinner("Fetching available models..."):
         if want_claude:
-            claude_models, claude_reason = discover_claude_models(workspace, token)
+            claude_models, claude_allowed, claude_reason = discover_claude_models(
+                workspace, token, claude_model_prefix
+            )
         else:
             claude_models = {}
         if want_gemini:
@@ -196,7 +210,7 @@ def configure_shared_state(
             codex_models = []
     opencode_models: dict[str, list[str]] = {}
     if claude_models:
-        opencode_models["anthropic"] = list(claude_models.values())
+        opencode_models["anthropic"] = claude_allowed or list(claude_models.values())
     if gemini_models:
         opencode_models["gemini"] = gemini_models
 
@@ -210,6 +224,8 @@ def configure_shared_state(
     state["base_urls"] = build_shared_base_urls(workspace)
     if want_claude:
         state["claude_models"] = claude_models
+        state["claude_allowed_models"] = claude_allowed
+        state["claude_model_prefix"] = claude_model_prefix
     if want_gemini:
         state["gemini_models"] = gemini_models
     if want_codex:
@@ -236,13 +252,20 @@ def _configure_shared_workspace_states(
     tools: list[str] | None,
     *,
     force_login: bool,
+    model_prefix: str | None = None,
 ) -> list[dict]:
     if not workspaces:
         raise RuntimeError("At least one workspace must be provided.")
     states: list[dict] = []
     for workspace, profile in workspaces:
         states.append(
-            configure_shared_state(workspace, profile=profile, tools=tools, force_login=force_login)
+            configure_shared_state(
+                workspace,
+                profile=profile,
+                tools=tools,
+                force_login=force_login,
+                model_prefix=model_prefix,
+            )
         )
     return states
 
@@ -251,14 +274,17 @@ def configure_workspace_command(
     tool: str | None = None,
     selected_tools: list[str] | None = None,
     workspaces: list[tuple[str, str | None]] | None = None,
+    model_prefix: str | None = None,
 ) -> int:
     if tool is not None and selected_tools is not None:
         raise RuntimeError("Use either --agent or --agents, not both.")
 
     workspace_entries = workspaces or [_prompt_for_configuration(tool)]
 
     if tool is not None:
-        states = _configure_shared_workspace_states(workspace_entries, [tool], force_login=True)
+        states = _configure_shared_workspace_states(
+            workspace_entries, [tool], force_login=True, model_prefix=model_prefix
+        )
         state = states[0]
         state = configure_single_tool(tool, state)
         spec = TOOL_SPECS[tool]
@@ -285,7 +311,9 @@ def configure_workspace_command(
             raise RuntimeError(f"{spec['display']} validation failed — config reverted.")
         return 0
 
-    states = _configure_shared_workspace_states(workspace_entries, selected_tools, force_login=True)
+    states = _configure_shared_workspace_states(
+        workspace_entries, selected_tools, force_login=True, model_prefix=model_prefix
+    )
     state = states[0]
     save_state(state)
 
@@ -362,6 +390,23 @@ def status() -> int:
     if profile:
         print_kv("CLI profile", profile)
 
+    claude_models = state.get("claude_models") or {}
+    if claude_models:
+        print_heading("Claude Models")
+        prefix = state.get("claude_model_prefix")
+        if isinstance(prefix, str) and prefix and prefix != DEFAULT_CLAUDE_MODEL_PREFIX:
+            print_kv("Endpoint prefix", prefix)
+        allowed = [m for m in (state.get("claude_allowed_models") or []) if isinstance(m, str)]
+        for family in CLAUDE_FAMILIES:
+            default = claude_models.get(family)
+            if not default:
+                continue
+            extras = [m for m in allowed if m != default and f"-{family}-" in m]
+            label = f"{default}  (+{len(extras)} more)" if extras else default
+            print_kv(family.capitalize(), label)
+        if allowed:
+            print_kv("Selectable models", str(len(allowed)))
+
     print_heading("Coding Agents")
     for tool, spec in TOOL_SPECS.items():
         configured = tool in configured_tools
@@ -591,6 +636,17 @@ def configure(
             help="Configure a comma-separated list of workspaces without prompting.",
         ),
     ] = None,
+    model_prefix: Annotated[
+        str | None,
+        typer.Option(
+            "--model-prefix",
+            help=(
+                "Prefix for custom Claude endpoint names before the family token "
+                "(e.g. acme-bedrock-claude-). Scopes discovery to those endpoints and "
+                "excludes hosted databricks-claude-* models. Persisted for later runs."
+            ),
+        ),
+    ] = None,
 ) -> None:
     """Configure workspace URL and AI Gateway."""
     if ctx.invoked_subcommand is not None:
@@ -605,24 +661,30 @@ def configure(
             tool = normalize_tool(agent)
             install_tool_binary(tool, strict=True, update_existing=True)
             if workspace_entries is None:
-                configure_workspace_command(tool)
+                configure_workspace_command(tool, model_prefix=model_prefix)
             else:
-                configure_workspace_command(tool, workspaces=workspace_entries)
+                configure_workspace_command(
+                    tool, workspaces=workspace_entries, model_prefix=model_prefix
+                )
         elif agents is not None:
             selected_tools = _parse_agents_option(agents)
             if workspace_entries is None:
-                configure_workspace_command(selected_tools=selected_tools)
+                configure_workspace_command(
+                    selected_tools=selected_tools, model_prefix=model_prefix
+                )
             else:
                 configure_workspace_command(
-                    selected_tools=selected_tools, workspaces=workspace_entries
+                    selected_tools=selected_tools,
+                    workspaces=workspace_entries,
+                    model_prefix=model_prefix,
                 )
         else:
             # Tool binaries are installed after the user picks which agents
             # they want, in configure_workspace_command.
             if workspace_entries is None:
-                configure_workspace_command()
+                configure_workspace_command(model_prefix=model_prefix)
             else:
-                configure_workspace_command(workspaces=workspace_entries)
+                configure_workspace_command(workspaces=workspace_entries, model_prefix=model_prefix)
     except RuntimeError as exc:
         print_err(str(exc))
         raise typer.Exit(1) from None