fix imports

asheshvidyut · asheshvidyut · commit 4377d41eb524 · 2025-11-14T06:27:36.000Z
diff --git a/examples/snippets/clients/display_utilities.py b/examples/snippets/clients/display_utilities.py
@@ -6,9 +6,8 @@
 import asyncio
 import os
 
-from mcp import ClientSession, StdioServerParameters
+from mcp import ClientSession, ClientTransportSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
-from mcp.client.transport_session import ClientTransportSession
 from mcp.shared.metadata_utils import get_display_name
 
 # Create server parameters for stdio connection
diff --git a/examples/snippets/clients/stdio_client.py b/examples/snippets/clients/stdio_client.py
@@ -8,8 +8,7 @@
 
 from pydantic import AnyUrl
 
-from mcp import ClientSession, StdioServerParameters, types
-from mcp.client.session import ClientTransportSession
+from mcp import ClientSession, StdioServerParameters, types, ClientTransportSession
 from mcp.client.stdio import stdio_client
 from mcp.shared.context import RequestContext
 
diff --git a/examples/snippets/servers/elicitation.py b/examples/snippets/servers/elicitation.py
@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerTransportSession
+from mcp.server.transport_session import ServerTransportSession
 
 mcp = FastMCP(name="Elicitation Example")
 
diff --git a/examples/snippets/servers/lifespan_example.py b/examples/snippets/servers/lifespan_example.py
@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerTransportSession
+from mcp.server.transport_session import ServerTransportSession
 
 
 # Mock database class for example
diff --git a/examples/snippets/servers/notifications.py b/examples/snippets/servers/notifications.py
@@ -1,5 +1,5 @@
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerTransportSession
+from mcp.server.transport_session import ServerTransportSession
 
 mcp = FastMCP(name="Notifications Example")
 
diff --git a/examples/snippets/servers/tool_progress.py b/examples/snippets/servers/tool_progress.py
@@ -1,5 +1,5 @@
 from mcp.server.fastmcp import Context, FastMCP
-from mcp.server.session import ServerTransportSession
+from mcp.server.transport_session import ServerTransportSession
 
 mcp = FastMCP(name="Progress Example")
 
diff --git a/src/mcp/server/elicitation.py b/src/mcp/server/elicitation.py
@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 
-from mcp.server.session import ServerTransportSession
+from mcp.server.transport_session import ServerTransportSession
 from mcp.types import RequestId
 
 ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
@@ -54,7 +54,8 @@
 from mcp.server.lowlevel.server import LifespanResultT
 from mcp.server.lowlevel.server import Server as MCPServer
 from mcp.server.lowlevel.server import lifespan as default_lifespan
-from mcp.server.session import ServerSessionT, ServerTransportSession
+from mcp.server.session import ServerSessionT
+from mcp.server.transport_session import ServerTransportSession
 from mcp.server.sse import SseServerTransport
 from mcp.server.stdio import stdio_server
 from mcp.server.streamable_http import EventStore
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -85,7 +85,8 @@ async def main():
 from mcp.server.lowlevel.func_inspection import create_call_wrapper
 from mcp.server.lowlevel.helper_types import ReadResourceContents
 from mcp.server.models import InitializationOptions
-from mcp.server.session import ServerSession, ServerTransportSession
+from mcp.server.session import ServerSession
+from mcp.server.transport_session import ServerTransportSession
 from mcp.shared.context import RequestContext
 from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py
@@ -1,17 +1,16 @@
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Generic
+from typing import Any, Generic
 
 from typing_extensions import TypeVar
 
 from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParams
 
-if TYPE_CHECKING:
-    from mcp.client.session import ClientTransportSession
-    from mcp.server.session import ServerTransportSession
+from mcp.client.transport_session import ClientTransportSession
+from mcp.server.transport_session import ServerTransportSession
 
 SessionT = TypeVar(
-    "SessionT", bound=BaseSession[Any, Any, Any, Any, Any] | "ClientTransportSession" | "ServerTransportSession"
+    "SessionT", bound=BaseSession[Any, Any, Any, Any, Any] | ClientTransportSession | ServerTransportSession
 )
 LifespanContextT = TypeVar("LifespanContextT")
 RequestT = TypeVar("RequestT", default=Any)
diff --git a/src/mcp/shared/memory.py b/src/mcp/shared/memory.py
@@ -15,13 +15,13 @@
 import mcp.types as types
 from mcp.client.session import (
     ClientSession,
-    ClientTransportSession,
     ElicitationFnT,
     ListRootsFnT,
     LoggingFnT,
     MessageHandlerFnT,
     SamplingFnT,
 )
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server import Server
 from mcp.server.fastmcp import FastMCP
 from mcp.shared.message import SessionMessage
diff --git a/tests/client/test_list_roots_callback.py b/tests/client/test_list_roots_callback.py
@@ -1,7 +1,7 @@
 import pytest
 from pydantic import FileUrl
 
-from mcp.client.session import ClientTransportSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server.fastmcp.server import Context
 from mcp.server.session import ServerSession
 from mcp.shared.context import RequestContext
diff --git a/tests/client/test_sampling_callback.py b/tests/client/test_sampling_callback.py
@@ -1,6 +1,6 @@
 import pytest
 
-from mcp.client.session import ClientTransportSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server.session import ServerSession
 from mcp.shared.context import RequestContext
 from mcp.shared.memory import (
diff --git a/tests/client/test_session.py b/tests/client/test_session.py
@@ -4,7 +4,8 @@
 import pytest
 
 import mcp.types as types
-from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession, ClientTransportSession
+from mcp.client.session import DEFAULT_CLIENT_INFO, ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
diff --git a/tests/server/fastmcp/test_elicitation.py b/tests/server/fastmcp/test_elicitation.py
@@ -7,7 +7,8 @@
 import pytest
 from pydantic import BaseModel, Field
 
-from mcp.client.session import ClientTransportSession, ElicitationFnT
+from mcp.client.session import ElicitationFnT
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server.fastmcp import Context, FastMCP
 from mcp.server.session import ServerSession
 from mcp.shared.context import RequestContext
diff --git a/tests/server/fastmcp/test_integration.py b/tests/server/fastmcp/test_integration.py
@@ -32,7 +32,8 @@
     structured_output,
     tool_progress,
 )
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import GetSessionIdCallback, streamablehttp_client
 from mcp.shared.context import RequestContext
diff --git a/tests/shared/test_memory.py b/tests/shared/test_memory.py
@@ -2,7 +2,8 @@
 from pydantic import AnyUrl
 from typing_extensions import AsyncGenerator
 
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server import Server
 from mcp.shared.memory import create_connected_server_and_client_session
 from mcp.types import EmptyResult, Resource
diff --git a/tests/shared/test_session.py b/tests/shared/test_session.py
@@ -5,7 +5,8 @@
 import pytest
 
 import mcp.types as types
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.server.lowlevel.server import Server
 from mcp.shared.exceptions import McpError
 from mcp.shared.memory import create_client_server_memory_streams, create_connected_server_and_client_session
diff --git a/tests/shared/test_sse.py b/tests/shared/test_sse.py
@@ -17,7 +17,8 @@
 from starlette.routing import Mount, Route
 
 import mcp.types as types
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.client.sse import sse_client
 from mcp.server import Server
 from mcp.server.sse import SseServerTransport
diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py
@@ -21,7 +21,8 @@
 from starlette.routing import Mount
 
 import mcp.types as types
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.client.streamable_http import streamablehttp_client
 from mcp.server import Server
 from mcp.server.session import ServerSession
diff --git a/tests/shared/test_ws.py b/tests/shared/test_ws.py
@@ -12,7 +12,8 @@
 from starlette.routing import WebSocketRoute
 from starlette.websockets import WebSocket
 
-from mcp.client.session import ClientSession, ClientTransportSession
+from mcp.client.session import ClientSession
+from mcp.client.transport_session import ClientTransportSession
 from mcp.client.websocket import websocket_client
 from mcp.server import Server
 from mcp.server.websocket import websocket_server
