wehub-resource-sync
225 lines
9.0 KiB
Python
225 lines
9.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Minimal Resumable MCP Client
|
|
|
|
Simple client that demonstrates session resumption with the long_running_agent tool.
|
|
- If resumption token exists, resumes the session
|
|
- If no token exists, creates new session and saves token
|
|
- Clears token when task completes
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, Any, Optional
|
|
|
|
from mcp import ClientSession, types
|
|
from mcp.client.streamable_http import streamablehttp_client
|
|
from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, MCP_PROTOCOL_VERSION_HEADER
|
|
from mcp.shared.message import ClientMessageMetadata
|
|
from mcp.types import TextContent
|
|
|
|
# Configure logging to suppress noisy SSE warnings
|
|
logging.basicConfig(level=logging.WARNING)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Suppress specific noisy loggers
|
|
logging.getLogger('httpx').setLevel(logging.WARNING)
|
|
logging.getLogger('mcp.client.streamable_http').setLevel(logging.ERROR)
|
|
logging.getLogger('mcp').setLevel(logging.WARNING)
|
|
logging.getLogger('pydantic_core').setLevel(logging.ERROR)
|
|
|
|
# Filter out specific SSE JSON parsing errors
|
|
class SSEFilter(logging.Filter):
|
|
def filter(self, record):
|
|
return not (
|
|
"Error parsing SSE message" in record.getMessage() or
|
|
"Invalid JSON: EOF while parsing" in record.getMessage()
|
|
)
|
|
|
|
# Apply filter to relevant loggers
|
|
for logger_name in ['mcp.client.streamable_http', 'mcp', 'pydantic_core']:
|
|
log = logging.getLogger(logger_name)
|
|
log.addFilter(SSEFilter())
|
|
|
|
# Token file location
|
|
TOKEN_FILE = Path(".mcp_resumption_token.json")
|
|
|
|
|
|
def load_resumption_tokens() -> Optional[Dict[str, Any]]:
|
|
"""Load resumption tokens from file."""
|
|
if TOKEN_FILE.exists():
|
|
try:
|
|
with open(TOKEN_FILE, 'r') as f:
|
|
tokens = json.load(f)
|
|
logger.info(f"🔄 Loaded resumption tokens: session_id={tokens.get('session_id')}")
|
|
return tokens
|
|
except Exception as e:
|
|
logger.error(f"Error loading tokens: {e}")
|
|
return None
|
|
|
|
|
|
def save_resumption_tokens(session_id: str, resumption_token: str, protocol_version: str = "2025-03-26"):
|
|
"""Save resumption tokens to file."""
|
|
tokens = {
|
|
"session_id": session_id,
|
|
"resumption_token": resumption_token,
|
|
"protocol_version": protocol_version,
|
|
"tool_name": "long_running_agent",
|
|
"tool_args": {}
|
|
}
|
|
try:
|
|
with open(TOKEN_FILE, 'w') as f:
|
|
json.dump(tokens, f, indent=2)
|
|
logger.info(f"💾 Saved resumption tokens: session_id={session_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error saving tokens: {e}")
|
|
|
|
|
|
def clear_resumption_tokens():
|
|
"""Clear resumption tokens file."""
|
|
if TOKEN_FILE.exists():
|
|
try:
|
|
TOKEN_FILE.unlink()
|
|
logger.info("🗑️ Cleared resumption tokens")
|
|
except Exception as e:
|
|
logger.error(f"Error clearing tokens: {e}")
|
|
|
|
|
|
async def run_long_running_task(server_url: str):
|
|
"""Run the long_running_agent with resumption support."""
|
|
|
|
# Check for existing tokens
|
|
existing_tokens = load_resumption_tokens()
|
|
|
|
# Prepare headers for resumption if tokens exist
|
|
headers = {}
|
|
if existing_tokens:
|
|
headers[MCP_SESSION_ID_HEADER] = existing_tokens["session_id"]
|
|
headers[MCP_PROTOCOL_VERSION_HEADER] = existing_tokens.get("protocol_version", "2025-03-26")
|
|
print("🔄 Resuming existing session...")
|
|
else:
|
|
print("🆕 Creating new session...")
|
|
|
|
try:
|
|
async with streamablehttp_client(
|
|
server_url,
|
|
headers=headers if headers else None,
|
|
terminate_on_close=False # Enable resumption
|
|
) as (read_stream, write_stream, get_session_id):
|
|
|
|
# Message handler for real-time notifications (standardized to use log messages)
|
|
async def message_handler(message) -> None:
|
|
try:
|
|
if isinstance(message, types.ServerNotification):
|
|
if isinstance(message.root, types.LoggingMessageNotification):
|
|
log_data = message.root.params
|
|
# Format log messages similar to progress notifications for consistency
|
|
print(f"📡 [{log_data.logger}] {log_data.data}")
|
|
elif isinstance(message.root, types.ProgressNotification):
|
|
# Keep support for progress notifications in case other tools use them
|
|
progress = message.root.params
|
|
print(f"📊 Progress: {progress.progress}/{progress.total} - {progress.message}")
|
|
elif isinstance(message.root, types.ResourceUpdatedNotification):
|
|
print(f"🔄 Resource updated: {message.root.params.uri}")
|
|
except Exception:
|
|
# Silently ignore message handler errors to avoid breaking the flow
|
|
pass
|
|
|
|
async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
|
|
|
|
# Only initialize if this is a new session (no existing tokens)
|
|
if not existing_tokens:
|
|
result = await session.initialize()
|
|
print(f"✅ Session initialized: {result.serverInfo.name}")
|
|
else:
|
|
print("✅ Using existing session (no initialization needed)")
|
|
|
|
# Set up resumption token callback (matching test signature)
|
|
async def on_resumption_token_update(token: str) -> None:
|
|
print(f"💾 Resumption token received: {token[:20]}...")
|
|
session_id = get_session_id()
|
|
if session_id:
|
|
protocol_version = getattr(session, 'negotiated_protocol_version', '2025-03-26')
|
|
save_resumption_tokens(session_id, token, protocol_version)
|
|
print(f"💾 Resumption token saved")
|
|
|
|
# Prepare metadata with resumption support
|
|
if existing_tokens and existing_tokens.get("resumption_token"):
|
|
# Resume existing task
|
|
print("🔄 Resuming long-running task...")
|
|
metadata = ClientMessageMetadata(
|
|
resumption_token=existing_tokens["resumption_token"]
|
|
)
|
|
else:
|
|
# Start new task
|
|
print("🚀 Starting long-running task...")
|
|
metadata = ClientMessageMetadata(
|
|
on_resumption_token_update=on_resumption_token_update
|
|
)
|
|
|
|
# Execute the long_running_agent tool
|
|
try:
|
|
print("metadata:", metadata)
|
|
result = await session.send_request(
|
|
types.ClientRequest(
|
|
types.CallToolRequest(
|
|
method="tools/call",
|
|
params=types.CallToolRequestParams(
|
|
name="long_running_agent",
|
|
arguments={}
|
|
),
|
|
)
|
|
),
|
|
types.CallToolResult,
|
|
metadata=metadata,
|
|
)
|
|
|
|
# Task completed successfully
|
|
content_text = "Task completed"
|
|
if result.content and len(result.content) > 0:
|
|
first_content = result.content[0]
|
|
if isinstance(first_content, TextContent):
|
|
content_text = first_content.text
|
|
|
|
print(f"✅ Task completed: {content_text}")
|
|
clear_resumption_tokens()
|
|
|
|
except Exception as e:
|
|
print(f"❌ Task failed: {e}")
|
|
# Keep tokens in case we want to retry
|
|
|
|
except Exception as e:
|
|
print(f"❌ Connection error: {e}")
|
|
|
|
|
|
async def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="Minimal Resumable MCP Client")
|
|
parser.add_argument("--url", default="http://127.0.0.1:8006/mcp", help="MCP server URL")
|
|
parser.add_argument("--clear-tokens", action="store_true", help="Clear resumption tokens and exit")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.clear_tokens:
|
|
clear_resumption_tokens()
|
|
print("🗑️ Resumption tokens cleared")
|
|
return
|
|
|
|
print("🤖 Minimal Resumable MCP Client")
|
|
print(f"🔗 Connecting to: {args.url}")
|
|
|
|
await run_long_running_task(args.url)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
print("\n👋 Interrupted by user")
|
|
except Exception as e:
|
|
print(f"💥 Fatal error: {e}")
|