Feature/stream interrupt

src/utils/agents/streaming.py

@@ -86,7 +86,6 @@
     build_interrupted_response,
     deregister_stream,
     persist_interrupted_turn,
-    register_interrupt_callback,
 )
 from utils.token_counter import TokenCounter
 
@@ -236,13 +235,7 @@ async def generate_agent_response(  # pylint: disable=too-many-statements
     media_type = context.query_request.media_type or MEDIA_TYPE_JSON
     if turn is None:
         turn = PendingTurn.for_request(context.client, responses_params)
-    persist_guard = register_interrupt_callback(
-        context,
-        responses_params,
-        turn_summary,
-        background_topic_summary_tasks,
-        turn,
-    )
+    persist_guard: list[bool] = [False]
     stream_completed = False
     if emit_start:
         yield serialize_event(

src/utils/stream_interrupts.py

@@ -2,11 +2,12 @@
 
 import asyncio
 import datetime
-from collections.abc import Callable, Coroutine
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from enum import StrEnum
 from threading import Lock
-from typing import Any, Optional
+from typing import Optional
+
+from pydantic_ai import CancellationToken
 
 from constants import (
     INTERRUPTED_RESPONSE_MESSAGE,
@@ -31,20 +32,14 @@ class ActiveStream:
 
     Attributes:
         user_id: Owner of the streaming request.
-        task: Asyncio task producing the stream response.
-        on_interrupt: Optional async callback invoked when the stream
-            is cancelled, scheduled as a separate task so it runs
-            regardless of where the ``CancelledError`` lands.
+        token: Cancellation token of the agent run producing the stream.
         conversation_id: Conversation the streaming request belongs to,
             surfaced for observability of interrupt requests.
 
     """
 
     user_id: str
-    task: asyncio.Task[None]
-    on_interrupt: Optional[Callable[[], Coroutine[Any, Any, None]]] = field(
-        default=None, repr=False
-    )
+    token: CancellationToken
     conversation_id: Optional[str] = None
 
 
@@ -58,7 +53,7 @@ class CancelStreamResult(StrEnum):
 
 
 class StreamInterruptRegistry(metaclass=Singleton):
-    """Registry for active streaming tasks keyed by request ID."""
+    """Ownership and cancellation index of active streams keyed by request ID."""
 
     def __init__(self) -> None:
         """Initialize an empty registry with a lock for thread-safety."""
@@ -69,42 +64,40 @@ def register_stream(
         self,
         request_id: str,
         user_id: str,
-        task: asyncio.Task[None],
-        on_interrupt: Optional[Callable[[], Coroutine[Any, Any, None]]] = None,
+        token: CancellationToken,
         conversation_id: Optional[str] = None,
     ) -> None:
-        """Register an active stream task for interrupt support.
+        """Register an active stream for interrupt support.
 
         Parameters:
         ----------
             request_id: Unique streaming request identifier.
             user_id: User identifier that owns the stream.
-            task: Asyncio task associated with the stream.
-            on_interrupt: Optional async callback to run when the stream
-                is cancelled, executed in a separate task.
+            token: Cancellation token passed to the agent run of the stream.
             conversation_id: Conversation the stream belongs to, retained
                 for observability of subsequent interrupt requests.
 
         """
         with self._lock:
             self._streams[request_id] = ActiveStream(
                 user_id=user_id,
-                task=task,
-                on_interrupt=on_interrupt,
+                token=token,
                 conversation_id=conversation_id,
             )
 
     def cancel_stream(self, request_id: str, user_id: str) -> CancelStreamResult:
-        """Cancel an active stream owned by user.
+        """Cancel the token of an active stream owned by user.
 
         The entire lookup-check-cancel sequence is performed under the
         lock so that a concurrent ``deregister_stream`` cannot remove
-        the entry between the ownership check and the cancel call.
-
-        When an ``on_interrupt`` callback was registered, it is
-        scheduled as a **separate** asyncio task after the cancel so
-        persistence runs regardless of where the ``CancelledError``
-        is raised (inside the generator or in Starlette's send).
+        the entry between the ownership check and the cancel call, and
+        so that only one caller sees a token flip to cancelled.
+        Cancelling the token asks the agent run to stop; the streaming
+        request itself finalizes the interrupted turn. PAI ignores a
+        cancel once the run has finished, so an entry that is still
+        registered after its run completed also returns ``CANCELLED``;
+        the streaming request should deregister as soon as its agent
+        run ends.
 
         Parameters:
         ----------
@@ -113,10 +106,11 @@ def cancel_stream(self, request_id: str, user_id: str) -> CancelStreamResult:
 
         Returns:
         -------
-            CancelStreamResult: Structured cancellation result.
+            CancelStreamResult: ``NOT_FOUND`` for an unknown request,
+            ``FORBIDDEN`` when another user owns it, ``ALREADY_DONE``
+            when its token was already cancelled, ``CANCELLED`` otherwise.
 
         """
-        on_interrupt = None
         with self._lock:
             stream = self._streams.get(request_id)
             if stream is None:
@@ -128,18 +122,13 @@ def cancel_stream(self, request_id: str, user_id: str) -> CancelStreamResult:
                     request_id,
                 )
                 return CancelStreamResult.FORBIDDEN
-            if stream.task.done():
+            if stream.token.cancelled:
                 return CancelStreamResult.ALREADY_DONE
-            stream.task.cancel()
-            on_interrupt = stream.on_interrupt
-
-        if on_interrupt is not None:
-            asyncio.get_running_loop().create_task(on_interrupt())
-
+            stream.token.cancel()
         return CancelStreamResult.CANCELLED
 
     def deregister_stream(self, request_id: str) -> None:
-        """Remove stream task from registry once completed/cancelled.
+        """Remove stream from registry once completed/cancelled.
 
         Parameters:
         ----------
@@ -325,71 +314,3 @@ async def persist_interrupted_turn(
             "Failed to store interrupted query results for request %s",
             context.request_id,
         )
-
-
-def register_interrupt_callback(
-    context: ResponseGeneratorContext,
-    responses_params: ResponsesApiParams,
-    turn_summary: TurnSummary,
-    background_topic_summary_tasks: list[asyncio.Task[None]],
-    turn: Optional[PendingTurn] = None,
-) -> list[bool]:
-    """Build an interrupt callback and register the stream for cancellation.
-
-    The callback is invoked by ``cancel_stream`` when the client
-    interrupts, so persistence runs regardless of where the
-    ``CancelledError`` is raised in the ASGI stack.
-
-    A mutable one-element list is used as a shared guard so the
-    callback and the in-generator ``CancelledError`` handler never
-    both persist the same turn.
-
-    Parameters:
-    ----------
-        context: The response generator context.
-        responses_params: The Responses API parameters.
-        turn_summary: TurnSummary populated during streaming.
-        background_topic_summary_tasks: Mutable list tracking fire-and-forget
-            topic summary tasks for graceful shutdown.
-        turn: The pending turn of the request; required in compacted mode.
-
-    Returns:
-    -------
-        A mutable list ``[False]`` used as a persist-done guard; the
-        caller should check ``guard[0]`` before persisting and set
-        it to ``True`` afterwards.
-
-    """
-    guard: list[bool] = [False]
-
-    async def _on_interrupt() -> None:
-        if guard[0]:
-            return
-        guard[0] = True
-        full_text, _ = build_interrupted_response(turn_summary.partial_tokens)
-        turn_summary.llm_response = full_text
-        await persist_interrupted_turn(
-            context,
-            responses_params,
-            turn_summary,
-            background_topic_summary_tasks,
-            turn,
-        )
-
-    current_task = asyncio.current_task()
-    if current_task is not None:
-        get_stream_interrupt_registry().register_stream(
-            request_id=context.request_id,
-            user_id=context.user_id,
-            task=current_task,
-            on_interrupt=_on_interrupt,
-            conversation_id=context.conversation_id,
-        )
-    else:
-        logger.warning(
-            "No current asyncio task for request %s; "
-            "stream interruption will not be available",
-            context.request_id,
-        )
-
-    return guard

tests/integration/endpoints/test_stream_interrupt_integration.py

@@ -1,11 +1,11 @@
 """Integration tests for the streaming query interrupt lifecycle."""
 
-import asyncio
 from collections.abc import Generator
 
 import pytest
 from fastapi import HTTPException, status
 from fastapi.testclient import TestClient
+from pydantic_ai import CancellationToken
 
 from app.endpoints.stream_interrupt import stream_interrupt_endpoint_handler
 from configuration import AppConfig
@@ -35,11 +35,8 @@ async def test_stream_interrupt_full_round_trip(
     # test_config loads configuration so @authorize on the handler can resolve.
     _ = test_config
 
-    async def pending_stream() -> None:
-        await asyncio.sleep(10)
-
-    task = asyncio.create_task(pending_stream())
-    registry.register_stream(TEST_REQUEST_ID, OWNER_USER_ID, task)
+    token = CancellationToken()
+    registry.register_stream(TEST_REQUEST_ID, OWNER_USER_ID, token)
 
     assert registry.get_stream(TEST_REQUEST_ID) is not None
 
@@ -49,9 +46,7 @@ async def pending_stream() -> None:
         registry=registry,
     )
     assert response.interrupted is True
-
-    with pytest.raises(asyncio.CancelledError):
-        await task
+    assert token.cancelled is True
 
     completed_response = await stream_interrupt_endpoint_handler(
         interrupt_request=StreamingInterruptRequest(request_id=TEST_REQUEST_ID),

tests/integration/endpoints/test_turn_persistence.py

@@ -17,6 +17,7 @@
 # pylint: disable=too-many-lines
 
 import asyncio
+import contextlib
 import json
 from collections.abc import AsyncIterator
 from datetime import UTC, datetime
@@ -80,9 +81,7 @@
 )
 from utils.pending_turn import TurnNotStoredError
 from utils.stream_interrupts import (
-    CancelStreamResult,
     build_interrupted_response,
-    get_stream_interrupt_registry,
 )
 from utils.token_estimator import extract_message_text
 
@@ -354,18 +353,6 @@ async def __aexit__(self, *_args: object) -> None:
     return _RunStreamCtx()
 
 
-def _request_id_of(chunks: list[str]) -> str:
-    """Return the request id the stream announced in its start event."""
-    for chunk in chunks:
-        for line in chunk.splitlines():
-            if not line.startswith("data: "):
-                continue
-            event = json.loads(line[len("data: ") :])
-            if event.get("event") == "start":
-                return event["data"]["request_id"]
-    raise AssertionError(f"no start event in {chunks!r}")
-
-
 async def _send_streaming_query(test_request: Request, test_auth: AuthTuple) -> Any:
     """Send the new query on the stored conversation."""
     return await streaming_query_endpoint_handler(
@@ -391,11 +378,11 @@ async def _consume() -> None:
 
     consumer = asyncio.create_task(_consume())
     await asyncio.wait_for(first_token.wait(), timeout=5)
-    result = get_stream_interrupt_registry().cancel_stream(
-        _request_id_of(chunks), test_auth[0]
-    )
-    assert result == CancelStreamResult.CANCELLED
-    await asyncio.wait_for(consumer, timeout=5)
+    # The registry holds no cancellation token for the run until LCORE-4733,
+    # so the stream is cancelled the way a client disconnect cancels it.
+    consumer.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await asyncio.wait_for(consumer, timeout=5)
     # The interrupt callback runs as a task of its own; let it finish.
     pending = [
         task