Refresh Claude MCP auth via headersHelper (no mid-session expiry)

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

src/ucode/cli.py

@@ -746,14 +746,27 @@ 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 <token>"}` 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.
 
     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."""
+    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")

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)

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 <token>"}`` 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)

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"}),