[AIGTWY-4880] CUJ5: Verify budget recommendations and agent startup

.github/workflows/integration.yml

@@ -431,6 +431,8 @@ jobs:
       UV_INDEX_URL: https://pypi.org/simple
       UG_CUJ_SP_CLIENT_ID: ${{ secrets.UG_CUJ_SP_CLIENT_ID }}
       UG_CUJ_SP_CLIENT_SECRET: ${{ secrets.UG_CUJ_SP_CLIENT_SECRET }}
+      UG_BUDGET_CUJ_SP_CLIENT_ID: ${{ secrets.UG_BUDGET_CUJ_SP_CLIENT_ID }}
+      UG_BUDGET_CUJ_SP_CLIENT_SECRET: ${{ secrets.UG_BUDGET_CUJ_SP_CLIENT_SECRET }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020  # v4

tests/README.md

@@ -7,7 +7,7 @@ mocks, monkeypatching, fake binaries/services, or fabricated ug state.
 | Category | Location | What it proves |
 | --- | --- | --- |
 | Unit/component | Existing `test_*.py` files | Individual behavior; dependencies may be mocked |
-| Dedicated-workspace CUJ | `e2e_cuj/test_cuj_*.py` | Real user journeys against read-only, preconfigured workspaces |
+| Dedicated-workspace CUJ | `e2e_cuj/test_*.py` | Real user journeys against read-only, preconfigured workspaces |
 | Existing e2e | `test_e2e*.py` | Real workspace behavior with some patched setup/internal calls |
 | Integration CUJs | `integration/test_*.py` | Public configure, TUI, script, command, protocol, and lifecycle journeys |
 | Installation | `integration/test_installation.py` | Fresh installed package, CLI, and local helpers without credentials on Linux and advisory native Windows |
@@ -181,6 +181,7 @@ integration utilities; only CUJ-specific evidence correlation stays in a test fi
 | `test_ug_configure_claude_cleans_stale_skills_mcp_on_workspace_switch` | Configure the first workspace, register its skills MCP, switch to a second real workspace, and use Claude | Old registration removed from Claude and the new workspace state; old workspace bucket preserved; repeat configure stays clean; real file task completes on the second workspace |
 | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup |
 | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch. The Codex case also checks the shared catalog pointer, restart guidance, and a fresh bare app-server's visible model list |
+| `e2e_cuj/test_ug_budget_defaults.py` | Launch bare `ug` with separate low-spend and above-tier principals against the fixed 1% tier (`spending_percentage=0.01`); explicitly launch `ug claude` above the tier | Bare launches check Claude/Sonnet below the tier and the Codex/Luna recommendation above it; `ug usage` agrees with backend spend, threshold, and percentage. Explicit `ug claude` displays the backend's Codex/Luna recommendation while its generated model setting and native header select Sonnet. No budget writes or inference tasks. Run in the shared `E2E CUJs` job with both credential pairs documented in `integration/README.md`. |
 | `test_case_01_*` | Launch managed Claude without defaults after configure and from fresh state | Claude receives the admin MPS header; its gateway cache and replacement picker match the independently fetched provider model IDs; catalog labels are preserved and a model appears in a numbered picker row |
 | `test_case_03_*`, `test_case_05_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state |
 | `test_case_02_*` | Launch managed Codex after configure and from fresh state | The scoped and stable catalogs, ug-launched app server, and fresh bare app server match the independently fetched admin MPS model IDs. The configured case uses real `ug revert` to remove ug's shared pointer and stable file while preserving a user setting |

tests/e2e_cuj/AGENTS.md

@@ -3,6 +3,8 @@
 Each CUJ owns a separate workspace. Subclass `BaseCujTest` from `base.py` and set
 `WORKSPACE_URL`. Setup provides `self.workspace`, a Databricks SDK client using
 `UG_CUJ_SP_CLIENT_ID` and `UG_CUJ_SP_CLIENT_SECRET` with OAuth M2M authentication.
+A CUJ that needs another identity overrides `CLIENT_ID_ENV` and `CLIENT_SECRET_ENV`
+with its credential environment-variable names.
 The workspaces share one metastore, so metastore-scoped fixture names belong in
 `helpers/constants.py`; keep each workspace URL in its test class.
 All workspace configurations are READ ONLY. Never create, publish, update, delete,

tests/e2e_cuj/base.py

@@ -9,6 +9,8 @@
 
 class BaseCujTest:
     WORKSPACE_URL: ClassVar[str] = ""
+    CLIENT_ID_ENV: ClassVar[str] = "UG_CUJ_SP_CLIENT_ID"
+    CLIENT_SECRET_ENV: ClassVar[str] = "UG_CUJ_SP_CLIENT_SECRET"
     workspace: WorkspaceClient
 
     @pytest.fixture(scope="class", autouse=True)
@@ -17,10 +19,12 @@ def setup_workspace(self, request):
         if not workspace_url:
             pytest.fail("Set WORKSPACE_URL on your CUJ test class.", pytrace=False)
 
-        client_id = os.environ.get("UG_CUJ_SP_CLIENT_ID", "")
-        client_secret = os.environ.get("UG_CUJ_SP_CLIENT_SECRET", "")
+        client_id_env = request.cls.CLIENT_ID_ENV
+        client_secret_env = request.cls.CLIENT_SECRET_ENV
+        client_id = os.environ.get(client_id_env, "")
+        client_secret = os.environ.get(client_secret_env, "")
         if not client_id or not client_secret:
-            pytest.fail("Set UG_CUJ_SP_CLIENT_ID and UG_CUJ_SP_CLIENT_SECRET.", pytrace=False)
+            pytest.fail(f"Set {client_id_env} and {client_secret_env}.", pytrace=False)
 
         request.cls.workspace = WorkspaceClient(
             host=workspace_url,

tests/e2e_cuj/test_ug_budget_defaults.py

@@ -0,0 +1,238 @@
+"""Live budget recommendation and default-model checks for a fixed workspace tier."""
+
+import json
+import re
+from decimal import ROUND_HALF_UP, Decimal
+
+from tests.integration.utils.terminal import AgentTerminal
+
+from .base import BaseCujTest
+from .helpers.constants import CLAUDE, CODEX, CodingAgent
+
+
+class _BudgetDefaultsBase(BaseCujTest):
+    WORKSPACE_URL = "https://dbc-497adeef-62c0.cloud.databricks.com"
+    BUDGET_ID = "25c4ce5c-fcd6-4d00-8266-69030cc1a236"
+    SONNET_MODEL = "system.ai.claude-sonnet-4-6"
+    SOL_MODEL = "system.ai.gpt-5-6-sol"
+    LUNA_MODEL = "system.ai.gpt-5-6-luna"
+
+
+class TestCujBudgetDefaultsBelowTier(_BudgetDefaultsBase):
+    CLIENT_ID_ENV = "UG_BUDGET_CUJ_SP_CLIENT_ID"
+    CLIENT_SECRET_ENV = "UG_BUDGET_CUJ_SP_CLIENT_SECRET"
+
+    def test_bare_ug_uses_claude_default_below_budget_tier(self, cuj):
+        """Scenario: launch bare ``ug`` as a fresh principal below the fixed 1% tier.
+
+        Expected: the published config keeps Claude/Sonnet as the default, the real recommendation
+        also selects Claude/Sonnet, and the real Claude TUI starts on that model without a prompt,
+        inference request, budget write, or account-login flow. A fresh principal may have no
+        spend counters yet; when the backend returns an effective threshold, omitted spend is
+        treated as zero and must remain below the 1% tier. If both figures are absent, this case
+        verifies default selection without a numeric spend assertion.
+        """
+        session, workspace, _ = cuj
+        config_path = "/api/ai-gateway/v2/coding-agent-configs"
+        payload = workspace.client.api_client.do("GET", path=config_path)
+        config = payload["coding_agent_configs"][0]
+        assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config
+        claude = next(
+            agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CLAUDE_CODE
+        )
+        claude_defaults = claude["config"]["default_models"]
+        assert claude_defaults["default_model"] == self.SONNET_MODEL, config
+        assert config["smart_defaults"]["budget_id"] == self.BUDGET_ID, config
+        assert config["smart_defaults"]["tiers"] == [
+            {
+                "spending_percentage": 0.01,
+                "recommended_agent": CodingAgent.CODEX,
+                "recommended_model": self.LUNA_MODEL,
+            }
+        ], config
+
+        session.configure(
+            [
+                "configure",
+                "--workspace",
+                self.WORKSPACE_URL,
+                "--skip-upgrade",
+                "--disable-databricks-ai-tools",
+            ]
+        )
+
+        recommendation = workspace.client.api_client.do(
+            "POST", path=config_path + ":recommendModel", body={}
+        )
+        assert recommendation.get("recommended_agent") == CodingAgent.CLAUDE_CODE, recommendation
+        assert recommendation.get("recommended_model") == self.SONNET_MODEL, recommendation
+
+        threshold = recommendation.get("effective_threshold")
+        spend = recommendation.get("current_spend")
+        if threshold is not None:
+            threshold = Decimal(str(threshold))
+            assert threshold > 0, recommendation
+            spend = Decimal("0") if spend is None else Decimal(str(spend))
+            assert spend >= 0, recommendation
+            assert spend / threshold < Decimal("0.01"), recommendation
+        else:
+            # A fresh principal can have no usage counter; that does not establish a numeric ratio.
+            assert spend is None, recommendation
+
+        with AgentTerminal(session, CLAUDE, [str(session.binary)], "budget-default-claude") as tui:
+            tui.boot(timeout=240)
+            claude_settings = session.home / ".claude" / "ucode-settings.json"
+            settings = json.loads(claude_settings.read_text())
+            assert settings["env"]["ANTHROPIC_MODEL"] == self.SONNET_MODEL, settings
+            tui.wait_for(
+                lambda text: re.search(r"Claude Sonnet 4\.6\s*·", text),
+                "Claude's native Sonnet model header",
+            )
+            tui.exit_normally()
+
+
+class TestCujBudgetDefaults(_BudgetDefaultsBase):
+    def test_bare_ug_shows_luna_recommendation_and_starts_codex(self, cuj):
+        """Scenario: launch bare ``ug`` with the fixed 1% Luna budget tier active.
+
+        Expected: the read-only `usage` command reports the backend's spend and threshold, then
+        the real backend recommends Codex/Luna, the launch panel displays that recommendation, and
+        the native Codex TUI starts and exits normally, without budget writes or an inference task.
+        Applying the recommended Luna model over the managed Sol default is outside this coverage.
+        """
+        session, workspace, _ = cuj
+        config_path = "/api/ai-gateway/v2/coding-agent-configs"
+        payload = workspace.client.api_client.do("GET", path=config_path)
+        config = payload["coding_agent_configs"][0]
+        assert config["default_agent"] == CodingAgent.CLAUDE_CODE, config
+        codex_agent = next(
+            agent for agent in config["enabled_agents"] if agent["agent"] == CodingAgent.CODEX
+        )
+        live_default_model = (
+            codex_agent.get("config", {}).get("default_models", {}).get("default_model")
+        )
+        assert live_default_model == self.SOL_MODEL, (
+            f"Live Codex config model mismatch: expected {self.SOL_MODEL!r}, "
+            f"got {live_default_model!r}"
+        )
+        assert config["smart_defaults"]["budget_id"] == self.BUDGET_ID, config
+        assert config["smart_defaults"]["tiers"] == [
+            {
+                "spending_percentage": 0.01,
+                "recommended_agent": CodingAgent.CODEX,
+                "recommended_model": self.LUNA_MODEL,
+            }
+        ], config
+
+        session.configure(
+            [
+                "configure",
+                "--workspace",
+                self.WORKSPACE_URL,
+                "--skip-upgrade",
+                "--disable-databricks-ai-tools",
+            ]
+        )
+
+        before_usage = workspace.client.api_client.do(
+            "POST", path=config_path + ":recommendModel", body={}
+        )
+        usage = session.run("usage", timeout=120)
+        after_usage = workspace.client.api_client.do(
+            "POST", path=config_path + ":recommendModel", body={}
+        )
+        for recommendation in (before_usage, after_usage):
+            assert recommendation.get("recommended_agent") == CodingAgent.CODEX, recommendation
+            assert recommendation.get("recommended_model") == self.LUNA_MODEL, recommendation
+
+        match = re.search(
+            r"Budget spend:\s*\$([\d,]+\.\d{2})\s+of\s+\$([\d,]+\.\d{2})\s+\((\d+)%\)",
+            session.redact(usage.stdout + usage.stderr),
+        )
+        assert match, usage.stdout + usage.stderr
+        displayed_spend, displayed_threshold, displayed_percent = match.groups()
+        displayed_spend = Decimal(displayed_spend.replace(",", ""))
+        displayed_threshold = Decimal(displayed_threshold.replace(",", ""))
+        cent = Decimal("0.01")
+
+        before_spend = Decimal(str(before_usage["current_spend"]))
+        after_spend = Decimal(str(after_usage["current_spend"]))
+        before_threshold = Decimal(str(before_usage["effective_threshold"]))
+        after_threshold = Decimal(str(after_usage["effective_threshold"]))
+        assert before_spend <= after_spend, (before_usage, after_usage)
+        assert before_threshold == after_threshold, (before_usage, after_usage)
+        assert displayed_threshold == before_threshold.quantize(cent, rounding=ROUND_HALF_UP)
+        assert (
+            before_spend.quantize(cent, rounding=ROUND_HALF_UP)
+            <= displayed_spend
+            <= after_spend.quantize(cent, rounding=ROUND_HALF_UP)
+        ), (before_usage, after_usage, displayed_spend)
+        assert (
+            int(format(float(before_spend / before_threshold), ".0%").removesuffix("%"))
+            <= int(displayed_percent)
+            <= int(format(float(after_spend / after_threshold), ".0%").removesuffix("%"))
+        ), (before_usage, after_usage, usage)
+
+        with AgentTerminal(session, CODEX, [str(session.binary)], "budget-luna") as tui:
+            tui.boot(timeout=240)
+            transcript = session.redact(
+                "".join(tui.output) + "\n" + tui.visible,
+                strip_ansi=True,
+            )
+            assert re.search(
+                rf"Recommended agent is Codex with model {re.escape(self.LUNA_MODEL)}\.",
+                transcript,
+            ), (
+                "Codex launch did not display the backend recommendation "
+                f"for {self.LUNA_MODEL!r}:\n{transcript}"
+            )
+            tui.exit_normally()
+
+
+class TestCujBudgetDefaultsExplicitClaude(_BudgetDefaultsBase):
+    def test_ug_claude_starts_on_sonnet_despite_codex_luna_recommendation(self, cuj):
+        """Scenario: launch explicit ``ug claude`` for an above-tier principal.
+
+        Expected: the live backend recommends Codex/Luna, and the Claude launch displays
+        that recommendation while its generated model setting and native header select Sonnet.
+        No task is submitted; this checks startup selection, not the model used for inference.
+        """
+        session, workspace, _ = cuj
+        session.configure(
+            [
+                "configure",
+                "--workspace",
+                self.WORKSPACE_URL,
+                "--skip-upgrade",
+                "--disable-databricks-ai-tools",
+            ]
+        )
+        recommendation = workspace.client.api_client.do(
+            "POST", path="/api/ai-gateway/v2/coding-agent-configs:recommendModel", body={}
+        )
+        assert recommendation.get("recommended_agent") == CodingAgent.CODEX, recommendation
+        assert recommendation.get("recommended_model") == self.LUNA_MODEL, recommendation
+
+        with AgentTerminal(
+            session, CLAUDE, [str(session.binary), CLAUDE], "explicit-claude"
+        ) as tui:
+            tui.boot(timeout=240)
+            transcript = session.redact(
+                "".join(tui.output) + "\n" + tui.visible,
+                strip_ansi=True,
+            )
+            assert re.search(
+                rf"Recommended agent is Codex with model {re.escape(self.LUNA_MODEL)}\.",
+                transcript,
+            ), (
+                "Claude launch did not display the backend recommendation "
+                f"for {self.LUNA_MODEL!r}:\n{transcript}"
+            )
+            claude_settings = session.home / ".claude" / "ucode-settings.json"
+            settings = json.loads(claude_settings.read_text())
+            assert settings["env"]["ANTHROPIC_MODEL"] == self.SONNET_MODEL, settings
+            tui.wait_for(
+                lambda text: re.search(r"Claude Sonnet 4\.6\s*·", text),
+                "Claude's native Sonnet model header",
+            )
+            tui.exit_normally()

tests/integration/README.md

@@ -32,6 +32,32 @@ The `smart_defaults` wire schema, legacy `spend_tiers` cache reads, and recommen
 request gating are covered by unit/component tests listed in `../README.md`. This suite
 does not yet assert live `recommendModel` request counts for configs with and without tiers.
 
+## Dedicated-workspace CUJ5
+
+[`../e2e_cuj/test_ug_budget_defaults.py`](../e2e_cuj/test_ug_budget_defaults.py) runs
+three budget cases in the shared `E2E CUJs` job against the fixed 1% tier. The above-tier
+cases use `UG_CUJ_SP_CLIENT_ID` / `UG_CUJ_SP_CLIENT_SECRET`; the below-tier case uses
+`UG_BUDGET_CUJ_SP_CLIENT_ID` / `UG_BUDGET_CUJ_SP_CLIENT_SECRET`.
+
+Run from a clean POSIX runner with both credential pairs set:
+
+```bash
+UG_CUJ_SP_CLIENT_ID=... UG_CUJ_SP_CLIENT_SECRET=... \
+UG_BUDGET_CUJ_SP_CLIENT_ID=... UG_BUDGET_CUJ_SP_CLIENT_SECRET=... \
+uv run --with pexpect==4.9.0 --with pyte==0.8.2 \
+  pytest --confcutdir=tests/e2e_cuj tests/e2e_cuj -v
+```
+
+All three cases submit no inference tasks. The above-tier bare-launch case checks
+`ug usage`, the Luna recommendation, and Codex startup; it does not assert that Codex
+applies Luna. The below-tier case verifies Claude/Sonnet selection and checks the
+spend percentage when the backend supplies budget figures.
+
+The explicit `ug claude` case verifies that startup displays the backend's Codex/Luna
+recommendation while its generated `ANTHROPIC_MODEL` setting and native header select
+Sonnet, then exits. It does not verify Claude's inference model. All three cases leave
+workspace and budget configuration unchanged.
+
 Shared subprocess command resolution is covered by `../test_subprocess_cross_os.py` and
 enforced by Ruff. These component checks do not establish native Windows coverage
 for every agent; the Windows journeys below validate their explicitly selected paths.

tests/test_cuj_session.py

@@ -16,6 +16,8 @@ def test_cuj_session_does_not_inherit_agent_credentials_or_configuration(tmp_pat
         "CODEX_HOME",
         "UG_CUJ_SP_CLIENT_ID",
         "UG_CUJ_SP_CLIENT_SECRET",
+        "UG_BUDGET_CUJ_SP_CLIENT_ID",
+        "UG_BUDGET_CUJ_SP_CLIENT_SECRET",
     ):
         monkeypatch.setenv(name, "developer-value")
     session = UserSession(tmp_path, Path("/installed/ug"), tmp_path / "artifacts", "test-token")
@@ -27,6 +29,8 @@ def test_cuj_session_does_not_inherit_agent_credentials_or_configuration(tmp_pat
     assert "UCODE_MANAGED_CONFIG_STUB" not in session.env
     assert "UG_CUJ_SP_CLIENT_ID" not in session.env
     assert "UG_CUJ_SP_CLIENT_SECRET" not in session.env
+    assert "UG_BUDGET_CUJ_SP_CLIENT_ID" not in session.env
+    assert "UG_BUDGET_CUJ_SP_CLIENT_SECRET" not in session.env
     assert session.env["CODEX_HOME"] == str(session.home / ".codex")
     assert not list(session.home.iterdir()), "Fixtures must not manufacture ug/agent state"