diff --git a/api/config.py b/api/config.py index e49ba23a..45694f40 100644 --- a/api/config.py +++ b/api/config.py @@ -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 diff --git a/api/streaming.py b/api/streaming.py index 62a034a8..e6c6700c 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -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']*>.*', '', _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). diff --git a/tests/test_issue1361_cancel_data_loss.py b/tests/test_issue1361_cancel_data_loss.py new file mode 100644 index 00000000..c7f08d52 --- /dev/null +++ b/tests/test_issue1361_cancel_data_loss.py @@ -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}"