ucode: add --no-preflight to skip per-launch auth/gateway re-validation

Every `ucode <agent>` 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

Author: Edwinhe03

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,44 +230,60 @@ 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()
     previous_workspace = prior_state.get("workspace")
     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 <name>`."
-            )
-        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 = <PAT>` "
-                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 <name>`."
+                )
+            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 = <PAT>` "
+                    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.<model-name>` 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.<model-name>` 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)