diff --git a/scripts/repro_custom_oauth_lock_hang.py b/scripts/repro_custom_oauth_lock_hang.py new file mode 100644 index 00000000..4c26159a --- /dev/null +++ b/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() diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index c6b8ac7e..39eb219a 100644 --- a/src/ucode/custom_oauth.py +++ b/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 " diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index a95a1fe6..21038377 100644 --- a/tests/test_agent_codex.py +++ b/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") diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index ed0d945a..c1d13b65 100644 --- a/tests/test_custom_oauth.py +++ b/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) diff --git a/tests/test_state.py b/tests/test_state.py index 33dd6aba..b02436f6 100644 --- a/tests/test_state.py +++ b/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 # ---------------------------------------------------------------------------