composiohq--composio/python/tests/test_tool_router.py

"""Test ToolRouter functionality."""

import typing as t
from unittest.mock import MagicMock, patch

import httpx
import pytest
from composio_client import omit
from pydantic import BaseModel, Field

from composio.client import HttpClient
from composio.core.models.experimental import ExperimentalAPI
from composio.core.models.tool_router import (
    SESSION_PRESET_DIRECT_TOOLS,
    ToolkitConnectionsDetails,
    ToolkitConnectionState,
    ToolRouter,
    ToolRouterExperimentalConfig,
    ToolRouterManageConnectionsConfig,
    ToolRouterMCPServerConfig,
    ToolRouterMCPServerType,
    ToolRouterSessionExperimental,
    ToolRouterToolkitsDisableConfig,
    ToolRouterToolkitsEnableConfig,
)
from composio.core.models.tool_router_session import (
    DIRECT_CUSTOM_TOOL_DESCRIPTION_PREFIX,
    ToolRouterSession,
    ToolRouterSessionWithMcp,
)
from composio.exceptions import InvalidParams, ValidationError

experimental_api = ExperimentalAPI()


class GrepInput(BaseModel):
    pattern: str = Field(description="Pattern to search for")


@pytest.fixture
def mock_client():
    """Create a mock HTTP client."""
    client = MagicMock()
    client.api_key = "test-api-key"

    # Mock session responses
    mock_session_response = MagicMock()
    mock_session_response.session_id = "session_123"
    mock_session_response.mcp = MagicMock()
    mock_session_response.mcp.type = "http"
    mock_session_response.mcp.url = "https://mcp.example.com/session_123"
    mock_session_response.tool_router_tools = [
        "GMAIL_FETCH_EMAILS",
        "SLACK_SEND_MESSAGE",
        "GITHUB_CREATE_ISSUE",
    ]
    mock_session_response.config = MagicMock()
    mock_session_response.config.user_id = "user_123"
    mock_session_response.config.preload = MagicMock()
    mock_session_response.config.preload.tools = []
    mock_session_response.experimental = None  # Default to None

    client.tool_router.session.create.return_value = mock_session_response
    client.tool_router.session.retrieve.return_value = mock_session_response
    client.tool_router.session.attach.return_value = mock_session_response

    # Mock link response
    mock_link_response = MagicMock()
    mock_link_response.connected_account_id = "conn_456"
    mock_link_response.redirect_url = "https://composio.dev/auth/redirect"
    client.tool_router.session.link.return_value = mock_link_response

    # Mock toolkits response
    mock_toolkit_item_active = MagicMock()
    mock_toolkit_item_active.slug = "gmail"
    mock_toolkit_item_active.name = "Gmail"
    mock_toolkit_item_active.meta.logo = "https://example.com/gmail-logo.png"
    mock_toolkit_item_active.is_no_auth = False
    mock_toolkit_item_active.connected_account = MagicMock()
    mock_toolkit_item_active.connected_account.id = "conn_123"
    mock_toolkit_item_active.connected_account.status = "ACTIVE"
    mock_toolkit_item_active.connected_account.auth_config = MagicMock()
    mock_toolkit_item_active.connected_account.auth_config.id = "auth_config_123"
    mock_toolkit_item_active.connected_account.auth_config.auth_scheme = "OAUTH2"
    mock_toolkit_item_active.connected_account.auth_config.is_composio_managed = True

    mock_toolkit_item_inactive = MagicMock()
    mock_toolkit_item_inactive.slug = "github"
    mock_toolkit_item_inactive.name = "GitHub"
    mock_toolkit_item_inactive.meta.logo = "https://example.com/github-logo.png"
    mock_toolkit_item_inactive.is_no_auth = False
    mock_toolkit_item_inactive.connected_account = None

    mock_toolkits_response = MagicMock()
    mock_toolkits_response.items = [
        mock_toolkit_item_active,
        mock_toolkit_item_inactive,
    ]
    mock_toolkits_response.next_cursor = "cursor_789"
    mock_toolkits_response.total_pages = 2
    client.tool_router.session.toolkits.return_value = mock_toolkits_response

    return client


@pytest.fixture
def mock_provider():
    """Create a mock provider."""
    provider = MagicMock()
    return provider


@pytest.fixture
def tool_router(mock_client, mock_provider):
    """Create a ToolRouter instance with mocked client."""
    return ToolRouter(client=mock_client, provider=mock_provider)


class TestToolRouter:
    """Test cases for ToolRouter class."""

    def test_constructor(self, tool_router, mock_client, mock_provider):
        """Test that ToolRouter initializes correctly."""
        assert tool_router._client == mock_client
        assert tool_router._provider == mock_provider

    def test_create_basic_session(self, tool_router, mock_client):
        """Test creating a basic session with minimal configuration."""
        session = tool_router.create(user_id="user_123")

        # Verify session properties
        assert session.session_id == "session_123"
        assert session.mcp.type == ToolRouterMCPServerType.HTTP
        assert session.mcp.url == "https://mcp.example.com/session_123"
        assert session.mcp.headers == {"x-api-key": "test-api-key"}
        assert callable(session.tools)
        assert callable(session.authorize)
        assert callable(session.toolkits)

        # Verify API was called
        mock_client.tool_router.session.create.assert_called_once()

    def test_create_session_default_returns_base_session(self, tool_router):
        """Default create() (mcp omitted) returns the base ToolRouterSession.

        MCP is opt-in: the type must NOT be the mcp-surfacing subclass.
        """
        session = tool_router.create(user_id="user_123")

        assert type(session) is ToolRouterSession
        assert not isinstance(session, ToolRouterSessionWithMcp)

    def test_create_session_with_mcp_true_returns_with_mcp(self, tool_router):
        """create(..., mcp=True) returns ToolRouterSessionWithMcp with mcp surfaced."""
        session = tool_router.create(user_id="user_123", mcp=True)

        assert isinstance(session, ToolRouterSessionWithMcp)
        # The MCP endpoint is populated on the runtime object either way.
        assert session.mcp.type == ToolRouterMCPServerType.HTTP
        assert session.mcp.url == "https://mcp.example.com/session_123"

    def test_use_session_default_returns_base_session(self, tool_router):
        """Default use() (mcp omitted) returns the base ToolRouterSession."""
        session = tool_router.use("session_123")

        assert type(session) is ToolRouterSession
        assert not isinstance(session, ToolRouterSessionWithMcp)

    def test_use_session_with_mcp_true_returns_with_mcp(self, tool_router):
        """use(..., mcp=True) returns ToolRouterSessionWithMcp."""
        session = tool_router.use("session_123", mcp=True)

        assert isinstance(session, ToolRouterSessionWithMcp)
        assert session.mcp.url == "https://mcp.example.com/session_123"

    def test_delete_session_by_id(self, tool_router, mock_client):
        """delete() removes a session by ID."""
        mock_client.delete.return_value = {
            "session_id": "session_123",
            "deleted": True,
        }

        result = tool_router.delete("session_123")

        mock_client.delete.assert_called_once_with(
            "/api/v3.1/tool_router/session/session_123",
            cast_to=t.Dict[str, t.Any],
        )
        assert result == {
            "session_id": "session_123",
            "deleted": True,
        }

    def test_delete_session_by_id_escapes_session_id(self, tool_router, mock_client):
        """delete() URL-encodes session IDs before calling the raw endpoint."""
        mock_client.delete.return_value = {
            "session_id": "session/with spaces",
            "deleted": True,
        }

        tool_router.delete("session/with spaces")

        mock_client.delete.assert_called_once_with(
            "/api/v3.1/tool_router/session/session%2Fwith%20spaces",
            cast_to=t.Dict[str, t.Any],
        )

    def test_delete_session_by_id_calls_raw_http_endpoint(self):
        """delete() works with the low-level Stainless HTTP client."""
        requests: t.List[httpx.Request] = []

        def handler(request: httpx.Request) -> httpx.Response:
            requests.append(request)
            return httpx.Response(
                200,
                json={
                    "session_id": "session/with spaces",
                    "deleted": True,
                },
            )

        client = HttpClient(
            provider="test",
            api_key="sk-test",
            base_url="https://backend.invalid",
            http_client=httpx.Client(transport=httpx.MockTransport(handler)),
        )
        tool_router = ToolRouter(client=client, provider=MagicMock())

        result = tool_router.delete("session/with spaces")

        assert result == {
            "session_id": "session/with spaces",
            "deleted": True,
        }
        assert len(requests) == 1
        assert requests[0].method == "DELETE"
        assert (
            str(requests[0].url)
            == "https://backend.invalid/api/v3.1/tool_router/session/session%2Fwith%20spaces"
        )

    def test_create_session_with_toolkits_list(self, tool_router, mock_client):
        """Test creating a session with toolkits as a list."""
        session = tool_router.create(user_id="user_123", toolkits=["github", "slack"])

        assert session.session_id == "session_123"

        # Verify the API was called with correct parameters
        call_args = mock_client.tool_router.session.create.call_args
        assert call_args is not None
        kwargs = call_args.kwargs
        assert "toolkits" in kwargs
        assert kwargs["toolkits"] == {"enable": ["github", "slack"]}

    def test_create_session_with_disable_toolkits(self, tool_router, mock_client):
        """Test creating a session with disable toolkits."""
        session = tool_router.create(
            user_id="user_123", toolkits={"disable": ["linear", "notion"]}
        )

        assert session.session_id == "session_123"

        # Verify the API was called with correct parameters
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["toolkits"] == {"disable": ["linear", "notion"]}

    def test_create_session_with_tools_config(self, tool_router, mock_client):
        """Test creating a session with per-toolkit tool configuration."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": ["GMAIL_SEND_EMAIL", "GMAIL_SEARCH"],
                "github": {"enable": ["GITHUB_CREATE_ISSUE"]},
            },
        )

        assert session.session_id == "session_123"

        # Verify the API was called with tools config
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tools" in kwargs
        assert "gmail" in kwargs["tools"]
        assert kwargs["tools"]["gmail"] == {
            "enable": ["GMAIL_SEND_EMAIL", "GMAIL_SEARCH"]
        }
        assert kwargs["tools"]["github"] == {"enable": ["GITHUB_CREATE_ISSUE"]}

    def test_create_session_with_global_tags(self, tool_router, mock_client):
        """Test creating a session with global tag filtering (array format)."""
        session = tool_router.create(
            user_id="user_123", tags=["readOnlyHint", "idempotentHint"]
        )

        assert session.session_id == "session_123"

        # Verify the API was called with tags transformed to enable format
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tags" in kwargs
        assert kwargs["tags"] == {"enable": ["readOnlyHint", "idempotentHint"]}

    def test_create_session_with_global_tags_object_enable(
        self, tool_router, mock_client
    ):
        """Test creating a session with global tags using object format with enable."""
        session = tool_router.create(
            user_id="user_123", tags={"enable": ["readOnlyHint", "idempotentHint"]}
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tags" in kwargs
        assert kwargs["tags"] == {"enable": ["readOnlyHint", "idempotentHint"]}

    def test_create_session_with_global_tags_object_disable(
        self, tool_router, mock_client
    ):
        """Test creating a session with global tags using object format with disable."""
        session = tool_router.create(
            user_id="user_123", tags={"disable": ["destructiveHint"]}
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tags" in kwargs
        assert kwargs["tags"] == {"disable": ["destructiveHint"]}

    def test_create_session_with_global_tags_object_both(
        self, tool_router, mock_client
    ):
        """Test creating a session with global tags using object format with both enable and disable."""
        session = tool_router.create(
            user_id="user_123",
            tags={
                "enable": ["readOnlyHint", "idempotentHint"],
                "disable": ["destructiveHint"],
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tags" in kwargs
        assert kwargs["tags"] == {
            "enable": ["readOnlyHint", "idempotentHint"],
            "disable": ["destructiveHint"],
        }

    def test_create_session_with_toolkit_specific_tags(self, tool_router, mock_client):
        """Test creating a session with per-toolkit tag filtering (array format)."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": {"tags": ["readOnlyHint"]},
                "github": {"tags": ["readOnlyHint", "idempotentHint"]},
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tools" in kwargs
        assert kwargs["tools"]["gmail"] == {"tags": {"enable": ["readOnlyHint"]}}
        assert kwargs["tools"]["github"] == {
            "tags": {"enable": ["readOnlyHint", "idempotentHint"]}
        }

    def test_create_session_with_toolkit_specific_tags_object_enable(
        self, tool_router, mock_client
    ):
        """Test creating a session with per-toolkit tags using object format with enable."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": {"tags": {"enable": ["readOnlyHint", "idempotentHint"]}},
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tools" in kwargs
        assert kwargs["tools"]["gmail"] == {
            "tags": {"enable": ["readOnlyHint", "idempotentHint"]}
        }

    def test_create_session_with_toolkit_specific_tags_object_disable(
        self, tool_router, mock_client
    ):
        """Test creating a session with per-toolkit tags using object format with disable."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": {"tags": {"disable": ["destructiveHint"]}},
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tools" in kwargs
        assert kwargs["tools"]["gmail"] == {"tags": {"disable": ["destructiveHint"]}}

    def test_create_session_with_toolkit_specific_tags_object_both(
        self, tool_router, mock_client
    ):
        """Test creating a session with per-toolkit tags using object format with both enable and disable."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": {
                    "tags": {
                        "enable": ["readOnlyHint", "idempotentHint"],
                        "disable": ["destructiveHint"],
                    }
                },
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tools" in kwargs
        assert kwargs["tools"]["gmail"] == {
            "tags": {
                "enable": ["readOnlyHint", "idempotentHint"],
                "disable": ["destructiveHint"],
            }
        }

    def test_create_session_with_open_world_hint_tag(self, tool_router, mock_client):
        """Test creating a session with openWorldHint tag."""
        session = tool_router.create(
            user_id="user_123",
            tags=["openWorldHint", "readOnlyHint"],
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "tags" in kwargs
        assert kwargs["tags"] == {"enable": ["openWorldHint", "readOnlyHint"]}

    def test_create_session_with_mixed_tools_config(self, tool_router, mock_client):
        """Test creating a session with mixed tool configuration (enable, disable, tags)."""
        session = tool_router.create(
            user_id="user_123",
            tools={
                "gmail": ["GMAIL_SEND_EMAIL"],  # List shorthand
                "slack": {"disable": ["SLACK_DELETE_MESSAGE"]},
                "github": {"tags": ["readOnlyHint"]},
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["tools"]["gmail"] == {"enable": ["GMAIL_SEND_EMAIL"]}
        assert kwargs["tools"]["slack"] == {"disable": ["SLACK_DELETE_MESSAGE"]}
        assert kwargs["tools"]["github"] == {"tags": {"enable": ["readOnlyHint"]}}

    def test_create_session_with_global_and_toolkit_tags(
        self, tool_router, mock_client
    ):
        """Test creating a session with both global tags and toolkit-specific tags."""
        session = tool_router.create(
            user_id="user_123",
            tags=["readOnlyHint"],  # Global tags
            tools={
                "gmail": {
                    "tags": {
                        "enable": ["idempotentHint"],
                        "disable": ["destructiveHint"],
                    }
                },
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["tags"] == {"enable": ["readOnlyHint"]}
        assert kwargs["tools"]["gmail"] == {
            "tags": {"enable": ["idempotentHint"], "disable": ["destructiveHint"]}
        }

    def test_create_session_with_empty_tags_array(self, tool_router, mock_client):
        """Test creating a session with empty tags array."""
        session = tool_router.create(
            user_id="user_123",
            tags=[],
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["tags"] == {"enable": []}

    def test_create_session_with_all_tag_types(self, tool_router, mock_client):
        """Test creating a session with all available tag types."""
        session = tool_router.create(
            user_id="user_123",
            tags=[
                "readOnlyHint",
                "destructiveHint",
                "idempotentHint",
                "openWorldHint",
            ],
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["tags"] == {
            "enable": [
                "readOnlyHint",
                "destructiveHint",
                "idempotentHint",
                "openWorldHint",
            ]
        }

    def test_create_session_with_manage_connections_boolean(
        self, tool_router, mock_client
    ):
        """Test creating a session with manage_connections as boolean."""
        session = tool_router.create(user_id="user_123", manage_connections=True)

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["manage_connections"]["enable"] is True

    def test_create_session_with_manage_connections_config(
        self, tool_router, mock_client
    ):
        """Test creating a session with manage_connections as config object."""
        session = tool_router.create(
            user_id="user_123",
            manage_connections={
                "enable": True,
                "callback_url": "https://example.com/callback",
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["manage_connections"]["enable"] is True
        assert (
            kwargs["manage_connections"]["callback_url"]
            == "https://example.com/callback"
        )

    def test_create_session_with_wait_for_connections(self, tool_router, mock_client):
        """Test creating a session with wait_for_connections in manage_connections config."""
        session = tool_router.create(
            user_id="user_123",
            manage_connections={
                "enable": True,
                "callback_url": "https://example.com/callback",
                "wait_for_connections": True,
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["manage_connections"]["enable"] is True
        assert (
            kwargs["manage_connections"]["callback_url"]
            == "https://example.com/callback"
        )
        assert kwargs["manage_connections"]["enable_wait_for_connections"] is True

    def test_create_session_with_auth_configs(self, tool_router, mock_client):
        """Test creating a session with auth configs."""
        session = tool_router.create(
            user_id="user_123", auth_configs={"github": "ac_xxx", "slack": "ac_yyy"}
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["auth_configs"] == {"github": "ac_xxx", "slack": "ac_yyy"}

    def test_create_session_with_connected_accounts(self, tool_router, mock_client):
        """Test creating a session with connected accounts."""
        session = tool_router.create(
            user_id="user_123",
            connected_accounts={"github": "ca_xxx", "slack": "ca_yyy"},
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["connected_accounts"] == {
            "github": ["ca_xxx"],
            "slack": ["ca_yyy"],
        }

    def test_create_session_with_sandbox_config(self, tool_router, mock_client):
        """Test creating a session with preferred sandbox configuration."""
        session = tool_router.create(
            user_id="user_123",
            sandbox={
                "enable_proxy_execution": False,
                "auto_offload_threshold": 300,
                "sandbox_size": "large",
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "workbench" in kwargs
        assert kwargs["workbench"] == {
            "enable": True,
            "enable_proxy_execution": False,
            "auto_offload_threshold": 300,
            "sandbox_size": "large",
        }

    def test_create_session_rejects_sandbox_and_workbench(
        self, tool_router, mock_client
    ):
        """Test creating a session rejects both sandbox and workbench."""
        with pytest.raises(InvalidParams):
            tool_router.create(
                user_id="user_123",
                sandbox={"enable": True},
                workbench={"enable": True},
            )

        mock_client.tool_router.session.create.assert_not_called()

    def test_create_session_with_workbench_config(self, tool_router, mock_client):
        """Test creating a session with backwards-compatible workbench configuration."""
        session = tool_router.create(
            user_id="user_123",
            workbench={"enable_proxy_execution": False, "auto_offload_threshold": 300},
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "workbench" in kwargs
        assert kwargs["workbench"]["enable"] is True
        assert kwargs["workbench"]["enable_proxy_execution"] is False
        assert kwargs["workbench"]["auto_offload_threshold"] == 300

    def test_create_session_with_workbench_disabled(self, tool_router, mock_client):
        """Test creating a session with workbench entirely disabled."""
        session = tool_router.create(
            user_id="user_123",
            workbench={"enable": False},
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "workbench" in kwargs
        assert kwargs["workbench"]["enable"] is False

    def test_create_session_with_workbench_enabled_and_configured(
        self, tool_router, mock_client
    ):
        """Test creating a session with workbench explicitly enabled and configured."""
        session = tool_router.create(
            user_id="user_123",
            workbench={
                "enable": True,
                "enable_proxy_execution": False,
                "auto_offload_threshold": 20000,
            },
        )

        assert session.session_id == "session_123"

        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "workbench" in kwargs
        assert kwargs["workbench"]["enable"] is True
        assert kwargs["workbench"]["enable_proxy_execution"] is False
        assert kwargs["workbench"]["auto_offload_threshold"] == 20000

    def test_create_session_with_workbench_sandbox_size(self, tool_router, mock_client):
        """sandbox_size is forwarded on the wire when set, omitted otherwise."""
        tool_router.create(
            user_id="user_123",
            workbench={"sandbox_size": "large"},
        )
        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["workbench"]["sandbox_size"] == "large"

        tool_router.create(
            user_id="user_123",
            workbench={"enable_proxy_execution": True},
        )
        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert "sandbox_size" not in kwargs["workbench"]

    def test_create_session_with_multi_account(self, tool_router, mock_client):
        """multi_account is forwarded on the wire when set, omitted otherwise."""
        # Forwarded as-is when provided (TypedDict already uses snake_case keys).
        tool_router.create(
            user_id="user_123",
            multi_account={
                "enable": True,
                "max_accounts_per_toolkit": 3,
                "require_explicit_selection": True,
            },
        )
        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["multi_account"] == {
            "enable": True,
            "max_accounts_per_toolkit": 3,
            "require_explicit_selection": True,
        }

        # Omitted entirely when not provided.
        tool_router.create(user_id="user_123")
        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert "multi_account" not in kwargs

    def test_create_session_with_preload(self, tool_router, mock_client):
        """preload is forwarded and exposed on the session."""
        mock_client.tool_router.session.create.return_value.config.preload.tools = [
            "GMAIL_FETCH_EMAILS"
        ]

        session = tool_router.create(
            user_id="user_123",
            preload={"tools": ["GMAIL_FETCH_EMAILS"]},
        )

        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["preload"] == {"tools": ["GMAIL_FETCH_EMAILS"]}
        assert session.preload.tools == ["GMAIL_FETCH_EMAILS"]

    def test_create_session_with_preload_all(self, tool_router, mock_client):
        """preload tools='all' is forwarded and exposed on the session."""
        mock_client.tool_router.session.create.return_value.config.preload.tools = "all"

        session = tool_router.create(
            user_id="user_123",
            preload={"tools": "all"},
        )

        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["preload"] == {"tools": "all"}
        assert session.preload.tools == "all"

    def test_create_session_rejects_custom_slug_in_top_level_preload(
        self, tool_router, mock_client
    ):
        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
        )
        def grep(input: GrepInput, ctx):
            return {"matches": []}

        with pytest.raises(
            ValidationError,
            match="Set preload=True on the SDK custom tool",
        ):
            tool_router.create(
                user_id="user_123",
                preload={"tools": ["GREP"]},
                experimental={"custom_tools": [grep]},
            )

        mock_client.tool_router.session.create.assert_not_called()

    def test_create_session_rejects_bare_string_preload_tools(
        self, tool_router, mock_client
    ):
        with pytest.raises(
            ValidationError,
            match='preload.tools must be a list of Composio tool slugs or "all"',
        ):
            tool_router.create(
                user_id="user_123",
                preload={"tools": "GMAIL_FETCH_EMAILS"},
            )

        mock_client.tool_router.session.create.assert_not_called()

    def test_create_session_with_direct_tools_preset(self, tool_router, mock_client):
        """direct_tools applies SDK-side defaults for direct tool exposure."""
        mock_client.tool_router.session.create.return_value.config.preload.tools = "all"

        session = tool_router.create(
            user_id="user_123",
            session_preset=SESSION_PRESET_DIRECT_TOOLS,
            toolkits=["github"],
        )

        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["toolkits"] == {"enable": ["github"]}
        assert kwargs["manage_connections"] == {"enable": False}
        assert kwargs["workbench"] == {"enable": False}
        assert kwargs["preload"] == {"tools": "all"}
        assert kwargs["search"] == {"enable": False}
        assert kwargs["execute"] == {"enable_multi_execute": False}
        assert session.preload.tools == "all"

    def test_create_session_with_direct_tools_preset_overrides(
        self, tool_router, mock_client
    ):
        """Explicit user fields override direct_tools defaults."""
        mock_client.tool_router.session.create.return_value.config.preload.tools = [
            "GITHUB_CREATE_ISSUE"
        ]

        session = tool_router.create(
            user_id="user_123",
            session_preset=SESSION_PRESET_DIRECT_TOOLS,
            toolkits=["github"],
            manage_connections=True,
            workbench={"enable": True},
            preload={"tools": ["GITHUB_CREATE_ISSUE"]},
        )

        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["manage_connections"] == {"enable": True}
        assert kwargs["workbench"] == {"enable": True}
        assert kwargs["preload"] == {"tools": ["GITHUB_CREATE_ISSUE"]}
        assert kwargs["search"] == {"enable": False}
        assert kwargs["execute"] == {"enable_multi_execute": False}
        assert session.preload.tools == ["GITHUB_CREATE_ISSUE"]

    def test_direct_tools_does_not_default_custom_preload_when_preload_overridden(
        self, tool_router, mock_client
    ):
        """Explicit preload config also overrides the custom-tool preload default."""

        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
        )
        def grep(input: GrepInput, ctx):
            return {"matches": []}

        mock_response = mock_client.tool_router.session.create.return_value
        mock_response.config.preload.tools = ["GITHUB_CREATE_ISSUE"]
        mock_response.experimental = MagicMock()
        mock_response.experimental.assistive_prompt = None
        mock_response.experimental.custom_toolkits = []
        mock_custom_tool = MagicMock()
        mock_custom_tool.slug = "SERVER_GREP"
        mock_custom_tool.original_slug = "GREP"
        mock_custom_tool.extends_toolkit = None
        mock_response.experimental.custom_tools = [mock_custom_tool]

        tool_router.create(
            user_id="user_123",
            session_preset=SESSION_PRESET_DIRECT_TOOLS,
            toolkits=["github"],
            preload={"tools": ["GITHUB_CREATE_ISSUE"]},
            experimental={"custom_tools": [grep]},
        )

        kwargs = mock_client.tool_router.session.create.call_args.kwargs
        assert kwargs["preload"] == {"tools": ["GITHUB_CREATE_ISSUE"]}
        assert "preload" not in kwargs["experimental"]["custom_tools"][0]

    def test_create_session_complex_config(self, tool_router, mock_client):
        """Test creating a session with complex configuration."""
        session = tool_router.create(
            user_id="user_123",
            toolkits=["github", "slack"],
            tools={
                "gmail": ["GMAIL_SEND_EMAIL"],
            },
            tags=["readOnlyHint", "idempotentHint"],
            manage_connections={
                "enable": True,
                "callback_url": "https://example.com/callback",
            },
            auth_configs={"github": "ac_xxx"},
            connected_accounts={"slack": "ca_yyy"},
            workbench={"enable_proxy_execution": True, "auto_offload_threshold": 600},
            multi_account={"enable": True, "max_accounts_per_toolkit": 5},
        )

        assert session.session_id == "session_123"
        assert session.mcp.type == ToolRouterMCPServerType.HTTP

        # Verify all parameters were passed correctly
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert kwargs["user_id"] == "user_123"
        assert "toolkits" in kwargs
        assert "tools" in kwargs
        assert "tags" in kwargs
        assert "manage_connections" in kwargs
        assert "auth_configs" in kwargs
        assert "connected_accounts" in kwargs
        assert "workbench" in kwargs
        assert kwargs["multi_account"] == {
            "enable": True,
            "max_accounts_per_toolkit": 5,
            "require_explicit_selection": True,
        }

    def test_create_session_with_experimental_config(self, tool_router, mock_client):
        """Test creating a session with experimental configuration."""
        session = tool_router.create(
            user_id="user_123",
            experimental={
                "assistive_prompt": {
                    "user_timezone": "America/New_York",
                }
            },
        )

        assert session.session_id == "session_123"

        # Verify the API was called with experimental config transformed correctly
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "experimental" in kwargs
        assert kwargs["experimental"] == {
            "assistive_prompt_config": {
                "user_timezone": "America/New_York",
            }
        }

    def test_create_session_with_experimental_empty_config(
        self, tool_router, mock_client
    ):
        """Test creating a session with empty experimental configuration."""
        session = tool_router.create(
            user_id="user_123",
            experimental={},
        )

        assert session.session_id == "session_123"

        # Verify experimental is not included when empty
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "experimental" not in kwargs

    def test_create_session_with_experimental_no_timezone(
        self, tool_router, mock_client
    ):
        """Test creating a session with experimental config but no timezone."""
        session = tool_router.create(
            user_id="user_123",
            experimental={
                "assistive_prompt": {},
            },
        )

        assert session.session_id == "session_123"

        # Verify experimental is not included when timezone is not set
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "experimental" not in kwargs

    def test_create_session_with_experimental_empty_string_timezone(
        self, tool_router, mock_client
    ):
        """Test creating a session with empty string timezone is excluded.

        This ensures consistency with TypeScript SDK which uses truthy check.
        """
        session = tool_router.create(
            user_id="user_123",
            experimental={
                "assistive_prompt": {
                    "user_timezone": "",  # Empty string should be excluded
                },
            },
        )

        assert session.session_id == "session_123"

        # Verify experimental is not included when timezone is empty string
        call_args = mock_client.tool_router.session.create.call_args
        kwargs = call_args.kwargs
        assert "experimental" not in kwargs

    def test_create_session_experimental_response_transformation(
        self, tool_router, mock_client
    ):
        """Test that experimental response is transformed correctly."""
        # Setup mock response with experimental
        mock_response_with_experimental = MagicMock()
        mock_response_with_experimental.session_id = "session_exp"
        mock_response_with_experimental.mcp = MagicMock()
        mock_response_with_experimental.mcp.type = "http"
        mock_response_with_experimental.mcp.url = "https://mcp.example.com/session_exp"
        mock_response_with_experimental.tool_router_tools = ["TOOL_1"]
        mock_response_with_experimental.experimental = MagicMock()
        mock_response_with_experimental.experimental.assistive_prompt = (
            "You are a helpful assistant working in the America/New_York timezone."
        )

        mock_client.tool_router.session.create.return_value = (
            mock_response_with_experimental
        )

        session = tool_router.create(
            user_id="user_123",
            experimental={
                "assistive_prompt": {
                    "user_timezone": "America/New_York",
                }
            },
        )

        # Verify experimental response is transformed correctly
        assert session.experimental is not None
        assert isinstance(session.experimental, ToolRouterSessionExperimental)
        assert (
            session.experimental.assistive_prompt
            == "You are a helpful assistant working in the America/New_York timezone."
        )
        assert session.experimental.files is not None

    def test_create_session_without_experimental_response(
        self, tool_router, mock_client
    ):
        """Test that session without experimental response has experimental with files only."""
        session = tool_router.create(user_id="user_123")

        # Verify experimental is present with files, assistive_prompt is None
        assert session.experimental is not None
        assert session.experimental.assistive_prompt is None
        assert session.experimental.files is not None

    def test_create_session_raises_error_without_provider(self, mock_client):
        """Test that creating a session without provider raises an error when calling tools()."""
        tool_router = ToolRouter(client=mock_client, provider=None)

        # The error is raised during create() because it calls _create_tools_fn
        with pytest.raises(ValueError, match="Provider is required for tool router"):
            tool_router.create(user_id="user_123")

    def test_use_session_raises_error_without_provider(self, mock_client):
        """Test that use() raises when provider is None."""
        tool_router = ToolRouter(client=mock_client, provider=None)

        with pytest.raises(ValueError, match="Provider is required for tool router"):
            tool_router.use(session_id="session_123")

    def test_use_session(self, tool_router, mock_client):
        """Test retrieving an existing session."""
        mock_client.tool_router.session.retrieve.return_value.config.preload.tools = [
            "GMAIL_FETCH_EMAILS"
        ]

        session = tool_router.use(session_id="session_123")

        # Verify session properties
        assert session.session_id == "session_123"
        assert session.mcp.type == ToolRouterMCPServerType.HTTP
        assert session.mcp.url == "https://mcp.example.com/session_123"
        assert session.mcp.headers == {"x-api-key": "test-api-key"}
        assert session.experimental.files is not None
        assert hasattr(session.experimental.files, "list")
        assert hasattr(session.experimental.files, "upload")
        assert hasattr(session.experimental.files, "download")
        assert hasattr(session.experimental.files, "delete")
        assert callable(session.search)
        assert callable(session.execute)
        assert callable(session.tools)
        assert callable(session.authorize)
        assert callable(session.toolkits)
        assert callable(session.delete)
        assert session.preload.tools == ["GMAIL_FETCH_EMAILS"]

        mock_client.tool_router.session.retrieve.assert_called_once_with("session_123")
        mock_client.tool_router.session.attach.assert_not_called()
        mock_client.post.assert_not_called()

    def test_session_delete(self, tool_router, mock_client):
        """Session delete removes the current session."""
        mock_client.delete.return_value = {
            "session_id": "session_123",
            "deleted": True,
        }
        session = tool_router.use(session_id="session_123")

        result = session.delete()

        mock_client.delete.assert_called_once_with(
            "/api/v3.1/tool_router/session/session_123",
            cast_to=t.Dict[str, t.Any],
        )
        assert result == {
            "session_id": "session_123",
            "deleted": True,
        }

    def test_use_session_with_custom_tools(self, tool_router, mock_client):
        """Test attaching custom tools when reusing a session."""

        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
        )
        def grep(input: GrepInput, ctx):
            return {"matches": []}

        mock_response = mock_client.tool_router.session.attach.return_value
        mock_response.experimental = MagicMock()
        mock_response.experimental.assistive_prompt = None
        mock_response.experimental.custom_toolkits = []
        mock_custom_tool = MagicMock()
        mock_custom_tool.slug = "SERVER_GREP"
        mock_custom_tool.original_slug = "GREP"
        mock_custom_tool.extends_toolkit = None
        mock_response.experimental.custom_tools = [mock_custom_tool]

        session = tool_router.use(session_id="session_123", custom_tools=[grep])

        mock_client.tool_router.session.attach.assert_called_once()
        kwargs = mock_client.tool_router.session.attach.call_args.kwargs
        assert kwargs["experimental"]["custom_tools"][0]["slug"] == "GREP"
        assert "custom_toolkits" not in kwargs["experimental"]
        mock_client.tool_router.session.retrieve.assert_not_called()
        mock_client.post.assert_not_called()
        assert session.custom_tools()[0].slug == "SERVER_GREP"
        assert session.custom_tools()[0].name == "Grep"

    @patch("composio.core.models.tools.Tools")
    def test_use_session_attaches_preloaded_custom_tools(
        self,
        mock_tools_class,
        tool_router,
        mock_client,
        mock_provider,
    ):
        """use() attaches customs and exposes SDK-preloaded schemas locally."""

        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
            preload=True,
        )
        def grep(input: GrepInput, ctx):
            return {"matches": [input.pattern]}

        attach_response = mock_client.tool_router.session.attach.return_value
        attach_response.tool_router_tools = ["COMPOSIO_SEARCH_TOOLS"]
        attach_response.config.preload.tools = []
        attach_response.experimental = MagicMock()
        attach_response.experimental.assistive_prompt = None
        attach_response.experimental.custom_toolkits = []
        mock_custom_tool = MagicMock()
        mock_custom_tool.slug = "LOCAL_GREP"
        mock_custom_tool.original_slug = "GREP"
        mock_custom_tool.extends_toolkit = None
        attach_response.experimental.custom_tools = [mock_custom_tool]

        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "COMPOSIO_SEARCH_TOOLS"
        mock_tool.toolkit.slug = "composio"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_class.return_value = mock_tools_instance
        mock_provider.wrap_tools.return_value = "mocked-custom-tools"

        session = tool_router.use(session_id="session_123", custom_tools=[grep])
        session.tools()

        mock_client.tool_router.session.attach.assert_called_once()
        attach_kwargs = mock_client.tool_router.session.attach.call_args.kwargs
        assert attach_kwargs["experimental"]["custom_tools"][0]["slug"] == "GREP"
        assert attach_kwargs["experimental"]["custom_tools"][0]["preload"] is True
        mock_client.post.assert_not_called()
        mock_client.tool_router.session.retrieve.assert_not_called()

        wrapped_tools = mock_provider.wrap_tools.call_args.kwargs["tools"]
        assert [tool.slug for tool in wrapped_tools] == [
            "COMPOSIO_SEARCH_TOOLS",
            "LOCAL_GREP",
        ]

    @patch("composio.core.models.tools.Tools")
    def test_use_session_applies_preload_all_to_custom_tools(
        self,
        mock_tools_class,
        tool_router,
        mock_client,
        mock_provider,
    ):
        """preload.tools='all' on an existing session selects attached customs."""

        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
        )
        def grep(input: GrepInput, ctx):
            return {"matches": [input.pattern]}

        attach_response = mock_client.tool_router.session.attach.return_value
        attach_response.tool_router_tools = ["COMPOSIO_SEARCH_TOOLS"]
        attach_response.config.preload.tools = "all"
        attach_response.experimental = MagicMock()
        attach_response.experimental.assistive_prompt = None
        attach_response.experimental.custom_toolkits = []
        mock_custom_tool = MagicMock()
        mock_custom_tool.slug = "LOCAL_GREP"
        mock_custom_tool.original_slug = "GREP"
        mock_custom_tool.extends_toolkit = None
        attach_response.experimental.custom_tools = [mock_custom_tool]

        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "COMPOSIO_SEARCH_TOOLS"
        mock_tool.toolkit.slug = "composio"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_class.return_value = mock_tools_instance
        mock_provider.wrap_tools.return_value = "mocked-custom-tools"

        session = tool_router.use(session_id="session_123", custom_tools=[grep])
        session.tools()
        session.search(query="search local text")

        wrapped_tools = mock_provider.wrap_tools.call_args.kwargs["tools"]
        assert [tool.slug for tool in wrapped_tools] == [
            "COMPOSIO_SEARCH_TOOLS",
            "LOCAL_GREP",
        ]
        mock_client.tool_router.session.search.assert_called_once()
        search_kwargs = mock_client.tool_router.session.search.call_args.kwargs
        assert search_kwargs["experimental"]["custom_tools"][0]["slug"] == "GREP"
        assert search_kwargs["experimental"]["custom_tools"][0]["preload"] is True

    def test_use_session_with_different_user(self, tool_router, mock_client):
        """Test that use() extracts user_id from session config."""
        mock_attach_response = MagicMock()
        mock_attach_response.session_id = "session_456"
        mock_attach_response.mcp = MagicMock()
        mock_attach_response.mcp.type = "http"
        mock_attach_response.mcp.url = "https://mcp.example.com/session_456"
        mock_attach_response.tool_router_tools = ["TOOL_1"]
        mock_attach_response.config = MagicMock()
        mock_attach_response.config.user_id = "custom_user_789"
        mock_attach_response.config.preload = MagicMock()
        mock_attach_response.config.preload.tools = []
        mock_attach_response.experimental = None

        mock_client.tool_router.session.retrieve.return_value = mock_attach_response

        session = tool_router.use(session_id="session_456")

        assert session.session_id == "session_456"
        assert session.mcp.type == ToolRouterMCPServerType.HTTP

    def test_use_session_throws_error_on_failure(self, tool_router, mock_client):
        """Test that use() throws error if retrieve fails."""
        mock_client.tool_router.session.retrieve.side_effect = Exception(
            "Session not found"
        )

        with pytest.raises(Exception, match="Session not found"):
            tool_router.use(session_id="invalid_session")

    def test_session_execute_with_account(self, tool_router, mock_client):
        """session.execute forwards account for direct app tool execution."""
        mock_execute_response = MagicMock()
        mock_execute_response.data = {"ok": True}
        mock_execute_response.error = None
        mock_execute_response.log_id = "log_123"
        mock_client.tool_router.session.execute.return_value = mock_execute_response

        session = tool_router.create(user_id="user_123")
        session.execute(
            "GMAIL_FETCH_EMAILS",
            arguments={"max_results": 1},
            account="work",
        )

        mock_client.tool_router.session.execute.assert_called_once_with(
            session_id="session_123",
            tool_slug="GMAIL_FETCH_EMAILS",
            arguments={"max_results": 1},
            account="work",
            experimental=omit,
        )

    def test_authorize_function(self, tool_router, mock_client):
        """Test the authorize function returned by session."""
        session = tool_router.create(user_id="user_123")

        session.authorize("github")

        # Verify link was called
        mock_client.tool_router.session.link.assert_called_once()
        call_args = mock_client.tool_router.session.link.call_args
        assert call_args.kwargs["session_id"] == "session_123"
        assert call_args.kwargs["toolkit"] == "github"

    def test_authorize_function_with_callback_url(self, tool_router, mock_client):
        """Test the authorize function with callback URL."""
        session = tool_router.create(user_id="user_123")

        session.authorize("github", callback_url="https://myapp.com/callback")

        # Verify link was called with callback_url
        call_args = mock_client.tool_router.session.link.call_args
        assert call_args.kwargs.get("callback_url") == "https://myapp.com/callback"

    def test_authorize_forwards_experimental_block(self, tool_router, mock_client):
        """authorize() should pass the experimental options through verbatim."""
        session = tool_router.create(user_id="user_123")

        session.authorize(
            "github",
            experimental={
                "account_type": "SHARED",
                "acl_config_for_shared": {
                    "allow_all_users": True,
                    "not_allowed_user_ids": ["user_bob"],
                },
            },
        )

        call_kwargs = mock_client.tool_router.session.link.call_args.kwargs
        assert call_kwargs["experimental"] == {
            "account_type": "SHARED",
            "acl_config_for_shared": {
                "allow_all_users": True,
                "not_allowed_user_ids": ["user_bob"],
            },
        }

    def test_toolkits_function(self, tool_router, mock_client):
        """Test the toolkits function returned by session."""
        session = tool_router.create(user_id="user_123")

        toolkits_result = session.toolkits()

        # Verify toolkits was called
        mock_client.tool_router.session.toolkits.assert_called_once_with(
            session_id="session_123"
        )

        # Verify result structure
        assert isinstance(toolkits_result, ToolkitConnectionsDetails)
        assert len(toolkits_result.items) == 2
        assert toolkits_result.next_cursor == "cursor_789"
        assert toolkits_result.total_pages == 2

    def test_toolkits_function_with_pagination(self, tool_router, mock_client):
        """Test the toolkits function with pagination options."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(next_cursor="cursor_abc", limit=10)

        # Verify toolkits was called with pagination params
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("cursor") == "cursor_abc"
        assert call_args.kwargs.get("limit") == 10

    def test_toolkits_transform_active_connection(self, tool_router, mock_client):
        """Test that toolkits function correctly transforms active connections."""
        session = tool_router.create(user_id="user_123")

        toolkits_result = session.toolkits()

        gmail_toolkit = toolkits_result.items[0]
        assert gmail_toolkit.slug == "gmail"
        assert gmail_toolkit.name == "Gmail"
        assert gmail_toolkit.logo == "https://example.com/gmail-logo.png"
        assert gmail_toolkit.is_no_auth is False
        assert gmail_toolkit.connection.is_active is True
        assert gmail_toolkit.connection.auth_config is not None
        assert gmail_toolkit.connection.auth_config.id == "auth_config_123"
        assert gmail_toolkit.connection.auth_config.mode == "OAUTH2"
        assert gmail_toolkit.connection.auth_config.is_composio_managed is True
        assert gmail_toolkit.connection.connected_account is not None
        assert gmail_toolkit.connection.connected_account.id == "conn_123"
        assert gmail_toolkit.connection.connected_account.status == "ACTIVE"

    def test_toolkits_transform_no_connection(self, tool_router, mock_client):
        """Test that toolkits function correctly transforms toolkits with no connection."""
        session = tool_router.create(user_id="user_123")

        toolkits_result = session.toolkits()

        github_toolkit = toolkits_result.items[1]
        assert github_toolkit.slug == "github"
        assert github_toolkit.name == "GitHub"
        assert github_toolkit.connection.is_active is False
        assert github_toolkit.connection.auth_config is None
        assert github_toolkit.connection.connected_account is None

    def test_toolkits_function_with_search(self, tool_router, mock_client):
        """Test the toolkits function with search parameter."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(search="gmail")

        # Verify toolkits was called with search param
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == "gmail"
        assert call_args.kwargs["session_id"] == "session_123"

    def test_toolkits_function_with_search_and_toolkits_filter(
        self, tool_router, mock_client
    ):
        """Test the toolkits function with search and toolkits filter together."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(search="github", toolkits=["github", "slack"])

        # Verify toolkits was called with both search and toolkits params
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == "github"
        assert call_args.kwargs.get("toolkits") == ["github", "slack"]
        assert call_args.kwargs["session_id"] == "session_123"

    def test_toolkits_function_with_search_and_is_connected(
        self, tool_router, mock_client
    ):
        """Test the toolkits function with search and is_connected filter."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(search="mail", is_connected=True)

        # Verify toolkits was called with search and is_connected params
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == "mail"
        assert call_args.kwargs.get("is_connected") is True
        assert call_args.kwargs["session_id"] == "session_123"

    def test_toolkits_function_with_search_and_pagination(
        self, tool_router, mock_client
    ):
        """Test the toolkits function with search and pagination parameters."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(search="slack", next_cursor="cursor_xyz", limit=5)

        # Verify toolkits was called with search and pagination params
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == "slack"
        assert call_args.kwargs.get("cursor") == "cursor_xyz"
        assert call_args.kwargs.get("limit") == 5
        assert call_args.kwargs["session_id"] == "session_123"

    def test_toolkits_function_with_search_all_params(self, tool_router, mock_client):
        """Test the toolkits function with search and all other parameters."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(
            search="gmail",
            toolkits=["gmail", "github"],
            next_cursor="cursor_123",
            limit=20,
            is_connected=False,
        )

        # Verify toolkits was called with all params including search
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == "gmail"
        assert call_args.kwargs.get("toolkits") == ["gmail", "github"]
        assert call_args.kwargs.get("cursor") == "cursor_123"
        assert call_args.kwargs.get("limit") == 20
        assert call_args.kwargs.get("is_connected") is False
        assert call_args.kwargs["session_id"] == "session_123"

    def test_toolkits_function_with_empty_search(self, tool_router, mock_client):
        """Test the toolkits function with empty search string."""
        session = tool_router.create(user_id="user_123")

        session.toolkits(search="")

        # Verify toolkits was called with empty search string
        call_args = mock_client.tool_router.session.toolkits.call_args
        assert call_args.kwargs.get("search") == ""
        assert call_args.kwargs["session_id"] == "session_123"

    @patch("composio.core.models.tools.Tools")
    def test_tools_function(
        self, mock_tools_class, tool_router, mock_client, mock_provider
    ):
        """Test the tools function returned by session."""
        # Setup mock Tools instance
        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "GMAIL_FETCH_EMAILS"
        mock_tool.toolkit.slug = "gmail"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_instance._file_helper.process_file_uploadable_schema.return_value = {}
        mock_tools_class.return_value = mock_tools_instance

        # Setup provider mock
        mock_provider.wrap_tools.return_value = "mocked-wrapped-tools"

        session = tool_router.create(user_id="user_123")
        result = session.tools()

        # Verify Tools was instantiated (resolved flag; default ToolRouter is off)
        mock_tools_class.assert_called_once_with(
            client=mock_client,
            provider=mock_provider,
            dangerously_allow_auto_upload_download_files=False,
            sensitive_file_upload_protection=True,
            file_upload_path_deny_segments=None,
            file_upload_dirs=None,
        )

        # Verify get_raw_tool_router_meta_tools was called
        mock_tools_instance.get_raw_tool_router_meta_tools.assert_called_once()

        # Verify provider's wrap_tools was called
        mock_provider.wrap_tools.assert_called_once()
        wrapped_tools = mock_provider.wrap_tools.call_args.kwargs["tools"]
        assert [tool.slug for tool in wrapped_tools] == ["GMAIL_FETCH_EMAILS"]
        assert result == "mocked-wrapped-tools"

    @patch("composio.core.models.tools.Tools")
    def test_tools_function_includes_preloaded_custom_tools(
        self,
        mock_tools_class,
        tool_router,
        mock_client,
        mock_provider,
    ):
        """Custom tools with SDK preload enabled are exposed locally."""

        @experimental_api.tool(
            slug="GREP",
            name="Grep",
            description="Search local text",
            preload=True,
        )
        def grep(input: GrepInput, ctx):
            return {"matches": [input.pattern]}

        mock_response = mock_client.tool_router.session.create.return_value
        mock_response.tool_router_tools = ["COMPOSIO_SEARCH_TOOLS"]
        mock_response.experimental = MagicMock()
        mock_response.experimental.assistive_prompt = None
        mock_response.experimental.custom_toolkits = []
        mock_custom_tool = MagicMock()
        mock_custom_tool.slug = "SERVER_GREP"
        mock_custom_tool.original_slug = "GREP"
        mock_custom_tool.extends_toolkit = None
        mock_response.experimental.custom_tools = [mock_custom_tool]

        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "COMPOSIO_SEARCH_TOOLS"
        mock_tool.toolkit.slug = "composio"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_class.return_value = mock_tools_instance
        mock_provider.wrap_tools.return_value = "mocked-custom-tools"

        session = tool_router.create(
            user_id="user_123",
            experimental={"custom_tools": [grep]},
        )
        result = session.tools()

        wrapped_tools = mock_provider.wrap_tools.call_args.kwargs["tools"]
        assert [tool.slug for tool in wrapped_tools] == [
            "COMPOSIO_SEARCH_TOOLS",
            "SERVER_GREP",
        ]
        assert DIRECT_CUSTOM_TOOL_DESCRIPTION_PREFIX in wrapped_tools[1].description
        assert wrapped_tools[1].toolkit.slug == "custom"
        assert result == "mocked-custom-tools"

        execute_tool = mock_provider.wrap_tools.call_args.kwargs["execute_tool"]
        local_result = execute_tool(slug="SERVER_GREP", arguments={"pattern": "needle"})
        assert local_result == {
            "data": {"matches": ["needle"]},
            "error": None,
            "successful": True,
        }
        mock_client.tool_router.session.execute.assert_not_called()

    @patch("composio.core.models.tools.Tools")
    def test_tools_function_with_modifiers(
        self,
        mock_tools_class,
        tool_router,
        mock_client,
        mock_provider,
    ):
        """Test the tools function with modifiers."""
        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "GMAIL_FETCH_EMAILS"
        mock_tool.toolkit.slug = "gmail"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_instance._file_helper.process_file_uploadable_schema.return_value = {}
        mock_tools_class.return_value = mock_tools_instance

        # Setup provider mock
        mock_provider.wrap_tools.return_value = "mocked-wrapped-tools"

        session = tool_router.create(user_id="user_123")

        # Create a proper modifier structure
        from composio.core.models._modifiers import Modifier

        modifier = Modifier(modifier=None, type_="schema", tools=[], toolkits=[])
        modifiers = [modifier]
        result = session.tools(modifiers=modifiers)

        # Verify provider's wrap_tools was called
        mock_provider.wrap_tools.assert_called_once()
        # Verify get_raw_tool_router_meta_tools was called with modifiers
        call_args = mock_tools_instance.get_raw_tool_router_meta_tools.call_args
        assert call_args.kwargs.get("modifiers") == modifiers
        assert result == "mocked-wrapped-tools"

    def test_session_mcp_type_http(self, tool_router, mock_client):
        """Test that MCP type is correctly set to HTTP enum."""
        session = tool_router.create(user_id="user_123")

        assert session.mcp.type == ToolRouterMCPServerType.HTTP
        assert session.mcp.type.value == "http"

    def test_session_mcp_sse_type(self, tool_router, mock_client):
        """Test that SSE MCP type is handled correctly."""
        mock_sse_response = MagicMock()
        mock_sse_response.session_id = "session_sse"
        mock_sse_response.mcp = MagicMock()
        mock_sse_response.mcp.type = "sse"
        mock_sse_response.mcp.url = "https://mcp.example.com/sse/session_sse"
        mock_sse_response.tool_router_tools = ["TOOL_1"]

        mock_client.tool_router.session.create.return_value = mock_sse_response

        session = tool_router.create(user_id="user_123")

        assert session.mcp.type == ToolRouterMCPServerType.SSE
        assert session.mcp.type.value == "sse"
        assert session.mcp.url == "https://mcp.example.com/sse/session_sse"

    def test_multiple_sessions_independently(self, tool_router, mock_client):
        """Test that multiple sessions can be created independently."""
        mock_session_1 = MagicMock()
        mock_session_1.session_id = "session_1"
        mock_session_1.mcp = MagicMock()
        mock_session_1.mcp.type = "http"
        mock_session_1.mcp.url = "https://mcp.example.com/session_1"
        mock_session_1.tool_router_tools = ["TOOL_1"]

        mock_session_2 = MagicMock()
        mock_session_2.session_id = "session_2"
        mock_session_2.mcp = MagicMock()
        mock_session_2.mcp.type = "http"
        mock_session_2.mcp.url = "https://mcp.example.com/session_2"
        mock_session_2.tool_router_tools = ["TOOL_2"]

        mock_client.tool_router.session.create.side_effect = [
            mock_session_1,
            mock_session_2,
        ]

        session1 = tool_router.create(user_id="user_1")
        session2 = tool_router.create(user_id="user_2")

        assert session1.session_id == "session_1"
        assert session2.session_id == "session_2"
        assert mock_client.tool_router.session.create.call_count == 2

    def test_use_vs_create_independence(self, tool_router, mock_client):
        """Test that use() and create() are independent."""
        # Setup different responses for create and attach
        mock_create_response = MagicMock()
        mock_create_response.session_id = "created_session"
        mock_create_response.mcp = MagicMock()
        mock_create_response.mcp.type = "http"
        mock_create_response.mcp.url = "https://mcp.example.com/created"
        mock_create_response.tool_router_tools = ["TOOL_1"]

        mock_attach_response = MagicMock()
        mock_attach_response.session_id = "retrieved_session"
        mock_attach_response.mcp = MagicMock()
        mock_attach_response.mcp.type = "http"
        mock_attach_response.mcp.url = "https://mcp.example.com/retrieved"
        mock_attach_response.tool_router_tools = ["TOOL_2"]
        mock_attach_response.config = MagicMock()
        mock_attach_response.config.user_id = "user_456"
        mock_attach_response.config.preload = MagicMock()
        mock_attach_response.config.preload.tools = []
        mock_attach_response.experimental = None

        mock_client.tool_router.session.create.return_value = mock_create_response
        mock_client.tool_router.session.retrieve.return_value = mock_attach_response

        # Create a new session
        created_session = tool_router.create(user_id="user_123")
        assert created_session.session_id == "created_session"

        # Use an existing session
        retrieved_session = tool_router.use(session_id="retrieved_session")
        assert retrieved_session.session_id == "retrieved_session"

        # Verify both methods were called
        mock_client.tool_router.session.create.assert_called_once()
        mock_client.tool_router.session.retrieve.assert_called_once_with(
            "retrieved_session"
        )
        mock_client.post.assert_not_called()


class TestToolRouterTypes:
    """Test type definitions for ToolRouter."""

    def test_toolkit_enable_config_type(self):
        """Test ToolRouterToolkitsEnableConfig type."""
        config: ToolRouterToolkitsEnableConfig = {"enable": ["github", "slack"]}
        assert "enable" in config
        assert config["enable"] == ["github", "slack"]

    def test_toolkit_disable_config_type(self):
        """Test ToolRouterToolkitsDisableConfig type."""
        config: ToolRouterToolkitsDisableConfig = {"disable": ["linear"]}
        assert "disable" in config

    def test_manage_connections_config_type(self):
        """Test ToolRouterManageConnectionsConfig type."""
        config: ToolRouterManageConnectionsConfig = {
            "enable": True,
            "callback_url": "https://example.com/callback",
        }
        assert config["enable"] is True
        assert config["callback_url"] == "https://example.com/callback"

    def test_experimental_config_type(self):
        """Test ToolRouterExperimentalConfig type."""
        config: ToolRouterExperimentalConfig = {
            "assistive_prompt": {
                "user_timezone": "America/New_York",
            }
        }
        assert "assistive_prompt" in config
        assert config["assistive_prompt"]["user_timezone"] == "America/New_York"

    def test_session_experimental_dataclass(self):
        """Test ToolRouterSessionExperimental dataclass."""
        mock_files = MagicMock()
        experimental = ToolRouterSessionExperimental(
            files=mock_files,
            assistive_prompt="You are a helpful assistant.",
        )
        assert experimental.assistive_prompt == "You are a helpful assistant."
        assert experimental.files is mock_files

        # Test with None assistive_prompt
        mock_files_none = MagicMock()
        experimental_none = ToolRouterSessionExperimental(
            files=mock_files_none,
            assistive_prompt=None,
        )
        assert experimental_none.assistive_prompt is None
        assert experimental_none.files is mock_files_none

    def test_mcp_server_config(self):
        """Test ToolRouterMCPServerConfig dataclass."""
        mcp = ToolRouterMCPServerConfig(
            type=ToolRouterMCPServerType.HTTP, url="https://mcp.example.com"
        )
        assert mcp.type == ToolRouterMCPServerType.HTTP
        assert mcp.type.value == "http"
        assert mcp.url == "https://mcp.example.com"
        assert mcp.headers is None  # Headers are optional

    def test_mcp_server_config_with_headers(self):
        """Test ToolRouterMCPServerConfig dataclass with headers."""
        mcp = ToolRouterMCPServerConfig(
            type=ToolRouterMCPServerType.HTTP,
            url="https://mcp.example.com",
            headers={"x-api-key": "test-api-key"},
        )
        assert mcp.type == ToolRouterMCPServerType.HTTP
        assert mcp.url == "https://mcp.example.com"
        assert mcp.headers == {"x-api-key": "test-api-key"}

    def test_toolkit_connection_state(self):
        """Test ToolkitConnectionState dataclass."""
        from composio.core.models.tool_router import (
            ToolkitConnectedAccount,
            ToolkitConnection,
            ToolkitConnectionAuthConfig,
        )

        connection = ToolkitConnection(
            is_active=True,
            auth_config=ToolkitConnectionAuthConfig(
                id="ac_123", mode="OAUTH2", is_composio_managed=True
            ),
            connected_account=ToolkitConnectedAccount(id="ca_123", status="ACTIVE"),
        )

        state = ToolkitConnectionState(
            slug="github",
            name="GitHub",
            is_no_auth=False,
            connection=connection,
            logo="https://example.com/logo.png",
        )

        assert state.slug == "github"
        assert state.name == "GitHub"
        assert state.connection.is_active is True
        assert state.connection.auth_config is not None
        assert state.connection.auth_config.id == "ac_123"


class TestToolRouterExecution:
    """Tests for tool router execution functionality."""

    @patch("composio.core.models.tools.Tools")
    def test_wrap_execute_tool_for_tool_router_called(
        self, mock_tools_class, tool_router, mock_client, mock_provider
    ):
        """Test that _wrap_execute_tool_for_tool_router is called when creating tools."""
        # Setup
        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "GMAIL_SEND_EMAIL"
        mock_tool.toolkit.slug = "gmail"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_instance._file_helper.process_file_uploadable_schema.return_value = {}
        mock_tools_instance._wrap_execute_tool_for_tool_router = MagicMock(
            return_value=lambda slug, args: {
                "data": {},
                "error": None,
                "successful": True,
            }
        )
        mock_tools_class.return_value = mock_tools_instance

        # Setup provider mock
        mock_provider.wrap_tools.return_value = ["wrapped_tool"]

        # Create session and get tools
        session = tool_router.create(user_id="user_123")
        session.tools()

        # Verify _wrap_execute_tool_for_tool_router was called with correct session_id
        mock_tools_instance._wrap_execute_tool_for_tool_router.assert_called_once()
        call_args = mock_tools_instance._wrap_execute_tool_for_tool_router.call_args
        assert call_args.kwargs["session_id"] == "session_123"

    def test_execute_endpoint_called_with_auto_offload_opt_in(
        self, tool_router, mock_client, mock_provider
    ):
        """Test that session execute is called for agentic tool execution."""
        mock_execute_response = MagicMock()
        mock_execute_response.data = {"result": "success"}
        mock_execute_response.error = None
        mock_client.tool_router.session.execute.return_value = mock_execute_response

        # Create a real Tools instance to test the execute function
        from composio.core.models.tools import Tools as RealTools

        real_tools = RealTools(
            client=mock_client,
            provider=mock_provider,
            dangerously_allow_auto_upload_download_files=False,
        )
        execute_fn = real_tools._wrap_execute_tool_for_tool_router(
            session_id="session_123"
        )

        # Execute the tool
        result = execute_fn("GMAIL_SEND_EMAIL", {"to": "test@example.com"})

        # Verify execute was called with direct offload opt-in for agentic calls
        mock_client.tool_router.session.execute.assert_called_once_with(
            session_id="session_123",
            tool_slug="GMAIL_SEND_EMAIL",
            arguments={"to": "test@example.com"},
            enable_auto_workbench_offload=True,
            experimental=omit,
        )

        # Verify result format
        assert result["data"] == {"result": "success"}
        assert result["error"] is None
        assert result["successful"] is True

    def test_execute_endpoint_passes_inline_custom_tools(
        self, tool_router, mock_client, mock_provider
    ):
        """Provider-wrapped session tools pass inline custom definitions."""
        mock_execute_response = MagicMock()
        mock_execute_response.data = {"result": "success"}
        mock_execute_response.error = None
        mock_client.tool_router.session.execute.return_value = mock_execute_response

        from composio.core.models.tools import Tools as RealTools

        inline_payload = {
            "custom_tools": [
                {
                    "slug": "GREP",
                    "name": "Grep",
                    "description": "Search local text",
                    "input_schema": {"type": "object", "properties": {}},
                }
            ]
        }
        real_tools = RealTools(
            client=mock_client,
            provider=mock_provider,
            dangerously_allow_auto_upload_download_files=False,
        )
        execute_fn = real_tools._wrap_execute_tool_for_tool_router(
            session_id="session_123",
            inline_custom_tools_payload=inline_payload,
        )

        execute_fn("GMAIL_SEND_EMAIL", {"to": "test@example.com"})

        mock_client.tool_router.session.execute.assert_called_once_with(
            session_id="session_123",
            tool_slug="GMAIL_SEND_EMAIL",
            arguments={"to": "test@example.com"},
            enable_auto_workbench_offload=True,
            experimental=inline_payload,
        )

    def test_modifiers_applied_in_tool_router_execution(
        self, tool_router, mock_client, mock_provider
    ):
        """Test that modifiers are applied before and after tool execution."""
        mock_execute_response = MagicMock()
        mock_execute_response.data = {"result": "success"}
        mock_execute_response.error = None
        mock_client.tool_router.session.execute.return_value = mock_execute_response

        # Create modifier functions
        def before_modifier(tool, toolkit, params):
            """Modifier to change arguments before execution."""
            params["arguments"] = {"to": "modified@example.com"}
            return params

        def after_modifier(tool, toolkit, response):
            """Modifier to change response after execution."""
            response["data"] = {"modified": True}
            return response

        # Create a real execute function with modifiers
        from composio.core.models._modifiers import Modifier
        from composio.core.models.tools import Tools as RealTools

        real_tools = RealTools(
            client=mock_client,
            provider=mock_provider,
            dangerously_allow_auto_upload_download_files=False,
        )

        modifiers = [
            Modifier(
                modifier=before_modifier, type_="before_execute", tools=[], toolkits=[]
            ),
            Modifier(
                modifier=after_modifier, type_="after_execute", tools=[], toolkits=[]
            ),
        ]

        execute_fn = real_tools._wrap_execute_tool_for_tool_router(
            session_id="session_123", modifiers=modifiers
        )

        # Execute the tool
        result = execute_fn("GMAIL_SEND_EMAIL", {"to": "test@example.com"})

        # Verify execute was called with modified arguments
        mock_client.tool_router.session.execute.assert_called_once()
        call_args = mock_client.tool_router.session.execute.call_args
        assert call_args.kwargs["arguments"] == {"to": "modified@example.com"}
        assert call_args.kwargs["enable_auto_workbench_offload"] is True

        # Verify result was modified by after_execute
        assert result["data"] == {"modified": True}

    def test_execute_tool_error_handling(self, tool_router, mock_client, mock_provider):
        """Test that tool execution errors are handled correctly."""
        mock_execute_response = MagicMock()
        mock_execute_response.data = {}
        mock_execute_response.error = "Authentication failed"
        mock_client.tool_router.session.execute.return_value = mock_execute_response

        # Create a real execute function
        from composio.core.models.tools import Tools as RealTools

        real_tools = RealTools(
            client=mock_client,
            provider=mock_provider,
            dangerously_allow_auto_upload_download_files=False,
        )
        execute_fn = real_tools._wrap_execute_tool_for_tool_router(
            session_id="session_123"
        )

        # Execute the tool
        result = execute_fn("GMAIL_SEND_EMAIL", {"to": "test@example.com"})

        # Verify error is returned correctly
        assert result["error"] == "Authentication failed"
        assert result["successful"] is False
        assert result["data"] == {}


class TestToolRouterIntegration:
    """Integration tests for ToolRouter."""

    @patch("composio.core.models.tools.Tools")
    def test_full_session_workflow(
        self, mock_tools_class, tool_router, mock_client, mock_provider
    ):
        """Test a complete session workflow."""
        # Setup
        mock_tools_instance = MagicMock()
        mock_tool = MagicMock()
        mock_tool.slug = "TOOL1"
        mock_tool.toolkit.slug = "github"
        mock_tool.input_parameters = {}
        mock_tools_instance.get_raw_tool_router_meta_tools.return_value = [mock_tool]
        mock_tools_instance._file_helper.enhance_schema_descriptions.return_value = {}
        mock_tools_instance._file_helper.process_file_uploadable_schema.return_value = {}
        mock_tools_class.return_value = mock_tools_instance

        # Setup provider mock
        mock_provider.wrap_tools.return_value = ["tool1", "tool2"]

        # Create session
        session = tool_router.create(
            user_id="user_123", toolkits=["github", "slack"], manage_connections=True
        )

        assert session.session_id == "session_123"

        # Use authorize
        session.authorize("github")
        assert mock_client.tool_router.session.link.called

        # Use toolkits
        toolkits = session.toolkits()
        assert len(toolkits.items) == 2

        # Use tools
        tools = session.tools()
        assert tools == ["tool1", "tool2"]

    def test_create_and_use_same_session(self, tool_router, mock_client):
        """Test creating a session and then attaching it."""
        # Create session
        created_session = tool_router.create(user_id="user_123")
        created_session_id = created_session.session_id

        # Setup retrieve to return the same session
        mock_client.tool_router.session.retrieve.return_value = (
            mock_client.tool_router.session.create.return_value
        )

        # Use the session
        retrieved_session = tool_router.use(session_id=created_session_id)

        assert created_session.session_id == retrieved_session.session_id
        assert created_session.mcp.url == retrieved_session.mcp.url