Avoid blocking OAuth flow lock during Codex start
This commit is contained in:
1 parent
a700e8e0d3
commit
3c1b7392ba
2 files changed
+64
-8
No files matched your search
+37
-8
@@ -56,9 +56,27 @@ ANTHROPIC_PUBLIC_LINK_ERROR = "Claude Code credential linking failed. Check serv
|
||||
|
||||
_OAUTH_FLOWS: dict[str, dict[str, Any]] = {}
|
||||
_OAUTH_FLOWS_LOCK = threading.Lock()
|
||||
_OAUTH_START_LOCKS: dict[tuple[str, str], threading.Lock] = {}
|
||||
_OAUTH_START_LOCKS_LOCK = threading.Lock()
|
||||
_ANTHROPIC_ENV_KEYS = ("ANTHROPIC_TOKEN", "ANTHROPIC_API_KEY")
|
||||
|
||||
|
||||
def _oauth_start_key(provider: str, hermes_home: Path) -> tuple[str, str]:
|
||||
"""Return the canonical single-flight key for onboarding OAuth starts."""
|
||||
return (_normalize_onboarding_oauth_provider(provider), str(Path(hermes_home)))
|
||||
|
||||
|
||||
def _oauth_start_lock(provider: str, hermes_home: Path) -> threading.Lock:
|
||||
"""Return the per provider/profile lock that serializes OAuth flow creation."""
|
||||
key = _oauth_start_key(provider, hermes_home)
|
||||
with _OAUTH_START_LOCKS_LOCK:
|
||||
lock = _OAUTH_START_LOCKS.get(key)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_OAUTH_START_LOCKS[key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
def _pending_oauth_flow_for_locked(provider: str, hermes_home: Path) -> tuple[str, dict[str, Any]] | None:
|
||||
"""Return an existing live onboarding OAuth flow for the provider/profile.
|
||||
|
||||
@@ -753,11 +771,16 @@ def start_onboarding_oauth_flow(body: dict[str, Any] | None) -> dict[str, Any]:
|
||||
|
||||
# Codex flow
|
||||
hermes_home = _get_active_hermes_home()
|
||||
with _OAUTH_FLOWS_LOCK:
|
||||
existing = _pending_oauth_flow_for_locked("openai-codex", hermes_home)
|
||||
if existing is not None:
|
||||
flow_id, flow = existing
|
||||
return _public_start_payload(flow_id, flow)
|
||||
# Serialize check -> device-code request -> flow insertion -> worker spawn
|
||||
# for this provider/profile. The global flow lock is still held only for
|
||||
# short in-memory checks/inserts, so unrelated polling/status/cleanup paths
|
||||
# are not blocked by the slow provider request.
|
||||
with _oauth_start_lock("openai-codex", hermes_home):
|
||||
with _OAUTH_FLOWS_LOCK:
|
||||
existing = _pending_oauth_flow_for_locked("openai-codex", hermes_home)
|
||||
if existing is not None:
|
||||
flow_id, flow = existing
|
||||
return _public_start_payload(flow_id, flow)
|
||||
|
||||
try:
|
||||
device = _request_codex_user_code()
|
||||
@@ -784,9 +807,15 @@ def start_onboarding_oauth_flow(body: dict[str, Any] | None) -> dict[str, Any]:
|
||||
"created_at": time.time(),
|
||||
"updated_at": time.time(),
|
||||
}
|
||||
_OAUTH_FLOWS[flow_id] = flow
|
||||
_spawn_codex_oauth_worker(flow_id)
|
||||
return _public_start_payload(flow_id, flow)
|
||||
with _OAUTH_FLOWS_LOCK:
|
||||
existing = _pending_oauth_flow_for_locked("openai-codex", hermes_home)
|
||||
if existing is not None:
|
||||
flow_id, flow = existing
|
||||
return _public_start_payload(flow_id, flow)
|
||||
_OAUTH_FLOWS[flow_id] = flow
|
||||
|
||||
_spawn_codex_oauth_worker(flow_id)
|
||||
return _public_start_payload(flow_id, flow)
|
||||
|
||||
|
||||
def poll_onboarding_oauth_flow(flow_id: str) -> dict[str, Any]:
|
||||
|
||||
@@ -9,11 +9,15 @@ import api.oauth as oauth
|
||||
def setup_function():
|
||||
with oauth._OAUTH_FLOWS_LOCK:
|
||||
oauth._OAUTH_FLOWS.clear()
|
||||
with oauth._OAUTH_START_LOCKS_LOCK:
|
||||
oauth._OAUTH_START_LOCKS.clear()
|
||||
|
||||
|
||||
def teardown_function():
|
||||
with oauth._OAUTH_FLOWS_LOCK:
|
||||
oauth._OAUTH_FLOWS.clear()
|
||||
with oauth._OAUTH_START_LOCKS_LOCK:
|
||||
oauth._OAUTH_START_LOCKS.clear()
|
||||
|
||||
|
||||
def _run_two_concurrent(fn):
|
||||
@@ -100,6 +104,29 @@ def test_codex_onboarding_oauth_start_reuses_pending_flow_before_requesting_code
|
||||
assert len(pending) == 1
|
||||
|
||||
|
||||
def test_codex_onboarding_oauth_start_does_not_hold_lock_while_requesting_code(monkeypatch, tmp_path):
|
||||
"""The slow Codex device-code request must not block unrelated flow operations."""
|
||||
lock_available_during_request = []
|
||||
hermes_home = tmp_path / "hermes-home"
|
||||
|
||||
monkeypatch.setattr(oauth, "_get_active_hermes_home", lambda: Path(hermes_home))
|
||||
monkeypatch.setattr(oauth, "_spawn_codex_oauth_worker", lambda flow_id: None)
|
||||
|
||||
def fake_request_code():
|
||||
acquired = oauth._OAUTH_FLOWS_LOCK.acquire(blocking=False)
|
||||
lock_available_during_request.append(acquired)
|
||||
if acquired:
|
||||
oauth._OAUTH_FLOWS_LOCK.release()
|
||||
return {"user_code": "ABCD-EFGH", "device_auth_id": "device-1", "interval": 5, "expires_in": 900}
|
||||
|
||||
monkeypatch.setattr(oauth, "_request_codex_user_code", fake_request_code)
|
||||
|
||||
result = oauth.start_onboarding_oauth_flow({"provider": "openai-codex"})
|
||||
|
||||
assert result["status"] == "pending"
|
||||
assert lock_available_during_request == [True]
|
||||
|
||||
|
||||
def test_codex_onboarding_oauth_start_single_flight_is_thread_safe(monkeypatch, tmp_path):
|
||||
"""Concurrent Codex starts must request one device code and spawn one worker."""
|
||||
spawned = []
|
||||
|
||||
Reference in new issue
Block a user