Implement Stream Cancellation #2901

willgdjones · 2025-09-15T08:40:32Z

⚠️ (Generated by Cursor - in the process of being edited and refined)

Pydantic AI Stream Cancellation

This implementation adds stream cancellation functionality, allowing users to cancel streaming responses when clients disconnect or explicitly request cancellation.

🎯 Problem Solved

Previously, when users broke early from a streaming loop, Pydantic AI would continue consuming the entire response in the background to ensure proper usage tracking. This led to:

Wasted Resources: Unnecessary compute and network usage
Poor User Experience: No way to stop long-running streams
Memory Issues: Streams continuing even after clients disconnect

✨ Features

Explicit Cancellation API: Call await stream.cancel() to stop streaming
Automatic HTTP Disconnect Handling: Streams cancel when web clients disconnect
Partial Usage Tracking: Accurate token counts for cancelled streams
Exception Safety: Multiple cancel calls are safe and idempotent
OpenAI Support: Initial implementation supports OpenAI models

Basic Usage

import asyncio
from pydantic_ai import Agent
from pydantic_ai.exceptions import StreamCancelled

async def basic_cancellation():
    agent = Agent("openai:gpt-4o-mini")
    
    try:
        async with agent.run_stream("Tell me a long story") as result:
            chunk_count = 0
            async for content in result.stream_text(delta=True):
                print(content)
                chunk_count += 1
                
                # Cancel after 3 chunks
                if chunk_count >= 3:
                    await result.cancel()
                    
    except StreamCancelled as e:
        print(f"Stream cancelled: {e}")
        print(f"Partial usage: {result.usage()}")

asyncio.run(basic_cancellation())

📚 API Reference

AgentStream.cancel()

async def cancel(self) -> None:
    """Cancel the streaming response.
    
    This will close the underlying network connection and cause any active iteration
    over the stream to raise a StreamCancelled exception.
    
    Subsequent calls to cancel() are safe and will not raise additional exceptions.
    """

StreamCancelled Exception

class StreamCancelled(Exception):
    """Exception raised when a streaming response is cancelled."""
    
    def __init__(self, message: str = "Stream was cancelled"):
        self.message = message
        super().__init__(message)

🏗️ Implementation Details

Architecture

The implementation consists of several components:

StreamCancelled Exception (exceptions.py)
- New exception type for cancelled streams
AgentStream.cancel() (result.py)
- Public API for cancelling streams
- Sets internal cancellation flag
- Delegates to underlying StreamedResponse
StreamedResponse.cancel() (models/__init__.py)
- Abstract base method for model-specific cancellation
- Default no-op implementation
OpenAIStreamedResponse.cancel() (models/openai.py)
- OpenAI-specific cancellation implementation
- Marks stream as cancelled, causing iterator to raise exception
Cancellation-Aware Iterator (result.py)
- Checks cancellation flag before yielding events
- Raises StreamCancelled when cancelled
Agent Graph Updates (_agent_graph.py)
- Handles StreamCancelled in automatic consumption logic

Usage Tracking

Partial Usage: Cancelled streams report accurate token usage up to cancellation point
No Double Counting: Usage is accumulated as chunks are processed
Metadata: Usage objects can indicate partial/cancelled state

Error Handling

Idempotent Cancellation: Multiple cancel() calls are safe
Exception Propagation: StreamCancelled bubbles up through iteration
Resource Cleanup: Network connections are properly closed

DouweM

@willgdjones Thanks Will. I get where Cursor is going with this, but I'm not sure if it's doing too little (shouldn't we be calling openai.AgentStream.close() at some point?) and/or too much (do we need the multiple canceled booleans and exception, if the wrapped stream could just stop yielding events). Would be good to get your (human, not AI!) take :)

DouweM · 2025-09-15T17:24:01Z

pydantic_ai_slim/pydantic_ai/exceptions.py

+    """Exception raised when a streaming response is cancelled."""
+
+    def __init__(self, message: str = 'Stream was cancelled'):
+        self.message = message


No need to have a message as it's not used anywhere

DouweM · 2025-09-15T17:24:45Z

pydantic_ai_slim/pydantic_ai/models/__init__.py

+        This should close the underlying network connection and cause any active iteration
+        to raise a StreamCancelled exception. The default implementation is a no-op.
+        """
+        pass


I think this should raise NotImplementedError to not silently keep the stream going when the user thought they canceled it

DouweM · 2025-09-15T17:25:08Z

pydantic_ai_slim/pydantic_ai/models/openai.py


    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        async for chunk in self._response:
+            # Check for cancellation before processing each chunk
+            if self._cancelled:


Shouldn't we do this after recording the usage?

DouweM · 2025-09-15T17:25:29Z

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -1418,6 +1423,14 @@ def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._timestamp

+    async def cancel(self) -> None:


This doesn't need to be async if the recommended behavior is to always just set a flag and then cancel on the next iteration

DouweM · 2025-09-15T17:29:57Z

pydantic_ai_slim/pydantic_ai/models/openai.py


    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        async for chunk in self._response:
+            # Check for cancellation before processing each chunk
+            if self._cancelled:
+                raise StreamCancelled('OpenAI stream was cancelled')


Will this actually cause OpenAI to cleanly close the stream? Shouldn't we call await AsyncStream.close() or something?

They also have their own # Ensure the entire stream is consumed:

https://github.com/openai/openai-python/blob/4756247cee3d9548397b26a29109e76cc9522379/src/openai/_streaming.py#L216-L222

Note that right now, OpenAIStreamedResponse only has access to AsyncIterable[ChatCompletionChunk], but that's derived from the AsyncStream[ChatCompletionChunk]:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/openai.py

Lines 578 to 582 in f903d5b

async def _process_streamed_response(

self, response: AsyncStream[ChatCompletionChunk], model_request_parameters: ModelRequestParameters

) -> OpenAIStreamedResponse:

"""Process a streamed response, and prepare a streaming response to return."""

peekable_response = _utils.PeekableAsyncStream(response)

So in order to access that cancel method, we may need to put it on _utils.PeekableAsyncStream as well, and then forward it to the underlying stream.

DouweM · 2025-09-15T17:32:02Z

pydantic_ai_slim/pydantic_ai/result.py

+        async for item in stream_response:
+            # Check for cancellation first
+            if is_cancelled():
+                raise exceptions.StreamCancelled()


Why do we have to raise it here and inside the StreamedResponse?

Wouldn't the stream_response just stop yielding once it's cancelled and there are no more messages, meaning we shouldn't need to do anythings special here?

DouweM · 2025-09-15T17:33:44Z

pydantic_ai_slim/pydantic_ai/result.py

-            self._agent_stream_iterator = _get_usage_checking_stream_response(
-                self._raw_stream_response, self._usage_limits, self.usage
+            self._agent_stream_iterator = _get_cancellation_aware_stream_response(
+                self._raw_stream_response, self._usage_limits, self.usage, lambda: self._cancelled


Related to what I wrote below, I don't understand why we need this lambda, instead of just pushing the cancelation down to the wrapped StreamedResponse, and then relying on that to stop yielding events

DouweM · 2025-09-15T17:34:57Z

pydantic_ai_slim/pydantic_ai/_agent_graph.py

+            try:
+                async for _ in agent_stream:
+                    pass
+            except exceptions.StreamCancelled:


Do we need this exception at all?

Like I wrote below: Wouldn't the StreamedResponse just stop yielding once it's been cancelled and there are no more messages, meaning we shouldn't need to do anything special here?

DouweM · 2025-09-15T17:38:03Z

tests/models/test_openai.py

+    m = OpenAIChatModel('gpt-4o-mini', provider=OpenAIProvider(openai_client=mock_client))
+    agent = Agent(m)
+
+    async with agent.run_stream('Hello world') as result:


Can we add a test that cancels streaming in the middle of an unfinished tool call, like in the example at https://ai.pydantic.dev/agents/#streaming-events-and-final-output after a ToolCallPartDelta, and then see what the final ModelResponse in result.all_messages() looks like? I imagine it would have an incomplete ToolCallPart. I wonder if we should indicate on the ModelResponse somehow that it's incomplete because it's been canceled, and cannot be used as message_history, for example.

DouweM · 2025-09-15T17:38:17Z

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -641,6 +641,14 @@ def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        raise NotImplementedError()

+    async def cancel(self) -> None:


We should document this feature in the Streaming docs

willgdjones · 2025-09-16T14:15:57Z

Thank you for this detailed critique! I will address these points when I get the chance to.

willgdjones force-pushed the feature/halt-streaming branch 3 times, most recently from e7dab8d to 8fda0ab Compare September 15, 2025 08:58

Implement stream cancellation

fc05434

willgdjones force-pushed the feature/halt-streaming branch from 8fda0ab to fc05434 Compare September 15, 2025 09:06

willgdjones mentioned this pull request Sep 15, 2025

should cancel the response when user stop consuming #1524

Open

2 tasks

DouweM self-assigned this Sep 15, 2025

DouweM requested changes Sep 15, 2025

View reviewed changes

DouweM added the awaiting author revision label Sep 15, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Implement Stream Cancellation #2901

Implement Stream Cancellation #2901

willgdjones commented Sep 15, 2025 •

edited

Loading

Uh oh!

DouweM left a comment

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

DouweM Sep 15, 2025

Uh oh!

willgdjones commented Sep 16, 2025

Uh oh!

Uh oh!

	async def _process_streamed_response(
	self, response: AsyncStream[ChatCompletionChunk], model_request_parameters: ModelRequestParameters
	) -> OpenAIStreamedResponse:
	"""Process a streamed response, and prepare a streaming response to return."""
	peekable_response = _utils.PeekableAsyncStream(response)

Implement Stream Cancellation #2901

Are you sure you want to change the base?

Implement Stream Cancellation #2901

Conversation

willgdjones commented Sep 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pydantic AI Stream Cancellation

🎯 Problem Solved

✨ Features

Basic Usage

📚 API Reference

AgentStream.cancel()

StreamCancelled Exception

🏗️ Implementation Details

Architecture

Usage Tracking

Error Handling

Uh oh!

DouweM left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

willgdjones commented Sep 16, 2025

Uh oh!

Uh oh!

willgdjones commented Sep 15, 2025 •

edited

Loading