mirror of
https://github.com/MervinPraison/PraisonAI.git
synced 2026-11-07 04:35:58 -10:00
Release v2.3.10
This commit is contained in:
1 parent
cf2a5a135f
commit
67dbd1007a
36 files changed
+3843
-62
No files matched your search
@@ -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
@@ -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
|
||||
|
||||
@@ -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]"
|
||||
|
||||
|
||||
@@ -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]" \
|
||||
|
||||
@@ -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()}")
|
||||
@@ -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")
|
||||
@@ -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.")
|
||||
@@ -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")
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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"])
|
||||
Generated
+4
-1
@@ -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"]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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]"
|
||||
|
||||
@@ -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 @@
|
||||
__version__ = "2.3.9"
|
||||
__version__ = "2.3.10"
|
||||
@@ -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",
|
||||
|
||||
Generated
+4
-4
@@ -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]]
|
||||
|
||||
Reference in new issue
Block a user