diff --git a/api/gateway_chat.py b/api/gateway_chat.py index 0086e055..22c4ea0d 100644 --- a/api/gateway_chat.py +++ b/api/gateway_chat.py @@ -583,6 +583,35 @@ def _relay_gateway_run_approval(session_id, run_id, payload, base_url, api_key, put_gateway_event("approval", {**(head or approval_data), "pending_count": total}) +def _admit_gateway_run(url_runs, headers, run_body, stream_id) -> str: + """POST /v1/runs; the same stream and body return the originally admitted run id.""" + req = urllib.request.Request( + url_runs, + data=json.dumps(run_body).encode("utf-8"), + # Durable run record on the gateway: GET /v1/runs/{id} survives either side restarting. + headers={**headers, "Idempotency-Key": f"webui-{stream_id}"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=30) as resp: + run_data = json.loads(resp.read(65536)) + run_id = str(run_data.get("run_id") or run_data.get("id") or "").strip() + if not run_id: + raise ValueError(f"Gateway runs API returned no run_id: {run_data!r}") + return run_id + + +def _gateway_run_headers(session_id, api_key) -> dict: + headers = { + "Content-Type": "application/json", + "Accept": "application/json", + "X-Hermes-Session-Id": session_id, + } + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + headers["X-Hermes-Session-Key"] = f"webui:{session_id}" + return headers + + def _run_gateway_runs_api_streaming( session_id, msg_text, model, workspace, stream_id, base_url, api_key, prefill_messages, body_extras, @@ -594,14 +623,7 @@ def _run_gateway_runs_api_streaming( """Submit via POST /v1/runs and relay SSE events including approval.""" try: url_runs = f"{base_url.rstrip('/')}/v1/runs" - headers = { - "Content-Type": "application/json", - "Accept": "application/json", - "X-Hermes-Session-Id": session_id, - } - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - headers["X-Hermes-Session-Key"] = f"webui:{session_id}" + headers = _gateway_run_headers(session_id, api_key) message_content: Any = str(msg_text or "") if attachments: try: @@ -654,19 +676,11 @@ def _run_gateway_runs_api_streaming( run_body["instructions"] = "\n\n".join(part for part in instructions_parts if part) if conversation_history: run_body["conversation_history"] = conversation_history - req = urllib.request.Request( - url_runs, - data=json.dumps(run_body).encode("utf-8"), - # Durable run record on the gateway: GET /v1/runs/{id} survives either side restarting. - headers={**headers, "Idempotency-Key": f"webui-{stream_id}"}, - method="POST", - ) update_active_run(stream_id, phase="gateway-request") - with urllib.request.urlopen(req, timeout=30) as resp: - run_data = json.loads(resp.read(65536)) - run_id = str(run_data.get("run_id") or run_data.get("id") or "").strip() - if not run_id: - raise ValueError(f"Gateway runs API returned no run_id: {run_data!r}") + # Persist the exact body first: a restart before the run id is saved replays this admission. + if on_run_id is not None: + on_run_id("", request=run_body) + run_id = _admit_gateway_run(url_runs, headers, run_body, stream_id) except Exception: _finish_gateway_run_starting(stream_id) raise @@ -841,14 +855,16 @@ 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: - """Persist the gateway run id on the pending turn so a restarted WebUI can reattach.""" +def _record_gateway_run(session_id: str, stream_id: str, run_id: str, request=None, **extra) -> None: + """Persist the run id (or, before admission, the exact request) so a restarted WebUI can reattach.""" try: with _get_session_agent_lock(session_id): session = get_session(session_id) if not _stream_writeback_is_current(session, stream_id): return session.gateway_run = {"run_id": run_id, "stream_id": stream_id, **extra} + if not run_id: + session.gateway_run["request"] = request session.save(touch_updated_at=False) except Exception: logger.warning("Failed to persist gateway run %s for session %s", run_id, session_id, exc_info=True) @@ -926,11 +942,8 @@ def resume_gateway_runs_after_restart() -> list[str]: Call before serving so stale-pending repair does not mark these turns interrupted. """ from api import models as _models - from api.config import get_config try: - if not webui_gateway_chat_enabled(get_config()): - return [] candidates = _sidecars_with_active_stream(_models.SESSION_DIR) except Exception: logger.warning("gateway reattach: could not scan session sidecars", exc_info=True) @@ -969,8 +982,8 @@ def _resume_gateway_run_for_session(session) -> bool: run = (session.gateway_run if session is not None else None) or {} stream_id = str(run.get("stream_id") or "") - run_id = str(run.get("run_id") or "").strip() - if not stream_id or not run_id or session.active_stream_id != stream_id: + # A persisted run is the authority whatever the process default backend is. + if not stream_id or not (run.get("run_id") or run.get("request")) or session.active_stream_id != stream_id: return False sid = session.session_id # URL and key come together from the session's profile, never the process default. @@ -991,7 +1004,7 @@ def _resume_gateway_run_for_session(session) -> bool: "model_provider": session.model_provider, "goal_related": bool(run.get("goal_related")), "regeneration": bool(run.get("regeneration")), - "reattach_run_id": run_id, + "reattach_run": run, "reattach_endpoint": endpoint, }, name=f"gateway-reattach-{stream_id[:12]}", @@ -1108,7 +1121,7 @@ def _run_gateway_chat_streaming( model_provider=None, goal_related=False, regeneration=False, - reattach_run_id=None, + reattach_run=None, reattach_endpoint=None, ): """Bridge a WebUI chat turn through Hermes Gateway's API server. @@ -1207,7 +1220,7 @@ def _run_gateway_chat_streaming( except Exception: _gw_overrides = {} _runs_api_enabled = _gateway_use_runs_api_enabled(cfg) - _use_runs_api = bool(reattach_run_id) or (_runs_api_enabled and gateway_supports_approval(base_url, api_key)) + _use_runs_api = bool(reattach_run) or (_runs_api_enabled and gateway_supports_approval(base_url, api_key)) if not _use_runs_api and runs_api_pending_marked: _finish_gateway_run_starting(stream_id, result="fallback") runs_api_pending_marked = False @@ -1257,10 +1270,22 @@ def _run_gateway_chat_streaming( body_extras["reasoning_effort"] = reasoning_effort if _gw_overrides.get("service_tier"): body_extras["service_tier"] = _gw_overrides["service_tier"] + record_run = lambda run_id, request=None: _record_gateway_run( + session_id, stream_id, run_id, request=request, + regeneration=bool(regeneration), goal_related=bool(goal_related), + ) try: - if reattach_run_id: + if reattach_run: + run_id = str(reattach_run.get("run_id") or "").strip() + if not run_id: + # Crashed between admission and saving the id: the same key returns the original run. + run_id = _admit_gateway_run( + f"{base_url.rstrip('/')}/v1/runs", _gateway_run_headers(session_id, api_key), + reattach_run["request"], stream_id, + ) + record_run(run_id) final_text, usage = _await_gateway_run_result( - session_id, stream_id, reattach_run_id, base_url, api_key, + session_id, stream_id, run_id, base_url, api_key, put_gateway_event=put_gateway_event, cancel_event=cancel_event, ) @@ -1274,11 +1299,7 @@ def _run_gateway_chat_streaming( cfg=cfg, session=s, active_provider=(model_provider or ""), - on_run_id=lambda run_id: _record_gateway_run( - session_id, stream_id, run_id, - regeneration=bool(regeneration), - goal_related=bool(goal_related), - ), + on_run_id=record_run, ) except Exception as exc: error_payload = _settle_gateway_terminal_error( diff --git a/tests/test_gateway_run_reattach_after_restart.py b/tests/test_gateway_run_reattach_after_restart.py index a6a38cc6..5ffb9426 100644 --- a/tests/test_gateway_run_reattach_after_restart.py +++ b/tests/test_gateway_run_reattach_after_restart.py @@ -177,15 +177,11 @@ def test_reattach_surfaces_pending_approval_once(isolated_sessions, monkeypatch) assert _saved(sid)["messages"][-1]["content"] == "approved and done" -@pytest.mark.parametrize("case", ["legacy_backend", "no_gateway_run"]) -def test_nothing_to_reattach_is_left_to_stale_pending_repair(isolated_sessions, monkeypatch, case): +def test_nothing_to_reattach_is_left_to_stale_pending_repair(isolated_sessions, monkeypatch): sid, stream_id = _orphaned_gateway_turn() - if case == "legacy_backend": - monkeypatch.delenv("HERMES_WEBUI_CHAT_BACKEND") - else: - s = models.Session.load(sid) - s.gateway_run = None - s.save() + s = models.Session.load(sid) + s.gateway_run = None + s.save() monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", lambda *a, **k: pytest.fail("must not poll")) assert gateway_chat.resume_gateway_runs_after_restart() == [] assert stream_id not in STREAMS @@ -269,8 +265,12 @@ def test_reattach_skips_idle_sidecars_without_parsing_them(isolated_sessions, mo 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.""" +@pytest.mark.parametrize("default_backend", ["gateway", "local"]) +def test_reattach_resolves_gateway_from_the_session_profile(isolated_sessions, monkeypatch, default_backend): + """The process may restart under another default profile, even a local one; poll the session profile's gateway.""" + if default_backend == "local": + monkeypatch.delenv("HERMES_WEBUI_CHAT_BACKEND") + monkeypatch.delenv("HERMES_WEBUI_GATEWAY_USE_RUNS_API") sid, stream_id = _orphaned_gateway_turn() s = models.Session.load(sid) s.profile = "work" @@ -295,8 +295,66 @@ def test_reattach_resolves_gateway_from_the_session_profile(isolated_sessions, m _wait_for_reattach_threads() assert scopes == ["work"] - assert {(b, k) for b, k, _ in seen} == {("http://work-gateway:8642", "work-key")} - assert _saved(sid)["messages"][-1]["content"] == "answer" + assert {(b, k, r) for b, k, r in seen} == {("http://work-gateway:8642", "work-key", "run_survivor")} + saved = _saved(sid) + assert saved["messages"][-1]["content"] == "answer" + assert saved["active_stream_id"] is None and saved["gateway_run"] is None + + +def test_crash_after_admission_before_run_id_saved_replays_the_same_admission(isolated_sessions, monkeypatch): + s = new_session() + stream_id = "stream-admitted" + s.active_stream_id = stream_id + s.pending_user_message = "long task" + s.pending_attachments = [] + s.pending_started_at = 1.0 + s.save() + posts = [] + + class Crash(BaseException): + pass + + def fake_urlopen(req, timeout=None): + assert req.get_method() == "POST" + posts.append((req.get_header("Idempotency-key"), req.data)) + return io.BytesIO(b'{"run_id":"run_original","replayed":%s}' % (b"true" if len(posts) > 1 else b"false")) + + real_record = gateway_chat._record_gateway_run + + def record_then_die(session_id, stream_id, run_id, **kw): + if run_id: + raise Crash # process killed after the POST response, before the id is saved + real_record(session_id, stream_id, run_id, **kw) + + monkeypatch.setattr(gateway_chat, "gateway_supports_approval", lambda *a, **k: True) + monkeypatch.setattr(gateway_chat.urllib.request, "urlopen", fake_urlopen) + monkeypatch.setattr(gateway_chat, "_record_gateway_run", record_then_die) + with STREAMS_LOCK: + STREAMS[stream_id] = create_stream_channel() + with pytest.raises(Crash): + gateway_chat._run_gateway_runs_api_streaming( + s.session_id, "long task", "test-model", "/tmp", stream_id, "http://gateway.local", "", [], {}, + put_gateway_event=lambda *a: None, cancel_event=threading.Event(), session=s, + on_run_id=lambda run_id, **kw: gateway_chat._record_gateway_run(s.session_id, stream_id, run_id, **kw), + ) + monkeypatch.setattr(gateway_chat, "_record_gateway_run", real_record) + models.SESSIONS.clear() + with STREAMS_LOCK: + STREAMS.pop(stream_id, None) + ACTIVE_RUNS.pop(stream_id, None) + assert _saved(s.session_id)["gateway_run"]["run_id"] == "" + seen = _poll_until_completed(monkeypatch) + + assert gateway_chat.resume_gateway_runs_after_restart() == [s.session_id] + _wait_for_reattach_threads() + + # Replayed with the same key and byte-identical body, so the gateway returns the original run. + assert len(posts) == 2 and posts[0] == posts[1] == (f"webui-{stream_id}", posts[0][1]) + assert {r for _b, _k, r in seen} == {"run_original"} + saved = _saved(s.session_id) + assert saved["messages"][-1]["content"] == "answer" + assert not any(m.get("_error") for m in saved["messages"]) + assert saved["active_stream_id"] is None and saved["gateway_run"] is None @pytest.mark.parametrize("code, message", [