route launches by budget tier and show workspaces spend

src/ucode/cli.py

@@ -33,7 +33,7 @@
 )
 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 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,
@@ -56,17 +56,25 @@
     resolve_pat_token,
     run_databricks_login,
 )
+from ucode.managed_budget import (
+    budget_usage_percent,
+    recommendation_line,
+    render_budget_panel,
+)
 from ucode.managed_config import (
+    get_model_recommendation,
     load_managed_state,
     managed_agent_config_enabled,
     refresh_managed_config,
 )
 from ucode.managed_resolve import (
     managed_default_model,
     managed_enabled_tools,
+    managed_launch_model,
     managed_provider_service,
     managed_supplies_models,
     managed_unservable_models,
+    recommended_agent,
     resolve_state,
 )
 from ucode.managed_wizard import setup_command, show_command
@@ -1281,6 +1289,76 @@ def _fetch_managed_config(state: dict, *, skip_preflight: bool) -> dict | None:
         return refresh_managed_config(state)
 
 
+def _note_recommended_agent(recommendation: dict | None, tool: str) -> None:
+    """Say when the budget tier points at a different agent than the one being launched.
+
+    Launching any enabled agent is allowed, so this informs rather than blocks — and explains why
+    the session is not on the tier's model.
+    """
+    # The tier's own agent, not `recommended_agent`'s default_agent fallback: there is nothing to
+    # say when the config's baseline simply differs from what the developer asked for.
+    agent = (recommendation or {}).get("agent")
+    if agent == tool or agent not in TOOL_SPECS:
+        return
+    model = (recommendation or {}).get("model")
+    suffix = f" with {model}" if isinstance(model, str) and model else ""
+    print_note(
+        f"Your budget tier recommends {TOOL_SPECS[agent]['display']}{suffix}; "
+        f"launching {TOOL_SPECS[tool]['display']} as requested."
+    )
+
+
+def _fetch_budget_recommendation(
+    state: dict, managed: dict | None, *, skip_preflight: bool
+) -> dict | None:
+    """The agent and model the caller's budget tier allows, or None when there is no budget to read.
+
+    Enforcement is server-side, so a failed read only costs the recommendation: the config's own
+    ``default_model`` still applies and the launch proceeds.
+    """
+    # --dry-run resolves the agent from the last saved config alone, so it must not reach the
+    # control plane — mirroring the managed-config read, which is likewise skipped under --dry-run.
+    if managed is None or skip_preflight or is_dry_run():
+        return None
+    reason: str | None = None
+    recommendation = None
+    with spinner("Checking your budget..."):
+        try:
+            recommendation, reason = get_model_recommendation(
+                state["workspace"],
+                get_databricks_token(state["workspace"], state.get("profile")),
+            )
+        except (RuntimeError, OSError) as exc:
+            # A token that lapsed since the config refresh — or a Databricks CLI that isn't
+            # installed or reachable — must not block the launch; the config's default_model stands.
+            reason = str(exc)
+    if reason is not None:
+        print_warning(
+            f"Could not check your budget ({reason}); "
+            "using the default model from your workspace's config."
+        )
+    return recommendation
+
+
+def _print_budget_panel(recommendation: dict, tool: str, managed: dict | None = None) -> None:
+    """Show the workspace budget this launch spends against, when one is configured."""
+    agent = recommendation.get("agent")
+    display_agent = TOOL_SPECS[agent]["display"] if agent in TOOL_SPECS else None
+    percent = budget_usage_percent(
+        float(recommendation.get("current_spend") or 0.0),
+        float(recommendation.get("effective_threshold") or 0.0),
+    )
+    line = recommendation_line(display_agent, recommendation.get("model"), percent)
+    panel = render_budget_panel(
+        recommendation,
+        title=f"ucode with {TOOL_SPECS[tool]['display']}",
+        extra_lines=[line] if line else None,
+        managed=managed,
+    )
+    if panel is not None:
+        console.print(panel)
+
+
 def _launch_tool(
     tool_name: str,
     ctx: typer.Context,
@@ -1289,6 +1367,7 @@ def _launch_tool(
     workspace: str | None = None,
     enable_smart_routing_flag: bool = False,
     managed: dict | None = None,
+    recommendation: dict | None = None,
 ) -> None:
     try:
         tool = normalize_tool(tool_name)
@@ -1341,6 +1420,12 @@ def _launch_tool(
         # An admin-published managed config wins over the developer's own settings. Layered on after
         # `configure_shared_state`, whose returned state it overrides, and before the provider and
         # model are settled below — the two state files are never merged on disk.
+        # Bare `ucode` already read one to choose the agent; refetching would double the round trip.
+        if recommendation is None:
+            recommendation = _fetch_budget_recommendation(
+                state, managed, skip_preflight=skip_preflight
+            )
+        _note_recommended_agent(recommendation, tool)
         if managed is not None:
             state = resolve_state(managed, state, tool)
             print_success("Applied your workspace's managed coding agent config")
@@ -1413,7 +1498,9 @@ def _launch_tool(
             # in as the explicit model rather than being applied afterwards: for codex the proto has
             # no model list at all, so passing it here is the only way a launch succeeds when the
             # workspace's own discovery turned up nothing.
-            managed_model = managed_default_model(managed, tool) if managed is not None else None
+            managed_model = (
+                managed_launch_model(managed, recommendation, tool) if managed is not None else None
+            )
             state, resolved_model = resolve_launch_model(tool, state, managed_model)
             if routing_agent is not None and routing_agent.smart_routing_enabled(state):
                 display = TOOL_SPECS[tool]["display"]
@@ -1473,6 +1560,8 @@ def _launch_tool(
                 f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically "
                 f"every 30 minutes while the session is running."
             )
+        if recommendation is not None:
+            _print_budget_panel(recommendation, tool, managed)
         print_success(f"Starting {TOOL_SPECS[tool]['display']}")
         launch_agent(tool, state, ctx.args)
     except RuntimeError as exc:
@@ -1597,7 +1686,12 @@ def _launch_managed_default(
             return
         _print_no_managed_config_guidance(current, state.get("profile"))
         return
-    tool = managed.get("default_agent") or next(iter(managed.get("enabled_agents") or {}), None)
+    # The budget tier can move the org to a cheaper agent, so it outranks the config's
+    # default_agent. Fetched here and handed to _launch_tool so it is read once per launch.
+    recommendation = _fetch_budget_recommendation(state, managed, skip_preflight=skip_preflight)
+    tool = recommended_agent(recommendation, managed) or next(
+        iter(managed.get("enabled_agents") or {}), None
+    )
     if not isinstance(tool, str) or not tool:
         raise RuntimeError(
             "Your workspace's managed config names no agent to launch. Ask an admin to set a "
@@ -1610,6 +1704,7 @@ def _launch_managed_default(
         skip_preflight=skip_preflight,
         workspace=workspace,
         managed=managed,
+        recommendation=recommendation,
     )
 
 

src/ucode/databricks.py

@@ -1559,6 +1559,22 @@ def fetch_managed_coding_agent_configs(workspace: str, token: str) -> tuple[list
     return [c for c in configs if isinstance(c, dict)], None
 
 
+def fetch_model_recommendation(workspace: str, token: str) -> tuple[dict, str | None]:
+    """Ask the AI Gateway which agent and model the caller's budget tier allows.
+
+    The request takes no parameters: the server matches the caller's live spend against the managed
+    config's budget tiers and resolves the agent first, then that agent's model.
+    """
+    hostname = workspace_hostname(workspace)
+    url = f"https://{hostname}{_CODING_AGENT_CONFIGS_API_PATH}:recommendModel"
+    payload, reason = _http_post_json(url, token, {}, timeout=30)
+    if reason is not None:
+        return {}, reason
+    if not isinstance(payload, dict):
+        return {}, "recommendModel returned an unexpected response shape"
+    return payload, None
+
+
 # --- MCP services (parallel to model services) -----------------------------
 
 

src/ucode/managed_budget.py

@@ -0,0 +1,131 @@
+"""Render the workspace budget a managed coding-agent config is spending against.
+
+The figures come from the AI Gateway's ``:recommendModel`` response (``current_spend`` and
+``effective_threshold``), normalized by :func:`ucode.managed_config.get_model_recommendation`.
+Enforcement is entirely server-side — the gateway rejects over-budget requests — so this module
+only reports spend and never blocks a launch.
+"""
+
+from __future__ import annotations
+
+from rich.panel import Panel
+from rich.text import Text
+
+# Fallback amber point for a workspace whose policy defines no tier to derive one from.
+BUDGET_WARN_AT = 0.8
+
+_BUDGET_STATE_STYLE = {"ok": "green", "warn": "yellow", "exceeded": "red"}
+
+
+def budget_warn_fraction(managed: dict | None) -> float:
+    """The spend fraction at which to warn: the admin's lowest activating tier, else 0.8.
+
+    Tiers at 0 are skipped — they activate from the first dollar, so warning on one would leave the
+    panel permanently amber.
+    """
+    policy = (managed or {}).get("budget_policy")
+    tiers = policy.get("tiers") if isinstance(policy, dict) else None
+    fractions = [
+        float(pct)
+        for tier in (tiers if isinstance(tiers, list) else [])
+        if isinstance(tier, dict)
+        # `bool` is an `int` subclass, so it would otherwise read as a 0/1 fraction.
+        if isinstance(pct := tier.get("spending_percentage"), int | float)
+        and not isinstance(pct, bool)
+        and 0 < pct <= 1
+    ]
+    return min(fractions) if fractions else BUDGET_WARN_AT
+
+
+def budget_state(spend: float, threshold: float, warn_at: float = BUDGET_WARN_AT) -> str:
+    """Classify spend against its threshold as ``ok``, ``warn``, or ``exceeded``."""
+    if threshold <= 0:
+        return "ok"
+    if spend >= threshold:
+        return "exceeded"
+    if spend >= threshold * warn_at:
+        return "warn"
+    return "ok"
+
+
+def budget_usage_percent(spend: float, threshold: float) -> int:
+    """Spend as a whole percentage of its threshold, rounded half-up and floored at 0."""
+    if threshold <= 0:
+        return 0
+    return max(int(((spend / threshold) * 100) + 0.5), 0)
+
+
+def _bar_markup(percent: int, color: str, *, width: int = 28) -> str:
+    """A Rich-markup fill bar: the filled portion in the state color, the remainder dimmed."""
+    capped = min(max(percent, 0), 100)
+    filled = min(max((capped * width + 50) // 100, 0), width)
+    return f"[{color}]{'█' * filled}[/{color}][dim]{'░' * (width - filled)}[/dim]"
+
+
+def _header_line(state: str) -> str | None:
+    """The status callout shown above the bar for non-ok states."""
+    if state == "exceeded":
+        return "[bold red]⛔ Workspace budget exceeded[/bold red]"
+    if state == "warn":
+        return "[bold yellow]⚠️  Nearing workspace budget[/bold yellow]"
+    return None
+
+
+def render_budget_panel(
+    recommendation: dict,
+    *,
+    title: str | None = None,
+    extra_lines: list[str] | None = None,
+    managed: dict | None = None,
+) -> Panel | None:
+    """Render the managed budget as a bordered panel with a color-coded fill bar.
+
+    Returns None when the recommendation carries no threshold to measure against, so a workspace
+    with no budget policy shows nothing rather than an empty bar. ``managed`` supplies the admin's
+    tiers, so the amber point matches where their policy actually starts stepping down.
+    """
+    threshold = recommendation.get("effective_threshold")
+    spend = recommendation.get("current_spend")
+    if not isinstance(threshold, int | float) or isinstance(threshold, bool) or threshold <= 0:
+        return None
+    spend = float(spend) if isinstance(spend, int | float) and not isinstance(spend, bool) else 0.0
+    threshold = float(threshold)
+
+    state = budget_state(spend, threshold, budget_warn_fraction(managed))
+    color = _BUDGET_STATE_STYLE.get(state, "green")
+    percent = budget_usage_percent(spend, threshold)
+
+    lines: list[str] = []
+    header = _header_line(state)
+    if header:
+        lines.append(header)
+        lines.append("")
+    lines.append(
+        f"[bold]${spend:,.2f}[/bold] / ${threshold:,.2f}    [{color}]{percent}% used[/{color}]"
+    )
+    lines.append(_bar_markup(percent, color))
+    lines.append("")
+    if extra_lines:
+        lines.extend(extra_lines)
+        lines.append("")
+    lines.append(f"[bold]Remaining[/bold]   ${max(threshold - spend, 0.0):,.2f}")
+
+    return Panel(
+        Text.from_markup("\n".join(lines)),
+        title=Text(title or "Workspace Budget", style=f"bold {color}"),
+        border_style=color,
+        expand=False,
+        padding=(1, 2, 0, 2),
+    )
+
+
+def recommendation_line(display_agent: str | None, model: str | None, percent: int) -> str | None:
+    """The "you've used N%, recommended is X" sentence shown inside the panel."""
+    if not display_agent and not model:
+        return None
+    used = f"You've used [bold]{percent}%[/bold] of the workspace budget. "
+    if display_agent and model:
+        return f"{used}Recommended agent is [bold]{display_agent}[/bold] with model [bold]{model}[/bold]."
+    if display_agent:
+        return f"{used}Recommended agent is [bold]{display_agent}[/bold]."
+    return f"{used}Recommended model is [bold]{model}[/bold]."

src/ucode/managed_config.py

@@ -23,7 +23,11 @@
 from typing import cast
 
 import ucode.config_io as config_io
-from ucode.databricks import fetch_managed_coding_agent_configs, get_databricks_token
+from ucode.databricks import (
+    fetch_managed_coding_agent_configs,
+    fetch_model_recommendation,
+    get_databricks_token,
+)
 from ucode.ui import console, print_warning
 
 MANAGED_STATE_PATH = config_io.APP_DIR / "managed-state.json"
@@ -243,6 +247,42 @@ def normalize_managed_config(raw: dict) -> dict:
     return result
 
 
+def _decimal(value: object) -> float | None:
+    """Parse one of the API's decimal-string money fields, or None when absent/unparseable."""
+    text = _str(value)
+    if text is None:
+        return None
+    try:
+        return float(text)
+    except ValueError:
+        return None
+
+
+def get_model_recommendation(workspace: str, token: str) -> tuple[dict | None, str | None]:
+    """Fetch the agent and model the caller's budget tier allows, normalized for the launch path.
+
+    Returns ``(recommendation, reason)`` where the recommendation is ``{"agent", "model",
+    "current_spend", "effective_threshold"}``. Every field is optional server-side, so each is
+    normalized independently: an agent this build doesn't recognize is dropped rather than failing
+    the read, and a model can arrive without an agent.
+    """
+    payload, reason = fetch_model_recommendation(workspace, token)
+    if reason is not None:
+        return None, reason
+    agent = AGENT_ENUM_TO_TOOL.get(_str(payload.get("recommended_agent")) or "")
+    model = _str(payload.get("recommended_model"))
+    spend = _decimal(payload.get("current_spend"))
+    threshold = _decimal(payload.get("effective_threshold"))
+    if agent is None and model is None and spend is None and threshold is None:
+        return None, None
+    return {
+        "agent": agent,
+        "model": model,
+        "current_spend": spend,
+        "effective_threshold": threshold,
+    }, None
+
+
 def get_managed_config(workspace: str, token: str) -> tuple[dict | None, str | None]:
     """Fetch and normalize the workspace's managed config.