google--adk-python/tests/unittests/tools/mcp_tool/test_mcp_session_manager.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import hashlib
import json
import sys
from unittest.mock import ANY
from unittest.mock import AsyncMock
from unittest.mock import Mock
from unittest.mock import patch

from google.adk.platform import thread as platform_thread
from google.adk.tools.mcp_tool.mcp_session_manager import _DebugHttpxClientFactory
from google.adk.tools.mcp_tool.mcp_session_manager import _GoogleAuthAsyncByteStream
from google.adk.tools.mcp_tool.mcp_session_manager import _http_debug_var
from google.adk.tools.mcp_tool.mcp_session_manager import _RefreshableAsyncCredentials
from google.adk.tools.mcp_tool.mcp_session_manager import _SharedAsyncTransport
from google.adk.tools.mcp_tool.mcp_session_manager import create_mcp_http_client
from google.adk.tools.mcp_tool.mcp_session_manager import MCPSessionManager
from google.adk.tools.mcp_tool.mcp_session_manager import retry_on_errors
from google.adk.tools.mcp_tool.mcp_session_manager import SseConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
import httpx
from mcp import StdioServerParameters
import pytest

try:
  from google.auth.aio.transport.sessions import AsyncAuthorizedSession

  AIO_SUPPORTED = True
except ImportError:
  AIO_SUPPORTED = False


class MockClientSession:
  """Mock ClientSession for testing."""

  def __init__(self):
    self._read_stream = Mock()
    self._write_stream = Mock()
    self._read_stream._closed = False
    self._write_stream._closed = False
    self.initialize = AsyncMock()


class MockAsyncExitStack:
  """Mock AsyncExitStack for testing."""

  def __init__(self):
    self.aclose = AsyncMock()
    self.enter_async_context = AsyncMock()

  async def __aenter__(self):
    return self

  async def __aexit__(self, exc_type, exc_val, exc_tb):
    pass


class MockSessionContext:
  """Mock SessionContext for testing."""

  def __init__(self, session=None):
    """Initialize MockSessionContext.

    Args:
        session: The mock session to return from __aenter__ and session property.
    """
    self._session = session
    self._aenter_mock = AsyncMock(return_value=session)
    self._aexit_mock = AsyncMock(return_value=False)

  @property
  def session(self):
    """Get the mock session."""
    return self._session

  async def __aenter__(self):
    """Enter the async context manager."""
    return await self._aenter_mock()

  async def __aexit__(self, exc_type, exc_val, exc_tb):
    """Exit the async context manager."""
    return await self._aexit_mock(exc_type, exc_val, exc_tb)


class TestMCPSessionManager:
  """Test suite for MCPSessionManager class."""

  def setup_method(self):
    """Set up test fixtures."""
    self.mock_stdio_params = StdioServerParameters(
        command="test_command", args=[]
    )
    self.mock_stdio_connection_params = StdioConnectionParams(
        server_params=self.mock_stdio_params, timeout=5.0
    )

  def test_init_with_stdio_server_parameters(self):
    """Test initialization with StdioServerParameters (deprecated)."""
    with patch(
        "google.adk.tools.mcp_tool.mcp_session_manager.logger"
    ) as mock_logger:
      manager = MCPSessionManager(self.mock_stdio_params)

      # Should log deprecation warning
      mock_logger.warning.assert_called_once()
      assert "StdioServerParameters is not recommended" in str(
          mock_logger.warning.call_args
      )

      # Should convert to StdioConnectionParams
      assert isinstance(manager._connection_params, StdioConnectionParams)
      assert manager._connection_params.server_params == self.mock_stdio_params
      assert manager._connection_params.timeout == 5

  def test_init_with_stdio_connection_params(self):
    """Test initialization with StdioConnectionParams."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    assert manager._connection_params == self.mock_stdio_connection_params
    assert manager._errlog == sys.stderr
    assert manager._sessions == {}

  def test_init_with_sse_connection_params(self):
    """Test initialization with SseConnectionParams."""
    sse_params = SseConnectionParams(
        url="https://example.com/mcp",
        headers={"Authorization": "Bearer token"},
        timeout=10.0,
    )
    manager = MCPSessionManager(sse_params)

    assert manager._connection_params == sse_params

  @patch("google.adk.tools.mcp_tool.mcp_session_manager.sse_client")
  def test_init_with_sse_custom_httpx_factory(self, mock_sse_client):
    """Test that sse_client is called with custom httpx_client_factory."""
    custom_httpx_factory = Mock()

    sse_params = SseConnectionParams(
        url="https://example.com/mcp",
        timeout=10.0,
        httpx_client_factory=custom_httpx_factory,
    )
    manager = MCPSessionManager(sse_params)

    manager._create_client()

    mock_sse_client.assert_called_once()
    kwargs = mock_sse_client.call_args.kwargs
    assert kwargs["url"] == "https://example.com/mcp"
    assert kwargs["headers"] is None
    assert kwargs["timeout"] == 10.0
    assert kwargs["sse_read_timeout"] == 300.0
    factory = kwargs["httpx_client_factory"]
    assert isinstance(factory, _DebugHttpxClientFactory)
    assert factory._base_factory == custom_httpx_factory

  @patch("google.adk.tools.mcp_tool.mcp_session_manager.sse_client")
  def test_init_with_sse_default_httpx_factory(self, mock_sse_client):
    """Test that sse_client is called with default httpx_client_factory."""
    sse_params = SseConnectionParams(
        url="https://example.com/mcp",
        timeout=10.0,
    )
    manager = MCPSessionManager(sse_params)

    manager._create_client()

    mock_sse_client.assert_called_once()
    kwargs = mock_sse_client.call_args.kwargs
    assert kwargs["url"] == "https://example.com/mcp"
    assert kwargs["headers"] is None
    assert kwargs["timeout"] == 10.0
    assert kwargs["sse_read_timeout"] == 300.0
    factory = kwargs["httpx_client_factory"]
    assert isinstance(factory, _DebugHttpxClientFactory)
    assert (
        factory._base_factory
        == SseConnectionParams.model_fields[
            "httpx_client_factory"
        ].get_default()
    )

  def test_init_with_streamable_http_params(self):
    """Test initialization with StreamableHTTPConnectionParams."""
    http_params = StreamableHTTPConnectionParams(
        url="https://example.com/mcp", timeout=15.0
    )
    manager = MCPSessionManager(http_params)

    assert manager._connection_params == http_params

  @patch("google.adk.tools.mcp_tool.mcp_session_manager.streamable_http_client")
  def test_init_with_streamable_http_custom_httpx_factory(
      self, mock_streamable_http_client
  ):
    """Test that streamable_http_client is called with custom httpx_client_factory."""
    custom_httpx_factory = Mock()

    http_params = StreamableHTTPConnectionParams(
        url="https://example.com/mcp",
        timeout=15.0,
        httpx_client_factory=custom_httpx_factory,
    )
    manager = MCPSessionManager(http_params)

    manager._create_client()

    mock_streamable_http_client.assert_called_once()
    kwargs = mock_streamable_http_client.call_args.kwargs
    assert kwargs["url"] == "https://example.com/mcp"
    assert kwargs["terminate_on_close"] is True
    assert kwargs["http_client"] is not None
    custom_httpx_factory.assert_called_once()

  @patch("google.adk.tools.mcp_tool.mcp_session_manager.streamable_http_client")
  def test_init_with_streamable_http_default_httpx_factory(
      self, mock_streamable_http_client
  ):
    """Test that streamable_http_client is called with default httpx_client_factory."""
    http_params = StreamableHTTPConnectionParams(
        url="https://example.com/mcp", timeout=15.0
    )
    manager = MCPSessionManager(http_params)

    manager._create_client()

    mock_streamable_http_client.assert_called_once()
    kwargs = mock_streamable_http_client.call_args.kwargs
    assert kwargs["url"] == "https://example.com/mcp"
    assert kwargs["terminate_on_close"] is True
    assert isinstance(kwargs["http_client"], httpx.AsyncClient)

  @patch(
      "google.adk.tools.mcp_tool.mcp_session_manager.HTTPXClientInstrumentor",
      create=True,
  )
  @patch(
      "google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client"
  )
  @patch(
      "google.adk.tools.mcp_tool.mcp_session_manager._HAS_HTTPX_INSTRUMENTOR",
      True,
  )
  def test_default_httpx_factory_instruments_client_when_available(
      self, mock_base_factory, mock_instrumentor
  ):
    """Test default MCP HTTP factory instruments HTTPX client when available."""
    client = Mock()
    mock_base_factory.return_value = client

    result = create_mcp_http_client()

    assert result is client
    mock_instrumentor.instrument_client.assert_called_once_with(client)

  @patch(
      "google.adk.tools.mcp_tool.mcp_session_manager._create_mcp_http_client"
  )
  @patch(
      "google.adk.tools.mcp_tool.mcp_session_manager._HAS_HTTPX_INSTRUMENTOR",
      False,
  )
  def test_default_httpx_factory_handles_missing_opentelemetry(
      self, mock_base_factory
  ):
    """Test default MCP HTTP factory works without OTel instrumentation."""
    client = Mock()
    mock_base_factory.return_value = client

    result = create_mcp_http_client()

    assert result is client

  def test_generate_session_key_stdio(self):
    """Test session key generation for stdio connections."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # For stdio, headers should be ignored and return constant key
    key1 = manager._generate_session_key({"Authorization": "Bearer token"})
    key2 = manager._generate_session_key(None)

    assert key1 == "stdio_session"
    assert key2 == "stdio_session"
    assert key1 == key2

  def test_generate_session_key_sse(self):
    """Test session key generation for SSE connections."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    headers1 = {"Authorization": "Bearer token1"}
    headers2 = {"Authorization": "Bearer token2"}

    key1 = manager._generate_session_key(headers1)
    key2 = manager._generate_session_key(headers2)
    key3 = manager._generate_session_key(headers1)

    # Different headers should generate different keys
    assert key1 != key2
    # Same headers should generate same key
    assert key1 == key3

    # Should be deterministic hash
    headers_json = json.dumps(headers1, sort_keys=True)
    expected_hash = hashlib.md5(headers_json.encode()).hexdigest()
    assert key1 == f"session_{expected_hash}"

  def test_merge_headers_stdio(self):
    """Test header merging for stdio connections."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Stdio connections don't support headers
    headers = manager._merge_headers({"Authorization": "Bearer token"})
    assert headers is None

  def test_merge_headers_sse(self):
    """Test header merging for SSE connections."""
    base_headers = {"Content-Type": "application/json"}
    sse_params = SseConnectionParams(
        url="https://example.com/mcp", headers=base_headers
    )
    manager = MCPSessionManager(sse_params)

    # With additional headers
    additional = {"Authorization": "Bearer token"}
    merged = manager._merge_headers(additional)

    expected = {
        "Content-Type": "application/json",
        "Authorization": "Bearer token",
    }
    assert merged == expected

  def test_is_session_disconnected(self):
    """Test session disconnection detection."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Create mock session
    session = MockClientSession()

    # Not disconnected
    assert not manager._is_session_disconnected(session)

    # Disconnected - read stream closed
    session._read_stream._closed = True
    assert manager._is_session_disconnected(session)

  @pytest.mark.asyncio
  async def test_create_session_stdio_new(self):
    """Test creating a new stdio session."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    mock_exit_stack = MockAsyncExitStack()

    with patch(
        "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client"
    ) as mock_stdio:
      with patch(
          "google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack"
      ) as mock_exit_stack_class:
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.SessionContext"
        ) as mock_session_context_class:

          # Setup mocks
          mock_exit_stack_class.return_value = mock_exit_stack
          mock_stdio.return_value = AsyncMock()

          # Mock SessionContext using MockSessionContext
          # Create a mock session that will be returned by SessionContext
          mock_session = AsyncMock()
          mock_session_context = MockSessionContext(session=mock_session)
          mock_session_context_class.return_value = mock_session_context
          mock_exit_stack.enter_async_context.return_value = mock_session

          # Create session
          session = await manager.create_session()

          # Verify session creation
          assert session == mock_session
          assert len(manager._sessions) == 1
          assert "stdio_session" in manager._sessions
          session_data = manager._sessions["stdio_session"]
          assert len(session_data) == 3
          assert session_data[0] == mock_session
          assert session_data[2] == asyncio.get_running_loop()

          # Verify SessionContext was created
          mock_session_context_class.assert_called_once()
          # Verify enter_async_context was called (which internally calls __aenter__)
          mock_exit_stack.enter_async_context.assert_called_once()

  @pytest.mark.asyncio
  async def test_create_session_reuse_existing(self):
    """Test reusing an existing connected session."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Create mock existing session
    existing_session = MockClientSession()
    existing_exit_stack = MockAsyncExitStack()
    manager._sessions["stdio_session"] = (
        existing_session,
        existing_exit_stack,
        asyncio.get_running_loop(),
    )

    # Session is connected
    existing_session._read_stream._closed = False
    existing_session._write_stream._closed = False

    session = await manager.create_session()

    # Should reuse existing session
    assert session == existing_session
    assert len(manager._sessions) == 1

    # Should not create new session
    existing_session.initialize.assert_not_called()

  @pytest.mark.asyncio
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client")
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack")
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.SessionContext")
  async def test_create_session_timeout(
      self, mock_session_context_class, mock_exit_stack_class, mock_stdio
  ):
    """Test session creation timeout."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    mock_exit_stack = MockAsyncExitStack()

    mock_exit_stack_class.return_value = mock_exit_stack
    mock_stdio.return_value = AsyncMock()

    # Mock SessionContext
    mock_session_context = AsyncMock()
    mock_session_context.__aenter__ = AsyncMock(
        return_value=MockClientSession()
    )
    mock_session_context.__aexit__ = AsyncMock(return_value=False)
    mock_session_context_class.return_value = mock_session_context

    # Mock enter_async_context to raise TimeoutError (simulating asyncio.wait_for timeout)
    mock_exit_stack.enter_async_context = AsyncMock(
        side_effect=asyncio.TimeoutError("Test timeout")
    )

    # Expect ConnectionError due to timeout
    with pytest.raises(ConnectionError, match="Failed to create MCP session"):
      await manager.create_session()

    # Verify SessionContext was created
    mock_session_context_class.assert_called_once()
    # Verify session was not added to pool
    assert not manager._sessions
    # Verify cleanup was called
    mock_exit_stack.aclose.assert_called_once()

  @pytest.mark.asyncio
  async def test_close_success(self):
    """Test successful cleanup of all sessions."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Add mock sessions
    session1 = MockClientSession()
    exit_stack1 = MockAsyncExitStack()
    session2 = MockClientSession()
    exit_stack2 = MockAsyncExitStack()

    manager._sessions["session1"] = (
        session1,
        exit_stack1,
        asyncio.get_running_loop(),
    )
    manager._sessions["session2"] = (
        session2,
        exit_stack2,
        asyncio.get_running_loop(),
    )

    await manager.close()

    # All sessions should be closed
    exit_stack1.aclose.assert_called_once()
    exit_stack2.aclose.assert_called_once()
    assert len(manager._sessions) == 0

  @pytest.mark.asyncio
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.logger")
  async def test_close_with_errors(self, mock_logger):
    """Test cleanup when some sessions fail to close."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Add mock sessions
    session1 = MockClientSession()
    exit_stack1 = MockAsyncExitStack()
    exit_stack1.aclose.side_effect = Exception("Close error 1")

    session2 = MockClientSession()
    exit_stack2 = MockAsyncExitStack()

    manager._sessions["session1"] = (
        session1,
        exit_stack1,
        asyncio.get_running_loop(),
    )
    manager._sessions["session2"] = (
        session2,
        exit_stack2,
        asyncio.get_running_loop(),
    )

    # Should not raise exception
    await manager.close()

    # Good session should still be closed
    exit_stack2.aclose.assert_called_once()
    assert len(manager._sessions) == 0

    # Error should be logged via logger.warning
    mock_logger.warning.assert_called_once()
    args, kwargs = mock_logger.warning.call_args
    assert "Error during session cleanup for session1: Close error 1" in args[0]
    assert kwargs.get("exc_info")

  @pytest.mark.asyncio
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client")
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack")
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.SessionContext")
  async def test_create_and_close_session_in_different_tasks(
      self, mock_session_context_class, mock_exit_stack_class, mock_stdio
  ):
    """Test creating and closing a session in different tasks."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    mock_exit_stack_class.return_value = MockAsyncExitStack()
    mock_stdio.return_value = AsyncMock()

    # Mock SessionContext
    mock_session_context = AsyncMock()
    mock_session_context.__aenter__ = AsyncMock(
        return_value=MockClientSession()
    )
    mock_session_context.__aexit__ = AsyncMock(return_value=False)
    mock_session_context_class.return_value = mock_session_context

    # Create session in a new task
    await asyncio.create_task(manager.create_session())

    # Close session in another task
    await asyncio.create_task(manager.close())

    # Verify session was closed
    assert not manager._sessions

  @pytest.mark.asyncio
  async def test_session_lock_different_loops(self):
    """Verify that _session_lock returns different locks for different loops."""

    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Access in current loop
    lock1 = manager._session_lock
    assert isinstance(lock1, asyncio.Lock)

    # Access in a different loop (in a separate thread)
    lock_container = []

    def run_in_thread():
      loop2 = asyncio.new_event_loop()
      asyncio.set_event_loop(loop2)
      try:

        async def get_lock():
          return manager._session_lock

        lock_container.append(loop2.run_until_complete(get_lock()))
      finally:
        loop2.close()

    thread = platform_thread.create_thread(target=run_in_thread)
    thread.start()
    thread.join()

    assert lock_container
    lock2 = lock_container[0]
    assert isinstance(lock2, asyncio.Lock)
    assert lock1 is not lock2

  @pytest.mark.asyncio
  async def test_cleanup_session_cross_loop(self):
    """Verify that _cleanup_session uses run_coroutine_threadsafe for different loops."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)
    mock_exit_stack = MockAsyncExitStack()

    # Create a dummy loop that is "running" in another thread
    loop2 = asyncio.new_event_loop()
    try:
      with patch(
          "google.adk.tools.mcp_tool.mcp_session_manager.asyncio.run_coroutine_threadsafe"
      ) as mock_run_threadsafe:
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.logger"
        ) as mock_logger:
          # We need to mock the return value of run_coroutine_threadsafe to be a future
          mock_future = Mock()
          mock_run_threadsafe.return_value = mock_future

          await manager._cleanup_session("test_session", mock_exit_stack, loop2)

          # Verify run_coroutine_threadsafe was called
          # ANY is used because a new coroutine object is created each time
          mock_run_threadsafe.assert_called_once_with(ANY, loop2)

          mock_logger.info.assert_any_call(
              "Scheduling cleanup of session test_session on its original"
              " event loop."
          )
          mock_future.add_done_callback.assert_called_once()
    finally:
      loop2.close()

  @pytest.mark.asyncio
  async def test_create_session_cleans_up_without_aclose_if_loop_is_different(
      self,
  ):
    """Verify that sessions from different loops are cleaned up without calling aclose()."""
    from google.adk.features import FeatureName
    from google.adk.features._feature_registry import temporary_feature_override

    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # 1. Simulate a session created in a "different" loop
    mock_session = MockClientSession()
    mock_exit_stack = MockAsyncExitStack()
    # Use a dummy object as a different loop
    different_loop = Mock(spec=asyncio.AbstractEventLoop)

    manager._sessions["stdio_session"] = (
        mock_session,
        mock_exit_stack,
        different_loop,
    )

    # 2. Mock creation of a new session
    # We need to mock create_client, wait_for, and SessionContext
    with patch.object(manager, "_create_client") as mock_create_client:
      with patch(
          "google.adk.tools.mcp_tool.mcp_session_manager.asyncio.wait_for"
      ) as mock_wait_for:
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.SessionContext"
        ) as mock_session_context_class:
          # Setup mocks for new session creation
          mock_create_client.return_value = AsyncMock()
          new_session = MockClientSession()
          mock_wait_for.return_value = new_session
          mock_session_context_class.return_value = AsyncMock()

          # 3. Call create_session with flag off to hit wait_for branch
          with temporary_feature_override(
              FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False
          ):
            session = await manager.create_session()

          # 4. Verify results
          assert session == new_session
          assert len(manager._sessions) == 1
          # Verify that old exit_stack.aclose was NOT called since loop was different
          mock_exit_stack.aclose.assert_not_called()

  @pytest.mark.asyncio
  async def test_close_skips_aclose_for_different_loop_sessions(self):
    """Verify that close() skips aclose() for sessions from different loops."""
    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Add one session from same loop and one from different loop
    current_loop = asyncio.get_running_loop()
    different_loop = Mock(spec=asyncio.AbstractEventLoop)

    session1 = MockClientSession()
    exit_stack1 = MockAsyncExitStack()
    manager._sessions["session1"] = (session1, exit_stack1, current_loop)

    session2 = MockClientSession()
    exit_stack2 = MockAsyncExitStack()
    manager._sessions["session2"] = (session2, exit_stack2, different_loop)

    await manager.close()

    # exit_stack1 should be closed, exit_stack2 should be skipped
    exit_stack1.aclose.assert_called_once()
    exit_stack2.aclose.assert_not_called()
    assert len(manager._sessions) == 0

  @pytest.mark.asyncio
  async def test_pickle_mcp_session_manager(self):
    """Verify that MCPSessionManager can be pickled and unpickled."""
    import pickle

    manager = MCPSessionManager(self.mock_stdio_connection_params)

    # Access the lock to ensure it's initialized
    lock = manager._session_lock
    assert isinstance(lock, asyncio.Lock)

    # Add a mock session to verify it's cleared on pickling
    manager._sessions["test"] = (Mock(), Mock(), asyncio.get_running_loop())

    # Pickle and unpickle
    pickled = pickle.dumps(manager)
    unpickled = pickle.loads(pickled)

    # Verify basics are restored
    assert unpickled._connection_params == manager._connection_params

    # Verify transient/unpicklable members are re-initialized or cleared
    assert unpickled._sessions == {}
    assert unpickled._session_lock_map == {}
    assert isinstance(unpickled._lock_map_lock, type(manager._lock_map_lock))
    assert unpickled._lock_map_lock is not manager._lock_map_lock
    assert unpickled._errlog == sys.stderr

    # Verify we can still get a lock in the new instance
    new_lock = unpickled._session_lock
    assert isinstance(new_lock, asyncio.Lock)
    assert new_lock is not lock

  @pytest.mark.asyncio
  async def test_get_mtls_transport_flag_off(self):
    """Test that _get_mtls_transport returns None when flag is off."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)
    with patch.dict(
        "os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"}
    ):
      transport = await manager._get_mtls_transport()
      assert transport is None

  @pytest.mark.asyncio
  @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported")
  async def test_get_mtls_transport_success(self):
    """Test successful _GoogleAuthAsyncTransport creation with mTLS."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    mock_creds = Mock()
    mock_session = AsyncMock()
    mock_session.is_mtls = True
    mock_session.configure_mtls_channel = AsyncMock()

    with patch.dict(
        "os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}
    ):
      with patch("google.auth.default", return_value=(mock_creds, None)):
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.AsyncAuthorizedSession",
            return_value=mock_session,
        ):
          with patch(
              "google.adk.tools.mcp_tool.mcp_session_manager._GoogleAuthAsyncTransport"
          ) as mock_transport_class:
            mock_transport = Mock()
            mock_transport_class.return_value = mock_transport

            transport = await manager._get_mtls_transport()

            assert transport == mock_transport
            mock_session.configure_mtls_channel.assert_called_once()
            mock_transport_class.assert_called_once_with(mock_session)

            # Test caching
            transport2 = await manager._get_mtls_transport()
            assert transport2 == transport
            mock_session.configure_mtls_channel.assert_called_once()

  @pytest.mark.asyncio
  @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported")
  async def test_get_mtls_transport_failure_not_mtls(self):
    """Test that _get_mtls_transport returns None when channel is not mTLS."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    mock_creds = Mock()
    mock_session = AsyncMock()
    mock_session.is_mtls = False
    mock_session.configure_mtls_channel = AsyncMock()

    with patch.dict(
        "os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}
    ):
      with patch("google.auth.default", return_value=(mock_creds, None)):
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.AsyncAuthorizedSession",
            return_value=mock_session,
        ):
          transport = await manager._get_mtls_transport()
          assert transport is None

  @pytest.mark.asyncio
  @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported")
  async def test_get_mtls_transport_failure_exception(self):
    """Test that _get_mtls_transport returns None when exception occurs."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    with patch.dict(
        "os.environ", {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}
    ):
      with patch("google.auth.default", side_effect=Exception("auth error")):
        transport = await manager._get_mtls_transport()
        assert transport is None

  @patch("google.adk.tools.mcp_tool.mcp_session_manager.sse_client")
  def test_create_client_with_mtls_transport_sse(self, mock_sse_client):
    """Test that _create_client uses mtls_transport to create factory for SSE."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    mock_transport = Mock(spec=httpx.AsyncBaseTransport)

    manager._create_client(mtls_transport=mock_transport)

    mock_sse_client.assert_called_once()
    called_kwargs = mock_sse_client.call_args[1]
    factory = called_kwargs["httpx_client_factory"]

    # Verify the factory creates client with transport
    client = factory(headers={"a": "b"}, timeout=httpx.Timeout(10.0))
    assert isinstance(client, httpx.AsyncClient)
    assert isinstance(client._transport, _SharedAsyncTransport)
    assert client._transport._transport == mock_transport
    assert client.headers.get("a") == "b"
    assert client.timeout.read == 10.0

  @pytest.mark.asyncio
  async def test_google_auth_async_transport_handle_request(self):
    """Test that _GoogleAuthAsyncTransport correctly forwards request and returns response."""
    from google.adk.tools.mcp_tool.mcp_session_manager import _GoogleAuthAsyncTransport

    mock_session = AsyncMock()
    mock_auth_response = AsyncMock()
    mock_auth_response.status_code = 200
    mock_auth_response.headers = {"content-type": "application/json"}
    mock_auth_response.content = AsyncMock()

    mock_session.request.return_value = mock_auth_response

    transport = _GoogleAuthAsyncTransport(mock_session)

    request = httpx.Request(
        "GET", "https://example.com/api", headers={"x-test": "value"}
    )

    response = await transport.handle_async_request(request)

    assert response.status_code == 200
    assert response.headers["content-type"] == "application/json"

    mock_session.request.assert_called_once_with(
        method="GET",
        url="https://example.com/api",
        data=None,
        headers={"x-test": "value", "host": "example.com"},
        timeout=30.0,
    )


@pytest.mark.asyncio
async def test_retry_on_errors_decorator():
  """Test the retry_on_errors decorator."""

  call_count = 0

  @retry_on_errors
  async def mock_function(self):
    nonlocal call_count
    call_count += 1
    if call_count == 1:
      raise ConnectionError("Resource closed")
    return "success"

  mock_self = Mock()
  result = await mock_function(mock_self)

  assert result == "success"
  assert call_count == 2  # First call fails, second succeeds


@pytest.mark.asyncio
async def test_retry_on_errors_decorator_does_not_retry_cancelled_error():
  """Test the retry_on_errors decorator does not retry cancellation."""

  call_count = 0

  @retry_on_errors
  async def mock_function(self):
    nonlocal call_count
    call_count += 1
    raise asyncio.CancelledError()

  mock_self = Mock()
  with pytest.raises(asyncio.CancelledError):
    await mock_function(mock_self)

  assert call_count == 1


@pytest.mark.asyncio
async def test_retry_on_errors_decorator_does_not_retry_when_task_is_cancelling():
  """Test the retry_on_errors decorator does not retry when cancelling."""

  call_count = 0

  @retry_on_errors
  async def mock_function(self):
    nonlocal call_count
    call_count += 1
    raise ConnectionError("Resource closed")

  class _MockTask:

    def cancelling(self):
      return 1

  mock_self = Mock()
  with patch.object(asyncio, "current_task", return_value=_MockTask()):
    with pytest.raises(ConnectionError):
      await mock_function(mock_self)

  assert call_count == 1


@pytest.mark.asyncio
async def test_retry_on_errors_decorator_does_not_retry_exception_from_cancel():
  """Test the retry_on_errors decorator does not retry exceptions on cancel."""

  call_count = 0

  @retry_on_errors
  async def mock_function(self):
    nonlocal call_count
    call_count += 1
    try:
      raise asyncio.CancelledError()
    except asyncio.CancelledError:
      raise ConnectionError("Resource closed")

  mock_self = Mock()
  with pytest.raises(ConnectionError):
    await mock_function(mock_self)

  assert call_count == 1


class TestMCPSessionManagerGetSessionContext:
  """Tests for MCPSessionManager._get_session_context.

  This is the lookup that allows McpTool to obtain the SessionContext
  for the current session and call `_run_guarded` on it.
  """

  def setup_method(self):
    """Set up a manager with stdio params."""
    self.params = StdioServerParameters(command="echo", args=[])
    self.manager = MCPSessionManager(self.params)

  def test_returns_none_when_no_session_exists(self):
    """With an empty pool, _get_session_context returns None."""
    assert self.manager._get_session_context() is None
    assert self.manager._get_session_context(headers={"x": "y"}) is None

  def test_returns_stored_session_context_for_stdio(self):
    """A stored SessionContext is returned for the stdio session key."""
    fake_ctx = MockSessionContext()
    # Stdio uses a constant session key, so headers are ignored.
    self.manager._session_contexts["stdio_session"] = fake_ctx

    assert self.manager._get_session_context() is fake_ctx
    assert self.manager._get_session_context(headers={"x": "y"}) is fake_ctx

  def test_returns_correct_session_context_per_header_set(self):
    """Different header sets produce different keys, so different contexts."""
    sse_params = SseConnectionParams(url="https://example.com/mcp")
    manager = MCPSessionManager(sse_params)

    ctx_a = MockSessionContext()
    ctx_b = MockSessionContext()
    key_a = manager._generate_session_key(
        manager._merge_headers({"x-token": "a"})
    )
    key_b = manager._generate_session_key(
        manager._merge_headers({"x-token": "b"})
    )
    manager._session_contexts[key_a] = ctx_a
    manager._session_contexts[key_b] = ctx_b

    assert manager._get_session_context(headers={"x-token": "a"}) is ctx_a
    assert manager._get_session_context(headers={"x-token": "b"}) is ctx_b
    # Unknown header set returns None.
    assert manager._get_session_context(headers={"x-token": "c"}) is None

  def test_session_contexts_dict_is_independent_of_sessions_tuple(self):
    """Backward-compat guard: _sessions remains the tuple shape.

    Downstream tests poke at `_sessions` directly using tuple unpacking
    (`session, exit_stack, loop = manager._sessions[key]`). This test
    ensures we did not switch to a dataclass that would break that.
    """
    mock_session = MockClientSession()
    mock_exit_stack = MockAsyncExitStack()
    mock_loop = Mock()
    mock_ctx = MockSessionContext()

    self.manager._sessions["stdio_session"] = (
        mock_session,
        mock_exit_stack,
        mock_loop,
    )
    self.manager._session_contexts["stdio_session"] = mock_ctx

    # Should be unpackable as a 3-tuple.
    session, exit_stack, loop = self.manager._sessions["stdio_session"]
    assert session is mock_session
    assert exit_stack is mock_exit_stack
    assert loop is mock_loop

    # And the SessionContext is reachable independently.
    assert self.manager._get_session_context() is mock_ctx

  def test_pickling_round_trip_clears_runtime_state(self):
    """__getstate__/__setstate__ should drop runtime SessionContext refs."""
    import pickle

    self.manager._session_contexts["stdio_session"] = MockSessionContext()

    restored = pickle.loads(pickle.dumps(self.manager))

    # Runtime state must not survive pickling.
    assert restored._session_contexts == {}
    assert restored._sessions == {}


class TestMCPSessionManagerCreateSessionFlagOff:
  """Pin down that create_session does NOT consult task aliveness when off.

  Existing callers broke under an earlier unconditional version of this
  fix because the new `_is_task_alive` check caused session re-creation
  paths to fire when the test mocks did not have a live `_task`. The
  check must be gated behind the feature flag.
  """

  def setup_method(self):
    from google.adk.features import FeatureName  # noqa: F401

    self.params = StdioServerParameters(command="echo", args=[])
    self.manager = MCPSessionManager(self.params)

  @pytest.mark.asyncio
  async def test_existing_session_reused_when_flag_off_even_with_dead_ctx(
      self,
  ):
    """A 'dead' SessionContext does not invalidate the session when off."""
    from google.adk.features import FeatureName
    from google.adk.features._feature_registry import temporary_feature_override

    # Pre-populate a healthy-looking session and a SessionContext whose
    # _task looks dead.
    healthy_session = MockClientSession()
    dead_ctx = MockSessionContext()
    dead_ctx._is_task_alive = False  # pretend the task died
    self.manager._sessions["stdio_session"] = (
        healthy_session,
        MockAsyncExitStack(),
        asyncio.get_running_loop(),
    )
    self.manager._session_contexts["stdio_session"] = dead_ctx

    # With flag OFF, the create_session must reuse the existing session
    # rather than tearing it down because of the dead _task.
    with temporary_feature_override(
        FeatureName._MCP_GRACEFUL_ERROR_HANDLING, False
    ):
      returned = await self.manager.create_session()

    assert returned is healthy_session

  @pytest.mark.asyncio
  async def test_existing_session_recreated_when_flag_on_with_dead_ctx(
      self,
  ):
    """And confirm: with flag ON, the dead _task DOES trigger re-creation."""
    from google.adk.features import FeatureName
    from google.adk.features._feature_registry import temporary_feature_override

    healthy_session = MockClientSession()
    dead_ctx = MockSessionContext()
    dead_ctx._is_task_alive = False
    # Mark the existing exit_stack so we can confirm the new one is different.
    old_exit_stack = MockAsyncExitStack()
    self.manager._sessions["stdio_session"] = (
        healthy_session,
        old_exit_stack,
        asyncio.get_running_loop(),
    )
    self.manager._session_contexts["stdio_session"] = dead_ctx

    # Patch the SessionContext used inside create_session so we don't
    # actually try to launch a real subprocess. Mirrors the patching
    # pattern used by `test_create_session_stdio_new`.
    new_session = MockClientSession()
    mock_exit_stack = MockAsyncExitStack()
    mock_session_ctx = MockSessionContext(session=new_session)

    with temporary_feature_override(
        FeatureName._MCP_GRACEFUL_ERROR_HANDLING, True
    ):
      with patch("google.adk.tools.mcp_tool.mcp_session_manager.stdio_client"):
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack"
        ) as mock_exit_stack_class:
          with patch(
              "google.adk.tools.mcp_tool.mcp_session_manager.SessionContext"
          ) as mock_session_context_class:
            mock_exit_stack_class.return_value = mock_exit_stack
            mock_session_context_class.return_value = mock_session_ctx
            mock_exit_stack.enter_async_context.return_value = new_session

            returned = await self.manager.create_session()

    assert returned is new_session
    # The original 'healthy_session' was torn down because dead_ctx
    # told us the task was gone.
    assert returned is not healthy_session


class TestMCPGracefulErrorHandlingFlagContract:
  """Pin down the public contract that GE will rely on to enable the fix.

  GE will flip this fix on by setting an environment variable in their
  deployment config (per Sasha's confirmation: "environment variable, GE
  team is responsible for setting it"). The deployment expects:

    * `ADK_ENABLE_MCP_GRACEFUL_ERROR_HANDLING=1`  enables the fix
    * absence of the variable                      keeps it disabled
    * `ADK_DISABLE_MCP_GRACEFUL_ERROR_HANDLING=1` is the kill switch

  These tests are guards: if anyone refactors the feature-flag framework
  in a way that changes how the env var is read (renames it, caches the
  value at import time, requires a binary push, etc.), these tests fail
  loudly so we don't silently break GE's rollout.
  """

  def test_default_state_is_on(self):
    """The fix must be enabled by default."""
    import os

    from google.adk.features import FeatureName
    from google.adk.features import is_feature_enabled

    enable = "ADK_ENABLE_MCP_GRACEFUL_ERROR_HANDLING"
    disable = "ADK_DISABLE_MCP_GRACEFUL_ERROR_HANDLING"
    saved = {k: os.environ.pop(k) for k in (enable, disable) if k in os.environ}
    try:
      assert (
          is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING) is True
      )
    finally:
      os.environ.update(saved)

  def test_env_var_disable_flips_flag_off_at_runtime(self):
    """The env var must turn the fix off without a rebuild."""
    import os

    from google.adk.features import FeatureName
    from google.adk.features import is_feature_enabled

    disable = "ADK_DISABLE_MCP_GRACEFUL_ERROR_HANDLING"
    saved = os.environ.pop(disable, None)
    try:
      os.environ[disable] = "1"
      assert (
          is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING) is False
      )
      # And once it's removed, we revert. Confirms the value is read
      # live from os.environ on every call (no caching, no binary push).
      del os.environ[disable]
      assert (
          is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING) is True
      )
    finally:
      if saved is not None:
        os.environ[disable] = saved

  def test_env_var_disable_acts_as_kill_switch(self):
    """The disable env var lets consumers turn off without a rebuild."""
    import os

    from google.adk.features import FeatureName
    from google.adk.features import is_feature_enabled
    from google.adk.features._feature_registry import temporary_feature_override

    disable = "ADK_DISABLE_MCP_GRACEFUL_ERROR_HANDLING"
    enable = "ADK_ENABLE_MCP_GRACEFUL_ERROR_HANDLING"
    saved_disable = os.environ.pop(disable, None)
    saved_enable = os.environ.pop(enable, None)
    try:
      # If a future default flip ever turns this on by default, the
      # disable env var should still let consumers turn it back off
      # without a rebuild.
      os.environ[disable] = "1"
      assert (
          is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING) is False
      )
      # And confirm: a programmatic override takes precedence over the
      # disable env var (priority order documented in _feature_registry).
      with temporary_feature_override(
          FeatureName._MCP_GRACEFUL_ERROR_HANDLING, True
      ):
        assert (
            is_feature_enabled(FeatureName._MCP_GRACEFUL_ERROR_HANDLING) is True
        )
    finally:
      if saved_disable is not None:
        os.environ[disable] = saved_disable
      if saved_enable is not None:
        os.environ[enable] = saved_enable

  @pytest.mark.asyncio
  @patch("google.adk.tools.mcp_tool.mcp_session_manager.asyncio.wait_for")
  async def test_create_session_does_not_use_wait_for_when_ge_is_enabled(
      self, mock_wait_for
  ):
    """create_session must not wrap enter_async_context in asyncio.wait_for when GE is enabled."""
    from google.adk.features import FeatureName
    from google.adk.features._feature_registry import temporary_feature_override

    manager = MCPSessionManager(
        StdioConnectionParams(
            server_params=StdioServerParameters(command="dummy", args=[]),
            timeout=5.0,
        )
    )
    with temporary_feature_override(
        FeatureName._MCP_GRACEFUL_ERROR_HANDLING, True
    ):
      with patch(
          "google.adk.tools.mcp_tool.mcp_session_manager.AsyncExitStack"
      ) as mock_stack:
        mock_stack.return_value.enter_async_context = AsyncMock()
        with patch(
            "google.adk.tools.mcp_tool.mcp_session_manager.SessionContext"
        ):
          with patch(
              "google.adk.tools.mcp_tool.mcp_session_manager.stdio_client"
          ):
            await manager.create_session()

    mock_wait_for.assert_not_called()


class TestRefreshableAsyncCredentials:

  @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported")
  @pytest.mark.asyncio
  async def test_before_request_refreshes_and_injects_token(self):
    mock_creds = Mock()
    mock_creds.expired = True
    mock_creds.token = "new_token"

    # Mock creds.refresh to simulate refresh
    def mock_refresh(req):
      mock_creds.token = "refreshed_token"
      mock_creds.expired = False

    mock_creds.refresh = mock_refresh

    credentials = _RefreshableAsyncCredentials(mock_creds)
    headers = {}

    await credentials.before_request(None, "GET", "http://example.com", headers)

    assert headers["Authorization"] == "Bearer refreshed_token"

  @pytest.mark.skipif(not AIO_SUPPORTED, reason="google.auth.aio not supported")
  @pytest.mark.parametrize(
      "existing_header_key",
      ["Authorization", "authorization", "AUTHORIZATION", "authORIZATION"],
  )
  @pytest.mark.asyncio
  async def test_before_request_skips_refresh_if_authorization_header_exists_case_insensitive(
      self, existing_header_key
  ):
    mock_creds = Mock()
    mock_creds.expired = True
    mock_creds.token = "new_token"
    mock_creds.refresh = Mock()

    credentials = _RefreshableAsyncCredentials(mock_creds)
    headers = {existing_header_key: "Bearer existing_token"}

    await credentials.before_request(None, "GET", "http://example.com", headers)

    mock_creds.refresh.assert_not_called()
    assert headers == {existing_header_key: "Bearer existing_token"}


class TestGoogleAuthAsyncByteStream:

  @pytest.mark.asyncio
  async def test_iteration_yields_chunks(self):
    mock_auth_response = AsyncMock()

    async def mock_content():
      yield b"chunk1"
      yield b"chunk2"

    mock_auth_response.content = mock_content

    stream = _GoogleAuthAsyncByteStream(mock_auth_response)
    chunks = []
    async for chunk in stream:
      chunks.append(chunk)

    assert chunks == [b"chunk1", b"chunk2"]

  @pytest.mark.asyncio
  async def test_aclose_closes_response(self):
    mock_auth_response = AsyncMock()
    stream = _GoogleAuthAsyncByteStream(mock_auth_response)
    await stream.aclose()
    mock_auth_response.close.assert_called_once()


class TestDebugHttpxClientFactory:
  """Tests for _DebugHttpxClientFactory."""

  @pytest.mark.asyncio
  async def test_debug_factory_registers_hook(self):
    """Test that the debug factory registers the response hook on client creation."""
    base_client = httpx.AsyncClient()
    base_factory = Mock(return_value=base_client)
    debug_factory = _DebugHttpxClientFactory(base_factory)

    client = debug_factory()
    assert debug_factory._response_hook in client.event_hooks["response"]
    # Clean up client
    await base_client.aclose()

  @pytest.mark.asyncio
  async def test_response_hook_records_when_var_set(self):
    """Test that the response hook records HTTP info when _http_debug_var is set."""
    base_client = httpx.AsyncClient()
    base_factory = Mock(return_value=base_client)
    debug_factory = _DebugHttpxClientFactory(base_factory)

    # Mock httpx.Response
    mock_request = Mock(spec=httpx.Request)
    mock_request.method = "GET"
    mock_request.content = b"request body"
    mock_request.headers = httpx.Headers({"X-Req": "val"})

    mock_response = Mock(spec=httpx.Response)
    mock_response.url = httpx.URL("https://example.com/test")
    mock_response.status_code = 200
    mock_response.request = mock_request
    mock_response.headers = httpx.Headers({
        "content-type": "application/json",
        "X-Resp": "val",
    })
    mock_response.text = "response body"
    mock_response.aread = AsyncMock()

    debug_list = []
    token = _http_debug_var.set(debug_list)
    try:
      await debug_factory._response_hook(mock_response)
    finally:
      _http_debug_var.reset(token)

    assert len(debug_list) == 1
    record = debug_list[0]
    assert record["url"] == "https://example.com/test"
    assert record["status_code"] == 200
    assert record["method"] == "GET"
    assert record["request_body"] == "request body"
    assert record["response_body"] == "response body"
    assert record["request_headers"]["x-req"] == "val"
    assert record["response_headers"]["x-resp"] == "val"
    mock_response.aread.assert_called_once()
    await base_client.aclose()

  @pytest.mark.asyncio
  async def test_response_hook_does_not_record_when_var_not_set(self):
    """Test that the response hook does not record when _http_debug_var is not set."""
    base_client = httpx.AsyncClient()
    base_factory = Mock(return_value=base_client)
    debug_factory = _DebugHttpxClientFactory(base_factory)

    mock_response = Mock(spec=httpx.Response)
    mock_response.aread = AsyncMock()

    # _http_debug_var is not set (default None)
    await debug_factory._response_hook(mock_response)
    mock_response.aread.assert_not_called()
    await base_client.aclose()

  @pytest.mark.asyncio
  async def test_response_hook_skips_sse_body(self):
    """Test that the response hook avoids reading the body for SSE streams."""
    base_client = httpx.AsyncClient()
    base_factory = Mock(return_value=base_client)
    debug_factory = _DebugHttpxClientFactory(base_factory)

    mock_request = Mock(spec=httpx.Request)
    mock_request.method = "GET"
    mock_request.content = None
    mock_request.headers = httpx.Headers()

    mock_response = Mock(spec=httpx.Response)
    mock_response.url = httpx.URL("https://example.com/sse")
    mock_response.status_code = 200
    mock_response.request = mock_request
    mock_response.headers = httpx.Headers({"content-type": "text/event-stream"})
    mock_response.aread = AsyncMock()

    debug_list = []
    token = _http_debug_var.set(debug_list)
    try:
      await debug_factory._response_hook(mock_response)
    finally:
      _http_debug_var.reset(token)

    assert len(debug_list) == 1
    record = debug_list[0]
    assert record["response_body"] == "<SSE stream>"
    mock_response.aread.assert_not_called()
    await base_client.aclose()

  @pytest.mark.asyncio
  async def test_debug_factory_passes_keyword_arguments(self):
    """Test that the debug factory passes keyword arguments to base_factory."""
    base_client = httpx.AsyncClient()

    # A factory function that only accepts keyword arguments
    def keyword_only_factory(**kwargs) -> httpx.AsyncClient:
      assert "headers" in kwargs
      assert "timeout" in kwargs
      assert "auth" in kwargs
      return base_client

    debug_factory = _DebugHttpxClientFactory(keyword_only_factory)

    # Should work when called with positional arguments (which maps them to parameter names)
    client = debug_factory({"X-Test": "Val"}, None, None)
    assert client is base_client
    await base_client.aclose()

  @pytest.mark.asyncio
  async def test_response_hook_truncates_large_bodies(self):
    """Test that response hook truncates request and response bodies exceeding limit."""
    base_client = httpx.AsyncClient()
    base_factory = Mock(return_value=base_client)
    debug_factory = _DebugHttpxClientFactory(base_factory)

    # Mock request and response with large content
    large_req_body = b"a" * 1500
    large_resp_body = "b" * 1500

    mock_request = Mock(spec=httpx.Request)
    mock_request.method = "POST"
    mock_request.content = large_req_body
    mock_request.headers = httpx.Headers()

    mock_response = Mock(spec=httpx.Response)
    mock_response.url = httpx.URL("https://example.com/large")
    mock_response.status_code = 200
    mock_response.request = mock_request
    mock_response.headers = httpx.Headers({"content-type": "application/json"})
    mock_response.text = large_resp_body
    mock_response.aread = AsyncMock()

    debug_list = []
    token = _http_debug_var.set(debug_list)
    try:
      await debug_factory._response_hook(mock_response)
    finally:
      _http_debug_var.reset(token)

    assert len(debug_list) == 1
    record = debug_list[0]
    assert len(record["request_body"]) == 1015  # 1000 + len("... [truncated]")
    assert record["request_body"].endswith("... [truncated]")
    assert record["request_body"].startswith("a" * 1000)

    assert len(record["response_body"]) == 1015  # 1000 + len("... [truncated]")
    assert record["response_body"].endswith("... [truncated]")
    assert record["response_body"].startswith("b" * 1000)

    await base_client.aclose()