Mirrors/talemate
1
0
Fork 0
mirror of https://github.com/vegu-ai/talemate.git synced 2026-09-01 19:48:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor: streamline background task handling and enhance message processing in websocket handlers

Browse Source
This commit is contained in:
vegu-ai-tools
2026-05-06 21:40:42 +03:00
parent fe0826385a
commit ccc14dc401
7 changed files with 257 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
src/talemate/agents/editor/websocket_handler.py
View File

@@ -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)

89
src/talemate/server/websocket_plugin.py
View File

@@ -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)

217
src/talemate/server/world_state_manager/__init__.py
View File

@@ -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)

28
src/talemate/server/world_state_manager/history.py
View File

@@ -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):
"""

67
src/talemate/status.py
View File

@@ -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

4
talemate_frontend/src/components/WorldStateManagerCharacterReinforcements.vue
View File

@@ -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){

22
talemate_frontend/src/components/WorldStateManagerTemplateApplicator.vue
View File

@@ -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;
}