Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 2 additions & 1 deletion src/ucode/agents/claude.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
build_tool_base_url,
get_databricks_token,
)
from ucode.launcher import exec_or_spawn
from ucode.state import mark_tool_managed, save_state
from ucode.telemetry import agent_version, ucode_version
from ucode.tracing import tracing_env
Expand Down Expand Up @@ -467,7 +468,7 @@ def launch(state: dict, tool_args: list[str]) -> None:
workspace = state.get("workspace")
if workspace:
os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
os.execvp(binary, [binary, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args])
exec_or_spawn([binary, "--settings", str(CLAUDE_SETTINGS_PATH), *tool_args])


def validate_cmd(binary: str) -> list[str]:
Expand Down
13 changes: 8 additions & 5 deletions src/ucode/agents/codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,11 @@
write_toml_file,
)
from ucode.databricks import (
build_auth_shell_command,
build_auth_token_argv,
build_tool_base_url,
get_databricks_token,
)
from ucode.launcher import exec_or_spawn
from ucode.state import mark_tool_managed, save_state
from ucode.telemetry import agent_version, ucode_version

Expand Down Expand Up @@ -102,7 +103,7 @@ def _use_legacy_layout() -> bool:


def _provider_block(workspace: str, databricks_profile: str | None, use_pat: bool = False) -> dict:
auth_command = build_auth_shell_command(workspace, databricks_profile, use_pat=use_pat)
auth_argv = build_auth_token_argv(workspace, databricks_profile, use_pat=use_pat)
base_url = build_tool_base_url("codex", workspace)
return {
"name": "Databricks AI Gateway",
Expand All @@ -111,9 +112,11 @@ def _provider_block(workspace: str, databricks_profile: str | None, use_pat: boo
"http_headers": {
"User-Agent": f"ucode/{ucode_version()} codex/{agent_version('codex')}",
},
# Run the `ucode auth-token` executable directly (not via `sh -c`) so the
# helper works on Windows, where there is no POSIX shell (issue #116).
"auth": {
"command": "sh",
"args": ["-c", auth_command],
"command": auth_argv[0],
"args": auth_argv[1:],
"timeout_ms": 5000,
"refresh_interval_ms": 900000,
},
Expand Down Expand Up @@ -356,7 +359,7 @@ def launch(state: dict, tool_args: list[str]) -> None:
workspace = state.get("workspace")
if workspace:
os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile"))
os.execvp(binary, [binary, "--profile", CODEX_PROFILE_NAME, *tool_args])
exec_or_spawn([binary, "--profile", CODEX_PROFILE_NAME, *tool_args])


def validate_cmd(binary: str) -> list[str]:
Expand Down
60 changes: 56 additions & 4 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

from __future__ import annotations

import os
from typing import Annotated

import typer
Expand Down Expand Up @@ -38,6 +37,7 @@
discover_model_services,
ensure_ai_gateway_v2,
ensure_databricks_auth,
ensure_pat_bearer,
find_profile_name_for_host,
get_databricks_profiles,
get_databricks_token,
Expand Down Expand Up @@ -228,9 +228,11 @@ def configure_shared_state(
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; a bearer already in
# the environment wins.
os.environ.setdefault("DATABRICKS_BEARER", pat)
# 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)
Expand Down Expand Up @@ -597,6 +599,56 @@ def mcp_web_search_cmd() -> None:
serve()


@app.command("auth-token", hidden=True)
def auth_token_cmd(
host: Annotated[
str | None, typer.Option("--host", help="Workspace URL. Defaults to the saved workspace.")
] = None,
profile: Annotated[
str | None, typer.Option("--profile", help="Databricks CLI profile.")
] = None,
use_pat: Annotated[
bool, typer.Option("--use-pat", help="Read the profile's static PAT instead of OAuth.")
] = False,
) -> None:
"""Print a Databricks bearer token to stdout, then exit.

This is the cross-platform helper invoked by Claude Code's `apiKeyHelper`
and Codex's auth command on every token refresh. It is not meant for
interactive use. All token logic (DATABRICKS_BEARER short-circuit, PAT
profiles, OAuth refresh) lives in `get_databricks_token`, so the same
binary works on macOS, Linux, and Windows without any POSIX shell."""
import sys

state = load_state()
workspace = host or state.get("workspace")
if not workspace:
print_err("No workspace configured. Run `ucode configure` first.")
raise typer.Exit(1)
profile = profile or state.get("profile")
if use_pat or state.get("use_pat"):
# --use-pat explicitly means "serve the profile's static PAT". Fail
# closed if it can't be read rather than falling through to OAuth —
# `auth token` cannot serve a PAT-only profile, so that path would
# surface a misleading stale-login error instead of the real cause.
if not ensure_pat_bearer(profile):
print_err(
f"--use-pat: no personal access token available for profile "
f"'{profile or '<none>'}'. Add a `token = <PAT>` entry under "
f"[{profile or 'your-profile'}] in ~/.databrickscfg, or re-run "
"`ucode configure` without --use-pat to use OAuth."
)
raise typer.Exit(1)
try:
token = get_databricks_token(workspace, profile)
except RuntimeError as exc:
print_err(str(exc))
raise typer.Exit(1) from None
# Write the bare token (with trailing newline) to stdout — nothing else may
# land on stdout or the consuming agent will treat it as part of the token.
sys.stdout.write(token + "\n")


def _auto_configure_tool(tool: str) -> None:
"""First-time setup for a single tool — mirrors configure_workspace_command."""
existing = load_state()
Expand Down
91 changes: 60 additions & 31 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -704,19 +704,39 @@ def resolve_pat_token(profile: str | None) -> str | None:
return None


def ensure_pat_bearer(profile: str | None, pat: str | None = None) -> bool:
"""Ensure ``DATABRICKS_BEARER`` holds a usable token for a ``--use-pat`` profile.

If a non-empty bearer is already in the environment it wins (the CI escape
hatch). Otherwise the profile's static PAT is exported — callers that have
already resolved it (e.g. ``configure_shared_state``) pass it via ``pat`` to
skip a redundant ``~/.databrickscfg`` read; everyone else lets this resolve
it. An exported-but-*empty* ``DATABRICKS_BEARER`` is treated as absent —
matching ``get_databricks_token``'s own ``.strip()`` check — so a stray
``export DATABRICKS_BEARER=`` does not shadow the PAT and silently force the
OAuth path (which fails for PAT-only profiles).

Returns ``True`` iff a usable bearer is now present in the environment."""
if os.environ.get("DATABRICKS_BEARER", "").strip():
return True
pat = pat or resolve_pat_token(profile)
if pat:
os.environ["DATABRICKS_BEARER"] = pat
return True
return False


def apply_pat_environment(state: dict) -> None:
"""Export the configured profile's PAT as ``DATABRICKS_BEARER`` when the
workspace was configured with ``--use-pat``.

Every token fetch in this process (and in launched agent subprocesses,
which inherit the environment) then takes the existing static-bearer
short-circuit instead of the OAuth-only `databricks auth token` path.
A bearer already present in the environment is left untouched."""
A non-empty bearer already present in the environment is left untouched."""
if not state.get("use_pat"):
return
pat = resolve_pat_token(state.get("profile"))
if pat:
os.environ.setdefault("DATABRICKS_BEARER", pat)
ensure_pat_bearer(state.get("profile"))


def run_databricks_login(workspace: str, profile: str | None = None) -> None:
Expand Down Expand Up @@ -1026,36 +1046,45 @@ def list_databricks_apps(workspace: str, profile: str | None = None) -> list[dic
raise RuntimeError("Databricks apps listing returned invalid JSON.") from exc


def _ucode_binary() -> str:
"""Resolve the absolute path to the running `ucode` executable.

Agents persist the auth command into config files and re-run it on every
token refresh, possibly from launchers without a full PATH (desktop GUIs).
An absolute path keeps the helper working regardless of PATH. Falls back to
the bare name when resolution fails."""
return shutil.which("ucode") or "ucode"


def build_auth_token_argv(
workspace: str, profile: str | None = None, *, use_pat: bool = False
) -> list[str]:
"""Argv for the cross-platform token helper: `ucode auth-token ...`.

Unlike the previous POSIX `databricks ... | jq` pipeline, this is a single
executable with plain arguments — no `sh`, no `jq`, no shell quoting — so it
runs identically on macOS, Linux, and Windows (issue #116). The DATABRICKS_BEARER
short-circuit and the PAT path both live inside `auth-token` itself."""
argv = [_ucode_binary(), "auth-token", "--host", workspace.rstrip("/")]
if profile:
argv += ["--profile", profile]
if use_pat:
argv.append("--use-pat")
return argv


def build_auth_shell_command(
workspace: str, profile: str | None = None, *, use_pat: bool = False
) -> str:
workspace_arg = shlex.quote(workspace.rstrip("/"))
if use_pat and profile:
# --use-pat profiles have no OAuth cache for `auth token` to read, so
# the persisted command reads the profile's static token instead.
profile_arg = shlex.quote(profile)
cli_command = (
f"databricks auth describe --profile {profile_arg} --sensitive --output json "
"| jq -r '.details.configuration.token.value'"
)
elif profile:
profile_arg = shlex.quote(profile)
cli_command = (
f"databricks auth token --host {workspace_arg} "
f"--profile {profile_arg} --force-refresh --output json "
"| jq -r '.access_token'"
)
else:
cli_command = (
"env -u DATABRICKS_CONFIG_PROFILE "
f"databricks auth token --host {workspace_arg} --force-refresh --output json "
"| jq -r '.access_token'"
)
return (
'if [ -n "${DATABRICKS_BEARER:-}" ]; then '
'printf "%s\\n" "$DATABRICKS_BEARER"; '
f"else {cli_command}; fi"
)
"""Single-line, shell-quoted form of :func:`build_auth_token_argv`.

Used where a tool wants the helper as one command *string* (Claude Code's
`apiKeyHelper`). On every platform this resolves to the `ucode auth-token`
executable rather than a POSIX shell pipeline, so no `sh`/`jq` is required."""
argv = build_auth_token_argv(workspace, profile, use_pat=use_pat)
if platform.system() == "Windows":
return subprocess.list2cmdline(argv)
return shlex.join(argv)


# A model-service's `name` is `model-services/system.ai.<model-name>`; the
Expand Down
36 changes: 36 additions & 0 deletions src/ucode/launcher.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
"""Cross-platform process replacement for launching coding agents."""

from __future__ import annotations

import os
import signal
import subprocess
import sys


def exec_or_spawn(argv: list[str]) -> None:
"""Hand the terminal to ``argv``, then exit with its status.

On POSIX we ``os.execvp`` — the agent process *replaces* ucode, inheriting
the controlling terminal cleanly.

On Windows there is no real ``exec``: ``os.execvp`` spawns a *new* process
and immediately terminates the parent, so the launching shell resumes its
prompt and fights the agent for the console. That produces the garbled,
split-screen input reported in issue #173. Instead we spawn a child, wait
for it, and propagate its exit code — the same pattern the token-refreshing
agents (gemini/opencode/copilot/pi) already use.
"""
if os.name != "nt":
os.execvp(argv[0], argv)
return # unreachable on POSIX; keeps type-checkers happy

proc = subprocess.Popen(argv)
try:
returncode = proc.wait()
except KeyboardInterrupt:
# Ctrl-C is delivered to the whole console group; let the child handle
# it and report its own exit code rather than racing it.
proc.send_signal(signal.SIGINT)
returncode = proc.wait()
sys.exit(returncode)
14 changes: 10 additions & 4 deletions src/ucode/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@
import json

from ucode.config_io import APP_DIR, is_dry_run
from ucode.databricks import build_auth_shell_command, build_shared_base_urls
from ucode.databricks import (
build_auth_shell_command,
build_auth_token_argv,
build_shared_base_urls,
)

STATE_PATH = APP_DIR / "state.json"
STATE_VERSION = 3
Expand Down Expand Up @@ -124,7 +128,9 @@ def build_agent_state(state: dict) -> dict[str, dict]:
profile = state.get("profile") if isinstance(state.get("profile"), str) else None
base_urls_value = state.get("base_urls")
base_urls = base_urls_value if isinstance(base_urls_value, dict) else {}
auth_command = build_auth_shell_command(workspace, profile, use_pat=bool(state.get("use_pat")))
use_pat = bool(state.get("use_pat"))
auth_command = build_auth_shell_command(workspace, profile, use_pat=use_pat)
auth_argv = build_auth_token_argv(workspace, profile, use_pat=use_pat)
claude_models_value = state.get("claude_models")
claude_models: dict = claude_models_value if isinstance(claude_models_value, dict) else {}
codex_models_value = state.get("codex_models")
Expand Down Expand Up @@ -155,8 +161,8 @@ def build_agent_state(state: dict) -> dict[str, dict]:
"base_url": base_urls.get("codex"),
"auth_command": auth_command,
"auth": {
"command": "sh",
"args": ["-c", auth_command],
"command": auth_argv[0],
"args": auth_argv[1:],
"timeout_ms": AUTH_COMMAND_TIMEOUT_MS,
"refresh_interval_ms": AUTH_REFRESH_INTERVAL_MS,
},
Expand Down
9 changes: 6 additions & 3 deletions tests/test_agent_codex.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,11 +45,14 @@ def test_provider_wire_api(self):
provider = overlay["model_providers"]["ucode-databricks"]
assert provider["wire_api"] == "responses"

def test_auth_uses_sh(self):
def test_auth_runs_ucode_auth_token(self):
# The auth command runs the `ucode auth-token` executable directly
# (not `sh -c`), so it works on Windows where there is no POSIX shell.
overlay = codex.render_overlay(WS)
auth = overlay["model_providers"]["ucode-databricks"]["auth"]
assert auth["command"] == "sh"
assert "-c" in auth["args"]
assert auth["command"].endswith("ucode") or auth["command"] == "ucode"
assert auth["args"][0] == "auth-token"
assert auth["command"] != "sh"

def test_auth_contains_workspace(self):
overlay = codex.render_overlay(WS)
Expand Down
Loading
Loading