diff --git a/src/talemate/agents/editor/websocket_handler.py b/src/talemate/agents/editor/websocket_handler.py index b159c7d9..5be60561 100644 --- a/src/talemate/agents/editor/websocket_handler.py +++ b/src/talemate/agents/editor/websocket_handler.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING from talemate.instance import get_agent from talemate.server.websocket_plugin import Plugin from talemate.scene_message import CharacterMessage +from talemate.status import background_task from talemate.agents.editor.revision import RevisionContext, RevisionInformation if TYPE_CHECKING: @@ -32,6 +33,7 @@ class EditorWebsocketHandler(Plugin): def editor(self): return get_agent("editor") + @background_task("Revising message") async def handle_request_revision(self, data: dict): """ Generate clickable actions for the user @@ -54,20 +56,14 @@ class EditorWebsocketHandler(Plugin): if not message: raise Exception("Message not found") - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - with RevisionContext(message.id): - info = RevisionInformation( - text=message.message, - character=character, - ) - revised = await editor.revision_revise(info) - if isinstance(message, CharacterMessage): - if not revised.startswith(character.name + ":"): - revised = f"{character.name}: {revised}" + with RevisionContext(message.id): + info = RevisionInformation( + text=message.message, + character=character, + ) + revised = await editor.revision_revise(info) + if isinstance(message, CharacterMessage): + if not revised.startswith(character.name + ":"): + revised = f"{character.name}: {revised}" - scene.edit_message(message.id, revised) - - self.run_in_background(task_wrapper, "Revising message") + scene.edit_message(message.id, revised) diff --git a/src/talemate/server/websocket_plugin.py b/src/talemate/server/websocket_plugin.py index 3b2e8145..930a2716 100644 --- a/src/talemate/server/websocket_plugin.py +++ b/src/talemate/server/websocket_plugin.py @@ -1,8 +1,7 @@ import structlog -from typing import TYPE_CHECKING, Awaitable, Callable +from typing import TYPE_CHECKING, Callable from talemate.emit import emit from talemate.exceptions import GenerationCancelled -from talemate.status import set_loading import traceback import pydantic import asyncio @@ -52,43 +51,6 @@ class Plugin: if hasattr(cls, "sub_handlers"): cls.sub_handlers = {} - @staticmethod - def consume_task_exception(task: asyncio.Task) -> None: - """ - Mark a background task's exception as retrieved to suppress the - "Task exception was never retrieved" warning at GC time. The - exception itself is already logged by the set_loading wrapper - (or the underlying coroutine's own error handling). - """ - if task.cancelled(): - return - # Calling exception() is enough to mark it retrieved; the return value - # is intentionally discarded. - task.exception() - - def run_in_background( - self, - coro_fn: Callable[[], Awaitable[None]], - message: str, - *, - cancellable: bool = True, - set_error: bool = True, - ) -> asyncio.Task: - """ - Schedule coro_fn() as a background task wrapped in set_loading, - with consume_task_exception attached as a done-callback. - - The handler that calls this can return immediately, freeing the - websocket receive loop to dispatch follow-up messages such as the - cancel/retry/ignore dialog response from the LLM client. - """ - wrapped = set_loading( - message, cancellable=cancellable, set_error=set_error - )(coro_fn) - task = asyncio.create_task(wrapped()) - task.add_done_callback(self.consume_task_exception) - return task - async def signal_operation_failed(self, message: str, emit_status: bool = True): self.websocket_handler.queue_put( { @@ -175,6 +137,43 @@ class Plugin: return on_done + def _on_background_task_done(self, task: asyncio.Task) -> None: + """ + Done-callback for handlers decorated with @background_task. Posts an + operation_done envelope on the router for the cancel and error paths + so the frontend's per-component busy state clears even when the + handler bailed before reaching its trailing signal_operation_done(). + + - Success: the handler's own signal_operation_done() already ran, + nothing to do here. + - Cancel (GenerationCancelled): post a bare operation_done envelope. + - Error: post operation_done with an error envelope. set_loading + already emitted the "Failed" status snackbar (set_error=True), so + we skip the duplicate emit here. + + Implemented as direct queue_put rather than scheduling another task + so we don't orphan a follow-up coroutine inside a done-callback. + """ + if task.cancelled(): + return + exc = task.exception() + if exc is None: + # Success — handler called signal_operation_done() itself. + return + if isinstance(exc, GenerationCancelled): + self.websocket_handler.queue_put( + {"type": self.router, "action": "operation_done", "data": {}} + ) + return + # Real exception — error envelope. + self.websocket_handler.queue_put( + { + "type": self.router, + "action": "operation_done", + "error": {"message": str(exc)}, + } + ) + async def handle(self, data: dict): action: str = data.get("action") log.info(f"{self.router} action", action=action) @@ -196,7 +195,7 @@ class Plugin: return try: - await fn(data) + result = await fn(data) except Exception as e: action_name = data.get("action") log.error( @@ -206,3 +205,13 @@ class Plugin: traceback=traceback.format_exc(), ) await self.signal_operation_failed(f"Error during {action_name}: {e}") + return + + # Handlers decorated with @background_task return a Task. Attach a + # done-callback so cancellations and synchronous failures inside the + # task body (e.g. pydantic validation, missing-character lookups) + # still post an operation_done envelope on the router and the + # frontend can clear any local busy state it set when sending the + # request. + if isinstance(result, asyncio.Task): + result.add_done_callback(self._on_background_task_done) diff --git a/src/talemate/server/world_state_manager/__init__.py b/src/talemate/server/world_state_manager/__init__.py index 261af55c..d4dec627 100644 --- a/src/talemate/server/world_state_manager/__init__.py +++ b/src/talemate/server/world_state_manager/__init__.py @@ -11,7 +11,7 @@ from talemate.game.schema import ConditionGroup from talemate.export import ExportOptions, export from talemate.instance import get_agent from talemate.world_state.manager import WorldStateManager, Suggestion -from talemate.status import set_loading +from talemate.status import background_task, set_loading import talemate.game.focal as focal from talemate.config import save_config from talemate.server.websocket_plugin import Plugin @@ -408,6 +408,7 @@ class WorldStateManagerPlugin( await self.handle_get_character_details({"name": payload.name}) await self.signal_operation_done() + @background_task("Refreshing reinforcement") async def handle_run_character_detail_reinforcement(self, data): payload = CharacterDetailReinforcementPayload(**data) @@ -418,27 +419,21 @@ class WorldStateManagerPlugin( reset=payload.reset, ) - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - await self.world_state_manager.run_detail_reinforcement( - payload.name, payload.question, reset=payload.reset - ) + await self.world_state_manager.run_detail_reinforcement( + payload.name, payload.question, reset=payload.reset + ) - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "character_detail_reinforcement_run", - "data": payload.model_dump(), - } - ) + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "character_detail_reinforcement_run", + "data": payload.model_dump(), + } + ) - # resend character details - await self.handle_get_character_details({"name": payload.name}) - await self.signal_operation_done() - - self.run_in_background(task_wrapper, "Refreshing reinforcement") + # resend character details + await self.handle_get_character_details({"name": payload.name}) + await self.signal_operation_done() async def handle_delete_character_detail_reinforcement(self, data): payload = CharacterDetailReinforcementPayload(**data) @@ -612,42 +607,37 @@ class WorldStateManagerPlugin( await self.handle_get_world({}) await self.signal_operation_done() + @background_task("Refreshing world state") async def handle_run_world_state_reinforcement(self, data): payload = WorldEntryReinforcementPayload(**data) - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - await self.world_state_manager.run_detail_reinforcement( - None, payload.question, payload.reset - ) + await self.world_state_manager.run_detail_reinforcement( + None, payload.question, payload.reset + ) - ( - _, - reinforcement, - ) = await self.world_state_manager.world_state.find_reinforcement( - payload.question, None - ) + ( + _, + reinforcement, + ) = await self.world_state_manager.world_state.find_reinforcement( + payload.question, None + ) - if not reinforcement: - log.error("Reinforcement not found", question=payload.question) - await self.signal_operation_done() - return - - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "world_state_reinforcement_ran", - "data": reinforcement.model_dump(), - } - ) - - # resend world - await self.handle_get_world({}) + if not reinforcement: + log.error("Reinforcement not found", question=payload.question) await self.signal_operation_done() + return - self.run_in_background(task_wrapper, "Refreshing world state") + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "world_state_reinforcement_ran", + "data": reinforcement.model_dump(), + } + ) + + # resend world + await self.handle_get_world({}) + await self.signal_operation_done() async def handle_delete_world_state_reinforcement(self, data): payload = WorldEntryReinforcementPayload(**data) @@ -814,36 +804,31 @@ class WorldStateManagerPlugin( await self.scene.load_active_pins() self.scene.emit_status() + @background_task("Applying template") async def handle_apply_template(self, data): payload = ApplyWorldStateTemplatePayload(**data) log.debug("Apply world state template", payload=payload) - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - result = await self.world_state_manager.apply_template( - template=payload.template, - character_name=payload.character_name or "", - run_immediately=payload.run_immediately, - ) + result = await self.world_state_manager.apply_template( + template=payload.template, + character_name=payload.character_name or "", + run_immediately=payload.run_immediately, + ) - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "template_applied", - "status": "done", - "data": payload.model_dump(), - "result": result.model_dump() if result else None, - } - ) + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "template_applied", + "status": "done", + "data": payload.model_dump(), + "result": result.model_dump() if result else None, + } + ) - await self.handle_get_world({}) - await self.handle_get_templates({}) - await self.signal_operation_done() - - self.run_in_background(task_wrapper, "Applying template") + await self.handle_get_world({}) + await self.handle_get_templates({}) + await self.signal_operation_done() async def handle_save_template(self, data): payload = SaveWorldStateTemplatePayload(**data) @@ -886,6 +871,7 @@ class WorldStateManagerPlugin( await self.handle_get_templates({}) await self.signal_operation_done() + @background_task("Applying templates") async def handle_apply_templates(self, data): payload = ApplyWorldStateTemplatesPayload(**data) @@ -913,33 +899,27 @@ class WorldStateManagerPlugin( } ) - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - await self.world_state_manager.apply_templates( - payload.templates, - callback_start=callback_start, - callback_done=callback_done, - character_name=payload.character_name, - run_immediately=payload.run_immediately, - generation_options=payload.generation_options, - ) + await self.world_state_manager.apply_templates( + payload.templates, + callback_start=callback_start, + callback_done=callback_done, + character_name=payload.character_name, + run_immediately=payload.run_immediately, + generation_options=payload.generation_options, + ) - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "templates_applied", - "source": payload.source, - "data": payload.model_dump(), - } - ) + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "templates_applied", + "source": payload.source, + "data": payload.model_dump(), + } + ) - await self.handle_get_world({}) - await self.handle_get_templates({}) - await self.signal_operation_done() - - self.run_in_background(task_wrapper, "Applying templates") + await self.handle_get_world({}) + await self.handle_get_templates({}) + await self.signal_operation_done() async def handle_save_template_group(self, data): payload = SaveWorldStateTemplateGroupPayload(**data) @@ -978,6 +958,7 @@ class WorldStateManagerPlugin( await self.handle_get_templates({}) await self.signal_operation_done() + @background_task("Generating dialogue instructions") async def handle_generate_character_dialogue_instructions(self, data): payload = SelectiveCharacterPayload(**data) @@ -991,33 +972,27 @@ class WorldStateManagerPlugin( creator = get_agent("creator") - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - instructions = await creator.determine_character_dialogue_instructions( - character - ) + instructions = await creator.determine_character_dialogue_instructions( + character + ) - character.dialogue_instructions = instructions + character.dialogue_instructions = instructions - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "character_dialogue_instructions_generated", - "data": { - "name": payload.name, - "instructions": instructions, - }, - } - ) + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "character_dialogue_instructions_generated", + "data": { + "name": payload.name, + "instructions": instructions, + }, + } + ) - # signal_operation_done already emits scene status on the - # non-auto-save branch (and scene.save() emits it on the auto-save - # branch), so no trailing emit_status() is needed here. - await self.signal_operation_done() - - self.run_in_background(task_wrapper, "Generating dialogue instructions") + # signal_operation_done already emits scene status on the + # non-auto-save branch (and scene.save() emits it on the auto-save + # branch), so no trailing emit_status() is needed here. + await self.signal_operation_done() async def handle_delete_character(self, data): payload = SelectiveCharacterPayload(**data) diff --git a/src/talemate/server/world_state_manager/history.py b/src/talemate/server/world_state_manager/history.py index 57b14b00..3d27e73e 100644 --- a/src/talemate/server/world_state_manager/history.py +++ b/src/talemate/server/world_state_manager/history.py @@ -19,6 +19,7 @@ from talemate.history import ( ) from talemate.scene_message import TimePassageMessage from talemate.server.world_state_manager import world_state_templates +from talemate.status import background_task from talemate.util.time import ( amount_unit_to_iso8601_duration, iso8601_duration_to_human, @@ -152,6 +153,7 @@ class HistoryMixin: await self.signal_operation_done() + @background_task("Regenerating history entry") async def handle_regenerate_history_entry(self, data): """ Regenerate a single history entry. @@ -161,22 +163,16 @@ class HistoryMixin: log.debug("regenerate_history_entry", payload=payload) - # Run in a background task so the websocket receive loop stays free - # to dispatch the cancel/retry/ignore dialog response — otherwise - # awaiting the dialog future inside this handler deadlocks the loop. - async def task_wrapper(): - entry = await regenerate_history_entry(self.scene, payload.entry) - log.debug("regenerate_history_entry (done)", entry=entry) - self.websocket_handler.queue_put( - { - "type": "world_state_manager", - "action": "history_entry_regenerated", - "data": entry.model_dump(), - } - ) - await self.signal_operation_done() - - self.run_in_background(task_wrapper, "Regenerating history entry") + entry = await regenerate_history_entry(self.scene, payload.entry) + log.debug("regenerate_history_entry (done)", entry=entry) + self.websocket_handler.queue_put( + { + "type": "world_state_manager", + "action": "history_entry_regenerated", + "data": entry.model_dump(), + } + ) + await self.signal_operation_done() async def handle_inspect_history_entry(self, data): """ diff --git a/src/talemate/status.py b/src/talemate/status.py index 99a5f691..83fa5b50 100644 --- a/src/talemate/status.py +++ b/src/talemate/status.py @@ -1,5 +1,6 @@ import asyncio import structlog +from functools import wraps from talemate.emit import emit from talemate.exceptions import GenerationCancelled @@ -7,6 +8,7 @@ from talemate.context import handle_generation_cancelled __all__ = [ "set_loading", + "background_task", "LoadingStatus", ] @@ -54,6 +56,10 @@ class set_loading: # never resolves and the scene input stays disabled. emit("status", message="", status="idle") handle_generation_cancelled(e) + # Re-raise so callers (notably the @background_task decorator's + # done-callback in Plugin.handle) can post a router-level + # operation_done envelope and clear per-component busy state. + raise except Exception as e: log.error("Error in set_loading wrapper", error=e) if self.set_error: @@ -66,13 +72,72 @@ class set_loading: if self.as_async: async def async_wrapper(*args, **kwargs): - return asyncio.create_task(wrapper(*args, **kwargs)) + task = asyncio.create_task(wrapper(*args, **kwargs)) + # Mark exceptions as retrieved so cancellations / errors that + # propagate out of the wrapper don't surface as + # "Task exception was never retrieved" warnings. + task.add_done_callback(_consume_task_exception) + return task return async_wrapper return wrapper +def _consume_task_exception(task: asyncio.Task) -> None: + """ + Mark a background task's exception as retrieved to suppress the + "Task exception was never retrieved" warning at GC time. The + exception itself is already logged by the set_loading wrapper. + """ + if task.cancelled(): + return + # Calling exception() is enough to mark it retrieved; the return value + # is intentionally discarded. + task.exception() + + +def background_task( + message: str, + *, + cancellable: bool = True, + set_success: bool = False, + set_error: bool = True, +): + """ + Decorator: schedule the wrapped coroutine as a background asyncio task + with set_loading status emissions and exception cleanup. + + The wrapped function returns immediately with the task object — calling + code can ignore it. This is what frees the websocket receive loop so + follow-up messages (e.g. the cancel/retry/ignore dialog response from + the LLM client) can be dispatched while the work is still running. + + The frontend still observes a busy snackbar (with a cancel button when + cancellable=True) for the duration of the task; "background" here refers + to the backend handler returning before the work completes, not to the + UX being free to continue. + """ + + def decorator(fn): + wrapped = set_loading( + message, + cancellable=cancellable, + set_success=set_success, + set_error=set_error, + )(fn) + + @wraps(fn) + async def outer(*args, **kwargs): + task = asyncio.create_task(wrapped(*args, **kwargs)) + task.add_done_callback(_consume_task_exception) + return task + + return outer + + return decorator + + class LoadingStatus: def __init__(self, max_steps: int | None = None, cancellable: bool = False): self.max_steps = max_steps diff --git a/talemate_frontend/src/components/WorldStateManagerCharacterReinforcements.vue b/talemate_frontend/src/components/WorldStateManagerCharacterReinforcements.vue index 88b5a6c4..f5a692ef 100644 --- a/talemate_frontend/src/components/WorldStateManagerCharacterReinforcements.vue +++ b/talemate_frontend/src/components/WorldStateManagerCharacterReinforcements.vue @@ -432,7 +432,9 @@ export default { else if (message.action === 'operation_done') { // Clear busy on cancel/error too — the success-specific // character_detail_reinforcement_run message isn't sent - // when the operation fails or is cancelled. + // when the operation fails or is cancelled. The backend's + // @background_task done-callback always posts operation_done + // on cancel and error so this fires reliably. this.busy = false; } else if (message.action === 'template_applied' && message.source === this.source){ diff --git a/talemate_frontend/src/components/WorldStateManagerTemplateApplicator.vue b/talemate_frontend/src/components/WorldStateManagerTemplateApplicator.vue index c650ad8e..9e7b21f6 100644 --- a/talemate_frontend/src/components/WorldStateManagerTemplateApplicator.vue +++ b/talemate_frontend/src/components/WorldStateManagerTemplateApplicator.vue @@ -293,10 +293,26 @@ export default { }); }, handleMessage(message) { - if (message.type !== 'world_state_manager' || message.source !== this.source) { + if (message.type !== 'world_state_manager') { return; - } - else if (message.action === 'template_applying') { + } + if (message.action === 'operation_done') { + // Clear all local busy state on cancel/error — the success-only + // template_applied / templates_applied messages don't fire when + // the apply is interrupted, leaving spinners stuck on the + // currently-applying template and group. + if (this.busy || this.busyTemplateUID || this.busyGroupUID) { + this.busy = false; + this.busyTemplateUID = null; + this.busyGroupUID = null; + this.$emit('done'); + } + return; + } + if (message.source !== this.source) { + return; + } + if (message.action === 'template_applying') { this.busyTemplateUID = message.data.uid; this.busyGroupUID = message.data.group; }