From 00973abe0c2404e7df2cf65ec1f07f53c8e8b05a Mon Sep 17 00:00:00 2001 From: carlotestor Date: Thu, 24 Sep 2026 15:18:53 +0000 Subject: [PATCH] fix(gateway): reattach with the session profile's own endpoint (root too); keep the prompt when a replayed run ends cancelled --- api/gateway_chat.py | 32 +++++++-- ...test_gateway_run_reattach_after_restart.py | 72 +++++++++++++------ 2 files changed, 79 insertions(+), 25 deletions(-) diff --git a/api/gateway_chat.py b/api/gateway_chat.py index 1daa497f..a846704d 100644 --- a/api/gateway_chat.py +++ b/api/gateway_chat.py @@ -979,9 +979,19 @@ def _sidecars_with_active_stream(session_dir) -> list[str]: return ids -def _resume_gateway_run_for_session(session) -> bool: +def _gateway_endpoint_for_profile(profile_name) -> tuple[str, str]: + """URL and key of the session's own profile, root included, never the process-active profile.""" from api import profiles as _profiles - from api.config import create_stream_channel, get_config, register_session_writeback_owner, register_stream_owner + from api.config import get_config_for_profile_home + + home = _profiles.get_hermes_home_for_profile(str(profile_name or "").strip()) + environ = {k: v for k, v in os.environ.items() if k not in _profiles._loaded_profile_env_keys} + environ.update(_profiles.filter_runtime_env_for_gateway_parity(_profiles.get_profile_runtime_env(home))) + return _gateway_base_url(get_config_for_profile_home(home), environ), _gateway_api_key(environ) + + +def _resume_gateway_run_for_session(session) -> bool: + from api.config import create_stream_channel, register_session_writeback_owner, register_stream_owner run = (session.gateway_run if session is not None else None) or {} stream_id = str(run.get("stream_id") or "") @@ -989,9 +999,7 @@ def _resume_gateway_run_for_session(session) -> bool: 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. - with _profiles.profile_scope_for_detached_worker(session.profile, "gateway reattach", logger_override=logger): - endpoint = (_gateway_base_url(get_config()), _gateway_api_key()) + endpoint = _gateway_endpoint_for_profile(session.profile) with STREAMS_LOCK: if stream_id in STREAMS: return False @@ -1082,6 +1090,19 @@ def _settle_gateway_terminal_error(session_id, stream_id, workspace, model, mode return error_payload +def _settle_gateway_cancelled_turn(session_id, stream_id) -> None: + """Keep the prompt and a cancel marker when the run ended cancelled and Stop did not already settle it.""" + from api.streaming import _persist_cancelled_turn + + with _get_session_agent_lock(session_id): + session = get_session(session_id) + if not _stream_writeback_is_current(session, stream_id): + return + _persist_cancelled_turn(session, message="Cancelled by gateway") + session.gateway_run = None + session.save() + + def _stream_writeback_is_current(session: Any, stream_id: str) -> bool: return bool(stream_id and getattr(session, "active_stream_id", None) == stream_id) @@ -1318,6 +1339,7 @@ def _run_gateway_chat_streaming( put_gateway_event("apperror", error_payload) return if final_text is None: + _settle_gateway_cancelled_turn(session_id, stream_id) return else: # Legacy gateway path: emit unsupported approval notice once per session, diff --git a/tests/test_gateway_run_reattach_after_restart.py b/tests/test_gateway_run_reattach_after_restart.py index 7845ab53..4fef0461 100644 --- a/tests/test_gateway_run_reattach_after_restart.py +++ b/tests/test_gateway_run_reattach_after_restart.py @@ -1,5 +1,4 @@ """A gateway-owned run must survive a WebUI restart instead of being marked interrupted.""" -import contextlib from collections import OrderedDict import io import json @@ -7,7 +6,6 @@ import os import threading import urllib.error from email.message import Message -from unittest import mock import pytest @@ -266,36 +264,40 @@ def test_reattach_skips_idle_sidecars_without_parsing_them(isolated_sessions, mo @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.""" +@pytest.mark.parametrize("session_profile, process_profile", [("default", "work"), ("work", "default")]) +def test_reattach_resolves_gateway_from_the_session_profile( + isolated_sessions, tmp_path, monkeypatch, default_backend, session_profile, process_profile, +): + """Poll the session profile's own gateway, root included, whatever profile the process restarted under.""" + root = tmp_path / "hermes" + for name, home in (("default", root), ("work", root / "profiles" / "work")): + home.mkdir(parents=True, exist_ok=True) + (home / ".env").write_text( + f"HERMES_WEBUI_GATEWAY_BASE_URL=http://{name}-gateway:8642\nHERMES_WEBUI_GATEWAY_API_KEY={name}-key\n" + ) + monkeypatch.setattr(profiles, "_DEFAULT_HERMES_HOME", root) + monkeypatch.setattr(profiles, "_active_profile", process_profile) + monkeypatch.setattr(profiles, "_loaded_profile_env_keys", set()) + monkeypatch.delenv("HERMES_WEBUI_GATEWAY_BASE_URL") + monkeypatch.delenv("HERMES_WEBUI_GATEWAY_API_KEY", raising=False) + profiles._reload_dotenv(root if process_profile == "default" else root / "profiles" / "work") 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" + s.profile = session_profile 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) - with mock.patch.dict(os.environ, { - "HERMES_WEBUI_GATEWAY_BASE_URL": f"http://{profile_name}-gateway:8642", - "HERMES_WEBUI_GATEWAY_API_KEY": f"{profile_name}-key", - }): - yield - - monkeypatch.setattr(profiles, "profile_scope_for_detached_worker", fake_scope) - monkeypatch.setenv("HERMES_WEBUI_GATEWAY_API_KEY", "default-key") seen = _poll_until_completed(monkeypatch) + assert os.environ["HERMES_WEBUI_GATEWAY_API_KEY"] == f"{process_profile}-key" assert gateway_chat.resume_gateway_runs_after_restart() == [sid] _wait_for_reattach_threads() - assert scopes == ["work"] - assert {(b, k, r) for b, k, r in seen} == {("http://work-gateway:8642", "work-key", "run_survivor")} + assert {(b, k, r) for b, k, r in seen} == { + (f"http://{session_profile}-gateway:8642", f"{session_profile}-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 @@ -398,3 +400,33 @@ def test_failed_recovery_checkpoint_never_admits_the_run(isolated_sessions, monk gateway_chat._run_gateway_chat_streaming(s.session_id, "hi", "test-model", "/tmp", stream_id, []) assert stream_id not in STREAMS + + +@pytest.mark.parametrize("stopped_by", ["user", "gateway"]) +def test_cancelled_replayed_admission_keeps_the_prompt_and_is_not_reattached(isolated_sessions, monkeypatch, stopped_by): + sid, stream_id = _orphaned_gateway_turn(run_id="") + s = models.Session.load(sid) + s.gateway_run["request"] = {"input": "long task"} + s.save(touch_updated_at=False) + models.SESSIONS.clear() + monkeypatch.setattr(gateway_chat, "_admit_gateway_run", lambda *a, **k: "run_replayed") + state = "running" if stopped_by == "user" else "cancelled" + monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", lambda b, k, r: {"run_id": r, "status": state}) + monkeypatch.setattr(gateway_chat, "stop_gateway_run", lambda run_id: True) + + assert gateway_chat.resume_gateway_runs_after_restart() == [sid] + if stopped_by == "user": + assert gateway_chat.wait_for_gateway_run_id(stream_id, 5.0) == (True, "run_replayed") + assert streaming.cancel_stream(stream_id) + _wait_for_reattach_threads() + + saved = _saved(sid) + assert [(m["role"], m["content"]) for m in saved["messages"][:3]] == [ + ("user", "earlier"), ("assistant", "earlier reply"), ("user", "long task"), + ] + assert len(saved["messages"]) == 4 and saved["messages"][3]["_error"] is True + assert "Task cancelled" in saved["messages"][3]["content"] + assert saved["active_stream_id"] is None and saved["gateway_run"] is None + assert saved["pending_user_message"] is None + models.SESSIONS.clear() + assert gateway_chat.resume_gateway_runs_after_restart() == []