Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Edwinhe03 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 124 additions & 74 deletions src/ucode/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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
Expand All @@ -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())
Expand All @@ -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.
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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})
Expand All @@ -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)
Expand Down
Loading
Loading