fix(gateway): reattach under the session profile and recover runs admitted before the id was saved

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(

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", [