191 lines
7.3 KiB
Python
191 lines
7.3 KiB
Python
"""Utilities shared by MCP clients for the TradingAgents project.
|
||
|
||
This module contains a single public helper, :func:`perform_handshake`, which
|
||
wraps the Model Context Protocol session bootstrapping sequence. The helper
|
||
adds a thin compatibility layer around the official Python SDK so we can supply
|
||
client metadata, negotiate capabilities, and emit useful debug logs without
|
||
duplicating this logic in every integration.
|
||
|
||
The implementation favours graceful degradation – if the optional ``mcp``
|
||
package is not installed, or if the runtime SDK version does not implement a
|
||
specific API entry point, we simply skip that portion of the handshake while
|
||
providing informative logging. This keeps the behaviour predictable in
|
||
development environments where the dependency may not be available yet.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import importlib.metadata
|
||
import logging
|
||
from dataclasses import asdict, is_dataclass
|
||
from typing import Any, Dict, Optional
|
||
|
||
try: # pragma: no cover - optional dependency during linting
|
||
from mcp.client.session import ClientSession
|
||
except ImportError: # pragma: no cover - surfaced at runtime with helpful error
|
||
ClientSession = Any # type: ignore[misc,assignment]
|
||
|
||
try: # pragma: no cover - optional dependency during linting
|
||
from mcp.types import ClientCapabilities, Implementation
|
||
except ImportError: # pragma: no cover - surfaced at runtime when missing
|
||
ClientCapabilities = None # type: ignore[misc,assignment]
|
||
Implementation = None # type: ignore[misc,assignment]
|
||
|
||
|
||
_DEFAULT_PROTOCOL_VERSION = "2025-06-18"
|
||
|
||
|
||
def emit_console(level: str, message: str) -> None:
|
||
"""Mirror log output to stdout when no logging handlers are configured."""
|
||
|
||
root_logger = logging.getLogger()
|
||
if not root_logger.handlers:
|
||
print(f"{level.upper()}: {message}")
|
||
|
||
|
||
def _detect_package_version() -> str:
|
||
"""Return the installed TradingAgents version, falling back to ``0.0.0``."""
|
||
|
||
try:
|
||
return importlib.metadata.version("tradingagents")
|
||
except importlib.metadata.PackageNotFoundError:
|
||
return "0.0.0"
|
||
|
||
|
||
async def perform_handshake(
|
||
session: "ClientSession",
|
||
*,
|
||
client_label: str,
|
||
logger: logging.Logger,
|
||
capabilities: Optional[Any] = None,
|
||
protocol_version: str = _DEFAULT_PROTOCOL_VERSION,
|
||
) -> Dict[str, Any]:
|
||
"""Execute the MCP initialization handshake and emit diagnostic metadata.
|
||
|
||
Parameters
|
||
----------
|
||
session
|
||
A connected :class:`mcp.client.session.ClientSession` instance.
|
||
client_label
|
||
Human friendly tag used in log messages to identify the integration.
|
||
logger
|
||
Logger used for status updates. ``perform_handshake`` only emits at
|
||
``INFO`` and ``DEBUG`` levels, so callers can opt in/out via standard
|
||
logging configuration.
|
||
capabilities
|
||
Optional per-client capability overrides. When omitted we fall back to
|
||
an empty :class:`ClientCapabilities` object if the SDK exposes one. If
|
||
the caller supplies a plain mapping, the helper forwards it unchanged to
|
||
the ``initialize`` call – the Python SDK accepts either Pydantic models
|
||
or raw dictionaries.
|
||
protocol_version
|
||
Requested MCP protocol version. Defaults to the latest spec revision we
|
||
target in this repository.
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
A dictionary with the initial handshake result. The structure matches
|
||
the underlying ``InitializeResult`` object but is normalised to basic
|
||
Python types so it can be safely logged or inspected by tests if
|
||
desired.
|
||
"""
|
||
|
||
if ClientSession is Any: # pragma: no cover - defensive runtime guard
|
||
raise RuntimeError("perform_handshake cannot run without the 'mcp' package installed.")
|
||
|
||
initialize_kwargs: Dict[str, Any] = {}
|
||
|
||
if protocol_version:
|
||
initialize_kwargs["protocol_version"] = protocol_version
|
||
|
||
if Implementation is not None:
|
||
initialize_kwargs["client_info"] = Implementation(
|
||
name=f"TradingAgents::{client_label}",
|
||
version=_detect_package_version(),
|
||
)
|
||
else:
|
||
initialize_kwargs["client_info"] = {
|
||
"name": f"TradingAgents::{client_label}",
|
||
"version": _detect_package_version(),
|
||
}
|
||
|
||
if capabilities is not None:
|
||
initialize_kwargs["capabilities"] = capabilities
|
||
elif ClientCapabilities is not None:
|
||
initialize_kwargs["capabilities"] = ClientCapabilities()
|
||
|
||
try:
|
||
result = await session.initialize(**initialize_kwargs)
|
||
except TypeError:
|
||
# Older SDK versions did not support keyword arguments. Retry using the
|
||
# minimal signature, while still surfacing the original failure in debug
|
||
# logs so developers understand why metadata was omitted.
|
||
logger.debug(
|
||
"MCP initialize signature did not accept kwargs for %s, falling back to defaults.",
|
||
client_label,
|
||
exc_info=True,
|
||
)
|
||
result = await session.initialize()
|
||
|
||
# Normalise the result into a dictionary for consistent downstream usage.
|
||
payload: Dict[str, Any]
|
||
if hasattr(result, "model_dump"):
|
||
payload = result.model_dump() # type: ignore[assignment]
|
||
elif is_dataclass(result):
|
||
payload = asdict(result) # type: ignore[arg-type]
|
||
elif isinstance(result, dict):
|
||
payload = dict(result)
|
||
else:
|
||
payload = {
|
||
"protocolVersion": getattr(result, "protocolVersion", None)
|
||
or getattr(result, "protocol_version", None),
|
||
"capabilities": getattr(result, "capabilities", None),
|
||
"serverInfo": getattr(result, "serverInfo", None)
|
||
or getattr(result, "server_info", None),
|
||
"instructions": getattr(result, "instructions", None),
|
||
}
|
||
|
||
protocol = payload.get("protocolVersion") or payload.get("protocol_version")
|
||
server_info = payload.get("serverInfo") or payload.get("server_info")
|
||
msg = f"{client_label} MCP handshake complete (protocol={protocol or 'unknown'}, server={server_info or 'n/a'})"
|
||
logger.info(msg)
|
||
emit_console("INFO", msg)
|
||
|
||
instructions = payload.get("instructions")
|
||
if instructions:
|
||
instr_msg = f"{client_label} MCP server instructions: {instructions}"
|
||
logger.debug(instr_msg)
|
||
emit_console("DEBUG", instr_msg)
|
||
|
||
# Send notifications/initialized when the SDK exposes the helper. The
|
||
# attribute name changed between releases, so we probe the common forms.
|
||
notification_senders = [
|
||
getattr(session, "notify_initialized", None),
|
||
getattr(session, "notifications_initialized", None),
|
||
]
|
||
|
||
for sender in notification_senders:
|
||
if callable(sender):
|
||
try:
|
||
maybe_coro = sender()
|
||
if asyncio.iscoroutine(maybe_coro):
|
||
await maybe_coro
|
||
note = f"Sent notifications/initialized for {client_label}"
|
||
logger.debug(note)
|
||
emit_console("DEBUG", note)
|
||
break
|
||
except Exception as exc: # pragma: no cover - best-effort notification
|
||
fail_msg = (
|
||
f"Unable to send notifications/initialized for {client_label}: {exc}"
|
||
)
|
||
logger.debug(fail_msg, exc_info=True)
|
||
emit_console("DEBUG", fail_msg)
|
||
break
|
||
|
||
return payload
|
||
|
||
|
||
__all__ = ["perform_handshake", "emit_console"]
|