Add support for fable via ucode configure --enable-fable or ucode configure --disable-fable (#213)

Author: rohita5l

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

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(

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_<FAMILY>_MODEL). Add an entry to
+# support a new family in both discovery paths (`claude-<family>-*` via the
+# model-services listing and `databricks-claude-<family>-*` 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
@@ -1226,8 +1232,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 +1248,7 @@ def discover_model_services(
         return {}, [], [], [], reason
 
     claude_models: dict[str, str] = {}
-    for family in ("opus", "sonnet", "haiku"):
+    for family in ANTHROPIC_FAMILIES:
         candidates = sorted(
             [m for m in ids if f"claude-{family}-" in m],
             reverse=True,
@@ -1797,21 +1804,22 @@ 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 ANTHROPIC_FAMILIES:
         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:
         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-{{opus,sonnet,haiku}}-*` (got: {sample})"
+        f"`databricks-claude-{{{families}}}-*` (got: {sample})"
     )
 
 

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 (