Keep process-wakeup tags through cancel and recovery paths
This commit is contained in:
1 parent
2053825d74
commit
b09796375f
4 files changed
+55
-3
No files matched your search
@@ -1896,6 +1896,7 @@ def _apply_core_sync_or_error_marker(
|
|||||||
session.pending_user_message = None
|
session.pending_user_message = None
|
||||||
session.pending_attachments = []
|
session.pending_attachments = []
|
||||||
session.pending_started_at = None
|
session.pending_started_at = None
|
||||||
|
session.pending_user_source = None
|
||||||
session.save(touch_updated_at=touch_updated_at)
|
session.save(touch_updated_at=touch_updated_at)
|
||||||
logger.info(
|
logger.info(
|
||||||
"Session %s: cleared stale pending state for completed stream %s without error marker",
|
"Session %s: cleared stale pending state for completed stream %s without error marker",
|
||||||
@@ -1922,6 +1923,7 @@ def _apply_core_sync_or_error_marker(
|
|||||||
session.pending_user_message = None
|
session.pending_user_message = None
|
||||||
session.pending_attachments = []
|
session.pending_attachments = []
|
||||||
session.pending_started_at = None
|
session.pending_started_at = None
|
||||||
|
session.pending_user_source = None
|
||||||
session.messages.append(
|
session.messages.append(
|
||||||
_build_recovery_marker_with_retry_hook(
|
_build_recovery_marker_with_retry_hook(
|
||||||
recovered_output=recovered_output,
|
recovered_output=recovered_output,
|
||||||
@@ -1976,6 +1978,7 @@ def _apply_core_sync_or_error_marker(
|
|||||||
session.pending_user_message = None
|
session.pending_user_message = None
|
||||||
session.pending_attachments = []
|
session.pending_attachments = []
|
||||||
session.pending_started_at = None
|
session.pending_started_at = None
|
||||||
|
session.pending_user_source = None
|
||||||
if recovered_output:
|
if recovered_output:
|
||||||
session.messages.append(
|
session.messages.append(
|
||||||
_interrupted_recovery_marker(
|
_interrupted_recovery_marker(
|
||||||
@@ -2023,6 +2026,7 @@ def _apply_core_sync_or_error_marker(
|
|||||||
session.pending_user_message = None
|
session.pending_user_message = None
|
||||||
session.pending_attachments = []
|
session.pending_attachments = []
|
||||||
session.pending_started_at = None
|
session.pending_started_at = None
|
||||||
|
session.pending_user_source = None
|
||||||
session.messages.append(
|
session.messages.append(
|
||||||
_build_recovery_marker_with_retry_hook(
|
_build_recovery_marker_with_retry_hook(
|
||||||
recovered_output=recovered_output,
|
recovered_output=recovered_output,
|
||||||
|
|||||||
+11
-1
@@ -5008,6 +5008,9 @@ def _materialize_pending_user_turn_before_error(session) -> bool:
|
|||||||
'timestamp': recovered_ts,
|
'timestamp': recovered_ts,
|
||||||
'_recovered': True,
|
'_recovered': True,
|
||||||
}
|
}
|
||||||
|
pending_source = getattr(session, 'pending_user_source', None)
|
||||||
|
if pending_source and pending_source != 'webui':
|
||||||
|
recovered['_source'] = pending_source
|
||||||
pending_attachments = getattr(session, 'pending_attachments', None)
|
pending_attachments = getattr(session, 'pending_attachments', None)
|
||||||
if pending_attachments:
|
if pending_attachments:
|
||||||
recovered['attachments'] = list(pending_attachments)
|
recovered['attachments'] = list(pending_attachments)
|
||||||
@@ -8960,6 +8963,7 @@ def cancel_stream(stream_id: str) -> bool:
|
|||||||
# the cleanup.
|
# the cleanup.
|
||||||
try:
|
try:
|
||||||
_pending_user = getattr(_cs, 'pending_user_message', None)
|
_pending_user = getattr(_cs, 'pending_user_message', None)
|
||||||
|
_pending_source = getattr(_cs, 'pending_user_source', None)
|
||||||
_pending_atts_raw = getattr(_cs, 'pending_attachments', None)
|
_pending_atts_raw = getattr(_cs, 'pending_attachments', None)
|
||||||
_pending_atts = list(_pending_atts_raw) if isinstance(_pending_atts_raw, (list, tuple)) else []
|
_pending_atts = list(_pending_atts_raw) if isinstance(_pending_atts_raw, (list, tuple)) else []
|
||||||
_pending_started = getattr(_cs, 'pending_started_at', None) or 0
|
_pending_started = getattr(_cs, 'pending_started_at', None) or 0
|
||||||
@@ -8985,11 +8989,16 @@ def cancel_stream(stream_id: str) -> bool:
|
|||||||
if _pending_user == _last_content or _pending_user in _last_content:
|
if _pending_user == _last_content or _pending_user in _last_content:
|
||||||
_already_persisted = True
|
_already_persisted = True
|
||||||
if not _already_persisted:
|
if not _already_persisted:
|
||||||
|
_recovered_ts = int(time.time())
|
||||||
|
if isinstance(_pending_started, (int, float)) and _pending_started > 0:
|
||||||
|
_recovered_ts = int(_pending_started)
|
||||||
_user_turn: dict = {
|
_user_turn: dict = {
|
||||||
'role': 'user',
|
'role': 'user',
|
||||||
'content': _pending_user,
|
'content': _pending_user,
|
||||||
'timestamp': int(time.time()),
|
'timestamp': _recovered_ts,
|
||||||
}
|
}
|
||||||
|
if _pending_source and _pending_source != 'webui':
|
||||||
|
_user_turn['_source'] = _pending_source
|
||||||
if _pending_atts:
|
if _pending_atts:
|
||||||
_user_turn['attachments'] = _pending_atts
|
_user_turn['attachments'] = _pending_atts
|
||||||
_msgs_for_recovery.append(_user_turn)
|
_msgs_for_recovery.append(_user_turn)
|
||||||
@@ -9002,6 +9011,7 @@ def cancel_stream(stream_id: str) -> bool:
|
|||||||
_cs.pending_user_message = None
|
_cs.pending_user_message = None
|
||||||
_cs.pending_attachments = []
|
_cs.pending_attachments = []
|
||||||
_cs.pending_started_at = None
|
_cs.pending_started_at = None
|
||||||
|
_cs.pending_user_source = None
|
||||||
# Persist any partial assistant text that was streamed before cancel (#893).
|
# Persist any partial assistant text that was streamed before cancel (#893).
|
||||||
# Preserving partial content means the user sees what the agent had
|
# Preserving partial content means the user sees what the agent had
|
||||||
# produced rather than losing it entirely. The marker is _partial=True
|
# produced rather than losing it entirely. The marker is _partial=True
|
||||||
|
|||||||
@@ -1,9 +1,11 @@
|
|||||||
"""Test coverage for source field threading on process-wakeup synthetic turns."""
|
"""Test coverage for source field threading on process-wakeup synthetic turns."""
|
||||||
|
|
||||||
import time
|
import time
|
||||||
from api.models import Session, _append_recovered_pending_turn
|
from pathlib import Path
|
||||||
|
|
||||||
|
from api.models import Session, _append_recovered_pending_turn, _apply_core_sync_or_error_marker
|
||||||
from api.routes import _checkpoint_user_message_for_eager_session_save
|
from api.routes import _checkpoint_user_message_for_eager_session_save
|
||||||
from api.streaming import _merge_display_messages_after_agent_result
|
from api.streaming import _materialize_pending_user_turn_before_error, _merge_display_messages_after_agent_result
|
||||||
|
|
||||||
|
|
||||||
def test_append_recovered_pending_turn_stamps_process_wakeup_source():
|
def test_append_recovered_pending_turn_stamps_process_wakeup_source():
|
||||||
@@ -162,3 +164,29 @@ def test_merge_display_leaves_webui_user_turn_unmarked():
|
|||||||
|
|
||||||
assert merged[0]["role"] == "user"
|
assert merged[0]["role"] == "user"
|
||||||
assert "_source" not in merged[0]
|
assert "_source" not in merged[0]
|
||||||
|
|
||||||
|
|
||||||
|
def test_materialize_pending_user_turn_before_error_stamps_process_wakeup_source():
|
||||||
|
s = Session(
|
||||||
|
session_id="test-session-8",
|
||||||
|
pending_user_message="[IMPORTANT: Wakeup prompt]",
|
||||||
|
pending_user_source="process_wakeup",
|
||||||
|
)
|
||||||
|
s.messages = []
|
||||||
|
|
||||||
|
assert _materialize_pending_user_turn_before_error(s) is True
|
||||||
|
assert s.messages[0]["_source"] == "process_wakeup"
|
||||||
|
|
||||||
|
|
||||||
|
def test_apply_core_sync_or_error_marker_clears_pending_user_source(tmp_path):
|
||||||
|
s = Session(
|
||||||
|
session_id="test-session-9",
|
||||||
|
pending_user_message="[IMPORTANT: Wakeup prompt]",
|
||||||
|
pending_user_source="process_wakeup",
|
||||||
|
)
|
||||||
|
s.messages = [{"role": "assistant", "content": "done"}]
|
||||||
|
s.pending_started_at = time.time()
|
||||||
|
s.save = lambda *args, **kwargs: None
|
||||||
|
|
||||||
|
assert _apply_core_sync_or_error_marker(s, Path(tmp_path / "missing-core.json")) is True
|
||||||
|
assert s.pending_user_source is None
|
||||||
@@ -86,6 +86,8 @@ class TestCancelStreamEagerRelease:
|
|||||||
mock_session.pending_user_message = "hello"
|
mock_session.pending_user_message = "hello"
|
||||||
mock_session.pending_attachments = ["file.txt"]
|
mock_session.pending_attachments = ["file.txt"]
|
||||||
mock_session.pending_started_at = 1234567890.0
|
mock_session.pending_started_at = 1234567890.0
|
||||||
|
mock_session.pending_user_source = "process_wakeup"
|
||||||
|
mock_session.messages = []
|
||||||
|
|
||||||
STREAMS[stream_id] = queue.Queue()
|
STREAMS[stream_id] = queue.Queue()
|
||||||
CANCEL_FLAGS[stream_id] = threading.Event()
|
CANCEL_FLAGS[stream_id] = threading.Event()
|
||||||
@@ -102,6 +104,14 @@ class TestCancelStreamEagerRelease:
|
|||||||
"cancel_stream() should clear session.pending_attachments"
|
"cancel_stream() should clear session.pending_attachments"
|
||||||
assert mock_session.pending_started_at is None, \
|
assert mock_session.pending_started_at is None, \
|
||||||
"cancel_stream() should clear session.pending_started_at"
|
"cancel_stream() should clear session.pending_started_at"
|
||||||
|
assert mock_session.pending_user_source is None, \
|
||||||
|
"cancel_stream() should clear session.pending_user_source"
|
||||||
|
assert any(
|
||||||
|
isinstance(msg, dict)
|
||||||
|
and msg.get("role") == "user"
|
||||||
|
and msg.get("_source") == "process_wakeup"
|
||||||
|
for msg in mock_session.messages
|
||||||
|
), "cancel_stream() should preserve process_wakeup source on recovered user turns"
|
||||||
mock_session.save.assert_called_once()
|
mock_session.save.assert_called_once()
|
||||||
|
|
||||||
def test_cancel_without_agent_still_pops_streams(self):
|
def test_cancel_without_agent_still_pops_streams(self):
|
||||||
|
|||||||
Reference in new issue
Block a user