From c21003a1c9aa960dfec28005b8f5b1aa20711316 Mon Sep 17 00:00:00 2001 From: Edwin He Date: Wed, 8 Jul 2026 23:21:20 +0000 Subject: [PATCH] ucode: add --skip-preflight to skip per-launch auth/gateway re-validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every `ucode ` launch re-runs configure_shared_state, which does an auth login check, a token fetch, an AI Gateway v2 probe, and (outside provider mode) model discovery — ~5-10s of network round-trips — even when a prior `ucode configure` already validated all of it. Managed/headless launchers (e.g. omnigent) that run configure once and then launch repeatedly pay this on every turn. Add a launch-only `--skip-preflight` flag (distinct from the configure-only `--skip-validate`, which skips the post-configure model smoke test). It threads through _launch_tool as configure_shared_state(skip_preflight=True). configure_shared_state now assembles the shared workspace state (profile, base URLs, auth mode) up front and returns early when skip_preflight is set; otherwise it falls through into the preflight block (auth + AI Gateway validation, then model discovery). The PAT/bearer is already exported by apply_pat_environment and the gateway was verified by the earlier configure, so the skipped work is redundant, and the saved model lists are preserved. Wired into every launch command (codex/claude/gemini/opencode/copilot/pi), default off. Co-authored-by: Isaac --- src/ucode/cli.py | 138 ++++++++++++++++++++++++++++++++-------------- tests/test_cli.py | 105 ++++++++++++++++++++++++++++++++++- 2 files changed, 202 insertions(+), 41 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..9ebdae9 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -216,6 +216,7 @@ def configure_shared_state( force_login: bool = False, use_pat: bool | None = None, skip_model_discovery: bool = False, + skip_preflight: bool = False, ) -> dict: """Log into Databricks, enforce AI Gateway v2, fetch model lists, persist state. @@ -229,6 +230,12 @@ def configure_shared_state( ``--profile`` to every CLI invocation so ambiguous `~/.databrickscfg` entries (e.g. DEFAULT and a named profile both pointing at the same host) don't error out. If ``None``, we resolve it from the host after login. + If skip_preflight is True, skip the entire preflight block below — auth + validation, the AI Gateway probe, and model discovery — trusting a prior + ``ucode configure``. The PAT/bearer is already exported (``apply_pat_environment`` + 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. """ workspace = normalize_workspace_url(workspace) prior_state = load_state() @@ -236,6 +243,48 @@ def configure_shared_state( if use_pat is None: use_pat = bool(prior_state.get("use_pat")) and previous_workspace == workspace fetch_all = tools is None + + # Assemble the shared workspace state that doesn't depend on model discovery: + # workspace, profile, auth mode, base URLs. `profile` may still be None here; + # each path below resolves it once, where a host->profile lookup is reliable + # (the skip branch trusts the prior configure; the preflight resolves after + # login). --skip-preflight persists exactly this and returns, trusting a prior + # `ucode configure` — it already validated auth + the AI Gateway and saved the + # model lists (carried over by load_state, left untouched). + state = load_state() + state["workspace"] = workspace + if profile: + state["profile"] = profile + else: + state.pop("profile", None) + # UC discovery is now always-on; drop any flag persisted by older versions. + state.pop("uc_enabled", None) + # Persist the auth mode so launches rebuild the same (PAT-based) agent + # auth command; an explicit re-configure without --use-pat clears it. + if use_pat: + state["use_pat"] = True + else: + state.pop("use_pat", None) + state["base_urls"] = build_shared_base_urls(workspace) + + if skip_preflight: + # A prior `ucode configure` created the profile; resolve it locally (no + # login needed) and persist it so launches disambiguate. + if profile is None: + profile = find_profile_name_for_host(workspace) + if profile: + state["profile"] = profile + save_state(state) + # Scrub MCP entries ucode wrote for a previous workspace. + if previous_workspace and previous_workspace != workspace: + purge_cross_workspace_mcp_residue(state, workspace) + # Diagnostic reasons are transient (attached after save_state so they + # don't land on disk). No discovery ran, so there is nothing to report. + state["_discovery_reasons"] = {"claude": None, "gemini": None, "codex": None, "oss": None} + return state + + # ── Preflight (bypassed above under --skip-preflight): validate Databricks + # auth + the AI Gateway, then discover the available models. ── if use_pat: if not profile: raise RuntimeError( @@ -260,9 +309,11 @@ def configure_shared_state( else: ensure_databricks_auth(workspace, profile) # After login the profile exists in ~/.databrickscfg, so a host->profile - # lookup is reliable. Persist it so subsequent CLI calls disambiguate. + # lookup is reliable even when it returned nothing above. if profile is None: profile = find_profile_name_for_host(workspace) + if profile: + state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): token = get_databricks_token(workspace, profile) ensure_ai_gateway_v2(workspace, token) @@ -283,6 +334,7 @@ def configure_shared_state( gemini_models = [] codex_models = [] oss_models = [] + opencode_models: dict[str, list[str]] = {} web_search_model: str | None = None if skip_model_discovery: # Provider mode: the agent routes through a Model Provider Service and @@ -295,10 +347,11 @@ def configure_shared_state( if ws_models: web_search_model = ws_models[0] else: - # UC-first, best-effort: one UC model-services call yields all families as - # `system.ai.` ids, bucketed by name. If a family comes back - # empty (workspace without UC model-services, or the listing failed), fall - # back to the per-family AI Gateway listing for that family only. + # UC-first, best-effort: one UC model-services call yields all families + # as `system.ai.` ids, bucketed by name. If a family comes + # back empty (workspace without UC model-services, or the listing + # failed), fall back to the per-family AI Gateway listing for that + # family only. with spinner("Fetching available models..."): ms_claude, ms_codex, ms_gemini, ms_oss, ms_reason = discover_model_services( workspace, token @@ -317,30 +370,13 @@ def configure_shared_state( codex_models, codex_reason = discover_codex_models(workspace, token) if want_oss: oss_models, oss_reason = ms_oss, ms_reason - opencode_models: dict[str, list[str]] = {} - if claude_models: - opencode_models["anthropic"] = list(claude_models.values()) - if gemini_models: - opencode_models["gemini"] = gemini_models - if oss_models: - opencode_models["oss"] = oss_models - - # Merge into existing workspace state so prior tool configs are preserved. - state = load_state() - state["workspace"] = workspace - if profile: - state["profile"] = profile - else: - state.pop("profile", None) - # UC discovery is now always-on; drop any flag persisted by older versions. - state.pop("uc_enabled", None) - # Persist the auth mode so launches rebuild the same (PAT-based) agent - # auth command; an explicit re-configure without --use-pat clears it. - if use_pat: - state["use_pat"] = True - else: - state.pop("use_pat", None) - state["base_urls"] = build_shared_base_urls(workspace) + if claude_models: + opencode_models["anthropic"] = list(claude_models.values()) + if gemini_models: + opencode_models["gemini"] = gemini_models + if oss_models: + opencode_models["oss"] = oss_models + if skip_model_discovery: # Don't clobber any previously-discovered Databricks model lists; provider # mode just doesn't refresh or use them. Persist the web-search model so @@ -822,7 +858,12 @@ def _auto_configure_tool(tool: str) -> None: raise RuntimeError(f"{spec['display']} validation failed — config reverted.") -def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None) -> None: +def _launch_tool( + tool_name: str, + ctx: typer.Context, + provider: str | None = None, + skip_preflight: bool = False, +) -> None: try: tool = normalize_tool(tool_name) existing = load_state() @@ -860,6 +901,7 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None profile=state.get("profile"), tools=[tool], skip_model_discovery=bool(provider), + skip_preflight=skip_preflight, ) if provider: # Routing through a Model Provider Service pins no Databricks model; @@ -892,6 +934,20 @@ def _launch_tool(tool_name: str, ctx: typer.Context, provider: str | None = None raise typer.Exit(130) from None +# Launch-only escape hatch for managed/headless launchers (e.g. omnigent) that +# have already run `ucode configure`: skip the ~5-10s per-launch auth + AI +# Gateway re-validation. Distinct from the configure-only `--skip-validate`, +# which skips the model smoke test. +SkipPreflightOption = Annotated[ + bool, + typer.Option( + "--skip-preflight", + help="Skip the per-launch Databricks auth + AI Gateway re-validation, trusting a " + "prior `ucode configure`.", + ), +] + + @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def codex_cmd( ctx: typer.Context, @@ -904,9 +960,10 @@ def codex_cmd( "before any `--` separator.", ), ] = None, + skip_preflight: SkipPreflightOption = False, ) -> None: """Launch Codex via Databricks.""" - _launch_tool("codex", ctx, provider=provider) + _launch_tool("codex", ctx, provider=provider, skip_preflight=skip_preflight) @app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -921,35 +978,36 @@ def claude_cmd( "before any `--` separator.", ), ] = None, + skip_preflight: SkipPreflightOption = False, ) -> None: """Launch Claude Code via Databricks.""" - _launch_tool("claude", ctx, provider=provider) + _launch_tool("claude", ctx, provider=provider, skip_preflight=skip_preflight) @app.command("gemini", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def gemini_cmd(ctx: typer.Context) -> None: +def gemini_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: """Launch Gemini CLI via Databricks.""" - _launch_tool("gemini", ctx) + _launch_tool("gemini", ctx, skip_preflight=skip_preflight) @app.command( "opencode", context_settings={"allow_extra_args": True, "ignore_unknown_options": True} ) -def opencode_cmd(ctx: typer.Context) -> None: +def opencode_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: """Launch OpenCode via Databricks.""" - _launch_tool("opencode", ctx) + _launch_tool("opencode", ctx, skip_preflight=skip_preflight) @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def copilot_cmd(ctx: typer.Context) -> None: +def copilot_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: """Launch GitHub Copilot CLI via Databricks.""" - _launch_tool("copilot", ctx) + _launch_tool("copilot", ctx, skip_preflight=skip_preflight) @app.command("pi", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) -def pi_cmd(ctx: typer.Context) -> None: +def pi_cmd(ctx: typer.Context, skip_preflight: SkipPreflightOption = False) -> None: """Launch Pi coding agent via Databricks.""" - _launch_tool("pi", ctx) + _launch_tool("pi", ctx, skip_preflight=skip_preflight) @configure_app.callback(invoke_without_command=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..7bdb0e4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,9 +2,10 @@ from __future__ import annotations +import contextlib import os import re -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from typer.testing import CliRunner @@ -1313,3 +1314,105 @@ def _boom(*a, **k): assert state["web_search_model"] == "databricks-gpt-5" # Existing model list preserved, not overwritten to {}. assert state["claude_models"] == {"opus": "databricks-claude-opus-4-8"} + + +class TestConfigureSharedStateSkipPreflight: + """With skip_preflight (--skip-preflight), a prior configure is trusted: + no auth login, token fetch, gateway probe, or model discovery runs — but the + profile and base URLs are still resolved and state is persisted.""" + + WS = "https://cfg.databricks.com" + + @staticmethod + def _stub(monkeypatch): + import ucode.cli as cli_mod + + def _boom(name): + def _f(*a, **k): + raise AssertionError(f"{name} must not run under skip_preflight") + + return _f + + monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) + # Any network round-trip is a hard failure in this mode. + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", _boom("ensure_databricks_auth")) + monkeypatch.setattr(cli_mod, "run_databricks_login", _boom("run_databricks_login")) + monkeypatch.setattr(cli_mod, "ensure_pat_bearer", _boom("ensure_pat_bearer")) + monkeypatch.setattr(cli_mod, "get_databricks_token", _boom("get_databricks_token")) + monkeypatch.setattr(cli_mod, "ensure_ai_gateway_v2", _boom("ensure_ai_gateway_v2")) + monkeypatch.setattr(cli_mod, "discover_model_services", _boom("discover_model_services")) + monkeypatch.setattr(cli_mod, "discover_codex_models", _boom("discover_codex_models")) + monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: "resolved") + monkeypatch.setattr(cli_mod, "build_shared_base_urls", lambda w: {"codex": "u/codex"}) + saved: list[dict] = [] + monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) + return cli_mod, saved + + def test_skips_auth_gateway_and_discovery_but_persists(self, monkeypatch): + cli_mod, saved = self._stub(monkeypatch) + monkeypatch.setattr( + cli_mod, + "load_state", + lambda: {"workspace": self.WS, "codex_models": ["databricks-gpt-5"]}, + ) + + state = cli_mod.configure_shared_state( + self.WS, profile="DEFAULT", tools=["codex"], skip_preflight=True + ) + + # base_urls rebuilt and state saved, but the prior model list is left intact. + assert state["base_urls"] == {"codex": "u/codex"} + assert state["codex_models"] == ["databricks-gpt-5"] + assert saved and saved[-1]["base_urls"] == {"codex": "u/codex"} + + def test_resolves_profile_locally_when_missing(self, monkeypatch): + cli_mod, _ = self._stub(monkeypatch) + monkeypatch.setattr(cli_mod, "load_state", lambda: {"workspace": self.WS}) + + state = cli_mod.configure_shared_state(self.WS, profile=None, skip_preflight=True) + + # find_profile_name_for_host is a local ~/.databrickscfg lookup (no network). + assert state["profile"] == "resolved" + + +class TestSkipPreflightFlag: + """`--skip-preflight` on a launch command threads through _launch_tool to + configure_shared_state as skip_preflight.""" + + LAUNCH_TOOLS = ["codex", "claude", "gemini", "opencode", "copilot", "pi"] + + @staticmethod + def _patches(cfg): + return [ + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli._auto_configure_tool"), + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state", cfg), + patch( + "ucode.cli.resolve_launch_model", + return_value=(MINIMAL_STATE, "databricks-claude-sonnet-4"), + ), + patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE), + patch("ucode.cli.launch_agent"), + ] + + @pytest.mark.parametrize("tool", LAUNCH_TOOLS) + def test_flag_sets_skip_preflight_true(self, tool): + cfg = MagicMock(return_value=MINIMAL_STATE) + with contextlib.ExitStack() as stack: + for p in self._patches(cfg): + stack.enter_context(p) + result = runner.invoke(app, [tool, "--skip-preflight"]) + assert result.exit_code == 0, result.output + assert cfg.call_args.kwargs["skip_preflight"] is True + + @pytest.mark.parametrize("tool", ["codex", "gemini"]) + def test_absent_flag_defaults_false(self, tool): + cfg = MagicMock(return_value=MINIMAL_STATE) + with contextlib.ExitStack() as stack: + for p in self._patches(cfg): + stack.enter_context(p) + result = runner.invoke(app, [tool]) + assert result.exit_code == 0, result.output + assert cfg.call_args.kwargs["skip_preflight"] is False