feat: add ability to return Command to interceptors and add end-to-end test (#372)

Allow returning Command to interceptors and add end-to-end test

Author: eyurtsev

langchain_mcp_adapters/interceptors.py

@@ -16,10 +16,24 @@
 from mcp.types import CallToolResult
 from typing_extensions import NotRequired, TypedDict, Unpack
 
+try:
+    # langgraph installed
+    import langgraph
+
+    LANGGRAPH_PRESENT = True
+except ImportError:
+    LANGGRAPH_PRESENT = False
+
+
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
-MCPToolCallResult = CallToolResult | ToolMessage
+if LANGGRAPH_PRESENT:
+    from langgraph.types import Command
+
+    MCPToolCallResult = CallToolResult | ToolMessage | Command
+else:
+    MCPToolCallResult = CallToolResult | ToolMessage
 
 
 class _MCPToolCallRequestOverrides(TypedDict, total=False):

langchain_mcp_adapters/tools.py

@@ -36,24 +36,41 @@
 )
 from langchain_mcp_adapters.sessions import Connection, create_session
 
+try:
+    # langgraph installed
+    import langgraph
+    from langgraph.types import Command
+
+    LANGGRAPH_PRESENT = True
+except ImportError:
+    LANGGRAPH_PRESENT = False
+
 NonTextContent = ImageContent | AudioContent | ResourceLink | EmbeddedResource
+
+# Conditional type based on langgraph availability
+if LANGGRAPH_PRESENT:
+    ConvertedToolResult = str | list[str] | ToolMessage | Command
+else:
+    ConvertedToolResult = str | list[str] | ToolMessage
+
 MAX_ITERATIONS = 1000
 
 
 def _convert_call_tool_result(
     call_tool_result: MCPToolCallResult,
-) -> tuple[str | list[str] | ToolMessage, list[NonTextContent] | None]:
+) -> tuple[ConvertedToolResult, list[NonTextContent] | None]:
     """Convert MCP MCPToolCallResult to LangChain tool result format.
 
     Args:
         call_tool_result: The result from calling an MCP tool. Can be either
-            a CallToolResult (MCP format) or a ToolMessage (LangChain format).
+            a CallToolResult (MCP format), a ToolMessage (LangChain format),
+            or a Command (LangGraph format, if langgraph is installed).
 
     Returns:
-        A tuple containing the text content (which may be a ToolMessage) and any
-        non-text content. When a ToolMessage is returned by an interceptor, it's
-        placed in the first position of the tuple as the content, with None as
-        the artifact.
+        A tuple containing the text content (which may be a ToolMessage or Command)
+        and any non-text content. When a ToolMessage or Command is returned by an
+        interceptor, it's placed in the first position of the tuple as the content,
+        with None as the artifact.
 
     Raises:
         ToolException: If the tool call resulted in an error.
@@ -63,6 +80,10 @@ def _convert_call_tool_result(
     if isinstance(call_tool_result, ToolMessage):
         return call_tool_result, None
 
+    # If the interceptor returned a Command (LangGraph), return it directly
+    if LANGGRAPH_PRESENT and isinstance(call_tool_result, Command):
+        return call_tool_result, None
+
     # Otherwise, convert from CallToolResult
     text_contents: list[TextContent] = []
     non_text_contents = []
@@ -188,7 +209,7 @@ def convert_mcp_tool_to_langchain_tool(
     async def call_tool(
         runtime: Any = None,  # noqa: ANN401
         **arguments: dict[str, Any],
-    ) -> tuple[str | list[str] | ToolMessage, list[NonTextContent] | None]:
+    ) -> tuple[ConvertedToolResult, list[NonTextContent] | None]:
         """Execute tool call with interceptor chain and return formatted result.
 
         Args:
@@ -197,7 +218,8 @@ async def call_tool(
 
         Returns:
             A tuple of (text_content, non_text_content), where text_content may be
-            a ToolMessage if an interceptor returned one directly.
+            a ToolMessage or Command (if langgraph is installed) if an interceptor
+            returned one directly.
         """
         mcp_callbacks = (
             callbacks.to_mcp_format(

pyproject.toml

@@ -19,6 +19,9 @@ dependencies = [
 ]
 
 [dependency-groups]
+integration = [
+    "langchain>=1.0.8",
+]
 test = [
     "pytest>=8.0.0",
     "ruff>=0.9.4",

tests/test_tools.py

@@ -1,10 +1,14 @@
-from typing import Annotated
+import typing
+from collections.abc import Callable, Sequence
+from typing import Annotated, Any
 from unittest.mock import AsyncMock, MagicMock
 
 import httpx
 import pytest
 from langchain_core.callbacks import CallbackManagerForToolRun
-from langchain_core.messages import ToolMessage
+from langchain_core.language_models import LanguageModelInput
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langchain_core.runnables import Runnable
 from langchain_core.tools import BaseTool, InjectedToolArg, ToolException, tool
 from mcp.server import FastMCP
 from mcp.types import (
@@ -19,6 +23,7 @@
 from pydantic import BaseModel
 
 from langchain_mcp_adapters.client import MultiServerMCPClient
+from langchain_mcp_adapters.interceptors import MCPToolCallRequest
 from langchain_mcp_adapters.tools import (
     _convert_call_tool_result,
     convert_mcp_tool_to_langchain_tool,
@@ -526,3 +531,141 @@ async def test_convert_mcp_tool_metadata_variants():
         "openWorldHint": None,
         "_meta": {"flag": True},
     }
+
+
+def _create_increment_server():
+    server = FastMCP(port=8183)
+
+    @server.tool()
+    def increment(value: int) -> str:
+        """Increment a counter"""
+        return f"Incremented to {value + 1}"
+
+    return server
+
+
+try:
+    import langchain
+
+    LANGCHAIN_INSTALLED = True
+except ImportError:
+    LANGCHAIN_INSTALLED = False
+
+from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
+
+
+class FixedGenericFakeChatModel(GenericFakeChatModel):
+    def bind_tools(
+        self,
+        tools: Sequence[
+            typing.Dict[str, Any] | type | Callable | BaseTool  # noqa: UP006
+        ],
+        *,
+        tool_choice: str | None = None,
+        **kwargs: Any,
+    ) -> Runnable[LanguageModelInput, AIMessage]:
+        """Override bind-tools."""
+        return self
+
+
+@pytest.mark.skipif(not LANGCHAIN_INSTALLED, reason="langchain not installed")
+async def test_mcp_tools_with_agent_and_command_interceptor(socket_enabled) -> None:
+    """Test Command objects from interceptors work end-to-end with create_agent.
+
+    This test verifies that:
+    1. MCP tools can be used with create_agent
+    2. Interceptors can return Command objects to short-circuit execution
+    3. Commands can update custom agent state
+    """
+    from langchain.agents import AgentState, create_agent
+    from langchain.tools import ToolRuntime
+    from langgraph.checkpoint.memory import MemorySaver
+    from langgraph.types import Command
+
+    from langchain_mcp_adapters.interceptors import MCPToolCallResult
+
+    # Interceptor that returns Command to update state
+    async def counter_interceptor(
+        request: MCPToolCallRequest,
+        handler: Callable[[MCPToolCallRequest], typing.Awaitable[MCPToolCallResult]],
+    ) -> Command:
+        # Instead of calling the tool, return a Command that updates state
+        tool_runtime: ToolRuntime = request.runtime
+        assert tool_runtime.tool_call_id == "call_1"
+        return Command(
+            update={
+                "counter": 42,
+                "messages": [
+                    ToolMessage(
+                        content="Counter updated!",
+                        tool_call_id=tool_runtime.tool_call_id,
+                    ),
+                    AIMessage(content="hello"),
+                ],
+            },
+            goto="__end__",
+        )
+
+    with run_streamable_http(_create_increment_server, 8183):
+        # Initialize client and connect to server
+        client = MultiServerMCPClient(
+            {
+                "increment": {
+                    "url": "http://localhost:8183/mcp",
+                    "transport": "streamable_http",
+                }
+            },
+            tool_interceptors=[counter_interceptor],
+        )
+
+        # Get tools from the server
+        tools = await client.get_tools(server_name="increment")
+        assert len(tools) == 1
+        original_tool = tools[0]
+        assert original_tool.name == "increment"
+
+        # Custom state schema with counter field
+        class CustomAgentState(AgentState):
+            counter: typing.NotRequired[int]
+
+        model = FixedGenericFakeChatModel(
+            messages=iter(
+                [
+                    AIMessage(
+                        content="",
+                        tool_calls=[
+                            {
+                                "name": "increment",
+                                "args": {"value": 1},
+                                "id": "call_1",
+                                "type": "tool_call",
+                            }
+                        ],
+                    ),
+                    AIMessage(
+                        content="The counter has been incremented.",
+                    ),
+                ]
+            )
+        )
+        # Create agent with custom state
+        agent = create_agent(
+            model,
+            tools,
+            state_schema=CustomAgentState,
+            checkpointer=MemorySaver(),
+        )
+
+        # Run agent
+        result = await agent.ainvoke(
+            {"messages": [HumanMessage(content="Increment the counter")], "counter": 0},
+            {"configurable": {"thread_id": "test_1"}},
+        )
+
+        # Verify Command updated the state
+        assert result["counter"] == 42
+        # Verify the Command's message was added
+        assert any(
+            isinstance(msg, ToolMessage) and msg.content == "Counter updated!"
+            for msg in result["messages"]
+        )