github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2026-07-14 22:27:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Apply patch [skip ci]

Browse Source
This commit is contained in:
open-swe[bot]
2025-11-14 14:57:46 +00:00
parent b86bde538b
commit 157ef67a32
5 changed files with 396 additions and 387 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
libs/langchain_v1/examples/mcp_stateful_browser_example.py
View File

@@ -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:")

47
libs/langchain_v1/langchain/agents/factory.py
View File

@@ -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",
]

191
libs/langchain_v1/langchain/agents/mcp_utils.py
View File

@@ -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:

233
libs/langchain_v1/tests/integration_tests/agents/test_mcp_playwright_session.py
View File

@@ -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

245
libs/langchain_v1/tests/unit_tests/agents/test_mcp_utils.py
View File

@@ -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