fix(gateway): reattach runs-API turns after a WebUI restart instead of marking them interrupted
A gateway-backed turn on the runs API executes in the Hermes Gateway, but the
run id only lived in process memory, so a WebUI restart orphaned the turn and
stale-pending repair appended 'Response interrupted' while the run finished
on the gateway.
- persist session.gateway_run {run_id, stream_id} as soon as POST /v1/runs returns
- send Idempotency-Key webui-<stream_id> so the gateway keeps a durable run record
- on startup (after bind, before serving) re-register the stream and follow
GET /v1/runs/{id} until terminal; write back the real output, surface a
parked approval once, keep Stop working, clear the record on every settlement
This commit is contained in:
1 parent
2ba3e40210
commit
7ba246bfc0
5 files changed
+500
-28
No files matched your search
+228
-28
@@ -564,12 +564,32 @@ def _gateway_runs_approval_event(payload: dict) -> dict | None:
|
||||
}
|
||||
|
||||
|
||||
def _relay_gateway_run_approval(session_id, run_id, payload, base_url, api_key, *, put_gateway_event) -> None:
|
||||
"""Auto-approve or surface one runs-API approval request as a WebUI approval card."""
|
||||
approval_data = _gateway_runs_approval_event(payload)
|
||||
if not approval_data:
|
||||
return
|
||||
approval_data["run_id"] = run_id
|
||||
from api.config import gateway_supports_approval_identity_v1
|
||||
identity_v1 = bool(approval_data.get("_gateway_raw_approval_id_present")) and gateway_supports_approval_identity_v1(base_url, api_key)
|
||||
approval_data["_gateway_agent_identity_v1"] = identity_v1
|
||||
auto_approved, head, total = _settle_gateway_run_approval(
|
||||
session_id,
|
||||
approval_data,
|
||||
base_url,
|
||||
api_key,
|
||||
)
|
||||
if not auto_approved:
|
||||
put_gateway_event("approval", {**(head or approval_data), "pending_count": total})
|
||||
|
||||
|
||||
def _run_gateway_runs_api_streaming(
|
||||
session_id, msg_text, model, workspace, stream_id,
|
||||
base_url, api_key, prefill_messages, body_extras,
|
||||
*, put_gateway_event, cancel_event,
|
||||
attachments=None, cfg=None, session=None,
|
||||
active_provider: str = "",
|
||||
on_run_id=None,
|
||||
):
|
||||
"""Submit via POST /v1/runs and relay SSE events including approval."""
|
||||
try:
|
||||
@@ -582,6 +602,8 @@ def _run_gateway_runs_api_streaming(
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
headers["X-Hermes-Session-Key"] = f"webui:{session_id}"
|
||||
# Durable run record on the gateway: GET /v1/runs/{id} keeps answering after either side restarts.
|
||||
post_headers = {**headers, "Idempotency-Key": f"webui-{stream_id}"}
|
||||
message_content: Any = str(msg_text or "")
|
||||
if attachments:
|
||||
try:
|
||||
@@ -637,7 +659,7 @@ def _run_gateway_runs_api_streaming(
|
||||
req = urllib.request.Request(
|
||||
url_runs,
|
||||
data=json.dumps(run_body).encode("utf-8"),
|
||||
headers=headers,
|
||||
headers=post_headers,
|
||||
method="POST",
|
||||
)
|
||||
update_active_run(stream_id, phase="gateway-request")
|
||||
@@ -652,6 +674,8 @@ def _run_gateway_runs_api_streaming(
|
||||
|
||||
usage: dict = {}
|
||||
_publish_gateway_run_id(stream_id, run_id)
|
||||
if on_run_id is not None:
|
||||
on_run_id(run_id)
|
||||
|
||||
url_events = f"{base_url.rstrip('/')}/v1/runs/{run_id}/events"
|
||||
headers_sse = dict(headers)
|
||||
@@ -682,22 +706,10 @@ def _run_gateway_runs_api_streaming(
|
||||
continue
|
||||
payload_event = str(payload.get("event") or payload.get("type") or sse_event).strip() or "message"
|
||||
if payload_event == "approval.request":
|
||||
approval_data = _gateway_runs_approval_event(payload)
|
||||
if approval_data:
|
||||
approval_data["run_id"] = run_id
|
||||
from api.config import gateway_supports_approval_identity_v1
|
||||
identity_v1 = bool(approval_data.get("_gateway_raw_approval_id_present")) and gateway_supports_approval_identity_v1(base_url, api_key)
|
||||
approval_data["_gateway_agent_identity_v1"] = identity_v1
|
||||
auto_approved, head, total = _settle_gateway_run_approval(
|
||||
session_id,
|
||||
approval_data,
|
||||
base_url,
|
||||
api_key,
|
||||
)
|
||||
if auto_approved:
|
||||
sse_event = "message"
|
||||
continue
|
||||
put_gateway_event("approval", {**(head or approval_data), "pending_count": total})
|
||||
_relay_gateway_run_approval(
|
||||
session_id, run_id, payload, base_url, api_key,
|
||||
put_gateway_event=put_gateway_event,
|
||||
)
|
||||
sse_event = "message"
|
||||
continue
|
||||
if payload_event in {"tool.started", "tool.completed", "reasoning.available"}:
|
||||
@@ -823,6 +835,172 @@ def stop_gateway_run(run_id: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
_GATEWAY_RUN_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "interrupted"})
|
||||
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
|
||||
|
||||
|
||||
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."""
|
||||
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}
|
||||
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)
|
||||
|
||||
|
||||
def _get_gateway_run_status(base_url: str, api_key: str, run_id: str) -> dict:
|
||||
headers = {"Accept": "application/json"}
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
req = urllib.request.Request(
|
||||
f"{base_url.rstrip('/')}/v1/runs/{urllib.parse.quote(run_id, safe='')}",
|
||||
headers=headers,
|
||||
method="GET",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||||
payload = json.loads(resp.read(1024 * 1024) or b"{}")
|
||||
return payload if isinstance(payload, dict) else {}
|
||||
|
||||
|
||||
def _await_gateway_run_result(
|
||||
session_id, stream_id, run_id, base_url, api_key,
|
||||
*, put_gateway_event, cancel_event,
|
||||
):
|
||||
"""Follow a gateway run started by a previous WebUI process until it settles.
|
||||
|
||||
Returns the same ``(final_text | None, usage)`` contract as
|
||||
``_run_gateway_runs_api_streaming``; token/tool events emitted before the
|
||||
restart are not replayed, only the terminal result is written back.
|
||||
"""
|
||||
from api.route_approvals import settle_gateway_pending_run
|
||||
|
||||
_publish_gateway_run_id(stream_id, run_id)
|
||||
update_active_run(stream_id, phase="gateway-reattached")
|
||||
surfaced_approval_ids: set[str] = set()
|
||||
failures = 0
|
||||
while True:
|
||||
if cancel_event.is_set():
|
||||
put_gateway_event("cancel", {"message": "Cancelled by user"})
|
||||
return None, {}
|
||||
try:
|
||||
status = _get_gateway_run_status(base_url, api_key, run_id)
|
||||
failures = 0
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 404:
|
||||
raise RuntimeError("Gateway no longer has this run; its result could not be recovered after the WebUI restart.") from exc
|
||||
failures += 1
|
||||
status = None
|
||||
except (urllib.error.URLError, OSError, ValueError):
|
||||
failures += 1
|
||||
status = None
|
||||
if status is None:
|
||||
if failures >= GATEWAY_REATTACH_MAX_POLL_FAILURES:
|
||||
raise RuntimeError("Gateway became unreachable while waiting for a run that outlived the WebUI restart.")
|
||||
else:
|
||||
state = str(status.get("status") or "").strip().lower()
|
||||
if state == "waiting_for_approval" and isinstance(status.get("approval"), dict):
|
||||
approval = status["approval"]
|
||||
approval_key = str(approval.get("approval_id") or approval.get("id") or approval.get("timestamp") or "")
|
||||
if approval_key not in surfaced_approval_ids:
|
||||
surfaced_approval_ids.add(approval_key)
|
||||
_relay_gateway_run_approval(
|
||||
session_id, run_id, approval, base_url, api_key,
|
||||
put_gateway_event=put_gateway_event,
|
||||
)
|
||||
if state in _GATEWAY_RUN_TERMINAL_STATUSES:
|
||||
settle_gateway_pending_run(
|
||||
session_id,
|
||||
run_id,
|
||||
reason=f"Gateway run {state} before approval resolution",
|
||||
)
|
||||
if state == "cancelled":
|
||||
put_gateway_event("cancel", {"message": "Cancelled by gateway"})
|
||||
return None, {}
|
||||
if state != "completed":
|
||||
raise RuntimeError(str(status.get("error") or f"Gateway run {state}"))
|
||||
usage = {k: v for k, v in _gateway_stream_usage(status).items() if v}
|
||||
output = str(status.get("output") or "")
|
||||
if output and stream_id in STREAM_PARTIAL_TEXT:
|
||||
STREAM_PARTIAL_TEXT[stream_id] = output
|
||||
return output, usage
|
||||
cancel_event.wait(GATEWAY_REATTACH_POLL_INTERVAL)
|
||||
|
||||
|
||||
def resume_gateway_runs_after_restart() -> list[str]:
|
||||
"""Reattach pending turns whose gateway run outlived the previous WebUI process.
|
||||
|
||||
Must run before the server accepts requests: registering the stream first
|
||||
keeps stale-pending repair from marking these still-running turns interrupted.
|
||||
Never raises; startup must not depend on the gateway being reachable.
|
||||
"""
|
||||
from api import models as _models
|
||||
from api.config import get_config
|
||||
|
||||
try:
|
||||
if not webui_gateway_chat_enabled(get_config()):
|
||||
return []
|
||||
entries = json.loads(_models.SESSION_INDEX_FILE.read_bytes())
|
||||
except Exception:
|
||||
logger.debug("gateway reattach: no readable session index", 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 "")
|
||||
try:
|
||||
if _resume_gateway_run_for_session(_models.Session.load_metadata_only(sid)):
|
||||
resumed.append(sid)
|
||||
except Exception:
|
||||
logger.warning("gateway reattach failed for session %s", sid, exc_info=True)
|
||||
if resumed:
|
||||
logger.info("Reattached %d gateway run(s) that outlived the previous WebUI process", len(resumed))
|
||||
return resumed
|
||||
|
||||
|
||||
def _resume_gateway_run_for_session(session) -> bool:
|
||||
from api.config import create_stream_channel, register_session_writeback_owner, register_stream_owner
|
||||
|
||||
run = getattr(session, "gateway_run", None) if session is not None else None
|
||||
stream_id = str(getattr(session, "active_stream_id", None) or "")
|
||||
run_id = str((run or {}).get("run_id") or "").strip()
|
||||
if not stream_id or not run_id or run.get("stream_id") != stream_id:
|
||||
return False
|
||||
sid = session.session_id
|
||||
with STREAMS_LOCK:
|
||||
if stream_id in STREAMS:
|
||||
return False
|
||||
STREAMS[stream_id] = create_stream_channel()
|
||||
register_stream_owner(stream_id, sid)
|
||||
register_session_writeback_owner(sid, stream_id)
|
||||
_mark_gateway_run_starting(stream_id)
|
||||
threading.Thread(
|
||||
target=_run_gateway_chat_streaming,
|
||||
args=(
|
||||
sid,
|
||||
session.pending_user_message or "",
|
||||
session.model,
|
||||
session.workspace,
|
||||
stream_id,
|
||||
list(session.pending_attachments or []),
|
||||
),
|
||||
kwargs={
|
||||
"model_provider": session.model_provider,
|
||||
"goal_related": bool(run.get("goal_related")),
|
||||
"regeneration": bool(run.get("regeneration")),
|
||||
"reattach_run_id": run_id,
|
||||
},
|
||||
name=f"gateway-reattach-{stream_id[:12]}",
|
||||
daemon=True,
|
||||
).start()
|
||||
return True
|
||||
|
||||
def _settle_gateway_terminal_error(session_id, stream_id, workspace, model, model_provider, terminal_error):
|
||||
from api.streaming import (
|
||||
_classify_provider_error,
|
||||
@@ -846,6 +1024,7 @@ def _settle_gateway_terminal_error(session_id, stream_id, workspace, model, mode
|
||||
turn_duration = _terminal_turn_duration(session)
|
||||
_materialize_pending_user_turn_before_error(session)
|
||||
session.active_stream_id = None
|
||||
session.gateway_run = None
|
||||
session.pending_user_message = None
|
||||
session.pending_attachments = []
|
||||
session.pending_started_at = None
|
||||
@@ -895,8 +1074,13 @@ def _stream_writeback_is_current(session: Any, stream_id: str) -> bool:
|
||||
|
||||
def _clear_gateway_pending_state(session: Any, stream_id: str) -> None:
|
||||
if not _stream_writeback_is_current(session, stream_id):
|
||||
# Cancel clears active_stream_id eagerly; still drop this stream's run record.
|
||||
if session is not None and (getattr(session, "gateway_run", None) or {}).get("stream_id") == stream_id:
|
||||
session.gateway_run = None
|
||||
session.save(touch_updated_at=False)
|
||||
return
|
||||
session.active_stream_id = None
|
||||
session.gateway_run = None
|
||||
session.pending_user_message = None
|
||||
session.pending_attachments = None
|
||||
session.pending_started_at = None
|
||||
@@ -926,6 +1110,7 @@ def _run_gateway_chat_streaming(
|
||||
model_provider=None,
|
||||
goal_related=False,
|
||||
regeneration=False,
|
||||
reattach_run_id=None,
|
||||
):
|
||||
"""Bridge a WebUI chat turn through Hermes Gateway's API server.
|
||||
|
||||
@@ -1024,7 +1209,9 @@ def _run_gateway_chat_streaming(
|
||||
except Exception:
|
||||
_gw_overrides = {}
|
||||
_runs_api_enabled = _gateway_use_runs_api_enabled(cfg)
|
||||
_use_runs_api = _runs_api_enabled and gateway_supports_approval(base_url, api_key)
|
||||
_use_runs_api = bool(reattach_run_id) 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
|
||||
@@ -1075,16 +1262,28 @@ def _run_gateway_chat_streaming(
|
||||
if _gw_overrides.get("service_tier"):
|
||||
body_extras["service_tier"] = _gw_overrides["service_tier"]
|
||||
try:
|
||||
final_text, usage = _run_gateway_runs_api_streaming(
|
||||
session_id, msg_text, model, workspace, stream_id,
|
||||
base_url, api_key, prefill_messages, body_extras,
|
||||
put_gateway_event=put_gateway_event,
|
||||
cancel_event=cancel_event,
|
||||
attachments=attachments,
|
||||
cfg=cfg,
|
||||
session=s,
|
||||
active_provider=(model_provider or ""),
|
||||
)
|
||||
if reattach_run_id:
|
||||
final_text, usage = _await_gateway_run_result(
|
||||
session_id, stream_id, reattach_run_id, base_url, api_key,
|
||||
put_gateway_event=put_gateway_event,
|
||||
cancel_event=cancel_event,
|
||||
)
|
||||
else:
|
||||
final_text, usage = _run_gateway_runs_api_streaming(
|
||||
session_id, msg_text, model, workspace, stream_id,
|
||||
base_url, api_key, prefill_messages, body_extras,
|
||||
put_gateway_event=put_gateway_event,
|
||||
cancel_event=cancel_event,
|
||||
attachments=attachments,
|
||||
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),
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
error_payload = _settle_gateway_terminal_error(
|
||||
session_id,
|
||||
@@ -1371,6 +1570,7 @@ def _run_gateway_chat_streaming(
|
||||
display = display[:-1]
|
||||
s.messages = display + [user_msg, assistant_msg]
|
||||
s.active_stream_id = None
|
||||
s.gateway_run = None
|
||||
s.pending_user_message = None
|
||||
s.pending_attachments = None
|
||||
s.pending_started_at = None
|
||||
|
||||
@@ -1319,6 +1319,7 @@ class Session:
|
||||
process_wakeup_pause=None,
|
||||
share_token=None,
|
||||
share_created_at=None,
|
||||
gateway_run=None,
|
||||
**kwargs):
|
||||
self.session_id = session_id or uuid.uuid4().hex[:12]
|
||||
self.title = title
|
||||
@@ -1424,6 +1425,8 @@ class Session:
|
||||
self.process_wakeup_pause = process_wakeup_pause if isinstance(process_wakeup_pause, dict) else {}
|
||||
self.share_token = str(share_token).strip() if share_token else None
|
||||
self.share_created_at = share_created_at
|
||||
# Gateway /v1/runs id for the turn named by stream_id; lets a restarted WebUI reattach.
|
||||
self.gateway_run = gateway_run if isinstance(gateway_run, dict) else None
|
||||
# #5854: a compact fingerprint of anchor_activity_scenes ({scene_key:
|
||||
# updated_at}) persisted BEFORE the messages array so the sidebar-poll
|
||||
# freshness check can compare scene freshness without parsing the full
|
||||
@@ -1496,6 +1499,7 @@ class Session:
|
||||
'enabled_toolsets', 'composer_draft',
|
||||
'process_wakeup_pause',
|
||||
'share_token', 'share_created_at',
|
||||
'gateway_run',
|
||||
]
|
||||
meta = {k: getattr(self, k, None) for k in METADATA_FIELDS}
|
||||
# #5854: message_count and a compact anchor-scene fingerprint go in the
|
||||
|
||||
@@ -119,6 +119,18 @@ HERMES_WEBUI_GATEWAY_USE_RUNS_API=true \
|
||||
|
||||
Use this when the connected gateway advertises approval support and you want tool approval cards to appear in WebUI. Without `HERMES_WEBUI_GATEWAY_USE_RUNS_API=true`, gateway chat stays on the legacy chat-completions transport and approval-capable commands can remain pending in the agent without a WebUI approval card.
|
||||
|
||||
On the runs API path the turn is executed by the Gateway, so restarting WebUI does
|
||||
not stop it. WebUI stores the Gateway `run_id` on the pending turn (and submits it
|
||||
with an `Idempotency-Key` so the Gateway keeps a durable run record). On startup,
|
||||
WebUI reattaches to every such run by polling `GET /v1/runs/{run_id}` until it
|
||||
settles, then writes the real final answer into the session instead of a
|
||||
"Response interrupted" marker. Stop still cancels a reattached run, and a pending
|
||||
approval is shown again. Token-by-token output from before the restart is not
|
||||
replayed; the reattached turn shows only the final answer. If the Gateway no longer knows the run
|
||||
(for example, it restarted too and the run was interrupted), the turn ends with an error message
|
||||
instead. The legacy chat-completions transport cannot reattach: its turn ends when the
|
||||
WebUI process that holds the HTTP stream exits.
|
||||
|
||||
When YOLO is enabled for a gateway-backed browser session, WebUI approves every
|
||||
approval already parked for that session: Runs API prompts are relayed by their
|
||||
exact `run_id` and mirror token, and local/no-run waiters are all released. It
|
||||
|
||||
@@ -669,6 +669,7 @@ def main() -> None:
|
||||
|
||||
_abort_if_already_serving(HOST, PORT)
|
||||
httpd = QuietHTTPServer((HOST, PORT), Handler)
|
||||
from api.gateway_chat import resume_gateway_runs_after_restart; resume_gateway_runs_after_restart() # bound, not yet serving
|
||||
|
||||
from api.config import TLS_ENABLED, TLS_CERT, TLS_KEY
|
||||
scheme = 'https' if TLS_ENABLED else 'http'
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
"""A gateway-owned run must survive a WebUI restart instead of being marked interrupted."""
|
||||
from collections import OrderedDict
|
||||
import io
|
||||
import json
|
||||
import threading
|
||||
import urllib.error
|
||||
from email.message import Message
|
||||
|
||||
import pytest
|
||||
|
||||
import api.gateway_chat as gateway_chat
|
||||
import api.models as models
|
||||
import api.streaming as streaming
|
||||
from api.config import ACTIVE_RUNS, STREAMS, STREAMS_LOCK, create_stream_channel
|
||||
from api.models import new_session
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_sessions(tmp_path, monkeypatch):
|
||||
session_dir = tmp_path / "sessions"
|
||||
session_dir.mkdir()
|
||||
monkeypatch.setattr(models, "SESSION_DIR", session_dir)
|
||||
monkeypatch.setattr(models, "SESSION_INDEX_FILE", session_dir / "_index.json")
|
||||
monkeypatch.setattr(models, "SESSIONS", OrderedDict())
|
||||
monkeypatch.setenv("HERMES_WEBUI_CHAT_BACKEND", "gateway")
|
||||
monkeypatch.setenv("HERMES_WEBUI_GATEWAY_USE_RUNS_API", "1")
|
||||
monkeypatch.setenv("HERMES_WEBUI_GATEWAY_BASE_URL", "http://gateway.local")
|
||||
monkeypatch.setattr(gateway_chat, "GATEWAY_REATTACH_POLL_INTERVAL", 0.01)
|
||||
monkeypatch.setattr(gateway_chat, "_gateway_reasoning_effort_for_request", lambda *a, **k: None)
|
||||
monkeypatch.setattr(streaming, "_load_webui_prefill_context", lambda cfg: {
|
||||
"status": "not_configured", "source": "none", "label": "", "message_count": 0, "messages": [],
|
||||
})
|
||||
monkeypatch.setattr(streaming, "_prefill_messages_with_webui_context", lambda ctx, cfg: [])
|
||||
return session_dir
|
||||
|
||||
|
||||
def _orphaned_gateway_turn(run_id="run_survivor", stream_id="stream-before-restart", started_at=1.0):
|
||||
"""Persist the sidecar exactly as a WebUI process leaves it when killed mid-run."""
|
||||
s = new_session()
|
||||
s.messages = [
|
||||
{"role": "user", "content": "earlier", "timestamp": 0.5},
|
||||
{"role": "assistant", "content": "earlier reply", "timestamp": 0.6},
|
||||
]
|
||||
s.active_stream_id = stream_id
|
||||
s.pending_user_message = "long task"
|
||||
s.pending_attachments = []
|
||||
s.pending_started_at = started_at
|
||||
s.pending_user_source = "webui"
|
||||
s.gateway_run = {"run_id": run_id, "stream_id": stream_id, "regeneration": False, "goal_related": False}
|
||||
s.save()
|
||||
# Simulate the new process: nothing of the old run is in memory.
|
||||
models.SESSIONS.clear()
|
||||
with STREAMS_LOCK:
|
||||
STREAMS.pop(stream_id, None)
|
||||
ACTIVE_RUNS.pop(stream_id, None)
|
||||
return s.session_id, stream_id
|
||||
|
||||
|
||||
def _wait_for_reattach_threads(timeout=10.0):
|
||||
for thread in threading.enumerate():
|
||||
if thread.name.startswith("gateway-reattach-"):
|
||||
thread.join(timeout)
|
||||
assert not thread.is_alive(), "reattach worker did not settle"
|
||||
|
||||
|
||||
def test_runs_api_start_sends_idempotency_key_and_persists_run_id(isolated_sessions, monkeypatch):
|
||||
s = new_session()
|
||||
stream_id = "stream-live"
|
||||
s.active_stream_id = stream_id
|
||||
s.pending_user_message = "hi"
|
||||
s.pending_attachments = []
|
||||
s.pending_started_at = 123.0
|
||||
s.save()
|
||||
captured = {}
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, body=None, lines=None):
|
||||
self._body, self._lines = body, lines or []
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def read(self, *_a):
|
||||
return self._body
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self._lines)
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
if req.get_method() == "POST":
|
||||
captured["post_headers"] = dict(req.header_items())
|
||||
return _Resp(body=b'{"run_id":"run_live"}')
|
||||
# The run id is durable before the first event is relayed.
|
||||
captured["persisted_at_events"] = json.loads(
|
||||
(isolated_sessions / f"{s.session_id}.json").read_text()
|
||||
).get("gateway_run")
|
||||
return _Resp(lines=[
|
||||
b'data: {"event":"message.delta","delta":"done"}\n',
|
||||
b'data: {"event":"run.completed","output":"done"}\n',
|
||||
b"data: [DONE]\n",
|
||||
])
|
||||
|
||||
monkeypatch.setattr(gateway_chat, "gateway_supports_approval", lambda *a, **k: True)
|
||||
monkeypatch.setattr(gateway_chat.urllib.request, "urlopen", fake_urlopen)
|
||||
with STREAMS_LOCK:
|
||||
STREAMS[stream_id] = create_stream_channel()
|
||||
|
||||
gateway_chat._run_gateway_chat_streaming(s.session_id, "hi", "test-model", "/tmp", stream_id, [])
|
||||
|
||||
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
|
||||
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
|
||||
assert saved["messages"][-1]["content"] == "done"
|
||||
|
||||
|
||||
def test_restart_reattaches_and_writes_back_real_answer(isolated_sessions, monkeypatch):
|
||||
sid, stream_id = _orphaned_gateway_turn()
|
||||
release = threading.Event()
|
||||
polls = []
|
||||
|
||||
def fake_status(base_url, api_key, run_id):
|
||||
polls.append(run_id)
|
||||
if not release.is_set():
|
||||
return {"run_id": run_id, "status": "running"}
|
||||
return {
|
||||
"run_id": run_id,
|
||||
"status": "completed",
|
||||
"output": "finished after the restart",
|
||||
"usage": {"input_tokens": 7, "output_tokens": 3},
|
||||
}
|
||||
|
||||
monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", fake_status)
|
||||
|
||||
assert gateway_chat.resume_gateway_runs_after_restart() == [sid]
|
||||
# While the run is still going, the turn is live: no stale-pending repair, reconnect works.
|
||||
assert stream_id in STREAMS
|
||||
assert gateway_chat.wait_for_gateway_run_id(stream_id, 5.0) == (True, "run_survivor")
|
||||
live = models.get_session(sid)
|
||||
assert live.active_stream_id == stream_id
|
||||
assert not any(m.get("_error") for m in live.messages)
|
||||
|
||||
release.set()
|
||||
_wait_for_reattach_threads()
|
||||
|
||||
saved = json.loads((isolated_sessions / f"{sid}.json").read_text())
|
||||
assert [m["role"] for m in saved["messages"]] == ["user", "assistant", "user", "assistant"]
|
||||
assert saved["messages"][2]["content"] == "long task"
|
||||
assert saved["messages"][3]["content"] == "finished after the restart"
|
||||
assert not any(m.get("_error") for m in saved["messages"])
|
||||
assert saved["active_stream_id"] is None
|
||||
assert saved["pending_user_message"] is None
|
||||
assert saved["gateway_run"] is None
|
||||
assert set(polls) == {"run_survivor"}
|
||||
assert stream_id not in STREAMS
|
||||
|
||||
|
||||
def test_reattach_reports_run_the_gateway_no_longer_knows(isolated_sessions, monkeypatch):
|
||||
sid, _stream_id = _orphaned_gateway_turn(run_id="run_gone")
|
||||
|
||||
def fake_status(base_url, api_key, run_id):
|
||||
raise urllib.error.HTTPError("http://gateway.local", 404, "not found", Message(), io.BytesIO(b"{}"))
|
||||
|
||||
monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", fake_status)
|
||||
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"][-2]["content"] == "long task"
|
||||
assert saved["messages"][-1]["_error"] is True
|
||||
assert "could not be recovered after the WebUI restart" in saved["messages"][-1]["content"]
|
||||
assert saved["active_stream_id"] is None
|
||||
assert saved["gateway_run"] is None
|
||||
|
||||
|
||||
def test_reattach_surfaces_pending_approval_once(isolated_sessions, monkeypatch):
|
||||
sid, stream_id = _orphaned_gateway_turn(run_id="run_parked")
|
||||
relayed = []
|
||||
calls = {"n": 0}
|
||||
|
||||
def fake_status(base_url, api_key, run_id):
|
||||
calls["n"] += 1
|
||||
if calls["n"] < 4:
|
||||
return {"run_id": run_id, "status": "waiting_for_approval", "approval": {
|
||||
"event": "approval.request", "approval_id": "appr-1", "command": "rm -rf build", "description": "delete",
|
||||
}}
|
||||
return {"run_id": run_id, "status": "completed", "output": "approved and done"}
|
||||
|
||||
monkeypatch.setattr(gateway_chat, "_get_gateway_run_status", fake_status)
|
||||
monkeypatch.setattr(
|
||||
gateway_chat, "_relay_gateway_run_approval",
|
||||
lambda session_id, run_id, payload, *a, **k: relayed.append((session_id, run_id, payload["approval_id"])),
|
||||
)
|
||||
gateway_chat.resume_gateway_runs_after_restart()
|
||||
_wait_for_reattach_threads()
|
||||
|
||||
assert relayed == [(sid, "run_parked", "appr-1")]
|
||||
saved = json.loads((isolated_sessions / f"{sid}.json").read_text())
|
||||
assert saved["messages"][-1]["content"] == "approved and done"
|
||||
|
||||
|
||||
def test_reattach_is_default_off_without_gateway_backend(isolated_sessions, monkeypatch):
|
||||
_orphaned_gateway_turn()
|
||||
monkeypatch.delenv("HERMES_WEBUI_CHAT_BACKEND")
|
||||
monkeypatch.setattr(
|
||||
gateway_chat, "_get_gateway_run_status",
|
||||
lambda *a, **k: pytest.fail("legacy backend must not poll the gateway"),
|
||||
)
|
||||
assert gateway_chat.resume_gateway_runs_after_restart() == []
|
||||
|
||||
|
||||
def test_sidecar_without_gateway_run_is_left_to_stale_pending_repair(isolated_sessions, monkeypatch):
|
||||
sid, stream_id = _orphaned_gateway_turn()
|
||||
s = models.Session.load(sid)
|
||||
assert s is not None
|
||||
s.gateway_run = None
|
||||
s.save()
|
||||
monkeypatch.setattr(
|
||||
gateway_chat, "_get_gateway_run_status",
|
||||
lambda *a, **k: pytest.fail("no recorded run id: nothing to reattach"),
|
||||
)
|
||||
assert gateway_chat.resume_gateway_runs_after_restart() == []
|
||||
assert stream_id not in STREAMS
|
||||
|
||||
|
||||
def test_cancel_still_stops_a_reattached_run(isolated_sessions, monkeypatch):
|
||||
sid, stream_id = _orphaned_gateway_turn(run_id="run_to_stop")
|
||||
stopped = []
|
||||
monkeypatch.setattr(
|
||||
gateway_chat, "_get_gateway_run_status",
|
||||
lambda base_url, api_key, run_id: {"run_id": run_id, "status": "running"},
|
||||
)
|
||||
monkeypatch.setattr(gateway_chat, "stop_gateway_run", lambda run_id: stopped.append(run_id) or True)
|
||||
gateway_chat.resume_gateway_runs_after_restart()
|
||||
|
||||
# Same two steps /api/chat/cancel performs for a gateway-backed stream.
|
||||
_structured, run_id = gateway_chat.wait_for_gateway_run_id(stream_id, 5.0)
|
||||
assert gateway_chat.stop_gateway_run(run_id)
|
||||
for _ in range(500):
|
||||
if stream_id in gateway_chat.CANCEL_FLAGS:
|
||||
break
|
||||
threading.Event().wait(0.01)
|
||||
assert streaming.cancel_stream(stream_id)
|
||||
_wait_for_reattach_threads()
|
||||
|
||||
assert stopped == ["run_to_stop"]
|
||||
saved = json.loads((isolated_sessions / f"{sid}.json").read_text())
|
||||
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"])
|
||||
Reference in new issue
Block a user