diff --git a/README.md b/README.md index 0420190..59a33dd 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,13 @@ Options are shown in this order: Discovered external MCP connections are listed directly. MCP auth uses a Databricks token that `ucode` sets when launching each tool. +For **Claude Code**, the MCP server is registered with a `headersHelper` command +(`ucode auth-token --mcp-header`) that Claude re-runs on every connect/reconnect and on a +401/403 retry, so the Databricks OAuth token is refreshed automatically and never goes stale +mid-session. Other tools (Codex, Gemini CLI, OpenCode, GitHub Copilot CLI) receive the token via +the `OAUTH_TOKEN` env var set at launch; those clients expose no per-request auth hook for MCP, so +a session running past the token lifetime needs to be relaunched to pick up a fresh token. + --- ## Other Commands diff --git a/src/ucode/cli.py b/src/ucode/cli.py index b2f23be..d4c702a 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -746,6 +746,14 @@ def auth_token_cmd( use_pat: Annotated[ bool, typer.Option("--use-pat", help="Read the profile's static PAT instead of OAuth.") ] = False, + mcp_header: Annotated[ + bool, + typer.Option( + "--mcp-header", + help='Emit a JSON `{"Authorization": "Bearer "}` object instead of the ' + "bare token, for use as a Claude Code MCP `headersHelper`.", + ), + ] = False, ) -> None: """Print a Databricks bearer token to stdout, then exit. @@ -753,7 +761,12 @@ def auth_token_cmd( 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.""" + binary works on macOS, Linux, and Windows without any POSIX shell. + + With ``--mcp-header`` the output is the JSON header object Claude Code's MCP + ``headersHelper`` expects; Claude runs this fresh on every connect/reconnect + and on a 401/403 retry, so the MCP bearer never goes stale mid-session.""" + import json as _json import sys state = load_state() @@ -780,6 +793,11 @@ def auth_token_cmd( except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None + if mcp_header: + # Claude Code's MCP headersHelper merges the JSON on stdout into the + # connection headers. Nothing else may land on stdout. + sys.stdout.write(_json.dumps({"Authorization": f"Bearer {token}"}) + "\n") + return # 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") diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 9f89c6b..9bad356 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -1057,31 +1057,37 @@ def _ucode_binary() -> str: def build_auth_token_argv( - workspace: str, profile: str | None = None, *, use_pat: bool = False + workspace: str, profile: str | None = None, *, use_pat: bool = False, mcp_header: 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.""" + short-circuit and the PAT path both live inside `auth-token` itself. + + With ``mcp_header`` the helper emits the JSON ``{"Authorization": "Bearer …"}`` + object Claude Code's MCP ``headersHelper`` expects rather than the bare token.""" argv = [_ucode_binary(), "auth-token", "--host", workspace.rstrip("/")] if profile: argv += ["--profile", profile] if use_pat: argv.append("--use-pat") + if mcp_header: + argv.append("--mcp-header") return argv def build_auth_shell_command( - workspace: str, profile: str | None = None, *, use_pat: bool = False + workspace: str, profile: str | None = None, *, use_pat: bool = False, mcp_header: bool = False ) -> str: """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) + `apiKeyHelper`, or an MCP `headersHelper` with ``mcp_header=True``). 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, mcp_header=mcp_header) if platform.system() == "Windows": return subprocess.list2cmdline(argv) return shlex.join(argv) diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index 6257458..2e8479e 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -28,6 +28,7 @@ from ucode.config_io import restore_file from ucode.databricks import ( apply_pat_environment, + build_auth_shell_command, build_mcp_service_url, ensure_databricks_auth, get_databricks_token, @@ -106,6 +107,23 @@ def build_mcp_http_entry(url: str) -> dict: } +def build_claude_mcp_entry(url: str, headers_helper: str) -> dict: + """Claude Code HTTP MCP entry that authenticates via a ``headersHelper`` + command instead of a static bearer. + + Claude runs ``headersHelper`` fresh on every connect/reconnect and re-runs it + on a 401/403 retry, so the Databricks OAuth token never goes stale + mid-session — unlike the ``Bearer ${OAUTH_TOKEN}`` env-var header, which + Claude freezes at config-load time. The helper (``ucode auth-token + --mcp-header``) emits the ``{"Authorization": "Bearer "}`` JSON Claude + merges into the connection headers.""" + return { + "type": "http", + "url": url, + "headersHelper": headers_helper, + } + + def add_claude_mcp_server(name: str, entry: dict, scope: str = MCP_USER_SCOPE) -> None: try: subprocess.run( @@ -256,12 +274,19 @@ def configured_mcp_clients(state: dict, installed_clients: list[str]) -> list[st ] -def configure_client_mcp_server(client: str, name: str, url: str, entry: dict) -> list[str]: +def configure_client_mcp_server( + client: str, name: str, url: str, entry: dict, claude_headers_helper: str | None = None +) -> list[str]: if client == "claude": removed_scopes = [ scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope) ] - add_claude_mcp_server(name, entry, MCP_USER_SCOPE) + # Prefer a self-refreshing headersHelper over the static ${OAUTH_TOKEN} + # header so the MCP bearer doesn't expire mid-session. + claude_entry = ( + build_claude_mcp_entry(url, claude_headers_helper) if claude_headers_helper else entry + ) + add_claude_mcp_server(name, claude_entry, MCP_USER_SCOPE) return removed_scopes if client == "codex": removed = remove_codex_mcp_server(name) @@ -1019,6 +1044,7 @@ def apply_mcp_server_changes( original_servers: list[dict], working_servers: list[dict], clients: list[str], + claude_headers_helper: str | None = None, ) -> bool: original_by_name = _servers_by_name(original_servers) working_by_name = _servers_by_name(working_servers) @@ -1039,7 +1065,7 @@ def apply_mcp_server_changes( continue entry = build_mcp_http_entry(url) for client in clients: - configure_client_mcp_server(client, name, url, entry) + configure_client_mcp_server(client, name, url, entry, claude_headers_helper) changed = True return changed @@ -1223,6 +1249,17 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None apply_pat_environment(state) ensure_databricks_auth(workspace, profile) + # Claude Code registers MCP servers with a self-refreshing headersHelper + # (`ucode auth-token --mcp-header`) so the bearer never expires mid-session. + # Other clients keep the launch-set ${OAUTH_TOKEN} header (see build_mcp_http_entry). + claude_headers_helper = ( + build_auth_shell_command( + workspace, profile, use_pat=bool(state.get("use_pat")), mcp_header=True + ) + if "claude" in clients + else None + ) + print_section("MCP Servers") client_names = ", ".join(str(MCP_CLIENTS[client]["display"]) for client in clients) print_note(f"Configuring for: {client_names}") @@ -1238,7 +1275,7 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None workspace, profile, clients, location, original_mcp_servers_for_location, services ) changed = apply_mcp_server_changes( - original_mcp_servers_for_location, working_mcp_servers, clients + original_mcp_servers_for_location, working_mcp_servers, clients, claude_headers_helper ) if changed or original_mcp_servers_for_location != working_mcp_servers: state["mcp_servers"] = working_mcp_servers @@ -1326,7 +1363,9 @@ def configure_mcp_command(location: str | None = None, services: set[str] | None ) working_names.add(entry_name) - changed = apply_mcp_server_changes(original_mcp_servers, working_mcp_servers, clients) + changed = apply_mcp_server_changes( + original_mcp_servers, working_mcp_servers, clients, claude_headers_helper + ) if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers save_state(state) diff --git a/tests/test_cli.py b/tests/test_cli.py index 3c5d404..70e9668 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import os import re from unittest.mock import patch @@ -182,6 +183,16 @@ def test_prints_only_the_token_to_stdout(self): assert result.stdout == "tok-123\n" fetch.assert_called_once_with("https://ws", None) + def test_mcp_header_emits_json_authorization_object(self): + with ( + patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), + patch("ucode.cli.get_databricks_token", return_value="tok-123"), + ): + result = runner.invoke(app, ["auth-token", "--mcp-header"]) + assert result.exit_code == 0 + # Claude's MCP headersHelper consumes JSON on stdout; nothing else may land there. + assert json.loads(result.stdout) == {"Authorization": "Bearer tok-123"} + def test_host_and_profile_override_state(self): with ( patch("ucode.cli.load_state", return_value={"workspace": "https://saved"}), diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 81333c8..5bd8dbe 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -26,6 +26,60 @@ def test_uses_oauth_token_header_reference(self): assert "headersHelper" not in entry +class TestBuildClaudeMcpEntry: + def test_uses_headers_helper_not_static_header(self): + entry = mcp.build_claude_mcp_entry( + f"{WS}/api/2.0/mcp/functions/system/ai", "ucode auth-token --host x --mcp-header" + ) + assert entry == { + "type": "http", + "url": f"{WS}/api/2.0/mcp/functions/system/ai", + "headersHelper": "ucode auth-token --host x --mcp-header", + } + # No frozen static token — the helper is the whole point. + assert "headers" not in entry + + +class TestConfigureClaudeUsesHeadersHelper: + def test_claude_branch_registers_headers_helper_entry_when_provided(self, monkeypatch): + added: list[dict] = [] + monkeypatch.setattr(mcp, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr( + mcp, "add_claude_mcp_server", lambda name, entry, scope: added.append(entry) + ) + + url = f"{WS}/api/2.0/mcp/functions/system/ai" + mcp.configure_client_mcp_server( + "claude", + "databricks-system-ai", + url, + mcp.build_mcp_http_entry(url), + claude_headers_helper="ucode auth-token --host x --mcp-header", + ) + + assert added == [ + { + "type": "http", + "url": url, + "headersHelper": "ucode auth-token --host x --mcp-header", + } + ] + + def test_claude_branch_falls_back_to_static_entry_without_helper(self, monkeypatch): + added: list[dict] = [] + monkeypatch.setattr(mcp, "remove_claude_mcp_server", lambda name, scope: False) + monkeypatch.setattr( + mcp, "add_claude_mcp_server", lambda name, entry, scope: added.append(entry) + ) + + url = f"{WS}/api/2.0/mcp/external/github" + entry = mcp.build_mcp_http_entry(url) + mcp.configure_client_mcp_server("claude", "github", url, entry) + + assert added == [entry] + assert added[0]["headers"]["Authorization"] == "Bearer ${OAUTH_TOKEN}" + + class TestAddClaudeMcpServer: def test_adds_user_scoped_json(self, monkeypatch): calls: list[dict] = [] @@ -545,7 +599,7 @@ def test_registers_discovered_external_server(self, monkeypatch): monkeypatch.setattr(mcp, "discover_app_mcp_servers", lambda workspace, profile=None: []) _patch_mcp_choices(monkeypatch, f"{mcp.MCP_ADD_PREFIX}external:github-mcp") - def fake_configure_client_mcp_server(client, name, url, entry): + def fake_configure_client_mcp_server(client, name, url, entry, *a, **kw): configured.append((client, name, url, entry)) return [] @@ -607,7 +661,9 @@ def test_registers_discovered_genie_space_server(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -660,7 +716,9 @@ def test_registers_discovered_vector_search_server(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -714,7 +772,9 @@ def test_registers_discovered_uc_functions_server(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -768,7 +828,9 @@ def test_registers_discovered_app_mcp_server(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -999,7 +1061,9 @@ def test_warns_when_app_selection_is_no_longer_discoverable(self, monkeypatch, c monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1027,7 +1091,9 @@ def test_warns_for_unrecognized_selection_prefix(self, monkeypatch, capsys): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1071,7 +1137,9 @@ def test_continues_when_optional_discovery_fails(self, monkeypatch, capsys): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1142,7 +1210,9 @@ def test_configures_only_ucode_configured_clients(self, monkeypatch, capsys): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1176,7 +1246,9 @@ def test_registers_databricks_sql_server(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1315,7 +1387,9 @@ def fake_list(workspace, token, parent): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1362,7 +1436,9 @@ def test_replaces_servers_outside_location(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr( mcp, @@ -1412,7 +1488,9 @@ def test_existing_entry_gets_reconfigured_for_newly_added_clients(self, monkeypa monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1447,7 +1525,9 @@ def test_configures_only_the_requested_subset(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: saved_states.append(state.copy())) @@ -1476,7 +1556,9 @@ def test_matches_bare_short_names(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: None) @@ -1496,7 +1578,9 @@ def test_unknown_requested_service_warns_and_skips(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: None) monkeypatch.setattr(mcp, "print_warning", lambda msg: warnings.append(msg)) @@ -1531,7 +1615,9 @@ def test_empty_services_removes_everything(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr( mcp, @@ -1576,7 +1662,9 @@ def test_adds_and_removes_to_match_new_selection(self, monkeypatch): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr( mcp, @@ -1613,7 +1701,9 @@ def fake_list(workspace, token, parent): monkeypatch.setattr( mcp, "configure_client_mcp_server", - lambda client, name, url, entry: configured.append((client, name, url, entry)) or [], + lambda client, name, url, entry, *a, **kw: ( + configured.append((client, name, url, entry)) or [] + ), ) monkeypatch.setattr(mcp, "save_state", lambda state: None)