DasPoschi/hermes-webui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(#1361): preserve reasoning, tool calls, and partial output on Stop/Cancel (#1375)

Browse Source 
Three distinct data-loss paths fixed:

§A — Reasoning text was accumulated in a thread-local _reasoning_text
inside _run_agent_streaming. cancel_stream() never saw it because it
went out of scope when the thread was interrupted. Now mirrored to a
new shared dict STREAM_REASONING_TEXT keyed by stream_id, populated
in on_reasoning() and the reasoning branch of on_tool(), read in
cancel_stream().

§B — Live tool calls in thread-local _live_tool_calls were similarly
invisible to cancel_stream(). Now mirrored to STREAM_LIVE_TOOL_CALLS
on tool.started + tool.completed.

§C — Reasoning-only streams produced no partial message because the
thinking-block regex strip returned empty string and the `if _stripped:`
guard skipped the append. Now appends the partial message when EITHER
content text, reasoning trace, OR tool calls exist.

Mirrors the existing STREAM_PARTIAL_TEXT pattern from #893 exactly:
same dict creation in _run_agent_streaming, same _live_config fallback
in cancel_stream, same cleanup in _periodic_checkpoint.

8 regression tests in tests/test_issue1361_cancel_data_loss.py
covering all three sections plus tools+text combinations.

Co-authored-by: bergeouss <bergeouss@users.noreply.github.com>
This commit is contained in:
bergeouss
2026-04-30 23:24:29 +00:00
committed by nesquena-hermes
parent 63251ad206
commit c5f4f569d6
3 changed files with 382 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
api/config.py
View File

@@ -2115,6 +2115,8 @@ STREAMS_LOCK = threading.Lock()
CANCEL_FLAGS: dict = {}
AGENT_INSTANCES: dict = {} # stream_id -> AIAgent instance for interrupt propagation
STREAM_PARTIAL_TEXT: dict = {} # stream_id -> partial assistant text accumulated during streaming
STREAM_REASONING_TEXT: dict = {} # stream_id -> reasoning trace accumulated during streaming (#1361 §A)
STREAM_LIVE_TOOL_CALLS: dict = {} # stream_id -> live tool calls accumulated during streaming (#1361 §B)
SERVER_START_TIME = time.time()
# Agent cache: reuse AIAgent across messages in the same WebUI session so that

74
api/streaming.py
View File

@@ -21,6 +21,7 @@ logger = logging.getLogger(__name__)
from api.config import (
STREAMS, STREAMS_LOCK, CANCEL_FLAGS, AGENT_INSTANCES, STREAM_PARTIAL_TEXT,
STREAM_REASONING_TEXT, STREAM_LIVE_TOOL_CALLS,
LOCK, SESSIONS, SESSION_DIR,
_get_session_agent_lock, _set_thread_env, _clear_thread_env,
SESSION_AGENT_LOCKS, SESSION_AGENT_LOCKS_LOCK,
@@ -1373,6 +1374,8 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
with STREAMS_LOCK:
CANCEL_FLAGS[stream_id] = cancel_event
STREAM_PARTIAL_TEXT[stream_id] = '' # start accumulating partial text (#893)
STREAM_REASONING_TEXT[stream_id] = '' # start accumulating reasoning trace (#1361 §A)
STREAM_LIVE_TOOL_CALLS[stream_id] = [] # start accumulating tool calls (#1361 §B)
# Register this stream with the global streaming meter
meter().begin_session(stream_id)
@@ -1575,6 +1578,9 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
if text is None:
return
_reasoning_text += str(text)
# Mirror to shared dict so cancel_stream() can persist it (#1361 §A)
if stream_id in STREAM_REASONING_TEXT:
STREAM_REASONING_TEXT[stream_id] += str(text)
put('reasoning', {'text': str(text)})
# Track reasoning tokens in the meter so TPS reflects all AI output
meter().record_reasoning(stream_id, len(_reasoning_text))
@@ -1607,6 +1613,9 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
reason_text = preview if event_type == 'reasoning.available' else name
if reason_text:
_reasoning_text += str(reason_text)
# Mirror to shared dict so cancel_stream() can persist it (#1361 §A)
if stream_id in STREAM_REASONING_TEXT:
STREAM_REASONING_TEXT[stream_id] += str(reason_text)
put('reasoning', {'text': str(reason_text)})
meter().record_reasoning(stream_id, len(_reasoning_text))
_emit_metering()
@@ -1623,6 +1632,13 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
'name': name,
'args': args if isinstance(args, dict) else {},
})
# Mirror to shared dict so cancel_stream() can persist it (#1361 §B)
if stream_id in STREAM_LIVE_TOOL_CALLS:
STREAM_LIVE_TOOL_CALLS[stream_id].append({
'name': name,
'args': args if isinstance(args, dict) else {},
'done': False,
})
put('tool', {
'event_type': event_type or 'tool.started',
'name': name,
@@ -1651,6 +1667,16 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
live_tc['duration'] = cb_kwargs.get('duration')
live_tc['is_error'] = bool(cb_kwargs.get('is_error', False))
break
# Mirror done state to shared dict (#1361 §B)
if stream_id in STREAM_LIVE_TOOL_CALLS:
for shared_tc in reversed(STREAM_LIVE_TOOL_CALLS[stream_id]):
if shared_tc.get('done'):
continue
if not name or shared_tc.get('name') == name:
shared_tc['done'] = True
shared_tc['duration'] = cb_kwargs.get('duration')
shared_tc['is_error'] = bool(cb_kwargs.get('is_error', False))
break
# Signal the checkpoint thread that new work has completed (Issue #765).
# Each completed tool call is a meaningful unit of progress worth persisting.
_checkpoint_activity[0] += 1
@@ -2443,6 +2469,8 @@ def _run_agent_streaming(session_id, msg_text, model, workspace, stream_id, atta
CANCEL_FLAGS.pop(stream_id, None)
AGENT_INSTANCES.pop(stream_id, None) # Clean up agent instance reference
STREAM_PARTIAL_TEXT.pop(stream_id, None) # Clean up partial text buffer (#893)
STREAM_REASONING_TEXT.pop(stream_id, None) # Clean up reasoning trace (#1361 §A)
STREAM_LIVE_TOOL_CALLS.pop(stream_id, None) # Clean up tool calls (#1361 §B)
# ============================================================
# SECTION: HTTP Request Handler
@@ -2630,6 +2658,17 @@ def cancel_stream(stream_id: str) -> bool:
live_partials = getattr(_live_config, 'STREAM_PARTIAL_TEXT', partial_texts)
if live_partials is not partial_texts:
_cancel_partial_text = live_partials.get(stream_id, '')
# Capture reasoning trace and live tool calls (#1361 §A + §B)
_cancel_reasoning = STREAM_REASONING_TEXT.get(stream_id, '')
if not _cancel_reasoning:
live_reasoning = getattr(_live_config, 'STREAM_REASONING_TEXT', STREAM_REASONING_TEXT)
if live_reasoning is not STREAM_REASONING_TEXT:
_cancel_reasoning = live_reasoning.get(stream_id, '')
_cancel_tool_calls = STREAM_LIVE_TOOL_CALLS.get(stream_id, [])
if not _cancel_tool_calls:
live_tools = getattr(_live_config, 'STREAM_LIVE_TOOL_CALLS', STREAM_LIVE_TOOL_CALLS)
if live_tools is not STREAM_LIVE_TOOL_CALLS:
_cancel_tool_calls = live_tools.get(stream_id, [])
# Session cleanup outside STREAMS_LOCK to preserve lock ordering.
# Acquire the per-session _agent_lock too, mirroring every other session
@@ -2708,7 +2747,13 @@ def cancel_stream(stream_id: str) -> bool:
# content IS kept in the history sent to the agent on the next user
# message, letting the model continue from where it was cut off.
# See the inner comment on the append call below for the rationale.
#
# #1361: Also persist reasoning trace and live tool calls that were
# accumulated in thread-local variables but invisible to the cancel path.
# This prevents paid-token data loss when cancelling mid-reasoning or
# mid-tool-execution.
partial_text = _cancel_partial_text.strip() if _cancel_partial_text else ''
_stripped = ''
if partial_text:
import re as _re
# Strip thinking/reasoning markup from partial content before saving.
@@ -2721,17 +2766,24 @@ def cancel_stream(stream_id: str) -> bool:
_stripped = _re.sub(r'<think(?:ing)?\b[^>]*>.*',
'', _stripped,
flags=_re.DOTALL | _re.IGNORECASE).strip()
if _stripped:
# Mark _partial=True for session/UI identification only.
# Deliberately NOT _error=True — the partial content is real model
# output and should be visible in conversation history so the model
# can continue from it on the next turn (#893).
_cs.messages.append({
'role': 'assistant',
'content': _stripped,
'_partial': True,
'timestamp': int(time.time()),
})
# Determine whether there is anything to preserve beyond just the
# cancel marker. Content text, reasoning trace, or tool calls all
# count (#1361 §C — previously only _stripped was checked, so a
# reasoning-only or tool-only stream produced NO partial message).
_has_reasoning = bool(_cancel_reasoning and _cancel_reasoning.strip())
_has_tools = bool(_cancel_tool_calls)
if _stripped or _has_reasoning or _has_tools:
_partial_msg: dict = {
'role': 'assistant',
'content': _stripped, # may be empty for reasoning/tool-only turns
'_partial': True,
'timestamp': int(time.time()),
}
if _has_reasoning:
_partial_msg['reasoning'] = _cancel_reasoning.strip()
if _has_tools:
_partial_msg['tool_calls'] = list(_cancel_tool_calls)
_cs.messages.append(_partial_msg)
# Cancel marker — flagged _error=True so it is stripped from conversation
# history on the next turn (prevents model from seeing "Task cancelled."
# as a prior assistant reply).

317
tests/test_issue1361_cancel_data_loss.py Normal file
View File

@@ -0,0 +1,317 @@
"""Regression tests for #1361 — Stop/Cancel discards already-streamed content.
Three distinct data-loss paths on cancel:
§A Reasoning text accumulated in a thread-local `_reasoning_text` is never
visible to cancel_stream(), so it's lost on cancel.
§B Live tool calls accumulated in thread-local `_live_tool_calls` are lost
on cancel — only STREAM_PARTIAL_TEXT is captured.
§C When the entire streamed output is reasoning (no visible tokens),
_stripped is empty after regex cleanup, so NO partial assistant message
is appended — only the *Task cancelled.* marker survives.
All three fix the same "tokens-paid-for-data-loss" class of bug.
"""
import pathlib
import queue
import re
import threading
from unittest.mock import Mock, patch
import pytest
import api.config as config
import api.models as models
import api.streaming as streaming
from api.models import Session
from api.streaming import cancel_stream
REPO_ROOT = pathlib.Path(__file__).parent.parent.resolve()
# ── Fixtures ────────────────────────────────────────────────────────────────
@pytest.fixture(autouse=True)
def _isolate_session_dir(tmp_path, monkeypatch):
"""Redirect SESSION_DIR / SESSION_INDEX_FILE to an isolated temp dir."""
session_dir = tmp_path / "sessions"
session_dir.mkdir()
index_file = session_dir / "_index.json"
monkeypatch.setattr(models, "SESSION_DIR", session_dir)
monkeypatch.setattr(models, "SESSION_INDEX_FILE", index_file)
models.SESSIONS.clear()
yield
models.SESSIONS.clear()
@pytest.fixture(autouse=True)
def _isolate_stream_state():
"""Clear all shared streaming dicts before/after each test."""
config.STREAMS.clear()
config.CANCEL_FLAGS.clear()
config.AGENT_INSTANCES.clear()
config.STREAM_PARTIAL_TEXT.clear()
# New shared dicts for §A and §B
if hasattr(config, 'STREAM_REASONING_TEXT'):
config.STREAM_REASONING_TEXT.clear()
if hasattr(config, 'STREAM_LIVE_TOOL_CALLS'):
config.STREAM_LIVE_TOOL_CALLS.clear()
yield
config.STREAMS.clear()
config.CANCEL_FLAGS.clear()
config.AGENT_INSTANCES.clear()
config.STREAM_PARTIAL_TEXT.clear()
if hasattr(config, 'STREAM_REASONING_TEXT'):
config.STREAM_REASONING_TEXT.clear()
if hasattr(config, 'STREAM_LIVE_TOOL_CALLS'):
config.STREAM_LIVE_TOOL_CALLS.clear()
@pytest.fixture(autouse=True)
def _isolate_agent_locks():
config.SESSION_AGENT_LOCKS.clear()
yield
config.SESSION_AGENT_LOCKS.clear()
def _make_session(session_id="cancel_sid_1361",
pending_msg="Help me debug this",
messages=None):
"""Build a session in mid-stream state."""
s = Session(
session_id=session_id,
title="Test Session",
messages=messages or [],
)
s.pending_user_message = pending_msg
s.pending_attachments = []
s.pending_started_at = None
s.active_stream_id = "stream_1361"
s.save()
models.SESSIONS[session_id] = s
return s
def _setup_cancel_state(session_id, stream_id="stream_1361"):
"""Wire up STREAMS/CANCEL_FLAGS/AGENT_INSTANCES for cancel_stream()."""
config.STREAMS[stream_id] = queue.Queue()
config.CANCEL_FLAGS[stream_id] = threading.Event()
mock_agent = Mock()
mock_agent.session_id = session_id
mock_agent.interrupt = Mock()
config.AGENT_INSTANCES[stream_id] = mock_agent
return stream_id, mock_agent
# ── §A: Reasoning text lost on cancel ───────────────────────────────────────
class TestCancelPreservesReasoningText:
"""§A: _reasoning_text is thread-local and invisible to cancel_stream().
After fix: reasoning text should be persisted in a shared dict
(STREAM_REASONING_TEXT) keyed by stream_id, and cancel_stream()
should append it as a 'reasoning' field on the partial assistant message.
"""
def test_cancel_with_reasoning_only_preserves_reasoning(self):
"""Cancel during reasoning phase (no visible tokens) should persist reasoning."""
sid = "test_1361_a1"
stream_id = "stream_a1"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
# Simulate: reasoning was accumulated but no visible tokens
reasoning = "Let me think about this step by step..."
config.STREAM_PARTIAL_TEXT[stream_id] = "" # no visible tokens
if hasattr(config, 'STREAM_REASONING_TEXT'):
config.STREAM_REASONING_TEXT[stream_id] = reasoning
cancel_stream(stream_id)
# Reload and check
s2 = models.SESSIONS[sid]
msgs = s2.messages
# There should be a partial assistant message with reasoning
assistant_msgs = [m for m in msgs if isinstance(m, dict) and m.get('role') == 'assistant']
has_reasoning = any(m.get('reasoning') for m in assistant_msgs)
assert has_reasoning, \
f"Expected reasoning field on partial assistant msg after cancel. Got messages: {assistant_msgs}"
def test_cancel_with_reasoning_and_partial_tokens_preserves_both(self):
"""Cancel mid-stream with both reasoning and some visible tokens."""
sid = "test_1361_a2"
stream_id = "stream_a2"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
reasoning = "Let me analyze the code..."
partial_text = "Based on my analysis, the bug is in the"
config.STREAM_PARTIAL_TEXT[stream_id] = partial_text
if hasattr(config, 'STREAM_REASONING_TEXT'):
config.STREAM_REASONING_TEXT[stream_id] = reasoning
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
# Should have partial content
partial_msgs = [m for m in assistant_msgs if m.get('_partial')]
has_content = any(m.get('content') for m in partial_msgs)
assert has_content, \
f"Expected partial assistant content after cancel. Got: {partial_msgs}"
def test_cancel_without_reasoning_dict_works_as_before(self):
"""If STREAM_REASONING_TEXT doesn't exist yet (pre-fix), cancel still works."""
sid = "test_1361_a3"
stream_id = "stream_a3"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
config.STREAM_PARTIAL_TEXT[stream_id] = "Some partial text"
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
msgs = s2.messages
# Should have the cancel marker
has_cancel = any(
isinstance(m, dict) and m.get('role') == 'assistant' and m.get('_error')
for m in msgs
)
assert has_cancel, "Cancel marker should always be present"
# ── §B: Tool calls lost on cancel ───────────────────────────────────────────
class TestCancelPreservesToolCalls:
"""§B: _live_tool_calls is thread-local and invisible to cancel_stream().
After fix: tool calls should be persisted in a shared dict
(STREAM_LIVE_TOOL_CALLS) keyed by stream_id, and cancel_stream()
should append them as tool_call entries on the partial assistant message.
"""
def test_cancel_with_tool_calls_preserves_tools(self):
"""Cancel after tool execution should preserve the tool call info."""
sid = "test_1361_b1"
stream_id = "stream_b1"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
config.STREAM_PARTIAL_TEXT[stream_id] = ""
if hasattr(config, 'STREAM_LIVE_TOOL_CALLS'):
config.STREAM_LIVE_TOOL_CALLS[stream_id] = [
{"name": "read_file", "args": {"path": "/tmp/test.py"}, "done": True},
{"name": "terminal", "args": {"command": "ls"}, "done": False},
]
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
has_tools = any(m.get('tool_calls') or m.get('tools') for m in assistant_msgs)
assert has_tools, \
f"Expected tool_calls on partial assistant msg after cancel. Got: {assistant_msgs}"
def test_cancel_with_tools_and_text_preserves_both(self):
"""Cancel after tools + partial text should keep both."""
sid = "test_1361_b2"
stream_id = "stream_b2"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
config.STREAM_PARTIAL_TEXT[stream_id] = "Here's what I found:"
if hasattr(config, 'STREAM_LIVE_TOOL_CALLS'):
config.STREAM_LIVE_TOOL_CALLS[stream_id] = [
{"name": "web_search", "args": {"query": "test"}, "done": True},
]
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
partial_msgs = [m for m in assistant_msgs if m.get('_partial')]
has_content = any(m.get('content') for m in partial_msgs)
assert has_content, \
f"Expected partial content with tools after cancel. Got: {partial_msgs}"
# ── §C: Empty _stripped skips entire append ─────────────────────────────────
class TestCancelWithReasoningOnlyNoText:
"""§C: When streaming was 100% reasoning (no visible tokens), _stripped is
empty after regex cleanup, so no partial assistant message is appended.
After fix: even when _stripped is empty, if reasoning or tool calls exist,
a partial assistant message should be appended (with no content, but with
reasoning and/or tool_calls fields).
"""
def test_reasoning_only_creates_partial_message(self):
"""Cancel after reasoning-only output should still create a partial msg."""
sid = "test_1361_c1"
stream_id = "stream_c1"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
# Only reasoning, no visible tokens at all
config.STREAM_PARTIAL_TEXT[stream_id] = ""
if hasattr(config, 'STREAM_REASONING_TEXT'):
config.STREAM_REASONING_TEXT[stream_id] = "Deep reasoning here..."
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
# Should NOT be only the cancel marker — there should be a partial msg
partial_msgs = [m for m in assistant_msgs if m.get('_partial')]
assert len(partial_msgs) > 0, \
f"Expected at least one partial assistant msg for reasoning-only cancel. Got: {assistant_msgs}"
def test_tools_only_creates_partial_message(self):
"""Cancel after tool-only output (no text, no reasoning) should still create a partial msg."""
sid = "test_1361_c2"
stream_id = "stream_c2"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
config.STREAM_PARTIAL_TEXT[stream_id] = ""
if hasattr(config, 'STREAM_LIVE_TOOL_CALLS'):
config.STREAM_LIVE_TOOL_CALLS[stream_id] = [
{"name": "read_file", "args": {"path": "/tmp/x"}, "done": True},
]
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
partial_msgs = [m for m in assistant_msgs if m.get('_partial')]
assert len(partial_msgs) > 0, \
f"Expected at least one partial assistant msg for tools-only cancel. Got: {assistant_msgs}"
def test_no_reasoning_no_tools_no_partial(self):
"""Cancel with no reasoning and no tools and no text = only cancel marker (no change)."""
sid = "test_1361_c3"
stream_id = "stream_c3"
s = _make_session(session_id=sid)
_setup_cancel_state(sid, stream_id)
config.STREAM_PARTIAL_TEXT[stream_id] = ""
cancel_stream(stream_id)
s2 = models.SESSIONS[sid]
assistant_msgs = [m for m in s2.messages if isinstance(m, dict) and m.get('role') == 'assistant']
# Should only have the cancel marker, no partial messages
partial_msgs = [m for m in assistant_msgs if m.get('_partial')]
cancel_msgs = [m for m in assistant_msgs if m.get('_error')]
assert len(partial_msgs) == 0, \
f"Expected no partial msg when nothing was streamed. Got partials: {partial_msgs}"
assert len(cancel_msgs) == 1, \
f"Expected exactly 1 cancel marker. Got: {cancel_msgs}"