fix(loop): preserve partial context when /stop cancels a task

When a user sends /stop to interrupt an active agent turn, the task is
cancelled via CancelledError. Previously, the cancellation handler just
logged and re-raised, discarding any tool results and assistant messages
accumulated during the interrupted turn.

The runtime checkpoint mechanism already persists partial turn state
(assistant messages, completed tool results, pending tool calls) into
session metadata via _emit_checkpoint. However, this checkpoint was only
materialized into session history on the NEXT incoming message via
_restore_runtime_checkpoint — not at cancellation time.

Now the CancelledError handler in _dispatch calls
_restore_runtime_checkpoint immediately, so the partial context is
preserved in session history. This means the next message the user sends
will see all the work that was done before /stop, rather than starting
from scratch.

Fixes #2966

Includes 3 tests verifying checkpoint restoration on cancellation.
This commit is contained in:
hussein1362
2026-04-21 01:14:41 +08:00
committed by Xubin Ren
parent 899a9073ce
commit 847c50b2de
2 changed files with 107 additions and 0 deletions
+23
View File
@@ -579,6 +579,29 @@ class AgentLoop:
))
except asyncio.CancelledError:
logger.info("Task cancelled for session {}", session_key)
# Preserve partial context from the interrupted turn so
# the user does not lose tool results and assistant
# messages accumulated before /stop. The checkpoint was
# already persisted to session metadata by
# _emit_checkpoint during tool execution; materializing
# it into session history now makes it visible in the
# next conversation turn.
try:
key = self._effective_session_key(msg)
session = self.sessions.get_or_create(key)
if self._restore_runtime_checkpoint(session):
self._clear_pending_user_turn(session)
self.sessions.save(session)
logger.info(
"Restored partial context for cancelled session {}",
key,
)
except Exception:
logger.debug(
"Could not restore checkpoint for cancelled session {}",
session_key,
exc_info=True,
)
raise
except Exception:
logger.exception("Error processing message for session {}", session_key)