mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
fix(hanzoai): restore core SDK modules wiped by #45 regen — unbreaks hanzo-mcp (3.0.1)
#45 regenerated hanzoai as a pure OpenAPI client and DELETED the hand-written core-SDK modules (config/mcp/session/protocols/…) that live in the same package but aren't part of API generation. hanzo-mcp imports all four → the published hanzoai 3.0.0 crashed EVERY hanzo-mcp consumer on startup ('No module named hanzoai.protocols') → the MCP server wouldn't boot → agent couldn't reconnect. Restore the 4 self-contained modules hanzo-mcp needs (verified: zero imports of the removed Stainless internals); they coexist with the generated api/models client. Left agents/auth/cluster/llm_client OUT — they import the removed ._client and need migration to the new client (separate follow-up, no reconnect-path consumer). Bump 3.0.0→3.0.1. Boot-proven: hanzo-mcp starts, 26 tool packages / 29 tools, MCP initialize OK.
This commit is contained in:
@@ -0,0 +1,397 @@
|
||||
"""Hierarchical configuration discovery and merging.
|
||||
|
||||
Three layers (User < Project < Local). Later sources override earlier ones.
|
||||
MCP server blocks replace entirely. Missing files silently skipped.
|
||||
|
||||
Security: Project-level configs (.hanzo/settings.json in a repo) are untrusted.
|
||||
A cloned repo could contain a malicious config that defines mcpServers with
|
||||
stdio transport (arbitrary command execution) or overrides oauth URLs (credential
|
||||
phishing). By default, project and local configs are restricted to safe keys
|
||||
only. Pass trust_project_mcp=True to override (requires explicit user approval).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from enum import IntEnum
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConfigError(Exception):
|
||||
"""Raised when a config file exists but cannot be parsed."""
|
||||
|
||||
def __init__(self, path: Path, reason: str) -> None:
|
||||
self.path = path
|
||||
self.reason = reason
|
||||
super().__init__(f"{path}: {reason}")
|
||||
|
||||
|
||||
class ConfigSource(IntEnum):
|
||||
"""Config layer precedence. Higher value = higher priority."""
|
||||
|
||||
User = 0
|
||||
Project = 1
|
||||
Local = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConfigEntry:
|
||||
"""A discovered config file location."""
|
||||
|
||||
source: ConfigSource
|
||||
path: Path
|
||||
|
||||
|
||||
class McpTransport(IntEnum):
|
||||
Stdio = 0
|
||||
Sse = 1
|
||||
Http = 2
|
||||
Ws = 3
|
||||
Sdk = 4
|
||||
ClaudeAiProxy = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpStdioConfig:
|
||||
command: str
|
||||
args: list[str] = field(default_factory=list)
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def transport(self) -> McpTransport:
|
||||
return McpTransport.Stdio
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpOAuthConfig:
|
||||
client_id: str | None = None
|
||||
callback_port: int | None = None
|
||||
auth_server_metadata_url: str | None = None
|
||||
xaa: bool | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpRemoteConfig:
|
||||
url: str
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
headers_helper: str | None = None
|
||||
oauth: McpOAuthConfig | None = None
|
||||
transport: McpTransport = McpTransport.Http
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpWsConfig:
|
||||
url: str
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
headers_helper: str | None = None
|
||||
|
||||
@property
|
||||
def transport(self) -> McpTransport:
|
||||
return McpTransport.Ws
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpSdkConfig:
|
||||
name: str
|
||||
|
||||
@property
|
||||
def transport(self) -> McpTransport:
|
||||
return McpTransport.Sdk
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class McpClaudeAiProxyConfig:
|
||||
url: str
|
||||
id: str
|
||||
|
||||
@property
|
||||
def transport(self) -> McpTransport:
|
||||
return McpTransport.ClaudeAiProxy
|
||||
|
||||
|
||||
McpServerConfig = McpStdioConfig | McpRemoteConfig | McpWsConfig | McpSdkConfig | McpClaudeAiProxyConfig
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthConfig:
|
||||
client_id: str
|
||||
authorize_url: str
|
||||
token_url: str
|
||||
scopes: list[str] = field(default_factory=list)
|
||||
callback_port: int | None = None
|
||||
manual_redirect_url: str | None = None
|
||||
|
||||
|
||||
def _deep_merge(base: dict, override: dict) -> dict:
|
||||
"""Merge *override* into *base*, recursing into nested dicts.
|
||||
|
||||
Non-dict values in *override* replace those in *base*.
|
||||
"""
|
||||
merged = dict(base)
|
||||
for key, val in override.items():
|
||||
if key in merged and isinstance(merged[key], dict) and isinstance(val, dict):
|
||||
merged[key] = _deep_merge(merged[key], val)
|
||||
else:
|
||||
merged[key] = val
|
||||
return merged
|
||||
|
||||
|
||||
_TRANSPORT_MAP: dict[str, McpTransport] = {
|
||||
"stdio": McpTransport.Stdio,
|
||||
"sse": McpTransport.Sse,
|
||||
"http": McpTransport.Http,
|
||||
"ws": McpTransport.Ws,
|
||||
"sdk": McpTransport.Sdk,
|
||||
"claudeai-proxy": McpTransport.ClaudeAiProxy,
|
||||
}
|
||||
|
||||
|
||||
def _parse_mcp_server(name: str, raw: dict) -> McpServerConfig:
|
||||
transport_str = raw.get("transport", "stdio")
|
||||
transport = _TRANSPORT_MAP.get(transport_str)
|
||||
if transport is None:
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: unknown transport {transport_str!r}",
|
||||
)
|
||||
|
||||
if transport == McpTransport.Stdio:
|
||||
command = raw.get("command")
|
||||
if not command or not isinstance(command, str):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: stdio transport requires 'command' string",
|
||||
)
|
||||
args = raw.get("args", [])
|
||||
if not isinstance(args, list):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: 'args' must be a list",
|
||||
)
|
||||
env = raw.get("env", {})
|
||||
if not isinstance(env, dict):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: 'env' must be an object",
|
||||
)
|
||||
return McpStdioConfig(command=command, args=args, env=env)
|
||||
|
||||
if transport == McpTransport.Sdk:
|
||||
sdk_name = raw.get("name")
|
||||
if not sdk_name or not isinstance(sdk_name, str):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: sdk transport requires 'name' string",
|
||||
)
|
||||
return McpSdkConfig(name=sdk_name)
|
||||
|
||||
if transport == McpTransport.ClaudeAiProxy:
|
||||
url = raw.get("url")
|
||||
if not url or not isinstance(url, str):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: claudeai-proxy transport requires 'url' string",
|
||||
)
|
||||
server_id = raw.get("id")
|
||||
if not server_id or not isinstance(server_id, str):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: claudeai-proxy transport requires 'id' string",
|
||||
)
|
||||
return McpClaudeAiProxyConfig(url=url, id=server_id)
|
||||
|
||||
# Shared validation for url-based transports: sse, http, ws
|
||||
url = raw.get("url")
|
||||
if not url or not isinstance(url, str):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: {transport_str} transport requires 'url' string",
|
||||
)
|
||||
headers = raw.get("headers", {})
|
||||
if not isinstance(headers, dict):
|
||||
raise ConfigError(
|
||||
Path("<memory>"),
|
||||
f"mcpServers.{name}: 'headers' must be an object",
|
||||
)
|
||||
headers_helper = raw.get("headersHelper") or raw.get("headers_helper")
|
||||
|
||||
if transport == McpTransport.Ws:
|
||||
return McpWsConfig(url=url, headers=headers, headers_helper=headers_helper)
|
||||
|
||||
# sse, http -> McpRemoteConfig
|
||||
oauth_raw = raw.get("oauth")
|
||||
oauth: McpOAuthConfig | None = None
|
||||
if isinstance(oauth_raw, dict):
|
||||
oauth = McpOAuthConfig(
|
||||
client_id=oauth_raw.get("clientId") or oauth_raw.get("client_id"),
|
||||
callback_port=oauth_raw.get("callbackPort") or oauth_raw.get("callback_port"),
|
||||
auth_server_metadata_url=oauth_raw.get("authServerMetadataUrl") or oauth_raw.get("auth_server_metadata_url"),
|
||||
xaa=oauth_raw.get("xaa"),
|
||||
)
|
||||
return McpRemoteConfig(url=url, headers=headers, headers_helper=headers_helper, oauth=oauth, transport=transport)
|
||||
|
||||
|
||||
def _sanitize_project_config(data: dict, entry: ConfigEntry) -> dict:
|
||||
"""Remove dangerous keys from an untrusted project/local config.
|
||||
|
||||
Strips:
|
||||
- oauth (credential phishing via overridden auth URLs)
|
||||
- mcpServers entries with stdio transport (arbitrary command execution)
|
||||
|
||||
Non-stdio mcpServers (http, sse, ws, sdk) are kept -- they connect to
|
||||
remote URLs which are visible and auditable, not local binary execution.
|
||||
"""
|
||||
sanitized = dict(data)
|
||||
|
||||
if "oauth" in sanitized:
|
||||
logger.warning(
|
||||
"Ignoring 'oauth' from project config %s — "
|
||||
"project configs cannot override OAuth URLs (phishing risk). "
|
||||
"Move to user config (~/.hanzo/settings.json) instead.",
|
||||
entry.path,
|
||||
)
|
||||
del sanitized["oauth"]
|
||||
|
||||
mcp_raw = sanitized.get("mcpServers")
|
||||
if isinstance(mcp_raw, dict):
|
||||
safe_servers: dict = {}
|
||||
for name, cfg in mcp_raw.items():
|
||||
transport = cfg.get("transport", "stdio") if isinstance(cfg, dict) else "stdio"
|
||||
if transport == "stdio":
|
||||
logger.warning(
|
||||
"Ignoring mcpServers.%s (stdio) from project config %s — "
|
||||
"project configs cannot define stdio MCP servers (arbitrary "
|
||||
"command execution risk). Move to user config or pass "
|
||||
"trust_project_mcp=True.",
|
||||
name,
|
||||
entry.path,
|
||||
)
|
||||
else:
|
||||
safe_servers[name] = cfg
|
||||
if safe_servers:
|
||||
sanitized["mcpServers"] = safe_servers
|
||||
else:
|
||||
del sanitized["mcpServers"]
|
||||
|
||||
return sanitized
|
||||
|
||||
|
||||
@dataclass
|
||||
class RuntimeConfig:
|
||||
"""Merged configuration from all discovered layers."""
|
||||
|
||||
merged: dict = field(default_factory=dict)
|
||||
loaded_entries: list[ConfigEntry] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def empty(cls) -> RuntimeConfig:
|
||||
return cls()
|
||||
|
||||
def get(self, key: str) -> Any | None:
|
||||
return self.merged.get(key)
|
||||
|
||||
def mcp_servers(self) -> dict[str, McpServerConfig]:
|
||||
raw = self.merged.get("mcpServers")
|
||||
if not isinstance(raw, dict):
|
||||
return {}
|
||||
return {name: _parse_mcp_server(name, cfg) for name, cfg in raw.items()}
|
||||
|
||||
def oauth(self) -> OAuthConfig | None:
|
||||
raw = self.merged.get("oauth")
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
client_id = raw.get("clientId") or raw.get("client_id")
|
||||
if not client_id or not isinstance(client_id, str):
|
||||
return None
|
||||
authorize_url = raw.get("authorizeUrl") or raw.get("authorize_url")
|
||||
if not authorize_url or not isinstance(authorize_url, str):
|
||||
return None
|
||||
token_url = raw.get("tokenUrl") or raw.get("token_url")
|
||||
if not token_url or not isinstance(token_url, str):
|
||||
return None
|
||||
callback_port = raw.get("callbackPort") or raw.get("callback_port")
|
||||
return OAuthConfig(
|
||||
client_id=client_id,
|
||||
authorize_url=authorize_url,
|
||||
token_url=token_url,
|
||||
scopes=raw.get("scopes", []),
|
||||
callback_port=callback_port,
|
||||
manual_redirect_url=raw.get("manualRedirectUrl") or raw.get("manual_redirect_url"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConfigLoader:
|
||||
"""Discovers and loads hierarchical config files."""
|
||||
|
||||
cwd: Path
|
||||
config_home: Path
|
||||
|
||||
@classmethod
|
||||
def default_for(cls, cwd: Path) -> ConfigLoader:
|
||||
config_home = Path(os.environ.get("HANZO_CONFIG_HOME", Path.home() / ".hanzo"))
|
||||
return cls(cwd=cwd, config_home=config_home)
|
||||
|
||||
def discover(self) -> list[ConfigEntry]:
|
||||
"""Return potential config file entries in precedence order (low to high)."""
|
||||
return [
|
||||
ConfigEntry(ConfigSource.User, self.config_home / "settings.json"),
|
||||
ConfigEntry(ConfigSource.Project, self.cwd / ".hanzo" / "settings.json"),
|
||||
ConfigEntry(ConfigSource.Local, self.cwd / ".hanzo" / "settings.local.json"),
|
||||
]
|
||||
|
||||
# Keys that are safe to accept from project-level (untrusted) configs.
|
||||
# mcpServers with stdio transport allows arbitrary command execution.
|
||||
# oauth allows credential phishing by overriding auth URLs.
|
||||
# Both are blocked from project/local configs unless explicitly trusted.
|
||||
_PROJECT_SAFE_KEYS = frozenset({"model", "env", "hooks"})
|
||||
|
||||
def load(self, *, trust_project_mcp: bool = False) -> RuntimeConfig:
|
||||
"""Read and merge all existing config files.
|
||||
|
||||
Missing files are silently skipped. Invalid JSON raises ConfigError.
|
||||
The top-level value in each file must be a JSON object.
|
||||
|
||||
Args:
|
||||
trust_project_mcp: When False (default), project-level and
|
||||
local-level configs cannot define mcpServers with stdio
|
||||
transport or oauth blocks. This prevents a malicious cloned
|
||||
repo from executing arbitrary commands or phishing credentials
|
||||
via .hanzo/settings.json. Set True only after explicit user
|
||||
approval.
|
||||
"""
|
||||
merged: dict = {}
|
||||
loaded: list[ConfigEntry] = []
|
||||
|
||||
for entry in self.discover():
|
||||
if not entry.path.is_file():
|
||||
continue
|
||||
|
||||
text = entry.path.read_text(encoding="utf-8")
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise ConfigError(entry.path, f"invalid JSON: {exc}") from exc
|
||||
|
||||
if not isinstance(data, dict):
|
||||
raise ConfigError(entry.path, "top-level value must be a JSON object")
|
||||
|
||||
# R-06: Project/local configs are untrusted. Strip dangerous keys
|
||||
# unless the caller has explicitly opted in.
|
||||
if entry.source in (ConfigSource.Project, ConfigSource.Local) and not trust_project_mcp:
|
||||
data = _sanitize_project_config(data, entry)
|
||||
|
||||
# MCP servers from later sources replace entirely (not merged).
|
||||
if "mcpServers" in data and "mcpServers" in merged:
|
||||
del merged["mcpServers"]
|
||||
|
||||
merged = _deep_merge(merged, data)
|
||||
loaded.append(entry)
|
||||
|
||||
return RuntimeConfig(merged=merged, loaded_entries=loaded)
|
||||
@@ -0,0 +1,433 @@
|
||||
"""Hanzo AI MCP (Model Context Protocol) module.
|
||||
|
||||
This module provides access to the hanzo-mcp SDK for building
|
||||
MCP servers and clients with extensive tool support.
|
||||
|
||||
Includes JSON-RPC 2.0 stdio and HTTP transport clients for
|
||||
communicating with MCP servers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
try:
|
||||
# Try to import hanzo-mcp if installed
|
||||
from hanzo_mcp import (
|
||||
# Base classes
|
||||
BaseTool,
|
||||
ToolRegistry,
|
||||
# Server creation
|
||||
HanzoMCPServer,
|
||||
# Permissions
|
||||
PermissionManager,
|
||||
# Version
|
||||
__version__ as mcp_version,
|
||||
create_server,
|
||||
get_git_tools,
|
||||
get_agent_tools,
|
||||
get_shell_tools,
|
||||
get_memory_tools,
|
||||
get_search_tools,
|
||||
get_jupyter_tools,
|
||||
# Tool registration
|
||||
register_all_tools,
|
||||
# Tool categories
|
||||
get_filesystem_tools,
|
||||
register_all_prompts,
|
||||
)
|
||||
|
||||
MCP_AVAILABLE = True
|
||||
except ImportError:
|
||||
MCP_AVAILABLE = False
|
||||
|
||||
# Provide a helpful error message
|
||||
def _mcp_not_installed(*args, **kwargs):
|
||||
raise ImportError(
|
||||
"hanzo-mcp is not installed. Install it with: pip install hanzo-mcp"
|
||||
)
|
||||
|
||||
# Create placeholder classes/functions
|
||||
HanzoMCPServer = create_server = _mcp_not_installed
|
||||
register_all_tools = register_all_prompts = _mcp_not_installed
|
||||
get_filesystem_tools = get_shell_tools = get_agent_tools = _mcp_not_installed
|
||||
get_search_tools = get_jupyter_tools = get_git_tools = _mcp_not_installed
|
||||
get_memory_tools = _mcp_not_installed
|
||||
BaseTool = ToolRegistry = _mcp_not_installed
|
||||
PermissionManager = _mcp_not_installed
|
||||
mcp_version = "not installed"
|
||||
|
||||
|
||||
def create_mcp_server(name: str = "hanzo-mcp", allowed_paths: list = None, **kwargs):
|
||||
"""Create a new MCP server.
|
||||
|
||||
Args:
|
||||
name: Name of the server
|
||||
allowed_paths: List of allowed file paths
|
||||
**kwargs: Additional server options
|
||||
|
||||
Returns:
|
||||
HanzoMCPServer instance
|
||||
"""
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"hanzo-mcp is not installed. Install it with: pip install hanzo-mcp"
|
||||
)
|
||||
|
||||
return create_server(name=name, allowed_paths=allowed_paths, **kwargs)
|
||||
|
||||
|
||||
def run_mcp_server(name: str = "hanzo-mcp", transport: str = "stdio", **kwargs):
|
||||
"""Run an MCP server.
|
||||
|
||||
Args:
|
||||
name: Name of the server
|
||||
transport: Transport protocol ("stdio" or "sse")
|
||||
**kwargs: Additional server options
|
||||
"""
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"hanzo-mcp is not installed. Install it with: pip install hanzo-mcp"
|
||||
)
|
||||
|
||||
server = create_server(name=name, **kwargs)
|
||||
server.run(transport=transport)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP tool name normalization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NORM_RE = re.compile(r"[^a-zA-Z0-9_\-]")
|
||||
|
||||
|
||||
def _collapse_underscores(value: str) -> str:
|
||||
"""Collapse consecutive underscores into a single underscore."""
|
||||
out: list[str] = []
|
||||
last_was_underscore = False
|
||||
for ch in value:
|
||||
if ch == "_":
|
||||
if not last_was_underscore:
|
||||
out.append(ch)
|
||||
last_was_underscore = True
|
||||
else:
|
||||
out.append(ch)
|
||||
last_was_underscore = False
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def normalize_mcp_name(name: str) -> str:
|
||||
"""Normalize an MCP server/tool name for the mcp__ prefix convention.
|
||||
|
||||
Exact port of claw-code's ``normalize_name_for_mcp`` (mcp.rs):
|
||||
- Replace any char that is not alphanumeric, underscore, or hyphen with ``_``.
|
||||
- If *name* starts with ``"claude.ai "``, also collapse consecutive
|
||||
underscores and strip leading/trailing underscores.
|
||||
- Non-claude.ai names keep hyphens and trailing underscores as-is.
|
||||
"""
|
||||
normalized = _NORM_RE.sub("_", name)
|
||||
if name.startswith("claude.ai "):
|
||||
normalized = _collapse_underscores(normalized).strip("_")
|
||||
return normalized
|
||||
|
||||
|
||||
def mcp_tool_name(server_name: str, tool_name: str) -> str:
|
||||
"""Build the canonical mcp__server__tool name."""
|
||||
return f"mcp__{normalize_mcp_name(server_name)}__{normalize_mcp_name(tool_name)}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON-RPC 2.0 protocol types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MCPClientError(Exception):
|
||||
"""Raised on MCP client transport or protocol errors."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class JsonRpcError:
|
||||
code: int
|
||||
message: str
|
||||
data: Any = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict) -> JsonRpcError:
|
||||
return cls(code=d["code"], message=d["message"], data=d.get("data"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class JsonRpcRequest:
|
||||
method: str
|
||||
params: dict[str, Any] = field(default_factory=dict)
|
||||
id: int = 0
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"jsonrpc": "2.0", "method": self.method, "params": self.params, "id": self.id}
|
||||
|
||||
def to_line(self) -> bytes:
|
||||
return json.dumps(self.to_dict(), separators=(",", ":")).encode() + b"\n"
|
||||
|
||||
|
||||
@dataclass
|
||||
class JsonRpcResponse:
|
||||
id: int
|
||||
result: Any = None
|
||||
error: JsonRpcError | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict) -> JsonRpcResponse:
|
||||
err = None
|
||||
if "error" in d and d["error"] is not None:
|
||||
err = JsonRpcError.from_dict(d["error"])
|
||||
return cls(id=d.get("id", 0), result=d.get("result"), error=err)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stdio transport MCP client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MCP_PROTOCOL_VERSION = "2024-11-05"
|
||||
|
||||
_INIT_PARAMS: dict[str, Any] = {
|
||||
"protocolVersion": _MCP_PROTOCOL_VERSION,
|
||||
"capabilities": {},
|
||||
"clientInfo": {"name": "hanzoai-python", "version": "1.0.0"},
|
||||
}
|
||||
|
||||
|
||||
class MCPClient:
|
||||
"""MCP client that communicates with a server subprocess over stdio JSON-RPC."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
server_command: list[str],
|
||||
env: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self.server_command = server_command
|
||||
self.env = env
|
||||
self._process: asyncio.subprocess.Process | None = None
|
||||
self._next_id: int = 1
|
||||
self.server_info: dict[str, Any] | None = None
|
||||
self.capabilities: dict[str, Any] | None = None
|
||||
|
||||
def _require_connected(self) -> None:
|
||||
if self._process is None or self._process.stdin is None:
|
||||
raise MCPClientError("not connected - call connect() first")
|
||||
|
||||
async def _send(self, method: str, params: dict[str, Any] | None = None) -> Any:
|
||||
self._require_connected()
|
||||
req = JsonRpcRequest(method=method, params=params or {}, id=self._next_id)
|
||||
self._next_id += 1
|
||||
|
||||
stdin = self._process.stdin
|
||||
stdout = self._process.stdout
|
||||
assert stdin is not None and stdout is not None
|
||||
|
||||
stdin.write(req.to_line())
|
||||
await stdin.drain()
|
||||
|
||||
try:
|
||||
raw = await asyncio.wait_for(stdout.readline(), timeout=30.0)
|
||||
except asyncio.TimeoutError:
|
||||
raise MCPClientError("server did not respond within 30 seconds") from None
|
||||
if not raw:
|
||||
raise MCPClientError("server closed stdout unexpectedly")
|
||||
if len(raw) > 10_485_760: # 10 MB guard
|
||||
raise MCPClientError("server response exceeded 10 MB limit")
|
||||
|
||||
resp = JsonRpcResponse.from_dict(json.loads(raw))
|
||||
if resp.error is not None:
|
||||
raise MCPClientError(f"JSON-RPC error {resp.error.code}: {resp.error.message}")
|
||||
return resp.result
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Spawn the server process and perform the MCP initialize handshake."""
|
||||
try:
|
||||
self._process = await asyncio.create_subprocess_exec(
|
||||
*self.server_command,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
env=self.env,
|
||||
)
|
||||
except (FileNotFoundError, PermissionError, OSError) as exc:
|
||||
raise MCPClientError(f"failed to start server process: {exc}") from exc
|
||||
|
||||
result = await self._send("initialize", _INIT_PARAMS)
|
||||
self.server_info = result.get("serverInfo")
|
||||
self.capabilities = result.get("capabilities")
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
"""Terminate the server process."""
|
||||
proc = self._process
|
||||
if proc is None:
|
||||
return
|
||||
self._process = None
|
||||
if proc.stdin is not None:
|
||||
proc.stdin.close()
|
||||
try:
|
||||
proc.terminate()
|
||||
await asyncio.wait_for(proc.wait(), timeout=5.0)
|
||||
except (ProcessLookupError, asyncio.TimeoutError):
|
||||
proc.kill()
|
||||
await proc.wait()
|
||||
|
||||
async def list_tools(self) -> list[dict[str, Any]]:
|
||||
"""Fetch all tools, following pagination cursors."""
|
||||
self._require_connected()
|
||||
tools: list[dict[str, Any]] = []
|
||||
cursor: str | None = None
|
||||
while True:
|
||||
params: dict[str, Any] = {}
|
||||
if cursor is not None:
|
||||
params["cursor"] = cursor
|
||||
result = await self._send("tools/list", params)
|
||||
tools.extend(result.get("tools", []))
|
||||
cursor = result.get("nextCursor")
|
||||
if not cursor:
|
||||
break
|
||||
return tools
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
"""Call a tool on the server."""
|
||||
self._require_connected()
|
||||
params: dict[str, Any] = {"name": name}
|
||||
if arguments:
|
||||
params["arguments"] = arguments
|
||||
return await self._send("tools/call", params)
|
||||
|
||||
async def list_resources(self) -> list[dict[str, Any]]:
|
||||
"""Fetch all resources."""
|
||||
self._require_connected()
|
||||
result = await self._send("resources/list")
|
||||
return result.get("resources", [])
|
||||
|
||||
async def __aenter__(self) -> MCPClient:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
await self.disconnect()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP transport MCP client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MCPHttpClient:
|
||||
"""MCP client that communicates with a server over HTTP POST (JSON-RPC)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url: str,
|
||||
headers: dict[str, str] | None = None,
|
||||
timeout: float = 30.0,
|
||||
) -> None:
|
||||
self.url = url
|
||||
self._client = httpx.AsyncClient(
|
||||
headers=headers or {},
|
||||
timeout=httpx.Timeout(timeout),
|
||||
)
|
||||
self._next_id: int = 1
|
||||
self.server_info: dict[str, Any] | None = None
|
||||
self.capabilities: dict[str, Any] | None = None
|
||||
|
||||
async def _send(self, method: str, params: dict[str, Any] | None = None) -> Any:
|
||||
req = JsonRpcRequest(method=method, params=params or {}, id=self._next_id)
|
||||
self._next_id += 1
|
||||
|
||||
resp = await self._client.post(self.url, json=req.to_dict())
|
||||
resp.raise_for_status()
|
||||
body = resp.json()
|
||||
|
||||
rpc_resp = JsonRpcResponse.from_dict(body)
|
||||
if rpc_resp.error is not None:
|
||||
raise MCPClientError(f"JSON-RPC error {rpc_resp.error.code}: {rpc_resp.error.message}")
|
||||
return rpc_resp.result
|
||||
|
||||
async def connect(self) -> None:
|
||||
"""Perform the MCP initialize handshake over HTTP."""
|
||||
result = await self._send("initialize", _INIT_PARAMS)
|
||||
self.server_info = result.get("serverInfo")
|
||||
self.capabilities = result.get("capabilities")
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
"""Close the HTTP client."""
|
||||
await self._client.aclose()
|
||||
|
||||
async def list_tools(self) -> list[dict[str, Any]]:
|
||||
tools: list[dict[str, Any]] = []
|
||||
cursor: str | None = None
|
||||
while True:
|
||||
params: dict[str, Any] = {}
|
||||
if cursor is not None:
|
||||
params["cursor"] = cursor
|
||||
result = await self._send("tools/list", params)
|
||||
tools.extend(result.get("tools", []))
|
||||
cursor = result.get("nextCursor")
|
||||
if not cursor:
|
||||
break
|
||||
return tools
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
params: dict[str, Any] = {"name": name}
|
||||
if arguments:
|
||||
params["arguments"] = arguments
|
||||
return await self._send("tools/call", params)
|
||||
|
||||
async def list_resources(self) -> list[dict[str, Any]]:
|
||||
result = await self._send("resources/list")
|
||||
return result.get("resources", [])
|
||||
|
||||
async def __aenter__(self) -> MCPHttpClient:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
await self.disconnect()
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Server
|
||||
"HanzoMCPServer",
|
||||
"create_server",
|
||||
"create_mcp_server",
|
||||
"run_mcp_server",
|
||||
# Tools
|
||||
"register_all_tools",
|
||||
"register_all_prompts",
|
||||
"get_filesystem_tools",
|
||||
"get_shell_tools",
|
||||
"get_agent_tools",
|
||||
"get_search_tools",
|
||||
"get_jupyter_tools",
|
||||
"get_git_tools",
|
||||
"get_memory_tools",
|
||||
# Base classes
|
||||
"BaseTool",
|
||||
"ToolRegistry",
|
||||
# Permissions
|
||||
"PermissionManager",
|
||||
# Protocol types
|
||||
"JsonRpcRequest",
|
||||
"JsonRpcResponse",
|
||||
"JsonRpcError",
|
||||
"MCPClientError",
|
||||
# Clients
|
||||
"MCPClient",
|
||||
"MCPHttpClient",
|
||||
# Name helpers
|
||||
"normalize_mcp_name",
|
||||
"mcp_tool_name",
|
||||
# Status
|
||||
"MCP_AVAILABLE",
|
||||
"mcp_version",
|
||||
]
|
||||
@@ -0,0 +1,437 @@
|
||||
"""Protocol abstractions for pluggable tool execution, permissions, and conversation runtime.
|
||||
|
||||
Ported from claw-code's trait-based architecture to enable dependency injection
|
||||
and testing without network calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum, auto
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool execution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@runtime_checkable
|
||||
class ToolExecutor(Protocol):
|
||||
"""Execute a named tool with a JSON-encoded input string."""
|
||||
|
||||
def execute(self, tool_name: str, input: str) -> str: ...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Permissions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class PermissionMode(Enum):
|
||||
Allow = auto()
|
||||
Deny = auto()
|
||||
Prompt = auto()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PermissionRequest:
|
||||
tool_name: str
|
||||
input: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PermissionOutcome:
|
||||
allowed: bool
|
||||
reason: str = ""
|
||||
|
||||
@classmethod
|
||||
def allow(cls) -> PermissionOutcome:
|
||||
return cls(allowed=True)
|
||||
|
||||
@classmethod
|
||||
def deny(cls, reason: str) -> PermissionOutcome:
|
||||
return cls(allowed=False, reason=reason)
|
||||
|
||||
|
||||
# Named constructors matching Rust enum variants.
|
||||
PERMISSION_ALLOW = PermissionOutcome(allowed=True)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class PermissionPrompter(Protocol):
|
||||
"""Interactively decide whether a tool invocation is allowed."""
|
||||
|
||||
def decide(self, request: PermissionRequest) -> PermissionOutcome: ...
|
||||
|
||||
|
||||
class PermissionPolicy:
|
||||
"""Default-mode + per-tool override permission policy.
|
||||
|
||||
Mirrors the Rust ``PermissionPolicy`` struct: a default mode with an
|
||||
ordered map of per-tool overrides.
|
||||
"""
|
||||
|
||||
__slots__ = ("default_mode", "_tool_modes")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_mode: PermissionMode = PermissionMode.Prompt,
|
||||
tool_modes: dict[str, PermissionMode] | None = None,
|
||||
) -> None:
|
||||
self.default_mode = default_mode
|
||||
# Sorted dict to match Rust BTreeMap ordering.
|
||||
self._tool_modes: dict[str, PermissionMode] = dict(
|
||||
sorted((tool_modes or {}).items())
|
||||
)
|
||||
|
||||
def with_tool_mode(self, name: str, mode: PermissionMode) -> PermissionPolicy:
|
||||
"""Fluent builder: set a per-tool override and return self for chaining."""
|
||||
self._tool_modes[name] = mode
|
||||
self._tool_modes = dict(sorted(self._tool_modes.items()))
|
||||
return self
|
||||
|
||||
def mode_for(self, tool_name: str) -> PermissionMode:
|
||||
return self._tool_modes.get(tool_name, self.default_mode)
|
||||
|
||||
def authorize(
|
||||
self,
|
||||
tool_name: str,
|
||||
input: str,
|
||||
prompter: PermissionPrompter | None = None,
|
||||
) -> PermissionOutcome:
|
||||
mode = self.mode_for(tool_name)
|
||||
if mode is PermissionMode.Allow:
|
||||
return PermissionOutcome.allow()
|
||||
if mode is PermissionMode.Deny:
|
||||
return PermissionOutcome.deny(f"tool '{tool_name}' denied by policy")
|
||||
# Prompt
|
||||
if prompter is None:
|
||||
return PermissionOutcome.deny("no prompter available for interactive decision")
|
||||
return prompter.decide(PermissionRequest(tool_name=tool_name, input=input))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TokenUsage:
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
|
||||
def __iadd__(self, other: TokenUsage) -> TokenUsage:
|
||||
self.input_tokens += other.input_tokens
|
||||
self.output_tokens += other.output_tokens
|
||||
self.cache_creation_input_tokens += other.cache_creation_input_tokens
|
||||
self.cache_read_input_tokens += other.cache_read_input_tokens
|
||||
return self
|
||||
|
||||
def total_tokens(self) -> int:
|
||||
return (
|
||||
self.input_tokens
|
||||
+ self.output_tokens
|
||||
+ self.cache_creation_input_tokens
|
||||
+ self.cache_read_input_tokens
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"cache_creation_input_tokens": self.cache_creation_input_tokens,
|
||||
"cache_read_input_tokens": self.cache_read_input_tokens,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> TokenUsage:
|
||||
return TokenUsage(
|
||||
input_tokens=d.get("input_tokens", 0),
|
||||
output_tokens=d.get("output_tokens", 0),
|
||||
cache_creation_input_tokens=d.get("cache_creation_input_tokens", 0),
|
||||
cache_read_input_tokens=d.get("cache_read_input_tokens", 0),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelPricing:
|
||||
"""Per-token prices in USD."""
|
||||
input_price_per_token: float
|
||||
output_price_per_token: float
|
||||
cache_creation_price_per_token: float = 0.0
|
||||
cache_read_price_per_token: float = 0.0
|
||||
|
||||
def cost(self, usage: TokenUsage) -> float:
|
||||
return (
|
||||
usage.input_tokens * self.input_price_per_token
|
||||
+ usage.output_tokens * self.output_price_per_token
|
||||
+ usage.cache_creation_input_tokens * self.cache_creation_price_per_token
|
||||
+ usage.cache_read_input_tokens * self.cache_read_price_per_token
|
||||
)
|
||||
|
||||
|
||||
class UsageTracker:
|
||||
"""Accumulates per-turn and cumulative token usage."""
|
||||
|
||||
__slots__ = ("_turns",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._turns: list[TokenUsage] = []
|
||||
|
||||
def record(self, usage: TokenUsage) -> None:
|
||||
self._turns.append(usage)
|
||||
|
||||
@property
|
||||
def turns(self) -> int:
|
||||
"""Number of recorded turns. Matches Rust ``pub fn turns(&self) -> u32``."""
|
||||
return len(self._turns)
|
||||
|
||||
def turns_list(self) -> list[TokenUsage]:
|
||||
"""Return a copy of per-turn usage records."""
|
||||
return list(self._turns)
|
||||
|
||||
def cumulative_usage(self) -> TokenUsage:
|
||||
total = TokenUsage()
|
||||
for t in self._turns:
|
||||
total += t
|
||||
return total
|
||||
|
||||
def cost(self, pricing: ModelPricing) -> float:
|
||||
return pricing.cost(self.cumulative_usage())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# API transport types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AssistantEvent:
|
||||
"""Event from the assistant stream.
|
||||
|
||||
``kind`` uses string values matching Rust enum variant names:
|
||||
``"text_delta"``, ``"tool_use"``, ``"usage"``, ``"message_stop"``.
|
||||
"""
|
||||
kind: str
|
||||
text: str = ""
|
||||
tool_use_id: str = ""
|
||||
tool_name: str = ""
|
||||
tool_input: str = ""
|
||||
usage: TokenUsage | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ApiRequest:
|
||||
system_prompt: list[str]
|
||||
messages: list[dict]
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ApiClient(Protocol):
|
||||
"""Send a conversation request and receive a stream of events."""
|
||||
|
||||
def stream(self, request: ApiRequest) -> list[AssistantEvent]: ...
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Turn summary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass(slots=True)
|
||||
class TurnSummary:
|
||||
assistant_messages: list[dict] = field(default_factory=list)
|
||||
tool_results: list[dict] = field(default_factory=list)
|
||||
usage: TokenUsage = field(default_factory=TokenUsage)
|
||||
iterations: int = 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static tool executor (test helper)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class StaticToolExecutor:
|
||||
"""Register handlers by name. Useful for testing."""
|
||||
|
||||
__slots__ = ("_handlers",)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._handlers: dict[str, Callable[[str], str]] = {}
|
||||
|
||||
def register(self, name: str, handler: Callable[[str], str]) -> None:
|
||||
self._handlers[name] = handler
|
||||
|
||||
def execute(self, tool_name: str, input: str) -> str:
|
||||
handler = self._handlers.get(tool_name)
|
||||
if handler is None:
|
||||
raise KeyError(f"no handler registered for tool '{tool_name}'")
|
||||
return handler(input)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Conversation runtime
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DEFAULT_MAX_ITERATIONS = 16
|
||||
|
||||
|
||||
class ConversationRuntime:
|
||||
"""Orchestrate turn execution over an API client and tool executor.
|
||||
|
||||
Generic over any ``ApiClient`` and ``ToolExecutor`` implementations,
|
||||
mirroring the Rust ``ConversationRuntime<A, T>`` generic struct.
|
||||
"""
|
||||
|
||||
__slots__ = (
|
||||
"session",
|
||||
"api_client",
|
||||
"tool_executor",
|
||||
"permission_policy",
|
||||
"system_prompt",
|
||||
"max_iterations",
|
||||
"usage_tracker",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session: list[dict],
|
||||
api_client: ApiClient,
|
||||
tool_executor: ToolExecutor,
|
||||
permission_policy: PermissionPolicy,
|
||||
system_prompt: list[str] | None = None,
|
||||
max_iterations: int = _DEFAULT_MAX_ITERATIONS,
|
||||
) -> None:
|
||||
self.session = session
|
||||
self.api_client = api_client
|
||||
self.tool_executor = tool_executor
|
||||
self.permission_policy = permission_policy
|
||||
self.system_prompt = system_prompt or []
|
||||
self.max_iterations = max_iterations
|
||||
self.usage_tracker = UsageTracker()
|
||||
|
||||
# ---- public API -------------------------------------------------------
|
||||
|
||||
def run_turn(
|
||||
self,
|
||||
user_input: str,
|
||||
prompter: PermissionPrompter | None = None,
|
||||
) -> TurnSummary:
|
||||
"""Execute one conversational turn, looping on tool use.
|
||||
|
||||
Appends the user message, then repeatedly:
|
||||
1. Build ``ApiRequest`` with system_prompt + messages.
|
||||
2. Call ``api_client.stream()`` to get assistant events.
|
||||
3. Assemble the assistant message from events.
|
||||
4. If no tool_use blocks are present, break.
|
||||
5. For each tool_use, authorize via permission_policy then execute.
|
||||
6. Append results to session.
|
||||
7. Guard against ``max_iterations``.
|
||||
"""
|
||||
self.session.append({"role": "user", "content": user_input})
|
||||
|
||||
summary = TurnSummary()
|
||||
|
||||
for iteration in range(1, self.max_iterations + 1):
|
||||
summary.iterations = iteration
|
||||
|
||||
request = ApiRequest(
|
||||
system_prompt=self.system_prompt,
|
||||
messages=list(self.session),
|
||||
)
|
||||
events = self.api_client.stream(request)
|
||||
assistant_msg, tool_uses, turn_usage = _build_assistant_message(events)
|
||||
|
||||
self.session.append(assistant_msg)
|
||||
summary.assistant_messages.append(assistant_msg)
|
||||
summary.usage += turn_usage
|
||||
self.usage_tracker.record(turn_usage)
|
||||
|
||||
if not tool_uses:
|
||||
break
|
||||
|
||||
for tu in tool_uses:
|
||||
outcome = self.permission_policy.authorize(
|
||||
tu.tool_name, tu.tool_input, prompter,
|
||||
)
|
||||
|
||||
if outcome.allowed:
|
||||
try:
|
||||
output = self.tool_executor.execute(tu.tool_name, tu.tool_input)
|
||||
is_error = False
|
||||
except Exception as exc:
|
||||
output = str(exc)
|
||||
is_error = True
|
||||
else:
|
||||
output = outcome.reason
|
||||
is_error = True
|
||||
|
||||
result_msg = {
|
||||
"role": "tool",
|
||||
"tool_use_id": tu.tool_use_id,
|
||||
"tool_name": tu.tool_name,
|
||||
"output": output,
|
||||
"is_error": is_error,
|
||||
}
|
||||
self.session.append(result_msg)
|
||||
summary.tool_results.append(result_msg)
|
||||
|
||||
return summary
|
||||
|
||||
|
||||
def _build_assistant_message(
|
||||
events: list[AssistantEvent],
|
||||
) -> tuple[dict, list[AssistantEvent], TokenUsage]:
|
||||
"""Parse stream events into an assistant message dict, tool uses, and usage."""
|
||||
text_parts: list[str] = []
|
||||
tool_uses: list[AssistantEvent] = []
|
||||
usage = TokenUsage()
|
||||
blocks: list[dict] = []
|
||||
|
||||
for event in events:
|
||||
if event.kind == "text_delta":
|
||||
text_parts.append(event.text)
|
||||
elif event.kind == "tool_use":
|
||||
# Flush accumulated text
|
||||
if text_parts:
|
||||
blocks.append({"type": "text", "text": "".join(text_parts)})
|
||||
text_parts.clear()
|
||||
blocks.append({
|
||||
"type": "tool_use",
|
||||
"id": event.tool_use_id,
|
||||
"name": event.tool_name,
|
||||
"input": event.tool_input,
|
||||
})
|
||||
tool_uses.append(event)
|
||||
elif event.kind == "usage" and event.usage is not None:
|
||||
usage += event.usage
|
||||
|
||||
if text_parts:
|
||||
blocks.append({"type": "text", "text": "".join(text_parts)})
|
||||
|
||||
msg = {"role": "assistant", "blocks": blocks}
|
||||
return msg, tool_uses, usage
|
||||
|
||||
|
||||
__all__ = [
|
||||
# Tool execution
|
||||
"ToolExecutor",
|
||||
"StaticToolExecutor",
|
||||
# Permissions
|
||||
"PermissionMode",
|
||||
"PermissionRequest",
|
||||
"PermissionOutcome",
|
||||
"PermissionPrompter",
|
||||
"PermissionPolicy",
|
||||
# API client
|
||||
"ApiClient",
|
||||
"ApiRequest",
|
||||
"AssistantEvent",
|
||||
# (AssistantEvent uses string 'kind' values, no separate enum)
|
||||
# Usage
|
||||
"TokenUsage",
|
||||
"ModelPricing",
|
||||
"UsageTracker",
|
||||
# Conversation
|
||||
"ConversationRuntime",
|
||||
"TurnSummary",
|
||||
]
|
||||
@@ -0,0 +1,395 @@
|
||||
"""Session management with compaction for Hanzo AI SDK.
|
||||
|
||||
Ported from claw-code's session.rs and compact.rs. Provides conversation
|
||||
message types, session persistence, and context window compaction.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Union
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Message primitives
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MessageRole(Enum):
|
||||
System = "system"
|
||||
User = "user"
|
||||
Assistant = "assistant"
|
||||
Tool = "tool"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextBlock:
|
||||
text: str
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"type": "text", "text": self.text}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> TextBlock:
|
||||
return TextBlock(text=d["text"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolUseBlock:
|
||||
id: str
|
||||
name: str
|
||||
input: str # Raw JSON string, matching Rust's String type
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {"type": "tool_use", "id": self.id, "name": self.name, "input": self.input}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> ToolUseBlock:
|
||||
raw = d["input"]
|
||||
# Backwards compat: accept dict from old serialized data, convert to JSON string
|
||||
if isinstance(raw, dict):
|
||||
raw = json.dumps(raw, separators=(",", ":"))
|
||||
return ToolUseBlock(id=d["id"], name=d["name"], input=raw)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResultBlock:
|
||||
tool_use_id: str
|
||||
tool_name: str
|
||||
output: str
|
||||
is_error: bool = False
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "tool_result",
|
||||
"tool_use_id": self.tool_use_id,
|
||||
"tool_name": self.tool_name,
|
||||
"output": self.output,
|
||||
"is_error": self.is_error,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> ToolResultBlock:
|
||||
return ToolResultBlock(
|
||||
tool_use_id=d["tool_use_id"],
|
||||
tool_name=d["tool_name"],
|
||||
output=d["output"],
|
||||
is_error=d.get("is_error", False),
|
||||
)
|
||||
|
||||
|
||||
ContentBlock = Union[TextBlock, ToolUseBlock, ToolResultBlock]
|
||||
|
||||
_BLOCK_PARSERS = {
|
||||
"text": TextBlock.from_dict,
|
||||
"tool_use": ToolUseBlock.from_dict,
|
||||
"tool_result": ToolResultBlock.from_dict,
|
||||
}
|
||||
|
||||
|
||||
def _block_from_dict(d: dict[str, Any]) -> ContentBlock:
|
||||
parser = _BLOCK_PARSERS.get(d["type"])
|
||||
if parser is None:
|
||||
raise ValueError(f"unknown block type: {d['type']}")
|
||||
return parser(d)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token usage — canonical definition lives in protocols.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
from hanzoai.protocols import TokenUsage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConversationMessage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class ConversationMessage:
|
||||
role: MessageRole
|
||||
blocks: list[ContentBlock]
|
||||
usage: TokenUsage | None = None
|
||||
|
||||
# -- Factory methods --
|
||||
|
||||
@staticmethod
|
||||
def user_text(text: str) -> ConversationMessage:
|
||||
return ConversationMessage(role=MessageRole.User, blocks=[TextBlock(text=text)])
|
||||
|
||||
@staticmethod
|
||||
def assistant(blocks: list[ContentBlock]) -> ConversationMessage:
|
||||
return ConversationMessage(role=MessageRole.Assistant, blocks=blocks)
|
||||
|
||||
@staticmethod
|
||||
def assistant_with_usage(blocks: list[ContentBlock], usage: TokenUsage) -> ConversationMessage:
|
||||
return ConversationMessage(role=MessageRole.Assistant, blocks=blocks, usage=usage)
|
||||
|
||||
@staticmethod
|
||||
def tool_result(tool_use_id: str, tool_name: str, output: str, is_error: bool = False) -> ConversationMessage:
|
||||
block = ToolResultBlock(tool_use_id=tool_use_id, tool_name=tool_name, output=output, is_error=is_error)
|
||||
return ConversationMessage(role=MessageRole.Tool, blocks=[block])
|
||||
|
||||
# -- Serialization --
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d: dict[str, Any] = {
|
||||
"role": self.role.value,
|
||||
"blocks": [b.to_dict() for b in self.blocks],
|
||||
}
|
||||
if self.usage is not None:
|
||||
d["usage"] = self.usage.to_dict()
|
||||
return d
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> ConversationMessage:
|
||||
return ConversationMessage(
|
||||
role=MessageRole(d["role"]),
|
||||
blocks=[_block_from_dict(b) for b in d["blocks"]],
|
||||
usage=TokenUsage.from_dict(d["usage"]) if "usage" in d else None,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
version: int = 1
|
||||
messages: list[ConversationMessage] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"version": self.version,
|
||||
"messages": [m.to_dict() for m in self.messages],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: dict[str, Any]) -> Session:
|
||||
return Session(
|
||||
version=d.get("version", 1),
|
||||
messages=[ConversationMessage.from_dict(m) for m in d.get("messages", [])],
|
||||
)
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
Path(path).write_text(json.dumps(self.to_dict(), indent=2), encoding="utf-8")
|
||||
|
||||
@staticmethod
|
||||
def load(path: str | Path) -> Session:
|
||||
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
return Session.from_dict(data)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Compaction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class CompactionConfig:
|
||||
preserve_recent_messages: int = 4
|
||||
max_estimated_tokens: int = 10_000
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionResult:
|
||||
summary: str
|
||||
formatted_summary: str
|
||||
compacted_session: Session
|
||||
removed_message_count: int
|
||||
|
||||
|
||||
_FILE_EXTENSIONS = frozenset({
|
||||
".py", ".rs", ".go", ".ts", ".tsx", ".js", ".jsx", ".json", ".yaml", ".yml",
|
||||
".toml", ".md", ".sh", ".sql", ".html", ".css", ".c", ".h", ".cpp", ".hpp",
|
||||
})
|
||||
|
||||
_PENDING_KEYWORDS = frozenset({"todo", "next", "pending", "follow up", "remaining"})
|
||||
|
||||
|
||||
def estimate_session_tokens(session: Session) -> int:
|
||||
"""Estimate token count: len(text)/4 + 1 per block."""
|
||||
total = 0
|
||||
for msg in session.messages:
|
||||
for block in msg.blocks:
|
||||
if isinstance(block, TextBlock):
|
||||
total += len(block.text) // 4 + 1
|
||||
elif isinstance(block, ToolUseBlock):
|
||||
total += len(block.input) // 4 + 1
|
||||
elif isinstance(block, ToolResultBlock):
|
||||
total += len(block.output) // 4 + 1
|
||||
return total
|
||||
|
||||
|
||||
def should_compact(session: Session, config: CompactionConfig) -> bool:
|
||||
return (
|
||||
len(session.messages) > config.preserve_recent_messages
|
||||
and estimate_session_tokens(session) >= config.max_estimated_tokens
|
||||
)
|
||||
|
||||
|
||||
def compact_session(session: Session, config: CompactionConfig) -> CompactionResult:
|
||||
"""Compact a session by summarizing old messages and keeping recent ones."""
|
||||
if not should_compact(session, config):
|
||||
return CompactionResult(
|
||||
summary="",
|
||||
formatted_summary="",
|
||||
compacted_session=session,
|
||||
removed_message_count=0,
|
||||
)
|
||||
|
||||
msgs = session.messages
|
||||
keep_count = min(config.preserve_recent_messages, len(msgs))
|
||||
split = len(msgs) - keep_count
|
||||
removed = msgs[:split]
|
||||
kept = msgs[split:]
|
||||
|
||||
summary = _build_summary(removed)
|
||||
formatted = format_compact_summary(summary)
|
||||
|
||||
continuation_text = (
|
||||
"This session is being continued from a previous conversation that ran "
|
||||
"out of context. The summary below covers the earlier portion of the "
|
||||
"conversation.\n\n"
|
||||
f"{formatted}\n\n"
|
||||
"Recent messages are preserved verbatim.\n"
|
||||
"Continue the conversation from where it left off without asking the "
|
||||
"user any further questions. Resume directly \u2014 do not acknowledge the "
|
||||
"summary, do not recap what was happening, and do not preface with "
|
||||
"continuation text."
|
||||
)
|
||||
|
||||
system_msg = ConversationMessage(
|
||||
role=MessageRole.System,
|
||||
blocks=[TextBlock(text=continuation_text)],
|
||||
)
|
||||
|
||||
compacted = Session(
|
||||
version=session.version,
|
||||
messages=[system_msg] + kept,
|
||||
)
|
||||
|
||||
return CompactionResult(
|
||||
summary=summary,
|
||||
formatted_summary=formatted,
|
||||
compacted_session=compacted,
|
||||
removed_message_count=len(removed),
|
||||
)
|
||||
|
||||
|
||||
def _build_summary(messages: list[ConversationMessage]) -> str:
|
||||
"""Build a structured summary of removed messages, matching Rust summarize_messages()."""
|
||||
role_counts: dict[str, int] = {}
|
||||
tool_names: list[str] = []
|
||||
user_texts: list[str] = []
|
||||
pending_items: list[str] = []
|
||||
file_candidates: set[str] = set()
|
||||
last_text = ""
|
||||
timeline: list[tuple[str, str]] = [] # (role, block_summary)
|
||||
|
||||
for msg in messages:
|
||||
role_counts[msg.role.value] = role_counts.get(msg.role.value, 0) + 1
|
||||
|
||||
block_parts: list[str] = []
|
||||
for block in msg.blocks:
|
||||
if isinstance(block, ToolUseBlock):
|
||||
if block.name not in tool_names:
|
||||
tool_names.append(block.name)
|
||||
block_parts.append(f"[tool_use: {block.name}]")
|
||||
elif isinstance(block, ToolResultBlock):
|
||||
if block.tool_name not in tool_names:
|
||||
tool_names.append(block.tool_name)
|
||||
status = "error" if block.is_error else "ok"
|
||||
block_parts.append(f"[tool_result: {block.tool_name} ({status})]")
|
||||
|
||||
if isinstance(block, TextBlock):
|
||||
if msg.role == MessageRole.User:
|
||||
user_texts.append(block.text)
|
||||
last_text = block.text
|
||||
|
||||
text_lower = block.text.lower()
|
||||
for kw in _PENDING_KEYWORDS:
|
||||
if kw in text_lower:
|
||||
for line in block.text.splitlines():
|
||||
if kw in line.lower():
|
||||
pending_items.append(line.strip())
|
||||
break
|
||||
break
|
||||
|
||||
for token in block.text.split():
|
||||
if "/" in token and _has_interesting_extension(token):
|
||||
cleaned = token.strip("(),;:\"'`")
|
||||
if "/" in cleaned:
|
||||
file_candidates.add(cleaned)
|
||||
|
||||
truncated = block.text[:80] + "..." if len(block.text) > 80 else block.text
|
||||
block_parts.append(truncated.replace("\n", " "))
|
||||
|
||||
if block_parts:
|
||||
timeline.append((msg.role.value, "; ".join(block_parts)))
|
||||
|
||||
recent_requests = user_texts[-3:] if user_texts else []
|
||||
|
||||
u = role_counts.get("user", 0)
|
||||
a = role_counts.get("assistant", 0)
|
||||
t = role_counts.get("tool", 0)
|
||||
tools_str = ", ".join(tool_names) if tool_names else "none"
|
||||
files_str = ", ".join(sorted(file_candidates)[:20]) if file_candidates else "none"
|
||||
last_text_truncated = last_text[:100].replace("\n", " ") if last_text else "n/a"
|
||||
|
||||
lines: list[str] = ["<summary>", "Conversation summary:"]
|
||||
lines.append(f"- Scope: {len(messages)} earlier messages compacted (user={u}, assistant={a}, tool={t}).")
|
||||
lines.append(f"- Tools mentioned: {tools_str}.")
|
||||
|
||||
if recent_requests:
|
||||
lines.append("- Recent user requests:")
|
||||
for req in recent_requests:
|
||||
truncated = req[:200] + "..." if len(req) > 200 else req
|
||||
lines.append(f" - {truncated}")
|
||||
|
||||
if pending_items:
|
||||
lines.append("- Pending work:")
|
||||
for item in pending_items[:5]:
|
||||
lines.append(f" - {item}")
|
||||
|
||||
lines.append(f"- Key files referenced: {files_str}.")
|
||||
lines.append(f"- Current work: {last_text_truncated}")
|
||||
lines.append("- Key timeline:")
|
||||
for role, block_summary in timeline[-10:]:
|
||||
lines.append(f" - {role}: {block_summary}")
|
||||
lines.append("</summary>")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _has_interesting_extension(token: str) -> bool:
|
||||
"""Check if a token looks like a file path with a known extension."""
|
||||
for ext in _FILE_EXTENSIONS:
|
||||
if token.endswith(ext):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def format_compact_summary(summary: str) -> str:
|
||||
"""Strip analysis tags, extract summary content, collapse blank lines.
|
||||
|
||||
Matches Rust's format_compact_summary:
|
||||
- Strip <analysis>...</analysis> block
|
||||
- Replace <summary>content</summary> with "Summary:\\ncontent"
|
||||
- Collapse consecutive blank lines
|
||||
"""
|
||||
# Strip <analysis>...</analysis> block
|
||||
result = re.sub(r"<analysis>.*?</analysis>\s*", "", summary, flags=re.DOTALL)
|
||||
# Replace <summary>content</summary> with "Summary:\ncontent"
|
||||
result = re.sub(
|
||||
r"<summary>(.*?)</summary>",
|
||||
lambda m: "Summary:\n" + m.group(1).strip(),
|
||||
result,
|
||||
flags=re.DOTALL,
|
||||
)
|
||||
# Collapse consecutive blank lines
|
||||
result = re.sub(r"\n{3,}", "\n\n", result)
|
||||
return result.strip()
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hanzoai"
|
||||
version = "3.0.0"
|
||||
version = "3.0.1"
|
||||
description = "The official Python library for the Hanzo API"
|
||||
dynamic = ["readme"]
|
||||
license = "BSD-3-Clause"
|
||||
|
||||
Reference in New Issue
Block a user