@@ -40,51 +40,105 @@ def __init__(self, name: str, config: dict):
4040 self .session : ClientSession | None = None
4141 self ._cleanup_lock = asyncio .Lock ()
4242
43- async def initialize (self ) -> None :
44- transport_type = self .config .get ("type" , "stdio" )
45- logger .info (f"[MCP-INIT] Server '{ self .name } ': transport='{ transport_type } '" )
43+ async def _init_stdio (self ):
44+ """Set up stdio transport and return (read, write) streams."""
45+ raw_command = self .config .get ("command" )
46+ if not isinstance (raw_command , str ) or not raw_command .strip ():
47+ raise ValueError (
48+ f"MCP server '{ self .name } ' with type 'stdio' "
49+ f"requires a non-empty 'command' string"
50+ )
51+ command = shutil .which (
52+ "npx" if raw_command == "npx" else raw_command
53+ )
54+ if not command :
55+ raise ValueError (
56+ f"Command '{ raw_command } ' not found "
57+ f"for MCP server '{ self .name } '"
58+ )
59+ logger .info (
60+ f"[MCP-INIT] Server '{ self .name } ': "
61+ f"command='{ command } ', args={ self .config .get ('args' , [])} "
62+ )
63+ params = StdioServerParameters (
64+ command = command ,
65+ args = self .config .get ("args" , []),
66+ env = (
67+ {** os .environ , ** self .config .get ("env" , {})}
68+ if self .config .get ("env" ) else None
69+ ),
70+ )
71+ transport = await self .exit_stack .enter_async_context (
72+ stdio_client (params )
73+ )
74+ return transport
4675
47- if transport_type == "stdio" :
48- raw_command = self .config .get ("command" )
49- if not isinstance (raw_command , str ) or not raw_command .strip ():
50- raise ValueError (f"MCP server '{ self .name } ' with type 'stdio' requires a non-empty 'command' string" )
51- command = shutil .which ("npx" if raw_command == "npx" else raw_command )
52- if not command :
53- raise ValueError (f"Command '{ raw_command } ' not found for MCP server '{ self .name } '" )
54- logger .info (f"[MCP-INIT] Server '{ self .name } ': command='{ command } ', args={ self .config .get ('args' , [])} " )
55- params = StdioServerParameters (
56- command = command ,
57- args = self .config .get ("args" , []),
58- env = {** os .environ , ** self .config .get ("env" , {})} if self .config .get ("env" ) else None ,
76+ async def _init_sse (self ):
77+ """Set up SSE transport and return (read, write) streams."""
78+ url = self .config .get ("url" )
79+ if not url :
80+ raise ValueError (
81+ f"MCP server '{ self .name } ' with type 'sse' "
82+ f"requires a 'url'"
5983 )
60- transport = await self .exit_stack .enter_async_context (stdio_client (params ))
61- read , write = transport
84+ headers = self .config .get ("headers" )
85+ logger .info (
86+ f"[MCP-INIT] Server '{ self .name } ': "
87+ f"url='{ _sanitize_url (url )} '"
88+ )
89+ transport = await self .exit_stack .enter_async_context (
90+ sse_client (url = url , headers = headers )
91+ )
92+ return transport
6293
63- elif transport_type == "sse" :
64- url = self .config .get ("url" )
65- if not url :
66- raise ValueError (f"MCP server '{ self .name } ' with type 'sse' requires a 'url'" )
67- headers = self .config .get ("headers" )
68- logger .info (f"[MCP-INIT] Server '{ self .name } ': url='{ _sanitize_url (url )} '" )
69- transport = await self .exit_stack .enter_async_context (sse_client (url = url , headers = headers ))
70- read , write = transport
94+ async def _init_streamable_http (self ):
95+ """Set up streamable-http transport and return (read, write)."""
96+ url = self .config .get ("url" )
97+ if not url :
98+ raise ValueError (
99+ f"MCP server '{ self .name } ' with type "
100+ f"'streamable-http' requires a 'url'"
101+ )
102+ headers = self .config .get ("headers" )
103+ logger .info (
104+ f"[MCP-INIT] Server '{ self .name } ': "
105+ f"url='{ _sanitize_url (url )} '"
106+ )
107+ transport = await self .exit_stack .enter_async_context (
108+ streamablehttp_client (url = url , headers = headers )
109+ )
110+ return transport [0 ], transport [1 ]
71111
72- elif transport_type == "streamable-http" :
73- url = self .config .get ("url" )
74- if not url :
75- raise ValueError (f"MCP server '{ self .name } ' with type 'streamable-http' requires a 'url'" )
76- headers = self .config .get ("headers" )
77- logger .info (f"[MCP-INIT] Server '{ self .name } ': url='{ _sanitize_url (url )} '" )
78- transport = await self .exit_stack .enter_async_context (streamablehttp_client (url = url , headers = headers ))
79- read , write = transport [0 ], transport [1 ]
112+ async def initialize (self ) -> None :
113+ transport_type = self .config .get ("type" , "stdio" )
114+ logger .info (
115+ f"[MCP-INIT] Server '{ self .name } ': "
116+ f"transport='{ transport_type } '"
117+ )
80118
81- else :
82- raise ValueError (f"Unknown transport type '{ transport_type } ' for MCP server '{ self .name } '. Supported types: stdio, sse, streamable-http" )
119+ init_methods = {
120+ "stdio" : self ._init_stdio ,
121+ "sse" : self ._init_sse ,
122+ "streamable-http" : self ._init_streamable_http ,
123+ }
124+ init_fn = init_methods .get (transport_type )
125+ if not init_fn :
126+ raise ValueError (
127+ f"Unknown transport type '{ transport_type } ' for "
128+ f"MCP server '{ self .name } '. "
129+ f"Supported types: stdio, sse, streamable-http"
130+ )
83131
84- session = await self .exit_stack .enter_async_context (ClientSession (read , write ))
132+ read , write = await init_fn ()
133+ session = await self .exit_stack .enter_async_context (
134+ ClientSession (read , write )
135+ )
85136 await session .initialize ()
86137 self .session = session
87- logger .info (f"[MCP-INIT] Server '{ self .name } ' session initialized successfully" )
138+ logger .info (
139+ f"[MCP-INIT] Server '{ self .name } ' "
140+ f"session initialized successfully"
141+ )
88142
89143 async def list_tools (self ) -> List [Dict [str , Any ]]:
90144 """List available tools with metadata from the server."""
0 commit comments