diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index 6192487f1..738eaa8d2 100644 --- a/src/utils/agents/streaming.py +++ b/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( diff --git a/src/utils/stream_interrupts.py b/src/utils/stream_interrupts.py index 6b4b11a1b..76f6fa85e 100644 --- a/src/utils/stream_interrupts.py +++ b/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,19 +64,16 @@ 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. @@ -89,22 +81,23 @@ def register_stream( 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 diff --git a/tests/integration/endpoints/test_stream_interrupt_integration.py b/tests/integration/endpoints/test_stream_interrupt_integration.py index caadd8099..c05df745f 100644 --- a/tests/integration/endpoints/test_stream_interrupt_integration.py +++ b/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), diff --git a/tests/integration/endpoints/test_turn_persistence.py b/tests/integration/endpoints/test_turn_persistence.py index dcb933b45..d01f45fc4 100644 --- a/tests/integration/endpoints/test_turn_persistence.py +++ b/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 diff --git a/tests/unit/app/endpoints/test_stream_interrupt.py b/tests/unit/app/endpoints/test_stream_interrupt.py index 0c9f005f2..5353873c4 100644 --- a/tests/unit/app/endpoints/test_stream_interrupt.py +++ b/tests/unit/app/endpoints/test_stream_interrupt.py @@ -1,6 +1,5 @@ """Unit tests for streaming query interrupt endpoint.""" -import asyncio import threading from collections.abc import Generator from typing import Any @@ -11,6 +10,7 @@ InMemorySpanExporter, ) from opentelemetry.trace import StatusCode +from pydantic_ai import CancellationToken from pytest_mock import MockerFixture from app.endpoints.stream_interrupt import stream_interrupt_endpoint_handler @@ -56,12 +56,8 @@ async def test_stream_interrupt_endpoint_success( registry: StreamInterruptRegistry, ) -> None: """Interrupt endpoint cancels an active stream for the same user.""" - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) - registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, task) + token = CancellationToken() + registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, token) response = await stream_interrupt_endpoint_handler( interrupt_request=StreamingInterruptRequest(request_id=REQUEST_ID_SUCCESS), @@ -72,9 +68,7 @@ async def pending_stream() -> None: assert isinstance(response, StreamingInterruptResponse) assert response.request_id == REQUEST_ID_SUCCESS assert response.interrupted is True - - with pytest.raises(asyncio.CancelledError): - await task + assert token.cancelled is True @pytest.mark.asyncio @@ -104,15 +98,11 @@ async def test_stream_interrupt_endpoint_wrong_user( registry: StreamInterruptRegistry, ) -> None: """Interrupt endpoint does not cancel streams owned by other users.""" - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) + token = CancellationToken() registry.register_stream( request_id=REQUEST_ID_WRONG_USER, user_id=OWNER_USER_ID, - task=task, + token=token, ) with pytest.raises(HTTPException) as exc_info: @@ -130,25 +120,17 @@ async def pending_stream() -> None: ) assert exc_info.value.status_code == 403 - assert task.done() is False - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task + assert token.cancelled is False @pytest.mark.asyncio async def test_stream_interrupt_endpoint_already_completed( registry: StreamInterruptRegistry, ) -> None: - """Interrupt endpoint reports already-completed streams without error.""" - - async def completed_stream() -> None: - return None - - task = asyncio.create_task(completed_stream()) - await task - registry.register_stream(REQUEST_ID_ALREADY_COMPLETED, OWNER_USER_ID, task) + """Interrupt endpoint reports already-cancelled streams without error.""" + token = CancellationToken() + token.cancel() + registry.register_stream(REQUEST_ID_ALREADY_COMPLETED, OWNER_USER_ID, token) response = await stream_interrupt_endpoint_handler( interrupt_request=StreamingInterruptRequest( @@ -163,17 +145,11 @@ async def completed_stream() -> None: assert response.interrupted is False -@pytest.mark.asyncio -async def test_stream_interrupt_registry_concurrent_cancel_and_deregister( +def test_stream_interrupt_registry_concurrent_cancel_and_deregister( registry: StreamInterruptRegistry, ) -> None: """Concurrent cancel and deregister do not raise under the registry lock.""" - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) - registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, task) + registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, CancellationToken()) barrier = threading.Barrier(2) errors: list[Exception] = [] @@ -201,23 +177,14 @@ def deregister_in_thread() -> None: CancelStreamResult.NOT_FOUND, ) - if not task.done(): - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - @pytest.mark.asyncio async def test_stream_interrupt_endpoint_double_interrupt( registry: StreamInterruptRegistry, ) -> None: """Second interrupt on the same stream returns interrupted=False.""" - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) - registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, task) + token = CancellationToken() + registry.register_stream(REQUEST_ID_SUCCESS, OWNER_USER_ID, token) first_response = await stream_interrupt_endpoint_handler( interrupt_request=StreamingInterruptRequest(request_id=REQUEST_ID_SUCCESS), @@ -225,9 +192,7 @@ async def pending_stream() -> None: registry=registry, ) assert first_response.interrupted is True - - with pytest.raises(asyncio.CancelledError): - await task + assert token.cancelled is True second_response = await stream_interrupt_endpoint_handler( interrupt_request=StreamingInterruptRequest(request_id=REQUEST_ID_SUCCESS), @@ -250,15 +215,11 @@ async def test_interrupt_success_emits_span_with_attributes( """A successful interrupt emits one span with the expected attributes.""" tracer, exporter = otel mocker.patch("app.endpoints.stream_interrupt.tracer", tracer) - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) + token = CancellationToken() registry.register_stream( REQUEST_ID_SUCCESS, OWNER_USER_ID, - task, + token, conversation_id=CONVERSATION_ID, ) @@ -278,9 +239,7 @@ async def pending_stream() -> None: assert attrs["stream.conversation.id"] == CONVERSATION_ID assert attrs["interrupt.result"] == CancelStreamResult.CANCELLED.value assert span.status.status_code == StatusCode.OK - - with pytest.raises(asyncio.CancelledError): - await task + assert token.cancelled is True @pytest.mark.asyncio async def test_interrupt_not_found_sets_error_status( @@ -322,15 +281,11 @@ async def test_interrupt_wrong_user_sets_error_and_hides_conversation( """A non-owner interrupt records forbidden and does not leak conversation id.""" tracer, exporter = otel mocker.patch("app.endpoints.stream_interrupt.tracer", tracer) - - async def pending_stream() -> None: - await asyncio.sleep(10) - - task = asyncio.create_task(pending_stream()) + token = CancellationToken() registry.register_stream( request_id=REQUEST_ID_WRONG_USER, user_id=OWNER_USER_ID, - task=task, + token=token, conversation_id=CONVERSATION_ID, ) @@ -351,10 +306,7 @@ async def pending_stream() -> None: # A non-owner must not see the stream's conversation id. assert "stream.conversation.id" not in attrs assert span.status.status_code == StatusCode.ERROR - - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task + assert token.cancelled is False @pytest.mark.asyncio async def test_interrupt_already_done_records_ok_status( @@ -363,19 +315,15 @@ async def test_interrupt_already_done_records_ok_status( registry: StreamInterruptRegistry, otel: tuple[Any, InMemorySpanExporter], ) -> None: - """An already-completed stream records already_done with an OK status.""" + """An already-cancelled stream records already_done with an OK status.""" tracer, exporter = otel mocker.patch("app.endpoints.stream_interrupt.tracer", tracer) - - async def completed_stream() -> None: - return None - - task = asyncio.create_task(completed_stream()) - await task + token = CancellationToken() + token.cancel() registry.register_stream( REQUEST_ID_ALREADY_COMPLETED, OWNER_USER_ID, - task, + token, conversation_id=CONVERSATION_ID, ) diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index ed08037f0..1c9113b48 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -216,14 +216,9 @@ def patch_streaming_configuration_fixture(mocker: MockerFixture) -> None: @pytest.fixture(autouse=True, name="stream_interrupt_mocks") def stream_interrupt_mocks_fixture(mocker: MockerFixture) -> dict[str, Any]: - """Patch stream interrupt registry and deregister for wrapper tests.""" - registry = mocker.Mock() - mocker.patch( - "utils.stream_interrupts.get_stream_interrupt_registry", - return_value=registry, - ) + """Patch deregister_stream for wrapper tests.""" deregister = mocker.patch("utils.agents.streaming.deregister_stream") - return {"registry": registry, "deregister": deregister} + return {"deregister": deregister} class TestSerializeEvent: @@ -740,10 +735,6 @@ async def inner() -> AsyncIterator[str]: "utils.agents.streaming.persist_interrupted_turn", new=mocker.AsyncMock(), ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) result = [ event @@ -788,10 +779,6 @@ async def inner() -> AsyncIterator[str]: "utils.agents.error_handler.map_agent_inference_error", return_value=mock_error, ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) result = [ event @@ -809,48 +796,6 @@ async def inner() -> AsyncIterator[str]: assert _sse_event_types(result) == ["token", "error"] stream_interrupt_mocks["deregister"].assert_called_once_with(context.request_id) - @pytest.mark.asyncio - async def test_interrupt_guard_skips_duplicate_persist( - self, - mocker: MockerFixture, - make_generator_context: Callable[..., ResponseGeneratorContext], - responses_params: ResponsesApiParams, - ) -> None: - """Test persist guard prevents double persistence when already handled.""" - context = make_generator_context() - turn_summary = TurnSummary() - - async def inner() -> AsyncIterator[str]: - yield serialize_event( - TokenStreamPayload.create(chunk_id=0, token="partial"), - MEDIA_TYPE_JSON, - ) - raise asyncio.CancelledError() - - persist_mock = mocker.patch( - "utils.agents.streaming.persist_interrupted_turn", - new=mocker.AsyncMock(), - ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[True], - ) - - result = [ - event - async for event in generate_agent_response( - inner(), - context, - responses_params, - turn_summary, - [], - root_span=_dummy_root_span(), - ) - ] - - assert _sse_event_types(result) == ["start", "token", "token", "interrupted"] - persist_mock.assert_not_awaited() - class TestGenerateAgentResponseOtel: """Tests for OTEL instrumentation in generate_agent_response.""" @@ -1021,7 +966,6 @@ async def test_sets_tool_call_span_attributes( @pytest.mark.asyncio async def test_span_ended_on_stream_error( self, - mocker: MockerFixture, make_generator_context: Callable[..., ResponseGeneratorContext], responses_params: ResponsesApiParams, otel: tuple[Any, InMemorySpanExporter], @@ -1038,11 +982,6 @@ async def inner() -> AsyncIterator[str]: ) raise AgentRunError("inference failure") - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) - [ event async for event in generate_agent_response( @@ -1136,10 +1075,6 @@ async def inner() -> AsyncIterator[str]: "utils.agents.streaming.persist_interrupted_turn", new=mocker.AsyncMock(), ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) [ event @@ -1467,10 +1402,6 @@ async def __aexit__(self, *_args: object) -> None: "utils.agents.streaming.persist_interrupted_turn", new=mocker.AsyncMock(), ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) inner = agent_response_generator( mock_agent, @@ -1537,10 +1468,6 @@ async def inner() -> AsyncIterator[str]: "utils.agents.streaming.persist_interrupted_turn", new=mocker.AsyncMock(), ) - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", - return_value=[False], - ) result = [ event @@ -1709,61 +1636,6 @@ async def inner() -> AsyncIterator[str]: assert stored == [] - @pytest.mark.asyncio - async def test_interrupt_guard_prevents_double_persistence( - self, - mocker: MockerFixture, - make_generator_context: Callable[..., ResponseGeneratorContext], - make_responses_params: Callable[..., ResponsesApiParams], - ) -> None: - """An interrupt that already persisted the turn blocks the success path. - - Both protections agree here: the guard shared with the interrupt path - is set, and the turn is settled. The next test opens the guard. - """ - context = make_generator_context() - compacted_params = make_responses_params(omit_conversation=True) - stored = capture_conversation_writes(context) - self._patch_finalizers(mocker) - # The interrupt path has already persisted this turn. - turn = PendingTurn.for_request( - context.client, compacted_params, "the original question" - ) - await turn.store_interrupted("The ans") - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", return_value=[True] - ) - - turn_summary = TurnSummary() - turn_summary.token_usage = TokenCounter(input_tokens=1, output_tokens=1) - turn_summary.output_items = [ - OpenAIResponseMessage(role="assistant", content="The answer.") - ] - - async def inner() -> AsyncIterator[str]: - yield serialize_event( - TokenStreamPayload.create(chunk_id=0, token="x"), MEDIA_TYPE_JSON - ) - - _ = [ - event - async for event in generate_agent_response( - inner(), - context, - compacted_params, - turn_summary, - [], - turn=turn, - root_span=_dummy_root_span(), - ) - ] - - texts = [str(item) for item in stored] - assert len(stored) == 2, f"expected the interrupted turn only, got {texts}" - assert "the original question" in texts[0] - assert "The ans" in texts[1] - assert "The answer." not in texts[1] - @pytest.mark.asyncio async def test_settled_turn_is_not_stored_again_at_the_end_of_the_stream( self, @@ -1785,9 +1657,6 @@ async def test_settled_turn_is_not_stored_again_at_the_end_of_the_stream( context.client, compacted_params, "the original question" ) await turn.store_interrupted("The ans") - mocker.patch( - "utils.agents.streaming.register_interrupt_callback", return_value=[False] - ) turn_summary = TurnSummary() turn_summary.token_usage = TokenCounter(input_tokens=1, output_tokens=1) diff --git a/tests/unit/utils/test_stream_interrupts.py b/tests/unit/utils/test_stream_interrupts.py index 62707b922..85541cf6e 100644 --- a/tests/unit/utils/test_stream_interrupts.py +++ b/tests/unit/utils/test_stream_interrupts.py @@ -1,10 +1,12 @@ """Unit tests for stream interrupt registry and persistence utilities.""" import asyncio +from collections.abc import Generator from typing import Any import pytest from ogx_api.openai_responses import OpenAIResponseMessage +from pydantic_ai import CancellationToken from pytest_mock import MockerFixture from constants import INTERRUPTED_RESPONSE_MESSAGE @@ -15,15 +17,95 @@ from models.common.turn_summary import TurnSummary from utils.pending_turn import PendingTurn from utils.stream_interrupts import ( + CancelStreamResult, StreamInterruptRegistry, build_interrupted_response, persist_interrupted_turn, - register_interrupt_callback, ) from utils.token_estimator import extract_message_text INTERRUPTED_INDICATOR = f"\n\n*{INTERRUPTED_RESPONSE_MESSAGE}*" +REQUEST_ID = "123e4567-e89b-12d3-a456-426614174005" +OWNER_USER_ID = "00000001-0001-0001-0001-000000000001" +NON_OWNER_USER_ID = "00000001-0001-0001-0001-000000000999" + + +@pytest.fixture(name="registry") +def registry_fixture() -> Generator[StreamInterruptRegistry, None, None]: + """Provide the singleton registry without this module's test stream.""" + registry = StreamInterruptRegistry() + registry.deregister_stream(REQUEST_ID) + yield registry + registry.deregister_stream(REQUEST_ID) + + +@pytest.mark.parametrize( + ("state", "caller", "expected_result", "expected_cancelled"), + [ + pytest.param( + "absent", + OWNER_USER_ID, + CancelStreamResult.NOT_FOUND, + False, + id="not_registered", + ), + pytest.param( + "deregistered", + OWNER_USER_ID, + CancelStreamResult.NOT_FOUND, + False, + id="deregistered", + ), + pytest.param( + "registered", + NON_OWNER_USER_ID, + CancelStreamResult.FORBIDDEN, + False, + id="other_user", + ), + pytest.param( + "cancelled", + NON_OWNER_USER_ID, + CancelStreamResult.FORBIDDEN, + True, + id="other_user_already_cancelled", + ), + pytest.param( + "cancelled", + OWNER_USER_ID, + CancelStreamResult.ALREADY_DONE, + True, + id="already_cancelled", + ), + pytest.param( + "registered", + OWNER_USER_ID, + CancelStreamResult.CANCELLED, + True, + id="owner", + ), + ], +) +def test_cancel_stream_outcomes( + registry: StreamInterruptRegistry, + state: str, + caller: str, + expected_result: CancelStreamResult, + expected_cancelled: bool, +) -> None: + """cancel_stream checks ownership and cancels the registered token once.""" + token = CancellationToken() + if state == "cancelled": + token.cancel() + if state != "absent": + registry.register_stream(REQUEST_ID, OWNER_USER_ID, token) + if state == "deregistered": + registry.deregister_stream(REQUEST_ID) + + assert registry.cancel_stream(REQUEST_ID, caller) == expected_result + assert token.cancelled is expected_cancelled + def _params( conversation: str, @@ -185,54 +267,6 @@ async def test_persist_interrupted_turn_schedules_background_topic_summary( ) -def test_register_interrupt_callback_registers_current_task( - mocker: MockerFixture, -) -> None: - """register_interrupt_callback binds the current asyncio task to the registry.""" - registry = mocker.Mock(spec=StreamInterruptRegistry) - mocker.patch( - "utils.stream_interrupts.get_stream_interrupt_registry", - return_value=registry, - ) - persist_mock = mocker.patch( - "utils.stream_interrupts.persist_interrupted_turn", - new=mocker.AsyncMock(), - ) - - context = mocker.Mock(spec=ResponseGeneratorContext) - context.request_id = "req-1" - context.user_id = "user_1" - context.conversation_id = "conv-1" - responses_params = mocker.Mock(spec=ResponsesApiParams) - turn_summary = TurnSummary() - background_tasks: list[asyncio.Task[None]] = [] - - async def run() -> list[bool]: - return register_interrupt_callback( - context, - responses_params, - turn_summary, - background_tasks, - ) - - guard = asyncio.run(run()) - - assert guard == [False] - registry.register_stream.assert_called_once() - assert registry.register_stream.call_args.kwargs["request_id"] == "req-1" - assert registry.register_stream.call_args.kwargs["user_id"] == "user_1" - assert registry.register_stream.call_args.kwargs["conversation_id"] == "conv-1" - assert persist_mock.await_count == 0 - - on_interrupt = registry.register_stream.call_args.kwargs["on_interrupt"] - - async def invoke_callback() -> None: - await on_interrupt() - - asyncio.run(invoke_callback()) - persist_mock.assert_awaited_once() - - class TestBuildInterruptedResponse: """Tests for build_interrupted_response helper."""