letta-server/tests/managers/test_mcp_manager.py

import uuid
from unittest.mock import AsyncMock, patch

import pytest

# Import shared fixtures and constants from conftest
from letta.constants import (
    MCP_TOOL_TAG_NAME_PREFIX,
)
from letta.functions.mcp_client.types import MCPTool
from letta.schemas.enums import (
    ToolType,
)
from letta.server.db import db_registry
from letta.settings import settings

# ======================================================================================================================
# MCPManager Tests
# ======================================================================================================================


@pytest.mark.asyncio
@patch("letta.services.mcp_manager.MCPManager.get_mcp_client")
async def test_create_mcp_server(mock_get_client, server, default_user):
    from letta.schemas.mcp import MCPServer, MCPServerType, SSEServerConfig, StdioServerConfig
    from letta.settings import tool_settings

    if tool_settings.mcp_read_from_config:
        return

    # create mock client with required methods
    mock_client = AsyncMock()
    mock_client.connect_to_server = AsyncMock()
    mock_client.list_tools = AsyncMock(
        return_value=[
            MCPTool(
                name="get_simple_price",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "ids": {"type": "string"},
                        "vs_currencies": {"type": "string"},
                        "include_market_cap": {"type": "boolean"},
                        "include_24hr_vol": {"type": "boolean"},
                        "include_24hr_change": {"type": "boolean"},
                    },
                    "required": ["ids", "vs_currencies"],
                    "additionalProperties": False,
                },
            )
        ]
    )
    mock_client.execute_tool = AsyncMock(
        return_value=(
            '{"bitcoin": {"usd": 50000, "usd_market_cap": 900000000000, "usd_24h_vol": 30000000000, "usd_24h_change": 2.5}}',
            True,
        )
    )
    mock_get_client.return_value = mock_client

    # Test with a valid StdioServerConfig
    server_config = StdioServerConfig(
        server_name="test_server", type=MCPServerType.STDIO, command="echo 'test'", args=["arg1", "arg2"], env={"ENV1": "value1"}
    )
    mcp_server = MCPServer(server_name="test_server", server_type=MCPServerType.STDIO, stdio_config=server_config)
    created_server = await server.mcp_manager.create_or_update_mcp_server(mcp_server, actor=default_user)
    print(created_server)
    assert created_server.server_name == server_config.server_name
    assert created_server.server_type == server_config.type

    # Test with a valid SSEServerConfig
    mcp_server_name = "coingecko"
    server_url = "https://mcp.api.coingecko.com/sse"
    sse_mcp_config = SSEServerConfig(server_name=mcp_server_name, server_url=server_url)
    mcp_sse_server = MCPServer(server_name=mcp_server_name, server_type=MCPServerType.SSE, server_url=server_url)
    created_server = await server.mcp_manager.create_or_update_mcp_server(mcp_sse_server, actor=default_user)
    print(created_server)
    assert created_server.server_name == mcp_server_name
    assert created_server.server_type == MCPServerType.SSE

    # list mcp servers
    servers = await server.mcp_manager.list_mcp_servers(actor=default_user)
    print(servers)
    assert len(servers) > 0, "No MCP servers found"

    # list tools from sse server
    tools = await server.mcp_manager.list_mcp_server_tools(created_server.server_name, actor=default_user)
    print(tools)

    # call a tool from the sse server
    tool_name = "get_simple_price"
    tool_args = {
        "ids": "bitcoin",
        "vs_currencies": "usd",
        "include_market_cap": True,
        "include_24hr_vol": True,
        "include_24hr_change": True,
    }
    result = await server.mcp_manager.execute_mcp_server_tool(
        created_server.server_name, tool_name=tool_name, tool_args=tool_args, actor=default_user, environment_variables={}
    )
    print(result)

    # add a tool
    tool = await server.mcp_manager.add_tool_from_mcp_server(created_server.server_name, tool_name, actor=default_user)
    print(tool)
    assert tool.name == tool_name
    assert f"mcp:{created_server.server_name}" in tool.tags, f"Expected tag {f'mcp:{created_server.server_name}'}, got {tool.tags}"
    print("TAGS", tool.tags)


@patch("letta.services.mcp_manager.MCPManager.get_mcp_client")
async def test_create_mcp_server_with_tools(mock_get_client, server, default_user):
    """Test that creating an MCP server automatically syncs and persists its tools."""
    from letta.functions.mcp_client.types import MCPToolHealth
    from letta.schemas.mcp import MCPServer, MCPServerType
    from letta.settings import tool_settings

    if tool_settings.mcp_read_from_config:
        return

    # Create mock tools with different health statuses
    mock_tools = [
        MCPTool(
            name="valid_tool_1",
            description="A valid tool",
            inputSchema={
                "type": "object",
                "properties": {
                    "param1": {"type": "string"},
                },
                "required": ["param1"],
            },
            health=MCPToolHealth(status="VALID", reasons=[]),
        ),
        MCPTool(
            name="valid_tool_2",
            description="Another valid tool",
            inputSchema={
                "type": "object",
                "properties": {
                    "param2": {"type": "number"},
                },
            },
            health=MCPToolHealth(status="VALID", reasons=[]),
        ),
        MCPTool(
            name="invalid_tool",
            description="An invalid tool that should be skipped",
            inputSchema={
                "type": "invalid_type",  # Invalid schema
            },
            health=MCPToolHealth(status="INVALID", reasons=["Invalid schema type"]),
        ),
        MCPTool(
            name="warning_tool",
            description="A tool with warnings but should still be persisted",
            inputSchema={
                "type": "object",
                "properties": {},
            },
            health=MCPToolHealth(status="WARNING", reasons=["No properties defined"]),
        ),
    ]

    # Create mock client
    mock_client = AsyncMock()
    mock_client.connect_to_server = AsyncMock()
    mock_client.list_tools = AsyncMock(return_value=mock_tools)
    mock_client.cleanup = AsyncMock()
    mock_get_client.return_value = mock_client

    # Create MCP server config
    server_name = f"test_server_{uuid.uuid4().hex[:8]}"
    server_url = "https://test-with-tools.example.com/sse"
    mcp_server = MCPServer(server_name=server_name, server_type=MCPServerType.SSE, server_url=server_url)

    # Create server with tools using the new method
    created_server = await server.mcp_manager.create_mcp_server_with_tools(mcp_server, actor=default_user)

    # Verify server was created
    assert created_server.server_name == server_name
    assert created_server.server_type == MCPServerType.SSE
    assert created_server.server_url == server_url

    # Verify tools were persisted (all except the invalid one)
    # Get all tools and filter by checking metadata
    all_tools = await server.tool_manager.list_tools_async(
        actor=default_user, names=["valid_tool_1", "valid_tool_2", "warning_tool", "invalid_tool"]
    )

    # Filter tools that belong to our MCP server
    persisted_tools = [
        tool
        for tool in all_tools
        if tool.metadata_
        and MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_
        and tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX].get("server_name") == server_name
    ]

    # Should have 3 tools (2 valid + 1 warning, but not the invalid one)
    assert len(persisted_tools) == 3, f"Expected 3 tools, got {len(persisted_tools)}"

    # Check tool names
    tool_names = {tool.name for tool in persisted_tools}
    assert "valid_tool_1" in tool_names
    assert "valid_tool_2" in tool_names
    assert "warning_tool" in tool_names
    assert "invalid_tool" not in tool_names  # Invalid tool should be filtered out

    # Verify each tool has correct metadata
    for tool in persisted_tools:
        assert tool.metadata_ is not None
        assert MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_
        assert tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX]["server_name"] == server_name
        assert tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX]["server_id"] == created_server.id
        assert tool.tool_type == ToolType.EXTERNAL_MCP

    # Clean up - delete the server
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)

    # Verify tools were also deleted (cascade) by trying to get them again
    remaining_tools = await server.tool_manager.list_tools_async(actor=default_user, names=["valid_tool_1", "valid_tool_2", "warning_tool"])

    # Filter to see if any still belong to our deleted server
    remaining_mcp_tools = [
        tool
        for tool in remaining_tools
        if tool.metadata_
        and MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_
        and tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX].get("server_name") == server_name
    ]
    assert len(remaining_mcp_tools) == 0, "Tools should be deleted when server is deleted"


@pytest.mark.asyncio
@patch("letta.services.mcp_manager.MCPManager.get_mcp_client")
async def test_complex_schema_normalization(mock_get_client, server, default_user):
    """Test that complex MCP schemas with nested objects are normalized and accepted."""
    from letta.functions.mcp_client.types import MCPTool, MCPToolHealth
    from letta.schemas.mcp import MCPServer, MCPServerType
    from letta.settings import tool_settings

    if tool_settings.mcp_read_from_config:
        return

    # Create mock tools with complex schemas that would normally be INVALID
    # These schemas have: nested $defs, $ref references, missing additionalProperties
    mock_tools = [
        # 1. Nested object with $ref (like create_person)
        MCPTool(
            name="create_person",
            description="Create a person with nested address",
            inputSchema={
                "$defs": {
                    "Address": {
                        "type": "object",
                        "properties": {
                            "street": {"type": "string"},
                            "city": {"type": "string"},
                            "zip_code": {"type": "string"},
                        },
                        "required": ["street", "city", "zip_code"],
                    },
                    "Person": {
                        "type": "object",
                        "properties": {
                            "name": {"type": "string"},
                            "age": {"type": "integer"},
                            "address": {"$ref": "#/$defs/Address"},
                        },
                        "required": ["name", "age"],
                    },
                },
                "type": "object",
                "properties": {"person": {"$ref": "#/$defs/Person"}},
                "required": ["person"],
            },
            health=MCPToolHealth(
                status="INVALID",
                reasons=["root: 'additionalProperties' not explicitly set", "root.properties.person: Missing 'type'"],
            ),
        ),
        # 2. List of objects (like manage_tasks)
        MCPTool(
            name="manage_tasks",
            description="Manage multiple tasks",
            inputSchema={
                "$defs": {
                    "TaskItem": {
                        "type": "object",
                        "properties": {
                            "title": {"type": "string"},
                            "priority": {"type": "integer", "default": 1},
                            "completed": {"type": "boolean", "default": False},
                            "tags": {"type": "array", "items": {"type": "string"}},
                        },
                        "required": ["title"],
                    }
                },
                "type": "object",
                "properties": {
                    "tasks": {
                        "type": "array",
                        "items": {"$ref": "#/$defs/TaskItem"},
                    }
                },
                "required": ["tasks"],
            },
            health=MCPToolHealth(
                status="INVALID",
                reasons=["root: 'additionalProperties' not explicitly set", "root.properties.tasks.items: Missing 'type'"],
            ),
        ),
        # 3. Complex filter object with optional fields
        MCPTool(
            name="search_with_filters",
            description="Search with complex filters",
            inputSchema={
                "$defs": {
                    "SearchFilter": {
                        "type": "object",
                        "properties": {
                            "keywords": {"type": "array", "items": {"type": "string"}},
                            "min_score": {"type": "number"},
                            "categories": {"type": "array", "items": {"type": "string"}},
                        },
                        "required": ["keywords"],
                    }
                },
                "type": "object",
                "properties": {
                    "query": {"type": "string"},
                    "filters": {"$ref": "#/$defs/SearchFilter"},
                },
                "required": ["query", "filters"],
            },
            health=MCPToolHealth(
                status="INVALID",
                reasons=["root: 'additionalProperties' not explicitly set", "root.properties.filters: Missing 'type'"],
            ),
        ),
    ]

    # Create mock client
    mock_client = AsyncMock()
    mock_client.connect_to_server = AsyncMock()
    mock_client.list_tools = AsyncMock(return_value=mock_tools)
    mock_client.cleanup = AsyncMock()
    mock_get_client.return_value = mock_client

    # Create MCP server
    server_name = f"test_complex_schema_{uuid.uuid4().hex[:8]}"
    server_url = "https://test-complex.example.com/sse"
    mcp_server = MCPServer(server_name=server_name, server_type=MCPServerType.SSE, server_url=server_url)

    try:
        # Create server (this will auto-sync tools)
        created_server = await server.mcp_manager.create_mcp_server_with_tools(mcp_server, actor=default_user)

        assert created_server.server_name == server_name

        # Now attempt to add each tool - they should be normalized from INVALID to acceptable
        # The normalization happens in add_tool_from_mcp_server

        # Test 1: create_person should normalize successfully
        person_tool = await server.mcp_manager.add_tool_from_mcp_server(server_name, "create_person", actor=default_user)
        assert person_tool is not None
        assert person_tool.name == "create_person"
        # Verify the schema has additionalProperties set
        assert person_tool.json_schema["parameters"]["additionalProperties"] == False
        # Verify nested $defs have additionalProperties
        if "$defs" in person_tool.json_schema["parameters"]:
            for def_name, def_schema in person_tool.json_schema["parameters"]["$defs"].items():
                if def_schema.get("type") == "object":
                    assert "additionalProperties" in def_schema, f"$defs.{def_name} missing additionalProperties after normalization"

        # Test 2: manage_tasks should normalize successfully
        tasks_tool = await server.mcp_manager.add_tool_from_mcp_server(server_name, "manage_tasks", actor=default_user)
        assert tasks_tool is not None
        assert tasks_tool.name == "manage_tasks"
        # Verify array items have explicit type
        tasks_prop = tasks_tool.json_schema["parameters"]["properties"]["tasks"]
        assert "items" in tasks_prop
        assert "type" in tasks_prop["items"], "Array items should have explicit type after normalization"

        # Test 3: search_with_filters should normalize successfully
        search_tool = await server.mcp_manager.add_tool_from_mcp_server(server_name, "search_with_filters", actor=default_user)
        assert search_tool is not None
        assert search_tool.name == "search_with_filters"

        # Verify all tools were persisted
        all_tools = await server.tool_manager.list_tools_async(
            actor=default_user, names=["create_person", "manage_tasks", "search_with_filters"]
        )

        # Filter to tools from our MCP server
        mcp_tools = [
            tool
            for tool in all_tools
            if tool.metadata_
            and MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_
            and tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX].get("server_name") == server_name
        ]

        # All 3 complex schema tools should have been normalized and persisted
        assert len(mcp_tools) == 3, f"Expected 3 normalized tools, got {len(mcp_tools)}"

        # Verify they all have the correct MCP metadata
        for tool in mcp_tools:
            assert tool.tool_type == ToolType.EXTERNAL_MCP
            assert f"mcp:{server_name}" in tool.tags

    finally:
        # Clean up
        await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)


@patch("letta.services.mcp_manager.MCPManager.get_mcp_client")
async def test_create_mcp_server_with_tools_connection_failure(mock_get_client, server, default_user):
    """Test that MCP server creation succeeds even when tool sync fails (optimistic approach)."""
    from letta.schemas.mcp import MCPServer, MCPServerType
    from letta.settings import tool_settings

    if tool_settings.mcp_read_from_config:
        return

    # Create mock client that fails to connect
    mock_client = AsyncMock()
    mock_client.connect_to_server = AsyncMock(side_effect=Exception("Connection failed"))
    mock_client.cleanup = AsyncMock()
    mock_get_client.return_value = mock_client

    # Create MCP server config
    server_name = f"test_server_fail_{uuid.uuid4().hex[:8]}"
    server_url = "https://test-fail.example.com/sse"
    mcp_server = MCPServer(server_name=server_name, server_type=MCPServerType.SSE, server_url=server_url)

    # Create server with tools - should succeed despite connection failure
    created_server = await server.mcp_manager.create_mcp_server_with_tools(mcp_server, actor=default_user)

    # Verify server was created successfully
    assert created_server.server_name == server_name
    assert created_server.server_type == MCPServerType.SSE
    assert created_server.server_url == server_url

    # Verify no tools were persisted (due to connection failure)
    # Try to get tools by the names we would have expected
    all_tools = await server.tool_manager.list_tools_async(
        actor=default_user,
        names=["tool1", "tool2", "tool3"],  # Generic names since we don't know what tools would have been listed
    )

    # Filter to see if any belong to our server (there shouldn't be any)
    persisted_tools = [
        tool
        for tool in all_tools
        if tool.metadata_
        and MCP_TOOL_TAG_NAME_PREFIX in tool.metadata_
        and tool.metadata_[MCP_TOOL_TAG_NAME_PREFIX].get("server_name") == server_name
    ]
    assert len(persisted_tools) == 0, "No tools should be persisted when connection fails"

    # Clean up
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)


async def test_get_mcp_servers_by_ids(server, default_user):
    from letta.schemas.mcp import MCPServer, MCPServerType, SSEServerConfig, StdioServerConfig
    from letta.settings import tool_settings

    if tool_settings.mcp_read_from_config:
        return

    # Create multiple MCP servers for testing
    servers_data = [
        {
            "name": "test_server_1",
            "config": StdioServerConfig(
                server_name="test_server_1", type=MCPServerType.STDIO, command="echo 'test1'", args=["arg1"], env={"ENV1": "value1"}
            ),
            "type": MCPServerType.STDIO,
        },
        {
            "name": "test_server_2",
            "config": SSEServerConfig(server_name="test_server_2", server_url="https://test2.example.com/sse"),
            "type": MCPServerType.SSE,
        },
        {
            "name": "test_server_3",
            "config": SSEServerConfig(server_name="test_server_3", server_url="https://test3.example.com/mcp"),
            "type": MCPServerType.STREAMABLE_HTTP,
        },
    ]

    created_servers = []
    for server_data in servers_data:
        if server_data["type"] == MCPServerType.STDIO:
            mcp_server = MCPServer(server_name=server_data["name"], server_type=server_data["type"], stdio_config=server_data["config"])
        else:
            mcp_server = MCPServer(
                server_name=server_data["name"], server_type=server_data["type"], server_url=server_data["config"].server_url
            )

        created = await server.mcp_manager.create_or_update_mcp_server(mcp_server, actor=default_user)
        created_servers.append(created)

    # Test fetching multiple servers by IDs
    server_ids = [s.id for s in created_servers]
    fetched_servers = await server.mcp_manager.get_mcp_servers_by_ids(server_ids, actor=default_user)

    assert len(fetched_servers) == len(created_servers)
    fetched_ids = {s.id for s in fetched_servers}
    expected_ids = {s.id for s in created_servers}
    assert fetched_ids == expected_ids

    # Test fetching subset of servers
    subset_ids = server_ids[:2]
    subset_servers = await server.mcp_manager.get_mcp_servers_by_ids(subset_ids, actor=default_user)
    assert len(subset_servers) == 2
    assert all(s.id in subset_ids for s in subset_servers)

    # Test fetching with empty list
    empty_result = await server.mcp_manager.get_mcp_servers_by_ids([], actor=default_user)
    assert empty_result == []

    # Test fetching with non-existent ID mixed with valid IDs
    mixed_ids = [server_ids[0], "non-existent-id", server_ids[1]]
    mixed_result = await server.mcp_manager.get_mcp_servers_by_ids(mixed_ids, actor=default_user)
    # Should only return the existing servers
    assert len(mixed_result) == 2
    assert all(s.id in server_ids for s in mixed_result)

    # Test that servers from different organizations are not returned
    # This would require creating another user/org, but for now we'll just verify
    # that the function respects the actor's organization
    all_servers = await server.mcp_manager.list_mcp_servers(actor=default_user)
    all_server_ids = [s.id for s in all_servers]
    bulk_fetched = await server.mcp_manager.get_mcp_servers_by_ids(all_server_ids, actor=default_user)

    # All fetched servers should belong to the same organization
    assert all(s.organization_id == default_user.organization_id for s in bulk_fetched)


# Additional MCPManager OAuth session tests
@pytest.mark.asyncio
async def test_mcp_server_deletion_cascades_oauth_sessions(server, default_organization, default_user):
    """Deleting an MCP server deletes associated OAuth sessions (same user + URL)."""

    from letta.schemas.mcp import MCPOAuthSessionCreate, MCPServer as PydanticMCPServer, MCPServerType

    test_server_url = "https://test.example.com/mcp"

    # Create orphaned OAuth sessions (no server id) for same user and URL
    created_session_ids: list[str] = []
    for i in range(3):
        session = await server.mcp_manager.create_oauth_session(
            MCPOAuthSessionCreate(
                server_url=test_server_url,
                server_name=f"test_mcp_server_{i}",
                user_id=default_user.id,
                organization_id=default_organization.id,
            ),
            actor=default_user,
        )
        created_session_ids.append(session.id)

    # Create the MCP server with the same URL
    created_server = await server.mcp_manager.create_mcp_server(
        PydanticMCPServer(
            server_name=f"test_mcp_server_{str(uuid.uuid4().hex[:8])}",  # ensure unique name
            server_type=MCPServerType.SSE,
            server_url=test_server_url,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Now delete the server via manager
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)

    # Verify all sessions are gone
    for sid in created_session_ids:
        session = await server.mcp_manager.get_oauth_session_by_id(sid, actor=default_user)
        assert session is None, f"OAuth session {sid} should be deleted"


@pytest.mark.asyncio
async def test_oauth_sessions_with_different_url_persist(server, default_organization, default_user):
    """Sessions with different URL should not be deleted when deleting the server for another URL."""

    from letta.schemas.mcp import MCPOAuthSessionCreate, MCPServer as PydanticMCPServer, MCPServerType

    server_url = "https://test.example.com/mcp"
    other_url = "https://other.example.com/mcp"

    # Create a session for other_url (should persist)
    other_session = await server.mcp_manager.create_oauth_session(
        MCPOAuthSessionCreate(
            server_url=other_url,
            server_name="standalone_oauth",
            user_id=default_user.id,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Create the MCP server at server_url
    created_server = await server.mcp_manager.create_mcp_server(
        PydanticMCPServer(
            server_name=f"test_mcp_server_{str(uuid.uuid4().hex[:8])}",
            server_type=MCPServerType.SSE,
            server_url=server_url,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Delete the server at server_url
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)

    # Verify the session at other_url still exists
    persisted = await server.mcp_manager.get_oauth_session_by_id(other_session.id, actor=default_user)
    assert persisted is not None, "OAuth session with different URL should persist"


@pytest.mark.asyncio
async def test_mcp_server_creation_links_orphaned_sessions(server, default_organization, default_user):
    """Creating a server should link any existing orphaned sessions (same user + URL)."""

    from letta.schemas.mcp import MCPOAuthSessionCreate, MCPServer as PydanticMCPServer, MCPServerType

    server_url = "https://test-atomic-create.example.com/mcp"

    # Pre-create orphaned sessions (no server_id) for same user + URL
    orphaned_ids: list[str] = []
    for i in range(3):
        session = await server.mcp_manager.create_oauth_session(
            MCPOAuthSessionCreate(
                server_url=server_url,
                server_name=f"atomic_session_{i}",
                user_id=default_user.id,
                organization_id=default_organization.id,
            ),
            actor=default_user,
        )
        orphaned_ids.append(session.id)

    # Create server
    created_server = await server.mcp_manager.create_mcp_server(
        PydanticMCPServer(
            server_name=f"test_atomic_server_{str(uuid.uuid4().hex[:8])}",
            server_type=MCPServerType.SSE,
            server_url=server_url,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Sessions should still be retrievable via manager API
    for sid in orphaned_ids:
        s = await server.mcp_manager.get_oauth_session_by_id(sid, actor=default_user)
        assert s is not None

    # Indirect verification: deleting the server removes sessions for that URL+user
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)
    for sid in orphaned_ids:
        assert await server.mcp_manager.get_oauth_session_by_id(sid, actor=default_user) is None


@pytest.mark.asyncio
async def test_mcp_server_delete_removes_all_sessions_for_url_and_user(server, default_organization, default_user):
    """Deleting a server removes both linked and orphaned sessions for same user+URL."""

    from letta.schemas.mcp import MCPOAuthSessionCreate, MCPServer as PydanticMCPServer, MCPServerType

    server_url = "https://test-atomic-cleanup.example.com/mcp"

    # Create orphaned session
    orphaned = await server.mcp_manager.create_oauth_session(
        MCPOAuthSessionCreate(
            server_url=server_url,
            server_name="orphaned",
            user_id=default_user.id,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Create server
    created_server = await server.mcp_manager.create_mcp_server(
        PydanticMCPServer(
            server_name=f"cleanup_server_{str(uuid.uuid4().hex[:8])}",
            server_type=MCPServerType.SSE,
            server_url=server_url,
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )

    # Delete server
    await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)

    # Both orphaned and any linked sessions for that URL+user should be gone
    assert await server.mcp_manager.get_oauth_session_by_id(orphaned.id, actor=default_user) is None


@pytest.mark.asyncio
async def test_mcp_server_resync_tools(server, default_user, default_organization):
    """Test that resyncing MCP server tools correctly handles added, deleted, and updated tools."""
    from unittest.mock import AsyncMock, patch

    from letta.functions.mcp_client.types import MCPTool, MCPToolHealth
    from letta.schemas.mcp import MCPServer as PydanticMCPServer, MCPServerType
    from letta.schemas.tool import ToolCreate

    # Create MCP server
    mcp_server = await server.mcp_manager.create_mcp_server(
        PydanticMCPServer(
            server_name=f"test_resync_{uuid.uuid4().hex[:8]}",
            server_type=MCPServerType.SSE,
            server_url="https://test-resync.example.com/mcp",
            organization_id=default_organization.id,
        ),
        actor=default_user,
    )
    mcp_server_id = mcp_server.id

    try:
        # Create initial persisted tools (simulating previously added tools)
        # Use sync method like in the existing mcp_tool fixture
        tool1_create = ToolCreate.from_mcp(
            mcp_server_name=mcp_server.server_name,
            mcp_tool=MCPTool(
                name="tool1",
                description="Tool 1",
                inputSchema={"type": "object", "properties": {"param1": {"type": "string"}}},
            ),
        )
        tool1 = await server.tool_manager.create_or_update_mcp_tool_async(
            tool_create=tool1_create,
            mcp_server_name=mcp_server.server_name,
            mcp_server_id=mcp_server_id,
            actor=default_user,
        )

        tool2_create = ToolCreate.from_mcp(
            mcp_server_name=mcp_server.server_name,
            mcp_tool=MCPTool(
                name="tool2",
                description="Tool 2 to be deleted",
                inputSchema={"type": "object", "properties": {"param2": {"type": "number"}}},
            ),
        )
        tool2 = await server.tool_manager.create_or_update_mcp_tool_async(
            tool_create=tool2_create,
            mcp_server_name=mcp_server.server_name,
            mcp_server_id=mcp_server_id,
            actor=default_user,
        )

        # Mock the list_mcp_server_tools to return updated tools from server
        # tool1 is updated, tool2 is deleted, tool3 is added
        updated_tools = [
            MCPTool(
                name="tool1",
                description="Tool 1 Updated",
                inputSchema={"type": "object", "properties": {"param1": {"type": "string"}, "param1b": {"type": "boolean"}}},
                health=MCPToolHealth(status="VALID", reasons=[]),
            ),
            MCPTool(
                name="tool3",
                description="Tool 3 New",
                inputSchema={"type": "object", "properties": {"param3": {"type": "array"}}},
                health=MCPToolHealth(status="VALID", reasons=[]),
            ),
        ]

        with patch.object(server.mcp_manager, "list_mcp_server_tools", new_callable=AsyncMock) as mock_list_tools:
            mock_list_tools.return_value = updated_tools

            # Run resync
            result = await server.mcp_manager.resync_mcp_server_tools(
                mcp_server_name=mcp_server.server_name,
                actor=default_user,
            )

        # Verify the resync result
        assert len(result.deleted) == 1
        assert "tool2" in result.deleted

        assert len(result.updated) == 1
        assert "tool1" in result.updated

        assert len(result.added) == 1
        assert "tool3" in result.added

        # Verify tool2 was actually deleted
        try:
            deleted_tool = await server.tool_manager.get_tool_by_id_async(tool_id=tool2.id, actor=default_user)
            assert False, "Tool2 should have been deleted"
        except Exception:
            pass  # Expected - tool should be deleted

        # Verify tool1 was updated with new schema
        updated_tool1 = await server.tool_manager.get_tool_by_id_async(tool_id=tool1.id, actor=default_user)
        assert "param1b" in updated_tool1.json_schema["parameters"]["properties"]

        # Verify tool3 was added
        tools = await server.tool_manager.list_tools_async(actor=default_user, names=["tool3"])
        assert len(tools) == 1
        assert tools[0].name == "tool3"

    finally:
        # Clean up
        await server.mcp_manager.delete_mcp_server_by_id(mcp_server_id, actor=default_user)


# ======================================================================================================================
# MCPManager Tests - Encryption
# ======================================================================================================================


@pytest.fixture
def encryption_key():
    """Fixture to ensure encryption key is set for tests."""
    original_key = settings.encryption_key
    # Set a test encryption key if not already set
    if not settings.encryption_key:
        settings.encryption_key = "test-encryption-key-32-bytes!!"
    yield settings.encryption_key
    # Restore original
    settings.encryption_key = original_key


@pytest.mark.asyncio
async def test_mcp_server_token_encryption_on_create(server, default_user, encryption_key):
    """Test that creating an MCP server encrypts the token in the database."""
    from letta.functions.mcp_client.types import MCPServerType
    from letta.orm.mcp_server import MCPServer as MCPServerModel
    from letta.schemas.mcp import MCPServer
    from letta.schemas.secret import Secret

    # Create MCP server with token
    mcp_server = MCPServer(
        server_name="test-encrypted-server",
        server_type=MCPServerType.STREAMABLE_HTTP,
        server_url="https://api.example.com/mcp",
        token="sk-test-secret-token-12345",
    )

    created_server = await server.mcp_manager.create_mcp_server(mcp_server, actor=default_user)

    try:
        # Verify server was created
        assert created_server is not None
        assert created_server.server_name == "test-encrypted-server"

        # Verify plaintext token field is NOT set (no dual-write)
        assert created_server.token is None

        # Verify token_enc is a Secret object and decrypts correctly
        assert created_server.token_enc is not None
        assert isinstance(created_server.token_enc, Secret)
        assert created_server.token_enc.get_plaintext() == "sk-test-secret-token-12345"

        # Read directly from database to verify encryption
        async with db_registry.async_session() as session:
            server_orm = await MCPServerModel.read_async(
                db_session=session,
                identifier=created_server.id,
                actor=default_user,
            )

            # Verify encrypted column is populated and different from plaintext
            assert server_orm.token_enc is not None
            assert server_orm.token_enc != "sk-test-secret-token-12345"
            # Encrypted value should be longer
            assert len(server_orm.token_enc) > len("sk-test-secret-token-12345")

    finally:
        # Clean up
        await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)


@pytest.mark.asyncio
async def test_mcp_server_token_decryption_on_read(server, default_user, encryption_key):
    """Test that reading an MCP server decrypts the token correctly."""
    from letta.functions.mcp_client.types import MCPServerType
    from letta.schemas.mcp import MCPServer
    from letta.schemas.secret import Secret

    # Create MCP server
    mcp_server = MCPServer(
        server_name="test-decrypt-server",
        server_type=MCPServerType.STREAMABLE_HTTP,
        server_url="https://api.example.com/mcp",
        token="sk-test-decrypt-token-67890",
    )

    created_server = await server.mcp_manager.create_mcp_server(mcp_server, actor=default_user)
    server_id = created_server.id

    try:
        # Read the server back
        retrieved_server = await server.mcp_manager.get_mcp_server_by_id_async(server_id, actor=default_user)

        # Verify plaintext token field is NOT set (no dual-write)
        assert retrieved_server.token is None

        # Verify the token is decrypted correctly via token_enc
        assert retrieved_server.token_enc is not None
        assert retrieved_server.token_enc.get_plaintext() == "sk-test-decrypt-token-67890"

        # Verify we can get the decrypted token through the secret getter
        token_secret = retrieved_server.get_token_secret()
        assert isinstance(token_secret, Secret)
        decrypted_token = token_secret.get_plaintext()
        assert decrypted_token == "sk-test-decrypt-token-67890"

    finally:
        # Clean up
        await server.mcp_manager.delete_mcp_server_by_id(server_id, actor=default_user)


@pytest.mark.asyncio
async def test_mcp_server_custom_headers_encryption(server, default_user, encryption_key):
    """Test that custom headers are encrypted as JSON strings."""
    from letta.functions.mcp_client.types import MCPServerType
    from letta.orm.mcp_server import MCPServer as MCPServerModel
    from letta.schemas.mcp import MCPServer
    from letta.schemas.secret import Secret

    # Create MCP server with custom headers
    custom_headers = {"Authorization": "Bearer token123", "X-API-Key": "secret-key-456"}
    mcp_server = MCPServer(
        server_name="test-headers-server",
        server_type=MCPServerType.STREAMABLE_HTTP,
        server_url="https://api.example.com/mcp",
        custom_headers=custom_headers,
    )

    created_server = await server.mcp_manager.create_mcp_server(mcp_server, actor=default_user)

    try:
        # Verify plaintext custom_headers field is NOT set (no dual-write)
        assert created_server.custom_headers is None

        # Verify custom_headers are accessible via encrypted field
        assert created_server.get_custom_headers_dict() == custom_headers

        # Verify custom_headers_enc is a Secret object (stores JSON string)
        assert created_server.custom_headers_enc is not None
        assert isinstance(created_server.custom_headers_enc, Secret)

        # Verify the getter method returns a Secret (JSON string)
        headers_secret = created_server.get_custom_headers_secret()
        assert isinstance(headers_secret, Secret)
        # Verify the Secret contains JSON string
        json_str = headers_secret.get_plaintext()
        assert json_str is not None
        import json

        assert json.loads(json_str) == custom_headers

        # Verify the convenience method returns dict directly
        headers_dict = created_server.get_custom_headers_dict()
        assert headers_dict == custom_headers

        # Read from DB to verify encryption
        async with db_registry.async_session() as session:
            server_orm = await MCPServerModel.read_async(
                db_session=session,
                identifier=created_server.id,
                actor=default_user,
            )

            # Verify encrypted column contains encrypted JSON string
            assert server_orm.custom_headers_enc is not None
            # Decrypt and verify it's valid JSON matching original headers
            decrypted_json = Secret.from_encrypted(server_orm.custom_headers_enc).get_plaintext()
            import json

            decrypted_headers = json.loads(decrypted_json)
            assert decrypted_headers == custom_headers

    finally:
        # Clean up
        await server.mcp_manager.delete_mcp_server_by_id(created_server.id, actor=default_user)


@pytest.mark.asyncio
async def test_oauth_session_tokens_encryption(server, default_user, encryption_key):
    """Test that OAuth session tokens are encrypted in the database."""
    from letta.orm.mcp_oauth import MCPOAuth as MCPOAuthModel
    from letta.schemas.mcp import MCPOAuthSessionCreate, MCPOAuthSessionUpdate
    from letta.schemas.secret import Secret

    # Create OAuth session
    session_create = MCPOAuthSessionCreate(
        server_url="https://oauth.example.com",
        server_name="test-oauth-server",
        organization_id=default_user.organization_id,
        user_id=default_user.id,
    )

    created_session = await server.mcp_manager.create_oauth_session(session_create, actor=default_user)
    session_id = created_session.id

    try:
        # Update with OAuth tokens
        session_update = MCPOAuthSessionUpdate(
            access_token="access-token-abc123",
            refresh_token="refresh-token-xyz789",
            client_secret="client-secret-def456",
            authorization_code="auth-code-ghi012",
        )

        updated_session = await server.mcp_manager.update_oauth_session(session_id, session_update, actor=default_user)

        # Verify tokens are accessible
        assert updated_session.access_token == "access-token-abc123"
        assert updated_session.refresh_token == "refresh-token-xyz789"
        assert updated_session.client_secret == "client-secret-def456"
        assert updated_session.authorization_code == "auth-code-ghi012"

        # Verify encrypted fields are Secret objects
        assert isinstance(updated_session.access_token_enc, Secret)
        assert isinstance(updated_session.refresh_token_enc, Secret)
        assert isinstance(updated_session.client_secret_enc, Secret)
        assert isinstance(updated_session.authorization_code_enc, Secret)

        # Read from DB to verify all tokens are encrypted
        async with db_registry.async_session() as session:
            oauth_orm = await MCPOAuthModel.read_async(
                db_session=session,
                identifier=session_id,
                actor=default_user,
            )

            # Verify all encrypted columns are populated and encrypted
            assert oauth_orm.access_token_enc is not None
            assert oauth_orm.refresh_token_enc is not None
            assert oauth_orm.client_secret_enc is not None
            assert oauth_orm.authorization_code_enc is not None

            # Decrypt and verify
            assert Secret.from_encrypted(oauth_orm.access_token_enc).get_plaintext() == "access-token-abc123"
            assert Secret.from_encrypted(oauth_orm.refresh_token_enc).get_plaintext() == "refresh-token-xyz789"
            assert Secret.from_encrypted(oauth_orm.client_secret_enc).get_plaintext() == "client-secret-def456"
            assert Secret.from_encrypted(oauth_orm.authorization_code_enc).get_plaintext() == "auth-code-ghi012"

    finally:
        # Clean up
        await server.mcp_manager.delete_oauth_session(session_id, actor=default_user)