refactor: share mcp session opener

src/wf_mcp/runtime/factory.py

@@ -4,13 +4,12 @@ import asyncio
 from contextlib import AsyncExitStack
 from dataclasses import dataclass, field
 
-import httpx
 from mcp.client.session import ClientSession
-from mcp.client.stdio import StdioServerParameters, stdio_client
-from mcp.client.streamable_http import streamable_http_client
 from mcp.types import CallToolResult
 
-from wf_sources_mcp.auth import AuthRecord, mcp_auth_env, mcp_auth_headers
+from wf_sources_mcp.auth import AuthRecord
+from wf_sources_mcp.client import open_mcp_session
+from wf_sources_mcp.connections import mcp_source_connection_from_connection_config
 
 from ..models import ConnectionConfig
 from .session import PersistentMcpSession
@@ -46,48 +45,11 @@ class PersistentSessionFactory:
         connection: ConnectionConfig,
         auth: AuthRecord | None,
     ) -> ClientSession:
-        transport = connection.metadata.get("transport", "stdio")
-        if transport == "stdio":
-            env = connection.metadata.get("env")
-            auth_env = mcp_auth_env(auth)
-            if auth_env:
-                env = {**(env or {}), **auth_env}
-            params = StdioServerParameters(
-                command=connection.metadata["command"],
-                args=list(connection.metadata.get("args", [])),
-                env=env,
-                cwd=connection.metadata.get("cwd"),
-            )
-            read_stream, write_stream = await stack.enter_async_context(
-                stdio_client(params)
-            )
-            session = await stack.enter_async_context(
-                ClientSession(read_stream, write_stream)
-            )
-            await session.initialize()
-            return session
-
-        if transport == "streamable_http":
-            http_client = await stack.enter_async_context(
-                httpx.AsyncClient(headers=mcp_auth_headers(auth) or None)
-            )
-            (
-                read_stream,
-                write_stream,
-                _get_session_id,
-            ) = await stack.enter_async_context(
-                streamable_http_client(
-                    connection.metadata["url"],
-                    http_client=http_client,
-                )
-            )
-            session = await stack.enter_async_context(
-                ClientSession(read_stream, write_stream)
-            )
-            await session.initialize()
-            return session
-
-        raise ValueError(f"unsupported MCP transport {transport!r}")
+        source_connection = mcp_source_connection_from_connection_config(connection)
+        session = await stack.enter_async_context(
+            open_mcp_session(source_connection, auth)
+        )
+        return session
 
 
 @dataclass(slots=True)

src/wf_mcp/sdk/adapter.py

@@ -3,11 +3,7 @@ from __future__ import annotations
 from contextlib import asynccontextmanager
 from typing import Any
 
-import httpx
 from mcp import ClientResult
-from mcp.client.session import ClientSession
-from mcp.client.stdio import StdioServerParameters, stdio_client
-from mcp.client.streamable_http import streamable_http_client
 from mcp.types import (
     ClientNotification,
     ClientRequest,
@@ -17,8 +13,9 @@ from mcp.types import (
 )
 from pydantic import AnyUrl
 
-from wf_sources_mcp.auth import AuthRecord, mcp_auth_env, mcp_auth_headers
+from wf_sources_mcp.auth import AuthRecord
 from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
+from wf_sources_mcp.client import open_mcp_session
 from wf_sources_mcp.connections import McpSourceConnection
 from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult
 from wf_sources_mcp.sdk.converters import (
@@ -27,7 +24,6 @@ from wf_sources_mcp.sdk.converters import (
     tool_result_to_call_result,
     tool_to_discovered,
 )
-from wf_sources_mcp.transports import HttpSourceTransport, StdioSourceTransport
 
 
 class McpSdkAdapter(BackendAdapter):
@@ -37,39 +33,8 @@ class McpSdkAdapter(BackendAdapter):
         connection: McpSourceConnection,
         auth: AuthRecord | None,
     ):
-        transport = connection.transport
-        if transport is None:
-            raise ValueError(f"connection {connection.id!r} requires metadata.transport")
-        if isinstance(transport, StdioSourceTransport):
-            auth_env = mcp_auth_env(auth)
-            env = dict(transport.env)
-            if auth_env:
-                env = {**env, **auth_env}
-            params = StdioServerParameters(
-                command=transport.command,
-                args=list(transport.args),
-                env=env,
-            )
-            async with stdio_client(params) as (read_stream, write_stream):
-                async with ClientSession(read_stream, write_stream) as session:
-                    await session.initialize()
-                    yield session
-            return
-
-        if isinstance(transport, HttpSourceTransport):
-            headers = mcp_auth_headers(auth)
-            http_client = httpx.AsyncClient(headers=headers or None)
-            async with http_client:
-                async with streamable_http_client(
-                    str(transport.url),
-                    http_client=http_client,
-                ) as (read_stream, write_stream, _get_session_id):
-                    async with ClientSession(read_stream, write_stream) as session:
-                        await session.initialize()
-                        yield session
-            return
-
-        raise ValueError(f"unsupported MCP transport {type(transport).__name__}")
+        async with open_mcp_session(connection, auth) as session:
+            yield session
 
     async def list_tools(
         self,

src/wf_sources_mcp/client/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .transport import open_mcp_session
+
+__all__ = ["open_mcp_session"]

src/wf_sources_mcp/client/transport.py

@@ -0,0 +1,68 @@
+"""Shared MCP session opener for one-shot and persistent runtimes."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import httpx
+from mcp.client.session import ClientSession
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from mcp.client.streamable_http import streamable_http_client
+
+from wf_sources_mcp.auth import AuthRecord, mcp_auth_env, mcp_auth_headers
+from wf_sources_mcp.connections import McpSourceConnection
+from wf_sources_mcp.transports import HttpSourceTransport, StdioSourceTransport
+
+
+@asynccontextmanager
+async def open_mcp_session(
+    connection: McpSourceConnection,
+    auth: AuthRecord | None,
+) -> AsyncIterator[ClientSession]:
+    """Open and initialize an MCP client session for the given connection.
+
+    For stdio transports, merges transport env with auth env (auth wins on
+    duplicate keys) and passes command, args, env, and cwd to
+    StdioServerParameters.
+
+    For HTTP transports, creates an httpx.AsyncClient with auth headers and
+    enters streamable_http_client.
+
+    Yields an initialized ClientSession. Caller owns the session lifetime.
+    """
+    transport = connection.transport
+    if transport is None:
+        raise ValueError(f"connection {connection.id!r} requires metadata.transport")
+
+    if isinstance(transport, StdioSourceTransport):
+        auth_env = mcp_auth_env(auth)
+        env = dict(transport.env)
+        if auth_env:
+            env = {**env, **auth_env}
+        params = StdioServerParameters(
+            command=transport.command,
+            args=list(transport.args),
+            env=env,
+            cwd=transport.cwd,
+        )
+        async with stdio_client(params) as (read_stream, write_stream):
+            async with ClientSession(read_stream, write_stream) as session:
+                await session.initialize()
+                yield session
+        return
+
+    if isinstance(transport, HttpSourceTransport):
+        headers = mcp_auth_headers(auth)
+        http_client = httpx.AsyncClient(headers=headers or None)
+        async with http_client:
+            async with streamable_http_client(
+                str(transport.url),
+                http_client=http_client,
+            ) as (read_stream, write_stream, _get_session_id):
+                async with ClientSession(read_stream, write_stream) as session:
+                    await session.initialize()
+                    yield session
+        return
+
+    raise ValueError(f"unsupported MCP transport {transport.kind!r}")