Release v2.3.10

This commit is contained in:
Mervin Praison committed 2025-12-16 09:54:57 +00:00
1 parent cf2a5a135f
commit 67dbd1007a
36 files changed
+3843 -62

No files matched your search

+4
View File
@@ -64,6 +64,10 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- 📋 **Planning Mode** (Plan before execution - Agent & Multi-Agent)
- 🔧 **Planning Tools** (Research with tools during planning)
- 🧠 **Planning Reasoning** (Chain-of-thought planning)
- 🔌 **MCP Transports** (stdio, Streamable HTTP, WebSocket, SSE - Protocol 2025-11-25)
- 🌐 **WebSocket MCP** (Real-time bidirectional connections with auto-reconnect)
- 🔐 **MCP Security** (Origin validation, DNS rebinding prevention, secure sessions)
- 🔄 **MCP Resumability** (SSE stream recovery via Last-Event-ID)
## Using Python Code
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
flask \
"praisonai>=2.3.9" \
"praisonai>=2.3.10" \
"praisonai[api]" \
gunicorn \
markdown
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.3.9" \
"praisonai>=2.3.10" \
"praisonai[chat]" \
"embedchain[github,youtube]"
+1 -1
View File
@@ -20,7 +20,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.3.9" \
"praisonai>=2.3.10" \
"praisonai[ui]" \
"praisonai[chat]" \
"praisonai[realtime]" \
+1 -1
View File
@@ -16,7 +16,7 @@ RUN mkdir -p /root/.praison
# Install Python packages (using latest versions)
RUN pip install --no-cache-dir \
praisonai_tools \
"praisonai>=2.3.9" \
"praisonai>=2.3.10" \
"praisonai[ui]" \
"praisonai[crewai]"
@@ -0,0 +1,83 @@
"""
MCP Backward Compatibility Example
Demonstrates backward compatibility with older MCP servers
using the deprecated HTTP+SSE transport (protocol version 2024-11-05).
Features:
- Legacy SSE endpoint detection (/sse suffix)
- Automatic transport negotiation
- Protocol version support (2024-11-05 to 2025-11-25)
Protocol: MCP 2025-11-25 with 2024-11-05 compatibility
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # Legacy SSE server (auto-detected via /sse suffix)
# agent_legacy = Agent(
# name="Legacy SSE Assistant",
# tools=MCP("http://localhost:8080/sse")
# )
#
# # Modern Streamable HTTP server
# agent_modern = Agent(
# name="Modern HTTP Assistant",
# tools=MCP("http://localhost:8080/mcp")
# )
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_transport import get_transport_type
from praisonaiagents.mcp.mcp_compat import (
detect_transport_support,
is_legacy_sse_url,
TransportNegotiator
)
from praisonaiagents.mcp.mcp_session import VALID_PROTOCOL_VERSIONS
print("MCP Backward Compatibility Example")
print("=" * 50)
# URL detection
print("\n1. URL-based Transport Detection:")
urls = [
"http://localhost:8080/sse",
"https://api.example.com/sse",
"http://localhost:8080/mcp",
"https://api.example.com/mcp",
]
for url in urls:
transport = get_transport_type(url)
legacy = is_legacy_sse_url(url)
print(f" {url}")
print(f" Transport: {transport}, Legacy: {legacy}")
# Transport negotiation
print("\n2. Transport Negotiation (HTTP response based):")
responses = [
(200, "application/json", "Streamable HTTP"),
(200, "text/event-stream", "SSE stream"),
(400, None, "Fallback to SSE"),
(404, None, "Fallback to SSE"),
(405, None, "Fallback to SSE"),
]
for status, content_type, desc in responses:
detected = detect_transport_support(status, content_type)
print(f" HTTP {status} ({desc}): {detected}")
# Protocol versions
print("\n3. Supported Protocol Versions:")
for version in sorted(VALID_PROTOCOL_VERSIONS):
print(f" - {version}")
# Negotiator
print("\n4. Transport Negotiator:")
negotiator = TransportNegotiator("https://api.example.com/mcp")
print(f" Base URL: {negotiator.base_url}")
print(f" Try Streamable first: {negotiator.should_try_streamable_first()}")
negotiator_sse = TransportNegotiator("https://api.example.com/sse")
print(f" SSE URL: {negotiator_sse.base_url}")
print(f" Try Streamable first: {negotiator_sse.should_try_streamable_first()}")
+110
View File
@@ -0,0 +1,110 @@
"""
MCP Custom Transport Example
Demonstrates how to create custom MCP transports by extending
the BaseTransport class.
Features:
- BaseTransport abstract class
- TransportRegistry for registration
- TransportConfig for configuration
- Pluggable transport mechanism
Protocol: MCP 2025-11-25
"""
from praisonaiagents.mcp.mcp_transport import (
BaseTransport,
TransportConfig,
TransportRegistry,
get_default_registry
)
# Custom transport implementation
class MyCustomTransport(BaseTransport):
"""Example custom transport implementation."""
def __init__(self, url: str, config: TransportConfig = None):
self.url = url
self.config = config or TransportConfig()
self._connected = False
async def connect(self):
"""Establish connection."""
print(f"Connecting to {self.url}...")
self._connected = True
print("Connected!")
async def send(self, message: dict):
"""Send JSON-RPC message."""
if not self._connected:
raise RuntimeError("Not connected")
print(f"Sending: {message}")
async def receive(self) -> dict:
"""Receive JSON-RPC message."""
if not self._connected:
raise RuntimeError("Not connected")
return {"jsonrpc": "2.0", "result": "ok", "id": 1}
async def close(self):
"""Close connection."""
print("Closing connection...")
self._connected = False
@property
def is_connected(self) -> bool:
return self._connected
if __name__ == "__main__":
import asyncio
print("MCP Custom Transport Example")
print("=" * 50)
# 1. Transport configuration
print("\n1. Transport Configuration:")
config = TransportConfig(
timeout=60,
debug=True,
retry_count=3,
retry_delay=1.0,
auth_token="my-token"
)
print(f" Timeout: {config.timeout}s")
print(f" Debug: {config.debug}")
print(f" Retry count: {config.retry_count}")
# 2. Default registry
print("\n2. Default Transport Registry:")
registry = get_default_registry()
print(f" Registered transports: {registry.list_transports()}")
# 3. Register custom transport
print("\n3. Registering Custom Transport:")
custom_registry = TransportRegistry()
custom_registry.register("my_custom", MyCustomTransport)
print(f" Registered: {custom_registry.list_transports()}")
# 4. Use custom transport
print("\n4. Using Custom Transport:")
async def demo():
transport = MyCustomTransport("custom://localhost:8080", config)
async with transport:
await transport.send({"jsonrpc": "2.0", "method": "test", "id": 1})
response = await transport.receive()
print(f" Response: {response}")
asyncio.run(demo())
# 5. BaseTransport interface
print("\n5. BaseTransport Interface:")
print(" Required methods:")
print(" - async connect()")
print(" - async send(message: dict)")
print(" - async receive() -> dict")
print(" - async close()")
print(" - @property is_connected -> bool")
+47
View File
@@ -0,0 +1,47 @@
"""
MCP Resumability Example
Demonstrates SSE stream resumability for MCP transports.
Resumability allows recovering from disconnections without losing messages.
Features:
- Event ID tracking per stream
- Last-Event-ID header on reconnection
- Retry delay handling from server
- Message replay on reconnection
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Agent with resumability support (automatic)
agent = Agent(
name="Resumable Assistant",
instructions="You are a helpful assistant with resumable connections.",
tools=MCP("https://api.example.com/mcp")
)
# Resumability is automatic - the MCP client:
# 1. Tracks event IDs from SSE streams (id: field)
# 2. Stores last event ID per stream
# 3. Includes Last-Event-ID header on reconnection
# 4. Respects retry delay from server (retry: field)
if __name__ == "__main__":
print("MCP Resumability Example")
print("=" * 40)
print()
print("SSE Event Format:")
print(" id: event-123")
print(" retry: 3000")
print(" data: {\"jsonrpc\": \"2.0\", ...}")
print()
print("Resumability Flow:")
print("1. Client receives SSE events with id: field")
print("2. Client stores last event ID")
print("3. On disconnection, client reconnects with Last-Event-ID header")
print("4. Server replays messages after that event ID")
print("5. Client respects retry: field for reconnection delay")
print()
print("This prevents message loss during network interruptions.")
+79
View File
@@ -0,0 +1,79 @@
"""
MCP Security Example
Demonstrates security features for MCP transports.
These features help prevent attacks like DNS rebinding.
Features:
- Origin header validation
- DNS rebinding prevention
- Localhost binding for local servers
- Authentication header support
- Secure session ID generation
Protocol: MCP 2025-11-25
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # Agent with authentication
# agent = Agent(
# name="Secure Assistant",
# tools=MCP(
# "https://api.example.com/mcp",
# headers={"Authorization": "Bearer your-token"}
# )
# )
#
# # WebSocket with auth token
# agent_ws = Agent(
# name="Secure WebSocket Assistant",
# tools=MCP(
# "wss://api.example.com/mcp",
# auth_token="Bearer your-secret-token"
# )
# )
# Using security utilities directly (no server required)
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_security import (
is_valid_origin,
create_auth_header,
generate_secure_session_id,
SecurityConfig
)
print("MCP Security Example")
print("=" * 40)
# Origin validation (DNS rebinding prevention)
print("\n1. Origin Validation:")
allowed = ["localhost", "127.0.0.1", "example.com"]
print(f" Allowed origins: {allowed}")
print(f" 'http://localhost:8080' valid: {is_valid_origin('http://localhost:8080', allowed)}")
print(f" 'https://evil.com' valid: {is_valid_origin('https://evil.com', allowed)}")
# Authentication headers
print("\n2. Authentication Headers:")
bearer = create_auth_header("my-token", auth_type="bearer")
print(f" Bearer: {bearer}")
basic = create_auth_header("user:pass", auth_type="basic")
print(f" Basic: {basic}")
# Secure session IDs
print("\n3. Secure Session IDs:")
session_id = generate_secure_session_id()
print(f" Generated: {session_id}")
print(f" Length: {len(session_id)} chars")
# Security config
print("\n4. Security Configuration:")
config = SecurityConfig(
allowed_origins=["localhost", "example.com"],
require_auth=True
)
print(f" Validate origin: {config.validate_origin}")
print(f" Require auth: {config.require_auth}")
print(f" Bind address: {config.get_bind_address()}")
@@ -0,0 +1,44 @@
"""
MCP Session Management Example
Demonstrates session management features for MCP transports.
Sessions allow stateful interactions with MCP servers.
Features:
- Automatic session ID handling
- Session expiration detection (HTTP 404)
- Session termination via HTTP DELETE
- Protocol version negotiation
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Basic agent with automatic session management
agent = Agent(
name="Session Assistant",
instructions="You are a helpful assistant with session support.",
tools=MCP("https://api.example.com/mcp")
)
# Session management is automatic - the MCP client:
# 1. Receives Mcp-Session-Id from server during initialization
# 2. Includes session ID in all subsequent requests
# 3. Handles HTTP 404 (session expired) by reinitializing
# 4. Can terminate session explicitly when done
if __name__ == "__main__":
print("MCP Session Management Example")
print("=" * 40)
print()
print("Session Lifecycle:")
print("1. Client sends InitializeRequest")
print("2. Server returns InitializeResult with Mcp-Session-Id header")
print("3. Client includes Mcp-Session-Id in all subsequent requests")
print("4. If HTTP 404 received, client reinitializes session")
print("5. Client can send HTTP DELETE to terminate session")
print()
print("Headers used:")
print("- Mcp-Session-Id: Session identifier")
print("- Mcp-Protocol-Version: Protocol version (e.g., 2025-03-26)")
@@ -0,0 +1,91 @@
"""
MCP Transports Overview Example
Demonstrates all MCP transport types and auto-detection.
The MCP class automatically selects the appropriate transport
based on the URL pattern.
Transport Types:
- stdio: Command strings (local subprocess)
- Streamable HTTP: http:// or https:// URLs
- WebSocket: ws:// or wss:// URLs
- SSE (Legacy): URLs ending in /sse
Protocol: MCP 2025-11-25
"""
# Example code (requires running MCP servers):
#
# from praisonaiagents import Agent, MCP
#
# # 1. stdio Transport - Local subprocess
# agent_stdio = Agent(
# name="Stdio Assistant",
# tools=MCP("npx @modelcontextprotocol/server-memory")
# )
#
# # 2. Streamable HTTP Transport
# agent_http = Agent(
# name="HTTP Assistant",
# tools=MCP("https://api.example.com/mcp")
# )
#
# # 3. WebSocket Transport
# agent_ws = Agent(
# name="WebSocket Assistant",
# tools=MCP("wss://api.example.com/mcp")
# )
#
# # 4. SSE Transport (Legacy)
# agent_sse = Agent(
# name="SSE Assistant",
# tools=MCP("http://localhost:8080/sse")
# )
if __name__ == "__main__":
from praisonaiagents.mcp.mcp_transport import get_transport_type
print("MCP Transports Overview")
print("=" * 50)
print()
print("Transport Auto-Detection:")
print("-" * 50)
examples = [
("npx @mcp/server-memory", "stdio"),
("python3 server.py", "stdio"),
("https://api.example.com/mcp", "http_stream"),
("http://localhost:8080/mcp", "http_stream"),
("ws://localhost:8080/mcp", "websocket"),
("wss://api.example.com/mcp", "websocket"),
("http://localhost:8080/sse", "sse"),
("https://api.example.com/sse", "sse"),
]
for url, expected in examples:
detected = get_transport_type(url)
status = "✓" if detected == expected else "✗"
print(f" {status} '{url}'")
print(f" → {detected}")
print()
print("Transport Features:")
print("-" * 50)
print(" stdio:")
print(" - Local subprocess communication")
print(" - Newline-delimited JSON-RPC")
print(" - Environment variable support")
print()
print(" Streamable HTTP:")
print(" - Session management (Mcp-Session-Id)")
print(" - Protocol versioning (Mcp-Protocol-Version)")
print(" - SSE streaming with resumability")
print()
print(" WebSocket:")
print(" - Bidirectional real-time")
print(" - Auto-reconnect with exponential backoff")
print(" - Authentication token support")
print()
print(" SSE (Legacy):")
print(" - Backward compatibility")
print(" - Protocol version 2024-11-05")
@@ -0,0 +1,56 @@
"""
MCP Streamable HTTP Transport Example
Demonstrates connecting to MCP servers via Streamable HTTP transport.
This is the current standard transport (Protocol Revision 2025-11-25).
Features:
- Single MCP endpoint for all communication
- Session management via Mcp-Session-Id header
- Protocol versioning via Mcp-Protocol-Version header
- SSE streaming with resumability
- Session termination via HTTP DELETE
Protocol: MCP 2025-11-25
"""
from praisonaiagents import Agent, MCP
# Basic Streamable HTTP connection
agent = Agent(
name="HTTP Stream Assistant",
instructions="You are a helpful assistant connected via Streamable HTTP.",
tools=MCP("https://api.example.com/mcp")
)
# With authentication headers
agent_auth = Agent(
name="Authenticated Assistant",
instructions="You are a helpful assistant with authenticated connection.",
tools=MCP(
"https://api.example.com/mcp",
headers={"Authorization": "Bearer your-token"}
)
)
# With timeout and debug
agent_debug = Agent(
name="Debug Assistant",
tools=MCP(
"https://api.example.com/mcp",
timeout=120,
debug=True
)
)
if __name__ == "__main__":
print("Streamable HTTP MCP Transport Example")
print("=" * 40)
print("This example requires a Streamable HTTP MCP server running.")
print("URL pattern: https://host:port/mcp")
print()
print("Features enabled:")
print("- Mcp-Protocol-Version header (default: 2025-03-26)")
print("- Mcp-Session-Id header (automatic)")
print("- SSE streaming support")
print("- Resumability via Last-Event-ID")
+55
View File
@@ -0,0 +1,55 @@
"""
MCP WebSocket Transport Example
Demonstrates connecting to MCP servers via WebSocket (ws:// or wss://).
WebSocket transport provides bidirectional real-time communication,
ideal for cloud deployments and long-lived connections.
Features:
- Auto-detection of ws:// and wss:// URLs
- Authentication token support
- Auto-reconnect with exponential backoff
- Ping/pong keepalive
Protocol: SEP-1288 (WebSocket Transport)
"""
from praisonaiagents import Agent, MCP
# Basic WebSocket connection
agent = Agent(
name="WebSocket Assistant",
instructions="You are a helpful assistant connected via WebSocket.",
tools=MCP("ws://localhost:8080/mcp")
)
# Secure WebSocket with authentication
agent_secure = Agent(
name="Secure WebSocket Assistant",
instructions="You are a helpful assistant with secure WebSocket connection.",
tools=MCP(
"wss://api.example.com/mcp",
auth_token="Bearer your-secret-token",
timeout=60
)
)
# WebSocket with custom options
agent_custom = Agent(
name="Custom WebSocket Assistant",
tools=MCP(
"wss://api.example.com/mcp",
auth_token="your-token",
timeout=120,
debug=True
)
)
if __name__ == "__main__":
# Note: Requires a running WebSocket MCP server
# Example server: https://github.com/modelcontextprotocol/servers
print("WebSocket MCP Transport Example")
print("=" * 40)
print("This example requires a WebSocket MCP server running.")
print("URL patterns: ws://host:port/path or wss://host:port/path")
+31 -4
View File
@@ -655,14 +655,23 @@ agent = Agent(tools=[my_tool])
## MCP (Model Context Protocol)
Connect to MCP servers for extended capabilities.
Connect to MCP servers for extended capabilities. Supports **Protocol Revision 2025-11-25**.
### MCP Transports
| Transport | URL Pattern | Use Case |
|-----------|-------------|----------|
| **stdio** | Command string | Local NPX/Python servers |
| **Streamable HTTP** | `http(s)://...` | Production deployments |
| **WebSocket** | `ws://` or `wss://` | Real-time bidirectional |
| **SSE (Legacy)** | `http(s)://.../sse` | Backward compatibility |
### MCP Client (Consume MCP Servers)
```python
from praisonaiagents import Agent, MCP
# Local MCP server
# stdio - Local MCP server
agent = Agent(
tools=MCP(
command="npx",
@@ -670,9 +679,19 @@ agent = Agent(
)
)
# Remote MCP server (SSE)
# Streamable HTTP - Production server
agent = Agent(
tools=MCP(url="http://localhost:8080/sse")
tools=MCP("https://api.example.com/mcp")
)
# WebSocket - Real-time bidirectional (NEW)
agent = Agent(
tools=MCP("wss://api.example.com/mcp", auth_token="your-token")
)
# SSE (Legacy) - Backward compatibility
agent = Agent(
tools=MCP("http://localhost:8080/sse")
)
# With environment variables
@@ -685,6 +704,14 @@ agent = Agent(
)
```
### MCP Features (Protocol 2025-11-25)
- **Session Management**: Automatic `Mcp-Session-Id` handling
- **Protocol Versioning**: `Mcp-Protocol-Version` header on all requests
- **Resumability**: SSE stream recovery via `Last-Event-ID`
- **Security**: Origin validation, DNS rebinding prevention
- **WebSocket**: Auto-reconnect with exponential backoff
### MCP Server (Expose Tools as MCP Server)
Expose your Python functions as MCP tools for external clients (Claude Desktop, Cursor, etc.).
@@ -1,40 +1,66 @@
"""
Model Context Protocol (MCP) integration for PraisonAI Agents.
This package provides classes and utilities for:
- Connecting to MCP servers (as client) using different transports (stdio, SSE, etc.)
- Exposing Python functions as MCP servers for external clients
This package provides classes and utilities for connecting to MCP servers
using different transport methods:
- stdio: Standard input/output (subprocess)
- sse: Server-Sent Events (legacy HTTP+SSE)
- http_stream: Streamable HTTP (current standard)
- websocket: WebSocket (SEP-1288)
Client Usage:
from praisonaiagents.mcp import MCP
agent = Agent(
tools=MCP("npx -y @modelcontextprotocol/server-weather")
)
Server Usage:
from praisonaiagents.mcp import ToolsMCPServer
def my_tool(query: str) -> str:
'''Search for information.'''
return f"Results for {query}"
server = ToolsMCPServer(name="my-tools")
server.register_tool(my_tool)
server.run() # Starts MCP server
Protocol Revision: 2025-11-25
"""
from .mcp import MCP
from .mcp_server import ToolsMCPServer, launch_tools_mcp_server
from .mcp_utils import function_to_mcp_schema, get_tool_metadata, python_type_to_json_schema
# Lazy imports for optional components
def __getattr__(name):
"""Lazy import for optional modules."""
if name == "WebSocketMCPClient":
from .mcp_websocket import WebSocketMCPClient
return WebSocketMCPClient
elif name == "SessionManager":
from .mcp_session import SessionManager
return SessionManager
elif name == "SecurityConfig":
from .mcp_security import SecurityConfig
return SecurityConfig
elif name == "BaseTransport":
from .mcp_transport import BaseTransport
return BaseTransport
elif name == "TransportConfig":
from .mcp_transport import TransportConfig
return TransportConfig
elif name == "ToolsMCPServer":
from .mcp_server import ToolsMCPServer
return ToolsMCPServer
elif name == "launch_tools_mcp_server":
from .mcp_server import launch_tools_mcp_server
return launch_tools_mcp_server
elif name == "function_to_mcp_schema":
from .mcp_utils import function_to_mcp_schema
return function_to_mcp_schema
elif name == "get_tool_metadata":
from .mcp_utils import get_tool_metadata
return get_tool_metadata
elif name == "python_type_to_json_schema":
from .mcp_utils import python_type_to_json_schema
return python_type_to_json_schema
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
# Client
"MCP",
# Server
# Server (lazy loaded)
"ToolsMCPServer",
"launch_tools_mcp_server",
# Utilities
# Utilities (lazy loaded)
"function_to_mcp_schema",
"get_tool_metadata",
"python_type_to_json_schema",
# Transport components (lazy loaded)
"WebSocketMCPClient",
"SessionManager",
"SecurityConfig",
"BaseTransport",
"TransportConfig",
]
@@ -227,6 +227,26 @@ class MCP:
self.timeout = timeout
self.debug = debug
# Check if this is a WebSocket URL (ws:// or wss://)
if isinstance(command_or_string, str) and re.match(r'^wss?://', command_or_string):
from .mcp_websocket import WebSocketMCPClient
# Extract auth token if provided
auth_token = kwargs.pop('auth_token', None)
self.websocket_client = WebSocketMCPClient(
command_or_string,
debug=debug,
timeout=timeout,
auth_token=auth_token,
options=kwargs
)
self._tools = list(self.websocket_client.tools)
self.is_sse = False
self.is_http_stream = False
self.is_websocket = True
self.is_npx = False
return
# Check if this is an HTTP URL
if isinstance(command_or_string, str) and re.match(r'^https?://', command_or_string):
# Determine transport type based on URL or kwargs
@@ -0,0 +1,150 @@
"""
MCP Backward Compatibility Module.
This module provides utilities for maintaining backward compatibility
with older MCP servers and the deprecated HTTP+SSE transport from
protocol version 2024-11-05.
Per MCP specification, clients wanting to support older servers should:
1. Attempt POST to server URL with InitializeRequest
2. If it fails with 400, 404, or 405, fall back to legacy SSE
3. Issue GET to server URL expecting SSE stream with endpoint event
"""
from typing import Optional
def detect_transport_support(
response_status: int,
content_type: Optional[str] = None
) -> str:
"""
Detect which transport a server supports based on response.
Per MCP spec backward compatibility guide:
- Success (200) with JSON indicates Streamable HTTP
- 400, 404, 405 indicates need for legacy SSE fallback
Args:
response_status: HTTP response status code
content_type: Response Content-Type header
Returns:
Transport type: "http_stream" or "sse"
"""
# Check for fallback indicators
if response_status in (400, 404, 405):
return "sse"
# Successful response indicates Streamable HTTP support
if response_status == 200:
return "http_stream"
# Default to SSE for unknown cases
return "sse"
def is_legacy_sse_url(url: str) -> bool:
"""
Check if URL indicates legacy SSE transport.
Legacy SSE URLs typically end with /sse.
Args:
url: Server URL
Returns:
True if URL indicates legacy SSE
"""
return url.rstrip('/').endswith('/sse')
def get_legacy_sse_endpoint(base_url: str) -> str:
"""
Get the legacy SSE endpoint URL.
Args:
base_url: Base server URL
Returns:
SSE endpoint URL
"""
base = base_url.rstrip('/')
if not base.endswith('/sse'):
return f"{base}/sse"
return base
def get_streamable_http_endpoint(base_url: str) -> str:
"""
Get the Streamable HTTP endpoint URL.
Args:
base_url: Base server URL
Returns:
MCP endpoint URL
"""
base = base_url.rstrip('/')
# Remove /sse suffix if present
if base.endswith('/sse'):
base = base[:-4]
if not base.endswith('/mcp'):
return f"{base}/mcp"
return base
class TransportNegotiator:
"""
Handles automatic transport negotiation with MCP servers.
This class implements the backward compatibility algorithm
from the MCP specification for detecting server capabilities.
"""
def __init__(self, base_url: str):
"""
Initialize negotiator.
Args:
base_url: Server base URL
"""
self.base_url = base_url
self._detected_transport: Optional[str] = None
@property
def detected_transport(self) -> Optional[str]:
"""Get the detected transport type."""
return self._detected_transport
def set_detected_transport(self, transport: str) -> None:
"""
Set the detected transport type.
Args:
transport: Transport type ("http_stream" or "sse")
"""
self._detected_transport = transport
def get_endpoint_url(self) -> str:
"""
Get the appropriate endpoint URL based on detected transport.
Returns:
Endpoint URL for the detected transport
"""
if self._detected_transport == "sse":
return get_legacy_sse_endpoint(self.base_url)
return get_streamable_http_endpoint(self.base_url)
def should_try_streamable_first(self) -> bool:
"""
Determine if Streamable HTTP should be tried first.
Per MCP spec, clients should try Streamable HTTP first
unless the URL explicitly indicates legacy SSE.
Returns:
True if should try Streamable HTTP first
"""
return not is_legacy_sse_url(self.base_url)
@@ -200,16 +200,31 @@ class HTTPStreamTransport:
This transport provides a single endpoint for all MCP communication,
supporting both batch (JSON) and streaming (SSE) response modes.
Per MCP Protocol Revision 2025-11-25:
- Includes MCP-Protocol-Version header on all requests
- Handles session management via Mcp-Session-Id header
- Supports SSE resumability via Last-Event-ID
"""
# Default protocol version for backward compatibility
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
def __init__(self, base_url: str, session_id: Optional[str] = None, options: Optional[Dict[str, Any]] = None):
self.base_url = base_url
self.session_id = session_id
self.options = options or {}
self.response_mode = self.options.get('responseMode', 'batch')
self.protocol_version = self.options.get('protocol_version', self.DEFAULT_PROTOCOL_VERSION)
# Track last event ID for resumability
self.last_event_id: Optional[str] = None
self._retry_delay_ms: int = 3000 # Default retry delay
self.headers = {
'Content-Type': 'application/json',
'Accept': 'application/json, text/event-stream'
'Accept': 'application/json, text/event-stream',
'Mcp-Protocol-Version': self.protocol_version # Required per spec
}
if session_id:
self.headers['Mcp-Session-Id'] = session_id
@@ -281,7 +296,7 @@ class HTTPStreamTransport:
raise
async def _process_sse_response(self, response):
"""Process SSE response stream."""
"""Process SSE response stream with resumability support."""
buffer = ""
async for chunk in response.content:
buffer += chunk.decode('utf-8')
@@ -291,12 +306,34 @@ class HTTPStreamTransport:
event, buffer = buffer.split("\n\n", 1)
lines = event.strip().split("\n")
# Parse SSE event
data = None
for line in lines:
if line.startswith("data: "):
data = line[6:] # Remove "data: " prefix
# Parse SSE event fields per spec
event_id = None
data_lines = []
retry_value = None
for line in lines:
if line.startswith("id:"):
# Track event ID for resumability
event_id = line[3:].strip()
elif line.startswith("data:"):
data_lines.append(line[5:].strip())
elif line.startswith("retry:"):
# Handle retry field per SSE spec
try:
retry_value = int(line[6:].strip())
except ValueError:
pass
# Update last event ID for resumability
if event_id:
self.last_event_id = event_id
# Update retry delay if provided
if retry_value is not None:
self._retry_delay_ms = retry_value
# Process data if present
data = '\n'.join(data_lines) if data_lines else None
if data:
try:
message = json.loads(data)
@@ -310,13 +347,49 @@ class HTTPStreamTransport:
except json.JSONDecodeError:
logger.error(f"Failed to parse SSE event: {data}")
async def terminate_session(self) -> bool:
"""
Terminate the session via HTTP DELETE.
Per MCP spec: Clients that no longer need a particular session
SHOULD send an HTTP DELETE to the MCP endpoint with the
Mcp-Session-Id header to explicitly terminate the session.
Returns:
True if session was terminated, False if server doesn't support it
"""
if not self._session or not self.session_id:
return False
try:
headers = {
'Mcp-Session-Id': self.session_id,
'Mcp-Protocol-Version': self.protocol_version
}
async with self._session.delete(self.base_url, headers=headers) as response:
if response.status == 405:
# Server doesn't allow client-initiated session termination
logger.debug("Server does not allow session termination via DELETE")
return False
elif response.status in (200, 202, 204):
# Session terminated successfully
self.session_id = None
self.headers.pop('Mcp-Session-Id', None)
return True
else:
logger.warning(f"Unexpected response to DELETE: {response.status}")
return False
except Exception as e:
logger.error(f"Error terminating session: {e}")
return False
async def start_sse_listener(self):
"""Start listening for SSE events from the server."""
if self._sse_task is None or self._sse_task.done():
self._sse_task = asyncio.create_task(self._sse_listener())
async def _sse_listener(self):
"""Background task to listen for SSE events."""
"""Background task to listen for SSE events with resumability support."""
while True:
try:
# Check if we should stop
@@ -324,18 +397,28 @@ class HTTPStreamTransport:
break
url = self.base_url
if self.session_id:
# Add session as query parameter for SSE connection
url = f"{url}?session={self.session_id}"
# Build headers per MCP spec
headers = {
'Accept': 'text/event-stream',
'Cache-Control': 'no-cache'
'Cache-Control': 'no-cache',
'Mcp-Protocol-Version': self.protocol_version # Required per spec
}
if self.session_id:
headers['Mcp-Session-Id'] = self.session_id
# Include Last-Event-ID for resumability per spec
if self.last_event_id:
headers['Last-Event-ID'] = self.last_event_id
async with self._session.get(url, headers=headers) as response:
# Handle session expiration (HTTP 404)
if response.status == 404:
logger.warning("Session expired (HTTP 404), need to reinitialize")
self.session_id = None
self.headers.pop('Mcp-Session-Id', None)
break
buffer = ""
async for chunk in response.content:
# Check if we should stop
@@ -349,12 +432,32 @@ class HTTPStreamTransport:
event, buffer = buffer.split("\n\n", 1)
lines = event.strip().split("\n")
# Parse SSE event
data = None
for line in lines:
if line.startswith("data: "):
data = line[6:] # Remove "data: " prefix
# Parse SSE event fields
event_id = None
data_lines = []
retry_value = None
for line in lines:
if line.startswith("id:"):
event_id = line[3:].strip()
elif line.startswith("data:"):
data_lines.append(line[5:].strip())
elif line.startswith("retry:"):
try:
retry_value = int(line[6:].strip())
except ValueError:
pass
# Update last event ID for resumability
if event_id:
self.last_event_id = event_id
# Update retry delay if provided
if retry_value is not None:
self._retry_delay_ms = retry_value
# Process data
data = '\n'.join(data_lines) if data_lines else None
if data:
try:
message = json.loads(data)
@@ -368,7 +471,8 @@ class HTTPStreamTransport:
except Exception as e:
if not (hasattr(self, '_closing') and self._closing):
logger.error(f"SSE listener error: {e}")
await asyncio.sleep(1) # Reconnect after 1 second
# Use retry delay from server or default
await asyncio.sleep(self._retry_delay_ms / 1000.0)
else:
break
@@ -0,0 +1,248 @@
"""
MCP Security Module.
This module provides security utilities for MCP transports,
implementing security best practices from MCP Protocol
Revision 2025-11-25.
Security features:
- Origin header validation (DNS rebinding prevention)
- Localhost binding recommendations
- Authentication header support
- Secure session ID generation
"""
import secrets
import base64
from typing import Optional, Dict, List, Set
from urllib.parse import urlparse
from dataclasses import dataclass, field
# Localhost addresses
LOCALHOST_ADDRESSES: Set[str] = {
'localhost',
'127.0.0.1',
'::1', # IPv6 localhost
}
def is_valid_origin(
origin: Optional[str],
allowed_origins: List[str],
allow_missing: bool = False
) -> bool:
"""
Validate Origin header against allowed origins.
Per MCP spec, servers MUST validate the Origin header on all
incoming connections to prevent DNS rebinding attacks.
Args:
origin: The Origin header value
allowed_origins: List of allowed origin hosts
allow_missing: Whether to allow missing Origin header
Returns:
True if origin is valid, False otherwise
"""
if not origin:
return allow_missing
# Extract host from origin
host = extract_origin_host(origin)
if not host:
return False
# Check against allowed origins
return host in allowed_origins
def extract_origin_host(origin: str) -> Optional[str]:
"""
Extract the host from an Origin header value.
Args:
origin: Origin header value (e.g., "https://example.com:443")
Returns:
Host portion of the origin, or None if invalid
"""
if not origin:
return None
try:
parsed = urlparse(origin)
return parsed.hostname
except Exception:
return None
def is_localhost_address(address: str) -> bool:
"""
Check if an address is a localhost address.
Args:
address: IP address or hostname
Returns:
True if localhost, False otherwise
"""
return address.lower() in LOCALHOST_ADDRESSES
def should_bind_localhost_only(is_local: bool = True) -> bool:
"""
Determine if server should bind to localhost only.
Per MCP spec, when running locally, servers SHOULD bind only
to localhost (127.0.0.1) rather than all network interfaces.
Args:
is_local: Whether this is a local development server
Returns:
True if should bind to localhost only
"""
return is_local
def is_potential_dns_rebinding(origin: str, server_host: str) -> bool:
"""
Detect potential DNS rebinding attack.
DNS rebinding occurs when an external origin tries to access
a local server through DNS manipulation.
Args:
origin: The Origin header value
server_host: The host the server is bound to
Returns:
True if potential DNS rebinding attempt
"""
origin_host = extract_origin_host(origin)
if not origin_host:
return True # Treat invalid origin as suspicious
# If server is on localhost, only localhost origins are safe
if is_localhost_address(server_host):
return not is_localhost_address(origin_host)
# For non-localhost servers, check if origin matches server
return origin_host != server_host
def create_auth_header(
credentials: str,
auth_type: str = "bearer",
header_name: str = "Authorization"
) -> Dict[str, str]:
"""
Create an authentication header.
Args:
credentials: The credentials/token
auth_type: Type of auth ("bearer", "basic", "custom")
header_name: Custom header name (for auth_type="custom")
Returns:
Dictionary with authentication header
"""
if auth_type.lower() == "bearer":
return {"Authorization": f"Bearer {credentials}"}
elif auth_type.lower() == "basic":
# Base64 encode for Basic auth
encoded = base64.b64encode(credentials.encode()).decode()
return {"Authorization": f"Basic {encoded}"}
elif auth_type.lower() == "custom":
return {header_name: credentials}
else:
raise ValueError(f"Unknown auth_type: {auth_type}")
def validate_auth_header(headers: Dict[str, str]) -> bool:
"""
Validate that authentication header is present.
Args:
headers: Request headers dictionary
Returns:
True if Authorization header is present
"""
# Check for Authorization header (case-insensitive)
for key in headers:
if key.lower() == 'authorization':
return bool(headers[key])
return False
def generate_secure_session_id(length: int = 32) -> str:
"""
Generate a cryptographically secure session ID.
Per MCP spec, session IDs SHOULD be globally unique and
cryptographically secure (e.g., securely generated UUID,
JWT, or cryptographic hash).
The ID MUST only contain visible ASCII characters (0x21-0x7E).
Args:
length: Desired length of session ID (minimum 32)
Returns:
Secure session ID string
"""
# Ensure minimum length for security
length = max(length, 32)
# Use URL-safe base64 which only uses visible ASCII
# Generate extra bytes to account for base64 expansion
num_bytes = (length * 3) // 4 + 1
random_bytes = secrets.token_bytes(num_bytes)
# Encode to URL-safe base64 (uses A-Z, a-z, 0-9, -, _)
session_id = base64.urlsafe_b64encode(random_bytes).decode()
# Trim to desired length
return session_id[:length]
@dataclass
class SecurityConfig:
"""
Security configuration for MCP transports.
This class holds security settings that can be applied
to MCP server implementations.
Attributes:
validate_origin: Whether to validate Origin header
allow_missing_origin: Whether to allow missing Origin
allowed_origins: List of allowed origin hosts
require_auth: Whether authentication is required
bind_localhost_only: Whether to bind to localhost only
"""
validate_origin: bool = True
allow_missing_origin: bool = False
allowed_origins: List[str] = field(default_factory=lambda: ["localhost", "127.0.0.1"])
require_auth: bool = False
bind_localhost_only: bool = True
def is_origin_allowed(self, origin: Optional[str]) -> bool:
"""Check if an origin is allowed by this config."""
if not self.validate_origin:
return True
return is_valid_origin(
origin,
self.allowed_origins,
self.allow_missing_origin
)
def get_bind_address(self) -> str:
"""Get recommended bind address based on config."""
return "127.0.0.1" if self.bind_localhost_only else "0.0.0.0"
@@ -0,0 +1,457 @@
"""
MCP Session Management Module.
This module provides session management utilities for MCP transports,
implementing the session management requirements from MCP Protocol
Revision 2025-11-25.
Features:
- Session ID validation (visible ASCII 0x21-0x7E)
- Session lifecycle management
- Protocol version header handling
- Session expiration detection
"""
import re
from typing import Optional, Dict, Any
# Valid protocol versions
VALID_PROTOCOL_VERSIONS = {
'2024-11-05', # Original HTTP+SSE
'2025-03-26', # Streamable HTTP introduction
'2025-06-18', # Updates
'2025-11-25', # Current revision
}
# Default protocol version for backward compatibility
DEFAULT_PROTOCOL_VERSION = '2025-03-26'
def is_valid_session_id(session_id: Optional[str]) -> bool:
"""
Validate MCP session ID per specification.
Session IDs MUST only contain visible ASCII characters (0x21 to 0x7E).
This excludes space (0x20), control characters (0x00-0x1F), and DEL (0x7F).
Args:
session_id: The session ID to validate
Returns:
True if valid, False otherwise
"""
if session_id is None or session_id == "":
return False
# Check each character is in visible ASCII range (0x21-0x7E)
for char in session_id:
code = ord(char)
if code < 0x21 or code > 0x7E:
return False
return True
def is_valid_protocol_version(version: Optional[str]) -> bool:
"""
Validate MCP protocol version string.
Args:
version: Protocol version string (e.g., '2025-11-25')
Returns:
True if valid, False otherwise
"""
if not version:
return False
# Check format: YYYY-MM-DD
if not re.match(r'^\d{4}-\d{2}-\d{2}$', version):
return False
return True
def is_session_expired(status_code: int) -> bool:
"""
Check if HTTP status code indicates session expiration.
Per MCP spec, HTTP 404 indicates the session has been terminated.
Args:
status_code: HTTP response status code
Returns:
True if session is expired, False otherwise
"""
return status_code == 404
def extract_session_id(headers: Dict[str, str]) -> Optional[str]:
"""
Extract session ID from HTTP response headers.
The header name is case-insensitive per HTTP spec.
Args:
headers: HTTP response headers dictionary
Returns:
Session ID if present, None otherwise
"""
# Check various case combinations
for key in headers:
if key.lower() == 'mcp-session-id':
return headers[key]
return None
class SessionManager:
"""
Manages MCP session state for HTTP transports.
This class handles:
- Session ID storage and validation
- Protocol version negotiation
- Session lifecycle (active, expired, terminated)
- Header generation for requests
Example:
```python
manager = SessionManager()
# After receiving InitializeResult with session ID
manager.set_session_id("abc123")
# Get headers for subsequent requests
headers = manager.get_headers()
# {'Mcp-Session-Id': 'abc123', 'Mcp-Protocol-Version': '2025-03-26'}
```
"""
def __init__(self, protocol_version: Optional[str] = None):
"""
Initialize session manager.
Args:
protocol_version: Initial protocol version (defaults to 2025-03-26)
"""
self._session_id: Optional[str] = None
self._protocol_version: str = protocol_version or DEFAULT_PROTOCOL_VERSION
self._is_expired: bool = False
self._metadata: Dict[str, Any] = {}
@property
def session_id(self) -> Optional[str]:
"""Get current session ID."""
return self._session_id
@property
def protocol_version(self) -> str:
"""Get current protocol version."""
return self._protocol_version
@property
def is_active(self) -> bool:
"""Check if session is active (has ID and not expired)."""
return self._session_id is not None and not self._is_expired
@property
def is_expired(self) -> bool:
"""Check if session has expired."""
return self._is_expired
def set_session_id(self, session_id: str) -> None:
"""
Set the session ID.
Args:
session_id: The session ID from server
Raises:
ValueError: If session ID is invalid
"""
if not is_valid_session_id(session_id):
raise ValueError("Invalid session ID: must contain only visible ASCII (0x21-0x7E)")
self._session_id = session_id
self._is_expired = False
def set_protocol_version(self, version: str) -> None:
"""
Set the negotiated protocol version.
Args:
version: Protocol version string
Raises:
ValueError: If version format is invalid
"""
if not is_valid_protocol_version(version):
raise ValueError(f"Invalid protocol version format: {version}")
self._protocol_version = version
def clear(self) -> None:
"""Clear the session (logout/disconnect)."""
self._session_id = None
self._is_expired = False
self._metadata = {}
def mark_expired(self) -> None:
"""Mark the session as expired (received HTTP 404)."""
self._is_expired = True
def get_headers(self) -> Dict[str, str]:
"""
Get HTTP headers for MCP requests.
Returns:
Dictionary of headers to include in requests
"""
headers = {
'Mcp-Protocol-Version': self._protocol_version
}
if self._session_id:
headers['Mcp-Session-Id'] = self._session_id
return headers
def update_from_response(self, headers: Dict[str, str], status_code: int = 200) -> None:
"""
Update session state from HTTP response.
Args:
headers: Response headers
status_code: HTTP status code
"""
# Check for session expiration
if is_session_expired(status_code):
self.mark_expired()
return
# Extract and update session ID if present
new_session_id = extract_session_id(headers)
if new_session_id:
self.set_session_id(new_session_id)
def set_metadata(self, key: str, value: Any) -> None:
"""Store arbitrary metadata with the session."""
self._metadata[key] = value
def get_metadata(self, key: str, default: Any = None) -> Any:
"""Retrieve metadata from the session."""
return self._metadata.get(key, default)
class ResumabilityManager:
"""
Manages SSE stream resumability for MCP transports.
This class handles:
- Event ID tracking per stream
- Last-Event-ID header generation
- Retry delay handling
Per MCP spec, event IDs should be assigned per-stream to act as
a cursor within that particular stream.
"""
def __init__(self):
"""Initialize resumability manager."""
self._last_event_ids: Dict[str, str] = {} # stream_id -> last_event_id
self._retry_delay: int = 3000 # Default retry delay in ms
def set_last_event_id(self, stream_id: str, event_id: str) -> None:
"""
Store the last received event ID for a stream.
Args:
stream_id: Identifier for the SSE stream
event_id: The event ID from the SSE event
"""
self._last_event_ids[stream_id] = event_id
def get_last_event_id(self, stream_id: str) -> Optional[str]:
"""
Get the last received event ID for a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
Last event ID if available, None otherwise
"""
return self._last_event_ids.get(stream_id)
def set_retry_delay(self, delay_ms: int) -> None:
"""
Set the retry delay from SSE retry field.
Args:
delay_ms: Delay in milliseconds
"""
self._retry_delay = delay_ms
def get_retry_delay(self) -> int:
"""Get current retry delay in milliseconds."""
return self._retry_delay
def get_resume_headers(self, stream_id: str) -> Dict[str, str]:
"""
Get headers for resuming a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
Headers dict with Last-Event-ID if available
"""
headers = {}
last_event_id = self.get_last_event_id(stream_id)
if last_event_id:
headers['Last-Event-ID'] = last_event_id
return headers
def clear_stream(self, stream_id: str) -> None:
"""Clear tracking for a specific stream."""
self._last_event_ids.pop(stream_id, None)
def clear_all(self) -> None:
"""Clear all stream tracking."""
self._last_event_ids.clear()
def parse_sse_event(event_text: str) -> Dict[str, Any]:
"""
Parse an SSE event into its components.
SSE events have the format:
id: <event-id>
retry: <milliseconds>
data: <data-line>
data: <data-line-2>
(blank line ends event)
Args:
event_text: Raw SSE event text
Returns:
Dictionary with 'id', 'retry', 'data', 'event' fields as available
"""
result: Dict[str, Any] = {}
data_lines = []
for line in event_text.strip().split('\n'):
if line.startswith('id:'):
result['id'] = line[3:].strip()
elif line.startswith('retry:'):
try:
result['retry'] = int(line[6:].strip())
except ValueError:
pass
elif line.startswith('data:'):
data_lines.append(line[5:].strip())
elif line.startswith('event:'):
result['event'] = line[6:].strip()
# Join multiple data lines with newlines
if data_lines:
result['data'] = '\n'.join(data_lines)
return result
class EventStore:
"""
Optional event store for message redelivery support.
This class stores events per stream to enable replay when
clients reconnect with Last-Event-ID header.
Note: This is an in-memory store suitable for single-server
deployments. For multi-node deployments, use an external
store (Redis, database, etc.).
"""
def __init__(self, max_events: int = 1000):
"""
Initialize event store.
Args:
max_events: Maximum events to store per stream
"""
self.max_events = max_events
self._streams: Dict[str, list] = {} # stream_id -> [(event_id, data), ...]
def store(self, stream_id: str, event_id: str, data: Any) -> None:
"""
Store an event for potential replay.
Args:
stream_id: Identifier for the SSE stream
event_id: The event ID
data: The event data
"""
if stream_id not in self._streams:
self._streams[stream_id] = []
self._streams[stream_id].append((event_id, data))
# Trim to max_events
if len(self._streams[stream_id]) > self.max_events:
self._streams[stream_id] = self._streams[stream_id][-self.max_events:]
def get_events_after(self, stream_id: str, after_event_id: str) -> list:
"""
Get events after a specific event ID.
Args:
stream_id: Identifier for the SSE stream
after_event_id: The Last-Event-ID from client
Returns:
List of event data after the specified event
"""
if stream_id not in self._streams:
return []
events = self._streams[stream_id]
# Find the index of the after_event_id
start_index = None
for i, (event_id, _) in enumerate(events):
if event_id == after_event_id:
start_index = i + 1
break
if start_index is None:
return []
return [data for _, data in events[start_index:]]
def get_all_events(self, stream_id: str) -> list:
"""
Get all stored events for a stream.
Args:
stream_id: Identifier for the SSE stream
Returns:
List of all event data
"""
if stream_id not in self._streams:
return []
return [data for _, data in self._streams[stream_id]]
def clear_stream(self, stream_id: str) -> None:
"""Clear events for a specific stream."""
self._streams.pop(stream_id, None)
def clear_all(self) -> None:
"""Clear all stored events."""
self._streams.clear()
@@ -0,0 +1,298 @@
"""
MCP Transport Abstraction Layer.
This module provides a common interface for all MCP transports,
enabling easy switching between transport types and custom
transport implementations.
Transport types:
- stdio: Standard input/output (subprocess)
- sse: Server-Sent Events (legacy HTTP+SSE)
- http_stream: Streamable HTTP (current standard)
- websocket: WebSocket (SEP-1288)
"""
import re
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, List, Type
from dataclasses import dataclass, field
@dataclass
class TransportConfig:
"""
Configuration for MCP transports.
Attributes:
timeout: Connection/operation timeout in seconds
debug: Enable debug logging
retry_count: Number of retry attempts
retry_delay: Base delay between retries in seconds
auth_token: Optional authentication token
headers: Additional headers for HTTP transports
"""
timeout: int = 60
debug: bool = False
retry_count: int = 3
retry_delay: float = 1.0
auth_token: Optional[str] = None
headers: Dict[str, str] = field(default_factory=dict)
class BaseTransport(ABC):
"""
Abstract base class for MCP transports.
All MCP transports must implement this interface to ensure
consistent behavior across transport types.
The transport is responsible for:
- Establishing and maintaining connections
- Sending JSON-RPC messages
- Receiving JSON-RPC messages
- Proper cleanup on close
"""
@abstractmethod
async def connect(self) -> None:
"""
Establish connection to the MCP server.
Raises:
ConnectionError: If connection fails
"""
pass
@abstractmethod
async def send(self, message: Dict[str, Any]) -> None:
"""
Send a JSON-RPC message to the server.
Args:
message: JSON-RPC message dictionary
Raises:
RuntimeError: If not connected
"""
pass
@abstractmethod
async def receive(self) -> Dict[str, Any]:
"""
Receive a JSON-RPC message from the server.
Returns:
JSON-RPC message dictionary
Raises:
RuntimeError: If not connected
"""
pass
@abstractmethod
async def close(self) -> None:
"""
Close the connection and cleanup resources.
"""
pass
@property
@abstractmethod
def is_connected(self) -> bool:
"""Check if transport is currently connected."""
pass
async def __aenter__(self):
"""Async context manager entry."""
await self.connect()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit."""
await self.close()
def get_transport_type(target: str) -> str:
"""
Determine the appropriate transport type for a target.
Args:
target: URL or command string
Returns:
Transport type string: "stdio", "sse", "http_stream", or "websocket"
"""
if not isinstance(target, str):
return "stdio"
# WebSocket URLs
if re.match(r'^wss?://', target, re.IGNORECASE):
return "websocket"
# HTTP URLs
if re.match(r'^https?://', target, re.IGNORECASE):
# Legacy SSE endpoint
if target.endswith('/sse'):
return "sse"
# Default to Streamable HTTP
return "http_stream"
# Everything else is stdio (command)
return "stdio"
class TransportRegistry:
"""
Registry for MCP transport implementations.
This allows registration of custom transports and provides
a central place to look up transport classes by name.
"""
def __init__(self):
"""Initialize empty registry."""
self._transports: Dict[str, Type[BaseTransport]] = {}
def register(self, name: str, transport_class: Type[BaseTransport]) -> None:
"""
Register a transport class.
Args:
name: Transport name (e.g., "websocket")
transport_class: Transport class (must extend BaseTransport)
"""
if not issubclass(transport_class, BaseTransport):
raise TypeError(f"{transport_class} must extend BaseTransport")
self._transports[name] = transport_class
def get(self, name: str) -> Optional[Type[BaseTransport]]:
"""
Get a registered transport class.
Args:
name: Transport name
Returns:
Transport class or None if not found
"""
return self._transports.get(name)
def list_transports(self) -> List[str]:
"""
List all registered transport names.
Returns:
List of transport names
"""
return list(self._transports.keys())
def unregister(self, name: str) -> None:
"""
Unregister a transport.
Args:
name: Transport name to remove
"""
self._transports.pop(name, None)
# Placeholder transport classes for registry
# These are minimal implementations to satisfy the registry
class StdioTransportPlaceholder(BaseTransport):
"""Placeholder for stdio transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class SSETransportPlaceholder(BaseTransport):
"""Placeholder for SSE transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class HTTPStreamTransportPlaceholder(BaseTransport):
"""Placeholder for HTTP Stream transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
class WebSocketTransportPlaceholder(BaseTransport):
"""Placeholder for WebSocket transport."""
async def connect(self): pass
async def send(self, message): pass
async def receive(self): return {}
async def close(self): pass
@property
def is_connected(self): return False
# Default registry instance
_default_registry: Optional[TransportRegistry] = None
def get_default_registry() -> TransportRegistry:
"""
Get the default transport registry with built-in transports.
Returns:
TransportRegistry with default transports registered
"""
global _default_registry
if _default_registry is None:
_default_registry = TransportRegistry()
# Register default transports
_default_registry.register("stdio", StdioTransportPlaceholder)
_default_registry.register("sse", SSETransportPlaceholder)
_default_registry.register("http_stream", HTTPStreamTransportPlaceholder)
_default_registry.register("websocket", WebSocketTransportPlaceholder)
return _default_registry
def create_transport(
target: str,
config: Optional[TransportConfig] = None,
registry: Optional[TransportRegistry] = None
) -> BaseTransport:
"""
Create a transport instance for the given target.
Args:
target: URL or command string
config: Optional transport configuration
registry: Optional custom registry (uses default if None)
Returns:
Transport instance
Raises:
ValueError: If no suitable transport found
"""
if registry is None:
registry = get_default_registry()
if config is None:
config = TransportConfig()
transport_type = get_transport_type(target)
transport_class = registry.get(transport_type)
if transport_class is None:
raise ValueError(f"No transport registered for type: {transport_type}")
return transport_class()
@@ -0,0 +1,596 @@
"""
WebSocket Transport for MCP (Model Context Protocol).
This module provides WebSocket-based transport for connecting to MCP servers.
WebSocket transport offers:
- Long-lived bidirectional communication
- Better cloud provider support (hibernatable connections)
- Lower protocol overhead compared to HTTP
- Session persistence across network interruptions
This is an optional transport - the 'websockets' package is only imported
when ws:// or wss:// URLs are used, ensuring no performance impact on
existing stdio/SSE/HTTP transports.
Based on SEP-1288: WebSocket Transport for MCP (in review).
Protocol Revision: 2025-11-25
"""
import asyncio
import json
import logging
import threading
import inspect
import re
from typing import Any, Dict, Optional, List
logger = logging.getLogger("mcp-websocket")
def is_websocket_url(url: str) -> bool:
"""
Check if a URL is a WebSocket URL.
Args:
url: The URL string to check
Returns:
True if the URL starts with ws:// or wss://, False otherwise
"""
if not isinstance(url, str):
return False
return bool(re.match(r'^wss?://', url, re.IGNORECASE))
def calculate_backoff(attempt: int, base_delay: float = 1.0, max_delay: float = 60.0) -> float:
"""
Calculate exponential backoff delay for reconnection.
Args:
attempt: The retry attempt number (0-indexed)
base_delay: Base delay in seconds
max_delay: Maximum delay cap in seconds
Returns:
Delay in seconds for this attempt
"""
delay = base_delay * (2 ** attempt)
return min(delay, max_delay)
# Global event loop for WebSocket operations (similar to other MCP transports)
_event_loop = None
_loop_lock = threading.Lock()
def get_event_loop():
"""Get or create a global event loop for WebSocket operations."""
global _event_loop
with _loop_lock:
if _event_loop is None or _event_loop.is_closed():
_event_loop = asyncio.new_event_loop()
return _event_loop
class WebSocketTransport:
"""
Low-level WebSocket transport for MCP communication.
This class handles the WebSocket connection and JSON-RPC message framing.
It supports:
- Secure (wss://) and insecure (ws://) connections
- Session ID management
- Authentication via subprotocol or headers
- Automatic reconnection with exponential backoff
"""
def __init__(
self,
url: str,
session_id: Optional[str] = None,
auth_token: Optional[str] = None,
max_retries: int = 3,
retry_delay: float = 1.0,
timeout: float = 30.0,
ping_interval: float = 30.0,
ping_timeout: float = 10.0
):
"""
Initialize WebSocket transport.
Args:
url: WebSocket URL (ws:// or wss://)
session_id: Optional MCP session ID
auth_token: Optional authentication token
max_retries: Maximum reconnection attempts
retry_delay: Base delay between retries (exponential backoff)
timeout: Connection timeout in seconds
ping_interval: Interval between ping messages (keepalive)
ping_timeout: Timeout for ping response
"""
self.url = url
self.session_id = session_id
self.auth_token = auth_token
self.max_retries = max_retries
self.retry_delay = retry_delay
self.timeout = timeout
self.ping_interval = ping_interval
self.ping_timeout = ping_timeout
self._ws = None
self._closed = False
self._message_queue: asyncio.Queue = asyncio.Queue()
self._receive_task = None
async def __aenter__(self):
"""Async context manager entry - establish connection."""
await self.connect()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
"""Async context manager exit - close connection."""
await self.close()
async def connect(self):
"""Establish WebSocket connection."""
try:
# Lazy import websockets
import websockets
except ImportError:
raise ImportError(
"websockets package is required for WebSocket transport. "
"Install it with: pip install websockets"
)
# Build connection options
extra_headers = {}
subprotocols = []
# Add session ID if present
if self.session_id:
extra_headers['Mcp-Session-Id'] = self.session_id
# Add auth token via subprotocol (SEP-1288 recommendation)
if self.auth_token:
# Smuggle auth token in subprotocol list
# Format: "mcp-auth-<base64-encoded-token>"
import base64
encoded_token = base64.b64encode(self.auth_token.encode()).decode()
subprotocols.append(f"mcp-auth-{encoded_token}")
# Always include MCP subprotocol
subprotocols.append("mcp")
# Connect with retry logic
last_error = None
for attempt in range(self.max_retries + 1):
try:
self._ws = await asyncio.wait_for(
websockets.connect(
self.url,
extra_headers=extra_headers if extra_headers else None,
subprotocols=subprotocols if subprotocols else None,
ping_interval=self.ping_interval,
ping_timeout=self.ping_timeout,
close_timeout=10
),
timeout=self.timeout
)
logger.debug(f"WebSocket connected to {self.url}")
# Start background receive task
self._receive_task = asyncio.create_task(self._receive_loop())
return
except Exception as e:
last_error = e
if attempt < self.max_retries:
delay = calculate_backoff(attempt, self.retry_delay)
logger.warning(f"WebSocket connection failed, retrying in {delay}s: {e}")
await asyncio.sleep(delay)
else:
raise ConnectionError(f"Failed to connect after {self.max_retries + 1} attempts: {last_error}")
async def _receive_loop(self):
"""Background task to receive messages and queue them."""
try:
async for message in self._ws:
try:
data = json.loads(message)
await self._message_queue.put(data)
except json.JSONDecodeError as e:
logger.error(f"Failed to parse WebSocket message: {e}")
except Exception as e:
if not self._closed:
logger.error(f"WebSocket receive error: {e}")
await self._message_queue.put({"error": str(e)})
async def send(self, message: Dict[str, Any]):
"""
Send a JSON-RPC message over WebSocket.
Args:
message: JSON-RPC message dictionary
"""
if self._ws is None:
raise RuntimeError("WebSocket not connected")
data = json.dumps(message)
await self._ws.send(data)
logger.debug(f"Sent: {message.get('method', message.get('id', 'response'))}")
async def receive(self) -> Dict[str, Any]:
"""
Receive a JSON-RPC message from WebSocket.
Returns:
Parsed JSON-RPC message dictionary
"""
message = await self._message_queue.get()
if "error" in message and len(message) == 1:
raise ConnectionError(message["error"])
return message
async def close(self):
"""Close WebSocket connection."""
self._closed = True
if self._receive_task:
self._receive_task.cancel()
try:
await self._receive_task
except asyncio.CancelledError:
pass
if self._ws:
await self._ws.close()
self._ws = None
logger.debug("WebSocket connection closed")
class WebSocketMCPTool:
"""
A wrapper for an MCP tool that can be used with praisonaiagents.
This class wraps a tool from a WebSocket MCP server and makes it
compatible with the Agent class tool interface.
"""
def __init__(
self,
name: str,
description: str,
session: Any, # ClientSession from mcp package
input_schema: Optional[Dict[str, Any]] = None,
timeout: int = 60
):
"""
Initialize WebSocket MCP tool wrapper.
Args:
name: Tool name
description: Tool description
session: MCP ClientSession for making calls
input_schema: JSON Schema for tool input
timeout: Timeout for tool calls in seconds
"""
self.name = name
self.__name__ = name
self.__qualname__ = name
self.__doc__ = description
self.description = description
self.session = session
self.input_schema = input_schema or {}
self.timeout = timeout
# Build function signature from input schema
self.__signature__ = self._build_signature()
def _build_signature(self) -> inspect.Signature:
"""Build function signature from input schema."""
params = []
if self.input_schema and 'properties' in self.input_schema:
required = self.input_schema.get('required', [])
for param_name, prop_schema in self.input_schema['properties'].items():
# Determine type annotation
prop_type = prop_schema.get('type', 'string') if isinstance(prop_schema, dict) else 'string'
annotation = self._json_type_to_python(prop_type)
# Determine default value
default = inspect.Parameter.empty if param_name in required else None
params.append(
inspect.Parameter(
name=param_name,
kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
default=default,
annotation=annotation
)
)
return inspect.Signature(params)
def _json_type_to_python(self, json_type: str) -> type:
"""Convert JSON Schema type to Python type."""
type_map = {
'string': str,
'integer': int,
'number': float,
'boolean': bool,
'array': list,
'object': dict
}
return type_map.get(json_type, Any)
def __call__(self, **kwargs):
"""
Call the tool synchronously.
This wraps the async call for synchronous usage.
"""
logger.debug(f"Tool {self.name} called with args: {kwargs}")
loop = get_event_loop()
future = asyncio.run_coroutine_threadsafe(self._async_call(**kwargs), loop)
try:
return future.result(timeout=self.timeout)
except Exception as e:
logger.error(f"Error calling tool {self.name}: {e}")
return f"Error: {str(e)}"
async def _async_call(self, **kwargs):
"""Call the tool asynchronously."""
logger.debug(f"Async calling tool {self.name} with args: {kwargs}")
try:
result = await self.session.call_tool(self.name, kwargs)
# Extract text from result
if hasattr(result, 'content') and result.content:
if hasattr(result.content[0], 'text'):
return result.content[0].text
return str(result.content[0])
return str(result)
except Exception as e:
logger.error(f"Error in _async_call for {self.name}: {e}")
raise
def _fix_array_schemas(self, schema: Dict[str, Any]) -> Dict[str, Any]:
"""
Fix array schemas by adding missing 'items' attribute.
OpenAI requires array types to specify item types.
"""
if not isinstance(schema, dict):
return schema
fixed_schema = schema.copy()
# Fix array types at current level
if fixed_schema.get("type") == "array" and "items" not in fixed_schema:
fixed_schema["items"] = {"type": "string"}
# Recursively fix nested schemas
if "properties" in fixed_schema:
fixed_properties = {}
for prop_name, prop_schema in fixed_schema["properties"].items():
fixed_properties[prop_name] = self._fix_array_schemas(prop_schema)
fixed_schema["properties"] = fixed_properties
if "items" in fixed_schema:
fixed_schema["items"] = self._fix_array_schemas(fixed_schema["items"])
return fixed_schema
def to_openai_tool(self) -> Dict[str, Any]:
"""Convert the tool to OpenAI function calling format."""
fixed_schema = self._fix_array_schemas(self.input_schema)
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": fixed_schema
}
}
class WebSocketMCPClient:
"""
High-level client for connecting to MCP servers over WebSocket.
This client handles:
- Connection establishment and lifecycle
- Tool discovery and wrapping
- Session management
- Automatic reconnection
Example:
```python
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPClient
client = WebSocketMCPClient("ws://localhost:8080/mcp")
for tool in client.tools:
print(tool.name)
```
"""
def __init__(
self,
server_url: str,
debug: bool = False,
timeout: int = 60,
auth_token: Optional[str] = None,
options: Optional[Dict[str, Any]] = None
):
"""
Initialize WebSocket MCP client.
Args:
server_url: WebSocket URL (ws:// or wss://)
debug: Enable debug logging
timeout: Timeout for operations in seconds
auth_token: Optional authentication token
options: Additional configuration options
"""
self.server_url = server_url
self.debug = debug
self.timeout = timeout
self.auth_token = auth_token
self.options = options or {}
self.session = None
self.transport = None
self.tools: List[WebSocketMCPTool] = []
self._closed = False
# Configure logging
if debug:
logger.setLevel(logging.DEBUG)
else:
logger.setLevel(logging.WARNING)
# Initialize connection
self._initialize()
def _initialize(self):
"""Initialize the connection and discover tools."""
loop = get_event_loop()
# Start event loop in background thread if not running
def run_event_loop():
asyncio.set_event_loop(loop)
loop.run_forever()
self.loop_thread = threading.Thread(target=run_event_loop, daemon=True)
self.loop_thread.start()
# Run initialization
future = asyncio.run_coroutine_threadsafe(self._async_initialize(), loop)
self.tools = future.result(timeout=self.timeout)
async def _async_initialize(self):
"""Asynchronously initialize connection and discover tools."""
logger.debug(f"Connecting to MCP server at {self.server_url}")
# Lazy import mcp package
try:
from mcp import ClientSession
except ImportError:
raise ImportError(
"mcp package is required. Install it with: pip install mcp"
)
# Create transport
self.transport = WebSocketTransport(
self.server_url,
auth_token=self.auth_token,
timeout=self.timeout
)
await self.transport.connect()
# Create read/write streams for ClientSession
# The ClientSession expects async iterators for read/write
read_stream = self._create_read_stream()
write_stream = self._create_write_stream()
# Create and initialize session
self._session_context = ClientSession(read_stream, write_stream)
self.session = await self._session_context.__aenter__()
await self.session.initialize()
# Discover tools
logger.debug("Listing tools...")
response = await self.session.list_tools()
tools_data = response.tools
logger.debug(f"Found {len(tools_data)} tools: {[t.name for t in tools_data]}")
# Create tool wrappers
tools = []
for tool in tools_data:
input_schema = tool.inputSchema if hasattr(tool, 'inputSchema') else None
wrapper = WebSocketMCPTool(
name=tool.name,
description=tool.description if hasattr(tool, 'description') else f"Call the {tool.name} tool",
session=self.session,
input_schema=input_schema,
timeout=self.timeout
)
tools.append(wrapper)
return tools
def _create_read_stream(self):
"""Create async iterator for reading from transport."""
async def read():
while not self._closed:
try:
message = await self.transport.receive()
yield message
except Exception as e:
if not self._closed:
logger.error(f"Read error: {e}")
break
return read()
def _create_write_stream(self):
"""Create write function for transport."""
async def write(message):
if hasattr(message, 'to_dict'):
message = message.to_dict()
await self.transport.send(message)
return write
def __iter__(self):
"""Return iterator over tools."""
return iter(self.tools)
def to_openai_tools(self) -> List[Dict[str, Any]]:
"""Convert all tools to OpenAI format."""
return [tool.to_openai_tool() for tool in self.tools]
async def aclose(self):
"""Async cleanup."""
if self._closed:
return
self._closed = True
try:
if hasattr(self, '_session_context') and self._session_context:
await self._session_context.__aexit__(None, None, None)
except Exception:
pass
try:
if self.transport:
await self.transport.close()
except Exception:
pass
def close(self):
"""Synchronous cleanup."""
if self._closed:
return
try:
loop = get_event_loop()
if not loop.is_closed():
future = asyncio.run_coroutine_threadsafe(self.aclose(), loop)
future.result(timeout=5)
except Exception:
self._closed = True
def __del__(self):
"""Cleanup on garbage collection."""
try:
if not self._closed:
self.close()
except Exception:
pass
+3 -2
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "praisonaiagents"
version = "0.0.188"
version = "0.0.189"
description = "Praison AI agents for completing complex tasks with Self Reflection Agents"
requires-python = ">=3.10"
authors = [
@@ -22,7 +22,8 @@ dependencies = [
mcp = [
"mcp>=1.6.0",
"fastapi>=0.115.0",
"uvicorn>=0.34.0"
"uvicorn>=0.34.0",
"websockets>=12.0"
]
memory = [
@@ -0,0 +1,122 @@
"""Tests for MCP Backward Compatibility - TDD approach.
These tests ensure backward compatibility with older MCP servers
and the deprecated HTTP+SSE transport from protocol version 2024-11-05.
Backward compatibility features:
- Legacy SSE endpoint support (/sse suffix)
- Fallback detection for older servers
- Automatic transport negotiation
"""
import pytest
class TestLegacySSESupport:
"""Test legacy SSE transport support."""
def test_sse_url_detection(self):
"""Test that /sse URLs are detected as legacy SSE."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/sse") == "sse"
assert get_transport_type("https://example.com/api/sse") == "sse"
def test_non_sse_http_uses_streamable(self):
"""Test that non-/sse HTTP URLs use Streamable HTTP."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/mcp") == "http_stream"
assert get_transport_type("https://example.com/api") == "http_stream"
class TestTransportNegotiation:
"""Test automatic transport negotiation."""
def test_detect_streamable_http_support(self):
"""Test detection of Streamable HTTP support."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
# Successful POST to MCP endpoint indicates Streamable HTTP
result = detect_transport_support(
response_status=200,
content_type="application/json"
)
assert result == "http_stream"
def test_detect_legacy_sse_fallback(self):
"""Test fallback to legacy SSE on specific error codes."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
# 400, 404, 405 indicate need for legacy SSE
assert detect_transport_support(response_status=400) == "sse"
assert detect_transport_support(response_status=404) == "sse"
assert detect_transport_support(response_status=405) == "sse"
def test_detect_sse_from_event_stream(self):
"""Test detection of SSE from content type."""
from praisonaiagents.mcp.mcp_compat import detect_transport_support
result = detect_transport_support(
response_status=200,
content_type="text/event-stream"
)
# SSE response indicates legacy transport
assert result in ["sse", "http_stream"]
class TestProtocolVersionCompatibility:
"""Test protocol version compatibility."""
def test_supports_2024_11_05(self):
"""Test support for protocol version 2024-11-05."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2024-11-05") is True
def test_supports_2025_03_26(self):
"""Test support for protocol version 2025-03-26."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2025-03-26") is True
def test_supports_2025_11_25(self):
"""Test support for protocol version 2025-11-25."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version("2025-11-25") is True
class TestMCPClassBackwardCompat:
"""Test MCP class backward compatibility."""
def test_mcp_accepts_sse_url(self):
"""Test MCP class accepts /sse URLs."""
# This is a design test - actual connection would require server
from praisonaiagents.mcp.mcp_transport import get_transport_type
# Verify URL would be routed to SSE transport
assert get_transport_type("http://localhost:8080/sse") == "sse"
def test_mcp_accepts_command_string(self):
"""Test MCP class accepts command strings (stdio)."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("python server.py") == "stdio"
assert get_transport_type("npx @mcp/server") == "stdio"
def test_mcp_accepts_http_url(self):
"""Test MCP class accepts HTTP URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/mcp") == "http_stream"
def test_mcp_accepts_websocket_url(self):
"""Test MCP class accepts WebSocket URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("ws://localhost:8080/mcp") == "websocket"
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,227 @@
"""Tests for MCP Resumability - TDD approach.
These tests define the expected behavior of SSE stream resumability
per MCP Protocol Revision 2025-11-25.
Resumability features:
- SSE event ID tracking
- Last-Event-ID header support
- Message replay on reconnection
- Retry field handling
"""
import pytest
class TestResumabilityManager:
"""Test ResumabilityManager class."""
def test_manager_creation(self):
"""Test ResumabilityManager can be created."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
assert manager.get_retry_delay() == 3000 # Default 3 seconds
def test_set_and_get_last_event_id(self):
"""Test setting and getting last event ID for a stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
# Initially no event ID
assert manager.get_last_event_id("stream-1") is None
# Set event ID
manager.set_last_event_id("stream-1", "event-123")
assert manager.get_last_event_id("stream-1") == "event-123"
def test_multiple_streams(self):
"""Test tracking multiple streams independently."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
assert manager.get_last_event_id("stream-1") == "event-a"
assert manager.get_last_event_id("stream-2") == "event-b"
def test_update_event_id(self):
"""Test updating event ID for a stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-1")
manager.set_last_event_id("stream-1", "event-2")
assert manager.get_last_event_id("stream-1") == "event-2"
def test_set_retry_delay(self):
"""Test setting retry delay from SSE retry field."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_retry_delay(5000)
assert manager.get_retry_delay() == 5000
def test_get_resume_headers_with_event_id(self):
"""Test getting resume headers when event ID is available."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-123")
headers = manager.get_resume_headers("stream-1")
assert headers == {'Last-Event-ID': 'event-123'}
def test_get_resume_headers_without_event_id(self):
"""Test getting resume headers when no event ID is stored."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
headers = manager.get_resume_headers("stream-1")
assert headers == {}
def test_clear_stream(self):
"""Test clearing tracking for a specific stream."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
manager.clear_stream("stream-1")
assert manager.get_last_event_id("stream-1") is None
assert manager.get_last_event_id("stream-2") == "event-b"
def test_clear_all(self):
"""Test clearing all stream tracking."""
from praisonaiagents.mcp.mcp_session import ResumabilityManager
manager = ResumabilityManager()
manager.set_last_event_id("stream-1", "event-a")
manager.set_last_event_id("stream-2", "event-b")
manager.clear_all()
assert manager.get_last_event_id("stream-1") is None
assert manager.get_last_event_id("stream-2") is None
class TestSSEEventParsing:
"""Test SSE event parsing utilities."""
def test_parse_sse_event_with_id(self):
"""Test parsing SSE event with ID field."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "id: event-123\ndata: {\"test\": 1}\n\n"
event = parse_sse_event(event_text)
assert event['id'] == 'event-123'
assert event['data'] == '{"test": 1}'
def test_parse_sse_event_with_retry(self):
"""Test parsing SSE event with retry field."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "retry: 5000\ndata: {}\n\n"
event = parse_sse_event(event_text)
assert event['retry'] == 5000
def test_parse_sse_event_data_only(self):
"""Test parsing SSE event with only data."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "data: {\"message\": \"hello\"}\n\n"
event = parse_sse_event(event_text)
assert event['data'] == '{"message": "hello"}'
assert event.get('id') is None
def test_parse_sse_event_multiline_data(self):
"""Test parsing SSE event with multiline data."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "data: line1\ndata: line2\n\n"
event = parse_sse_event(event_text)
assert event['data'] == 'line1\nline2'
def test_parse_sse_event_empty_data(self):
"""Test parsing SSE event with empty data (for priming)."""
from praisonaiagents.mcp.mcp_session import parse_sse_event
event_text = "id: prime-123\ndata: \n\n"
event = parse_sse_event(event_text)
assert event['id'] == 'prime-123'
assert event['data'] == ''
class TestEventStore:
"""Test optional event store for message redelivery."""
def test_event_store_creation(self):
"""Test EventStore can be created."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=100)
assert store.max_events == 100
def test_store_and_retrieve_event(self):
"""Test storing and retrieving events."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore()
store.store("stream-1", "event-1", {"data": "test1"})
store.store("stream-1", "event-2", {"data": "test2"})
events = store.get_events_after("stream-1", "event-1")
assert len(events) == 1
assert events[0] == {"data": "test2"}
def test_store_respects_max_events(self):
"""Test that store respects max_events limit."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore(max_events=3)
for i in range(5):
store.store("stream-1", f"event-{i}", {"data": i})
# Should only have last 3 events
all_events = store.get_all_events("stream-1")
assert len(all_events) <= 3
def test_store_per_stream_isolation(self):
"""Test that events are isolated per stream."""
from praisonaiagents.mcp.mcp_session import EventStore
store = EventStore()
store.store("stream-1", "event-a", {"data": "a"})
store.store("stream-2", "event-b", {"data": "b"})
events_1 = store.get_all_events("stream-1")
events_2 = store.get_all_events("stream-2")
assert len(events_1) == 1
assert len(events_2) == 1
assert events_1[0] == {"data": "a"}
assert events_2[0] == {"data": "b"}
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,224 @@
"""Tests for MCP Security Features - TDD approach.
These tests define the expected behavior of security features
per MCP Protocol Revision 2025-11-25.
Security features:
- Origin header validation (DNS rebinding prevention)
- Localhost binding recommendation
- Authentication header support
- Secure session ID handling
"""
import pytest
class TestOriginValidation:
"""Test Origin header validation for DNS rebinding prevention."""
def test_valid_localhost_origins(self):
"""Test that localhost origins are accepted."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
assert is_valid_origin("http://localhost:8080", allowed_origins=["localhost"]) is True
assert is_valid_origin("http://127.0.0.1:3000", allowed_origins=["127.0.0.1"]) is True
assert is_valid_origin("http://localhost", allowed_origins=["localhost"]) is True
def test_valid_custom_origins(self):
"""Test that configured custom origins are accepted."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
allowed = ["example.com", "api.example.com"]
assert is_valid_origin("https://example.com", allowed_origins=allowed) is True
assert is_valid_origin("https://api.example.com/path", allowed_origins=allowed) is True
def test_invalid_origins_rejected(self):
"""Test that invalid origins are rejected."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
allowed = ["localhost", "example.com"]
assert is_valid_origin("https://evil.com", allowed_origins=allowed) is False
assert is_valid_origin("https://example.com.evil.com", allowed_origins=allowed) is False
def test_missing_origin_handling(self):
"""Test handling of missing Origin header."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
# Missing origin should be rejected by default
assert is_valid_origin(None, allowed_origins=["localhost"]) is False
assert is_valid_origin("", allowed_origins=["localhost"]) is False
def test_allow_missing_origin_option(self):
"""Test option to allow missing Origin header."""
from praisonaiagents.mcp.mcp_security import is_valid_origin
# Can optionally allow missing origin for non-browser clients
assert is_valid_origin(None, allowed_origins=["localhost"], allow_missing=True) is True
class TestLocalhostBinding:
"""Test localhost binding recommendations."""
def test_is_localhost_address(self):
"""Test detection of localhost addresses."""
from praisonaiagents.mcp.mcp_security import is_localhost_address
assert is_localhost_address("127.0.0.1") is True
assert is_localhost_address("localhost") is True
assert is_localhost_address("::1") is True # IPv6 localhost
assert is_localhost_address("0.0.0.0") is False # All interfaces
assert is_localhost_address("192.168.1.1") is False
def test_should_bind_localhost_only(self):
"""Test recommendation for localhost-only binding."""
from praisonaiagents.mcp.mcp_security import should_bind_localhost_only
# Local servers should bind to localhost only
assert should_bind_localhost_only(is_local=True) is True
# Remote/production servers may bind to all interfaces
assert should_bind_localhost_only(is_local=False) is False
class TestAuthenticationHeaders:
"""Test authentication header support."""
def test_create_bearer_auth_header(self):
"""Test creating Bearer authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
header = create_auth_header("my-secret-token", auth_type="bearer")
assert header == {"Authorization": "Bearer my-secret-token"}
def test_create_basic_auth_header(self):
"""Test creating Basic authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
import base64
header = create_auth_header("user:pass", auth_type="basic")
expected = base64.b64encode(b"user:pass").decode()
assert header == {"Authorization": f"Basic {expected}"}
def test_create_custom_auth_header(self):
"""Test creating custom authentication header."""
from praisonaiagents.mcp.mcp_security import create_auth_header
header = create_auth_header("custom-value", auth_type="custom", header_name="X-API-Key")
assert header == {"X-API-Key": "custom-value"}
def test_validate_auth_header(self):
"""Test validating authentication header presence."""
from praisonaiagents.mcp.mcp_security import validate_auth_header
headers_with_auth = {"Authorization": "Bearer token"}
headers_without_auth = {"Content-Type": "application/json"}
assert validate_auth_header(headers_with_auth) is True
assert validate_auth_header(headers_without_auth) is False
class TestSecureSessionHandling:
"""Test secure session ID handling."""
def test_generate_secure_session_id(self):
"""Test generating cryptographically secure session ID."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
session_id = generate_secure_session_id()
# Should be non-empty
assert len(session_id) > 0
# Should be unique
session_id2 = generate_secure_session_id()
assert session_id != session_id2
# Should contain only valid characters (visible ASCII)
for char in session_id:
assert 0x21 <= ord(char) <= 0x7E
def test_session_id_length(self):
"""Test session ID has sufficient length for security."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
session_id = generate_secure_session_id()
# Should be at least 32 characters for security
assert len(session_id) >= 32
def test_session_id_entropy(self):
"""Test session ID has sufficient entropy."""
from praisonaiagents.mcp.mcp_security import generate_secure_session_id
# Generate multiple IDs and check they're all different
ids = [generate_secure_session_id() for _ in range(100)]
unique_ids = set(ids)
assert len(unique_ids) == 100 # All should be unique
class TestDNSRebindingPrevention:
"""Test DNS rebinding attack prevention."""
def test_extract_origin_host(self):
"""Test extracting host from Origin header."""
from praisonaiagents.mcp.mcp_security import extract_origin_host
assert extract_origin_host("http://localhost:8080") == "localhost"
assert extract_origin_host("https://example.com:443/path") == "example.com"
assert extract_origin_host("http://127.0.0.1") == "127.0.0.1"
def test_detect_dns_rebinding_attempt(self):
"""Test detection of potential DNS rebinding attempts."""
from praisonaiagents.mcp.mcp_security import is_potential_dns_rebinding
# External origin trying to access localhost server
assert is_potential_dns_rebinding(
origin="https://evil.com",
server_host="localhost"
) is True
# Same origin is safe
assert is_potential_dns_rebinding(
origin="http://localhost:8080",
server_host="localhost"
) is False
class TestSecurityConfig:
"""Test security configuration."""
def test_security_config_defaults(self):
"""Test default security configuration."""
from praisonaiagents.mcp.mcp_security import SecurityConfig
config = SecurityConfig()
# Secure defaults
assert config.validate_origin is True
assert config.allow_missing_origin is False
assert "localhost" in config.allowed_origins
assert "127.0.0.1" in config.allowed_origins
def test_security_config_custom(self):
"""Test custom security configuration."""
from praisonaiagents.mcp.mcp_security import SecurityConfig
config = SecurityConfig(
allowed_origins=["example.com"],
allow_missing_origin=True,
require_auth=True
)
assert "example.com" in config.allowed_origins
assert config.allow_missing_origin is True
assert config.require_auth is True
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,242 @@
"""Tests for MCP Session Management - TDD approach.
These tests define the expected behavior of session management for
Streamable HTTP transport per MCP Protocol Revision 2025-11-25.
Session management features:
- MCP-Session-Id header handling
- Session persistence across requests
- Session expiration handling (HTTP 404)
- Session termination via HTTP DELETE
- Session ID validation (visible ASCII 0x21-0x7E)
"""
import pytest
class TestSessionIdValidation:
"""Test session ID validation per MCP spec."""
def test_valid_session_id_ascii(self):
"""Test that valid ASCII session IDs are accepted."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
# Valid session IDs (visible ASCII 0x21-0x7E)
assert is_valid_session_id("abc123") is True
assert is_valid_session_id("session-id-12345") is True
assert is_valid_session_id("a1b2c3d4e5f6") is True
assert is_valid_session_id("!@#$%^&*()") is True # Special chars in range
assert is_valid_session_id("UUID-1234-5678-90ab") is True
def test_invalid_session_id_with_spaces(self):
"""Test that session IDs with spaces are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session id") is False
assert is_valid_session_id(" leading") is False
assert is_valid_session_id("trailing ") is False
def test_invalid_session_id_with_control_chars(self):
"""Test that session IDs with control characters are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session\x00id") is False # NULL
assert is_valid_session_id("session\nid") is False # Newline
assert is_valid_session_id("session\tid") is False # Tab
def test_invalid_session_id_empty(self):
"""Test that empty session IDs are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("") is False
assert is_valid_session_id(None) is False
def test_invalid_session_id_non_ascii(self):
"""Test that non-ASCII characters are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_session_id
assert is_valid_session_id("session\x7fid") is False # DEL char
assert is_valid_session_id("session\x80id") is False # Extended ASCII
class TestSessionManager:
"""Test SessionManager class for managing MCP sessions."""
def test_session_manager_creation(self):
"""Test SessionManager can be created."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
assert manager.session_id is None
assert manager.is_active is False
def test_session_manager_set_session_id(self):
"""Test setting session ID."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session-123")
assert manager.session_id == "test-session-123"
assert manager.is_active is True
def test_session_manager_rejects_invalid_id(self):
"""Test that invalid session IDs are rejected."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
with pytest.raises(ValueError, match="Invalid session ID"):
manager.set_session_id("invalid session") # Contains space
def test_session_manager_clear_session(self):
"""Test clearing session."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session-123")
manager.clear()
assert manager.session_id is None
assert manager.is_active is False
def test_session_manager_get_headers(self):
"""Test getting headers with session ID."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
# No session - no header
headers = manager.get_headers()
assert 'Mcp-Session-Id' not in headers
# With session - include header
manager.set_session_id("test-session-123")
headers = manager.get_headers()
assert headers['Mcp-Session-Id'] == "test-session-123"
class TestSessionFromResponse:
"""Test extracting session ID from HTTP response headers."""
def test_extract_session_id_from_headers(self):
"""Test extracting session ID from response headers."""
from praisonaiagents.mcp.mcp_session import extract_session_id
headers = {'Mcp-Session-Id': 'new-session-456'}
session_id = extract_session_id(headers)
assert session_id == 'new-session-456'
def test_extract_session_id_case_insensitive(self):
"""Test that header name is case-insensitive."""
from praisonaiagents.mcp.mcp_session import extract_session_id
# Various case combinations
assert extract_session_id({'mcp-session-id': 'test'}) == 'test'
assert extract_session_id({'MCP-SESSION-ID': 'test'}) == 'test'
assert extract_session_id({'Mcp-Session-Id': 'test'}) == 'test'
def test_extract_session_id_missing(self):
"""Test when session ID header is missing."""
from praisonaiagents.mcp.mcp_session import extract_session_id
headers = {'Content-Type': 'application/json'}
session_id = extract_session_id(headers)
assert session_id is None
class TestSessionExpiration:
"""Test handling of session expiration (HTTP 404)."""
def test_is_session_expired_404(self):
"""Test that HTTP 404 indicates session expiration."""
from praisonaiagents.mcp.mcp_session import is_session_expired
assert is_session_expired(404) is True
assert is_session_expired(200) is False
assert is_session_expired(400) is False
assert is_session_expired(500) is False
def test_session_manager_marks_expired(self):
"""Test that session manager can mark session as expired."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_session_id("test-session")
manager.mark_expired()
assert manager.is_expired is True
assert manager.is_active is False
class TestProtocolVersion:
"""Test MCP-Protocol-Version header handling."""
def test_default_protocol_version(self):
"""Test default protocol version is 2025-03-26 for compatibility."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
headers = manager.get_headers()
# Default version for backward compatibility
assert headers.get('Mcp-Protocol-Version') == '2025-03-26'
def test_set_protocol_version(self):
"""Test setting negotiated protocol version."""
from praisonaiagents.mcp.mcp_session import SessionManager
manager = SessionManager()
manager.set_protocol_version('2025-11-25')
headers = manager.get_headers()
assert headers['Mcp-Protocol-Version'] == '2025-11-25'
def test_valid_protocol_versions(self):
"""Test that valid protocol versions are accepted."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version('2025-03-26') is True
assert is_valid_protocol_version('2025-06-18') is True
assert is_valid_protocol_version('2025-11-25') is True
assert is_valid_protocol_version('2024-11-05') is True
def test_invalid_protocol_versions(self):
"""Test that invalid protocol versions are rejected."""
from praisonaiagents.mcp.mcp_session import is_valid_protocol_version
assert is_valid_protocol_version('invalid') is False
assert is_valid_protocol_version('2025') is False
assert is_valid_protocol_version('') is False
class TestHTTPStreamSessionIntegration:
"""Test session management integration with HTTPStreamTransport."""
def test_transport_uses_session_manager(self):
"""Test that HTTPStreamTransport uses SessionManager."""
# This will be implemented when we update HTTPStreamTransport
pass
def test_transport_handles_session_from_init_response(self):
"""Test that transport extracts session ID from InitializeResult."""
pass
def test_transport_includes_session_in_requests(self):
"""Test that transport includes session ID in subsequent requests."""
pass
def test_transport_handles_session_expiration(self):
"""Test that transport handles HTTP 404 session expiration."""
pass
def test_transport_can_terminate_session(self):
"""Test that transport can terminate session via HTTP DELETE."""
pass
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,167 @@
"""Tests for MCP Transport Abstraction Layer - TDD approach.
These tests define the expected behavior of the transport abstraction
layer that provides a common interface for all MCP transports.
Transport abstraction features:
- BaseTransport abstract class
- Common interface (connect, send, receive, close)
- Transport factory for automatic selection
"""
import pytest
from abc import ABC
class TestBaseTransport:
"""Test BaseTransport abstract class."""
def test_base_transport_is_abstract(self):
"""Test that BaseTransport cannot be instantiated directly."""
from praisonaiagents.mcp.mcp_transport import BaseTransport
with pytest.raises(TypeError):
BaseTransport()
def test_base_transport_defines_interface(self):
"""Test that BaseTransport defines required interface methods."""
from praisonaiagents.mcp.mcp_transport import BaseTransport
import inspect
# Check required abstract methods exist
assert hasattr(BaseTransport, 'connect')
assert hasattr(BaseTransport, 'send')
assert hasattr(BaseTransport, 'receive')
assert hasattr(BaseTransport, 'close')
assert hasattr(BaseTransport, 'is_connected')
def test_concrete_transport_must_implement_interface(self):
"""Test that concrete transports must implement all methods."""
from praisonaiagents.mcp.mcp_transport import BaseTransport
# Incomplete implementation should fail
class IncompleteTransport(BaseTransport):
pass
with pytest.raises(TypeError):
IncompleteTransport()
class TestTransportFactory:
"""Test transport factory for automatic selection."""
def test_factory_selects_stdio_for_command(self):
"""Test factory selects stdio transport for command strings."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("python /path/to/server.py") == "stdio"
assert get_transport_type("npx @modelcontextprotocol/server") == "stdio"
def test_factory_selects_websocket_for_ws_url(self):
"""Test factory selects WebSocket transport for ws:// URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("ws://localhost:8080/mcp") == "websocket"
assert get_transport_type("wss://example.com/mcp") == "websocket"
def test_factory_selects_sse_for_sse_url(self):
"""Test factory selects SSE transport for /sse URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/sse") == "sse"
assert get_transport_type("https://example.com/sse") == "sse"
def test_factory_selects_http_stream_for_http_url(self):
"""Test factory selects HTTP Stream transport for HTTP URLs."""
from praisonaiagents.mcp.mcp_transport import get_transport_type
assert get_transport_type("http://localhost:8080/mcp") == "http_stream"
assert get_transport_type("https://example.com/mcp") == "http_stream"
class TestTransportConfig:
"""Test transport configuration."""
def test_transport_config_creation(self):
"""Test TransportConfig can be created."""
from praisonaiagents.mcp.mcp_transport import TransportConfig
config = TransportConfig(
timeout=30,
debug=True,
retry_count=3
)
assert config.timeout == 30
assert config.debug is True
assert config.retry_count == 3
def test_transport_config_defaults(self):
"""Test TransportConfig has sensible defaults."""
from praisonaiagents.mcp.mcp_transport import TransportConfig
config = TransportConfig()
assert config.timeout == 60 # Default timeout
assert config.debug is False
assert config.retry_count == 3
class TestTransportRegistry:
"""Test transport registration mechanism."""
def test_register_custom_transport(self):
"""Test registering a custom transport."""
from praisonaiagents.mcp.mcp_transport import (
TransportRegistry, BaseTransport
)
class CustomTransport(BaseTransport):
async def connect(self): pass
async def send(self, message): pass
async def receive(self): pass
async def close(self): pass
@property
def is_connected(self): return False
registry = TransportRegistry()
registry.register("custom", CustomTransport)
assert "custom" in registry.list_transports()
def test_get_registered_transport(self):
"""Test getting a registered transport class."""
from praisonaiagents.mcp.mcp_transport import (
TransportRegistry, BaseTransport
)
class MyTransport(BaseTransport):
async def connect(self): pass
async def send(self, message): pass
async def receive(self): pass
async def close(self): pass
@property
def is_connected(self): return False
registry = TransportRegistry()
registry.register("my_transport", MyTransport)
transport_class = registry.get("my_transport")
assert transport_class is MyTransport
def test_default_transports_registered(self):
"""Test that default transports are pre-registered."""
from praisonaiagents.mcp.mcp_transport import get_default_registry
registry = get_default_registry()
transports = registry.list_transports()
# Default transports should be available
assert "stdio" in transports
assert "sse" in transports
assert "http_stream" in transports
assert "websocket" in transports
if __name__ == "__main__":
pytest.main([__file__, "-v"])
@@ -0,0 +1,296 @@
"""Tests for MCP WebSocket Transport - TDD approach.
These tests define the expected behavior of the WebSocket transport module.
WebSocket transport is an optional feature that requires the 'websockets' package.
Performance considerations:
- WebSocket is lazy-loaded only when ws:// or wss:// URLs are used
- No impact on existing stdio/SSE/HTTP transports
- Uses asyncio for non-blocking I/O
"""
import pytest
from unittest.mock import Mock
# Skip all tests if websockets is not installed
pytest.importorskip("websockets", reason="websockets package not installed")
class TestWebSocketURLDetection:
"""Test URL detection for WebSocket transport."""
def test_detects_ws_url(self):
"""Test that ws:// URLs are detected as WebSocket."""
from praisonaiagents.mcp.mcp_websocket import is_websocket_url
assert is_websocket_url("ws://localhost:8080") is True
assert is_websocket_url("ws://example.com/mcp") is True
assert is_websocket_url("ws://127.0.0.1:3000/ws") is True
def test_detects_wss_url(self):
"""Test that wss:// URLs are detected as WebSocket (secure)."""
from praisonaiagents.mcp.mcp_websocket import is_websocket_url
assert is_websocket_url("wss://localhost:8080") is True
assert is_websocket_url("wss://example.com/mcp") is True
assert is_websocket_url("wss://secure.server.com:443/ws") is True
def test_rejects_http_urls(self):
"""Test that HTTP URLs are not detected as WebSocket."""
from praisonaiagents.mcp.mcp_websocket import is_websocket_url
assert is_websocket_url("http://localhost:8080") is False
assert is_websocket_url("https://example.com/mcp") is False
assert is_websocket_url("http://localhost:8080/sse") is False
def test_rejects_non_url_strings(self):
"""Test that non-URL strings are not detected as WebSocket."""
from praisonaiagents.mcp.mcp_websocket import is_websocket_url
assert is_websocket_url("python") is False
assert is_websocket_url("/path/to/script.py") is False
assert is_websocket_url("npx @modelcontextprotocol/server") is False
class TestWebSocketMCPTool:
"""Test WebSocketMCPTool wrapper class."""
def test_tool_has_required_attributes(self):
"""Test that tool wrapper has all required attributes for Agent."""
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPTool
mock_session = Mock()
tool = WebSocketMCPTool(
name="test_tool",
description="A test tool",
session=mock_session,
input_schema={"type": "object", "properties": {"query": {"type": "string"}}}
)
assert tool.name == "test_tool"
assert tool.__name__ == "test_tool"
assert tool.__qualname__ == "test_tool"
assert tool.__doc__ == "A test tool"
assert tool.description == "A test tool"
assert hasattr(tool, '__signature__')
def test_tool_generates_correct_signature(self):
"""Test that tool generates correct function signature from schema."""
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPTool
import inspect
mock_session = Mock()
tool = WebSocketMCPTool(
name="search",
description="Search tool",
session=mock_session,
input_schema={
"type": "object",
"properties": {
"query": {"type": "string"},
"max_results": {"type": "integer"}
},
"required": ["query"]
}
)
sig = tool.__signature__
params = list(sig.parameters.keys())
assert "query" in params
assert "max_results" in params
assert sig.parameters["query"].default == inspect.Parameter.empty
assert sig.parameters["max_results"].default is None # Optional
def test_tool_converts_to_openai_format(self):
"""Test conversion to OpenAI tool format."""
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPTool
mock_session = Mock()
tool = WebSocketMCPTool(
name="search",
description="Search the web",
session=mock_session,
input_schema={
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"]
}
)
openai_tool = tool.to_openai_tool()
assert openai_tool["type"] == "function"
assert openai_tool["function"]["name"] == "search"
assert openai_tool["function"]["description"] == "Search the web"
assert "parameters" in openai_tool["function"]
def test_tool_fixes_array_schemas(self):
"""Test that array schemas without items are fixed for OpenAI compatibility."""
from praisonaiagents.mcp.mcp_websocket import WebSocketMCPTool
mock_session = Mock()
tool = WebSocketMCPTool(
name="process",
description="Process items",
session=mock_session,
input_schema={
"type": "object",
"properties": {
"items": {"type": "array"} # Missing 'items' attribute
}
}
)
openai_tool = tool.to_openai_tool()
items_schema = openai_tool["function"]["parameters"]["properties"]["items"]
assert "items" in items_schema # Should be added
assert items_schema["items"]["type"] == "string" # Default
class TestWebSocketTransport:
"""Test WebSocketTransport low-level implementation."""
def test_transport_handles_session_id(self):
"""Test that transport handles MCP session ID."""
from praisonaiagents.mcp.mcp_websocket import WebSocketTransport
transport = WebSocketTransport("ws://localhost:8080/mcp")
assert transport.session_id is None
transport.session_id = "test-session-123"
assert transport.session_id == "test-session-123"
def test_transport_initialization(self):
"""Test transport initializes with correct parameters."""
from praisonaiagents.mcp.mcp_websocket import WebSocketTransport
transport = WebSocketTransport(
"wss://example.com/mcp",
auth_token="secret",
max_retries=5,
timeout=30.0
)
assert transport.url == "wss://example.com/mcp"
assert transport.auth_token == "secret"
assert transport.max_retries == 5
assert transport.timeout == 30.0
def test_transport_send_requires_connection(self):
"""Test that send raises error when not connected."""
from praisonaiagents.mcp.mcp_websocket import WebSocketTransport
import asyncio
transport = WebSocketTransport("ws://localhost:8080/mcp")
with pytest.raises(RuntimeError, match="WebSocket not connected"):
asyncio.get_event_loop().run_until_complete(transport.send({"test": 1}))
class TestWebSocketReconnection:
"""Test WebSocket reconnection with exponential backoff."""
def test_exponential_backoff_calculation(self):
"""Test exponential backoff delay calculation."""
from praisonaiagents.mcp.mcp_websocket import calculate_backoff
# First retry: base delay
assert calculate_backoff(attempt=0, base_delay=1.0, max_delay=60.0) == 1.0
# Second retry: 2x
assert calculate_backoff(attempt=1, base_delay=1.0, max_delay=60.0) == 2.0
# Third retry: 4x
assert calculate_backoff(attempt=2, base_delay=1.0, max_delay=60.0) == 4.0
# Should cap at max_delay
assert calculate_backoff(attempt=10, base_delay=1.0, max_delay=60.0) == 60.0
class TestWebSocketMCPClient:
"""Test WebSocketMCPClient high-level interface."""
def test_client_initialization(self):
"""Test client can be initialized with WebSocket URL."""
# This test will be implemented after the basic structure is in place
pass
def test_client_lists_tools(self):
"""Test client can list available tools from server."""
pass
def test_client_calls_tool(self):
"""Test client can call a tool and get result."""
pass
def test_client_handles_timeout(self):
"""Test client handles timeout correctly."""
pass
def test_client_is_iterable(self):
"""Test client tools are iterable."""
pass
def test_client_converts_to_openai_tools(self):
"""Test client can convert all tools to OpenAI format."""
pass
class TestMCPClassWebSocketIntegration:
"""Test MCP class integration with WebSocket transport."""
def test_mcp_detects_ws_url(self):
"""Test MCP class detects ws:// URL and uses WebSocket transport."""
# Will be implemented after WebSocketMCPClient is ready
pass
def test_mcp_detects_wss_url(self):
"""Test MCP class detects wss:// URL and uses WebSocket transport."""
pass
def test_mcp_websocket_tools_work_with_agent(self):
"""Test WebSocket tools work correctly with Agent class."""
pass
class TestWebSocketKeepalive:
"""Test WebSocket keepalive/ping-pong handling."""
def test_ping_interval_configurable(self):
"""Test that ping interval is configurable."""
from praisonaiagents.mcp.mcp_websocket import WebSocketTransport
transport = WebSocketTransport(
"ws://localhost:8080/mcp",
ping_interval=60.0,
ping_timeout=20.0
)
assert transport.ping_interval == 60.0
assert transport.ping_timeout == 20.0
class TestWebSocketPerformance:
"""Test WebSocket transport performance characteristics."""
def test_lazy_import_websocket_module(self):
"""Test that mcp_websocket module can be imported."""
# Just verify the module imports without errors
from praisonaiagents.mcp.mcp_websocket import (
is_websocket_url,
calculate_backoff,
WebSocketTransport,
WebSocketMCPTool
)
# Verify functions exist
assert callable(is_websocket_url)
assert callable(calculate_backoff)
if __name__ == "__main__":
pytest.main([__file__, "-v"])
+4 -1
View File
@@ -2422,7 +2422,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.0.188"
version = "0.0.189"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },
@@ -2450,6 +2450,7 @@ all = [
{ name = "python-jose", extra = ["cryptography"] },
{ name = "python-multipart" },
{ name = "uvicorn" },
{ name = "websockets" },
]
api = [
{ name = "fastapi" },
@@ -2479,6 +2480,7 @@ mcp = [
{ name = "fastapi" },
{ name = "mcp" },
{ name = "uvicorn" },
{ name = "websockets" },
]
memory = [
{ name = "chromadb" },
@@ -2530,6 +2532,7 @@ requires-dist = [
{ name = "rich" },
{ name = "uvicorn", marker = "extra == 'api'", specifier = ">=0.34.0" },
{ name = "uvicorn", marker = "extra == 'mcp'", specifier = ">=0.34.0" },
{ name = "websockets", marker = "extra == 'mcp'", specifier = ">=12.0" },
]
provides-extras = ["mcp", "memory", "knowledge", "graph", "llm", "api", "telemetry", "mongodb", "auth", "all"]
+4
View File
@@ -64,6 +64,10 @@ PraisonAI is a production-ready Multi-AI Agents framework with self-reflection,
- 📋 **Planning Mode** (Plan before execution - Agent & Multi-Agent)
- 🔧 **Planning Tools** (Research with tools during planning)
- 🧠 **Planning Reasoning** (Chain-of-thought planning)
- 🔌 **MCP Transports** (stdio, Streamable HTTP, WebSocket, SSE - Protocol 2025-11-25)
- 🌐 **WebSocket MCP** (Real-time bidirectional connections with auto-reconnect)
- 🔐 **MCP Security** (Origin validation, DNS rebinding prevention, secure sessions)
- 🔄 **MCP Resumability** (SSE stream recovery via Last-Event-ID)
## Using Python Code
+2 -2
View File
@@ -3,8 +3,8 @@ class Praisonai < Formula
desc "AI tools for various AI applications"
homepage "https://github.com/MervinPraison/PraisonAI"
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.9.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.9.tar.gz | shasum -a 256`.split.first
url "https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.10.tar.gz"
sha256 `curl -sL https://github.com/MervinPraison/PraisonAI/archive/refs/tags/v2.3.10.tar.gz | shasum -a 256`.split.first
license "MIT"
depends_on "[email protected]"
+1 -1
View File
@@ -57,7 +57,7 @@ class CloudDeployer:
file.write("FROM python:3.11-slim\n")
file.write("WORKDIR /app\n")
file.write("COPY . .\n")
file.write("RUN pip install flask praisonai==2.3.9 gunicorn markdown\n")
file.write("RUN pip install flask praisonai==2.3.10 gunicorn markdown\n")
file.write("EXPOSE 8080\n")
file.write('CMD ["gunicorn", "-b", "0.0.0.0:8080", "api:app"]\n')
+1 -1
View File
@@ -1 +1 @@
__version__ = "2.3.9"
__version__ = "2.3.10"
+1 -1
View File
@@ -12,7 +12,7 @@ dependencies = [
"rich>=13.7",
"markdown>=3.5",
"pyparsing>=3.0.0",
"praisonaiagents>=0.0.188",
"praisonaiagents>=0.0.189",
"python-dotenv>=0.19.0",
"instructor>=1.3.3",
"PyYAML>=6.0",
+4 -4
View File
@@ -4197,7 +4197,7 @@ requires-dist = [
{ name = "praisonai-tools", marker = "extra == 'autogen'", specifier = ">=0.0.22" },
{ name = "praisonai-tools", marker = "extra == 'autogen-v4'", specifier = ">=0.0.22" },
{ name = "praisonai-tools", marker = "extra == 'crewai'", specifier = ">=0.0.22" },
{ name = "praisonaiagents", specifier = ">=0.0.188" },
{ name = "praisonaiagents", specifier = ">=0.0.189" },
{ name = "pyautogen", marker = "extra == 'autogen'", specifier = "==0.2.29" },
{ name = "pydantic", marker = "extra == 'chat'", specifier = "<=2.10.1" },
{ name = "pydantic", marker = "extra == 'code'", specifier = "<=2.10.1" },
@@ -4252,7 +4252,7 @@ wheels = [
[[package]]
name = "praisonaiagents"
version = "0.0.188"
version = "0.0.189"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "aiohttp" },
@@ -4261,9 +4261,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/76/b1/4ea337476bdb07ef504af6023e378c513a9c35414ff080a777221e70add5/praisonaiagents-0.0.188.tar.gz", hash = "sha256:ba498e6328e3a6184e3d99d2ae7a73473bd0aef747abe71e9a2c99d201e791f5", size = 438151, upload-time = "2025-12-16T08:17:22.931Z" }
sdist = { url = "https://files.pythonhosted.org/packages/62/ed/bd468ae67dea1489590fb306d50b02fbe2da5462745b2c8519a095b81fc1/praisonaiagents-0.0.189.tar.gz", hash = "sha256:a788b540b2b1f3418fb98013f1b2aba92d3010797d852507325735e2eb3fd944", size = 461258, upload-time = "2025-12-16T09:47:12.238Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/0d/3c/122ca9485b531a7307a44c1149f7b6e72b5a962b2244f8a8756f459f50e8/praisonaiagents-0.0.188-py3-none-any.whl", hash = "sha256:3ba651ca4e52de2b2f63101f265411a676f918ac308d4dbb95f2124cdc42ca6e", size = 420694, upload-time = "2025-12-16T08:17:20.78Z" },
{ url = "https://files.pythonhosted.org/packages/3a/70/2ea1d2a8baa02ae5a277616bbcc2f70c8a2a25a068d4cd5a8083eaf759cb/praisonaiagents-0.0.189-py3-none-any.whl", hash = "sha256:158693b6c68b9d23df64ee092648992222718279c47cac7b96e1007bf869ad8c", size = 437480, upload-time = "2025-12-16T09:47:09.846Z" },
]
[[package]]