diff --git a/hermes_cli/config.py b/hermes_cli/config.py index ff4629dacad7..56048c3b2bf4 100644 --- a/hermes_cli/config.py +++ b/hermes_cli/config.py @@ -27,6 +27,7 @@ import tempfile import threading import time +from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from typing import Dict, Any, Optional, List, Tuple, Set @@ -3862,6 +3863,32 @@ def _env_line_defines_key(line: str, key: str) -> bool: return stripped.startswith(f"{key}=") +# Cross-process advisory lock serializing .env read-modify-write cycles. +# save_env_value()/remove_env_value() each read the whole file, transform an +# in-memory snapshot, then atomic-replace it — atomic_replace() rules out a +# torn/partial file, but two concurrent callers reading the same snapshot +# each write their own version, silently losing whichever wrote first (lost +# update). This lock closes that window; it does not need to be reentrant +# across the two functions since neither calls the other. +_ENV_FILE_LOCK_HOLDER = threading.local() +_ENV_FILE_LOCK_TIMEOUT_SECONDS = 15.0 + + +@contextmanager +def _env_file_lock(): + from hermes_cli.auth import _file_lock # lazy: auth imports this module + + env_path = get_env_path() + lock_path = env_path.parent / (env_path.name + ".lock") + with _file_lock( + lock_path, + _ENV_FILE_LOCK_HOLDER, + _ENV_FILE_LOCK_TIMEOUT_SECONDS, + "Timed out waiting for .env file lock", + ): + yield + + def save_env_value(key: str, value: str): """Save or update a value in ~/.hermes/.env.""" if is_managed(): @@ -3894,63 +3921,64 @@ def save_env_value(key: str, value: str): read_kw = {"encoding": "utf-8-sig", "errors": "replace"} write_kw = {"encoding": "utf-8"} - lines = [] - if env_path.exists(): - with open(env_path, **read_kw) as f: - lines = f.readlines() - # Normalize safe line formatting without interpreting values as syntax. - lines = _sanitize_env_lines(lines) + with _env_file_lock(): + lines = [] + if env_path.exists(): + with open(env_path, **read_kw) as f: + lines = f.readlines() + # Normalize safe line formatting without interpreting values as syntax. + lines = _sanitize_env_lines(lines) + + serialized_value = _quote_env_value(value) + + # Find and update or append. Match both ``KEY=`` and the bash-compatible + # ``export KEY=`` form — load_env() parses export lines (#6659), so a + # user-added ``export GITHUB_TOKEN=...`` shows as set in every UI. If the + # writer didn't match it, a save would append a SECOND line and a later + # delete of that line would silently resurrect the old exported value + # (#40041: "token detected but cannot be replaced through the UI"). + found = False + for i, line in enumerate(lines): + if _env_line_defines_key(line, key): + lines[i] = f"{key}={serialized_value}\n" + found = True + break - serialized_value = _quote_env_value(value) - - # Find and update or append. Match both ``KEY=`` and the bash-compatible - # ``export KEY=`` form — load_env() parses export lines (#6659), so a - # user-added ``export GITHUB_TOKEN=...`` shows as set in every UI. If the - # writer didn't match it, a save would append a SECOND line and a later - # delete of that line would silently resurrect the old exported value - # (#40041: "token detected but cannot be replaced through the UI"). - found = False - for i, line in enumerate(lines): - if _env_line_defines_key(line, key): - lines[i] = f"{key}={serialized_value}\n" - found = True - break + if not found: + # Ensure there's a newline at the end of the file before appending + if lines and not lines[-1].endswith("\n"): + lines[-1] += "\n" + lines.append(f"{key}={serialized_value}\n") - if not found: - # Ensure there's a newline at the end of the file before appending - if lines and not lines[-1].endswith("\n"): - lines[-1] += "\n" - lines.append(f"{key}={serialized_value}\n") - - fd, tmp_path = tempfile.mkstemp(dir=str(env_path.parent), suffix='.tmp', prefix='.env_') - # Preserve original permissions so Docker volume mounts aren't clobbered. - original_mode = None - if env_path.exists(): - try: - original_mode = stat.S_IMODE(env_path.stat().st_mode) - except OSError: - pass - try: - with os.fdopen(fd, 'w', **write_kw) as f: - f.writelines(lines) - f.flush() - os.fsync(f.fileno()) - atomic_replace(tmp_path, env_path) - # Preserve the original file mode (e.g. 0640 for Docker volume mounts) - # instead of letting _secure_file unconditionally tighten to 0600. - if original_mode is not None: + fd, tmp_path = tempfile.mkstemp(dir=str(env_path.parent), suffix='.tmp', prefix='.env_') + # Preserve original permissions so Docker volume mounts aren't clobbered. + original_mode = None + if env_path.exists(): try: - os.chmod(env_path, original_mode) + original_mode = stat.S_IMODE(env_path.stat().st_mode) except OSError: pass - else: - _secure_file(env_path) - except BaseException: try: - os.unlink(tmp_path) - except OSError: - pass - raise + with os.fdopen(fd, 'w', **write_kw) as f: + f.writelines(lines) + f.flush() + os.fsync(f.fileno()) + atomic_replace(tmp_path, env_path) + # Preserve the original file mode (e.g. 0640 for Docker volume mounts) + # instead of letting _secure_file unconditionally tighten to 0600. + if original_mode is not None: + try: + os.chmod(env_path, original_mode) + except OSError: + pass + else: + _secure_file(env_path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise os.environ[key] = value invalidate_env_cache() @@ -3998,50 +4026,52 @@ def remove_env_value(key: str) -> bool: if not _ENV_VAR_NAME_RE.match(key): raise ValueError(f"Invalid environment variable name: {key!r}") env_path = get_env_path() - if not env_path.exists(): - os.environ.pop(key, None) - return False read_kw = {"encoding": "utf-8-sig", "errors": "replace"} write_kw = {"encoding": "utf-8"} - with open(env_path, **read_kw) as f: - lines = f.readlines() - lines = _sanitize_env_lines(lines) + with _env_file_lock(): + if not env_path.exists(): + os.environ.pop(key, None) + return False - new_lines = [line for line in lines if not _env_line_defines_key(line, key)] - found = len(new_lines) < len(lines) + with open(env_path, **read_kw) as f: + lines = f.readlines() + lines = _sanitize_env_lines(lines) - if found: - fd, tmp_path = tempfile.mkstemp(dir=str(env_path.parent), suffix='.tmp', prefix='.env_') - # Preserve original permissions so Docker volume mounts aren't clobbered. - original_mode = None - try: - original_mode = stat.S_IMODE(env_path.stat().st_mode) - except OSError: - pass - try: - with os.fdopen(fd, 'w', **write_kw) as f: - f.writelines(new_lines) - f.flush() - os.fsync(f.fileno()) - atomic_replace(tmp_path, env_path) - # Preserve the original file mode (e.g. 0640 for Docker volume - # mounts) instead of letting _secure_file unconditionally tighten - # to 0600. Mirrors save_env_value(). - if original_mode is not None: - try: - os.chmod(env_path, original_mode) - except OSError: - pass - else: - _secure_file(env_path) - except BaseException: + new_lines = [line for line in lines if not _env_line_defines_key(line, key)] + found = len(new_lines) < len(lines) + + if found: + fd, tmp_path = tempfile.mkstemp(dir=str(env_path.parent), suffix='.tmp', prefix='.env_') + # Preserve original permissions so Docker volume mounts aren't clobbered. + original_mode = None try: - os.unlink(tmp_path) + original_mode = stat.S_IMODE(env_path.stat().st_mode) except OSError: pass - raise + try: + with os.fdopen(fd, 'w', **write_kw) as f: + f.writelines(new_lines) + f.flush() + os.fsync(f.fileno()) + atomic_replace(tmp_path, env_path) + # Preserve the original file mode (e.g. 0640 for Docker volume + # mounts) instead of letting _secure_file unconditionally tighten + # to 0600. Mirrors save_env_value(). + if original_mode is not None: + try: + os.chmod(env_path, original_mode) + except OSError: + pass + else: + _secure_file(env_path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise os.environ.pop(key, None) invalidate_env_cache() diff --git a/plugins/dashboard_auth/basic/__init__.py b/plugins/dashboard_auth/basic/__init__.py index 12ec0fe51355..33b92461b272 100644 --- a/plugins/dashboard_auth/basic/__init__.py +++ b/plugins/dashboard_auth/basic/__init__.py @@ -64,7 +64,9 @@ import logging import os import secrets +import tempfile import time +from pathlib import Path from typing import Any, Optional from hermes_cli.dashboard_auth import ( @@ -107,6 +109,59 @@ LAST_SKIP_REASON: str = "" +# --------------------------------------------------------------------------- +# Session epoch (RAH-01: revoke_session()/password rotation invalidation) +# --------------------------------------------------------------------------- +# +# Access/refresh tokens are stateless HMAC blobs — verify_session() has no +# server-side session to check. revoke_session() alone therefore cannot +# invalidate anything already issued. To make "logout" and "password +# changed" actually reject prior tokens without turning this into a stateful +# session store, every minted token carries an "epoch" claim; verification +# rejects any token whose epoch doesn't match the provider's current one. +# The epoch is persisted (JSON, atomic replace) so it survives process +# restarts and is shared across multi-worker deployments that already share +# an explicit `secret` — the same file both workers' registrations read. + + +def _epoch_store_path() -> Path: + from hermes_constants import get_hermes_home + + return get_hermes_home() / "dashboard_auth_basic_session_epoch.json" + + +def _load_epoch_state(path: Path) -> dict: + try: + data = json.loads(path.read_text(encoding="utf-8")) + if isinstance(data, dict): + return { + "epoch": int(data.get("epoch", 0)), + "fingerprint": str(data.get("fingerprint", "")), + } + except (OSError, ValueError, TypeError): + pass + return {"epoch": 0, "fingerprint": ""} + + +def _save_epoch_state(path: Path, epoch: int, fingerprint: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp( + dir=str(path.parent), suffix=".tmp", prefix=".epoch_" + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump({"epoch": epoch, "fingerprint": fingerprint}, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path) + except BaseException: + try: + os.unlink(tmp_path) + except OSError: + pass + raise + + # --------------------------------------------------------------------------- # Password hashing (stdlib scrypt) # --------------------------------------------------------------------------- @@ -212,6 +267,8 @@ def __init__( password_hash: str, secret: bytes, ttl_seconds: int = _DEFAULT_TTL_SECONDS, + credential_fingerprint: Optional[str] = None, + epoch_store_path: Optional[Path] = None, ) -> None: if not username: raise ValueError("username must be non-empty") @@ -224,6 +281,36 @@ def __init__( self._secret = secret self._ttl = max(60, int(ttl_seconds)) + # credential_fingerprint is None for direct/test construction — the + # provider then behaves exactly as before (in-memory-only epoch, + # no disk state). register() always passes a real fingerprint + # derived from the credential source material (see module docstring + # above _epoch_store_path), so production instances get persisted, + # cross-restart/cross-worker session-epoch invalidation. + self._credential_fingerprint = credential_fingerprint + self._epoch_store_path = epoch_store_path or _epoch_store_path() + if credential_fingerprint is not None: + state = _load_epoch_state(self._epoch_store_path) + if state["fingerprint"] and state["fingerprint"] != credential_fingerprint: + # Credential source changed since the epoch file was last + # written — rotate the epoch so every previously issued + # token (signed under the old epoch) stops verifying. + self._epoch = state["epoch"] + 1 + _save_epoch_state( + self._epoch_store_path, self._epoch, credential_fingerprint + ) + elif state["fingerprint"] != credential_fingerprint: + # First time this credential is recorded (empty fingerprint + # on disk) — keep the existing epoch, just record it. + self._epoch = state["epoch"] + _save_epoch_state( + self._epoch_store_path, self._epoch, credential_fingerprint + ) + else: + self._epoch = state["epoch"] + else: + self._epoch = 0 + # ---- OAuth methods: not used (pure-password provider) ------------------ def start_login(self, *, redirect_uri: str) -> LoginStart: @@ -266,6 +353,7 @@ def verify_session(self, *, access_token: str) -> Optional[Session]: payload is None or payload.get("kind") != "access" or payload.get("exp", 0) <= int(time.time()) + or payload.get("epoch", -1) != self._epoch ): return None return self._session_from_payload(access_token, "", payload) @@ -278,14 +366,27 @@ def refresh_session(self, *, refresh_token: str) -> Session: payload is None or payload.get("kind") != "refresh" or payload.get("exp", 0) <= int(time.time()) + or payload.get("epoch", -1) != self._epoch ): raise RefreshExpiredError("refresh token expired or invalid") return self._mint_session(str(payload.get("sub", self._username))) def revoke_session(self, *, refresh_token: str) -> None: - # Stateless tokens — nothing to revoke server-side. The session - # expires within its TTL. Best-effort no-op, must not raise. + # Stateless tokens: there is no per-session server state to delete. + # Instead, bump the session epoch — every access/refresh token + # already issued (signed under the old epoch) stops verifying + # immediately, in this instance and (once persisted) in every other + # worker/process that shares this provider's epoch store. This is a + # "logout everywhere" operation; BasicAuthProvider has exactly one + # identity (single configured username), so that is the correct + # granularity — there is no per-session/per-device state to target + # more narrowly without turning this into a stateful session store. _ = refresh_token + self._epoch += 1 + if self._credential_fingerprint is not None: + _save_epoch_state( + self._epoch_store_path, self._epoch, self._credential_fingerprint + ) return None # ---- internals --------------------------------------------------------- @@ -294,10 +395,16 @@ def _mint_session(self, user_id: str) -> Session: now = int(time.time()) exp = now + self._ttl access_token = _sign( - {"sub": user_id, "kind": "access", "exp": exp}, self._secret + {"sub": user_id, "kind": "access", "exp": exp, "epoch": self._epoch}, + self._secret, ) refresh_token = _sign( - {"sub": user_id, "kind": "refresh", "exp": now + _REFRESH_TTL_SECONDS}, + { + "sub": user_id, + "kind": "refresh", + "exp": now + _REFRESH_TTL_SECONDS, + "epoch": self._epoch, + }, self._secret, ) return Session( @@ -451,19 +558,30 @@ def register(ctx) -> None: "HERMES_DASHBOARD_BASIC_AUTH_PASSWORD", "" ).strip() if plaintext_from_env: + # Fingerprint the plaintext itself, not the hash: hash_password() + # salts randomly, so the hash differs on every restart even when the + # password hasn't changed — hashing that would falsely look like a + # rotation and invalidate every session on every restart (RAH-01). + credential_source = plaintext_from_env password_hash = hash_password(plaintext_from_env) logger.info( "dashboard-auth-basic: hashed env-supplied password in-memory " "(overrides any config password_hash)." ) elif not password_hash: - # config-only plaintext password. + # config-only plaintext password — same salting concern as above. + credential_source = plaintext password_hash = hash_password(plaintext) logger.info( "dashboard-auth-basic: hashed plaintext password in-memory. " "For production, precompute dashboard.basic_auth.password_hash " "and remove the plaintext password from config." ) + else: + # A precomputed password_hash from config/env is stable across + # restarts already (the operator sets it once), so it's safe to + # fingerprint directly. + credential_source = password_hash secret = _resolve_secret(section) @@ -472,12 +590,20 @@ def register(ctx) -> None: except ValueError: ttl = _DEFAULT_TTL_SECONDS + # Binds the epoch-rotation fingerprint to *this* username too, so + # changing the configured username (a distinct identity) also + # invalidates prior sessions, not just a password change. + credential_fingerprint = hashlib.sha256( + f"{username}\x00{credential_source}".encode("utf-8") + ).hexdigest() + try: provider = BasicAuthProvider( username=username, password_hash=password_hash, secret=secret, ttl_seconds=ttl, + credential_fingerprint=credential_fingerprint, ) except ValueError as exc: LAST_SKIP_REASON = f"BasicAuthProvider construction failed: {exc}" diff --git a/tests/hermes_cli/test_config.py b/tests/hermes_cli/test_config.py index c2db04cd9661..502737518bf7 100644 --- a/tests/hermes_cli/test_config.py +++ b/tests/hermes_cli/test_config.py @@ -1,6 +1,7 @@ """Tests for hermes_cli configuration management.""" import os +import threading from pathlib import Path from unittest.mock import patch @@ -323,6 +324,41 @@ def test_save_env_value_already_quoted_input_is_not_double_wrapped_idempotently( assert load_env()["TERMINAL_SSH_KEY"] == raw +class TestSaveEnvValueConcurrency: + def test_concurrent_saves_of_distinct_keys_both_survive(self, tmp_path): + """Regression: two threads each read the same on-disk snapshot, + append their own key, and atomic-replace — without serialization, + the second write silently loses the first (lost update). Both keys + must be present after both threads finish. + """ + env_path = tmp_path / ".env" + env_path.write_text("BASE=kept\n") + + barrier = threading.Barrier(2) + errors = [] + + def _save(key, value): + try: + barrier.wait(timeout=5) + with patch.dict(os.environ, {"HERMES_HOME": str(tmp_path)}): + save_env_value(key, value) + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + t1 = threading.Thread(target=_save, args=("RACE_A", "a")) + t2 = threading.Thread(target=_save, args=("RACE_B", "b")) + t1.start() + t2.start() + t1.join(timeout=10) + t2.join(timeout=10) + + assert not errors + content = env_path.read_text(encoding="utf-8") + assert "BASE=kept" in content + assert "RACE_A=a" in content + assert "RACE_B=b" in content + + class TestRemoveEnvValue: def test_removes_key_from_env_file(self, tmp_path): env_path = tmp_path / ".env" diff --git a/tests/plugins/dashboard_auth/test_basic_provider.py b/tests/plugins/dashboard_auth/test_basic_provider.py index cbb4112c0517..50afe8513b0b 100644 --- a/tests/plugins/dashboard_auth/test_basic_provider.py +++ b/tests/plugins/dashboard_auth/test_basic_provider.py @@ -129,6 +129,76 @@ def test_revoke_is_silent(self, basic): p = self._make(basic) p.revoke_session(refresh_token="anything") # must not raise + def test_revoke_session_invalidates_prior_access_token(self, basic): + """RAH-01: revoke_session() must not be a pure no-op — a token minted + before the call must stop verifying after it, on the same instance.""" + p = self._make(basic) + s = p.complete_password_login(username="admin", password="hunter2") + assert p.verify_session(access_token=s.access_token) is not None + p.revoke_session(refresh_token=s.refresh_token) + assert p.verify_session(access_token=s.access_token) is None + with pytest.raises(RefreshExpiredError): + p.refresh_session(refresh_token=s.refresh_token) + + def test_password_rotation_invalidates_prior_session_across_restart( + self, basic, tmp_path + ): + """RAH-01: rotating the password (same explicit secret, simulating a + process restart via a fresh provider instance) must invalidate + sessions minted under the old password — not just the process that + revoked. credential_fingerprint + a shared epoch_store_path model + what register() does across two real process starts.""" + store = tmp_path / "epoch.json" + secret = secrets.token_bytes(32) + old_hash = basic.hash_password("old-password") + p_old = basic.BasicAuthProvider( + username="admin", + password_hash=old_hash, + secret=secret, + credential_fingerprint="fp-old-password", + epoch_store_path=store, + ) + s = p_old.complete_password_login(username="admin", password="old-password") + assert p_old.verify_session(access_token=s.access_token) is not None + + # Simulate a restart after rotating the password: a brand-new + # provider instance, same secret + epoch store, different fingerprint. + new_hash = basic.hash_password("new-password") + p_new = basic.BasicAuthProvider( + username="admin", + password_hash=new_hash, + secret=secret, + credential_fingerprint="fp-new-password", + epoch_store_path=store, + ) + assert p_new.verify_session(access_token=s.access_token) is None + + def test_same_password_across_restart_keeps_session_valid(self, basic, tmp_path): + """A restart with the SAME credential fingerprint (nothing rotated) + must NOT bump the epoch — otherwise every restart would silently log + everyone out, breaking the documented explicit-secret portability + contract.""" + store = tmp_path / "epoch.json" + secret = secrets.token_bytes(32) + h = basic.hash_password("hunter2") + p1 = basic.BasicAuthProvider( + username="admin", + password_hash=h, + secret=secret, + credential_fingerprint="fp-stable", + epoch_store_path=store, + ) + s = p1.complete_password_login(username="admin", password="hunter2") + + p2 = basic.BasicAuthProvider( + username="admin", + password_hash=h, + secret=secret, + credential_fingerprint="fp-stable", + epoch_store_path=store, + ) + assert p2.verify_session(access_token=s.access_token) is not None + def test_oauth_methods_raise_not_implemented(self, basic): p = self._make(basic) with pytest.raises(NotImplementedError): diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 9eb3468d5bef..3a7627ed0bc2 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -11675,6 +11675,11 @@ def test_model_save_key_uses_credential_lifecycle_and_picker_context(monkeypatch "name": "Test Provider", "models": ["test-model"], "total_models": 1, + # A real build_models_payload(picker_hints=True) call always sets + # this via _apply_picker_hints; the mock below bypasses that, so it + # must supply the field itself (RAH-04: model.save_key no longer + # force-overwrites whatever the real inventory computed). + "authenticated": True, } server._sessions["save-key-session"] = _session(agent=agent) monkeypatch.setattr( @@ -11723,6 +11728,55 @@ def test_model_save_key_uses_credential_lifecycle_and_picker_context(monkeypatch ) +def test_model_save_key_does_not_force_authenticated_when_inventory_says_false( + monkeypatch, +): + """RAH-04: a saved key that the rebuilt inventory still treats as an + unauthenticated skeleton row (e.g. the model-list probe failed) must be + reported that way, not silently upgraded to authenticated=True.""" + env_var = "TEST_PROVIDER_API_KEY" + agent = object() + provider = { + "slug": "test-provider", + "name": "Test Provider", + "models": [], + "total_models": 0, + "authenticated": False, + } + server._sessions["save-key-session-2"] = _session(agent=agent) + monkeypatch.setattr( + "hermes_cli.auth.PROVIDER_REGISTRY", + { + "test-provider": types.SimpleNamespace( + name="Test Provider", + auth_type="api_key", + api_key_env_vars=(env_var,), + ) + }, + ) + monkeypatch.setattr("hermes_cli.config.is_managed", lambda: False) + monkeypatch.setattr( + "hermes_cli.credential_lifecycle.save_provider_env_credential", Mock() + ) + monkeypatch.setattr(server, "_model_picker_context", Mock(return_value=object())) + monkeypatch.setattr( + "hermes_cli.inventory.build_models_payload", + Mock(return_value={"providers": [provider]}), + ) + + resp = server._methods["model.save_key"]( + 104, + { + "slug": "test-provider", + "api_key": "some-key", + "session_id": "save-key-session-2", + }, + ) + + assert "result" in resp, resp + assert resp["result"]["provider"]["authenticated"] is False + + # --------------------------------------------------------------------------- # prompt.submit — auto-title # --------------------------------------------------------------------------- @@ -14204,8 +14258,8 @@ def test_session_close_rpc_claims_then_tears_down(monkeypatch): def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): seen = [] monkeypatch.setattr( - server, "_close_session_by_id", - lambda sid, *, end_reason: bool(seen.append((sid, end_reason))) or True, + server, "_teardown_popped_session", + lambda session, *, end_reason: bool(seen.append((session["_sid"], end_reason))) or True, ) # Detached session "b" would schedule a real grace-reap threading.Timer that # outlives the test; grace=0 short-circuits it so no thread lingers. @@ -14217,11 +14271,79 @@ def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): try: server._close_sessions_for_transport(transport, end_reason="ws_disconnect") assert seen == [("a", "ws_disconnect")] # only the flagged one closed + assert "a" not in server._sessions # claimed/popped assert server._sessions["b"]["transport"] is server._detached_ws_transport # re-pointed finally: server._sessions.clear() +def test_close_sessions_for_transport_skips_session_reattached_mid_teardown(monkeypatch): + """Regression for the disconnect/reconnect race: if session.resume rebinds + a session onto a new (live) transport between the ownership snapshot and + this function's per-session claim, the old transport's teardown must not + close it or stomp its transport back to the detached sentinel. + + Unlike a naive version of this test that starts both sessions already on + ``new_transport`` (which never even enters ``owned_sids`` and exercises + nothing beyond the initial filter), this drives the actual interleaving + the fix revalidates against: the session starts on ``old_transport`` so + the snapshot captures it, and the reattach happens strictly between that + snapshot and this function's per-sid claim under ``_session_resume_lock`` + — the exact TOCTOU window closed by the WS disconnect/reconnect fix. A + ``_RaceLock`` stand-in for the module's real resume lock performs the + reattach the first time the loop acquires it, modeling session.resume + winning the lock race before teardown's revalidation runs. Against the + pre-fix implementation (no per-sid lock/revalidation at all) the injected + mutation never fires and the session is torn down/stomped regardless — + this test fails there and passes only once the race window is closed.""" + seen = [] + monkeypatch.setattr( + server, "_teardown_popped_session", + lambda session, *, end_reason: bool(seen.append((session["_sid"], end_reason))) or True, + ) + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0) + old_transport = object() + new_transport = object() + server._sessions.clear() + server._sessions["x"] = {"transport": old_transport, "close_on_disconnect": True} + + real_resume_lock = server._session_resume_lock + + class _RaceLock: + """Wraps the real resume lock. The first acquire simulates + session.resume winning the race: it rebinds "x" onto new_transport + right after the snapshot above already captured it as owned by + old_transport, but before this function's own per-sid claim (which + also needs this lock) gets to revalidate.""" + + def __init__(self): + self._fired = False + + def __enter__(self): + real_resume_lock.acquire() + if not self._fired: + self._fired = True + with server._sessions_lock: + server._sessions["x"]["transport"] = new_transport + return self + + def __exit__(self, *exc_info): + real_resume_lock.release() + return False + + monkeypatch.setattr(server, "_session_resume_lock", _RaceLock()) + try: + reaped, detached = server._close_sessions_for_transport( + old_transport, end_reason="ws_disconnect" + ) + assert reaped == 0 and detached == 0 + assert seen == [] # teardown must not have claimed the reattached session + assert "x" in server._sessions # not closed + assert server._sessions["x"]["transport"] is new_transport # not stomped back + finally: + server._sessions.clear() + + def test_session_create_records_close_on_disconnect_flag(monkeypatch): monkeypatch.setattr(server, "_start_agent_build", lambda sid, session: None) server._sessions.clear() diff --git a/tui_gateway/methods_complete.py b/tui_gateway/methods_complete.py index 701c11f0eaed..6ef6fb264d96 100644 --- a/tui_gateway/methods_complete.py +++ b/tui_gateway/methods_complete.py @@ -410,7 +410,12 @@ def _(rid, params: dict) -> dict: (p for p in payload["providers"] if p["slug"] == slug), None ) if provider_data is None: - # Key was saved but provider didn't appear — still return success. + # Key was saved but provider didn't appear in the rebuilt + # inventory — still return success. picker_hints sets + # `authenticated` from the row state for every row that DID + # come back from build_models_payload; this synthetic fallback + # is the only case that never goes through that path, so it's + # the only case that needs the field set explicitly here. provider_data = { "slug": slug, "name": pconfig.name, @@ -419,9 +424,6 @@ def _(rid, params: dict) -> dict: "total_models": 0, "authenticated": True, } - # picker_hints sets `authenticated` from the row state, but the - # synthetic fallback above doesn't go through that path. - provider_data["authenticated"] = True return _ok(rid, {"provider": provider_data}) except Exception as e: return _err(rid, 5034, str(e)) diff --git a/tui_gateway/server.py b/tui_gateway/server.py index d637fb2f3c1f..1e4bcb4ad3f1 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1066,9 +1066,8 @@ def _close_sessions_for_transport( transport, *, end_reason: str = "ws_disconnect" ) -> tuple[int, int]: """On transport disconnect, reap the sessions that opted into - close_on_disconnect (sidecar/dashboard) immediately via the unified - ``_close_session_by_id`` path, and re-point the rest back to stdio so later - emits don't hit a dead socket. + close_on_disconnect (sidecar/dashboard) immediately, and re-point the rest + back to the drop sentinel so later emits don't hit a dead socket. Non-flagged detached sessions are handed to the grace-windowed WS-orphan reaper (``_schedule_ws_orphan_reap``): a quick reconnect / session.resume @@ -1077,23 +1076,48 @@ def _close_sessions_for_transport( the single WS-disconnect teardown entry point — there is no second independent reap loop in ``handle_ws``. + The initial snapshot below is taken outside any lock, so each session's + close/detach decision is re-validated (transport still matches) under + ``_session_resume_lock`` immediately before acting. session.resume's + warm-reuse rebind (``_reuse_live_payload`` / ``_live_session_payload``) + takes the same lock to repoint ``session["transport"]``, so a reconnect + that wins the race is never silently closed or stomped back to the + detached sentinel by this (now-stale) transport's teardown. + Returns ``(reaped, detached)`` counts for disconnect-path observability.""" with _sessions_lock: - owned = [(sid, s) for sid, s in _sessions.items() if s.get("transport") is transport] + owned_sids = [sid for sid, s in _sessions.items() if s.get("transport") is transport] reaped = 0 detached = 0 - for sid, session in owned: - if session.get("close_on_disconnect"): - _close_session_by_id(sid, end_reason=end_reason) + for sid in owned_sids: + with _session_resume_lock: + with _sessions_lock: + session = _sessions.get(sid) + if session is None or session.get("transport") is not transport: + # Already torn down, or reattached to a new transport + # since the snapshot above — leave it alone. + continue + if session.get("close_on_disconnect"): + del _sessions[sid] + session["_sid"] = sid + to_teardown, to_detach = session, None + else: + # Point detached sessions at the drop sentinel (NOT real + # stdio) so _ws_session_is_orphaned recognizes them and the + # grace-reap can actually fire; a standalone `hermes --tui` + # keeps real _stdio. + session["transport"] = _detached_ws_transport + to_teardown, to_detach = None, sid + # Slow teardown/timer scheduling happens after releasing both locks — + # see the module note above _pop_session_by_id about keeping that work + # off _session_resume_lock. + if to_teardown is not None: + _teardown_popped_session(to_teardown, end_reason=end_reason) reaped += 1 else: - # Point detached sessions at the drop sentinel (NOT real stdio) so - # _ws_session_is_orphaned recognizes them and the grace-reap can - # actually fire; a standalone `hermes --tui` keeps real _stdio. - session["transport"] = _detached_ws_transport detached += 1 try: - _schedule_ws_orphan_reap(sid) + _schedule_ws_orphan_reap(to_detach) except Exception: pass return reaped, detached