From 0109b51d1bee2bf4c645837b9cd3cf99677ae2f2 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 15 Jul 2026 22:08:48 +0200 Subject: [PATCH 1/3] fable --- src/ucode/agents/claude.py | 12 ++++ src/ucode/cli.py | 46 +++++++++++++++ src/ucode/databricks.py | 13 +++-- tests/test_agent_claude.py | 31 ++++++++++ tests/test_cli.py | 116 +++++++++++++++++++++++++++++++++++++ tests/test_databricks.py | 18 +++++- 6 files changed, 229 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 422e763..f1e62e5 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -84,6 +84,8 @@ def _resolve_web_search_model(state: dict) -> str | None: # every launch, so stale values from older ucode versions never linger. CLAUDE_MANAGED_MODEL_ENV_KEYS = ( "ANTHROPIC_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL", + "ANTHROPIC_DEFAULT_FABLE_MODEL_NAME", "ANTHROPIC_DEFAULT_OPUS_MODEL", "ANTHROPIC_DEFAULT_OPUS_MODEL_NAME", "ANTHROPIC_DEFAULT_SONNET_MODEL", @@ -134,6 +136,7 @@ def render_overlay( use_pat: bool = False, provider: str | None = None, provider_models: dict[str, str] | None = None, + fable_enabled: bool = False, ) -> tuple[dict, list[list[str]]]: """Return (overlay, managed_key_paths) for Claude settings.json. @@ -192,6 +195,14 @@ def render_overlay( # so users can see which gateway-routable model is behind each shortcut. # We deliberately don't set the `_NAME` companion env vars — the raw id # is more useful than a friendly label for debugging gateway routing. + # + # Fable is opt-in only (`ucode configure --enable-fable`): it's a premium + # model, so we don't pin the family alias unless the user asked for it. + # When off, ANTHROPIC_DEFAULT_FABLE_MODEL is simply never written — and + # since it's in CLAUDE_MANAGED_MODEL_ENV_KEYS, any stale value from a + # prior `--enable-fable` run is pruned from settings.json on next launch. + if fable_enabled and claude_models.get("fable"): + env["ANTHROPIC_DEFAULT_FABLE_MODEL"] = claude_models["fable"] if claude_models.get("opus"): env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = _maybe_add_1m_suffix(claude_models["opus"]) if claude_models.get("sonnet"): @@ -296,6 +307,7 @@ def write_tool_config( use_pat=bool(state.get("use_pat")), provider=provider, provider_models=provider_models, + fable_enabled=bool(state.get("fable_enabled")), ) tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 9ebdae9..a0ef7c4 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -217,6 +217,7 @@ def configure_shared_state( use_pat: bool | None = None, skip_model_discovery: bool = False, skip_preflight: bool = False, + fable_enabled: bool | None = None, ) -> dict: """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. @@ -236,12 +237,18 @@ def configure_shared_state( in ``_launch_tool``) and the gateway was verified by that earlier configure. Only the local profile resolution and the shared state assembly still run; the saved model lists are preserved. + ``fable_enabled`` opts the premium Claude Fable family into Claude Code's + ``ANTHROPIC_DEFAULT_FABLE_MODEL`` pin (default off). ``None`` means "inherit": + a launch re-run keeps whatever the workspace was configured with; ``True``/ + ``False`` come from an explicit ``configure --enable-fable``/``--disable-fable``. """ workspace = normalize_workspace_url(workspace) prior_state = load_state() previous_workspace = prior_state.get("workspace") if use_pat is None: use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace + if fable_enabled is None: + fable_enabled = bool(prior_state.get("fable_enabled")) and previous_workspace == workspace fetch_all = tools is None # Assemble the shared workspace state that doesn't depend on model discovery: @@ -265,6 +272,12 @@ def configure_shared_state( state["use_pat"] = True else: state.pop("use_pat", None) + # Persist the Fable opt-in so launches keep pinning the family; an explicit + # `configure --disable-fable` (fable_enabled=False) clears it. + if fable_enabled: + state["fable_enabled"] = True + else: + state.pop("fable_enabled", None) state["base_urls"] = build_shared_base_urls(workspace) if skip_preflight: @@ -360,6 +373,12 @@ def configure_shared_state( claude_models, claude_reason = ms_claude, ms_reason if not claude_models: claude_models, claude_reason = discover_claude_models(workspace, token) + # Fable is opt-in (`configure --enable-fable`). Unless enabled, + # drop it from the discovered bundle entirely so it never becomes + # part of any agent's config — not claude's family pins, nor the + # opencode/pi/copilot model lists built from claude_models. + if not fable_enabled: + claude_models.pop("fable", None) if want_gemini: gemini_models, gemini_reason = ms_gemini, ms_reason if not gemini_models: @@ -416,6 +435,7 @@ def _configure_shared_workspace_states( *, force_login: bool, use_pat: bool = False, + fable_enabled: bool | None = None, ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") @@ -428,6 +448,7 @@ def _configure_shared_workspace_states( tools=tools, force_login=force_login, use_pat=use_pat, + fable_enabled=fable_enabled, ) ) return states @@ -507,6 +528,7 @@ def configure_workspace_command( prompt_optional_updates: bool = True, use_pat: bool = False, skip_validate: bool = False, + fable_enabled: bool | None = None, ) -> int: if tool is not None and selected_tools is not None: raise RuntimeError("Use either --agent or --agents, not both.") @@ -524,6 +546,7 @@ def configure_workspace_command( [tool], force_login=True, use_pat=use_pat, + fable_enabled=fable_enabled, ) state = states[0] state = configure_single_tool(tool, state) @@ -560,6 +583,7 @@ def configure_workspace_command( selected_tools, force_login=True, use_pat=use_pat, + fable_enabled=fable_enabled, ) state = states[0] save_state(state) @@ -1066,6 +1090,18 @@ def configure( "freshly discovered models.", ), ] = False, + enable_fable: Annotated[ + bool | None, + typer.Option( + "--enable-fable/--disable-fable", + help="Pin the premium Claude Fable family via ANTHROPIC_DEFAULT_FABLE_MODEL " + "for Claude Code (opt-in; off by default). Only takes effect when the " + "workspace's AI Gateway actually advertises a Claude Fable model. " + "--disable-fable clears a prior opt-in. Omitting both keeps the " + "workspace's existing setting. Passed on its own (no --agent/--agents), " + "it configures Claude Code directly since Fable is Claude-only.", + ), + ] = None, tracing: Annotated[ bool, typer.Option( @@ -1121,6 +1157,16 @@ def configure( skip_kwargs["use_pat"] = True if skip_validate: skip_kwargs["skip_validate"] = True + # Only forward the Fable opt-in when the user passed the flag; `None` + # (neither flag given) lets configure_shared_state inherit the prior + # workspace setting instead of clobbering it. + if enable_fable is not None: + skip_kwargs["fable_enabled"] = enable_fable + # Fable is a Claude-only model family, so `--enable-fable`/`--disable-fable` + # only makes sense for Claude Code. When passed on its own, implicitly + # target claude instead of dropping into the interactive agent picker. + if enable_fable is not None and agent is None and agents is None: + agent = "claude" if agent is not None: tool = normalize_tool(agent) install_tool_binary( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..091feb4 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1226,8 +1226,9 @@ def discover_model_services( Returns (claude_models, codex_models, gemini_models, oss_models, reason): - - ``claude_models`` maps ``opus``/``sonnet``/``haiku`` to the newest - matching ``system.ai.claude-*`` id (mirrors ``discover_claude_models``). + - ``claude_models`` maps ``fable``/``opus``/``sonnet``/``haiku`` to the + newest matching ``system.ai.claude-*`` id (mirrors + ``discover_claude_models``). - ``codex_models`` is the list of ``system.ai.*gpt-*`` ids. - ``gemini_models`` is the list of ``system.ai.*gemini-*`` ids, newest first. - ``oss_models`` is the list of OSS-model ``system.ai.*`` ids. @@ -1241,7 +1242,7 @@ def discover_model_services( return {}, [], [], [], reason claude_models: dict[str, str] = {} - for family in ("opus", "sonnet", "haiku"): + for family in ("fable", "opus", "sonnet", "haiku"): candidates = sorted( [m for m in ids if f"claude-{family}-" in m], reverse=True, @@ -1797,13 +1798,13 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], ] result: dict[str, str] = {} - for family, key in [("opus", "opus"), ("sonnet", "sonnet"), ("haiku", "haiku")]: + for family in ("fable", "opus", "sonnet", "haiku"): candidates = sorted( [m for m in raw_ids if f"databricks-claude-{family}-" in m], reverse=True, ) if candidates: - result[key] = candidates[0] + result[family] = candidates[0] if result: return result, None if not raw_ids: @@ -1811,7 +1812,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], sample = ", ".join(raw_ids[:5]) return {}, ( "AI Gateway returned model ids but none matched " - f"`databricks-claude-{{opus,sonnet,haiku}}-*` (got: {sample})" + f"`databricks-claude-{{fable,opus,sonnet,haiku}}-*` (got: {sample})" ) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index b2d3a27..c7b5b32 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -113,6 +113,37 @@ def test_model_overrides_not_set_when_no_models(self): env = overlay["env"] assert "ANTHROPIC_DEFAULT_SONNET_MODEL" not in env + def test_fable_not_pinned_by_default(self): + # Fable is opt-in: even when the workspace advertises it, the env var is + # absent unless fable_enabled is passed. + models = {"fable": "databricks-claude-fable-5", "opus": "databricks-claude-opus-4-8"} + overlay, _ = claude.render_overlay(WS, "s4", claude_models=models) + env = overlay["env"] + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in env + assert env["ANTHROPIC_DEFAULT_OPUS_MODEL"] == "databricks-claude-opus-4-8[1m]" + + def test_fable_pinned_when_enabled_and_discovered(self): + models = {"fable": "system.ai.claude-fable-5"} + overlay, _ = claude.render_overlay(WS, "s4", claude_models=models, fable_enabled=True) + env = overlay["env"] + # Fable 5 is 1M-context by default, so no `[1m]` suffix is appended. + assert env["ANTHROPIC_DEFAULT_FABLE_MODEL"] == "system.ai.claude-fable-5" + + def test_fable_not_pinned_when_enabled_but_not_discovered(self): + # --enable-fable is a no-op when the workspace advertises no fable model, + # mirroring the opus/sonnet/haiku "only if discovered" behavior. + models = {"opus": "databricks-claude-opus-4-8"} + overlay, _ = claude.render_overlay(WS, "s4", claude_models=models, fable_enabled=True) + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in overlay["env"] + + def test_fable_not_pinned_under_provider(self): + # A Model Provider Service routes by header and pins no Databricks model. + models = {"fable": "databricks-claude-fable-5"} + overlay, _ = claude.render_overlay( + WS, "s4", claude_models=models, fable_enabled=True, provider="main.x.claude-svc" + ) + assert "ANTHROPIC_DEFAULT_FABLE_MODEL" not in overlay["env"] + def test_provider_adds_routing_header(self): overlay, _ = claude.render_overlay(WS, "s4", provider="main.aarushi.aarushi-claude") assert ( diff --git a/tests/test_cli.py b/tests/test_cli.py index 7bdb0e4..17d9ef2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -594,6 +594,52 @@ def test_agent_flag_calls_configure_with_tool(self): ) mock_cfg.assert_called_once_with("claude") + def test_disable_fable_alone_implicitly_targets_claude(self): + # Fable is Claude-only, so `--disable-fable` on its own should configure + # claude directly instead of dropping into the interactive agent picker. + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary") as mock_install, + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke(app, ["configure", "--disable-fable"]) + assert result.exit_code == 0, result.output + mock_install.assert_called_once_with( + "claude", strict=True, update_existing=True, prompt_optional_updates=True + ) + mock_cfg.assert_called_once_with("claude", fable_enabled=False) + + def test_enable_fable_alone_implicitly_targets_claude(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary") as mock_install, + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke(app, ["configure", "--enable-fable"]) + assert result.exit_code == 0, result.output + mock_install.assert_called_once_with( + "claude", strict=True, update_existing=True, prompt_optional_updates=True + ) + mock_cfg.assert_called_once_with("claude", fable_enabled=True) + + def test_enable_fable_with_explicit_agents_does_not_override(self): + # An explicit --agents selection wins; the fable flag rides along without + # forcing the claude-only single-agent path. + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke( + app, ["configure", "--enable-fable", "--agents", "claude,codex"] + ) + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + selected_tools=["claude", "codex"], + prompt_optional_updates=True, + fable_enabled=True, + ) + def test_skip_upgrade_flag_disables_optional_update_prompt(self): with ( patch("ucode.cli.install_databricks_cli"), @@ -815,6 +861,7 @@ def fake_configure_shared_state( tools=None, force_login=False, use_pat=False, + fable_enabled=None, ): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) @@ -1140,6 +1187,75 @@ def test_uc_models_used_without_legacy_fallback(self, monkeypatch): assert legacy_called == [] assert "uc_enabled" not in state + def _stub_with_fable(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ( + {"fable": "system.ai.claude-fable-5", "opus": "system.ai.claude-opus-4-8"}, + [], + [], + [], + None, + ), + ) + return cli_mod + + def test_fable_stripped_from_discovery_when_not_enabled(self, monkeypatch): + # Discovery buckets fable, but without --enable-fable it's dropped from + # the persisted bundle so it never reaches any agent's config. + cli_mod = self._stub_with_fable(monkeypatch) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["claude_models"] == {"opus": "system.ai.claude-opus-4-8"} + assert "fable_enabled" not in state + + def test_fable_retained_and_persisted_when_enabled(self, monkeypatch): + cli_mod = self._stub_with_fable(monkeypatch) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT", fable_enabled=True) + + assert state["claude_models"]["fable"] == "system.ai.claude-fable-5" + assert state["fable_enabled"] is True + + def test_launch_inherits_persisted_fable_opt_in(self, monkeypatch): + # A launch re-run passes fable_enabled=None; the persisted opt-in for the + # same workspace applies, so fable stays in the discovered bundle. + cli_mod, *_ = self._stub_deps( + monkeypatch, + pat_token="dapi-pat", + existing_state={"workspace": self.WS, "profile": "DEFAULT", "fable_enabled": True}, + ) + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({"fable": "system.ai.claude-fable-5"}, [], [], [], None), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT") + + assert state["claude_models"]["fable"] == "system.ai.claude-fable-5" + assert state["fable_enabled"] is True + + def test_reconfigure_with_disable_fable_clears_opt_in(self, monkeypatch): + cli_mod, *_ = self._stub_deps( + monkeypatch, + pat_token="dapi-pat", + existing_state={"workspace": self.WS, "profile": "DEFAULT", "fable_enabled": True}, + ) + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ({"fable": "system.ai.claude-fable-5"}, [], [], [], None), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT", fable_enabled=False) + + assert "fable_enabled" not in state + assert "fable" not in state["claude_models"] + def test_falls_back_to_legacy_when_uc_empty(self, monkeypatch): # No UC model-services: each family falls back to the legacy listing. cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") diff --git a/tests/test_databricks.py b/tests/test_databricks.py index f59aa1d..40af048 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -134,6 +134,20 @@ def test_selects_opus_4_8_when_advertised(self, monkeypatch): assert reason is None assert models["opus"] == "databricks-claude-opus-4-8" + def test_buckets_fable_family(self, monkeypatch): + payload = { + "data": [ + {"id": "databricks-claude-fable-5"}, + {"id": "databricks-claude-opus-4-8"}, + ] + } + monkeypatch.setattr(db_mod, "_http_get_json", lambda url, token: (payload, None)) + + models, reason = db_mod.discover_claude_models(WS, "token") + + assert reason is None + assert models["fable"] == "databricks-claude-fable-5" + def _model_service(model_id: str) -> dict: """A model-services entry whose `name` strips to `model_id`.""" @@ -161,6 +175,7 @@ class TestDiscoverModelServices: def test_buckets_families_by_name(self, monkeypatch): payload = { "model_services": [ + _model_service("system.ai.claude-fable-5"), _model_service("system.ai.claude-opus-4-7"), _model_service("system.ai.claude-opus-4-8"), _model_service("system.ai.claude-sonnet-4-6"), @@ -179,8 +194,9 @@ def test_buckets_families_by_name(self, monkeypatch): claude, codex, gemini, oss, reason = db_mod.discover_model_services(WS, "token") assert reason is None - # Newest opus wins; sonnet bucketed; haiku absent. + # Fable bucketed; newest opus wins; sonnet bucketed; haiku absent. assert claude == { + "fable": "system.ai.claude-fable-5", "opus": "system.ai.claude-opus-4-8", "sonnet": "system.ai.claude-sonnet-4-6", } From 2762c9bf80cd720860ee373cb63add1ed6758084 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Thu, 16 Jul 2026 22:07:46 +0200 Subject: [PATCH 2/3] ruff --- tests/test_cli.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 17d9ef2..76a2ec5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -630,9 +630,7 @@ def test_enable_fable_with_explicit_agents_does_not_override(self): patch("ucode.cli.install_tool_binary"), patch("ucode.cli.configure_workspace_command") as mock_cfg, ): - result = runner.invoke( - app, ["configure", "--enable-fable", "--agents", "claude,codex"] - ) + result = runner.invoke(app, ["configure", "--enable-fable", "--agents", "claude,codex"]) assert result.exit_code == 0, result.output mock_cfg.assert_called_once_with( selected_tools=["claude", "codex"], From 55b890b4a101d50785163665e0344f6b41de07f8 Mon Sep 17 00:00:00 2001 From: Lilly Luo Date: Thu, 16 Jul 2026 20:24:58 +0000 Subject: [PATCH 3/3] databricks: hoist Claude family list into ANTHROPIC_FAMILIES const Both discovery paths (discover_model_services + discover_claude_models) and the "no models matched" error string hardcoded the same ("fable", "opus", "sonnet", "haiku") tuple. Add a single ANTHROPIC_FAMILIES const (mirroring _OSS_MODEL_FAMILIES) so a new Claude family is added in one place and the two loops can't drift apart. Co-authored-by: Isaac --- src/ucode/databricks.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 091feb4..b38bca9 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1100,6 +1100,12 @@ def build_auth_shell_command( # support a new family. _OSS_MODEL_FAMILIES = ("kimi-", "glm-") +# Claude model families ucode buckets, newest tier first. Each maps to a +# Claude Code family alias (ANTHROPIC_DEFAULT__MODEL). Add an entry to +# support a new family in both discovery paths (`claude--*` via the +# model-services listing and `databricks-claude--*` via the AI Gateway). +ANTHROPIC_FAMILIES = ("fable", "opus", "sonnet", "haiku") + # Per-family token limits (context window + max output tokens). These are a # property of the model + its `/ai-gateway/mlflow/v1` route (the gateway rejects # requests whose output exceeds the cap), not of any one agent — so every agent @@ -1242,7 +1248,7 @@ def discover_model_services( return {}, [], [], [], reason claude_models: dict[str, str] = {} - for family in ("fable", "opus", "sonnet", "haiku"): + for family in ANTHROPIC_FAMILIES: candidates = sorted( [m for m in ids if f"claude-{family}-" in m], reverse=True, @@ -1798,7 +1804,7 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], ] result: dict[str, str] = {} - for family in ("fable", "opus", "sonnet", "haiku"): + for family in ANTHROPIC_FAMILIES: candidates = sorted( [m for m in raw_ids if f"databricks-claude-{family}-" in m], reverse=True, @@ -1810,9 +1816,10 @@ def discover_claude_models(workspace: str, token: str) -> tuple[dict[str, str], if not raw_ids: return {}, "AI Gateway returned no Claude model ids" sample = ", ".join(raw_ids[:5]) + families = ",".join(ANTHROPIC_FAMILIES) return {}, ( "AI Gateway returned model ids but none matched " - f"`databricks-claude-{{fable,opus,sonnet,haiku}}-*` (got: {sample})" + f"`databricks-claude-{{{families}}}-*` (got: {sample})" )