feat: implementing support for callbacks (#328)

Introduces a callback system for monitoring MCP server events during
tool execution.

Callbacks contain context with server name and tool name and are applied
to all servers for a given client.

**New API:**
```python
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.callbacks import Callbacks

async def on_progress(progress, total, message, context):
    print(f"[{context.server_name}:{context.tool_name}] {progress}/{total}")

client = MultiServerMCPClient(connections={...}, callbacks=Callbacks(on_progress=on_progress))
```

`Callbacks` has the potential to be extended quite broadly to support a
variety of notification hooks. We simply match signatures from `mcp` +
add a final `context` arg which is easily extensible as well.

**What's Added:**
- `CallbackContext` - provides server/tool context to callbacks
- `Callbacks` - main configuration class 
- `LoggingMessageCallback` & `ProgressCallback` - callback protocols w/
context
- Optional `callbacks` parameter on `MultiServerMCPClient`

Author: sydney-runkle

langchain_mcp_adapters/callbacks.py

@@ -0,0 +1,92 @@
+"""Types for callbacks."""
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from mcp.client.session import LoggingFnT
+from mcp.shared.session import ProgressFnT
+from mcp.types import LoggingMessageNotificationParams
+
+
+@dataclass
+class CallbackContext:
+    """LangChain MCP client callback context."""
+
+    server_name: str
+    tool_name: str | None = None
+
+
+class LoggingMessageCallback(Protocol):
+    """Light wrapper around the mcp.client.session.LoggingFnT.
+
+    Injects callback context as the last argument.
+    """
+
+    async def __call__(
+        self,
+        params: LoggingMessageNotificationParams,
+        context: CallbackContext,
+    ) -> None:
+        """Execute callback on logging message notification."""
+        ...
+
+
+class ProgressCallback(Protocol):
+    """Light wrapper around the mcp.shared.session.ProgressFnT.
+
+    Injects callback context as the last argument.
+    """
+
+    async def __call__(
+        self,
+        progress: float,
+        total: float | None,
+        message: str | None,
+        context: CallbackContext,
+    ) -> None:
+        """Execute callback on progress notification."""
+        ...
+
+
+@dataclass
+class _MCPCallbacks:
+    """Callbacks compatible with the MCP SDK. For internal use only."""
+
+    logging_callback: LoggingFnT | None = None
+    progress_callback: ProgressFnT | None = None
+
+
+@dataclass
+class Callbacks:
+    """Callbacks for the LangChain MCP client."""
+
+    on_logging_message: LoggingMessageCallback | None = None
+    on_progress: ProgressCallback | None = None
+
+    def to_mcp_format(self, *, context: CallbackContext) -> _MCPCallbacks:
+        """Convert the LangChain MCP client callbacks to MCP SDK callbacks.
+
+        Injects the LangChain CallbackContext as the last argument.
+        """
+        if (on_logging_message := self.on_logging_message) is not None:
+
+            async def mcp_logging_callback(
+                params: LoggingMessageNotificationParams,
+            ) -> None:
+                await on_logging_message(params, context)
+        else:
+            mcp_logging_callback = None
+
+        if (on_progress := self.on_progress) is not None:
+
+            async def mcp_progress_callback(
+                progress: float, total: float | None, message: str | None
+            ) -> None:
+                await on_progress(progress, total, message, context)
+        else:
+            mcp_progress_callback = None
+
+        return _MCPCallbacks(
+            logging_callback=mcp_logging_callback,
+            progress_callback=mcp_progress_callback,
+        )

langchain_mcp_adapters/client.py

@@ -15,6 +15,7 @@
 from langchain_core.tools import BaseTool
 from mcp import ClientSession
 
+from langchain_mcp_adapters.callbacks import CallbackContext, Callbacks
 from langchain_mcp_adapters.prompts import load_mcp_prompt
 from langchain_mcp_adapters.resources import load_mcp_resources
 from langchain_mcp_adapters.sessions import (
@@ -46,12 +47,18 @@ class MultiServerMCPClient:
     Loads LangChain-compatible tools, prompts and resources from MCP servers.
     """
 
-    def __init__(self, connections: dict[str, Connection] | None = None) -> None:
+    def __init__(
+        self,
+        connections: dict[str, Connection] | None = None,
+        *,
+        callbacks: Callbacks | None = None,
+    ) -> None:
         """Initialize a MultiServerMCPClient with MCP servers connections.
 
         Args:
             connections: A dictionary mapping server names to connection configurations.
                 If None, no initial connections are established.
+            callbacks: Optional callbacks for handling notifications and events.
 
         Example: basic usage (starting a new session on each tool call)
 
@@ -87,11 +94,11 @@ def __init__(self, connections: dict[str, Connection] | None = None) -> None:
         async with client.session("math") as session:
             tools = await load_mcp_tools(session)
         ```
-
         """
         self.connections: dict[str, Connection] = (
             connections if connections is not None else {}
         )
+        self.callbacks = callbacks or Callbacks()
 
     @asynccontextmanager
     async def session(
@@ -120,7 +127,13 @@ async def session(
             )
             raise ValueError(msg)
 
-        async with create_session(self.connections[server_name]) as session:
+        mcp_callbacks = self.callbacks.to_mcp_format(
+            context=CallbackContext(server_name=server_name)
+        )
+
+        async with create_session(
+            self.connections[server_name], mcp_callbacks=mcp_callbacks
+        ) as session:
             if auto_initialize:
                 await session.initialize()
             yield session
@@ -145,13 +158,23 @@ async def get_tools(self, *, server_name: str | None = None) -> list[BaseTool]:
                     f"expected one of '{list(self.connections.keys())}'"
                 )
                 raise ValueError(msg)
-            return await load_mcp_tools(None, connection=self.connections[server_name])
+            return await load_mcp_tools(
+                None,
+                connection=self.connections[server_name],
+                callbacks=self.callbacks,
+                server_name=server_name,
+            )
 
         all_tools: list[BaseTool] = []
         load_mcp_tool_tasks = []
-        for connection in self.connections.values():
+        for name, connection in self.connections.items():
             load_mcp_tool_task = asyncio.create_task(
-                load_mcp_tools(None, connection=connection)
+                load_mcp_tools(
+                    None,
+                    connection=connection,
+                    callbacks=self.callbacks,
+                    server_name=name,
+                )
             )
             load_mcp_tool_tasks.append(load_mcp_tool_task)
         tools_list = await asyncio.gather(*load_mcp_tool_tasks)
@@ -218,6 +241,7 @@ def __aexit__(
 
 
 __all__ = [
+    "Callbacks",
     "McpHttpClientFactory",
     "MultiServerMCPClient",
     "SSEConnection",

langchain_mcp_adapters/sessions.py

@@ -22,6 +22,8 @@
 
     import httpx
 
+    from langchain_mcp_adapters.callbacks import _MCPCallbacks
+
 EncodingErrorHandler = Literal["strict", "ignore", "replace"]
 
 DEFAULT_ENCODING = "utf-8"
@@ -187,7 +189,7 @@ class WebsocketConnection(TypedDict):
 
 
 @asynccontextmanager
-async def _create_stdio_session(  # noqa: PLR0913
+async def _create_stdio_session(
     *,
     command: str,
     args: list[str],
@@ -233,7 +235,7 @@ async def _create_stdio_session(  # noqa: PLR0913
 
 
 @asynccontextmanager
-async def _create_sse_session(  # noqa: PLR0913
+async def _create_sse_session(
     *,
     url: str,
     headers: dict[str, Any] | None = None,
@@ -273,7 +275,7 @@ async def _create_sse_session(  # noqa: PLR0913
 
 
 @asynccontextmanager
-async def _create_streamable_http_session(  # noqa: PLR0913
+async def _create_streamable_http_session(
     *,
     url: str,
     headers: dict[str, Any] | None = None,
@@ -356,11 +358,14 @@ async def _create_websocket_session(
 
 
 @asynccontextmanager
-async def create_session(connection: Connection) -> AsyncIterator[ClientSession]:  # noqa: C901
+async def create_session(
+    connection: Connection, *, mcp_callbacks: _MCPCallbacks | None = None
+) -> AsyncIterator[ClientSession]:
     """Create a new session to an MCP server.
 
     Args:
         connection: Connection config to use to connect to the server
+        mcp_callbacks: mcp sdk compatible callbacks to use for the ClientSession
 
     Raises:
         ValueError: If transport is not recognized
@@ -381,6 +386,16 @@ async def create_session(connection: Connection) -> AsyncIterator[ClientSession]
     transport = connection["transport"]
     params = {k: v for k, v in connection.items() if k != "transport"}
 
+    if mcp_callbacks is not None:
+        params["session_kwargs"] = params.get("session_kwargs", {})
+        # right now the only callback supported on the ClientSession
+        # is the logging callback, but long term we'll also want to
+        # support sampling, elicitation, list roots, etc.
+        if mcp_callbacks.logging_callback is not None:
+            params["session_kwargs"]["logging_callback"] = (
+                mcp_callbacks.logging_callback
+            )
+
     if transport == "sse":
         if "url" not in params:
             msg = "'url' parameter is required for SSE connection"

langchain_mcp_adapters/tools.py

@@ -16,13 +16,21 @@
 from mcp import ClientSession
 from mcp.server.fastmcp.tools import Tool as FastMCPTool
 from mcp.server.fastmcp.utilities.func_metadata import ArgModelBase, FuncMetadata
-from mcp.types import CallToolResult, EmbeddedResource, ImageContent, TextContent
+from mcp.types import (
+    AudioContent,
+    CallToolResult,
+    EmbeddedResource,
+    ImageContent,
+    ResourceLink,
+    TextContent,
+)
 from mcp.types import Tool as MCPTool
 from pydantic import BaseModel, create_model
 
+from langchain_mcp_adapters.callbacks import CallbackContext, Callbacks, _MCPCallbacks
 from langchain_mcp_adapters.sessions import Connection, create_session
 
-NonTextContent = ImageContent | EmbeddedResource
+NonTextContent = ImageContent | AudioContent | ResourceLink | EmbeddedResource
 MAX_ITERATIONS = 1000
 
 
@@ -102,6 +110,8 @@ def convert_mcp_tool_to_langchain_tool(
     tool: MCPTool,
     *,
     connection: Connection | None = None,
+    callbacks: Callbacks | None = None,
+    server_name: str | None = None,
 ) -> BaseTool:
     """Convert an MCP tool to a LangChain tool.
 
@@ -112,6 +122,8 @@ def convert_mcp_tool_to_langchain_tool(
         tool: MCP tool to convert
         connection: Optional connection config to use to create a new session
                     if a `session` is not provided
+        callbacks: Optional callbacks for handling notifications and events
+        server_name: Name of the server this tool belongs to
 
     Returns:
         a LangChain tool
@@ -124,17 +136,37 @@ def convert_mcp_tool_to_langchain_tool(
     async def call_tool(
         **arguments: dict[str, Any],
     ) -> tuple[str | list[str], list[NonTextContent] | None]:
+        mcp_callbacks = (
+            callbacks.to_mcp_format(
+                context=CallbackContext(server_name=server_name, tool_name=tool.name)
+            )
+            if callbacks is not None
+            else _MCPCallbacks()
+        )
+
+        # Execute the tool call
         call_tool_result = None
         if session is None:
             # If a session is not provided, we will create one on the fly
-            async with create_session(connection) as tool_session:
+            if connection is None:
+                msg = "Either session or connection must be provided"
+                raise ValueError(msg)
+
+            async with create_session(
+                connection, mcp_callbacks=mcp_callbacks
+            ) as tool_session:
                 await tool_session.initialize()
                 call_tool_result = await cast("ClientSession", tool_session).call_tool(
                     tool.name,
                     arguments,
+                    progress_callback=mcp_callbacks.progress_callback,
                 )
         else:
-            call_tool_result = await session.call_tool(tool.name, arguments)
+            call_tool_result = await session.call_tool(
+                tool.name,
+                arguments,
+                progress_callback=mcp_callbacks.progress_callback,
+            )
 
         if call_tool_result is None:
             msg = (
@@ -148,6 +180,7 @@ async def call_tool(
         return _convert_call_tool_result(call_tool_result)
 
     meta = tool.meta if hasattr(tool, "meta") else None
+
     base = tool.annotations.model_dump() if tool.annotations is not None else {}
     meta = {"_meta": meta} if meta is not None else {}
     metadata = {**base, **meta} or None
@@ -166,12 +199,16 @@ async def load_mcp_tools(
     session: ClientSession | None,
     *,
     connection: Connection | None = None,
+    callbacks: Callbacks | None = None,
+    server_name: str | None = None,
 ) -> list[BaseTool]:
     """Load all available MCP tools and convert them to LangChain tools.
 
     Args:
         session: The MCP client session. If None, connection must be provided.
         connection: Connection config to create a new session if session is None.
+        callbacks: Optional callbacks for handling notifications and events.
+        server_name: Name of the server these tools belong to.
 
     Returns:
         List of LangChain tools. Tool annotations are returned as part
@@ -184,16 +221,33 @@ async def load_mcp_tools(
         msg = "Either a session or a connection config must be provided"
         raise ValueError(msg)
 
+    mcp_callbacks = (
+        callbacks.to_mcp_format(context=CallbackContext(server_name=server_name))
+        if callbacks is not None
+        else _MCPCallbacks()
+    )
+
     if session is None:
         # If a session is not provided, we will create one on the fly
-        async with create_session(connection) as tool_session:
+        if connection is None:
+            msg = "Either session or connection must be provided"
+            raise ValueError(msg)
+        async with create_session(
+            connection, mcp_callbacks=mcp_callbacks
+        ) as tool_session:
             await tool_session.initialize()
             tools = await _list_all_tools(tool_session)
     else:
         tools = await _list_all_tools(session)
 
     return [
-        convert_mcp_tool_to_langchain_tool(session, tool, connection=connection)
+        convert_mcp_tool_to_langchain_tool(
+            session,
+            tool,
+            connection=connection,
+            callbacks=callbacks,
+            server_name=server_name,
+        )
         for tool in tools
     ]