Merge pull request #73914 from JoaoMarcos44/fix/codex-oauth-cancel-race-ia01

fix(web_server): stop Codex OAuth worker from finishing after cancel
This commit is contained in:
Teknium 2026-07-31 22:36:44 -07:00 committed by GitHub
commit 75aeba09e0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 293 additions and 7 deletions

View File

@ -10677,13 +10677,23 @@ def _codex_full_login_worker(session_id: str) -> None:
sess["interval"] = poll_interval
sess["expires_in"] = 15 * 60 # OpenAI's effective limit
sess["expires_at"] = time.time() + sess["expires_in"]
# Captured now (not re-derived after cancel pops the session) so a
# cancelled session can never fall back to the caller's current
# profile scope at save time.
session_profile = sess.get("profile")
# Step 2: poll until authorized
deadline = time.monotonic() + sess["expires_in"]
code_resp = None
with httpx.Client(timeout=httpx.Timeout(15.0)) as client:
while time.monotonic() < deadline:
if sess.get("cancelled"):
_log.info("oauth/device: openai-codex login cancelled (session=%s)", session_id)
return
time.sleep(poll_interval)
if sess.get("cancelled"):
_log.info("oauth/device: openai-codex login cancelled (session=%s)", session_id)
return
poll = client.post(
f"{issuer}/api/accounts/deviceauth/token",
json={"device_auth_id": device_auth_id, "user_code": user_code},
@ -10702,6 +10712,10 @@ def _codex_full_login_worker(session_id: str) -> None:
sess["error_message"] = "Device code expired before approval"
return
if sess.get("cancelled"):
_log.info("oauth/device: openai-codex login cancelled before token exchange (session=%s)", session_id)
return
# Step 3: exchange authorization_code for tokens
authorization_code = code_resp.get("authorization_code", "")
code_verifier = code_resp.get("code_verifier", "")
@ -10729,12 +10743,23 @@ def _codex_full_login_worker(session_id: str) -> None:
from hermes_cli.auth import _save_codex_tokens
with _profile_scope(_oauth_session_profile(session_id)):
_save_codex_tokens({
"access_token": access_token,
"refresh_token": refresh_token,
})
# The cancellation check and the save must be one atomic critical
# section under the same lock cancel_oauth_session() uses. Checking
# "cancelled" and then saving as two separate steps left a window
# where DELETE could flip the flag between them and the worker would
# still persist tokens after the user believed the login was
# aborted. Holding the lock across both closes that window: DELETE
# either lands before this section (worker observes cancelled and
# returns) or blocks until this section (and the save) is done.
with _oauth_sessions_lock:
if sess.get("cancelled"):
_log.info("oauth/device: openai-codex login cancelled before token save (session=%s)", session_id)
return
with _profile_scope(session_profile):
_save_codex_tokens({
"access_token": access_token,
"refresh_token": refresh_token,
})
sess["status"] = "approved"
_log.info("oauth/device: openai-codex login completed (session=%s)", session_id)
except Exception as e:
@ -10832,10 +10857,20 @@ async def cancel_oauth_session(
request: Request,
profile: Optional[str] = None,
):
"""Cancel a pending OAuth session. Token-protected."""
"""Cancel a pending OAuth session. Token-protected.
Marks the session dict ``cancelled`` before popping it so any
background worker still holding a reference to that same dict (e.g.
the Codex device-code poller) observes the cancellation and stops
polling/exchanging/saving instead of completing the login after the
user believed it was aborted.
"""
_require_token(request)
with _oauth_sessions_lock:
sess = _oauth_sessions.pop(session_id, None)
sess = _oauth_sessions.get(session_id)
if sess is not None:
sess["cancelled"] = True
_oauth_sessions.pop(session_id, None)
if sess is None:
return {"ok": False, "message": "session not found"}
return {"ok": True, "session_id": session_id}

View File

@ -223,8 +223,259 @@ def test_codex_dashboard_start_rewords_device_authorization_error(monkeypatch):
ws._oauth_sessions.pop(sid, None)
def test_codex_dashboard_worker_stops_polling_after_cancel(tmp_path, monkeypatch):
"""A real DELETE mid-poll must stop the worker before it exchanges/saves tokens.
Regression for IA-01: cancelling only popped the session dict; the
background worker kept polling/exchanging/saving regardless, and once
the session was gone `_oauth_session_profile()` fell back to the
caller's current profile scope instead of the one the login started
in. The fix marks the dict `cancelled` before popping, and the worker
checks that flag before every remaining step.
Exercises the actual `DELETE /api/providers/oauth/sessions/{id}`
endpoint (rather than mutating the session dict directly) so the
endpoint/worker race and the real removal from `_oauth_sessions` are
both under test.
"""
from hermes_cli import auth as auth_mod
from hermes_cli import web_server as ws
class _Resp:
def __init__(self, status_code, payload):
self.status_code = status_code
self._payload = payload
def json(self):
return self._payload
class _Client:
def __init__(self, *args, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *args):
return False
def post(self, url, **kwargs):
if url.endswith("/deviceauth/usercode"):
return _Resp(200, {
"device_auth_id": "device-auth-id",
"interval": 3,
"user_code": "CODEX-1234",
})
raise AssertionError(
f"worker must stop before calling {url} once cancelled"
)
saved = []
monkeypatch.setenv("HERMES_HOME", str(tmp_path))
monkeypatch.setattr(httpx, "Client", _Client)
monkeypatch.setattr(auth_mod, "_save_codex_tokens", lambda tokens: saved.append(tokens))
sid, _ = ws._new_oauth_session("openai-codex", "device_code", profile="coder")
def fake_sleep(_interval):
# Simulate a real concurrent DELETE /api/providers/oauth/sessions/{sid}
# firing while the worker is asleep between polls.
resp = client.delete(f"/api/providers/oauth/sessions/{sid}", headers=HEADERS)
assert resp.status_code == 200, resp.text
monkeypatch.setattr(ws.time, "sleep", fake_sleep)
try:
ws._codex_full_login_worker(sid)
assert saved == []
assert sid not in ws._oauth_sessions
finally:
ws._oauth_sessions.pop(sid, None)
def test_codex_worker_final_save_is_atomic_with_cancel_delete(tmp_path, monkeypatch):
"""The final cancellation check and the token save must be one atomic
section under `_oauth_sessions_lock`.
Regression: checking `cancelled` and calling `_save_codex_tokens()` used
to be two separate steps with no lock held across them, so a DELETE
landing in that gap flipped the flag too late for the worker to see it
and the tokens were saved anyway. This drives a real DELETE from another
thread exactly while the worker holds the lock for its check+save, and
asserts DELETE stays blocked for the whole critical section instead of
slipping in between the check and the save.
"""
import threading
from hermes_cli import auth as auth_mod
from hermes_cli import web_server as ws
class _Resp:
def __init__(self, status_code, payload):
self.status_code = status_code
self._payload = payload
def json(self):
return self._payload
class _Client:
def __init__(self, *args, **kwargs):
pass
def __enter__(self):
return self
def __exit__(self, *args):
return False
def post(self, url, **kwargs):
if url.endswith("/deviceauth/usercode"):
return _Resp(200, {
"device_auth_id": "device-auth-id",
"interval": 0,
"user_code": "CODEX-1234",
})
return _Resp(200, {
"authorization_code": "auth-code",
"code_verifier": "verifier",
})
class _TokenClient(_Client):
def post(self, url, **kwargs):
return _Resp(200, {"access_token": "at", "refresh_token": "rt"})
clients = iter([_Client, _Client, _TokenClient])
_make_profile_home(tmp_path, monkeypatch, profile="coder")
monkeypatch.setattr(httpx, "Client", lambda *a, **k: next(clients)(*a, **k))
saved = []
delete_threads = []
delete_started = threading.Event()
delete_finished = threading.Event()
def fake_save(tokens):
# We are inside the worker's critical section right now (holding
# _oauth_sessions_lock). Fire a real DELETE from another thread and
# prove it cannot complete until this section releases the lock.
# Do NOT join the DELETE thread here: it is blocked on the very
# lock this section holds, so joining here would deadlock.
delete_thread = threading.Thread(target=_fire_delete, daemon=True)
delete_threads.append(delete_thread)
delete_thread.start()
delete_started.wait(timeout=2)
still_blocked = not delete_finished.wait(timeout=0.2)
saved.append((tokens, still_blocked))
def _fire_delete():
delete_started.set()
client.delete(f"/api/providers/oauth/sessions/{sid}", headers=HEADERS)
delete_finished.set()
monkeypatch.setattr(auth_mod, "_save_codex_tokens", fake_save)
monkeypatch.setattr(ws.time, "sleep", lambda *_a, **_k: None)
sid, _ = ws._new_oauth_session("openai-codex", "device_code", profile="coder")
ws._codex_full_login_worker(sid)
# The lock is released now (worker returned), so the DELETE thread can
# finally complete.
delete_threads[0].join(timeout=2)
assert len(saved) == 1
tokens, delete_was_still_blocked_during_save = saved[0]
assert tokens == {"access_token": "at", "refresh_token": "rt"}
assert delete_was_still_blocked_during_save, (
"DELETE must block until the worker's check+save critical section "
"finishes, not slip in between the check and the save"
)
# DELETE arrived after the point of no return (save already committed),
# so this is the legitimate too-late-to-cancel outcome: token saved,
# session subsequently removed by the now-unblocked DELETE.
assert sid not in ws._oauth_sessions
def test_cancel_oauth_session_marks_dict_cancelled_before_popping():
"""The DELETE endpoint must flag the session dict before removing it.
A background worker holds its own reference to the same dict object;
it can only observe cancellation if the flag is set on that shared
object prior to (or instead of) removal from the global session map.
"""
from hermes_cli import web_server as ws
session_id = "cancel-flag-test"
ws._oauth_sessions[session_id] = {
"session_id": session_id,
"provider": "openai-codex",
"flow": "device_code",
"profile": "coder",
"created_at": time.time(),
"status": "pending",
"error_message": None,
}
worker_ref = ws._oauth_sessions[session_id]
resp = client.delete(
f"/api/providers/oauth/sessions/{session_id}",
headers=HEADERS,
)
assert resp.status_code == 200, resp.text
assert resp.json() == {"ok": True, "session_id": session_id}
assert session_id not in ws._oauth_sessions
assert worker_ref["cancelled"] is True
def test_nous_dashboard_poller_preserves_effective_scope_when_token_omits_scope(monkeypatch):
from hermes_cli import auth as auth_mod
from hermes_cli import web_server as ws
session_id = "nous-effective-scope-test"
ws._oauth_sessions[session_id] = {
"session_id": session_id,
"provider": "nous",
"flow": "device_code",
"created_at": time.time(),
"status": "pending",
"error_message": None,
"portal_base_url": "https://portal.nousresearch.com",
"client_id": "hermes-cli",
"device_code": "device-code",
"interval": 5,
"expires_at": time.time() + 600,
"scope": auth_mod.DEFAULT_NOUS_SCOPE,
}
captured_state = {}
def fake_refresh_nous_oauth_from_state(state, **kwargs):
captured_state.update(state)
return {**state, "agent_key": "jwt-agent-key"}
monkeypatch.setattr(
auth_mod,
"_poll_for_token",
lambda **kwargs: {
"access_token": "access-token",
"refresh_token": "refresh-token",
"expires_in": 3600,
"token_type": "Bearer",
},
)
monkeypatch.setattr(
auth_mod,
"refresh_nous_oauth_from_state",
fake_refresh_nous_oauth_from_state,
)
monkeypatch.setattr(auth_mod, "persist_nous_credentials", lambda state: None)
try:
ws._nous_poller(session_id)
assert captured_state["scope"] == auth_mod.DEFAULT_NOUS_SCOPE
assert ws._oauth_sessions[session_id]["status"] == "approved"
finally:
ws._oauth_sessions.pop(session_id, None)