release lock on oauth

scripts/repro_custom_oauth_lock_hang.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+"""Reproduce the old OAuth lock hang or verify the owner-lease fix, offline."""
+
+import argparse
+import fcntl
+import multiprocessing
+import time
+from pathlib import Path
+
+from ucode.custom_oauth import CustomOAuthFlowTimeout, _custom_oauth_lock
+
+DEFAULT_LOCK = Path("/tmp/ug-oauth-lock-repro/ug-oauth-8020.lock")
+
+
+def _raw_owner(lock_path, ready, seconds):
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    with lock_path.open("a+b") as lock_file:
+        fcntl.flock(lock_file, fcntl.LOCK_EX)
+        ready.set()
+        time.sleep(seconds + 60)
+
+
+def _leased_owner(lock_path, ready, seconds):
+    try:
+        with _custom_oauth_lock(
+            lock_path.parent,
+            "http://localhost:8020/callback",
+            lease_seconds=seconds,
+        ):
+            ready.set()
+            time.sleep(seconds + 60)
+    except CustomOAuthFlowTimeout:
+        pass
+
+
+def _waiter(lock_path, acquired):
+    with _custom_oauth_lock(
+        lock_path.parent,
+        "http://localhost:8020/callback",
+        lease_seconds=5,
+    ):
+        acquired.set()
+
+
+def _stop(processes):
+    for process in processes:
+        if process.is_alive():
+            process.terminate()
+        process.join(5)
+
+
+def run(mode, seconds):
+    lock_path = DEFAULT_LOCK
+    context = multiprocessing.get_context("spawn")
+    ready = context.Event()
+    acquired = context.Event()
+    owner_target = _raw_owner if mode == "broken" else _leased_owner
+    owner = context.Process(target=owner_target, args=(lock_path, ready, seconds))
+    waiter = context.Process(target=_waiter, args=(lock_path, acquired))
+    processes = (owner, waiter)
+
+    try:
+        owner.start()
+        if not ready.wait(5):
+            raise SystemExit("Owner failed to acquire the lock")
+        waiter.start()
+        waiter.join(seconds + 2)
+
+        if mode == "broken":
+            if acquired.is_set():
+                raise SystemExit("Unexpectedly acquired the lock")
+            print("BROKEN REPRODUCED: the waiter is still blocked behind the hung owner.")
+        else:
+            if waiter.is_alive() or not acquired.is_set():
+                raise SystemExit("FIX FAILED: the waiter did not acquire the released lock")
+            print("FIX VERIFIED: the owner lease expired and the waiter acquired the lock.")
+    finally:
+        _stop(processes)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("mode", choices=("broken", "fixed"))
+    parser.add_argument("--seconds", type=float, required=True)
+    args = parser.parse_args()
+    if args.seconds <= 0:
+        parser.error("--seconds must be positive")
+    run(args.mode, args.seconds)
+
+
+if __name__ == "__main__":
+    main()

src/ucode/custom_oauth.py

@@ -4,9 +4,11 @@
 
 import platform
 import shlex
+import signal
 import subprocess
 from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
+from os import getpid
 from pathlib import Path
 from typing import TypedDict
 from urllib.parse import urlparse
@@ -18,8 +20,14 @@
 from ucode.ui import err_console, normalize_workspace_url, print_warning_err
 
 DEFAULT_REDIRECT_URL = f"http://{LOCALHOST}:8020"
-# Custom OAuth may need a human to finish browser consent, not just a token fetch.
-CUSTOM_OAUTH_TIMEOUT_MS = 180_000
+# Custom OAuth may need a human to finish browser consent. Codex's outer process timeout covers one
+# full wait followed by a fresh owner's full browser flow.
+CUSTOM_OAUTH_FLOW_TIMEOUT_SECONDS = 180.0
+CUSTOM_OAUTH_TIMEOUT_MS = 365_000
+
+
+class CustomOAuthFlowTimeout(RuntimeError):
+    """The custom-OAuth lock owner exceeded its authentication lease."""
 
 
 class CustomOAuthConfig(TypedDict):
@@ -94,7 +102,12 @@ def build_custom_auth_shell_command(workspace: str, config: CustomOAuthConfig) -
 
 
 @contextmanager
-def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]:
+def _custom_oauth_lock(
+    cache_dir: Path,
+    redirect_url: str,
+    *,
+    lease_seconds: float,
+) -> Iterator[None]:
     """Serialize helpers sharing a callback port with a POSIX file lock.
 
     Keep the lock file in place: unlinking it could let waiters lock different
@@ -104,14 +117,39 @@ def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]:
 
     cache_dir.mkdir(parents=True, exist_ok=True)
     port = urlparse(redirect_url).port
-    with (cache_dir / f"ug-oauth-{port}.lock").open("a+b") as lock_file:
+    lock_path = cache_dir / f"ug-oauth-{port}.lock"
+    with lock_path.open("a+b") as lock_file:
         fcntl.flock(lock_file, fcntl.LOCK_EX)
+        lock_file.seek(0)
+        lock_file.truncate()
+        lock_file.write(f"{getpid()}\n".encode())
+        lock_file.flush()
         try:
-            yield
+            with _custom_oauth_flow_deadline(lease_seconds):
+                yield
         finally:
             fcntl.flock(lock_file, fcntl.LOCK_UN)
 
 
+@contextmanager
+def _custom_oauth_flow_deadline(timeout_seconds: float) -> Iterator[None]:
+    """Interrupt a lock owner's OAuth work so its advisory lock cannot live forever."""
+
+    def expire(_signum: int, _frame: object) -> None:
+        raise CustomOAuthFlowTimeout(
+            f"Custom OAuth did not finish within {timeout_seconds:g}s; its lock was released. Retry "
+            "the coding agent to start a new authentication attempt."
+        )
+
+    previous_handler = signal.signal(signal.SIGALRM, expire)
+    signal.setitimer(signal.ITIMER_REAL, timeout_seconds)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, previous_handler)
+
+
 def get_custom_client_token(
     workspace: str,
     client_id: str,
@@ -132,7 +170,11 @@ def get_custom_client_token(
             redirect_url=config["redirect_url"],
             scopes=config["scopes"],
         )
-        with _custom_oauth_lock(Path(cache.filename).parent, config["redirect_url"]):
+        with _custom_oauth_lock(
+            Path(cache.filename).parent,
+            config["redirect_url"],
+            lease_seconds=CUSTOM_OAUTH_FLOW_TIMEOUT_SECONDS,
+        ):
             # Read only after acquiring the lock: another helper may have just
             # completed login or rotated the refresh token while we waited.
             credentials = cache.load()
@@ -168,6 +210,8 @@ def get_custom_client_token(
                 raise ValueError("OAuth returned no access token")
             cache.save(credentials)
             return token
+    except CustomOAuthFlowTimeout:
+        raise
     except Exception as exc:
         raise RuntimeError(
             "Custom-client OAuth failed. Check the workspace, client ID, and registered "

tests/test_agent_codex.py

@@ -145,7 +145,7 @@ def test_auth_timeout_allows_custom_browser_login_only(self, render, custom):
         )
         overlay = render(WS, custom_oauth=config)
         auth = overlay["model_providers"][codex.CODEX_MODEL_PROVIDER_NAME]["auth"]
-        assert auth["timeout_ms"] == (180_000 if custom else 5000)
+        assert auth["timeout_ms"] == (365_000 if custom else 5000)
 
     def test_provider_adds_routing_header(self):
         overlay = codex.render_overlay(WS, provider="main.aarushi.aarushi-openai")

tests/test_custom_oauth.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import time
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from unittest.mock import Mock, patch
@@ -15,7 +16,11 @@
 import ucode.cli as cli_mod
 import ucode.databricks as db_mod
 from ucode.cli import app
-from ucode.custom_oauth import _custom_oauth_lock, get_custom_client_token
+from ucode.custom_oauth import (
+    CustomOAuthFlowTimeout,
+    _custom_oauth_lock,
+    get_custom_client_token,
+)
 
 WS = "https://example.databricks.com"
 TEST_SCOPES = ("offline_access", "catalog.catalogs:read")
@@ -25,11 +30,34 @@
 class TestCustomOAuthLock:
     def test_releases_lock_when_login_fails(self, tmp_path):
         with pytest.raises(ValueError, match="login failed"):
-            with _custom_oauth_lock(tmp_path, "http://localhost:8020/callback"):
+            with _custom_oauth_lock(
+                tmp_path,
+                "http://localhost:8020/callback",
+                lease_seconds=1,
+            ):
                 raise ValueError("login failed")
-        with _custom_oauth_lock(tmp_path, "http://127.0.0.1:8020/other-callback"):
+        with _custom_oauth_lock(
+            tmp_path,
+            "http://127.0.0.1:8020/other-callback",
+            lease_seconds=1,
+        ):
             assert len(list(tmp_path.glob("*.lock"))) == 1
 
+    def test_owner_lease_interrupts_work_and_releases_lock(self, tmp_path):
+        with pytest.raises(CustomOAuthFlowTimeout, match="its lock was released"):
+            with _custom_oauth_lock(
+                tmp_path,
+                "http://localhost:8020/callback",
+                lease_seconds=0.01,
+            ):
+                time.sleep(1)
+        with _custom_oauth_lock(
+            tmp_path,
+            "http://localhost:8020/callback",
+            lease_seconds=1,
+        ):
+            pass
+
 
 class TestCustomClientToken:
     @pytest.fixture(autouse=True)

tests/test_state.py

@@ -307,7 +307,7 @@ def test_custom_oauth_applies_to_claude_and_codex(self):
 
         assert "--client-id custom-client" in result["claude"]["auth_command"]
         assert result["codex"]["auth"]["args"][-1] == "offline_access,model-serving"
-        assert result["codex"]["auth"]["timeout_ms"] == 180_000
+        assert result["codex"]["auth"]["timeout_ms"] == 365_000
 
 
 # ---------------------------------------------------------------------------