diff --git a/api/gateway_chat.py b/api/gateway_chat.py index d8470b08..f18ee584 100644 --- a/api/gateway_chat.py +++ b/api/gateway_chat.py @@ -839,6 +839,7 @@ _GATEWAY_RUN_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", GATEWAY_REATTACH_POLL_INTERVAL = 2.0 # Consecutive unreachable polls tolerated before the reattached turn is failed (~5 min at 2s). GATEWAY_REATTACH_MAX_POLL_FAILURES = 150 +_REATTACH_SCAN_HEAD_BYTES = 16 * 1024 def _record_gateway_run(session_id: str, stream_id: str, run_id: str, **extra) -> None: @@ -945,15 +946,12 @@ def resume_gateway_runs_after_restart() -> list[str]: try: if not webui_gateway_chat_enabled(get_config()): return [] - entries = json.loads(_models.SESSION_INDEX_FILE.read_bytes()) + candidates = _sidecars_with_active_stream(_models.SESSION_DIR) except Exception: - logger.debug("gateway reattach: no readable session index", exc_info=True) + logger.warning("gateway reattach: could not scan session sidecars", exc_info=True) return [] resumed: list[str] = [] - for entry in entries if isinstance(entries, list) else []: - if not isinstance(entry, dict) or not entry.get("active_stream_id"): - continue - sid = str(entry.get("session_id") or "") + for sid in candidates: try: if _resume_gateway_run_for_session(_models.Session.load_metadata_only(sid)): resumed.append(sid) @@ -964,6 +962,40 @@ def resume_gateway_runs_after_restart() -> list[str]: return resumed +def _sidecars_with_active_stream(session_dir) -> list[str]: + """Session ids whose sidecar may still carry a pending turn. + + Scans the sidecars rather than the derived index, which can lag a save. + A head read skips the common idle case without parsing the file. + """ + ids = [] + for path in sorted(session_dir.glob("*.json")): + if path.name.startswith("_"): + continue + try: + with path.open("rb") as fp: + head = fp.read(_REATTACH_SCAN_HEAD_BYTES) + except OSError: + continue + if b'"active_stream_id": null' not in head: + ids.append(path.stem) + return ids + + +def _gateway_endpoint_for_reattach(session, run: dict) -> tuple[str, str]: + """Resolve the gateway that owns ``run`` using the session's profile, not the process default.""" + from api import profiles as _profiles + from api.config import get_config + + with _profiles.profile_scope_for_detached_worker( + getattr(session, "profile", None), "gateway reattach", logger_override=logger, + ): + base_url = _gateway_base_url(get_config()) + api_key = _gateway_api_key() + # The run lives on the gateway that accepted it, even if config changed since. + return str(run.get("base_url") or base_url).rstrip("/"), api_key + + def _resume_gateway_run_for_session(session) -> bool: from api.config import create_stream_channel, register_session_writeback_owner, register_stream_owner @@ -973,6 +1005,7 @@ def _resume_gateway_run_for_session(session) -> bool: if not stream_id or not run_id or run.get("stream_id") != stream_id: return False sid = session.session_id + endpoint = _gateway_endpoint_for_reattach(session, run) with STREAMS_LOCK: if stream_id in STREAMS: return False @@ -995,6 +1028,7 @@ def _resume_gateway_run_for_session(session) -> bool: "goal_related": bool(run.get("goal_related")), "regeneration": bool(run.get("regeneration")), "reattach_run_id": run_id, + "reattach_endpoint": endpoint, }, name=f"gateway-reattach-{stream_id[:12]}", daemon=True, @@ -1111,6 +1145,7 @@ def _run_gateway_chat_streaming( goal_related=False, regeneration=False, reattach_run_id=None, + reattach_endpoint=None, ): """Bridge a WebUI chat turn through Hermes Gateway's API server. @@ -1197,8 +1232,11 @@ def _run_gateway_chat_streaming( model=model, model_provider=model_provider, ) - base_url = _gateway_base_url(cfg) - api_key = _gateway_api_key() + if reattach_endpoint: + base_url, api_key = reattach_endpoint + else: + base_url = _gateway_base_url(cfg) + api_key = _gateway_api_key() try: from api.config import _main_model_request_overrides _gw_overrides = _main_model_request_overrides( @@ -1280,6 +1318,7 @@ def _run_gateway_chat_streaming( active_provider=(model_provider or ""), on_run_id=lambda run_id: _record_gateway_run( session_id, stream_id, run_id, + base_url=base_url, regeneration=bool(regeneration), goal_related=bool(goal_related), ), diff --git a/tests/test_gateway_run_reattach_after_restart.py b/tests/test_gateway_run_reattach_after_restart.py index 2d459104..0ddf872a 100644 --- a/tests/test_gateway_run_reattach_after_restart.py +++ b/tests/test_gateway_run_reattach_after_restart.py @@ -113,6 +113,7 @@ def test_runs_api_start_sends_idempotency_key_and_persists_run_id(isolated_sessi assert captured["post_headers"]["Idempotency-key"] == f"webui-{stream_id}" assert captured["persisted_at_events"]["run_id"] == "run_live" assert captured["persisted_at_events"]["stream_id"] == stream_id + assert captured["persisted_at_events"]["base_url"] == "http://gateway.local" saved = json.loads((isolated_sessions / f"{s.session_id}.json").read_text()) assert saved["gateway_run"] is None assert saved["active_stream_id"] is None @@ -253,3 +254,110 @@ def test_cancel_still_stops_a_reattached_run(isolated_sessions, monkeypatch): assert saved["active_stream_id"] is None assert saved["gateway_run"] is None assert not any(m.get("content") == "long task" and m.get("role") == "assistant" for m in saved["messages"]) + + +def _poll_until_completed(monkeypatch, seen): + def fake_status(base_url, api_key, run_id): + seen.append((base_url, api_key, run_id)) + return {"run_id": run_id, "status": "completed", "output": "answer"} + + monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", fake_status) + + +@pytest.mark.parametrize("index_state", ["missing", "stale", "corrupt"]) +def test_reattach_does_not_depend_on_the_session_index(isolated_sessions, monkeypatch, index_state): + sid, _stream_id = _orphaned_gateway_turn() + index = isolated_sessions / "_index.json" + if index_state == "missing": + index.unlink(missing_ok=True) + elif index_state == "stale": + # Sidecar saved, index update lost: the row still shows no active stream. + index.write_text(json.dumps([{"session_id": sid, "active_stream_id": None}])) + else: + index.write_text("{not json") + seen = [] + _poll_until_completed(monkeypatch, seen) + + assert gateway_chat.resume_gateway_runs_after_restart() == [sid] + _wait_for_reattach_threads() + + saved = json.loads((isolated_sessions / f"{sid}.json").read_text()) + assert saved["messages"][-1]["content"] == "answer" + assert saved["active_stream_id"] is None + + +def test_reattach_skips_idle_sidecars_without_parsing_them(isolated_sessions, monkeypatch): + idle = new_session() + idle.messages = [{"role": "user", "content": "x", "timestamp": 1.0}] + idle.save() + sid, _ = _orphaned_gateway_turn() + loaded = [] + real_load = models.Session.load_metadata_only + monkeypatch.setattr( + models.Session, "load_metadata_only", + classmethod(lambda cls, s, **kw: loaded.append(s) or real_load.__func__(cls, s, **kw)), + ) + _poll_until_completed(monkeypatch, []) + + assert gateway_chat.resume_gateway_runs_after_restart() == [sid] + _wait_for_reattach_threads() + assert idle.session_id not in loaded + assert sid in loaded + + +def test_reattach_resolves_gateway_from_the_session_profile(isolated_sessions, monkeypatch): + """The process may restart under another default profile; poll the session profile's gateway.""" + import contextlib + import os + from api import profiles + + sid, stream_id = _orphaned_gateway_turn() + s = models.Session.load(sid) + s.profile = "work" + s.save(touch_updated_at=False) + models.SESSIONS.clear() + scopes = [] + + @contextlib.contextmanager + def fake_scope(profile_name, purpose="", logger_override=None): + scopes.append(profile_name) + env = {"work": ("http://work-gateway:8642", "work-key")}.get(profile_name) + old = {k: os.environ.get(k) for k in ("HERMES_WEBUI_GATEWAY_BASE_URL", "HERMES_WEBUI_GATEWAY_API_KEY")} + if env: + os.environ["HERMES_WEBUI_GATEWAY_BASE_URL"], os.environ["HERMES_WEBUI_GATEWAY_API_KEY"] = env + try: + yield + finally: + for k, v in old.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + monkeypatch.setattr(profiles, "profile_scope_for_detached_worker", fake_scope) + monkeypatch.setenv("HERMES_WEBUI_GATEWAY_API_KEY", "default-key") + seen = [] + _poll_until_completed(monkeypatch, seen) + + assert gateway_chat.resume_gateway_runs_after_restart() == [sid] + _wait_for_reattach_threads() + + assert scopes == ["work"] + assert {(b, k) for b, k, _ in seen} == {("http://work-gateway:8642", "work-key")} + saved = json.loads((isolated_sessions / f"{sid}.json").read_text()) + assert saved["messages"][-1]["content"] == "answer" + + +def test_reattach_polls_the_gateway_that_accepted_the_run(isolated_sessions, monkeypatch): + sid, stream_id = _orphaned_gateway_turn() + s = models.Session.load(sid) + s.gateway_run = {**s.gateway_run, "base_url": "http://accepted-gateway:8642/"} + s.save(touch_updated_at=False) + models.SESSIONS.clear() + seen = [] + _poll_until_completed(monkeypatch, seen) + + gateway_chat.resume_gateway_runs_after_restart() + _wait_for_reattach_threads() + + assert {b for b, _k, _r in seen} == {"http://accepted-gateway:8642"}