From 824b862bb061dd2087cc61327ab8ca078aecbf7e Mon Sep 17 00:00:00 2001 From: vegu-ai-tools <152010387+vegu-ai-tools@users.noreply.github.com> Date: Sat, 25 Apr 2026 20:12:30 +0300 Subject: [PATCH] feat: add Cancel button to contextual generation dialog and enable interrupt functionality --- CHANGELOG.yaml | 1 + src/talemate/server/assistant.py | 87 ++++++++++++++----- .../src/components/ContextualGenerate.vue | 21 ++++- 3 files changed, 84 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.yaml b/CHANGELOG.yaml index a1ab5ab1..3a3f3440 100644 --- a/CHANGELOG.yaml +++ b/CHANGELOG.yaml @@ -32,6 +32,7 @@ - "Encryption: Fixed API keys in agent action configs (e.g., OpenAI-compatible TTS, visual backends) being stored in plaintext. The encryption walker now handles the nested AgentActionConfig structure where the key is wrapped in a {value: ...} dict." - "Agent Config: Stopped persisting unified_api_key fields to config.yaml. These are static reference pointers (e.g., 'openai.api_key') defined in code and never change, so saving them was unnecessary." - "Contextual Generate: Fixed list-type generation prefilling an empty line after '1.', which caused some LLMs to produce empty lists." + - "Contextual Generate: Added a Cancel button to the generate dialog and moved generation to a background task so the interrupt signal is actually processed mid-generation (previously the websocket receive loop was blocked until the generation finished, defeating the purpose of any cancel)." - "Narrate Progress: Fixed narration sometimes generating screenplay-style dialogue entries instead of continuous narrative prose." - "Advance Time: Fixed toolbar time advancement being completely non-functional due to a missing websocket handler. Also fixed invalid ISO 8601 duration strings for the 1 week and 2 weeks options." - "Prompts UI: Fixed the template preview CodeMirror editor being clipped at the bottom across the Active, group, and LLM prompt template tabs, so the bottom scrollbar is fully visible." diff --git a/src/talemate/server/assistant.py b/src/talemate/server/assistant.py index 2efa4bc4..534b01d3 100644 --- a/src/talemate/server/assistant.py +++ b/src/talemate/server/assistant.py @@ -8,6 +8,7 @@ from talemate.agents.creator.assistant import ContentGenerationContext from talemate.client.context import ClientContext from talemate.context import RegenerationContext from talemate.emit import emit +from talemate.exceptions import GenerationCancelled from talemate.instance import get_agent from talemate.load.character_card import analyze_character_card from talemate.regenerate import ensure_regenerate_allowed, regenerate @@ -48,35 +49,73 @@ class AssistantPlugin(Plugin): self.websocket_handler = websocket_handler async def handle_contextual_generate(self, data: dict): + """ + Run contextual_generate as a background task so the websocket receive loop + stays free to process inbound messages (notably `interrupt`) while the + generation is in flight. + """ payload = ContentGenerationContext(**data) creator = get_agent("creator") - if payload.computed_context[0] == "acting_instructions": - content = await creator.determine_character_dialogue_instructions( - self.scene.get_character(payload.character), - instructions=payload.instructions, + async def _run() -> str: + if payload.computed_context[0] == "acting_instructions": + return await creator.determine_character_dialogue_instructions( + self.scene.get_character(payload.character), + instructions=payload.instructions, + ) + return await creator.contextual_generate(payload) + + # Inlined instead of using Plugin.create_task_done_callback because the + # success branch needs a uid-scoped payload + status emit. If a third + # handler ends up needing this same shape, generalize the helper to + # accept a custom success payload builder rather than copying again. + def _on_done(task: asyncio.Task): + try: + content = task.result() + except GenerationCancelled: + log.warning("contextual_generate cancelled", uid=payload.uid) + self.websocket_handler.queue_put( + { + "type": self.router, + "action": "contextual_generate_cancelled", + "data": {"uid": payload.uid}, + } + ) + return + except Exception as e: + log.error( + "Error running contextual_generate", + error=traceback.format_exc(), + ) + self.websocket_handler.queue_put( + { + "type": self.router, + "action": "contextual_generate_failed", + "data": {"uid": payload.uid, "message": str(e)}, + } + ) + return + + context_type, context_name = payload.computed_context + emit( + "status", + message=f"Generated {context_type}: {context_name}", + status="success", + ) + self.websocket_handler.queue_put( + { + "type": self.router, + "action": "contextual_generate_done", + "data": { + "generated_content": content, + "uid": payload.uid, + **payload.model_dump(), + }, + } ) - else: - content = await creator.contextual_generate(payload) - context_type, context_name = payload.computed_context - emit( - "status", - message=f"Generated {context_type}: {context_name}", - status="success", - ) - - self.websocket_handler.queue_put( - { - "type": self.router, - "action": "contextual_generate_done", - "data": { - "generated_content": content, - "uid": payload.uid, - **payload.model_dump(), - }, - } - ) + task = asyncio.create_task(_run()) + task.add_done_callback(_on_done) async def handle_autocomplete(self, data: dict): data = ContentGenerationContext(**data) diff --git a/talemate_frontend/src/components/ContextualGenerate.vue b/talemate_frontend/src/components/ContextualGenerate.vue index 748687a9..a75ac105 100644 --- a/talemate_frontend/src/components/ContextualGenerate.vue +++ b/talemate_frontend/src/components/ContextualGenerate.vue @@ -55,6 +55,7 @@ + Cancel Generate @@ -314,9 +315,27 @@ export default { })); }, + cancel() { + this.getWebsocket().send(JSON.stringify({ type: 'interrupt' })); + }, + handleMessage(message) { + if (message.type === "assistant" && message.action === "contextual_generate_cancelled") { + if(message.data.uid !== this.uid) + return; + this.busy = false; + this.dialog = false; + return; + } + if (message.type === "assistant" && message.action === "contextual_generate_failed") { + if(message.data.uid !== this.uid) + return; + this.busy = false; + this.dialog = false; + return; + } if (message.type === "assistant" && message.action === "contextual_generate_done") { - + if(message.data.uid !== this.uid) return;