Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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
Expand Down
46 changes: 46 additions & 0 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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.")
Expand All @@ -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
Expand Down Expand Up @@ -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.")
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
20 changes: 14 additions & 6 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -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,
Expand Down Expand Up @@ -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})"
)


Expand Down
31 changes: 31 additions & 0 deletions tests/test_agent_claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
Loading
Loading