from typing import Any

from unittest.mock import patch

import anyio

import pytest

from mcp.shared.message import SessionMessage

from mcp.shared.session import BaseSession, RequestId, SendResultT

from mcp.types import ClientNotification, ClientRequest, ClientResult, EmptyResult, ErrorData, PingRequest

@pytest.mark.anyio

async def test_send_request_stream_cleanup():

"""

Test that send_request properly cleans up streams when an exception occurs.

This test mocks out most of the session functionality to focus on stream cleanup.

"""

# Create a mock session with the minimal required functionality

class TestSession(BaseSession[ClientRequest, ClientNotification, ClientResult, Any, Any]):

async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:

pass

# Create streams

write_stream_send, write_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)

read_stream_send, read_stream_receive = anyio.create_memory_object_stream[SessionMessage](1)

# Create the session

session = TestSession(

read_stream_receive,

write_stream_send,

object, # Request type doesn't matter for this test

object, # Notification type doesn't matter for this test

)

# Create a test request

request = ClientRequest(PingRequest())

# Patch the _write_stream.send method to raise an exception

async def mock_send(*args: Any, **kwargs: Any):

raise RuntimeError("Simulated network error")

# Record the response streams before the test

initial_stream_count = len(session._response_streams)

# Run the test with the patched method

with patch.object(session._write_stream, "send", mock_send):

with pytest.raises(RuntimeError):

await session.send_request(request, EmptyResult)

# Verify that no response streams were leaked

assert len(session._response_streams) == initial_stream_count, (

f"Expected {initial_stream_count} response streams after request, but found {len(session._response_streams)}"

)

# Clean up

await write_stream_send.aclose()

await write_stream_receive.aclose()

await read_stream_send.aclose()

await read_stream_receive.aclose()