fix(databricks): serialize token --force-refresh with cross-process flock (#190)

Concurrent long-lived ucode sessions each ran `databricks auth token
--force-refresh` independently, racing on the shared OAuth cache and
redeeming (and revoking) the rotating refresh token concurrently.

Serialize --force-refresh calls per workspace+profile behind a
cross-process flock under APP_DIR, and coalesce redundant refreshes:
if a peer force-refreshed within TOKEN_REFRESH_COALESCE_SECONDS (60s),
drop --force-refresh and reuse the now-current cached token instead of
redeeming the refresh token again.

Degrades gracefully to no locking when fcntl is unavailable (Windows)
or lock acquisition fails for any reason -- a stuck/missing lock must
never block a token refresh. The non-force-refresh (cheap read) path
is unchanged and unlocked.

Author: jesco39

src/ucode/databricks.py

@@ -4,7 +4,10 @@
 from __future__ import annotations
 
 import configparser
+import contextlib
 import functools
+import hashlib
+import importlib
 import json
 import logging
 import logging.handlers
@@ -15,6 +18,7 @@
 import shutil
 import subprocess
 import time
+from collections.abc import Callable
 from concurrent.futures import (
     ThreadPoolExecutor,
     as_completed,
@@ -23,6 +27,7 @@
     TimeoutError as FutureTimeoutError,
 )
 from pathlib import Path
+from types import ModuleType
 from typing import Literal, cast, overload
 from urllib import error as urllib_error
 from urllib import request as urllib_request
@@ -42,6 +47,15 @@
     spinner,
 )
 
+# `fcntl` (POSIX file locking) doesn't exist on Windows. Load it dynamically
+# so we degrade to no cross-process locking there instead of failing to
+# import this module at all; see `_refresh_lock`.
+fcntl: ModuleType | None
+try:
+    fcntl = importlib.import_module("fcntl")
+except ImportError:
+    fcntl = None
+
 UNIX_DATABRICKS_INSTALL_URL = (
     "https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh"
 )
@@ -51,6 +65,13 @@
 AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta"
 MIN_DATABRICKS_CLI_VERSION = (0, 298, 0)
 TOKEN_REFRESH_INTERVAL_SECONDS = 1800
+# A `--force-refresh` within this window of a peer process's successful
+# force-refresh (same workspace+profile) is considered redundant: the OAuth
+# refresh token is rotating, so redeeming it twice in quick succession from
+# concurrent `ucode opencode` sessions can revoke the token family (#190).
+TOKEN_REFRESH_COALESCE_SECONDS = 60
+_LOCK_ACQUIRE_TIMEOUT_SECONDS = 15
+_LOCK_POLL_INTERVAL_SECONDS = 0.2
 
 
 def _debug_enabled() -> bool:
@@ -776,6 +797,116 @@ def ensure_databricks_auth(workspace: str, profile: str | None = None) -> None:
     run_databricks_login(workspace, profile)
 
 
+def _refresh_lock_paths(workspace: str, profile: str | None) -> tuple[Path, Path]:
+    """Per-workspace+profile lock file and sentinel path, both under APP_DIR.
+
+    Hashing the workspace+profile pair keeps concurrent refreshes for
+    *different* workspaces from serializing against each other, while still
+    giving every process refreshing the *same* workspace+profile a shared,
+    stable path to flock."""
+    digest = hashlib.sha256(f"{workspace}|{profile or ''}".encode()).hexdigest()[:16]
+    lock_path = APP_DIR / f".token-refresh-{digest}.lock"
+    sentinel_path = Path(f"{lock_path}.ts")
+    return lock_path, sentinel_path
+
+
+def _sentinel_is_fresh(sentinel_path: Path) -> bool:
+    """Whether `sentinel_path`'s mtime is within TOKEN_REFRESH_COALESCE_SECONDS."""
+    try:
+        mtime = sentinel_path.stat().st_mtime
+    except OSError:
+        return False
+    return (time.time() - mtime) < TOKEN_REFRESH_COALESCE_SECONDS
+
+
+def _touch_sentinel(sentinel_path: Path) -> None:
+    try:
+        sentinel_path.touch()
+    except OSError as exc:
+        _debug("refresh_lock", f"failed to touch sentinel {sentinel_path}: {exc}")
+
+
+@contextlib.contextmanager
+def _refresh_lock(lock_path: Path):
+    """Best-effort cross-process exclusive lock guarding a force-refresh.
+
+    Degrades to a no-op (yields ``False`` without locking anything) when
+    ``fcntl`` is unavailable (Windows) or the lock can't be opened or
+    acquired for any ``OSError`` — a missing or stuck lock must never
+    prevent a token refresh from proceeding. Acquisition is a bounded,
+    non-blocking retry loop (flock has no native timeout) capped at
+    ``_LOCK_ACQUIRE_TIMEOUT_SECONDS`` before falling back to proceeding
+    without the lock."""
+    if fcntl is None:
+        yield False
+        return
+
+    try:
+        APP_DIR.mkdir(parents=True, exist_ok=True)
+        handle = lock_path.open("a+")
+    except OSError as exc:
+        _debug("refresh_lock", f"could not open lock file {lock_path}: {exc}")
+        yield False
+        return
+
+    acquired = False
+    try:
+        deadline = time.monotonic() + _LOCK_ACQUIRE_TIMEOUT_SECONDS
+        while time.monotonic() < deadline:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                acquired = True
+                break
+            except OSError:
+                time.sleep(_LOCK_POLL_INTERVAL_SECONDS)
+        _debug("refresh_lock", f"acquired={acquired} path={lock_path}")
+        yield acquired
+    finally:
+        if acquired:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+            except OSError:
+                pass
+        handle.close()
+
+
+def _perform_force_refresh(
+    workspace: str,
+    profile: str | None,
+    cmd: list[str],
+    fetch: Callable[[], str],
+) -> str:
+    """Serialize `--force-refresh` across processes and coalesce redundant ones.
+
+    Holds a cross-process flock (see `_refresh_lock`) for the duration of an
+    actual force-refresh so two processes can never redeem the rotating
+    refresh token concurrently. If a peer already force-refreshed within
+    `TOKEN_REFRESH_COALESCE_SECONDS`, drops `--force-refresh` from `cmd` (in
+    place, so retries via the same `fetch` closure stay coalesced too) and
+    just fetches the now-current cached token instead."""
+    lock_path, sentinel_path = _refresh_lock_paths(workspace, profile)
+    with _refresh_lock(lock_path) as locked:
+        _debug(
+            "get_databricks_token",
+            f"force-refresh lock acquired={locked} workspace={workspace} profile={profile or '<none>'}",
+        )
+        if _sentinel_is_fresh(sentinel_path):
+            _debug(
+                "get_databricks_token",
+                f"coalescing --force-refresh: peer refreshed within "
+                f"{TOKEN_REFRESH_COALESCE_SECONDS}s",
+            )
+            if "--force-refresh" in cmd:
+                cmd.remove("--force-refresh")
+            return fetch()
+
+        _debug("get_databricks_token", "performing --force-refresh")
+        token = fetch()
+        if token:
+            _touch_sentinel(sentinel_path)
+        return token
+
+
 def get_databricks_token(
     workspace: str,
     profile: str | None = None,
@@ -835,7 +966,10 @@ def _fetch() -> str:
             _debug("auth token", f"exception: {type(exc).__name__}: {exc}")
         return ""
 
-    token = _fetch()
+    if force_refresh:
+        token = _perform_force_refresh(workspace, profile, cmd, _fetch)
+    else:
+        token = _fetch()
     if not token:
         # Session may have expired — attempt non-interactive re-auth and retry once.
         _debug("auth token", "empty on first fetch; attempting auth login --no-browser")

tests/test_databricks.py

@@ -5,6 +5,7 @@
 import json
 import os
 import subprocess
+import time
 
 import pytest
 
@@ -1163,6 +1164,116 @@ def test_error_suggests_logout_when_matching_profile_exists(self, tmp_path, monk
         assert f"databricks auth login --host {WS} --profile example-profile" in message
 
 
+class TestForceRefreshLockAndCoalescing:
+    """Axis A of #190: serialize `--force-refresh` across processes and
+    coalesce redundant refreshes within TOKEN_REFRESH_COALESCE_SECONDS."""
+
+    def _fake_databricks(self, tmp_path, script: str) -> dict:
+        fake = tmp_path / "databricks"
+        fake.write_text(f"#!/bin/sh\n{script}\n")
+        fake.chmod(0o755)
+        return {**os.environ, "PATH": f"{tmp_path}:{os.environ['PATH']}"}
+
+    def _argv_recording_env(self, tmp_path, argv_log) -> dict:
+        return self._fake_databricks(
+            tmp_path,
+            f'printf "%s\\n" "$@" >> {argv_log}\n'
+            'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'',
+        )
+
+    def test_coalesces_force_refresh_when_sentinel_is_fresh(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+        argv_log = tmp_path / "argv"
+        env = self._argv_recording_env(tmp_path, argv_log)
+        monkeypatch.setattr("os.environ", env)
+
+        _, sentinel_path = db_mod._refresh_lock_paths(WS, None)
+        sentinel_path.touch()  # fresh mtime == a peer just force-refreshed
+
+        token = get_databricks_token(WS, force_refresh=True)
+
+        assert token == "good-token"
+        argv = argv_log.read_text().splitlines()
+        assert "--force-refresh" not in argv
+
+    def test_performs_force_refresh_and_touches_sentinel_when_stale(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+        argv_log = tmp_path / "argv"
+        env = self._argv_recording_env(tmp_path, argv_log)
+        monkeypatch.setattr("os.environ", env)
+
+        _, sentinel_path = db_mod._refresh_lock_paths(WS, None)
+        assert not sentinel_path.exists()
+
+        token = get_databricks_token(WS, force_refresh=True)
+
+        assert token == "good-token"
+        argv = argv_log.read_text().splitlines()
+        assert "--force-refresh" in argv
+        assert sentinel_path.exists()
+
+    def test_performs_force_refresh_when_sentinel_is_stale(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+        argv_log = tmp_path / "argv"
+        env = self._argv_recording_env(tmp_path, argv_log)
+        monkeypatch.setattr("os.environ", env)
+
+        _, sentinel_path = db_mod._refresh_lock_paths(WS, None)
+        sentinel_path.touch()
+        old = time.time() - db_mod.TOKEN_REFRESH_COALESCE_SECONDS - 5
+        os.utime(sentinel_path, (old, old))
+
+        token = get_databricks_token(WS, force_refresh=True)
+
+        assert token == "good-token"
+        argv = argv_log.read_text().splitlines()
+        assert "--force-refresh" in argv
+
+    def test_degrades_to_no_lock_when_fcntl_missing(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+        monkeypatch.setattr(db_mod, "fcntl", None)
+        env = self._fake_databricks(
+            tmp_path,
+            'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'',
+        )
+        monkeypatch.setattr("os.environ", env)
+
+        token = get_databricks_token(WS, force_refresh=True)
+
+        assert token == "good-token"
+
+    def test_lock_file_path_differs_for_different_workspace_or_profile(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+
+        lock_a, sentinel_a = db_mod._refresh_lock_paths(WS, None)
+        lock_b, sentinel_b = db_mod._refresh_lock_paths("https://other.databricks.com", None)
+        lock_c, sentinel_c = db_mod._refresh_lock_paths(WS, "some-profile")
+
+        assert lock_a != lock_b
+        assert lock_a != lock_c
+        assert lock_b != lock_c
+        assert sentinel_a != sentinel_b
+        assert sentinel_a != sentinel_c
+        assert lock_a.parent == tmp_path
+
+    def test_non_force_refresh_path_ignores_lock_and_sentinel(self, tmp_path, monkeypatch):
+        # force_refresh=False must stay cheap/unlocked: no lock file should
+        # even be created.
+        monkeypatch.setattr(db_mod, "APP_DIR", tmp_path)
+        env = self._fake_databricks(
+            tmp_path,
+            'echo \'{"access_token": "good-token", "token_type": "Bearer"}\'',
+        )
+        monkeypatch.setattr("os.environ", env)
+
+        token = get_databricks_token(WS)
+
+        assert token == "good-token"
+        lock_path, sentinel_path = db_mod._refresh_lock_paths(WS, None)
+        assert not lock_path.exists()
+        assert not sentinel_path.exists()
+
+
 class TestListDatabricksConnections:
     def test_lists_paginated_connections_with_workspace_env(self, monkeypatch):
         calls: list[dict] = []