fix: stop Parlant session processing and preamble generation

.github/workflows/ci.yml

@@ -317,6 +317,7 @@ jobs:
           tests/integrations/parlant/test_customtools.py \
           tests/integrations/parlant/test_filetools.py \
           tests/integrations/parlant/test_guard.py \
+          tests/integrations/parlant/test_preambletask.py \
           tests/integrations/parlant/test_selection.py \
           tests/integrations/parlant/test_server.py \
           tests/integrations/parlant/test_tasktools.py \

pyproject.toml

@@ -310,6 +310,7 @@ dev-parlant = [
     "pydantic-settings>=2.0.0",
     "pytest>=7.0.0",
     "pytest-asyncio>=0.21.0",
+    "pytest-randomly>=5.0.0",
     "looptime>=0.7",
     "pytest-mock>=3.10.0",
     "pytest-cov>=4.0.0",

src/band/adapters/parlant/adapter.py

@@ -9,6 +9,7 @@
 
 import logging
 from collections.abc import Sequence
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, NoReturn
 
 from band_sdk_core import AgentFailure
@@ -18,13 +19,20 @@
 from band.adapters.parlant.guidelines import GuidelineLedger, GuidelineSpec
 from band.adapters.parlant.history import inject_history, name_this_agent_in_history
 from band.adapters.parlant.lifecycle import ServerLifecycle
-from band.adapters.parlant.messages import compose_user_message, post_user_message
+from band.adapters.parlant.messages import (
+    await_stopped_task,
+    compose_user_message,
+    live_preamble_tasks,
+    post_user_message,
+    process_session_tag,
+)
 from band.adapters.parlant.responses import relay_agent_response
 from band.adapters.parlant.rooms import RoomSessions
 from band.converters.parlant import ParlantHistoryConverter, ParlantMessages
 from band.core.delivery import DeliveryFailedError, reraise_delivery_cause
 from band.core.protocols import GENERIC_PROVIDER_FAILURE_MESSAGE, AgentToolsProtocol
 from band.core.simple_adapter import SimpleAdapter
+from band.core.turn_lifecycle import BackendTurns
 from band.core.types import Capability, Emit, FeatureKwargs, PlatformMessage
 from band.integrations.parlant.customschema import check_parlant_custom_tools
 from band.integrations.parlant.sessiontools import bound_session_tools
@@ -34,6 +42,7 @@
 if TYPE_CHECKING:
     import parlant.sdk as p
     from parlant.core.application import Application
+    from parlant.core.background_tasks import BackgroundTaskService
     from parlant.core.sessions import SessionId
 
 logger = logging.getLogger(__name__)
@@ -49,6 +58,17 @@
     "in configure=."
 )
 
+# Recorded by Parlant on the cancelled session task.
+STOP_REASON = "Band turn stopped"
+
+
+@dataclass(frozen=True)
+class ParlantSessionTurn:
+    """The engine's background processing of one session's response."""
+
+    background_tasks: BackgroundTaskService
+    session_id: SessionId
+
 
 class ParlantAdapter(SimpleAdapter[ParlantMessages]):
     """
@@ -165,7 +185,11 @@ def __init__(
         self._guidelines = GuidelineLedger()
         self._tools: list[Any] = []
         self._rooms: RoomSessions | None = None
+        self._background_tasks: BackgroundTaskService | None = None
         self._started = False
+        self._backend_turns = BackendTurns[ParlantSessionTurn](
+            self._cancel_processing, provider=PROVIDER
+        )
 
     def add_guideline(
         self,
@@ -284,7 +308,11 @@ async def _prepare_server(
             await self._configure(server, agent)
 
         from parlant.core.application import Application  # noqa: PLC0415
+        from parlant.core.background_tasks import (  # noqa: PLC0415
+            BackgroundTaskService,
+        )
 
+        self._background_tasks = server.container[BackgroundTaskService]
         return agent, server.container[Application]
 
     def _warn_if_custom_tools_unreachable(self) -> None:
@@ -305,6 +333,10 @@ def _forget_guidelines_if_agent_gone(self) -> None:
         if self._lifecycle.agent is None:
             self._guidelines.forget_applied()
 
+    @property
+    def backend_turns(self) -> BackendTurns[ParlantSessionTurn]:
+        return self._backend_turns
+
     async def on_message(
         self,
         msg: PlatformMessage,
@@ -410,30 +442,64 @@ async def _converse(
         tools: AgentToolsProtocol,
         sender_name: str,
     ) -> None:
-        """Post *user_message* as the customer, then relay the engine's reply."""
-        offset = await post_user_message(
-            app=app, session_id=session_id, message=user_message
-        )
-        await relay_agent_response(
-            app=app,
-            session_id=session_id,
-            room_id=room_id,
-            min_offset=offset,
-            tools=tools,
-            sender_name=sender_name,
-            timeout=self.config.response_timeout,
-            poll=self.config.response_poll,
+        """Post *user_message* as the customer, then relay the engine's reply.
+
+        The engine answers in its own background task, so ending this turn
+        early cancels that task too.
+        """
+        handle = (
+            None
+            if self._background_tasks is None
+            else ParlantSessionTurn(self._background_tasks, session_id)
         )
+        async with self.backend_turns.running(room_id, handle):
+            offset = await post_user_message(
+                app=app, session_id=session_id, message=user_message
+            )
+            gave_up = await relay_agent_response(
+                app=app,
+                session_id=session_id,
+                room_id=room_id,
+                min_offset=offset,
+                tools=tools,
+                sender_name=sender_name,
+                timeout=self.config.response_timeout,
+                poll=self.config.response_poll,
+            )
+            if gave_up:
+                await self.backend_turns.stop(room_id)
+
+    @staticmethod
+    async def _cancel_processing(turn: ParlantSessionTurn) -> None:
+        """Cancel the session's process-session task and its preamble child.
+
+        Parlant 3.3.2's cancel marks the parent and returns without waiting,
+        and the parent does not cancel the preamble task it already created.
+        Waiting for the parent first is what makes that child visible.
+        """
+        tag = process_session_tag(turn.session_id)
+        parent = turn.background_tasks._tasks.get(tag)
+        await turn.background_tasks.cancel(tag=tag, reason=STOP_REASON)
+        if parent is not None:
+            await await_stopped_task(parent)
+        children = live_preamble_tasks(str(turn.session_id))
+        for child in children:
+            child.cancel()
+        for child in children:
+            await await_stopped_task(child)
 
     async def on_cleanup(self, room_id: str) -> None:
         """Forget the room's Parlant session when the agent leaves it."""
+        await self.backend_turns.settled(room_id)
         if self._rooms is not None:
             self._rooms.forget(room_id)
         logger.debug("Room %s: Cleaned up Parlant session", room_id)
 
     async def cleanup_all(self) -> None:
         """Release all sessions and the owned Parlant server (call on stop)."""
+        await self.backend_turns.drain()
         self._rooms = None
+        self._background_tasks = None
         await self._lifecycle.release()
         self._forget_guidelines_if_agent_gone()
         self._started = False

src/band/adapters/parlant/messages.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import re
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any
@@ -71,6 +72,51 @@ def compose_user_message(
     return user_message
 
 
+def process_session_tag(session_id: SessionId | str) -> str:
+    """The tag Parlant runs a session's response under in its background tasks.
+
+    Mirrors Parlant's own tag (pinned by a parity test against the engine), so
+    cancelling it stops the engine's work on the session.
+    """
+    return f"process-session({session_id})"
+
+
+# Parlant 3.3.2 creates this child with asyncio.create_task inside
+# AlphaEngine._get_preamble_task and does not cancel it when process-session
+# is cancelled. The nested coroutine's qualname is the class that defines it.
+PREAMBLE_TASK_QUALNAME = "AlphaEngine._get_preamble_task.<locals>.preamble_task"
+
+
+def live_preamble_tasks(session_id: str) -> list[asyncio.Task[Any]]:
+    """Preamble generation still running for ``session_id``."""
+    tasks: list[asyncio.Task[Any]] = []
+    for task in asyncio.all_tasks():
+        if task.done():
+            continue
+        frame = next(
+            (
+                suspended
+                for suspended in task.get_stack()
+                if suspended.f_code.co_qualname == PREAMBLE_TASK_QUALNAME
+            ),
+            None,
+        )
+        if frame is None:
+            continue
+        context = frame.f_locals.get("context")
+        session = getattr(context, "session", None)
+        if str(getattr(session, "id", "")) == str(session_id):
+            tasks.append(task)
+    return tasks
+
+
+async def await_stopped_task(task: asyncio.Task[Any]) -> None:
+    """Wait until ``task`` finishes, including when it was cancelled."""
+    await asyncio.wait({task})
+    if not task.cancelled():
+        task.exception()
+
+
 async def post_user_message(
     *, app: Application, session_id: SessionId, message: str
 ) -> int:

src/band/adapters/parlant/responses.py

@@ -11,7 +11,6 @@
 
 import asyncio
 import logging
-import time
 from collections.abc import Sequence
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Any, NamedTuple, cast
@@ -119,18 +118,21 @@ async def relay_agent_response(
     sender_name: str,
     timeout: float,
     poll: float,
-) -> None:
+) -> bool:
     """Relay the engine's final reply to *sender_name*, waiting up to *timeout*.
 
     Waiting polls in *poll*-second windows and keeps going past empty ones,
     so a slow (cold-start) turn is still answered. If the budget elapses with
     no final message the turn is given up honestly: Parlant intermittently
     stalls after a preamble, and a preamble is not an answer.
+
+    Returns whether the turn was given up that way, with the engine possibly
+    still working on it.
     """
     offset = min_offset
-    # perf_counter keeps its resolution on Windows, where monotonic() is coarse.
-    deadline = time.perf_counter() + timeout
-    while (remaining := deadline - time.perf_counter()) > 0:
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + timeout
+    while (remaining := deadline - loop.time()) > 0:
         polled = await next_agent_events(
             app=app,
             session_id=session_id,
@@ -144,13 +146,13 @@ async def relay_agent_response(
                     "Room %s: stopped waiting for a Parlant reply because reading events failed",
                     room_id,
                 )
-                return
+                return False
             case PollMiss.NO_UPDATE if tools.turn.replied:
                 logger.info(
                     "Room %s: stopped waiting for a Parlant reply because the turn already replied or declined",
                     room_id,
                 )
-                return
+                return False
             case PollMiss():
                 await _back_off(deadline)
                 continue
@@ -169,15 +171,17 @@ async def relay_agent_response(
                     room_id,
                     final_text,
                 )
-            return
+            return False
         if batch.final_seen:
-            return
+            return False
         logger.debug("Room %s: only a preamble so far; still waiting", room_id)
 
-    if not tools.turn.replied:
-        logger.warning(
-            "Room %s: timed out after %ss waiting for agent response", room_id, timeout
-        )
+    if tools.turn.replied:
+        return False
+    logger.warning(
+        "Room %s: timed out after %ss waiting for agent response", room_id, timeout
+    )
+    return True
 
 
 def _message_parts(data: Any) -> tuple[str, list[str]]:
@@ -194,6 +198,8 @@ def _message_parts(data: Any) -> tuple[str, list[str]]:
 
 
 async def _back_off(deadline: float) -> None:
-    backoff = min(EMPTY_READ_BACKOFF_SECONDS, deadline - time.perf_counter())
+    backoff = min(
+        EMPTY_READ_BACKOFF_SECONDS, deadline - asyncio.get_running_loop().time()
+    )
     if backoff > 0:
         await asyncio.sleep(backoff)