From 696ffe1a3e4f42694e63b9fb53a584828cac0391 Mon Sep 17 00:00:00 2001 From: Anjali Sujithan Date: Fri, 7 Aug 2026 00:11:23 +0000 Subject: [PATCH 1/2] route launches by budget tier and show workspaces spend --- src/ucode/cli.py | 96 ++++++++++++++++++++++++- src/ucode/databricks.py | 16 +++++ src/ucode/managed_budget.py | 131 +++++++++++++++++++++++++++++++++ src/ucode/managed_config.py | 42 ++++++++++- src/ucode/managed_resolve.py | 27 +++++++ tests/test_cli.py | 132 ++++++++++++++++++++++++++++++++++ tests/test_managed_budget.py | 121 +++++++++++++++++++++++++++++++ tests/test_managed_config.py | 69 ++++++++++++++++++ tests/test_managed_resolve.py | 46 ++++++++++++ 9 files changed, 677 insertions(+), 3 deletions(-) create mode 100644 src/ucode/managed_budget.py create mode 100644 tests/test_managed_budget.py diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 3d4fb87..7bd2742 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -56,7 +56,13 @@ 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, @@ -64,9 +70,11 @@ 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.mcp import ( @@ -1276,6 +1284,73 @@ 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. + """ + if managed is None or skip_preflight: + 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 as exc: + # A token that lapsed since the config refresh must not block the launch. + 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, @@ -1284,6 +1359,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) @@ -1336,6 +1412,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") @@ -1408,7 +1490,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"] @@ -1468,6 +1552,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: @@ -1592,7 +1678,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 " @@ -1605,6 +1696,7 @@ def _launch_managed_default( skip_preflight=skip_preflight, workspace=workspace, managed=managed, + recommendation=recommendation, ) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index c912e2c..d2af124 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1450,6 +1450,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) ----------------------------- diff --git a/src/ucode/managed_budget.py b/src/ucode/managed_budget.py new file mode 100644 index 0000000..dd147f3 --- /dev/null +++ b/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]." diff --git a/src/ucode/managed_config.py b/src/ucode/managed_config.py index 90d4d58..00b0165 100644 --- a/src/ucode/managed_config.py +++ b/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" @@ -241,6 +245,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. diff --git a/src/ucode/managed_resolve.py b/src/ucode/managed_resolve.py index f6bd264..864c19f 100644 --- a/src/ucode/managed_resolve.py +++ b/src/ucode/managed_resolve.py @@ -185,6 +185,33 @@ def managed_default_model(managed: dict, tool: str) -> str | None: return _str(_agent_model_config(managed, tool).get("default_model")) +def recommended_agent(recommendation: dict | None, managed: dict) -> str | None: + """The agent the budget tier recommends, or the config's ``default_agent`` when it names none. + + The server resolves the agent before the model, so a tier can move a developer to a cheaper + agent without restating a model. + """ + agent = _str(_as_dict(recommendation).get("agent")) + return agent or _str(_as_dict(managed).get("default_agent")) + + +def managed_launch_model(managed: dict, recommendation: dict | None, tool: str) -> str | None: + """The model the admin's policy wants ``tool`` to start on, or None. + + A budget recommendation supersedes the config's own ``default_model``, since it additionally + reflects which spend tier the developer has reached — but only for the agent it was recommended + for. A tier that moves the org to another agent names that agent's model, which the one being + launched may not be able to serve. + """ + recommended = _as_dict(recommendation) + agent = _str(recommended.get("agent")) + if agent is None or agent == tool: + model = _str(recommended.get("model")) + if model: + return model + return managed_default_model(managed, tool) + + def resolve_state(managed: dict, state: dict, tool: str) -> dict: """Return a copy of ``state`` with ``tool``'s managed values layered on top. diff --git a/tests/test_cli.py b/tests/test_cli.py index eaa7551..c6d3608 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2419,3 +2419,135 @@ def test_subcommands_still_work(self, monkeypatch): monkeypatch.setattr("ucode.cli.load_state", lambda: {"workspace": "https://w"}) result = runner.invoke(app, ["status"]) assert result.exit_code == 0, result.output + + +class TestBudgetRecommendationAtLaunch: + """The budget read informs the launch; it never blocks it.""" + + @staticmethod + def _launch(monkeypatch, *, tool="claude", managed, recommendation=None, reason=None): + state = dict(MINIMAL_STATE) + calls: list[str] = [] + + def fake_recommendation(workspace, token): + calls.append(workspace) + return recommendation, reason + + monkeypatch.setattr("ucode.cli.get_model_recommendation", fake_recommendation) + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state) as cfg, + patch("ucode.cli.get_databricks_token", return_value="tok"), + patch("ucode.cli._fetch_managed_config", return_value=managed), + patch("ucode.cli.launch_agent"), + ): + result = runner.invoke(app, [tool]) + return result, calls, cfg + + def test_not_checked_without_a_managed_config(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + result, calls, _ = self._launch(monkeypatch, managed=None) + assert result.exit_code == 0, result.output + assert calls == [] + + def test_the_recommended_agent_gets_the_recommended_model(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + managed = { + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}} + } + } + _result, _calls, cfg = self._launch( + monkeypatch, + managed=managed, + recommendation={"agent": "claude", "model": "system.ai.claude-haiku-4-5"}, + ) + assert cfg.call_args.args[2] == "system.ai.claude-haiku-4-5" + + def test_another_agent_keeps_its_own_model_and_is_told_why(self, monkeypatch): + # A tier's model belongs to the tier's agent; pinning it on claude would land a Kimi id in + # ANTHROPIC_MODEL, which the Anthropic-dialect endpoint cannot serve. + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + managed = { + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, + "opencode": {}, + } + } + result, _calls, cfg = self._launch( + monkeypatch, + managed=managed, + recommendation={ + "agent": "opencode", + "model": "system.ai.kimi-k2-7-code", + "current_spend": 412.5, + "effective_threshold": 500.0, + }, + ) + assert result.exit_code == 0, result.output + assert cfg.call_args.args[2] == "system.ai.claude-opus-4-8" + assert "recommends OpenCode" in result.output + + def test_a_failed_read_does_not_block_the_launch(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + result, _calls, _cfg = self._launch( + monkeypatch, + managed={"enabled_agents": {"claude": {}}}, + recommendation=None, + reason="HTTP 500", + ) + assert result.exit_code == 0, result.output + assert "Could not check your budget" in result.output + + def test_a_token_failure_does_not_block_the_launch(self, monkeypatch): + # Auth can lapse between the config refresh and the budget check. + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + state = dict(MINIMAL_STATE) + monkeypatch.setattr("ucode.cli.get_model_recommendation", lambda ws, tok: (None, None)) + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=state), + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.configure_tool", return_value=state), + patch("ucode.cli.get_databricks_token", side_effect=RuntimeError("token expired")), + patch( + "ucode.cli._fetch_managed_config", + return_value={"enabled_agents": {"claude": {}}}, + ), + patch("ucode.cli.launch_agent"), + ): + result = runner.invoke(app, ["claude"]) + assert result.exit_code == 0, result.output + assert "Could not check your budget" in result.output + + def test_shows_the_budget_bar(self, monkeypatch): + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", "1") + result, _calls, _cfg = self._launch( + monkeypatch, + managed={"enabled_agents": {"claude": {}}}, + recommendation={ + "agent": "claude", + "model": "m", + "current_spend": 412.5, + "effective_threshold": 500.0, + }, + ) + assert result.exit_code == 0, result.output + assert "83% used" in result.output + assert "█" in result.output + + @pytest.mark.parametrize("env_value", [None, "", "0"]) + def test_not_checked_when_the_env_var_is_off(self, monkeypatch, env_value): + if env_value is None: + monkeypatch.delenv("ENABLE_MANAGED_AGENT_CONFIG", raising=False) + else: + monkeypatch.setenv("ENABLE_MANAGED_AGENT_CONFIG", env_value) + result, calls, _ = self._launch(monkeypatch, managed=None) + assert result.exit_code == 0, result.output + assert calls == [] diff --git a/tests/test_managed_budget.py b/tests/test_managed_budget.py new file mode 100644 index 0000000..0d60d2f --- /dev/null +++ b/tests/test_managed_budget.py @@ -0,0 +1,121 @@ +"""Tests for managed_budget.py — the workspace budget panel shown at launch.""" + +from __future__ import annotations + +import pytest + +from ucode.managed_budget import ( + BUDGET_WARN_AT, + budget_state, + budget_usage_percent, + budget_warn_fraction, + recommendation_line, + render_budget_panel, +) + + +class TestBudgetState: + @pytest.mark.parametrize( + ("spend", "threshold", "expected"), + [ + (0.0, 500.0, "ok"), + (399.0, 500.0, "ok"), + (400.0, 500.0, "warn"), + (499.0, 500.0, "warn"), + (500.0, 500.0, "exceeded"), + (750.0, 500.0, "exceeded"), + ], + ) + def test_classifies_spend(self, spend, threshold, expected): + assert budget_state(spend, threshold) == expected + + def test_no_threshold_is_never_over_budget(self): + # A workspace with no budget policy has nothing to exceed. + assert budget_state(100.0, 0.0) == "ok" + + def test_warn_point_follows_the_policy(self): + # At 60% spend a policy whose first tier is 50% is already stepping down, so warn. + assert budget_state(300.0, 500.0, 0.5) == "warn" + assert budget_state(300.0, 500.0, 0.8) == "ok" + + +class TestBudgetWarnFraction: + """The amber point comes from the admin's own tiers, not a fixed 80%.""" + + @staticmethod + def _policy(*percentages): + return {"budget_policy": {"tiers": [{"spending_percentage": p} for p in percentages]}} + + def test_uses_the_lowest_activating_tier(self): + assert budget_warn_fraction(self._policy(0.5, 0.8)) == 0.5 + + def test_ignores_tiers_at_zero(self): + # A 0.0 tier activates from the first dollar, so warning on it would be permanent amber. + assert budget_warn_fraction(self._policy(0.0, 0.5, 0.8)) == 0.5 + assert budget_warn_fraction(self._policy(0.0)) == BUDGET_WARN_AT + + def test_falls_back_without_a_policy(self): + assert budget_warn_fraction(None) == BUDGET_WARN_AT + assert budget_warn_fraction({}) == BUDGET_WARN_AT + assert budget_warn_fraction({"budget_policy": {"tiers": []}}) == BUDGET_WARN_AT + + @pytest.mark.parametrize("bad", [True, "0.5", None, 1.5, -0.2]) + def test_ignores_unusable_percentages(self, bad): + # `True` would otherwise read as a 1.0 fraction, and out-of-range values aren't meaningful. + assert budget_warn_fraction(self._policy(bad)) == BUDGET_WARN_AT + + +class TestBudgetUsagePercent: + @pytest.mark.parametrize( + ("spend", "threshold", "expected"), + [(412.5, 500.0, 83), (0.0, 500.0, 0), (500.0, 500.0, 100), (600.0, 500.0, 120)], + ) + def test_rounds_half_up(self, spend, threshold, expected): + assert budget_usage_percent(spend, threshold) == expected + + def test_zero_threshold_is_zero_percent(self): + assert budget_usage_percent(100.0, 0.0) == 0 + + +class TestRenderBudgetPanel: + def test_none_without_a_threshold(self): + # No budget policy on the workspace means no bar to draw. + assert render_budget_panel({"current_spend": 10.0}) is None + assert render_budget_panel({"current_spend": 10.0, "effective_threshold": 0.0}) is None + + def test_renders_the_bar_and_spend(self): + panel = render_budget_panel({"current_spend": 412.5, "effective_threshold": 500.0}) + assert panel is not None + body = panel.renderable.plain + assert "$412.50 / $500.00" in body + assert "83% used" in body + assert "█" in body and "░" in body + assert "Remaining" in body and "$87.50" in body + + def test_over_budget_remaining_never_goes_negative(self): + panel = render_budget_panel({"current_spend": 600.0, "effective_threshold": 500.0}) + assert panel is not None + assert "$0.00" in panel.renderable.plain + + def test_missing_spend_is_treated_as_zero(self): + panel = render_budget_panel({"effective_threshold": 500.0}) + assert panel is not None + assert "$0.00 / $500.00" in panel.renderable.plain + + +class TestRecommendationLine: + def test_names_both_agent_and_model(self): + line = recommendation_line("OpenCode", "system.ai.claude-haiku-4-5", 83) + assert "83%" in line + assert "OpenCode" in line and "system.ai.claude-haiku-4-5" in line + + def test_agent_only(self): + assert "OpenCode" in recommendation_line("OpenCode", None, 50) + + def test_model_only(self): + # The server can recommend a model without an agent. + line = recommendation_line(None, "system.ai.gpt-5", 50) + assert "system.ai.gpt-5" in line + + def test_none_when_nothing_recommended(self): + assert recommendation_line(None, None, 50) is None diff --git a/tests/test_managed_config.py b/tests/test_managed_config.py index 1f5c4a2..479fe97 100644 --- a/tests/test_managed_config.py +++ b/tests/test_managed_config.py @@ -392,3 +392,72 @@ def test_no_workspace_is_a_noop(self, monkeypatch): mc_mod, "get_managed_config", lambda ws, tok: pytest.fail("should not fetch") ) assert refresh_managed_config({}) is None + + +class TestGetModelRecommendation: + """The budget recommendation read. Every response field is optional server-side.""" + + @staticmethod + def _stub(monkeypatch, payload, reason=None): + monkeypatch.setattr(mc_mod, "fetch_model_recommendation", lambda ws, tok: (payload, reason)) + + def test_normalizes_agent_model_and_spend(self, monkeypatch): + self._stub( + monkeypatch, + { + "recommended_agent": "CODING_AGENT_OPENCODE", + "recommended_model": "system.ai.claude-haiku-4-5", + "current_spend": "412.50", + "effective_threshold": "500.00", + }, + ) + rec, reason = mc_mod.get_model_recommendation("https://w", "tok") + assert reason is None + assert rec == { + "agent": "opencode", + "model": "system.ai.claude-haiku-4-5", + "current_spend": 412.5, + "effective_threshold": 500.0, + } + + def test_model_without_an_agent(self, monkeypatch): + # A model-only tier with no default_agent recommends a model but no agent. + self._stub(monkeypatch, {"recommended_model": "system.ai.gpt-5", "current_spend": "1.00"}) + rec, _ = mc_mod.get_model_recommendation("https://w", "tok") + assert rec is not None and rec["agent"] is None and rec["model"] == "system.ai.gpt-5" + + def test_agent_without_a_model(self, monkeypatch): + self._stub(monkeypatch, {"recommended_agent": "CODING_AGENT_PI", "current_spend": "1.00"}) + rec, _ = mc_mod.get_model_recommendation("https://w", "tok") + assert rec is not None and rec["agent"] == "pi" and rec["model"] is None + + @pytest.mark.parametrize("agent_enum", ["CODING_AGENT_UNSPECIFIED", "CODING_AGENT_FUTURE", ""]) + def test_unknown_agent_is_dropped_not_fatal(self, monkeypatch, agent_enum): + self._stub( + monkeypatch, + {"recommended_agent": agent_enum, "recommended_model": "m", "current_spend": "1.00"}, + ) + rec, reason = mc_mod.get_model_recommendation("https://w", "tok") + assert reason is None + assert rec is not None and rec["agent"] is None and rec["model"] == "m" + + def test_threshold_alone_still_reports(self, monkeypatch): + # A budget with no spend yet still has a threshold worth showing. + self._stub(monkeypatch, {"effective_threshold": "500.00"}) + rec, _ = mc_mod.get_model_recommendation("https://w", "tok") + assert rec is not None and rec["effective_threshold"] == 500.0 + + def test_empty_response_is_no_recommendation(self, monkeypatch): + self._stub(monkeypatch, {}) + assert mc_mod.get_model_recommendation("https://w", "tok") == (None, None) + + def test_failed_read_surfaces_the_reason(self, monkeypatch): + self._stub(monkeypatch, {}, reason="HTTP 500") + assert mc_mod.get_model_recommendation("https://w", "tok") == (None, "HTTP 500") + + def test_unparseable_decimals_become_none(self, monkeypatch): + self._stub( + monkeypatch, {"recommended_agent": "CODING_AGENT_PI", "current_spend": "not-a-number"} + ) + rec, _ = mc_mod.get_model_recommendation("https://w", "tok") + assert rec is not None and rec["current_spend"] is None diff --git a/tests/test_managed_resolve.py b/tests/test_managed_resolve.py index 340a3fa..5dbf41d 100644 --- a/tests/test_managed_resolve.py +++ b/tests/test_managed_resolve.py @@ -13,10 +13,12 @@ from ucode.managed_resolve import ( managed_default_model, managed_enabled_tools, + managed_launch_model, managed_provider_service, managed_state_overrides, managed_supplies_models, managed_unservable_models, + recommended_agent, resolve_state, ) from ucode.state import MANAGED_OVERLAY_KEY @@ -534,3 +536,47 @@ def test_no_warning_when_anything_is_servable(self, tool, models): def test_agents_that_pass_models_through_never_warn(self): assert managed_unservable_models(self._managed("codex", ["anything"]), "codex") == [] + + +class TestRecommendedAgent: + """A tier can move the org to a cheaper agent; with none named, default_agent stands.""" + + def test_tier_agent_wins(self): + assert recommended_agent({"agent": "opencode"}, {"default_agent": "claude"}) == "opencode" + + def test_falls_back_to_default_agent(self): + assert recommended_agent({"agent": None}, {"default_agent": "claude"}) == "claude" + assert recommended_agent(None, {"default_agent": "claude"}) == "claude" + + def test_none_when_neither_is_set(self): + assert recommended_agent(None, {}) is None + + +class TestManagedLaunchModel: + """A tier's model supersedes the config default, but only for the tier's own agent.""" + + MANAGED = { + "enabled_agents": { + "claude": {"model_config": {"default_model": "system.ai.claude-opus-4-8"}}, + "opencode": {"model_config": {"default_model": "system.ai.claude-sonnet-4-6"}}, + } + } + + def test_the_recommended_agent_gets_the_recommended_model(self): + rec = {"agent": "opencode", "model": "system.ai.kimi-k2-7-code"} + assert managed_launch_model(self.MANAGED, rec, "opencode") == "system.ai.kimi-k2-7-code" + + def test_other_agents_keep_their_own_default(self): + # opencode's Kimi model is not servable by claude's Anthropic-dialect endpoint. + rec = {"agent": "opencode", "model": "system.ai.kimi-k2-7-code"} + assert managed_launch_model(self.MANAGED, rec, "claude") == "system.ai.claude-opus-4-8" + + def test_a_model_without_an_agent_applies_to_any_tool(self): + rec = {"agent": None, "model": "system.ai.claude-haiku-4-5"} + assert managed_launch_model(self.MANAGED, rec, "claude") == "system.ai.claude-haiku-4-5" + + def test_default_model_stands_without_a_recommendation(self): + assert managed_launch_model(self.MANAGED, None, "claude") == "system.ai.claude-opus-4-8" + + def test_none_when_neither_names_a_model(self): + assert managed_launch_model({}, None, "pi") is None From 8c26a83da3c4d04f9ad6d644e3aec2d40f66d23c Mon Sep 17 00:00:00 2001 From: AarushiShah-db Date: Fri, 7 Aug 2026 06:19:35 +0000 Subject: [PATCH 2/2] Don't fetch the budget recommendation under --dry-run, and never let its read block the launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bare `ucode --dry-run` resolved the agent from the last saved config but still called _fetch_budget_recommendation, which resolves a Databricks token and hits recommendModel — the control-plane fetch --dry-run promises to skip. Skip it under is_dry_run(), mirroring the managed-config read. Also broaden the read's except from RuntimeError to (RuntimeError, OSError): a Databricks CLI that isn't installed/reachable raises FileNotFoundError from the token subprocess, which was escaping and blocking the launch — contrary to the guarantee that a failed budget read never blocks. Co-authored-by: Isaac --- src/ucode/cli.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 7bd2742..96164a8 100644 --- a/src/ucode/cli.py +++ b/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, @@ -1311,7 +1311,9 @@ def _fetch_budget_recommendation( Enforcement is server-side, so a failed read only costs the recommendation: the config's own ``default_model`` still applies and the launch proceeds. """ - if managed is None or skip_preflight: + # --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 @@ -1321,8 +1323,9 @@ def _fetch_budget_recommendation( state["workspace"], get_databricks_token(state["workspace"], state.get("profile")), ) - except RuntimeError as exc: - # A token that lapsed since the config refresh must not block the launch. + 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(