diff --git a/api/oauth.py b/api/oauth.py index 9493de9f..9c645d49 100644 --- a/api/oauth.py +++ b/api/oauth.py @@ -59,6 +59,41 @@ _OAUTH_FLOWS_LOCK = threading.Lock() _ANTHROPIC_ENV_KEYS = ("ANTHROPIC_TOKEN", "ANTHROPIC_API_KEY") +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. + + Onboarding OAuth starts are unauthenticated while first-run auth is disabled. + Keep them single-flight per provider/profile so repeated start requests reuse + the existing worker instead of accumulating pending flows and daemon threads. + + The caller must hold _OAUTH_FLOWS_LOCK so the check can be paired atomically + with flow insertion on start paths. + """ + canonical_provider = _normalize_onboarding_oauth_provider(provider) + canonical_home = str(Path(hermes_home)) + now = time.time() + for flow_id, flow in _OAUTH_FLOWS.items(): + if flow.get("provider") != canonical_provider: + continue + if str(flow.get("hermes_home") or "") != canonical_home: + continue + if flow.get("status") != "pending": + continue + if float(flow.get("expires_at") or 0) <= now: + flow["status"] = "expired" + flow["updated_at"] = now + _drop_sensitive_flow_fields(flow) + continue + return flow_id, dict(flow) + return None + + +def _pending_oauth_flow_for(provider: str, hermes_home: Path) -> tuple[str, dict[str, Any]] | None: + """Return an existing live onboarding OAuth flow for the provider/profile, if any.""" + with _OAUTH_FLOWS_LOCK: + return _pending_oauth_flow_for_locked(provider, hermes_home) + + def _clear_process_anthropic_env_values() -> None: """Clear Anthropic process env fallbacks under the streaming env lock.""" from api.streaming import _ENV_LOCK @@ -687,6 +722,10 @@ def _start_anthropic_flow(hermes_home: Path) -> dict[str, Any]: "updated_at": time.time(), } with _OAUTH_FLOWS_LOCK: + existing = _pending_oauth_flow_for_locked("anthropic", hermes_home) + if existing is not None: + flow_id, flow = existing + return _public_start_payload(flow_id, flow) _OAUTH_FLOWS[flow_id] = flow _spawn_anthropic_credential_worker(flow_id) return _public_start_payload(flow_id, flow) @@ -714,32 +753,37 @@ def start_onboarding_oauth_flow(body: dict[str, Any] | None) -> dict[str, Any]: # Codex flow hermes_home = _get_active_hermes_home() - try: - device = _request_codex_user_code() - except Exception as exc: - raise RuntimeError(f"Failed to start Codex OAuth: {exc}") from exc - - user_code = str(device.get("user_code") or "").strip() - device_auth_id = str(device.get("device_auth_id") or "").strip() - if not user_code or not device_auth_id: - raise RuntimeError("Device code response missing required fields") - - interval = max(3, int(device.get("interval") or 5)) - expires_in = int(device.get("expires_in") or CODEX_FLOW_MAX_WAIT_SECONDS) - expires_at = time.time() + min(max(expires_in, 60), CODEX_FLOW_MAX_WAIT_SECONDS) - flow_id = uuid.uuid4().hex - flow = { - "provider": "openai-codex", - "status": "pending", - "device_auth_id": device_auth_id, - "user_code": user_code, - "expires_at": expires_at, - "poll_interval_seconds": interval, - "hermes_home": str(hermes_home), - "created_at": time.time(), - "updated_at": time.time(), - } 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() + except Exception as exc: + raise RuntimeError(f"Failed to start Codex OAuth: {exc}") from exc + + user_code = str(device.get("user_code") or "").strip() + device_auth_id = str(device.get("device_auth_id") or "").strip() + if not user_code or not device_auth_id: + raise RuntimeError("Device code response missing required fields") + + interval = max(3, int(device.get("interval") or 5)) + expires_in = int(device.get("expires_in") or CODEX_FLOW_MAX_WAIT_SECONDS) + expires_at = time.time() + min(max(expires_in, 60), CODEX_FLOW_MAX_WAIT_SECONDS) + flow_id = uuid.uuid4().hex + flow = { + "provider": "openai-codex", + "status": "pending", + "device_auth_id": device_auth_id, + "user_code": user_code, + "expires_at": expires_at, + "poll_interval_seconds": interval, + "hermes_home": str(hermes_home), + "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) diff --git a/tests/test_onboarding_oauth_single_flight.py b/tests/test_onboarding_oauth_single_flight.py new file mode 100644 index 00000000..9e6684b6 --- /dev/null +++ b/tests/test_onboarding_oauth_single_flight.py @@ -0,0 +1,129 @@ +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +import threading +import time + +import api.oauth as oauth + + +def setup_function(): + with oauth._OAUTH_FLOWS_LOCK: + oauth._OAUTH_FLOWS.clear() + + +def teardown_function(): + with oauth._OAUTH_FLOWS_LOCK: + oauth._OAUTH_FLOWS.clear() + + +def _run_two_concurrent(fn): + ready = threading.Barrier(2) + + def wrapped(): + ready.wait(timeout=5) + return fn() + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(wrapped), executor.submit(wrapped)] + return [future.result(timeout=5) for future in futures] + + +def test_anthropic_onboarding_oauth_start_reuses_pending_flow(monkeypatch, tmp_path): + """Repeated unauthenticated starts must not spawn unbounded pending workers.""" + spawned = [] + hermes_home = tmp_path / "hermes-home" + + monkeypatch.setattr(oauth, "_read_claude_code_credentials", lambda: None) + monkeypatch.setattr(oauth, "_spawn_anthropic_credential_worker", lambda flow_id: spawned.append(flow_id)) + + first = oauth._start_anthropic_flow(hermes_home) + second = oauth._start_anthropic_flow(hermes_home) + + assert first["status"] == "pending" + assert second["status"] == "pending" + assert second["flow_id"] == first["flow_id"] + assert spawned == [first["flow_id"]] + with oauth._OAUTH_FLOWS_LOCK: + pending = [flow for flow in oauth._OAUTH_FLOWS.values() if flow.get("status") == "pending"] + assert len(pending) == 1 + + +def test_anthropic_onboarding_oauth_start_single_flight_is_thread_safe(monkeypatch, tmp_path): + """Concurrent unauthenticated Anthropic starts must reserve one pending worker atomically.""" + spawned = [] + hermes_home = tmp_path / "hermes-home" + credential_reads = threading.Barrier(2) + + def no_credentials(): + credential_reads.wait(timeout=5) + return None + + monkeypatch.setattr(oauth, "_read_claude_code_credentials", no_credentials) + monkeypatch.setattr(oauth, "_spawn_anthropic_credential_worker", lambda flow_id: spawned.append(flow_id)) + + first, second = _run_two_concurrent(lambda: oauth._start_anthropic_flow(hermes_home)) + + assert first["status"] == "pending" + assert second["status"] == "pending" + assert second["flow_id"] == first["flow_id"] + assert spawned == [first["flow_id"]] + with oauth._OAUTH_FLOWS_LOCK: + pending = [flow for flow in oauth._OAUTH_FLOWS.values() if flow.get("status") == "pending"] + assert len(pending) == 1 + + +def test_codex_onboarding_oauth_start_reuses_pending_flow_before_requesting_code(monkeypatch, tmp_path): + """A second Codex start should reuse the live flow before creating a new device code or worker.""" + spawned = [] + requested = [] + hermes_home = tmp_path / "hermes-home" + + monkeypatch.setattr(oauth, "_get_active_hermes_home", lambda: Path(hermes_home)) + + def fake_request_code(): + requested.append(True) + 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) + monkeypatch.setattr(oauth, "_spawn_codex_oauth_worker", lambda flow_id: spawned.append(flow_id)) + + first = oauth.start_onboarding_oauth_flow({"provider": "openai-codex"}) + second = oauth.start_onboarding_oauth_flow({"provider": "openai-codex"}) + + assert first["status"] == "pending" + assert second["flow_id"] == first["flow_id"] + assert second["user_code"] == "ABCD-EFGH" + assert requested == [True] + assert spawned == [first["flow_id"]] + with oauth._OAUTH_FLOWS_LOCK: + pending = [flow for flow in oauth._OAUTH_FLOWS.values() if flow.get("status") == "pending"] + assert len(pending) == 1 + + +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 = [] + requested = [] + hermes_home = tmp_path / "hermes-home" + + monkeypatch.setattr(oauth, "_get_active_hermes_home", lambda: Path(hermes_home)) + + def fake_request_code(): + requested.append(True) + time.sleep(0.05) + 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) + monkeypatch.setattr(oauth, "_spawn_codex_oauth_worker", lambda flow_id: spawned.append(flow_id)) + + first, second = _run_two_concurrent(lambda: oauth.start_onboarding_oauth_flow({"provider": "openai-codex"})) + + assert first["status"] == "pending" + assert second["status"] == "pending" + assert second["flow_id"] == first["flow_id"] + assert second["user_code"] == "ABCD-EFGH" + assert requested == [True] + assert spawned == [first["flow_id"]] + with oauth._OAUTH_FLOWS_LOCK: + pending = [flow for flow in oauth._OAUTH_FLOWS.values() if flow.get("status") == "pending"] + assert len(pending) == 1