wehub-resource-sync
457 lines
24 KiB
Python
457 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Interactive MCP Client for Agent-to-Agent Communication Tutorial
|
|
"""
|
|
|
|
import asyncio
|
|
import argparse
|
|
import logging
|
|
from typing import Dict, Any
|
|
|
|
from mcp import ClientSession
|
|
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
|
|
import mcp.types as types
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
|
|
from .utils import TokenManager, cast_input_value
|
|
|
|
console = Console()
|
|
|
|
|
|
def display_tools(tools):
|
|
"""Display available tools as a simple list."""
|
|
console.print("[bold]Available Tools:[/bold]")
|
|
for tool in tools:
|
|
tool_type = "🤖" if any(word in tool.name.lower() for word in ["agent", "travel", "research"]) else "🔧"
|
|
console.print(f" {tool_type} [cyan]{tool.name}[/cyan]")
|
|
console.print()
|
|
console.print("[dim]Type a tool name to run it with default arguments[/dim]")
|
|
|
|
|
|
def extract_text_content(result) -> str:
|
|
"""Extract text content from tool result."""
|
|
for content in result.content:
|
|
if hasattr(content, 'text'):
|
|
return content.text
|
|
return "No text content available"
|
|
|
|
|
|
async def execute_tool_with_resumption(session, command: str, args: dict, get_session_id, on_resumption_token_update, existing_tokens=None, token_manager=None):
|
|
"""Execute a tool with resumption support using send_request."""
|
|
current_session_id = get_session_id()
|
|
if not current_session_id:
|
|
raise RuntimeError("No session ID available - resumption requires a valid session")
|
|
|
|
session_id = current_session_id
|
|
|
|
# If we have an existing resumption token, pass it for resumption
|
|
if existing_tokens and existing_tokens.get("resumption_token"):
|
|
metadata = ClientMessageMetadata(
|
|
resumption_token=existing_tokens["resumption_token"],
|
|
)
|
|
else:
|
|
# Create enhanced callback that saves tool context immediately when token is received
|
|
def enhanced_callback(token: str):
|
|
# Since callback fires immediately with the actual resumption token,
|
|
# save everything needed for resumption right away
|
|
protocol_version = getattr(session, 'protocol_version', None)
|
|
if token_manager:
|
|
token_manager.save_tokens(session_id, token, protocol_version, command, args)
|
|
# Also call the original callback
|
|
return on_resumption_token_update(session_id, token, command, args)
|
|
|
|
metadata = ClientMessageMetadata(
|
|
on_resumption_token_update=enhanced_callback,
|
|
)
|
|
|
|
result = await session.send_request(
|
|
types.ClientRequest(
|
|
types.CallToolRequest(
|
|
method="tools/call",
|
|
params=types.CallToolRequestParams(
|
|
name=command,
|
|
arguments=args
|
|
),
|
|
)
|
|
),
|
|
types.CallToolResult,
|
|
metadata=metadata,
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
async def interactive_mode(server_url: str):
|
|
"""Run interactive mode with tool exploration."""
|
|
# Configure logging to suppress noisy SSE parsing errors
|
|
logging.getLogger('mcp.client.streamable_http').setLevel(logging.ERROR)
|
|
|
|
# Filter out specific SSE JSON parsing errors
|
|
class SSEFilter(logging.Filter):
|
|
def filter(self, record):
|
|
# Suppress "Error parsing SSE message" and JSON validation errors
|
|
if ("Error parsing SSE message" in record.getMessage() or
|
|
"ValidationError" in record.getMessage() or
|
|
"EOF while parsing" in record.getMessage()):
|
|
return False
|
|
return True
|
|
|
|
# Apply filter to relevant loggers
|
|
for logger_name in ['mcp.client.streamable_http', 'mcp', 'pydantic_core']:
|
|
logger = logging.getLogger(logger_name)
|
|
logger.addFilter(SSEFilter())
|
|
|
|
console.print(Panel("[bold cyan]🎮 Interactive MCP Client[/bold cyan]", expand=False))
|
|
console.print("[dim]Explore MCP server tools interactively[/dim]\n")
|
|
|
|
# Check for existing resumption tokens
|
|
token_manager = TokenManager()
|
|
existing_tokens: Dict[str, Any] | None = token_manager.load_tokens()
|
|
|
|
# Prepare headers for resumption if tokens exist
|
|
headers = {}
|
|
if existing_tokens:
|
|
headers[MCP_SESSION_ID_HEADER] = existing_tokens["session_id"]
|
|
if "protocol_version" in existing_tokens:
|
|
headers[MCP_PROTOCOL_VERSION_HEADER] = existing_tokens["protocol_version"]
|
|
console.print(f"[cyan]🔄 Found existing session, attempting resumption...[/cyan]")
|
|
|
|
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):
|
|
|
|
# Create message handler for real-time notifications
|
|
async def message_handler(message) -> None:
|
|
try:
|
|
if isinstance(message, types.ServerNotification):
|
|
if isinstance(message.root, types.LoggingMessageNotification):
|
|
console.print(f"📡 [dim]{message.root.params.data}[/dim]")
|
|
elif isinstance(message.root, types.ProgressNotification):
|
|
progress = message.root.params
|
|
console.print(f"🔄 [yellow]{progress.message} ({progress.progress}/{progress.total})[/yellow]")
|
|
elif isinstance(message.root, types.ResourceUpdatedNotification):
|
|
console.print(f"� [blue]Resource updated: {message.root.params.uri}[/blue]")
|
|
except Exception:
|
|
# Silently ignore message handler errors
|
|
pass
|
|
|
|
# Add sampling callback for research agent
|
|
async def sampling_callback(context, params):
|
|
try:
|
|
message_text = params.messages[0].content.text if params.messages else 'No message'
|
|
console.print(f"\n🧠 [bold cyan]Server requested sampling:[/bold cyan]")
|
|
console.print(f" [yellow]{message_text}[/yellow]")
|
|
|
|
# Mock response instead of prompting user
|
|
mock_response = "Based on current research, MCP has evolved significantly with new features like resumable streams, elicitation, and sampling capabilities, enabling sophisticated agent-to-agent communication patterns."
|
|
|
|
console.print(f"[dim]🤖 Auto-responding with mock data:[/dim]")
|
|
console.print(f" [green]{mock_response[:80]}...[/green]")
|
|
|
|
return types.CreateMessageResult(
|
|
role="assistant",
|
|
content=types.TextContent(
|
|
type="text",
|
|
text=mock_response
|
|
),
|
|
model="interactive-client",
|
|
stopReason="endTurn"
|
|
)
|
|
except Exception as e:
|
|
return types.ErrorData(code=-1, message=str(e))
|
|
|
|
# Add elicitation callback for travel agent
|
|
async def elicitation_callback(context, params):
|
|
try:
|
|
console.print(f"\n💬 [bold yellow]Server is asking for confirmation:[/bold yellow]")
|
|
console.print(f" [cyan]{params.message}[/cyan]")
|
|
|
|
# Get user's decision
|
|
while True:
|
|
response = console.input("\n[bold]Do you accept? (y/n/details): [/bold]").strip().lower()
|
|
|
|
if response in ['y', 'yes', 'accept']:
|
|
user_notes = console.input("[dim]Any additional notes (optional): [/dim]").strip()
|
|
console.print(f"[dim]✅ Sending acceptance with notes: '{user_notes or 'Confirmed by user'}'[/dim]")
|
|
return types.ElicitResult(
|
|
action="accept",
|
|
content={
|
|
"confirm": True,
|
|
"notes": user_notes if user_notes else "Confirmed by user"
|
|
}
|
|
)
|
|
elif response in ['n', 'no', 'decline']:
|
|
reason = console.input("[dim]Reason for declining (optional): [/dim]").strip()
|
|
return types.ElicitResult(
|
|
action="decline",
|
|
content={
|
|
"confirm": False,
|
|
"notes": reason if reason else "Declined by user"
|
|
}
|
|
)
|
|
elif response in ['d', 'details']:
|
|
console.print(f"[dim]Schema: {params.requestedSchema if hasattr(params, 'requestedSchema') else 'Not specified'}[/dim]")
|
|
else:
|
|
console.print("[red]Please enter 'y' (yes), 'n' (no), or 'd' (details)[/red]")
|
|
|
|
except Exception as e:
|
|
console.print(f"[red]Error in elicitation: {e}[/red]")
|
|
return types.ElicitResult(action="decline", content={"confirm": False, "notes": f"Error: {e}"})
|
|
|
|
# Add resumption token callback for long-running tools
|
|
async def on_resumption_token_update(session_id: str, resumption_token: str, tool_name: str, tool_args: dict):
|
|
"""Callback for when resumption token is updated during long-running operations."""
|
|
# Save the updated token with tool information
|
|
protocol_version = existing_tokens.get("protocol_version") if existing_tokens else None
|
|
token_manager.save_tokens(session_id, resumption_token, protocol_version, tool_name, tool_args)
|
|
|
|
async with ClientSession(
|
|
read_stream,
|
|
write_stream,
|
|
message_handler=message_handler,
|
|
sampling_callback=sampling_callback,
|
|
elicitation_callback=elicitation_callback
|
|
) as session:
|
|
|
|
# Handle resumption vs new session based on tokens
|
|
if existing_tokens and existing_tokens.get("resumption_token") and existing_tokens.get("last_tool") and existing_tokens.get("last_args"):
|
|
# RESUMPTION FLOW: We have everything needed for resumption
|
|
console.print(f"[cyan]🔄 Resuming existing session...[/cyan]")
|
|
console.print(f"[cyan] Session ID: {existing_tokens['session_id']}[/cyan]")
|
|
console.print(f"[cyan] Resuming tool: {existing_tokens['last_tool']}[/cyan]")
|
|
console.print(f"[green]✅ Skipping initialization, directly resuming...[/green]")
|
|
|
|
try:
|
|
# Directly resume the tool execution with cached data
|
|
result = await execute_tool_with_resumption(
|
|
session,
|
|
existing_tokens["last_tool"],
|
|
existing_tokens["last_args"],
|
|
get_session_id,
|
|
on_resumption_token_update,
|
|
existing_tokens, # Pass existing tokens for resumption
|
|
token_manager
|
|
)
|
|
console.print(f"[green]✅ Resumed tool completed successfully![/green]")
|
|
console.print(f"[green]Result:[/green]")
|
|
console.print(f" {extract_text_content(result)}")
|
|
|
|
# Task completed successfully, clean up tokens
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
|
|
# After successful resumption, continue with normal interactive flow
|
|
console.print("\n[cyan]📋 Loading available tools for continued interaction...[/cyan]")
|
|
tools_result = await session.list_tools()
|
|
tools = tools_result.tools
|
|
tools_dict = {tool.name: tool for tool in tools}
|
|
display_tools(tools)
|
|
|
|
except Exception as resume_error:
|
|
error_msg = str(resume_error)
|
|
if "ValidationError" in error_msg and "JSON" in error_msg:
|
|
console.print(f"[yellow]⚠️ Resumed tool completed with connection issues, but likely succeeded[/yellow]")
|
|
# Even if there was a connection issue, task likely completed
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
|
|
# Load tools for continued interaction
|
|
console.print("\n[cyan]📋 Loading available tools for continued interaction...[/cyan]")
|
|
tools_result = await session.list_tools()
|
|
tools = tools_result.tools
|
|
tools_dict = {tool.name: tool for tool in tools}
|
|
display_tools(tools)
|
|
else:
|
|
console.print(f"[red]❌ Resume failed: {resume_error}[/red]")
|
|
console.print(f"[yellow]Falling back to normal initialization...[/yellow]")
|
|
# Fallback to normal initialization
|
|
result = await session.initialize()
|
|
console.print(f"[green]✅ Connected to: {result.serverInfo.name}[/green]")
|
|
tools_result = await session.list_tools()
|
|
tools = tools_result.tools
|
|
tools_dict = {tool.name: tool for tool in tools}
|
|
display_tools(tools)
|
|
else:
|
|
# NORMAL FLOW: No resumption tokens, do full initialization
|
|
console.print("[cyan]🔌 Connecting to server...[/cyan]")
|
|
result = await session.initialize()
|
|
|
|
console.print(f"[green]✅ Connected to: {result.serverInfo.name}[/green]")
|
|
console.print(f"[green]📋 Protocol: {result.protocolVersion}[/green]")
|
|
|
|
# Save protocol version for future resumption
|
|
current_session_id = get_session_id()
|
|
if current_session_id:
|
|
console.print(f"[green]🆔 Session: {current_session_id}[/green]")
|
|
|
|
|
|
# Load and display available tools
|
|
console.print("\n[cyan]📋 Loading available tools...[/cyan]")
|
|
tools_result = await session.list_tools()
|
|
tools = tools_result.tools
|
|
tools_dict = {tool.name: tool for tool in tools}
|
|
|
|
display_tools(tools)
|
|
|
|
# Interactive command loop
|
|
console.print("\n[dim]Commands: [tool_name], 'list', 'help', 'clean-tokens', 'quit'[/dim]")
|
|
|
|
while True:
|
|
try:
|
|
command = console.input("\n[bold blue]> [/bold blue]").strip()
|
|
|
|
if command.lower() in ['quit', 'exit', 'q']:
|
|
console.print("[yellow]👋 Goodbye![/yellow]")
|
|
break
|
|
|
|
elif command.lower() in ['list', 'l']:
|
|
display_tools(tools)
|
|
|
|
elif command.lower() in ['help', 'h']:
|
|
console.print("[bold]Available Commands:[/bold]")
|
|
console.print(" [cyan]list[/cyan] or [cyan]l[/cyan] - Show available tools")
|
|
console.print(" [cyan][tool_name][/cyan] - Execute a tool")
|
|
console.print(" [cyan]help[/cyan] or [cyan]h[/cyan] - Show this help")
|
|
console.print(" [cyan]clean-tokens[/cyan] - Delete resumption tokens")
|
|
console.print(" [cyan]quit[/cyan] or [cyan]q[/cyan] - Exit")
|
|
console.print("\n[bold]Agent tools (try these!):[/bold]")
|
|
for tool in tools:
|
|
if "agent" in tool.name.lower():
|
|
console.print(f" 🤖 [yellow]{tool.name}[/yellow]")
|
|
|
|
elif command.lower() in ['clean-tokens', 'clean']:
|
|
if token_manager.delete_tokens():
|
|
console.print("[green]✅ Resumption tokens cleared[/green]")
|
|
else:
|
|
console.print("[yellow]No tokens to clear[/yellow]")
|
|
|
|
elif command in tools_dict:
|
|
tool = tools_dict[command]
|
|
|
|
# Collect arguments
|
|
args = {}
|
|
if tool.inputSchema and tool.inputSchema.get("properties"):
|
|
console.print(f"[dim]Tool '{command}' parameters:[/dim]")
|
|
for prop_name, prop_info in tool.inputSchema["properties"].items():
|
|
required = prop_name in tool.inputSchema.get("required", [])
|
|
default = prop_info.get("default", "")
|
|
desc = prop_info.get("description", "")
|
|
|
|
prompt = f" {prop_name}"
|
|
if desc:
|
|
prompt += f" ({desc})"
|
|
if default:
|
|
prompt += f" [default: {default}]"
|
|
if required:
|
|
prompt += " [required]"
|
|
prompt += ": "
|
|
|
|
value = console.input(prompt).strip()
|
|
if value:
|
|
# Cast the input value to the correct type
|
|
args[prop_name] = cast_input_value(value, prop_info)
|
|
elif required and not default:
|
|
console.print(f"[red]❌ {prop_name} is required[/red]")
|
|
break
|
|
else:
|
|
# All required args collected, execute tool
|
|
console.print(f"[cyan]🔧 Executing {command}...[/cyan]")
|
|
|
|
try:
|
|
result = await execute_tool_with_resumption(
|
|
session, command, args, get_session_id, on_resumption_token_update, existing_tokens, token_manager
|
|
)
|
|
console.print(f"[green]✅ Result:[/green]")
|
|
console.print(f" {extract_text_content(result)}")
|
|
|
|
# Task completed successfully, clean up tokens
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
|
|
except Exception as tool_error:
|
|
error_msg = str(tool_error)
|
|
if "ValidationError" in error_msg and "JSON" in error_msg:
|
|
console.print(f"[yellow]⚠️ Tool completed with connection issues, but likely succeeded[/yellow]")
|
|
# Even if there was a connection issue, task likely completed
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
else:
|
|
console.print(f"[red]❌ Tool failed: {tool_error}[/red]")
|
|
else:
|
|
# No parameters needed, execute tool
|
|
console.print(f"[cyan]🔧 Executing {command}...[/cyan]")
|
|
|
|
try:
|
|
result = await execute_tool_with_resumption(
|
|
session, command, {}, get_session_id, on_resumption_token_update, existing_tokens, token_manager
|
|
)
|
|
console.print(f"[green]✅ Result:[/green]")
|
|
console.print(f" {extract_text_content(result)}")
|
|
|
|
# Task completed successfully, clean up tokens
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
|
|
except Exception as tool_error:
|
|
error_msg = str(tool_error)
|
|
if "ValidationError" in error_msg and "JSON" in error_msg:
|
|
console.print(f"[yellow]⚠️ Tool completed with connection issues, but likely succeeded[/yellow]")
|
|
# Even if there was a connection issue, task likely completed
|
|
if token_manager.delete_tokens():
|
|
console.print(f"[dim]🧹 Cleaned up resumption tokens[/dim]")
|
|
else:
|
|
console.print(f"[red]❌ Tool failed: {tool_error}[/red]")
|
|
|
|
else:
|
|
console.print(f"[red]❌ Unknown command: '{command}'[/red]")
|
|
console.print("[dim]Type 'list' to see available tools or 'help' for commands[/dim]")
|
|
|
|
except KeyboardInterrupt:
|
|
console.print("\n[yellow]👋 Interrupted![/yellow]")
|
|
break
|
|
|
|
except Exception as e:
|
|
console.print(f"[red]❌ Connection error: {e}[/red]")
|
|
|
|
|
|
async def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="Interactive MCP Client")
|
|
parser.add_argument("--url", default="http://127.0.0.1:8006/mcp", help="MCP server URL")
|
|
parser.add_argument("--clean-tokens", action="store_true", help="Delete existing resumption tokens and start fresh")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Handle token cleaning
|
|
if args.clean_tokens:
|
|
token_manager = TokenManager()
|
|
if token_manager.delete_tokens():
|
|
console.print("[green]✅ Resumption tokens cleared[/green]")
|
|
else:
|
|
console.print("[yellow]No tokens to clear[/yellow]")
|
|
return
|
|
|
|
await interactive_mode(args.url)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Configure logging to suppress noisy MCP warnings
|
|
logging.basicConfig(level=logging.WARNING)
|
|
|
|
# Suppress specific noisy loggers
|
|
logging.getLogger('mcp.client.streamable_http').setLevel(logging.ERROR)
|
|
logging.getLogger('httpx').setLevel(logging.WARNING)
|
|
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
console.print("\n[yellow]👋 Goodbye![/yellow]")
|
|
except Exception as e:
|
|
console.print(f"[red]💥 Fatal error: {e}[/red]")
|