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:
commit
75aeba09e0
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue