Fix task cancellation bug in graph beta API triggered by using MCPServerStreamableHTTP with agent.run_stream

[dmontagu]

Addresses an issue related to MCPServerStreamableHTTP cancellation during calls to agent.run_stream. @DouweM can add a test if he wants. It was painful to debug this issue.

The problem was that this script was erroring before this PR:

import asyncio

import logfire

from pydantic_ai import Agent
from pydantic_ai.mcp import MCPServerStreamableHTTP
from pydantic_ai.models.openai import OpenAIChatModel

logfire.configure()
logfire.instrument_pydantic_ai()
logfire.instrument_mcp()
logfire.instrument_httpx(capture_all=True)


model = OpenAIChatModel('gpt-4o')
agent = Agent(
    model=model,
    toolsets=[MCPServerStreamableHTTP(url='https://mcp.deepwiki.com/mcp')],
)


async def main():
    async with agent.run_stream('Tell me about the pydantic/pydantic-ai repo.') as result:
        output = await result.get_output()
    print(output)
    # async for event in agent.run_stream_events('Tell me about the pydantic/pydantic-ai repo.'):
    #     if isinstance(event, AgentRunResultEvent):
    #         print(event.result.output)


if __name__ == '__main__':
    asyncio.run(main())

This fixes the above script.

[github-actions]

Docs Preview

commit:	b39b168
Preview URL:	https://f703180f-pydantic-ai-previews.pydantic.workers.dev