From 7d9c331d125512ffe4d972cead11078bfa5b2335 Mon Sep 17 00:00:00 2001
From: vegu-ai-tools <152010387+vegu-ai-tools@users.noreply.github.com>
Date: Thu, 14 May 2026 22:35:07 +0300
Subject: [PATCH] regenerate - handle cancel generation

---
 src/talemate/regenerate.py | 44 ++++++++++++++++++++++++++++++--------
 tests/test_regenerate.py   | 39 +++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 9 deletions(-)

diff --git a/src/talemate/regenerate.py b/src/talemate/regenerate.py
index 17206bde..b9ce533b 100644
--- a/src/talemate/regenerate.py
+++ b/src/talemate/regenerate.py
@@ -6,6 +6,7 @@ import talemate.events as events
 import talemate.emit.async_signals as async_signals
 from talemate.agents.editor.revision import RevisionContext
 from talemate.context import regeneration_context
+from talemate.exceptions import GenerationCancelled
 from talemate.scene_message import (
     SceneMessage,
     CharacterMessage,
@@ -365,6 +366,27 @@ async def _regenerate_inplace(
     return new_text, mutations, canonical_mutation_source
 
 
+async def _restore_history_after_failed_regenerate(
+    scene: "Scene",
+    message: SceneMessage,
+    popped_reinforcement_messages: list[ReinforcementMessage],
+) -> None:
+    """
+    Put history back the way it was before the regenerate attempt and tell
+    the frontend to drop the spinner on the message's slot. Shared by the
+    failure and the user-cancellation paths of ``regenerate``.
+    """
+    scene.history.append(message)
+    emit(
+        "regenerate_failed",
+        "",
+        websocket_passthrough=True,
+        kwargs={"id": message.id},
+    )
+    for reinforcement_message in reversed(popped_reinforcement_messages):
+        await scene.push_history(reinforcement_message)
+
+
 async def regenerate(scene: "Scene") -> list[SceneMessage]:
     """
     In-place regenerate the most recent AI response (the tail of
@@ -414,27 +436,31 @@ async def regenerate(scene: "Scene") -> list[SceneMessage]:
 
     try:
         outcome = await _regenerate_inplace(message, scene)
+    except GenerationCancelled:
+        # User-initiated interrupt — not a failure. Restore history, report
+        # the cancellation as a normal (non-error) status, then re-raise so
+        # the task done-callback posts the regenerate_failed envelope and
+        # the scene's cancel flag is cleared.
+        log.warning("regenerate: Generation cancelled by user", message=message)
+        await _restore_history_after_failed_regenerate(
+            scene, message, popped_reinforcement_messages
+        )
+        emit("status", message="Regeneration cancelled.", status="idle")
+        raise
     except Exception as e:
         log.error("regenerate: Exception during regeneration", message=message, error=e)
         outcome = None
 
     if not outcome:
         log.error("No new message generated", message=message)
-        # Put the message back where it was; nothing changed.
-        scene.history.append(message)
-        emit(
-            "regenerate_failed",
-            "",
-            websocket_passthrough=True,
-            kwargs={"id": message.id},
+        await _restore_history_after_failed_regenerate(
+            scene, message, popped_reinforcement_messages
         )
         emit(
             "status",
             message="Could not regenerate message.",
             status="error",
         )
-        for reinforcement_message in reversed(popped_reinforcement_messages):
-            await scene.push_history(reinforcement_message)
         return regenerated_messages
 
     new_text, mutation_delta, canonical_mutation_source = outcome
diff --git a/tests/test_regenerate.py b/tests/test_regenerate.py
index 73137e71..c8b6cdf5 100644
--- a/tests/test_regenerate.py
+++ b/tests/test_regenerate.py
@@ -21,6 +21,7 @@ import pytest
 import talemate.instance as instance
 import talemate.regenerate as regenerate_mod
 from talemate.agents.conversation import ConversationAgent
+from talemate.exceptions import GenerationCancelled
 from talemate.regenerate import (
     can_regenerate,
     ensure_regenerate_allowed,
@@ -1021,3 +1022,41 @@ class TestRegenerate:
         names = [e["name"] for e in _silence_emit_and_signals]
         assert "status" in names
         assert "regenerate_failed" in names
+
+    @pytest.mark.asyncio
+    async def test_generation_cancelled_reraises_without_error_status(
+        self, scene, register_agent, _silence_emit_and_signals
+    ):
+        """A user interrupt (GenerationCancelled) is not a failure: history is
+        restored, the cancellation is reported as a non-error status, and the
+        exception propagates so the task done-callback can handle it."""
+
+        class _Cancelled:
+            async def progress_story(self, **kwargs):
+                raise GenerationCancelled("Generation cancelled")
+
+        register_agent("narrator", _Cancelled())
+
+        original = NarratorMessage(message="primary")
+        original.meta = {
+            "agent": "narrator",
+            "function": "progress_story",
+            "arguments": {},
+        }
+        reinforce = ReinforcementMessage(message="reinforce")
+        scene.history = [original, reinforce]
+
+        with pytest.raises(GenerationCancelled):
+            await regenerate(scene)
+
+        # History is restored: primary slot untouched, reinforcement re-pushed.
+        assert original in scene.history
+        assert original.message == "primary"
+        assert reinforce in scene.history
+        # The frontend gets a regenerate_failed event so the spinner clears.
+        events_by_name = {e["name"]: e for e in _silence_emit_and_signals}
+        assert "regenerate_failed" in events_by_name
+        # The status is a non-error cancellation notice, not "error".
+        status = events_by_name["status"]
+        assert status["kwargs"].get("status") != "error"
+        assert "cancel" in status["kwargs"].get("message", "").lower()