From 2ef7d984c4265a73df8cf89b9b032cee87d38056 Mon Sep 17 00:00:00 2001 From: Edwin He Date: Wed, 8 Jul 2026 23:21:20 +0000 Subject: [PATCH] ucode: add --no-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 per-launch 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 cost on every turn. Add a launch-only `--no-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(no_preflight=True), which skips the auth+gateway block and model discovery entirely — the PAT/bearer is already exported by apply_pat_environment and the gateway was verified by the earlier configure. The local profile resolution, base-URL rebuild, and state persistence still run, and previously-discovered model lists are preserved rather than clobbered with empties. Wired into every launch command (codex/claude/gemini/opencode/copilot/pi). Co-authored-by: Isaac --- src/ucode/cli.py | 198 +++++++++++++++++++++++++++++----------------- tests/test_cli.py | 105 +++++++++++++++++++++++- 2 files changed, 228 insertions(+), 75 deletions(-) diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..b1fd396 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, + no_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 no_preflight is True, trust a prior ``ucode configure`` and skip the + per-launch network round-trips entirely: no auth login, no token fetch, no + AI Gateway probe, and no model discovery. The PAT/bearer is already exported + (``apply_pat_environment`` in ``_launch_tool``) and the gateway was verified + by that earlier configure, so a managed/headless re-launch needs none of it. + The local profile/base-url derivation and state persistence still run. """ workspace = normalize_workspace_url(workspace) prior_state = load_state() @@ -236,37 +243,47 @@ 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 - if use_pat: - if not profile: - raise RuntimeError( - "--use-pat requires a Databricks CLI profile. Pass one via `--profiles `." - ) - pat = resolve_pat_token(profile) - if not pat: - raise RuntimeError( - f"--use-pat: profile '{profile}' has no personal access token in " - "~/.databrickscfg (its auth_type must be `pat`). Add a `token = ` " - f"entry under [{profile}], or re-run without --use-pat to use OAuth." - ) - # Export the PAT for this process and launched agent subprocesses so - # every token fetch takes the static-bearer path. ensure_pat_bearer - # keeps a non-empty pre-set bearer (CI escape hatch) but treats an - # empty one as absent, so it never shadows the PAT. Pass the validated - # token to avoid re-reading ~/.databrickscfg. - ensure_pat_bearer(profile, pat) - ensure_databricks_auth(workspace, profile) - elif force_login: - run_databricks_login(workspace, profile) + token: str | None = None + if no_preflight: + # A prior `ucode configure` already logged in and verified the AI Gateway, + # and the PAT/bearer is already exported (apply_pat_environment), so skip + # the auth login, token fetch, and gateway probe. Still resolve the + # profile locally from ~/.databrickscfg so downstream CLI calls + # disambiguate; no token is needed since model discovery is skipped too. + if profile is None: + profile = find_profile_name_for_host(workspace) 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. - if profile is None: - profile = find_profile_name_for_host(workspace) - with spinner("Verifying Unity AI Gateway..."): - token = get_databricks_token(workspace, profile) - ensure_ai_gateway_v2(workspace, token) - print_success("Unity AI Gateway detected") + if use_pat: + if not profile: + raise RuntimeError( + "--use-pat requires a Databricks CLI profile. Pass one via `--profiles `." + ) + pat = resolve_pat_token(profile) + if not pat: + raise RuntimeError( + f"--use-pat: profile '{profile}' has no personal access token in " + "~/.databrickscfg (its auth_type must be `pat`). Add a `token = ` " + f"entry under [{profile}], or re-run without --use-pat to use OAuth." + ) + # Export the PAT for this process and launched agent subprocesses so + # every token fetch takes the static-bearer path. ensure_pat_bearer + # keeps a non-empty pre-set bearer (CI escape hatch) but treats an + # empty one as absent, so it never shadows the PAT. Pass the validated + # token to avoid re-reading ~/.databrickscfg. + ensure_pat_bearer(profile, pat) + ensure_databricks_auth(workspace, profile) + elif force_login: + run_databricks_login(workspace, profile) + 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. + if profile is None: + profile = find_profile_name_for_host(workspace) + with spinner("Verifying Unity AI Gateway..."): + token = get_databricks_token(workspace, profile) + ensure_ai_gateway_v2(workspace, token) + print_success("Unity AI Gateway detected") want_claude = ( fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools @@ -284,39 +301,45 @@ def configure_shared_state( codex_models = [] oss_models = [] web_search_model: str | None = None - if skip_model_discovery: - # Provider mode: the agent routes through a Model Provider Service and - # pins no Databricks model, so the full family discovery is unused. Web - # search (claude only) still needs one Responses-capable model, so fetch - # just that with a single call. - if want_claude: - with spinner("Fetching web search model..."): - ws_models, _ = discover_codex_models(workspace, token) - 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. - with spinner("Fetching available models..."): - ms_claude, ms_codex, ms_gemini, ms_oss, ms_reason = discover_model_services( - workspace, token - ) + # Under no_preflight no discovery runs at all — the prior configure's + # model lists are trusted (and preserved below, not overwritten). + if not no_preflight: + # token is always set by the auth+gateway block above when validating. + assert token is not None + if skip_model_discovery: + # Provider mode: the agent routes through a Model Provider Service and + # pins no Databricks model, so the full family discovery is unused. Web + # search (claude only) still needs one Responses-capable model, so fetch + # just that with a single call. if want_claude: - claude_models, claude_reason = ms_claude, ms_reason - if not claude_models: - claude_models, claude_reason = discover_claude_models(workspace, token) - if want_gemini: - gemini_models, gemini_reason = ms_gemini, ms_reason - if not gemini_models: - gemini_models, gemini_reason = discover_gemini_models(workspace, token) - if want_codex: - codex_models, codex_reason = ms_codex, ms_reason - if not codex_models: - codex_models, codex_reason = discover_codex_models(workspace, token) - if want_oss: - oss_models, oss_reason = ms_oss, ms_reason + with spinner("Fetching web search model..."): + ws_models, _ = discover_codex_models(workspace, token) + 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. + with spinner("Fetching available models..."): + ms_claude, ms_codex, ms_gemini, ms_oss, ms_reason = discover_model_services( + workspace, token + ) + if want_claude: + claude_models, claude_reason = ms_claude, ms_reason + if not claude_models: + claude_models, claude_reason = discover_claude_models(workspace, token) + if want_gemini: + gemini_models, gemini_reason = ms_gemini, ms_reason + if not gemini_models: + gemini_models, gemini_reason = discover_gemini_models(workspace, token) + if want_codex: + codex_models, codex_reason = ms_codex, ms_reason + if not codex_models: + 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()) @@ -341,7 +364,11 @@ def configure_shared_state( else: state.pop("use_pat", None) state["base_urls"] = build_shared_base_urls(workspace) - if skip_model_discovery: + if no_preflight: + # No discovery ran; preserve the model lists the prior configure wrote + # (they'd otherwise be clobbered with empties below). + pass + elif 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 # claude's web_search MCP keeps working through the normal gateway. @@ -822,7 +849,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, + no_preflight: bool = False, +) -> None: try: tool = normalize_tool(tool_name) existing = load_state() @@ -860,6 +892,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), + no_preflight=no_preflight, ) if provider: # Routing through a Model Provider Service pins no Databricks model; @@ -892,6 +925,21 @@ 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. +NoPreflightOption = Annotated[ + bool, + typer.Option( + "--no-preflight", + help="Skip the per-launch Databricks auth + AI Gateway re-validation, trusting a " + "prior `ucode configure` (used by managed/headless launchers where configure " + "already ran).", + ), +] + + @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def codex_cmd( ctx: typer.Context, @@ -904,9 +952,10 @@ def codex_cmd( "before any `--` separator.", ), ] = None, + no_preflight: NoPreflightOption = False, ) -> None: """Launch Codex via Databricks.""" - _launch_tool("codex", ctx, provider=provider) + _launch_tool("codex", ctx, provider=provider, no_preflight=no_preflight) @app.command("claude", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) @@ -921,35 +970,36 @@ def claude_cmd( "before any `--` separator.", ), ] = None, + no_preflight: NoPreflightOption = False, ) -> None: """Launch Claude Code via Databricks.""" - _launch_tool("claude", ctx, provider=provider) + _launch_tool("claude", ctx, provider=provider, no_preflight=no_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, no_preflight: NoPreflightOption = False) -> None: """Launch Gemini CLI via Databricks.""" - _launch_tool("gemini", ctx) + _launch_tool("gemini", ctx, no_preflight=no_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, no_preflight: NoPreflightOption = False) -> None: """Launch OpenCode via Databricks.""" - _launch_tool("opencode", ctx) + _launch_tool("opencode", ctx, no_preflight=no_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, no_preflight: NoPreflightOption = False) -> None: """Launch GitHub Copilot CLI via Databricks.""" - _launch_tool("copilot", ctx) + _launch_tool("copilot", ctx, no_preflight=no_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, no_preflight: NoPreflightOption = False) -> None: """Launch Pi coding agent via Databricks.""" - _launch_tool("pi", ctx) + _launch_tool("pi", ctx, no_preflight=no_preflight) @configure_app.callback(invoke_without_command=True) diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..b6dfd46 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 TestConfigureSharedStateNoPreflight: + """With no_preflight (--no-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 no_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"], no_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, no_preflight=True) + + # find_profile_name_for_host is a local ~/.databrickscfg lookup (no network). + assert state["profile"] == "resolved" + + +class TestNoPreflightFlag: + """`--no-preflight` on a launch command threads through _launch_tool to + configure_shared_state as no_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_no_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, "--no-preflight"]) + assert result.exit_code == 0, result.output + assert cfg.call_args.kwargs["no_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["no_preflight"] is False