feat: hooks (#330)

ideas here:

1. Support updating headers and tool args via the `before_tool_call`
hook
2. Support returning a new `CallToolResult` in the `after_tool_call`
hook (can update tool result content)
3. Hooks add support for accessing server info, tool info, and
config/runtime when available

Author: sydney-runkle

README.md

@@ -126,7 +126,7 @@ client = MultiServerMCPClient(
         },
         "weather": {
             # Make sure you start your weather server on port 8000
-            "url": "http://localhost:8000/mcp/",
+            "url": "http://localhost:8000/mcp",
             "transport": "streamable_http",
         }
     }
@@ -172,7 +172,7 @@ from mcp.client.streamable_http import streamablehttp_client
 from langgraph.prebuilt import create_react_agent
 from langchain_mcp_adapters.tools import load_mcp_tools
 
-async with streamablehttp_client("http://localhost:3000/mcp/") as (read, write, _):
+async with streamablehttp_client("http://localhost:3000/mcp") as (read, write, _):
     async with ClientSession(read, write) as session:
         # Initialize the connection
         await session.initialize()
@@ -194,7 +194,7 @@ client = MultiServerMCPClient(
     {
         "math": {
             "transport": "streamable_http",
-            "url": "http://localhost:3000/mcp/"
+            "url": "http://localhost:3000/mcp"
         },
     }
 )
@@ -255,7 +255,7 @@ client = MultiServerMCPClient(
         },
         "weather": {
             # make sure you start your weather server on port 8000
-            "url": "http://localhost:8000/mcp/",
+            "url": "http://localhost:8000/mcp",
             "transport": "streamable_http",
         }
     }
@@ -304,7 +304,7 @@ async def make_graph():
             },
             "weather": {
                 # make sure you start your weather server on port 8000
-                "url": "http://localhost:8000/mcp/",
+                "url": "http://localhost:8000/mcp",
                 "transport": "streamable_http",
             }
         }

langchain_mcp_adapters/callbacks.py

@@ -3,9 +3,16 @@
 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
+from mcp.client.session import LoggingFnT as MCPLoggingFnT
+from mcp.shared.session import ProgressFnT as MCPProgressFnT
+from mcp.types import (
+    LoggingMessageNotificationParams as MCPLoggingMessageNotificationParams,
+)
+
+# Type aliases to avoid direct MCP type dependencies
+LoggingFnT = MCPLoggingFnT
+ProgressFnT = MCPProgressFnT
+LoggingMessageNotificationParams = MCPLoggingMessageNotificationParams
 
 
 @dataclass

langchain_mcp_adapters/client.py

@@ -16,6 +16,7 @@
 from mcp import ClientSession
 
 from langchain_mcp_adapters.callbacks import CallbackContext, Callbacks
+from langchain_mcp_adapters.hooks import Hooks
 from langchain_mcp_adapters.prompts import load_mcp_prompt
 from langchain_mcp_adapters.resources import load_mcp_resources
 from langchain_mcp_adapters.sessions import (
@@ -52,13 +53,15 @@ def __init__(
         connections: dict[str, Connection] | None = None,
         *,
         callbacks: Callbacks | None = None,
+        hooks: Hooks | 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.
+            hooks: Optional hooks for before/after tool call processing.
 
         Example: basic usage (starting a new session on each tool call)
 
@@ -99,6 +102,7 @@ def __init__(
             connections if connections is not None else {}
         )
         self.callbacks = callbacks or Callbacks()
+        self.hooks = hooks
 
     @asynccontextmanager
     async def session(
@@ -163,6 +167,7 @@ async def get_tools(self, *, server_name: str | None = None) -> list[BaseTool]:
                 connection=self.connections[server_name],
                 callbacks=self.callbacks,
                 server_name=server_name,
+                hooks=self.hooks,
             )
 
         all_tools: list[BaseTool] = []
@@ -174,6 +179,7 @@ async def get_tools(self, *, server_name: str | None = None) -> list[BaseTool]:
                     connection=connection,
                     callbacks=self.callbacks,
                     server_name=name,
+                    hooks=self.hooks,
                 )
             )
             load_mcp_tool_tasks.append(load_mcp_tool_task)

langchain_mcp_adapters/hooks.py

@@ -0,0 +1,112 @@
+"""Hook interfaces and types for MCP client lifecycle management.
+
+This module provides hook interfaces for intercepting and extending
+MCP client behavior before and after tool calls.
+
+In the future, we might add more hooks for other parts of the
+request / result lifecycle, for example to support elicitation.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol
+
+from mcp.types import CallToolResult as MCPCallToolResult
+from typing_extensions import NotRequired, TypedDict
+
+if TYPE_CHECKING:
+    from langchain_core.runnables import RunnableConfig
+
+# Type aliases to avoid direct MCP type dependencies
+CallToolResult = MCPCallToolResult
+
+
+@dataclass
+class ToolHookContext:
+    """Context object passed to hooks containing state and server information."""
+
+    server_name: str
+    tool_name: str
+
+    # we'll add state eventually when we have a context manager like get_state()
+    # state: object | None = None
+    config: RunnableConfig | None = None
+    runtime: object | None = None
+
+
+class CallToolRequestSpec(TypedDict, total=False):
+    """Result of before tool call hook."""
+
+    name: NotRequired[str]
+    args: NotRequired[dict[str, Any]]
+    headers: NotRequired[dict[str, Any]]
+
+
+class BeforeToolCallHook(Protocol):
+    """Protocol for before_tool_call hook functions.
+
+    Allows modification of tool call arguments and headers before execution.
+    Return None to proceed with original request.
+    """
+
+    async def __call__(
+        self,
+        request: CallToolRequestSpec,
+        context: ToolHookContext,
+    ) -> CallToolRequestSpec | None:
+        """Execute before tool call.
+
+        Args:
+            request: The original tool call request
+            context: Hook context with server/tool info and shared state
+
+        Returns:
+            Modified CallToolRequest or None to use original request
+        """
+        ...
+
+
+class AfterToolCallHook(Protocol):
+    """Protocol for after_tool_call hook functions.
+
+    Allows modification of tool call results after execution.
+    Return None to proceed with original result processing.
+    Return CallToolResult to use the modified result.
+    """
+
+    async def __call__(
+        self,
+        result: CallToolResult,
+        context: ToolHookContext,
+    ) -> CallToolResult | None:
+        """Execute after tool call.
+
+        Args:
+            result: The original tool call result
+            context: Hook context with server/tool info and shared state
+
+        Returns:
+            - CallToolResult to use the modified result
+            - None to use original result
+        """
+        ...
+
+
+class Hooks:
+    """Container for MCP client hook functions."""
+
+    def __init__(
+        self,
+        *,
+        before_tool_call: BeforeToolCallHook | None = None,
+        after_tool_call: AfterToolCallHook | None = None,
+    ) -> None:
+        """Initialize hooks.
+
+        Args:
+            before_tool_call: Hook called before tool execution
+            after_tool_call: Hook called after tool execution
+        """
+        self.before_tool_call = before_tool_call
+        self.after_tool_call = after_tool_call

langchain_mcp_adapters/tools.py

@@ -28,8 +28,23 @@
 from pydantic import BaseModel, create_model
 
 from langchain_mcp_adapters.callbacks import CallbackContext, Callbacks, _MCPCallbacks
+from langchain_mcp_adapters.hooks import CallToolRequestSpec, Hooks, ToolHookContext
 from langchain_mcp_adapters.sessions import Connection, create_session
 
+try:
+    from langgraph.config import get_config
+    from langgraph.runtime import get_runtime
+except ImportError:
+
+    def get_config() -> dict:
+        """no-op config getter."""
+        return {}
+
+    def get_runtime() -> None:
+        """no-op runtime getter."""
+        return
+
+
 NonTextContent = ImageContent | AudioContent | ResourceLink | EmbeddedResource
 MAX_ITERATIONS = 1000
 
@@ -111,6 +126,7 @@ def convert_mcp_tool_to_langchain_tool(
     *,
     connection: Connection | None = None,
     callbacks: Callbacks | None = None,
+    hooks: Hooks | None = None,
     server_name: str | None = None,
 ) -> BaseTool:
     """Convert an MCP tool to a LangChain tool.
@@ -123,6 +139,7 @@ def convert_mcp_tool_to_langchain_tool(
         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
+        hooks: Optional hooks for before/after tool call processing
         server_name: Name of the server this tool belongs to
 
     Returns:
@@ -144,27 +161,73 @@ async def call_tool(
             else _MCPCallbacks()
         )
 
+        tool_name = tool.name
+        tool_args = arguments
+        effective_connection = connection
+
+        # try to get config and runtime if we're in a langgraph context
+        try:
+            config = get_config()
+            runtime = get_runtime()
+        except Exception:  # noqa: BLE001
+            config = {}
+            runtime = None
+
+        hook_context = ToolHookContext(
+            server_name=server_name or "unknown",
+            tool_name=tool.name,
+            config=config,
+            runtime=runtime,
+        )
+
+        if hooks and hooks.before_tool_call:
+            tool_request_spec = CallToolRequestSpec(
+                name=tool.name,
+                args=arguments,
+            )
+
+            modified_request = await hooks.before_tool_call(
+                tool_request_spec, hook_context
+            )
+            if modified_request is not None:
+                tool_name = modified_request.get("name") or tool_name
+                tool_args = modified_request.get("args") or tool_args
+
+                # If headers were modified, create a new connection with updated headers
+                modified_headers = modified_request.get("headers")
+                if modified_headers is not None and connection is not None:
+                    # Create a new connection config with updated headers
+                    updated_connection = dict(connection)
+                    if connection["transport"] in ("sse", "streamable_http"):
+                        existing_headers = connection.get("headers", {})
+                        updated_connection["headers"] = {
+                            **existing_headers,
+                            **modified_headers,
+                        }
+                        effective_connection = updated_connection
+
         # 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
-            if connection is None:
+            if effective_connection is None:
                 msg = "Either session or connection must be provided"
                 raise ValueError(msg)
 
             async with create_session(
-                connection, mcp_callbacks=mcp_callbacks
+                effective_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,
+                    tool_name,
+                    tool_args,
                     progress_callback=mcp_callbacks.progress_callback,
                 )
         else:
             call_tool_result = await session.call_tool(
-                tool.name,
-                arguments,
+                tool_name,
+                tool_args,
                 progress_callback=mcp_callbacks.progress_callback,
             )
 
@@ -177,10 +240,14 @@ async def call_tool(
             )
             raise RuntimeError(msg)
 
-        return _convert_call_tool_result(call_tool_result)
+        if hooks and hooks.after_tool_call:
+            hook_result = await hooks.after_tool_call(call_tool_result, hook_context)
+            if hook_result is not None:
+                call_tool_result = hook_result
 
-    meta = tool.meta if hasattr(tool, "meta") else None
+        return _convert_call_tool_result(call_tool_result)
 
+    meta = getattr(tool, "meta", 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
@@ -200,6 +267,7 @@ async def load_mcp_tools(
     *,
     connection: Connection | None = None,
     callbacks: Callbacks | None = None,
+    hooks: Hooks | None = None,
     server_name: str | None = None,
 ) -> list[BaseTool]:
     """Load all available MCP tools and convert them to LangChain tools.
@@ -208,6 +276,7 @@ async def load_mcp_tools(
         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.
+        hooks: Optional hooks for before/after tool call processing.
         server_name: Name of the server these tools belong to.
 
     Returns:
@@ -246,6 +315,7 @@ async def load_mcp_tools(
             tool,
             connection=connection,
             callbacks=callbacks,
+            hooks=hooks,
             server_name=server_name,
         )
         for tool in tools

pyproject.toml

@@ -55,6 +55,8 @@ ignore = [
   "C901", # Too complex
   "PLR0913", # Too many arguments in function definition
   "PLR0912", # Too many branches
+  "PLR0915", # Too many statements
+  "ERA001", # Commented out code should be removed
 ]