diff --git a/libs/langchain_v1/examples/mcp_stateful_browser_example.py b/libs/langchain_v1/examples/mcp_stateful_browser_example.py index 0cd4211ec91..f8044bbb577 100644 --- a/libs/langchain_v1/examples/mcp_stateful_browser_example.py +++ b/libs/langchain_v1/examples/mcp_stateful_browser_example.py @@ -12,33 +12,31 @@ causing subsequent tool calls to fail with "Ref not found" errors. """ import asyncio -from typing import Any -from langchain.agents import create_agent from langchain_core.messages import HumanMessage from langgraph.checkpoint.memory import MemorySaver +from langchain.agents import create_agent + # Note: langchain-mcp-adapters must be installed separately # pip install langchain-mcp-adapters try: from langchain_mcp_adapters.client import MultiServerMCPClient from langchain_mcp_adapters.tools import load_mcp_tools except ImportError: - raise ImportError( - "Please install langchain-mcp-adapters: pip install langchain-mcp-adapters" - ) + raise ImportError("Please install langchain-mcp-adapters: pip install langchain-mcp-adapters") async def stateless_example_problematic(): """Problematic example showing stateless MCP usage that causes browser session termination. - + This demonstrates the INCORRECT way that leads to the reported issue. Each tool call creates a new browser session, causing state loss between calls. """ print("\n" + "=" * 80) print("PROBLEMATIC EXAMPLE: Stateless MCP Usage (Browser closes after each tool call)") print("=" * 80) - + # Initialize MCP client with Playwright server client = MultiServerMCPClient( { @@ -53,18 +51,18 @@ async def stateless_example_problematic(): } } ) - + # PROBLEM: Using get_tools() creates stateless tools # Each tool invocation will create a new session tools = await client.get_tools() print(f"✗ Loaded {len(tools)} stateless MCP tools") - + # Create model and agent from langchain_openai import ChatOpenAI - + model = ChatOpenAI(model="gpt-4", temperature=0) model_with_tools = model.bind_tools(tools) - + # Create agent with memory memory = MemorySaver() agent = create_agent( @@ -74,14 +72,14 @@ async def stateless_example_problematic(): system_prompt="You are a web testing engineer using Playwright tools.", checkpointer=memory, ) - + # This will fail on the second tool call because the browser session is lost test_prompt = """ 1. Navigate to https://example.com 2. Take a screenshot 3. Click on the 'More information...' link """ - + try: config = {"configurable": {"thread_id": "test-thread"}} response = await agent.ainvoke( @@ -96,14 +94,14 @@ async def stateless_example_problematic(): async def stateful_example_correct(): """Correct example showing stateful MCP usage that maintains browser sessions. - + This demonstrates the CORRECT way to use MCP tools for browser automation. A single browser session is maintained across all tool calls. """ print("\n" + "=" * 80) print("CORRECT EXAMPLE: Stateful MCP Usage (Browser session persists)") print("=" * 80) - + # Initialize MCP client with Playwright server client = MultiServerMCPClient( { @@ -118,20 +116,20 @@ async def stateful_example_correct(): } } ) - + # SOLUTION: Create a persistent session and load tools from it # This maintains the browser session across all tool calls async with client.session("playwright") as session: # Load tools with the persistent session tools = await load_mcp_tools(session) print(f"✓ Loaded {len(tools)} stateful MCP tools with persistent session") - + # Create model and agent from langchain_openai import ChatOpenAI - + model = ChatOpenAI(model="gpt-4", temperature=0) model_with_tools = model.bind_tools(tools) - + # Create agent with memory memory = MemorySaver() agent = create_agent( @@ -141,7 +139,7 @@ async def stateful_example_correct(): system_prompt="You are a web testing engineer using Playwright tools.", checkpointer=memory, ) - + # Complex multi-step browser interaction that requires session persistence test_prompt = """ Please perform the following browser automation tasks: @@ -152,7 +150,7 @@ async def stateful_example_correct(): 5. Go back to the previous page 6. Verify you're back on example.com """ - + try: config = {"configurable": {"thread_id": "test-thread"}} response = await agent.ainvoke( @@ -163,20 +161,20 @@ async def stateful_example_correct(): print("Response:", response["messages"][-1].content) except Exception as e: print(f"❌ Unexpected error: {e}") - + # Session automatically cleaned up when exiting the context manager print("✓ Browser session properly closed after all operations") async def advanced_stateful_example(): """Advanced example showing complex browser automation with form filling and validation. - + This demonstrates a real-world scenario where session persistence is critical. """ print("\n" + "=" * 80) print("ADVANCED EXAMPLE: Complex Browser Automation with Session State") print("=" * 80) - + client = MultiServerMCPClient( { "playwright": { @@ -189,15 +187,15 @@ async def advanced_stateful_example(): } } ) - + async with client.session("playwright") as session: tools = await load_mcp_tools(session) - + from langchain_openai import ChatOpenAI - + model = ChatOpenAI(model="gpt-4", temperature=0) model_with_tools = model.bind_tools(tools) - + memory = MemorySaver() agent = create_agent( model_with_tools, @@ -208,7 +206,7 @@ async def advanced_stateful_example(): Always maintain the browser session state across operations.""", checkpointer=memory, ) - + # Complex test scenario requiring persistent session test_scenario = """ Perform a comprehensive test of a web form: @@ -223,7 +221,7 @@ async def advanced_stateful_example(): 9. Submit the form 10. Verify the submission was successful by checking the response page """ - + try: config = {"configurable": {"thread_id": "form-test"}} response = await agent.ainvoke( @@ -238,7 +236,6 @@ async def advanced_stateful_example(): async def main(): """Run all examples to demonstrate the difference between stateless and stateful MCP usage.""" - print("\n" + "🔍" * 40) print("MCP STATEFUL BROWSER AUTOMATION EXAMPLES") print("🔍" * 40) @@ -255,19 +252,19 @@ The stateful approach is ESSENTIAL for: - Database connections that need transactions - Any tools requiring persistent state between calls """) - + # Run the problematic stateless example try: await stateless_example_problematic() except Exception as e: print(f"Stateless example failed (expected): {e}") - + # Run the correct stateful example await stateful_example_correct() - + # Run the advanced example await advanced_stateful_example() - + print("\n" + "=" * 80) print("SUMMARY: Always use stateful sessions for browser automation!") print("Pattern: async with client.session('server_name') as session:") diff --git a/libs/langchain_v1/langchain/agents/factory.py b/libs/langchain_v1/langchain/agents/factory.py index bcb4c25b7cb..daa2f663c78 100644 --- a/libs/langchain_v1/langchain/agents/factory.py +++ b/libs/langchain_v1/langchain/agents/factory.py @@ -66,17 +66,17 @@ STRUCTURED_OUTPUT_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes." class MCPSessionConfig(TypedDict): """Configuration for MCP session management in agents. - + When provided, enables stateful session management for MCP tools, maintaining persistent connections across multiple tool invocations. """ - + client: Required[Any] """The MCP client instance (e.g., MultiServerMCPClient).""" - + server_name: Required[str] """The name of the MCP server to connect to (e.g., 'playwright', 'database').""" - + auto_cleanup: NotRequired[bool] """Whether to automatically cleanup the session when the agent is destroyed. Defaults to True.""" @@ -653,12 +653,12 @@ def create_agent( # noqa: PLR0915 - `server_name`: The name of the MCP server (e.g., 'playwright') - `auto_cleanup`: Whether to auto-cleanup the session (default: True) - Example: + Example: ```python from langchain_mcp_adapters import MultiServerMCPClient - + mcp_client = MultiServerMCPClient() - + # Create agent with stateful MCP session agent = create_agent( model="gpt-4", @@ -667,7 +667,7 @@ def create_agent( # noqa: PLR0915 "client": mcp_client, "server_name": "playwright", "auto_cleanup": True, - } + }, ) ``` @@ -715,7 +715,7 @@ def create_agent( # noqa: PLR0915 # Handle tools being None or empty if tools is None: tools = [] - + # Handle MCP session configuration if provided # This enables stateful session management for MCP tools if mcp_session_config: @@ -724,11 +724,11 @@ def create_agent( # noqa: PLR0915 """Check if a tool is an MCP tool by examining its metadata.""" if not isinstance(tool, BaseTool): return False - + # Check for MCP-specific attributes or metadata if hasattr(tool, "__mcp_server__"): return True - + # Check tool metadata for MCP indicators metadata = getattr(tool, "metadata", {}) if isinstance(metadata, dict): @@ -737,20 +737,20 @@ def create_agent( # noqa: PLR0915 return True if metadata.get("source") == "mcp": return True - + # Check if tool name suggests MCP origin (common patterns) tool_name = getattr(tool, "name", "") mcp_prefixes = ["mcp_", "playwright_", "browser_", "puppeteer_"] if tool_name and any(tool_name.startswith(prefix) for prefix in mcp_prefixes): return True - + return False - + has_mcp_tools = any(is_mcp_tool(tool) for tool in tools) - + if has_mcp_tools: import warnings - + # Warn user about automatic session management warnings.warn( f"MCP tools detected. Enabling stateful session management for " @@ -758,14 +758,14 @@ def create_agent( # noqa: PLR0915 f"will share the same session. For explicit control, use " f"StatefulMCPAgentExecutor from langchain.agents.mcp_utils.", UserWarning, - stacklevel=2 + stacklevel=2, ) - + # Note: We don't actually create the session here during graph construction. # Instead, we'll wrap the tools with session-aware versions that will # create and manage the session at runtime. This maintains the lazy # initialization pattern and avoids creating sessions that might not be used. - + # Mark tools that should use the MCP session for tool in tools: if is_mcp_tool(tool) and isinstance(tool, BaseTool): @@ -774,7 +774,7 @@ def create_agent( # noqa: PLR0915 tool.metadata = {} elif not isinstance(tool.metadata, dict): tool.metadata = {"original_metadata": tool.metadata} - + tool.metadata["__mcp_session_config__"] = mcp_session_config # Convert response format and setup structured output tools @@ -1756,10 +1756,3 @@ def _add_middleware_edge( __all__ = [ "create_agent", ] - - - - - - - diff --git a/libs/langchain_v1/langchain/agents/mcp_utils.py b/libs/langchain_v1/langchain/agents/mcp_utils.py index 9c66cda0b30..b07383e85ef 100644 --- a/libs/langchain_v1/langchain/agents/mcp_utils.py +++ b/libs/langchain_v1/langchain/agents/mcp_utils.py @@ -9,11 +9,10 @@ stateful connections persist across multiple tool invocations. from __future__ import annotations import asyncio +from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager -from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, Literal +from typing import TYPE_CHECKING, Any -from langchain_core.messages import BaseMessage -from langchain_core.runnables import Runnable from langchain_core.tools import BaseTool if TYPE_CHECKING: @@ -22,40 +21,42 @@ if TYPE_CHECKING: from langchain_core.language_models import BaseChatModel from langgraph.graph.state import CompiledStateGraph - from langchain.agents.middleware.types import AgentMiddleware, AgentState + from langchain.agents.middleware.types import AgentMiddleware class StatefulMCPAgentExecutor: """Wrapper class that manages MCP session lifecycle for agents. - + This class ensures that MCP tools maintain persistent sessions across multiple invocations, solving the common issue of browser sessions terminating between tool calls. - + Example: ```python from langchain_mcp_adapters.client import MultiServerMCPClient from langchain.agents.mcp_utils import StatefulMCPAgentExecutor - - client = MultiServerMCPClient({ - "playwright": { - "command": "npx", - "args": ["@playwright/mcp@latest"], - "transport": "stdio", + + client = MultiServerMCPClient( + { + "playwright": { + "command": "npx", + "args": ["@playwright/mcp@latest"], + "transport": "stdio", + } } - }) - + ) + async with StatefulMCPAgentExecutor( client=client, server_name="playwright", model="gpt-4", ) as executor: - result = await executor.ainvoke({ - "messages": [{"role": "user", "content": "Navigate and interact with a webpage"}] - }) + result = await executor.ainvoke( + {"messages": [{"role": "user", "content": "Navigate and interact with a webpage"}]} + ) ``` """ - + def __init__( self, client: Any, # MultiServerMCPClient @@ -70,7 +71,7 @@ class StatefulMCPAgentExecutor: debug: bool = False, ) -> None: """Initialize the stateful MCP agent executor. - + Args: client: MultiServerMCPClient instance for MCP connections. server_name: Name of the MCP server to create a session for. @@ -91,34 +92,36 @@ class StatefulMCPAgentExecutor: self.interrupt_before = interrupt_before self.interrupt_after = interrupt_after self.debug = debug - + self._session = None self._agent: CompiledStateGraph | None = None self._tools: list[BaseTool] | None = None - + async def __aenter__(self) -> StatefulMCPAgentExecutor: """Async context manager entry: create session and initialize agent.""" try: # Import here to avoid circular dependencies from langchain_mcp_adapters.tools import load_mcp_tools + from langchain.agents import create_agent - + # Create persistent MCP session self._session = await self.client.session(self.server_name).__aenter__() - + # Load tools with the persistent session self._tools = await load_mcp_tools(self._session) - + # Initialize model if isinstance(self.model, str): from langchain.chat_models import init_chat_model + model = init_chat_model(self.model) else: model = self.model - + # Bind tools to model model_with_tools = model.bind_tools(self._tools) - + # Create agent with stateful tools self._agent = create_agent( model_with_tools, @@ -130,20 +133,20 @@ class StatefulMCPAgentExecutor: interrupt_after=self.interrupt_after, debug=self.debug, ) - + return self - + except Exception: # Clean up session if initialization fails if self._session: await self._session.__aexit__(None, None, None) raise - + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: """Async context manager exit: cleanup session.""" if self._session: await self._session.__aexit__(exc_type, exc_val, exc_tb) - + async def ainvoke( self, input: dict[str, Any], @@ -151,15 +154,15 @@ class StatefulMCPAgentExecutor: **kwargs: Any, ) -> dict[str, Any]: """Asynchronously invoke the agent with the given input. - + Args: input: Input dictionary containing messages. config: Optional configuration dictionary. **kwargs: Additional keyword arguments. - + Returns: Agent response dictionary. - + Raises: RuntimeError: If agent is not initialized (not in context manager). """ @@ -170,9 +173,9 @@ class StatefulMCPAgentExecutor: " result = await executor.ainvoke(...)" ) raise RuntimeError(msg) - + return await self._agent.ainvoke(input, config, **kwargs) - + def invoke( self, input: dict[str, Any], @@ -180,15 +183,15 @@ class StatefulMCPAgentExecutor: **kwargs: Any, ) -> dict[str, Any]: """Synchronously invoke the agent with the given input. - + Args: input: Input dictionary containing messages. config: Optional configuration dictionary. **kwargs: Additional keyword arguments. - + Returns: Agent response dictionary. - + Raises: RuntimeError: If agent is not initialized (not in context manager). """ @@ -199,9 +202,9 @@ class StatefulMCPAgentExecutor: " result = await executor.ainvoke(...)" ) raise RuntimeError(msg) - + return self._agent.invoke(input, config, **kwargs) - + async def astream( self, input: dict[str, Any], @@ -209,15 +212,15 @@ class StatefulMCPAgentExecutor: **kwargs: Any, ) -> AsyncIterator[dict[str, Any]]: """Stream agent responses asynchronously. - + Args: input: Input dictionary containing messages. config: Optional configuration dictionary. **kwargs: Additional keyword arguments. - + Yields: Agent response chunks. - + Raises: RuntimeError: If agent is not initialized (not in context manager). """ @@ -229,10 +232,10 @@ class StatefulMCPAgentExecutor: " ..." ) raise RuntimeError(msg) - + async for chunk in self._agent.astream(input, config, **kwargs): yield chunk - + def stream( self, input: dict[str, Any], @@ -240,15 +243,15 @@ class StatefulMCPAgentExecutor: **kwargs: Any, ) -> Iterator[dict[str, Any]]: """Stream agent responses synchronously. - + Args: input: Input dictionary containing messages. config: Optional configuration dictionary. **kwargs: Additional keyword arguments. - + Yields: Agent response chunks. - + Raises: RuntimeError: If agent is not initialized (not in context manager). """ @@ -260,15 +263,15 @@ class StatefulMCPAgentExecutor: " ..." ) raise RuntimeError(msg) - + for chunk in self._agent.stream(input, config, **kwargs): yield chunk - + @property def agent(self) -> CompiledStateGraph | None: """Get the underlying agent graph.""" return self._agent - + @property def tools(self) -> list[BaseTool] | None: """Get the loaded MCP tools.""" @@ -289,11 +292,11 @@ async def create_stateful_mcp_agent( auto_cleanup: bool = True, ) -> tuple[CompiledStateGraph, Any]: # (agent, session) """Factory function to create an agent with stateful MCP tools. - + This function creates an agent with MCP tools that maintain persistent sessions across multiple invocations. It returns both the agent and the session object for manual lifecycle management. - + Args: client: MultiServerMCPClient instance for MCP connections. server_name: Name of the MCP server to create a session for. @@ -305,24 +308,26 @@ async def create_stateful_mcp_agent( interrupt_after: Optional list of node names to interrupt after. debug: Whether to enable debug mode. auto_cleanup: If False, caller is responsible for session cleanup. - + Returns: Tuple of (agent, session). If auto_cleanup is False, the caller must call `await session.__aexit__(None, None, None)` when done. - + Example: ```python from langchain_mcp_adapters.client import MultiServerMCPClient from langchain.agents.mcp_utils import create_stateful_mcp_agent - - client = MultiServerMCPClient({ - "playwright": { - "command": "npx", - "args": ["@playwright/mcp@latest"], - "transport": "stdio", + + client = MultiServerMCPClient( + { + "playwright": { + "command": "npx", + "args": ["@playwright/mcp@latest"], + "transport": "stdio", + } } - }) - + ) + # With auto cleanup (recommended) agent, session = await create_stateful_mcp_agent( client=client, @@ -333,31 +338,33 @@ async def create_stateful_mcp_agent( result = await agent.ainvoke({"messages": [...]}) finally: await session.__aexit__(None, None, None) - + # Or use the StatefulMCPAgentExecutor context manager instead ``` """ # Import here to avoid circular dependencies from langchain_mcp_adapters.tools import load_mcp_tools + from langchain.agents import create_agent - + # Create persistent MCP session session = await client.session(server_name).__aenter__() - + try: # Load tools with the persistent session tools = await load_mcp_tools(session) - + # Initialize model if isinstance(model, str): from langchain.chat_models import init_chat_model + model_instance = init_chat_model(model) else: model_instance = model - + # Bind tools to model model_with_tools = model_instance.bind_tools(tools) - + # Create agent with stateful tools agent = create_agent( model_with_tools, @@ -369,11 +376,11 @@ async def create_stateful_mcp_agent( interrupt_after=interrupt_after, debug=debug, ) - + if auto_cleanup: # Wrap agent to auto-cleanup session on deletion original_del = getattr(agent, "__del__", None) - + def cleanup_on_del(self): # Schedule session cleanup try: @@ -384,14 +391,14 @@ async def create_stateful_mcp_agent( loop.run_until_complete(session.__aexit__(None, None, None)) except Exception: pass # Best effort cleanup - + if original_del: original_del() - + agent.__del__ = cleanup_on_del.__get__(agent, type(agent)) - + return agent, session - + except Exception: # Clean up session if initialization fails await session.__aexit__(None, None, None) @@ -412,10 +419,10 @@ async def mcp_agent_session( debug: bool = False, ) -> AsyncIterator[CompiledStateGraph]: """Context manager for creating an agent with stateful MCP tools. - + This is a convenience function that automatically manages the session lifecycle, ensuring proper cleanup even if an error occurs. - + Args: client: MultiServerMCPClient instance for MCP connections. server_name: Name of the MCP server to create a session for. @@ -426,31 +433,33 @@ async def mcp_agent_session( interrupt_before: Optional list of node names to interrupt before. interrupt_after: Optional list of node names to interrupt after. debug: Whether to enable debug mode. - + Yields: Configured agent with stateful MCP tools. - + Example: ```python from langchain_mcp_adapters.client import MultiServerMCPClient from langchain.agents.mcp_utils import mcp_agent_session - - client = MultiServerMCPClient({ - "playwright": { - "command": "npx", - "args": ["@playwright/mcp@latest"], - "transport": "stdio", + + client = MultiServerMCPClient( + { + "playwright": { + "command": "npx", + "args": ["@playwright/mcp@latest"], + "transport": "stdio", + } } - }) - + ) + async with mcp_agent_session( client=client, server_name="playwright", model="gpt-4", ) as agent: - result = await agent.ainvoke({ - "messages": [{"role": "user", "content": "Navigate to a webpage"}] - }) + result = await agent.ainvoke( + {"messages": [{"role": "user", "content": "Navigate to a webpage"}]} + ) ``` """ agent, session = await create_stateful_mcp_agent( @@ -465,7 +474,7 @@ async def mcp_agent_session( debug=debug, auto_cleanup=False, # We'll handle cleanup manually ) - + try: yield agent finally: diff --git a/libs/langchain_v1/tests/integration_tests/agents/test_mcp_playwright_session.py b/libs/langchain_v1/tests/integration_tests/agents/test_mcp_playwright_session.py index c2f995b8071..33d8908c3ff 100644 --- a/libs/langchain_v1/tests/integration_tests/agents/test_mcp_playwright_session.py +++ b/libs/langchain_v1/tests/integration_tests/agents/test_mcp_playwright_session.py @@ -25,12 +25,12 @@ from langchain.agents.mcp_utils import ( class MockPlaywrightServer: """Mock Playwright MCP server for testing browser automation.""" - + def __init__(self): self.browser_sessions = {} self.current_session_id = None self.session_counter = 0 - + def create_session(self, session_id: str) -> Dict[str, Any]: """Create a new browser session.""" self.session_counter += 1 @@ -45,11 +45,11 @@ class MockPlaywrightServer: self.browser_sessions[session_id] = session self.current_session_id = session_id return session - + def get_session(self, session_id: str) -> Dict[str, Any] | None: """Get an existing browser session.""" return self.browser_sessions.get(session_id) - + def close_session(self, session_id: str) -> None: """Close a browser session.""" if session_id in self.browser_sessions: @@ -60,20 +60,20 @@ class MockPlaywrightServer: class MockMCPSession: """Mock MCP session that simulates Playwright browser state.""" - + def __init__(self, server: MockPlaywrightServer, session_name: str): self.server = server self.session_name = session_name self.session_id = f"{session_name}_{id(self)}" self.is_open = False self.browser_session = None - + async def __aenter__(self): """Enter the session context and create browser session.""" self.is_open = True self.browser_session = self.server.create_session(self.session_id) return self - + async def __aexit__(self, exc_type, exc_val, exc_tb): """Exit the session context and close browser session.""" self.is_open = False @@ -84,11 +84,11 @@ class MockMCPSession: class MockMCPClient: """Mock MultiServerMCPClient for testing.""" - + def __init__(self, playwright_server: MockPlaywrightServer): self.playwright_server = playwright_server self.sessions_created = [] - + def session(self, server_name: str) -> MockMCPSession: """Create a mock session for the specified server.""" session = MockMCPSession(self.playwright_server, server_name) @@ -98,20 +98,20 @@ class MockMCPClient: def create_playwright_tools(session: MockMCPSession) -> List[BaseTool]: """Create mock Playwright tools that use the session's browser state.""" - + @tool def playwright_navigate(url: str) -> str: """Navigate to a URL in the browser. - + Args: url: The URL to navigate to. """ if not session.is_open or not session.browser_session: raise RuntimeError("No active browser session") - + if not session.browser_session["is_active"]: raise RuntimeError("Browser session has been closed") - + # Create a new page in the session page_id = f"page_{len(session.browser_session['pages'])}" session.browser_session["pages"][page_id] = { @@ -120,81 +120,83 @@ def create_playwright_tools(session: MockMCPSession) -> List[BaseTool]: "text_fields": [f"input_{i}" for i in range(2)], } session.browser_session["current_page"] = page_id - + return f"Navigated to {url} in session {session.session_id}" - + @tool def playwright_click(selector: str) -> str: """Click an element in the current page. - + Args: selector: The element selector to click. """ if not session.is_open or not session.browser_session: raise RuntimeError("No active browser session") - + if not session.browser_session["is_active"]: raise RuntimeError("Browser session has been closed") - + current_page_id = session.browser_session.get("current_page") if not current_page_id: raise RuntimeError("No page loaded - navigate to a URL first") - + page = session.browser_session["pages"][current_page_id] - + # Check if element exists if selector not in page["elements"]: raise RuntimeError(f"Element '{selector}' not found. Available: {page['elements']}") - + return f"Clicked {selector} on {page['url']} in session {session.session_id}" - + @tool def playwright_type(selector: str, text: str) -> str: """Type text into an input field. - + Args: selector: The input field selector. text: The text to type. """ if not session.is_open or not session.browser_session: raise RuntimeError("No active browser session") - + if not session.browser_session["is_active"]: raise RuntimeError("Browser session has been closed") - + current_page_id = session.browser_session.get("current_page") if not current_page_id: raise RuntimeError("No page loaded - navigate to a URL first") - + page = session.browser_session["pages"][current_page_id] - + # Check if text field exists if selector not in page["text_fields"]: - raise RuntimeError(f"Text field '{selector}' not found. Available: {page['text_fields']}") - + raise RuntimeError( + f"Text field '{selector}' not found. Available: {page['text_fields']}" + ) + # Store the typed text if "typed_text" not in page: page["typed_text"] = {} page["typed_text"][selector] = text - + return f"Typed '{text}' into {selector} on {page['url']} in session {session.session_id}" - + # Mark tools as MCP tools for tool_func in [playwright_navigate, playwright_click, playwright_type]: tool_func.metadata = {"mcp_server": "playwright", "source": "mcp"} - + return [playwright_navigate, playwright_click, playwright_type] class TestMCPPlaywrightSessionIntegration: """Integration tests for MCP Playwright session management.""" - + @pytest.mark.asyncio async def test_stateless_session_fails_on_multiple_operations(self): """Test that stateless MCP usage causes session termination between calls.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + # Simulate stateless tool usage (new session per tool call) async def simulate_stateless_tools(): # First tool call - navigate @@ -203,7 +205,7 @@ class TestMCPPlaywrightSessionIntegration: nav_tool = tools[0] result1 = nav_tool.invoke({"url": "https://example.com"}) assert "Navigated to https://example.com" in result1 - + # Second tool call - click (different session) async with mock_client.session("playwright") as session2: tools = create_playwright_tools(session2) @@ -211,77 +213,77 @@ class TestMCPPlaywrightSessionIntegration: # This should fail because it's a new session without navigation with pytest.raises(RuntimeError, match="No page loaded"): click_tool.invoke({"selector": "button_0"}) - + await simulate_stateless_tools() - + # Verify that multiple sessions were created assert len(mock_client.sessions_created) == 2 assert playwright_server.session_counter == 2 - + @pytest.mark.asyncio async def test_stateful_session_maintains_browser_state(self): """Test that stateful MCP session maintains browser state across operations.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + # Use stateful session management async with mock_client.session("playwright") as session: tools = create_playwright_tools(session) nav_tool, click_tool, type_tool = tools - + # Navigate to a page nav_result = nav_tool.invoke({"url": "https://example.com"}) assert "Navigated to https://example.com" in nav_result assert session.session_id in nav_result - + # Click an element - should work because session is maintained click_result = click_tool.invoke({"selector": "button_0"}) assert "Clicked button_0 on https://example.com" in click_result assert session.session_id in click_result - + # Type text - should also work in the same session type_result = type_tool.invoke({"selector": "input_0", "text": "Hello World"}) assert "Typed 'Hello World' into input_0" in type_result assert session.session_id in type_result - + # Verify all operations used the same session assert playwright_server.session_counter == 1 assert len(playwright_server.browser_sessions) == 1 - + # Verify browser state was maintained browser_session = playwright_server.get_session(session.session_id) assert browser_session is not None assert browser_session["is_active"] assert "page_0" in browser_session["pages"] assert browser_session["pages"]["page_0"]["typed_text"]["input_0"] == "Hello World" - + @pytest.mark.asyncio async def test_stateful_mcp_agent_executor(self): """Test StatefulMCPAgentExecutor maintains session across agent invocations.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: # Setup mock tool loading async def load_tools_with_session(session): return create_playwright_tools(session) - + mock_load_tools.side_effect = load_tools_with_session - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: # Create a mock agent that simulates tool usage mock_agent = MagicMock() - + async def simulate_agent_invoke(input_dict, config=None, **kwargs): messages = input_dict.get("messages", []) if not messages: return {"messages": [AIMessage(content="No input provided")]} - + last_message = messages[-1].content.lower() - + # Get the tools from the mock tools = mock_create_agent.call_args[1]["tools"] - + if "navigate" in last_message: result = tools[0].invoke({"url": "https://test.com"}) elif "click" in last_message: @@ -290,12 +292,12 @@ class TestMCPPlaywrightSessionIntegration: result = tools[2].invoke({"selector": "input_1", "text": "Test input"}) else: result = "Unknown command" - + return {"messages": [AIMessage(content=result)]} - + mock_agent.ainvoke = simulate_agent_invoke mock_create_agent.return_value = mock_agent - + # Use StatefulMCPAgentExecutor async with StatefulMCPAgentExecutor( client=mock_client, @@ -304,141 +306,142 @@ class TestMCPPlaywrightSessionIntegration: system_prompt="You are a browser automation assistant.", ) as executor: # Multiple operations in the same session - nav_result = await executor.ainvoke({ - "messages": [HumanMessage(content="Navigate to test.com")] - }) + nav_result = await executor.ainvoke( + {"messages": [HumanMessage(content="Navigate to test.com")]} + ) assert "Navigated to https://test.com" in nav_result["messages"][-1].content - - click_result = await executor.ainvoke({ - "messages": [HumanMessage(content="Click button_1")] - }) + + click_result = await executor.ainvoke( + {"messages": [HumanMessage(content="Click button_1")]} + ) assert "Clicked button_1" in click_result["messages"][-1].content - - type_result = await executor.ainvoke({ - "messages": [HumanMessage(content="Type some text")] - }) + + type_result = await executor.ainvoke( + {"messages": [HumanMessage(content="Type some text")]} + ) assert "Typed 'Test input'" in type_result["messages"][-1].content - + # Verify single session was used assert playwright_server.session_counter == 1 assert len(mock_client.sessions_created) == 1 - + @pytest.mark.asyncio async def test_mcp_agent_session_context_manager(self): """Test mcp_agent_session context manager for simplified usage.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: + async def load_tools_with_session(session): return create_playwright_tools(session) - + mock_load_tools.side_effect = load_tools_with_session - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() - mock_agent.ainvoke = AsyncMock(return_value={ - "messages": [AIMessage(content="Task completed")] - }) + mock_agent.ainvoke = AsyncMock( + return_value={"messages": [AIMessage(content="Task completed")]} + ) mock_create_agent.return_value = mock_agent - + async with mcp_agent_session( client=mock_client, server_name="playwright", model="gpt-4", ) as agent: # Use the agent - result = await agent.ainvoke({ - "messages": [HumanMessage(content="Automate browser")] - }) + result = await agent.ainvoke( + {"messages": [HumanMessage(content="Automate browser")]} + ) assert result["messages"][-1].content == "Task completed" - + # Verify session was created assert len(mock_client.sessions_created) == 1 assert playwright_server.session_counter == 1 - + # Verify session was closed after context exit session = mock_client.sessions_created[0] assert not session.is_open browser_session = playwright_server.get_session(session.session_id) assert browser_session is not None assert not browser_session["is_active"] - + @pytest.mark.asyncio async def test_element_references_remain_valid_across_calls(self): """Test that element references remain valid across multiple tool calls.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + async with mock_client.session("playwright") as session: tools = create_playwright_tools(session) nav_tool, click_tool, type_tool = tools - + # Navigate to create elements nav_tool.invoke({"url": "https://form.example.com"}) - + # Interact with multiple elements in sequence elements_to_click = ["button_0", "button_1", "button_2"] for element in elements_to_click: result = click_tool.invoke({"selector": element}) assert f"Clicked {element}" in result assert session.session_id in result - + # Type in multiple fields fields_to_fill = [("input_0", "First Name"), ("input_1", "Last Name")] for field, text in fields_to_fill: result = type_tool.invoke({"selector": field, "text": text}) assert f"Typed '{text}' into {field}" in result assert session.session_id in result - + # Verify all operations succeeded in the same session browser_session = playwright_server.get_session(session.session_id) assert browser_session["is_active"] page = browser_session["pages"]["page_0"] assert page["typed_text"]["input_0"] == "First Name" assert page["typed_text"]["input_1"] == "Last Name" - + @pytest.mark.asyncio async def test_session_cleanup_on_error(self): """Test that session is properly cleaned up even when errors occur.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + with pytest.raises(RuntimeError, match="Invalid selector"): async with mock_client.session("playwright") as session: tools = create_playwright_tools(session) nav_tool, click_tool, _ = tools - + # Navigate successfully nav_tool.invoke({"url": "https://example.com"}) - + # Try to click non-existent element click_tool.invoke({"selector": "non_existent_button"}) - + # Verify session was still cleaned up despite error assert len(mock_client.sessions_created) == 1 session = mock_client.sessions_created[0] assert not session.is_open - + browser_session = playwright_server.get_session(session.session_id) assert not browser_session["is_active"] - + @pytest.mark.asyncio async def test_create_agent_with_mcp_session_config(self): """Test create_agent with mcp_session_config parameter.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + # Create mock MCP tools mock_tools = [ Mock(spec=BaseTool, name="playwright_navigate", metadata={"mcp_server": "playwright"}), Mock(spec=BaseTool, name="playwright_click", metadata={"mcp_server": "playwright"}), ] - + with patch("langchain.chat_models.init_chat_model") as mock_init_model: mock_model = MagicMock() mock_init_model.return_value = mock_model - + # Create agent with MCP session config agent = create_agent( model="gpt-4", @@ -447,9 +450,9 @@ class TestMCPPlaywrightSessionIntegration: "client": mock_client, "server_name": "playwright", "auto_cleanup": True, - } + }, ) - + # Verify tools were marked with session config for tool in mock_tools: assert "__mcp_session_config__" in tool.metadata @@ -458,17 +461,17 @@ class TestMCPPlaywrightSessionIntegration: class TestBrowserSessionPersistence: """Specific tests for browser session persistence scenarios.""" - + @pytest.mark.asyncio async def test_complex_browser_workflow(self): """Test a complex browser workflow that requires session persistence.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + async with mock_client.session("playwright") as session: tools = create_playwright_tools(session) nav_tool, click_tool, type_tool = tools - + # Simulate a multi-step form submission workflow workflow_steps = [ ("navigate", {"url": "https://app.example.com/login"}), @@ -478,66 +481,68 @@ class TestBrowserSessionPersistence: ("navigate", {"url": "https://app.example.com/dashboard"}), ("click", {"selector": "button_1"}), # Some dashboard action ] - + results = [] for step_type, params in workflow_steps: if step_type == "navigate": result = nav_tool.invoke(params) elif step_type == "click": # For dashboard navigation, we need to navigate first - if params["selector"] == "button_1" and "dashboard" not in [p["url"] for p in session.browser_session["pages"].values()]: + if params["selector"] == "button_1" and "dashboard" not in [ + p["url"] for p in session.browser_session["pages"].values() + ]: nav_tool.invoke({"url": "https://app.example.com/dashboard"}) result = click_tool.invoke(params) elif step_type == "type": result = type_tool.invoke(params) else: result = "Unknown step" - + results.append(result) # Verify session ID is consistent assert session.session_id in result - + # Verify all steps completed successfully assert len(results) == len(workflow_steps) assert all(session.session_id in r for r in results) - + # Verify session remained active throughout assert playwright_server.session_counter == 1 browser_session = playwright_server.get_session(session.session_id) assert browser_session["is_active"] assert len(browser_session["pages"]) >= 2 # Login and dashboard pages - + @pytest.mark.asyncio async def test_session_persistence_prevents_ref_not_found_errors(self): """Test that session persistence prevents 'Ref not found' errors.""" playwright_server = MockPlaywrightServer() mock_client = MockMCPClient(playwright_server) - + # Track element references element_refs = {} - + async with mock_client.session("playwright") as session: tools = create_playwright_tools(session) nav_tool, click_tool, _ = tools - + # Navigate and store element references nav_tool.invoke({"url": "https://example.com"}) - + # Store available elements (simulating element discovery) browser_session = playwright_server.get_session(session.session_id) page = browser_session["pages"]["page_0"] for element in page["elements"]: element_refs[element] = session.session_id - + # Verify all element references remain valid for element_id in element_refs: result = click_tool.invoke({"selector": element_id}) assert "Clicked" in result assert element_refs[element_id] in result - + # Verify no "Ref not found" error occurred assert "not found" not in result.lower() assert "ref" not in result.lower() or "references" in result.lower() - + # Verify session remained consistent assert len(set(element_refs.values())) == 1 # All refs from same session diff --git a/libs/langchain_v1/tests/unit_tests/agents/test_mcp_utils.py b/libs/langchain_v1/tests/unit_tests/agents/test_mcp_utils.py index a4189780661..2669b8d55a8 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/test_mcp_utils.py +++ b/libs/langchain_v1/tests/unit_tests/agents/test_mcp_utils.py @@ -23,17 +23,17 @@ from langchain.agents.mcp_utils import ( class MockMCPSession: """Mock MCP session for testing.""" - + def __init__(self, name: str = "test_session"): self.name = name self.is_open = False self.call_count = 0 self.cleanup_called = False - + async def __aenter__(self): self.is_open = True return self - + async def __aexit__(self, exc_type, exc_val, exc_tb): self.is_open = False self.cleanup_called = True @@ -42,11 +42,11 @@ class MockMCPSession: class MockMCPClient: """Mock MultiServerMCPClient for testing.""" - + def __init__(self): self.sessions = {} self.session_create_count = 0 - + def session(self, server_name: str): """Create a mock session context manager.""" self.session_create_count += 1 @@ -75,24 +75,24 @@ def mock_browser_type(selector: str, text: str) -> str: class TestStatefulMCPAgentExecutor: """Test cases for StatefulMCPAgentExecutor class.""" - + @pytest.mark.asyncio async def test_session_persistence_across_tool_calls(self): """Test that session persists across multiple tool invocations.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate, mock_browser_click, mock_browser_type] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: # Create a mock agent that tracks tool calls mock_agent = MagicMock() - mock_agent.ainvoke = AsyncMock(return_value={ - "messages": [AIMessage(content="Task completed")] - }) + mock_agent.ainvoke = AsyncMock( + return_value={"messages": [AIMessage(content="Task completed")]} + ) mock_create_agent.return_value = mock_agent - + async with StatefulMCPAgentExecutor( client=mock_client, server_name="playwright", @@ -103,36 +103,38 @@ class TestStatefulMCPAgentExecutor: assert "playwright" in mock_client.sessions session = mock_client.sessions["playwright"] assert session.is_open - + # Make multiple invocations - await executor.ainvoke({"messages": [HumanMessage(content="Navigate to example.com")]}) + await executor.ainvoke( + {"messages": [HumanMessage(content="Navigate to example.com")]} + ) await executor.ainvoke({"messages": [HumanMessage(content="Click button")]}) await executor.ainvoke({"messages": [HumanMessage(content="Type text")]}) - + # Verify same session was used (only one session created) assert mock_client.session_create_count == 1 assert session.is_open # Session still open - + # Verify session was cleaned up after context exit assert session.cleanup_called assert not session.is_open - + @pytest.mark.asyncio async def test_session_cleanup_on_agent_termination(self): """Test that session is properly cleaned up when agent terminates.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() - mock_agent.ainvoke = AsyncMock(return_value={ - "messages": [AIMessage(content="Done")] - }) + mock_agent.ainvoke = AsyncMock( + return_value={"messages": [AIMessage(content="Done")]} + ) mock_create_agent.return_value = mock_agent - + # Create executor and verify cleanup async with StatefulMCPAgentExecutor( client=mock_client, @@ -142,26 +144,26 @@ class TestStatefulMCPAgentExecutor: session = mock_client.sessions["test_server"] assert session.is_open assert not session.cleanup_called - + # After context exit, session should be cleaned up assert session.cleanup_called assert not session.is_open - + @pytest.mark.asyncio async def test_error_handling_when_session_fails(self): """Test proper error handling when session creation fails.""" mock_client = MockMCPClient() - + # Make session creation fail original_session = mock_client.session - + def failing_session(server_name): if server_name == "failing_server": raise ConnectionError("Failed to connect to MCP server") return original_session(server_name) - + mock_client.session = failing_session - + with pytest.raises(ConnectionError, match="Failed to connect to MCP server"): async with StatefulMCPAgentExecutor( client=mock_client, @@ -169,15 +171,15 @@ class TestStatefulMCPAgentExecutor: model="gpt-4", ) as executor: pass # Should not reach here - + @pytest.mark.asyncio async def test_error_handling_during_tool_loading(self): """Test error handling when tool loading fails.""" mock_client = MockMCPClient() - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.side_effect = RuntimeError("Failed to load tools") - + with pytest.raises(RuntimeError, match="Failed to load tools"): async with StatefulMCPAgentExecutor( client=mock_client, @@ -185,79 +187,79 @@ class TestStatefulMCPAgentExecutor: model="gpt-4", ) as executor: pass # Should not reach here - + # Verify session was cleaned up even though initialization failed session = mock_client.sessions["test_server"] assert session.cleanup_called - + @pytest.mark.asyncio async def test_runtime_error_when_not_in_context(self): """Test that RuntimeError is raised when using executor outside context manager.""" mock_client = MockMCPClient() - + executor = StatefulMCPAgentExecutor( client=mock_client, server_name="test_server", model="gpt-4", ) - + # Try to use without entering context manager with pytest.raises(RuntimeError, match="Agent not initialized"): await executor.ainvoke({"messages": []}) - + with pytest.raises(RuntimeError, match="Agent not initialized"): executor.invoke({"messages": []}) - + with pytest.raises(RuntimeError, match="Agent not initialized"): async for _ in executor.astream({"messages": []}): pass - + with pytest.raises(RuntimeError, match="Agent not initialized"): for _ in executor.stream({"messages": []}): pass - + @pytest.mark.asyncio async def test_tools_maintain_state_between_invocations(self): """Test that tools maintain state between invocations.""" mock_client = MockMCPClient() - + # Create a stateful tool that maintains a counter class StatefulTool(BaseTool): name: str = "stateful_tool" description: str = "A tool that maintains state" - + def __init__(self): super().__init__() self.call_count = 0 self.session_id = None - + def _run(self, session_id: str) -> str: if self.session_id is None: self.session_id = session_id self.call_count += 1 return f"Call {self.call_count} in session {self.session_id}" - + async def _arun(self, session_id: str) -> str: return self._run(session_id) - + stateful_tool = StatefulTool() mock_tools = [stateful_tool] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: # Create mock agent that uses the stateful tool mock_agent = MagicMock() - + async def mock_ainvoke(input_dict, config=None, **kwargs): # Simulate tool usage result = stateful_tool._run("session_123") return {"messages": [AIMessage(content=result)]} - + mock_agent.ainvoke = mock_ainvoke mock_create_agent.return_value = mock_agent - + async with StatefulMCPAgentExecutor( client=mock_client, server_name="stateful_server", @@ -267,12 +269,12 @@ class TestStatefulMCPAgentExecutor: result1 = await executor.ainvoke({"messages": [HumanMessage(content="Call 1")]}) result2 = await executor.ainvoke({"messages": [HumanMessage(content="Call 2")]}) result3 = await executor.ainvoke({"messages": [HumanMessage(content="Call 3")]}) - + # Verify state was maintained assert "Call 1 in session session_123" in result1["messages"][-1].content assert "Call 2 in session session_123" in result2["messages"][-1].content assert "Call 3 in session session_123" in result3["messages"][-1].content - + # Verify same session was used throughout assert stateful_tool.call_count == 3 assert stateful_tool.session_id == "session_123" @@ -280,20 +282,20 @@ class TestStatefulMCPAgentExecutor: class TestCreateStatefulMCPAgent: """Test cases for create_stateful_mcp_agent factory function.""" - + @pytest.mark.asyncio async def test_agent_creation_with_session(self): """Test that agent is created with persistent session.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() mock_create_agent.return_value = mock_agent - + agent, session = await create_stateful_mcp_agent( client=mock_client, server_name="test_server", @@ -301,62 +303,62 @@ class TestCreateStatefulMCPAgent: system_prompt="Test prompt", auto_cleanup=False, ) - + # Verify agent was created assert agent == mock_agent - + # Verify session is open assert session.is_open assert not session.cleanup_called - + # Manual cleanup required when auto_cleanup=False await session.__aexit__(None, None, None) assert session.cleanup_called - + @pytest.mark.asyncio async def test_auto_cleanup_mode(self): """Test auto-cleanup mode with __del__ injection.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() mock_create_agent.return_value = mock_agent - + agent, session = await create_stateful_mcp_agent( client=mock_client, server_name="test_server", model="gpt-4", auto_cleanup=True, # Enable auto-cleanup ) - + # Verify __del__ was added to agent assert hasattr(agent, "__del__") - + # Session should still be open assert session.is_open - + # Manually cleanup for test await session.__aexit__(None, None, None) - + @pytest.mark.asyncio async def test_error_handling_during_creation(self): """Test that session is cleaned up if agent creation fails.""" mock_client = MockMCPClient() - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.side_effect = ValueError("Tool loading failed") - + with pytest.raises(ValueError, match="Tool loading failed"): await create_stateful_mcp_agent( client=mock_client, server_name="test_server", model="gpt-4", ) - + # Verify session was cleaned up session = mock_client.sessions["test_server"] assert session.cleanup_called @@ -364,23 +366,23 @@ class TestCreateStatefulMCPAgent: class TestMCPAgentSession: """Test cases for mcp_agent_session context manager.""" - + @pytest.mark.asyncio async def test_context_manager_lifecycle(self): """Test that context manager properly manages session lifecycle.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() - mock_agent.ainvoke = AsyncMock(return_value={ - "messages": [AIMessage(content="Done")] - }) + mock_agent.ainvoke = AsyncMock( + return_value={"messages": [AIMessage(content="Done")]} + ) mock_create_agent.return_value = mock_agent - + async with mcp_agent_session( client=mock_client, server_name="test_server", @@ -389,33 +391,33 @@ class TestMCPAgentSession: ) as agent: # Verify agent was created assert agent == mock_agent - + # Verify session is open session = mock_client.sessions["test_server"] assert session.is_open - + # Use the agent result = await agent.ainvoke({"messages": [HumanMessage(content="Test")]}) assert result["messages"][-1].content == "Done" - + # Verify session was cleaned up after context exit assert session.cleanup_called assert not session.is_open - + @pytest.mark.asyncio async def test_error_propagation_and_cleanup(self): """Test that errors are propagated and session is still cleaned up.""" mock_client = MockMCPClient() mock_tools = [mock_browser_navigate] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() mock_agent.ainvoke = AsyncMock(side_effect=RuntimeError("Agent failed")) mock_create_agent.return_value = mock_agent - + with pytest.raises(RuntimeError, match="Agent failed"): async with mcp_agent_session( client=mock_client, @@ -424,7 +426,7 @@ class TestMCPAgentSession: ) as agent: # This should raise an error await agent.ainvoke({"messages": []}) - + # Verify session was still cleaned up despite error session = mock_client.sessions["test_server"] assert session.cleanup_called @@ -433,19 +435,19 @@ class TestMCPAgentSession: class TestSessionStatePersistence: """Test cases specifically for verifying session state persistence.""" - + @pytest.mark.asyncio async def test_browser_session_persistence_simulation(self): """Simulate browser session persistence across navigation and interactions.""" mock_client = MockMCPClient() - + # Simulate browser state browser_state = { "current_url": None, "page_elements": [], "session_active": False, } - + @tool def browser_navigate_stateful(url: str) -> str: """Navigate to URL and maintain session.""" @@ -454,7 +456,7 @@ class TestSessionStatePersistence: browser_state["current_url"] = url browser_state["page_elements"] = [f"element_{i}" for i in range(5)] return f"Navigated to {url}, found {len(browser_state['page_elements'])} elements" - + @tool def browser_click_stateful(element_ref: str) -> str: """Click element in current session.""" @@ -463,82 +465,85 @@ class TestSessionStatePersistence: if element_ref not in browser_state["page_elements"]: raise RuntimeError(f"Element {element_ref} not found in current page") return f"Clicked {element_ref} on {browser_state['current_url']}" - + mock_tools = [browser_navigate_stateful, browser_click_stateful] - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: mock_load_tools.return_value = mock_tools - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: # Create mock agent that simulates tool usage mock_agent = MagicMock() - + async def simulate_browser_interaction(input_dict, config=None, **kwargs): message = input_dict["messages"][-1].content - + if "navigate" in message.lower(): result = browser_navigate_stateful.invoke({"url": "https://example.com"}) elif "click" in message.lower(): result = browser_click_stateful.invoke({"element_ref": "element_0"}) else: result = "Unknown command" - + return {"messages": [AIMessage(content=result)]} - + mock_agent.ainvoke = simulate_browser_interaction mock_create_agent.return_value = mock_agent - + async with StatefulMCPAgentExecutor( client=mock_client, server_name="playwright", model="gpt-4", ) as executor: # Navigate to page - nav_result = await executor.ainvoke({ - "messages": [HumanMessage(content="Navigate to example.com")] - }) + nav_result = await executor.ainvoke( + {"messages": [HumanMessage(content="Navigate to example.com")]} + ) assert "Navigated to https://example.com" in nav_result["messages"][-1].content assert browser_state["session_active"] - + # Click element - should work because session is maintained - click_result = await executor.ainvoke({ - "messages": [HumanMessage(content="Click element_0")] - }) - assert "Clicked element_0 on https://example.com" in click_result["messages"][-1].content - + click_result = await executor.ainvoke( + {"messages": [HumanMessage(content="Click element_0")]} + ) + assert ( + "Clicked element_0 on https://example.com" + in click_result["messages"][-1].content + ) + # Verify browser state was maintained throughout assert browser_state["current_url"] == "https://example.com" assert "element_0" in browser_state["page_elements"] assert browser_state["session_active"] - + @pytest.mark.asyncio async def test_multiple_server_sessions(self): """Test managing multiple MCP server sessions simultaneously.""" mock_client = MockMCPClient() - + # Track which sessions are used sessions_used = set() - + async def track_session_tool(session_name: str) -> str: sessions_used.add(session_name) return f"Used session: {session_name}" - + with patch("langchain.agents.mcp_utils.load_mcp_tools") as mock_load_tools: # Return different tools for different servers def get_tools_for_server(session): - if hasattr(session, 'name'): - if 'playwright' in session.name: + if hasattr(session, "name"): + if "playwright" in session.name: return [mock_browser_navigate] - elif 'database' in session.name: + elif "database" in session.name: return [mock_browser_click] # Different tool for database return [] - + mock_load_tools.side_effect = get_tools_for_server - + with patch("langchain.agents.mcp_utils.create_agent") as mock_create_agent: mock_agent = MagicMock() mock_create_agent.return_value = mock_agent - + # Create executors for different servers async with StatefulMCPAgentExecutor( client=mock_client, @@ -553,18 +558,18 @@ class TestSessionStatePersistence: # Verify different sessions were created assert "playwright" in mock_client.sessions assert "database" in mock_client.sessions - + playwright_session = mock_client.sessions["playwright"] db_session = mock_client.sessions["database"] - + # Verify both sessions are open assert playwright_session.is_open assert db_session.is_open - + # Verify sessions are different assert playwright_session != db_session assert playwright_session.name != db_session.name - + # Verify both sessions were cleaned up assert playwright_session.cleanup_called assert db_session.cleanup_called