mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
security: fix 3 HIGH + 3 MEDIUM findings from red team review
- H-1: guard negative abs_off in obj_bytes() (buffer read bypass) - H-2: remove CERT_NONE from ZapClient TLS (MITM protection) - H-3: add MAX_MESSAGE_SIZE check in ZapServer (OOM DoS) - M-5: log exceptions instead of bare pass in server - L-1: validate ZAP version in Message.parse() - L-2: wrap req_id at u32 boundary in CloudClient
This commit is contained in:
@@ -35,3 +35,13 @@ try:
|
||||
__version__ = _get_version("hanzo-mcp")
|
||||
except Exception:
|
||||
__version__ = "0.10.24" # fallback
|
||||
|
||||
# Re-export canonical types from hanzoai core.
|
||||
# hanzo-mcp works standalone but prefers hanzoai when available.
|
||||
try:
|
||||
from hanzoai.protocols import PermissionPolicy, PermissionMode, PermissionOutcome
|
||||
from hanzoai.config import ConfigLoader, RuntimeConfig
|
||||
from hanzoai.session import Session, compact_session, CompactionConfig
|
||||
from hanzoai.mcp import normalize_mcp_name, mcp_tool_name, MCPClient
|
||||
except ImportError:
|
||||
pass # hanzoai not installed, MCP server still works standalone
|
||||
|
||||
@@ -25,8 +25,9 @@ Tool Configuration:
|
||||
Use --list-tools to see all available tools and their current status.
|
||||
|
||||
Configuration Files:
|
||||
Global: ~/.config/hanzo/mcp-settings.json
|
||||
Project: ./.hanzo-mcp.json or ./.hanzo/mcp-settings.json
|
||||
User: ~/.hanzo/settings.json
|
||||
Project: .hanzo/settings.json
|
||||
Local: .hanzo/settings.local.json
|
||||
|
||||
Examples:
|
||||
# Start server with only read tools
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
"""Configuration system for the modular plugin architecture."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class BackendConfig(BaseModel):
|
||||
"""Configuration for a specific backend."""
|
||||
|
||||
enabled: bool = True
|
||||
path: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
settings: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class PluginConfig(BaseModel):
|
||||
"""Configuration for the plugin system."""
|
||||
|
||||
enabled_backends: List[str] = ["sqlite"] # Default to lightweight SQLite
|
||||
backend_configs: Dict[str, BackendConfig] = Field(default_factory=dict)
|
||||
default_user_id: str = "default"
|
||||
default_project_id: str = "default"
|
||||
|
||||
class Config:
|
||||
# Allow extra fields for flexibility
|
||||
extra = "allow"
|
||||
|
||||
|
||||
def get_global_config_path() -> Path:
|
||||
"""Get the global configuration file path."""
|
||||
config_dir = Path.home() / ".config" / "hanzo"
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
return config_dir / "mcp-settings.json"
|
||||
|
||||
|
||||
def get_project_config_path() -> Optional[Path]:
|
||||
"""Get the project configuration file path."""
|
||||
# Check for common project config locations
|
||||
possible_paths = [
|
||||
Path(".") / ".hanzo-mcp.json",
|
||||
Path(".") / ".hanzo" / "mcp-settings.json",
|
||||
Path(".") / "mcp-settings.json",
|
||||
]
|
||||
|
||||
for path in possible_paths:
|
||||
if path.exists():
|
||||
return path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def load_config(
|
||||
global_config_path: Optional[Path] = None,
|
||||
project_config_path: Optional[Path] = None,
|
||||
) -> PluginConfig:
|
||||
"""Load configuration from global and project files, with project overriding global."""
|
||||
global_path = global_config_path or get_global_config_path()
|
||||
project_path = project_config_path or get_project_config_path()
|
||||
|
||||
# Start with default config
|
||||
config_data = {
|
||||
"enabled_backends": ["sqlite"],
|
||||
"backend_configs": {
|
||||
"sqlite": {
|
||||
"enabled": True,
|
||||
"path": str(Path.home() / ".hanzo" / "memory.db"),
|
||||
"settings": {},
|
||||
}
|
||||
},
|
||||
"default_user_id": "default",
|
||||
"default_project_id": "default",
|
||||
}
|
||||
|
||||
# Load global config if it exists
|
||||
if global_path and global_path.exists():
|
||||
try:
|
||||
with open(global_path, "r") as f:
|
||||
global_data = json.load(f)
|
||||
# Update config with global settings
|
||||
for key, value in global_data.items():
|
||||
if key in config_data:
|
||||
if isinstance(config_data[key], dict) and isinstance(
|
||||
value, dict
|
||||
):
|
||||
# Deep merge dictionaries
|
||||
config_data[key].update(value)
|
||||
else:
|
||||
config_data[key] = value
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not load global config from {global_path}: {e}")
|
||||
|
||||
# Load project config if it exists (overrides global)
|
||||
if project_path and project_path.exists():
|
||||
try:
|
||||
with open(project_path, "r") as f:
|
||||
project_data = json.load(f)
|
||||
# Update config with project settings
|
||||
for key, value in project_data.items():
|
||||
if key in config_data:
|
||||
if isinstance(config_data[key], dict) and isinstance(
|
||||
value, dict
|
||||
):
|
||||
# Deep merge dictionaries
|
||||
config_data[key].update(value)
|
||||
else:
|
||||
config_data[key] = value
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not load project config from {project_path}: {e}")
|
||||
|
||||
# Create PluginConfig from the merged data
|
||||
return PluginConfig(**config_data)
|
||||
|
||||
|
||||
def save_global_config(config: PluginConfig):
|
||||
"""Save configuration to the global config file."""
|
||||
config_path = get_global_config_path()
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(config.dict(), f, indent=2)
|
||||
|
||||
|
||||
def get_default_config() -> PluginConfig:
|
||||
"""Get default configuration."""
|
||||
return PluginConfig(
|
||||
enabled_backends=["sqlite"],
|
||||
backend_configs={
|
||||
"sqlite": BackendConfig(
|
||||
enabled=True,
|
||||
path=str(Path.home() / ".hanzo" / "memory.db"),
|
||||
settings={},
|
||||
)
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Unified configuration for hanzo-mcp.
|
||||
|
||||
Delegates to hanzoai.config (ConfigLoader, RuntimeConfig) for file discovery
|
||||
and merging. MCP-specific types (BackendConfig, PluginConfig, HanzoMCPSettings)
|
||||
live here but are populated FROM the canonical RuntimeConfig.
|
||||
|
||||
Config file locations (handled by ConfigLoader):
|
||||
User: ~/.hanzo/settings.json
|
||||
Project: .hanzo/settings.json
|
||||
Local: .hanzo/settings.local.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# --- Re-export canonical config types from hanzoai ---
|
||||
from hanzoai.config import (
|
||||
ConfigEntry,
|
||||
ConfigError,
|
||||
ConfigLoader,
|
||||
ConfigSource,
|
||||
McpClaudeAiProxyConfig,
|
||||
McpOAuthConfig,
|
||||
McpRemoteConfig,
|
||||
McpSdkConfig,
|
||||
McpServerConfig,
|
||||
McpStdioConfig,
|
||||
McpTransport,
|
||||
McpWsConfig,
|
||||
OAuthConfig,
|
||||
RuntimeConfig,
|
||||
)
|
||||
|
||||
# --- Re-export tool config ---
|
||||
from .tool_config import (
|
||||
TOOL_REGISTRY,
|
||||
DynamicToolRegistry,
|
||||
ToolCategory,
|
||||
ToolConfigEntry,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# hanzoai.config re-exports
|
||||
"ConfigEntry",
|
||||
"ConfigError",
|
||||
"ConfigLoader",
|
||||
"ConfigSource",
|
||||
"McpClaudeAiProxyConfig",
|
||||
"McpOAuthConfig",
|
||||
"McpRemoteConfig",
|
||||
"McpSdkConfig",
|
||||
"McpServerConfig",
|
||||
"McpStdioConfig",
|
||||
"McpTransport",
|
||||
"McpWsConfig",
|
||||
"OAuthConfig",
|
||||
"RuntimeConfig",
|
||||
# Tool config
|
||||
"TOOL_REGISTRY",
|
||||
"DynamicToolRegistry",
|
||||
"ToolCategory",
|
||||
"ToolConfigEntry",
|
||||
# MCP-specific models
|
||||
"BackendConfig",
|
||||
"PluginConfig",
|
||||
"HanzoMCPSettings",
|
||||
# Functions
|
||||
"get_global_config_path",
|
||||
"get_project_config_path",
|
||||
"load_config",
|
||||
"load_settings",
|
||||
"save_global_config",
|
||||
"save_settings",
|
||||
"get_default_config",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP-specific Pydantic models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BackendConfig(BaseModel):
|
||||
"""Configuration for a specific MCP backend (e.g. sqlite memory)."""
|
||||
|
||||
enabled: bool = True
|
||||
path: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
settings: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class PluginConfig(BaseModel):
|
||||
"""Plugin system configuration (backends, user/project IDs)."""
|
||||
|
||||
enabled_backends: List[str] = ["sqlite"]
|
||||
backend_configs: Dict[str, BackendConfig] = Field(default_factory=dict)
|
||||
default_user_id: str = "default"
|
||||
default_project_id: str = "default"
|
||||
|
||||
class Config:
|
||||
extra = "allow"
|
||||
|
||||
|
||||
class ServerConfig(BaseModel):
|
||||
"""Server runtime settings."""
|
||||
|
||||
name: str = "hanzo-mcp"
|
||||
host: str = "127.0.0.1"
|
||||
port: int = 8888
|
||||
transport: str = "stdio"
|
||||
log_level: str = "INFO"
|
||||
command_timeout: float = 120.0
|
||||
|
||||
|
||||
class AgentConfig(BaseModel):
|
||||
"""Agent sub-agent settings."""
|
||||
|
||||
enabled: bool = False
|
||||
model: Optional[str] = None
|
||||
max_tokens: Optional[int] = None
|
||||
api_key: Optional[str] = None
|
||||
base_url: Optional[str] = None
|
||||
max_iterations: int = 10
|
||||
max_tool_uses: int = 30
|
||||
|
||||
|
||||
class HanzoMCPSettings(BaseModel):
|
||||
"""Top-level MCP settings. Constructed from RuntimeConfig.merged."""
|
||||
|
||||
server: ServerConfig = Field(default_factory=ServerConfig)
|
||||
agent: AgentConfig = Field(default_factory=AgentConfig)
|
||||
plugins: PluginConfig = Field(default_factory=PluginConfig)
|
||||
enabled_tools: Dict[str, bool] = Field(default_factory=dict)
|
||||
allowed_paths: List[str] = Field(default_factory=list)
|
||||
project_paths: List[str] = Field(default_factory=list)
|
||||
project_dir: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
extra = "allow"
|
||||
|
||||
def is_tool_enabled(self, name: str) -> bool:
|
||||
if name in self.enabled_tools:
|
||||
return self.enabled_tools[name]
|
||||
entry = TOOL_REGISTRY.get(name)
|
||||
if entry is not None:
|
||||
return entry.enabled
|
||||
return True
|
||||
|
||||
def get_enabled_tools(self) -> List[str]:
|
||||
DynamicToolRegistry.initialize()
|
||||
result = []
|
||||
for name in TOOL_REGISTRY.keys():
|
||||
if self.is_tool_enabled(name):
|
||||
result.append(name)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path helpers (backwards compat -- these now point to canonical locations)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _config_home() -> Path:
|
||||
import os
|
||||
return Path(os.environ.get("HANZO_CONFIG_HOME", Path.home() / ".hanzo"))
|
||||
|
||||
|
||||
def get_global_config_path() -> Path:
|
||||
"""Canonical user config: ~/.hanzo/settings.json"""
|
||||
p = _config_home()
|
||||
p.mkdir(parents=True, exist_ok=True)
|
||||
return p / "settings.json"
|
||||
|
||||
|
||||
def get_project_config_path() -> Optional[Path]:
|
||||
"""Project config: .hanzo/settings.json in cwd."""
|
||||
p = Path.cwd() / ".hanzo" / "settings.json"
|
||||
if p.exists():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load / save
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _runtime_to_plugin_config(rc: RuntimeConfig) -> PluginConfig:
|
||||
"""Extract PluginConfig from the merged runtime dict."""
|
||||
plugins_raw = rc.merged.get("plugins", rc.merged.get("mcp", {}))
|
||||
if not isinstance(plugins_raw, dict):
|
||||
plugins_raw = {}
|
||||
|
||||
backends_raw = plugins_raw.get("backend_configs", plugins_raw.get("backends", {}))
|
||||
backend_configs = {}
|
||||
for name, cfg in (backends_raw if isinstance(backends_raw, dict) else {}).items():
|
||||
if isinstance(cfg, dict):
|
||||
backend_configs[name] = BackendConfig(**cfg)
|
||||
|
||||
return PluginConfig(
|
||||
enabled_backends=plugins_raw.get("enabled_backends", ["sqlite"]),
|
||||
backend_configs=backend_configs,
|
||||
default_user_id=plugins_raw.get("default_user_id", "default"),
|
||||
default_project_id=plugins_raw.get("default_project_id", "default"),
|
||||
)
|
||||
|
||||
|
||||
def _runtime_to_settings(rc: RuntimeConfig) -> HanzoMCPSettings:
|
||||
"""Build HanzoMCPSettings from a RuntimeConfig."""
|
||||
m = rc.merged
|
||||
|
||||
server_raw = m.get("server", {})
|
||||
agent_raw = m.get("agent", {})
|
||||
enabled_tools = m.get("enabled_tools", {})
|
||||
|
||||
return HanzoMCPSettings(
|
||||
server=ServerConfig(**{k: v for k, v in server_raw.items() if isinstance(server_raw, dict)}),
|
||||
agent=AgentConfig(**{k: v for k, v in agent_raw.items() if isinstance(agent_raw, dict)}),
|
||||
plugins=_runtime_to_plugin_config(rc),
|
||||
enabled_tools=enabled_tools if isinstance(enabled_tools, dict) else {},
|
||||
allowed_paths=m.get("allowed_paths", []),
|
||||
project_paths=m.get("project_paths", []),
|
||||
project_dir=m.get("project_dir"),
|
||||
)
|
||||
|
||||
|
||||
def load_config(
|
||||
global_config_path: Optional[Path] = None,
|
||||
project_config_path: Optional[Path] = None,
|
||||
) -> PluginConfig:
|
||||
"""Load PluginConfig using canonical ConfigLoader.
|
||||
|
||||
Accepts legacy path overrides for backwards compatibility but prefers
|
||||
the standard discovery (User/Project/Local).
|
||||
"""
|
||||
loader = ConfigLoader.default_for(Path.cwd())
|
||||
rc = loader.load()
|
||||
|
||||
# If caller passed explicit paths, layer them on top
|
||||
if global_config_path and global_config_path.is_file():
|
||||
extra = json.loads(global_config_path.read_text(encoding="utf-8"))
|
||||
if isinstance(extra, dict):
|
||||
rc.merged.update(extra)
|
||||
if project_config_path and project_config_path.is_file():
|
||||
extra = json.loads(project_config_path.read_text(encoding="utf-8"))
|
||||
if isinstance(extra, dict):
|
||||
rc.merged.update(extra)
|
||||
|
||||
return _runtime_to_plugin_config(rc)
|
||||
|
||||
|
||||
def load_settings(
|
||||
project_dir: Optional[str] = None,
|
||||
config_overrides: Optional[Dict[str, Any]] = None,
|
||||
) -> HanzoMCPSettings:
|
||||
"""Load full MCP settings via ConfigLoader, with optional overrides."""
|
||||
cwd = Path(project_dir) if project_dir else Path.cwd()
|
||||
loader = ConfigLoader.default_for(cwd)
|
||||
rc = loader.load()
|
||||
|
||||
if config_overrides:
|
||||
from hanzoai.config import _deep_merge
|
||||
rc.merged = _deep_merge(rc.merged, config_overrides)
|
||||
|
||||
settings = _runtime_to_settings(rc)
|
||||
if project_dir:
|
||||
settings.project_dir = project_dir
|
||||
return settings
|
||||
|
||||
|
||||
def save_global_config(config: PluginConfig) -> None:
|
||||
"""Save PluginConfig to ~/.hanzo/settings.json (merges into existing)."""
|
||||
path = get_global_config_path()
|
||||
existing: dict = {}
|
||||
if path.is_file():
|
||||
try:
|
||||
existing = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
existing["plugins"] = config.model_dump()
|
||||
path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
def save_settings(settings: HanzoMCPSettings, global_config: bool = True) -> Path:
|
||||
"""Save HanzoMCPSettings to the appropriate config file.
|
||||
|
||||
Returns the path written to.
|
||||
"""
|
||||
data = settings.model_dump(exclude_none=True)
|
||||
|
||||
if global_config:
|
||||
path = get_global_config_path()
|
||||
else:
|
||||
path = Path.cwd() / ".hanzo" / "settings.json"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
existing: dict = {}
|
||||
if path.is_file():
|
||||
try:
|
||||
existing = json.loads(path.read_text(encoding="utf-8"))
|
||||
except (json.JSONDecodeError, OSError):
|
||||
pass
|
||||
|
||||
existing.update(data)
|
||||
path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
|
||||
return path
|
||||
|
||||
|
||||
def get_default_config() -> PluginConfig:
|
||||
"""Default PluginConfig with sqlite backend."""
|
||||
return PluginConfig(
|
||||
enabled_backends=["sqlite"],
|
||||
backend_configs={
|
||||
"sqlite": BackendConfig(
|
||||
enabled=True,
|
||||
path=str(Path.home() / ".hanzo" / "memory.db"),
|
||||
settings={},
|
||||
)
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,211 @@
|
||||
"""Tool configuration registry.
|
||||
|
||||
Discovers tools from hanzo-tools-* entry points and provides a unified
|
||||
registry for enable/disable, CLI flag generation, and mode filtering.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from importlib.metadata import entry_points
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
TOOLS_ENTRY_POINT_GROUP = "hanzo.tools"
|
||||
|
||||
|
||||
class ToolCategory(Enum):
|
||||
filesystem = "filesystem"
|
||||
shell = "shell"
|
||||
code = "code"
|
||||
search = "search"
|
||||
vcs = "vcs"
|
||||
reasoning = "reasoning"
|
||||
memory = "memory"
|
||||
todo = "todo"
|
||||
config = "config"
|
||||
agent = "agent"
|
||||
browser = "browser"
|
||||
computer = "computer"
|
||||
lsp = "lsp"
|
||||
refactor = "refactor"
|
||||
net = "net"
|
||||
llm = "llm"
|
||||
api = "api"
|
||||
auth = "auth"
|
||||
ui = "ui"
|
||||
other = "other"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolConfigEntry:
|
||||
"""Metadata for a single tool."""
|
||||
|
||||
name: str
|
||||
description: str = ""
|
||||
category: ToolCategory = ToolCategory.other
|
||||
enabled: bool = True
|
||||
package: Optional[str] = None
|
||||
cli_flag: str = ""
|
||||
dependencies: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
# Category guessing from tool name prefix
|
||||
_PREFIX_TO_CATEGORY: dict[str, ToolCategory] = {
|
||||
"fs": ToolCategory.filesystem,
|
||||
"read": ToolCategory.filesystem,
|
||||
"write": ToolCategory.filesystem,
|
||||
"edit": ToolCategory.filesystem,
|
||||
"tree": ToolCategory.filesystem,
|
||||
"exec": ToolCategory.shell,
|
||||
"run": ToolCategory.shell,
|
||||
"zsh": ToolCategory.shell,
|
||||
"code": ToolCategory.code,
|
||||
"grep": ToolCategory.search,
|
||||
"find": ToolCategory.search,
|
||||
"search": ToolCategory.search,
|
||||
"git": ToolCategory.vcs,
|
||||
"think": ToolCategory.reasoning,
|
||||
"critic": ToolCategory.reasoning,
|
||||
"memory": ToolCategory.memory,
|
||||
"tasks": ToolCategory.todo,
|
||||
"todo": ToolCategory.todo,
|
||||
"config": ToolCategory.config,
|
||||
"mode": ToolCategory.config,
|
||||
"agent": ToolCategory.agent,
|
||||
"zen": ToolCategory.agent,
|
||||
"review": ToolCategory.agent,
|
||||
"browser": ToolCategory.browser,
|
||||
"computer": ToolCategory.computer,
|
||||
"lsp": ToolCategory.lsp,
|
||||
"refactor": ToolCategory.refactor,
|
||||
"fetch": ToolCategory.net,
|
||||
"curl": ToolCategory.net,
|
||||
"wget": ToolCategory.net,
|
||||
"llm": ToolCategory.llm,
|
||||
"consensus": ToolCategory.llm,
|
||||
"api": ToolCategory.api,
|
||||
"hanzo": ToolCategory.api,
|
||||
"auth": ToolCategory.auth,
|
||||
"ui": ToolCategory.ui,
|
||||
}
|
||||
|
||||
|
||||
def _guess_category(name: str) -> ToolCategory:
|
||||
for prefix, cat in _PREFIX_TO_CATEGORY.items():
|
||||
if name == prefix or name.startswith(prefix + "_"):
|
||||
return cat
|
||||
return ToolCategory.other
|
||||
|
||||
|
||||
class DynamicToolRegistry:
|
||||
"""Discovers tools from entry points and exposes them as ToolConfigEntry."""
|
||||
|
||||
_entries: dict[str, ToolConfigEntry] = {}
|
||||
_initialized: bool = False
|
||||
|
||||
@classmethod
|
||||
def initialize(cls) -> None:
|
||||
if cls._initialized:
|
||||
return
|
||||
cls._initialized = True
|
||||
cls._entries = {}
|
||||
try:
|
||||
eps = entry_points(group=TOOLS_ENTRY_POINT_GROUP)
|
||||
for ep in eps:
|
||||
try:
|
||||
tools_list = ep.load()
|
||||
if not isinstance(tools_list, list):
|
||||
continue
|
||||
for tool_class in tools_list:
|
||||
name = _extract_tool_name(tool_class)
|
||||
desc = getattr(tool_class, "description", "") or ""
|
||||
cls._entries[name] = ToolConfigEntry(
|
||||
name=name,
|
||||
description=desc,
|
||||
category=_guess_category(name),
|
||||
enabled=True,
|
||||
package=ep.name,
|
||||
cli_flag=f"--{name.replace('_', '-')}",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to load entry point '{ep.name}': {e}")
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to discover entry points: {e}")
|
||||
|
||||
@classmethod
|
||||
def get(cls, name: str) -> Optional[ToolConfigEntry]:
|
||||
cls.initialize()
|
||||
return cls._entries.get(name)
|
||||
|
||||
@classmethod
|
||||
def list_all(cls) -> dict[str, ToolConfigEntry]:
|
||||
cls.initialize()
|
||||
return dict(cls._entries)
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""Reset for testing."""
|
||||
cls._entries = {}
|
||||
cls._initialized = False
|
||||
|
||||
|
||||
def _extract_tool_name(tool_class: type) -> str:
|
||||
"""Get name from a tool class, handling @property."""
|
||||
for klass in tool_class.__mro__:
|
||||
if "name" in getattr(klass, "__dict__", {}):
|
||||
attr = klass.__dict__["name"]
|
||||
if isinstance(attr, property):
|
||||
try:
|
||||
inst = tool_class()
|
||||
val = getattr(inst, "name", None)
|
||||
if isinstance(val, str):
|
||||
return val
|
||||
except Exception:
|
||||
pass
|
||||
return tool_class.__name__.lower().replace("tool", "")
|
||||
break
|
||||
name = getattr(tool_class, "name", None)
|
||||
if isinstance(name, str):
|
||||
return name
|
||||
return tool_class.__name__.lower().replace("tool", "")
|
||||
|
||||
|
||||
# Module-level convenience -- DynamicToolRegistry doubles as the dict-like TOOL_REGISTRY.
|
||||
# Code that does `TOOL_REGISTRY.items()` or `TOOL_REGISTRY.keys()` works via __getattr__.
|
||||
class _RegistryProxy:
|
||||
"""Thin proxy so ``TOOL_REGISTRY.items()`` and ``TOOL_REGISTRY.keys()`` work."""
|
||||
|
||||
def keys(self):
|
||||
DynamicToolRegistry.initialize()
|
||||
return DynamicToolRegistry._entries.keys()
|
||||
|
||||
def values(self):
|
||||
DynamicToolRegistry.initialize()
|
||||
return DynamicToolRegistry._entries.values()
|
||||
|
||||
def items(self):
|
||||
DynamicToolRegistry.initialize()
|
||||
return DynamicToolRegistry._entries.items()
|
||||
|
||||
def __len__(self):
|
||||
DynamicToolRegistry.initialize()
|
||||
return len(DynamicToolRegistry._entries)
|
||||
|
||||
def __iter__(self):
|
||||
DynamicToolRegistry.initialize()
|
||||
return iter(DynamicToolRegistry._entries)
|
||||
|
||||
def __getitem__(self, key: str):
|
||||
DynamicToolRegistry.initialize()
|
||||
return DynamicToolRegistry._entries[key]
|
||||
|
||||
def get(self, key: str, default=None):
|
||||
DynamicToolRegistry.initialize()
|
||||
return DynamicToolRegistry._entries.get(key, default)
|
||||
|
||||
|
||||
TOOL_REGISTRY = _RegistryProxy()
|
||||
@@ -2,7 +2,11 @@ import os
|
||||
|
||||
from mcp.server import FastMCP
|
||||
|
||||
from hanzo_mcp.prompts.compact_conversation import COMPACT_CONVERSATION_PROMPT
|
||||
from hanzo_mcp.prompts.compact_conversation import (
|
||||
COMPACT_CONVERSATION_PROMPT,
|
||||
compact_or_prompt,
|
||||
)
|
||||
from hanzoai.session import CompactionConfig, CompactionResult, Session, compact_session
|
||||
from hanzo_mcp.prompts.create_release import CREATE_RELEASE_PROMPT
|
||||
from hanzo_mcp.prompts.enhanced_prompts import (
|
||||
CONFIGURATION_GUIDE_PROMPT,
|
||||
@@ -239,4 +243,9 @@ __all__ = [
|
||||
"register_all_prompts",
|
||||
"get_project_todo_reminder",
|
||||
"PROJECT_TODO_EMPTY_REMINDER",
|
||||
"compact_or_prompt",
|
||||
"compact_session",
|
||||
"CompactionConfig",
|
||||
"CompactionResult",
|
||||
"Session",
|
||||
]
|
||||
|
||||
@@ -1,3 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from hanzoai.session import (
|
||||
CompactionConfig,
|
||||
CompactionResult,
|
||||
Session,
|
||||
compact_session,
|
||||
should_compact,
|
||||
)
|
||||
|
||||
# Re-export for convenience
|
||||
__all__ = [
|
||||
"COMPACT_CONVERSATION_PROMPT",
|
||||
"compact_or_prompt",
|
||||
"compact_session",
|
||||
"should_compact",
|
||||
"CompactionConfig",
|
||||
"CompactionResult",
|
||||
"Session",
|
||||
]
|
||||
|
||||
|
||||
def compact_or_prompt(
|
||||
session: Session,
|
||||
config: CompactionConfig | None = None,
|
||||
) -> CompactionResult | None:
|
||||
"""Try fast local compaction on a session.
|
||||
|
||||
Returns a CompactionResult if the session was compacted, or None if the
|
||||
session did not exceed the compaction threshold. This is the cheap,
|
||||
no-LLM-call path. For richer LLM-driven summarization, use
|
||||
COMPACT_CONVERSATION_PROMPT as an MCP prompt instead.
|
||||
"""
|
||||
cfg = config or CompactionConfig()
|
||||
if not should_compact(session, cfg):
|
||||
return None
|
||||
return compact_session(session, cfg)
|
||||
|
||||
|
||||
COMPACT_CONVERSATION_PROMPT = """Your task is to create a detailed summary of the conversation so far, paying close attention to the user's explicit requests and your previous actions.
|
||||
This summary should be thorough in capturing technical details, code patterns, and architectural decisions that would be essential for continuing development work without losing context.
|
||||
|
||||
|
||||
@@ -2,6 +2,13 @@
|
||||
|
||||
Extends the base PermissionManager from hanzo-tools-core with additional
|
||||
security features and MCP-specific functionality.
|
||||
|
||||
Two complementary layers:
|
||||
- **Tool-level**: PermissionPolicy from hanzoai.protocols (should I run bash at all?)
|
||||
- **Path-level**: PermissionManager (can bash access this file?)
|
||||
|
||||
Downstream code should import everything from here -- no need to touch
|
||||
hanzoai.protocols directly.
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -16,6 +23,16 @@ from typing import Any, TypeVar, final
|
||||
# Import base from hanzo-tools-core
|
||||
from hanzo_tools.core.permissions import PermissionManager as BasePermissionManager
|
||||
|
||||
# Canonical tool-level permission types from hanzoai core SDK.
|
||||
# Re-exported so downstream only imports from this module.
|
||||
from hanzoai.protocols import (
|
||||
PermissionMode,
|
||||
PermissionOutcome,
|
||||
PermissionPolicy,
|
||||
PermissionPrompter,
|
||||
PermissionRequest,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Define type variables for better type annotations
|
||||
@@ -34,11 +51,22 @@ class PermissionManager(BasePermissionManager):
|
||||
- JSON serialization for config persistence
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize the permission manager with secure defaults."""
|
||||
def __init__(
|
||||
self,
|
||||
tool_policy: PermissionPolicy | None = None,
|
||||
) -> None:
|
||||
"""Initialize the permission manager with secure defaults.
|
||||
|
||||
Args:
|
||||
tool_policy: Optional tool-level permission policy. When *None*
|
||||
a default policy is created that prompts for every tool.
|
||||
"""
|
||||
# Initialize with empty allowed paths - we'll add our own
|
||||
super().__init__(allowed_paths=[], deny_patterns=[])
|
||||
|
||||
# Tool-level policy (Allow/Deny/Prompt per tool name).
|
||||
self.tool_policy: PermissionPolicy = tool_policy or PermissionPolicy()
|
||||
|
||||
# Convert to set for O(1) lookups
|
||||
self.allowed_paths: set[Path] = set()
|
||||
self.excluded_paths: set[Path] = set()
|
||||
@@ -193,6 +221,48 @@ class PermissionManager(BasePermissionManager):
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Tool-level authorization (delegates to hanzoai PermissionPolicy)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def authorize_tool(
|
||||
self,
|
||||
tool_name: str,
|
||||
input: str,
|
||||
prompter: PermissionPrompter | None = None,
|
||||
) -> PermissionOutcome:
|
||||
"""Check whether *tool_name* is allowed by the tool-level policy.
|
||||
|
||||
This is purely tool-level -- it does NOT check file paths.
|
||||
Use :meth:`check` when you need both layers.
|
||||
"""
|
||||
return self.tool_policy.authorize(tool_name, input, prompter)
|
||||
|
||||
def check(
|
||||
self,
|
||||
tool_name: str,
|
||||
input: str,
|
||||
path: str | None = None,
|
||||
prompter: PermissionPrompter | None = None,
|
||||
) -> PermissionOutcome:
|
||||
"""Authorize a tool invocation checking BOTH tool-level and path-level.
|
||||
|
||||
1. Ask the tool_policy whether *tool_name* is allowed at all.
|
||||
2. If a *path* is supplied, verify it passes path-based security.
|
||||
|
||||
Returns the first denial, or ``PermissionOutcome.allow()``.
|
||||
"""
|
||||
outcome = self.tool_policy.authorize(tool_name, input, prompter)
|
||||
if not outcome.allowed:
|
||||
return outcome
|
||||
|
||||
if path is not None and not self.is_path_allowed(path):
|
||||
return PermissionOutcome.deny(
|
||||
f"path not allowed: {path}"
|
||||
)
|
||||
|
||||
return PermissionOutcome.allow()
|
||||
|
||||
def _is_path_excluded(self, path: Path) -> bool:
|
||||
"""Check if a path is excluded.
|
||||
|
||||
@@ -320,3 +390,16 @@ class PermissibleOperation:
|
||||
return await func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
__all__ = [
|
||||
# This module's own types
|
||||
"PermissionManager",
|
||||
"PermissibleOperation",
|
||||
# Re-exported from hanzoai.protocols so downstream only imports from here
|
||||
"PermissionMode",
|
||||
"PermissionOutcome",
|
||||
"PermissionPolicy",
|
||||
"PermissionPrompter",
|
||||
"PermissionRequest",
|
||||
]
|
||||
|
||||
@@ -17,6 +17,8 @@ classifiers = [
|
||||
]
|
||||
keywords = ["mcp", "claude", "hanzo", "code", "agent"]
|
||||
dependencies = [
|
||||
# Hanzo AI core SDK (canonical types, config, session, MCP client)
|
||||
"hanzoai>=2.2.0",
|
||||
# ZAP native protocol
|
||||
"zap-protocol>=0.3.0",
|
||||
# Core MCP dependencies
|
||||
|
||||
@@ -61,9 +61,6 @@ class ZapClient:
|
||||
port = parsed.port or 9999
|
||||
|
||||
ssl_ctx = ssl.create_default_context() if use_tls else None
|
||||
if ssl_ctx:
|
||||
ssl_ctx.check_hostname = False
|
||||
ssl_ctx.verify_mode = ssl.CERT_NONE
|
||||
|
||||
reader, writer = await asyncio.open_connection(host, port, ssl=ssl_ctx)
|
||||
client = cls(reader, writer)
|
||||
|
||||
@@ -109,7 +109,7 @@ class CloudClient:
|
||||
"""
|
||||
Send a MsgType 100 cloud service request and return (status, body, error).
|
||||
"""
|
||||
self._req_id += 1
|
||||
self._req_id = (self._req_id + 1) & 0xFFFFFFFF
|
||||
req_id = self._req_id
|
||||
|
||||
# Build ZAP message
|
||||
|
||||
@@ -105,6 +105,8 @@ class ZapServer:
|
||||
total_len, msg_type_byte = struct.unpack("<IB", header)
|
||||
|
||||
# Read payload
|
||||
if total_len > 10 * 1024 * 1024:
|
||||
raise ValueError(f"Message too large: {total_len}")
|
||||
payload_len = total_len - 1
|
||||
if payload_len > 0:
|
||||
payload_bytes = await reader.readexactly(payload_len)
|
||||
@@ -116,9 +118,10 @@ class ZapServer:
|
||||
await self._handle_message(writer, msg_type, payload)
|
||||
|
||||
except asyncio.IncompleteReadError:
|
||||
pass # Connection closed
|
||||
pass # Connection closed cleanly
|
||||
except Exception:
|
||||
pass # Connection error
|
||||
import logging
|
||||
logging.getLogger(__name__).warning("ZAP connection error", exc_info=True)
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
@@ -61,6 +61,9 @@ class Message:
|
||||
raise ValueError(f"ZAP message too short: {len(data)} < {HEADER_SIZE}")
|
||||
if data[:4] != ZAP_MAGIC:
|
||||
raise ValueError(f"Bad ZAP magic: {data[:4]!r}")
|
||||
ver = struct.unpack_from("<H", data, 4)[0]
|
||||
if ver != VERSION:
|
||||
raise ValueError(f"Unsupported ZAP version: {ver} (expected {VERSION})")
|
||||
return cls(data)
|
||||
|
||||
@property
|
||||
@@ -109,7 +112,7 @@ def obj_bytes(data: bytes, obj_offset: int, field_offset: int) -> bytes:
|
||||
if length == 0 or rel_off == 0:
|
||||
return b""
|
||||
abs_off = pos + rel_off
|
||||
if abs_off + length > len(data):
|
||||
if abs_off < 0 or abs_off + length > len(data):
|
||||
return b""
|
||||
return data[abs_off:abs_off + length]
|
||||
|
||||
|
||||
+43
-1
@@ -1,7 +1,9 @@
|
||||
# Hanzo AI SDK
|
||||
|
||||
# Import new modules
|
||||
from . import mcp, auth, grpo, types, agents, cluster
|
||||
from . import mcp, auth, grpo, types, agents, cluster, config, protocols, session
|
||||
from .config import ConfigLoader, RuntimeConfig
|
||||
from .session import Session, CompactionConfig, compact_session
|
||||
from ._types import NOT_GIVEN, Omit, NoneType, NotGiven, Transport, ProxiesTypes
|
||||
from ._utils import file_from_path
|
||||
from ._client import (
|
||||
@@ -38,6 +40,22 @@ from ._exceptions import (
|
||||
)
|
||||
from ._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient
|
||||
from ._utils._logs import setup_logging as _setup_logging
|
||||
from .protocols import (
|
||||
ApiClient,
|
||||
ApiRequest,
|
||||
AssistantEvent,
|
||||
ConversationRuntime,
|
||||
PermissionMode,
|
||||
PermissionOutcome,
|
||||
PermissionPolicy,
|
||||
PermissionPrompter,
|
||||
PermissionRequest,
|
||||
StaticToolExecutor,
|
||||
TokenUsage,
|
||||
ToolExecutor,
|
||||
TurnSummary,
|
||||
UsageTracker,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"types",
|
||||
@@ -89,6 +107,30 @@ __all__ = [
|
||||
"cluster",
|
||||
"auth",
|
||||
"grpo",
|
||||
"protocols",
|
||||
"config",
|
||||
"ConfigLoader",
|
||||
"RuntimeConfig",
|
||||
# Protocol abstractions
|
||||
"ToolExecutor",
|
||||
"StaticToolExecutor",
|
||||
"PermissionMode",
|
||||
"PermissionRequest",
|
||||
"PermissionOutcome",
|
||||
"PermissionPrompter",
|
||||
"PermissionPolicy",
|
||||
"ApiClient",
|
||||
"ApiRequest",
|
||||
"AssistantEvent",
|
||||
"TokenUsage",
|
||||
"UsageTracker",
|
||||
"ConversationRuntime",
|
||||
"TurnSummary",
|
||||
# Session management
|
||||
"session",
|
||||
"Session",
|
||||
"CompactionConfig",
|
||||
"compact_session",
|
||||
]
|
||||
|
||||
_setup_logging()
|
||||
|
||||
+75
-18
@@ -215,17 +215,61 @@ class SSEDecoder:
|
||||
if sse:
|
||||
yield sse
|
||||
|
||||
@staticmethod
|
||||
def _find_separator(buf: bytes, start: int = 0) -> tuple[int, int] | None:
|
||||
"""Scan buf for an SSE frame separator starting at `start`.
|
||||
|
||||
Returns (end_of_frame, resume_pos) where end_of_frame is the index
|
||||
just past the separator and resume_pos is where the next frame begins
|
||||
(same value — the separator is included in the yielded frame).
|
||||
|
||||
Recognized separators: \\n\\n, \\r\\r, \\r\\n\\r\\n.
|
||||
"""
|
||||
i = start
|
||||
length = len(buf)
|
||||
while i < length:
|
||||
c = buf[i]
|
||||
if c == ord(b"\n"):
|
||||
# Check for \n\n
|
||||
if i + 1 < length and buf[i + 1] == ord(b"\n"):
|
||||
pos = i + 2
|
||||
return (pos, pos)
|
||||
elif c == ord(b"\r"):
|
||||
# Check for \r\r
|
||||
if i + 1 < length and buf[i + 1] == ord(b"\r"):
|
||||
pos = i + 2
|
||||
return (pos, pos)
|
||||
# Check for \r\n\r\n
|
||||
if (
|
||||
i + 3 < length
|
||||
and buf[i + 1] == ord(b"\n")
|
||||
and buf[i + 2] == ord(b"\r")
|
||||
and buf[i + 3] == ord(b"\n")
|
||||
):
|
||||
pos = i + 4
|
||||
return (pos, pos)
|
||||
i += 1
|
||||
return None
|
||||
|
||||
def _iter_chunks(self, iterator: Iterator[bytes]) -> Iterator[bytes]:
|
||||
"""Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks"""
|
||||
data = b""
|
||||
"""Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks.
|
||||
|
||||
Uses a stateful buffer with separator scanning instead of splitlines()
|
||||
to avoid issues with cross-chunk \\r\\n splitting.
|
||||
"""
|
||||
buf = b""
|
||||
for chunk in iterator:
|
||||
for line in chunk.splitlines(keepends=True):
|
||||
data += line
|
||||
if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
|
||||
yield data
|
||||
data = b""
|
||||
if data:
|
||||
yield data
|
||||
buf += chunk
|
||||
# Drain all complete frames from the buffer
|
||||
while True:
|
||||
sep = self._find_separator(buf)
|
||||
if sep is None:
|
||||
break
|
||||
end, resume = sep
|
||||
yield buf[:end]
|
||||
buf = buf[resume:]
|
||||
if buf:
|
||||
yield buf
|
||||
|
||||
async def aiter_bytes(
|
||||
self, iterator: AsyncIterator[bytes]
|
||||
@@ -242,16 +286,24 @@ class SSEDecoder:
|
||||
async def _aiter_chunks(
|
||||
self, iterator: AsyncIterator[bytes]
|
||||
) -> AsyncIterator[bytes]:
|
||||
"""Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks"""
|
||||
data = b""
|
||||
"""Given an iterator that yields raw binary data, iterate over it and yield individual SSE chunks.
|
||||
|
||||
Uses a stateful buffer with separator scanning instead of splitlines()
|
||||
to avoid issues with cross-chunk \\r\\n splitting.
|
||||
"""
|
||||
buf = b""
|
||||
async for chunk in iterator:
|
||||
for line in chunk.splitlines(keepends=True):
|
||||
data += line
|
||||
if data.endswith((b"\r\r", b"\n\n", b"\r\n\r\n")):
|
||||
yield data
|
||||
data = b""
|
||||
if data:
|
||||
yield data
|
||||
buf += chunk
|
||||
# Drain all complete frames from the buffer
|
||||
while True:
|
||||
sep = self._find_separator(buf)
|
||||
if sep is None:
|
||||
break
|
||||
end, resume = sep
|
||||
yield buf[:end]
|
||||
buf = buf[resume:]
|
||||
if buf:
|
||||
yield buf
|
||||
|
||||
def decode(self, line: str) -> ServerSentEvent | None:
|
||||
# See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation # noqa: E501
|
||||
@@ -277,6 +329,11 @@ class SSEDecoder:
|
||||
self._data = []
|
||||
self._retry = None
|
||||
|
||||
# Skip [DONE] sentinel and ping events — don't surface them
|
||||
# as they are control signals, not parseable data.
|
||||
if sse.data == "[DONE]" or sse.event == "ping":
|
||||
return None
|
||||
|
||||
return sse
|
||||
|
||||
if line.startswith(":"):
|
||||
|
||||
+310
-59
@@ -7,18 +7,203 @@ Supports multiple authentication methods:
|
||||
4. MCP flow authentication
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import asyncio
|
||||
import webbrowser
|
||||
from dataclasses import dataclass
|
||||
from http.server import HTTPServer, BaseHTTPRequestHandler
|
||||
from threading import Thread
|
||||
from typing import Any, Dict, List, Optional
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlencode
|
||||
from urllib.parse import urlencode, parse_qs, urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PKCE (Proof Key for Code Exchange) with S256
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _base64url_encode(data: bytes) -> str:
|
||||
"""Base64url encode with no padding."""
|
||||
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PkceCodePair:
|
||||
"""PKCE verifier/challenge pair."""
|
||||
verifier: str
|
||||
challenge: str
|
||||
challenge_method: str = "S256"
|
||||
|
||||
|
||||
def generate_pkce_pair() -> PkceCodePair:
|
||||
"""Generate a PKCE code verifier and S256 challenge.
|
||||
|
||||
Uses 32 bytes of os.urandom for the verifier, SHA-256 for the challenge.
|
||||
Both are base64url encoded with no padding.
|
||||
"""
|
||||
verifier = _base64url_encode(os.urandom(32))
|
||||
challenge = _base64url_encode(hashlib.sha256(verifier.encode("ascii")).digest())
|
||||
return PkceCodePair(verifier=verifier, challenge=challenge)
|
||||
|
||||
|
||||
def generate_state() -> str:
|
||||
"""Generate a random state token (32 bytes, base64url, no padding)."""
|
||||
return _base64url_encode(os.urandom(32))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthAuthorizationRequest:
|
||||
"""OAuth authorization request with PKCE parameters."""
|
||||
authorize_url: str
|
||||
client_id: str
|
||||
redirect_uri: str
|
||||
scopes: List[str]
|
||||
state: str
|
||||
code_challenge: str
|
||||
code_challenge_method: str
|
||||
|
||||
def build_url(self) -> str:
|
||||
"""Build the full authorization URL with query parameters."""
|
||||
params = {
|
||||
"client_id": self.client_id,
|
||||
"redirect_uri": self.redirect_uri,
|
||||
"response_type": "code",
|
||||
"scope": " ".join(self.scopes),
|
||||
"state": self.state,
|
||||
"code_challenge": self.code_challenge,
|
||||
"code_challenge_method": self.code_challenge_method,
|
||||
}
|
||||
return f"{self.authorize_url}?{urlencode(params)}"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthTokenExchangeRequest:
|
||||
"""OAuth token exchange request with PKCE code_verifier."""
|
||||
code: str
|
||||
redirect_uri: str
|
||||
client_id: str
|
||||
code_verifier: str
|
||||
state: str
|
||||
grant_type: str = "authorization_code"
|
||||
|
||||
def form_params(self) -> Dict[str, str]:
|
||||
"""Return parameters for the token exchange POST body."""
|
||||
return {
|
||||
"grant_type": self.grant_type,
|
||||
"code": self.code,
|
||||
"redirect_uri": self.redirect_uri,
|
||||
"client_id": self.client_id,
|
||||
"code_verifier": self.code_verifier,
|
||||
"state": self.state,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class OAuthTokenSet:
|
||||
"""OAuth token set returned from token exchange."""
|
||||
access_token: str
|
||||
scopes: List[str]
|
||||
refresh_token: Optional[str] = None
|
||||
expires_at: Optional[int] = None
|
||||
|
||||
|
||||
class OAuthCredentialStore:
|
||||
"""Manages OAuth credentials in ~/.hanzo/credentials.json.
|
||||
|
||||
Respects HANZO_CONFIG_HOME env var. Uses atomic writes.
|
||||
Preserves other keys in the credentials file.
|
||||
"""
|
||||
|
||||
def __init__(self, config_home: Optional[str] = None):
|
||||
self._config_home = config_home
|
||||
|
||||
def _get_config_home(self) -> Path:
|
||||
if self._config_home:
|
||||
return Path(self._config_home)
|
||||
env = os.environ.get("HANZO_CONFIG_HOME")
|
||||
if env:
|
||||
return Path(env)
|
||||
return Path.home() / ".hanzo"
|
||||
|
||||
def credentials_path(self) -> Path:
|
||||
return self._get_config_home() / "credentials.json"
|
||||
|
||||
def save(self, token_set: OAuthTokenSet) -> None:
|
||||
"""Save OAuth token set. Atomic write, preserves other keys."""
|
||||
path = self.credentials_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Read existing data
|
||||
existing: Dict[str, Any] = {}
|
||||
if path.exists():
|
||||
with open(path, "r") as f:
|
||||
existing = json.load(f)
|
||||
|
||||
existing["oauth"] = {
|
||||
"access_token": token_set.access_token,
|
||||
"refresh_token": token_set.refresh_token,
|
||||
"expires_at": token_set.expires_at,
|
||||
"scopes": token_set.scopes,
|
||||
}
|
||||
|
||||
# Atomic write: write to .tmp then rename
|
||||
tmp_path = path.with_suffix(".json.tmp")
|
||||
with open(tmp_path, "w") as f:
|
||||
json.dump(existing, f, indent=2)
|
||||
|
||||
os.replace(tmp_path, path)
|
||||
|
||||
# Restrictive permissions
|
||||
if sys.platform != "win32":
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
def load(self) -> Optional[OAuthTokenSet]:
|
||||
"""Load OAuth token set from credentials file. Returns None if absent."""
|
||||
path = self.credentials_path()
|
||||
if not path.exists():
|
||||
return None
|
||||
|
||||
with open(path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
oauth = data.get("oauth")
|
||||
if not oauth:
|
||||
return None
|
||||
|
||||
return OAuthTokenSet(
|
||||
access_token=oauth["access_token"],
|
||||
refresh_token=oauth.get("refresh_token"),
|
||||
expires_at=oauth.get("expires_at"),
|
||||
scopes=oauth.get("scopes", []),
|
||||
)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Remove oauth key from credentials, preserving other keys."""
|
||||
path = self.credentials_path()
|
||||
if not path.exists():
|
||||
return
|
||||
|
||||
with open(path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
data.pop("oauth", None)
|
||||
|
||||
tmp_path = path.with_suffix(".json.tmp")
|
||||
with open(tmp_path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
os.replace(tmp_path, path)
|
||||
|
||||
if sys.platform != "win32":
|
||||
os.chmod(path, 0o600)
|
||||
|
||||
|
||||
class HanzoAuth:
|
||||
"""Hanzo authentication client."""
|
||||
|
||||
@@ -205,85 +390,151 @@ class HanzoAuth:
|
||||
|
||||
raise RuntimeError("Authentication timed out. Please try again.")
|
||||
|
||||
async def login_with_sso(self) -> Dict[str, Any]:
|
||||
async def login_with_sso(self) -> OAuthTokenSet:
|
||||
"""Login with SSO (browser-based).
|
||||
|
||||
Delegates to login_with_pkce with the SSO redirect port.
|
||||
|
||||
Returns:
|
||||
User information and tokens
|
||||
OAuthTokenSet with access token and metadata
|
||||
"""
|
||||
# Generate state for security
|
||||
import secrets
|
||||
return await self.login_with_pkce(redirect_port=8899)
|
||||
|
||||
state = secrets.token_urlsafe(32)
|
||||
async def login_with_pkce(
|
||||
self,
|
||||
authorize_url: Optional[str] = None,
|
||||
token_url: Optional[str] = None,
|
||||
client_id: str = "app-hanzo",
|
||||
scopes: Optional[List[str]] = None,
|
||||
redirect_port: int = 4545,
|
||||
) -> OAuthTokenSet:
|
||||
"""Login with PKCE (Proof Key for Code Exchange) flow.
|
||||
|
||||
# Build SSO URL
|
||||
params = {
|
||||
"client_id": "app-hanzo",
|
||||
"redirect_uri": "http://localhost:8899/callback",
|
||||
"response_type": "code",
|
||||
"scope": "openid profile email",
|
||||
"state": state,
|
||||
}
|
||||
Generates a PKCE pair, opens the browser, catches the callback
|
||||
on a local HTTP server, exchanges the code for tokens, and saves
|
||||
credentials.
|
||||
|
||||
sso_url = f"{self.base_url}/oauth/authorize?{urlencode(params)}"
|
||||
Args:
|
||||
authorize_url: Authorization endpoint (default: {base_url}/oauth/authorize)
|
||||
token_url: Token endpoint (default: {base_url}/oauth/token)
|
||||
client_id: OAuth client ID
|
||||
scopes: OAuth scopes
|
||||
redirect_port: Local port for the redirect callback server
|
||||
|
||||
Returns:
|
||||
OAuthTokenSet with access token and metadata
|
||||
"""
|
||||
if authorize_url is None:
|
||||
authorize_url = f"{self.base_url}/oauth/authorize"
|
||||
if token_url is None:
|
||||
token_url = f"{self.base_url}/oauth/token"
|
||||
if scopes is None:
|
||||
scopes = ["openid", "profile", "email"]
|
||||
|
||||
pkce = generate_pkce_pair()
|
||||
state = generate_state()
|
||||
redirect_uri = f"http://localhost:{redirect_port}/callback"
|
||||
|
||||
auth_req = OAuthAuthorizationRequest(
|
||||
authorize_url=authorize_url,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
scopes=scopes,
|
||||
state=state,
|
||||
code_challenge=pkce.challenge,
|
||||
code_challenge_method=pkce.challenge_method,
|
||||
)
|
||||
|
||||
# Capture the auth code from the callback
|
||||
result: Dict[str, Optional[str]] = {"code": None, "error": None}
|
||||
|
||||
class CallbackHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
qs = parse_qs(urlparse(self.path).query)
|
||||
|
||||
cb_state = qs.get("state", [None])[0]
|
||||
if cb_state != state:
|
||||
self.send_response(400)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Invalid state parameter.")
|
||||
result["error"] = "state mismatch"
|
||||
return
|
||||
|
||||
if "error" in qs:
|
||||
self.send_response(400)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
msg = qs["error"][0]
|
||||
self.wfile.write(f"Authorization error: {msg}".encode())
|
||||
result["error"] = msg
|
||||
return
|
||||
|
||||
result["code"] = qs.get("code", [None])[0]
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/html")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"Authentication successful. You can close this window.")
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass # Silence request logs
|
||||
|
||||
server = HTTPServer(("127.0.0.1", redirect_port), CallbackHandler)
|
||||
|
||||
# Open browser
|
||||
webbrowser.open(sso_url)
|
||||
url = auth_req.build_url()
|
||||
try:
|
||||
webbrowser.open(url)
|
||||
except Exception:
|
||||
pass
|
||||
print(f"\nOpen this URL to authenticate:\n {url}\n")
|
||||
|
||||
# Start local server to receive callback
|
||||
from aiohttp import web
|
||||
# Serve one request to catch the callback
|
||||
def serve_once():
|
||||
server.handle_request()
|
||||
server.server_close()
|
||||
|
||||
auth_code = None
|
||||
thread = Thread(target=serve_once, daemon=True)
|
||||
thread.start()
|
||||
# Wait for the callback (runs in background thread)
|
||||
await asyncio.get_event_loop().run_in_executor(None, thread.join, 300)
|
||||
|
||||
async def callback_handler(request):
|
||||
nonlocal auth_code
|
||||
if result["error"]:
|
||||
raise RuntimeError(f"PKCE authorization failed: {result['error']}")
|
||||
if not result["code"]:
|
||||
raise RuntimeError("No authorization code received.")
|
||||
|
||||
# Verify state
|
||||
if request.query.get("state") != state:
|
||||
return web.Response(text="Invalid state", status=400)
|
||||
# Exchange code for tokens
|
||||
exchange_req = OAuthTokenExchangeRequest(
|
||||
code=result["code"],
|
||||
redirect_uri=redirect_uri,
|
||||
client_id=client_id,
|
||||
code_verifier=pkce.verifier,
|
||||
state=state,
|
||||
)
|
||||
|
||||
auth_code = request.query.get("code")
|
||||
|
||||
return web.Response(
|
||||
text="Authentication successful! You can close this window.",
|
||||
content_type="text/html",
|
||||
)
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get("/callback", callback_handler)
|
||||
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "localhost", 8899)
|
||||
await site.start()
|
||||
|
||||
# Wait for callback
|
||||
while auth_code is None:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
await runner.cleanup()
|
||||
|
||||
# Exchange code for token
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/oauth/token",
|
||||
json={
|
||||
"client_id": "app-hanzo",
|
||||
"client_secret": "hanzo-cli-secret", # Public client
|
||||
"code": auth_code,
|
||||
"grant_type": "authorization_code",
|
||||
"redirect_uri": "http://localhost:8899/callback",
|
||||
},
|
||||
token_url,
|
||||
data=exchange_req.form_params(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
self._token = data.get("access_token")
|
||||
|
||||
# Get user info
|
||||
user_info = await self.get_user_info()
|
||||
token_set = OAuthTokenSet(
|
||||
access_token=data["access_token"],
|
||||
refresh_token=data.get("refresh_token"),
|
||||
expires_at=data.get("expires_at"),
|
||||
scopes=data.get("scope", " ".join(scopes)).split(),
|
||||
)
|
||||
|
||||
return {"token": self._token, **user_info}
|
||||
self._token = token_set.access_token
|
||||
|
||||
# Save credentials
|
||||
store = OAuthCredentialStore()
|
||||
store.save(token_set)
|
||||
|
||||
return token_set
|
||||
|
||||
async def logout(self):
|
||||
"""Logout and clear credentials."""
|
||||
|
||||
@@ -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)
|
||||
+305
-33
@@ -2,8 +2,21 @@
|
||||
|
||||
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 (
|
||||
@@ -87,49 +100,299 @@ def run_mcp_server(name: str = "hanzo-mcp", transport: str = "stdio", **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 for connecting to MCP servers."""
|
||||
"""MCP client that communicates with a server subprocess over stdio JSON-RPC."""
|
||||
|
||||
def __init__(self, server_command: list[str] = None):
|
||||
"""Initialize MCP client.
|
||||
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
|
||||
|
||||
Args:
|
||||
server_command: Command to start the server (e.g., ["python", "-m", "hanzo_mcp"])
|
||||
"""
|
||||
if not MCP_AVAILABLE:
|
||||
raise ImportError(
|
||||
"hanzo-mcp is not installed. Install it with: pip install hanzo-mcp"
|
||||
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
|
||||
|
||||
self.server_command = server_command or ["python", "-m", "hanzo_mcp"]
|
||||
self._client = None
|
||||
result = await self._send("initialize", _INIT_PARAMS)
|
||||
self.server_info = result.get("serverInfo")
|
||||
self.capabilities = result.get("capabilities")
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to the MCP server."""
|
||||
# This would use the MCP client SDK when available
|
||||
raise NotImplementedError("MCP client SDK integration pending")
|
||||
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 call_tool(self, tool_name: str, **kwargs):
|
||||
"""Call a tool on the server.
|
||||
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
|
||||
|
||||
Args:
|
||||
tool_name: Name of the tool to call
|
||||
**kwargs: Tool arguments
|
||||
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)
|
||||
|
||||
Returns:
|
||||
Tool result
|
||||
"""
|
||||
if not self._client:
|
||||
await self.connect()
|
||||
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", [])
|
||||
|
||||
# This would use the MCP client SDK when available
|
||||
raise NotImplementedError("MCP client SDK integration pending")
|
||||
async def __aenter__(self) -> MCPClient:
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from the server."""
|
||||
if self._client:
|
||||
# Cleanup would go here
|
||||
self._client = None
|
||||
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__ = [
|
||||
@@ -153,8 +416,17 @@ __all__ = [
|
||||
"ToolRegistry",
|
||||
# Permissions
|
||||
"PermissionManager",
|
||||
# Client
|
||||
# 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()
|
||||
@@ -0,0 +1,313 @@
|
||||
"""Tests for PKCE OAuth support in hanzoai.auth.
|
||||
|
||||
Run: .venv/bin/python -m pytest tests/test_auth_pkce.py -v
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from hanzoai.auth import (
|
||||
OAuthAuthorizationRequest,
|
||||
OAuthCredentialStore,
|
||||
OAuthTokenExchangeRequest,
|
||||
OAuthTokenSet,
|
||||
PkceCodePair,
|
||||
generate_pkce_pair,
|
||||
generate_state,
|
||||
)
|
||||
|
||||
|
||||
class TestPkceCodePair(unittest.TestCase):
|
||||
def test_fields(self):
|
||||
pair = PkceCodePair(verifier="v", challenge="c")
|
||||
self.assertEqual(pair.verifier, "v")
|
||||
self.assertEqual(pair.challenge, "c")
|
||||
self.assertEqual(pair.challenge_method, "S256")
|
||||
|
||||
def test_custom_method(self):
|
||||
pair = PkceCodePair(verifier="v", challenge="c", challenge_method="plain")
|
||||
self.assertEqual(pair.challenge_method, "plain")
|
||||
|
||||
|
||||
class TestGeneratePkcePair(unittest.TestCase):
|
||||
def test_returns_pkce_pair(self):
|
||||
pair = generate_pkce_pair()
|
||||
self.assertIsInstance(pair, PkceCodePair)
|
||||
self.assertEqual(pair.challenge_method, "S256")
|
||||
|
||||
def test_verifier_is_base64url_no_padding(self):
|
||||
pair = generate_pkce_pair()
|
||||
self.assertNotIn("=", pair.verifier)
|
||||
self.assertNotIn("+", pair.verifier)
|
||||
self.assertNotIn("/", pair.verifier)
|
||||
# 32 bytes -> 43 base64url chars (no padding)
|
||||
self.assertEqual(len(pair.verifier), 43)
|
||||
|
||||
def test_challenge_is_sha256_of_verifier(self):
|
||||
pair = generate_pkce_pair()
|
||||
# SHA-256 of the verifier string (ASCII), then base64url no padding
|
||||
expected_hash = hashlib.sha256(pair.verifier.encode("ascii")).digest()
|
||||
expected_challenge = base64.urlsafe_b64encode(expected_hash).rstrip(b"=").decode("ascii")
|
||||
self.assertEqual(pair.challenge, expected_challenge)
|
||||
|
||||
def test_challenge_length(self):
|
||||
pair = generate_pkce_pair()
|
||||
# SHA-256 is 32 bytes -> 43 base64url chars (no padding)
|
||||
self.assertEqual(len(pair.challenge), 43)
|
||||
|
||||
def test_pairs_are_unique(self):
|
||||
a = generate_pkce_pair()
|
||||
b = generate_pkce_pair()
|
||||
self.assertNotEqual(a.verifier, b.verifier)
|
||||
self.assertNotEqual(a.challenge, b.challenge)
|
||||
|
||||
def test_deterministic_with_known_input(self):
|
||||
"""Verify against a known PKCE test vector."""
|
||||
# Use a fixed 32-byte input to verify the encoding chain
|
||||
fixed_bytes = b"\x00" * 32
|
||||
verifier = base64.urlsafe_b64encode(fixed_bytes).rstrip(b"=").decode("ascii")
|
||||
challenge_hash = hashlib.sha256(verifier.encode("ascii")).digest()
|
||||
challenge = base64.urlsafe_b64encode(challenge_hash).rstrip(b"=").decode("ascii")
|
||||
|
||||
with patch("os.urandom", return_value=fixed_bytes):
|
||||
pair = generate_pkce_pair()
|
||||
|
||||
self.assertEqual(pair.verifier, verifier)
|
||||
self.assertEqual(pair.challenge, challenge)
|
||||
|
||||
|
||||
class TestGenerateState(unittest.TestCase):
|
||||
def test_returns_string(self):
|
||||
state = generate_state()
|
||||
self.assertIsInstance(state, str)
|
||||
|
||||
def test_base64url_no_padding(self):
|
||||
state = generate_state()
|
||||
self.assertNotIn("=", state)
|
||||
self.assertNotIn("+", state)
|
||||
self.assertNotIn("/", state)
|
||||
# 32 bytes -> 43 base64url chars
|
||||
self.assertEqual(len(state), 43)
|
||||
|
||||
def test_unique(self):
|
||||
a = generate_state()
|
||||
b = generate_state()
|
||||
self.assertNotEqual(a, b)
|
||||
|
||||
|
||||
class TestOAuthAuthorizationRequest(unittest.TestCase):
|
||||
def _make_req(self, **kwargs):
|
||||
defaults = dict(
|
||||
authorize_url="https://hanzo.id/oauth/authorize",
|
||||
client_id="app-hanzo",
|
||||
redirect_uri="http://localhost:4545/callback",
|
||||
scopes=["openid", "profile"],
|
||||
state="test-state",
|
||||
code_challenge="test-challenge",
|
||||
code_challenge_method="S256",
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return OAuthAuthorizationRequest(**defaults)
|
||||
|
||||
def test_build_url_contains_all_params(self):
|
||||
req = self._make_req()
|
||||
url = req.build_url()
|
||||
parsed = urlparse(url)
|
||||
params = parse_qs(parsed.query)
|
||||
|
||||
self.assertEqual(parsed.scheme, "https")
|
||||
self.assertEqual(parsed.netloc, "hanzo.id")
|
||||
self.assertEqual(parsed.path, "/oauth/authorize")
|
||||
self.assertEqual(params["client_id"], ["app-hanzo"])
|
||||
self.assertEqual(params["redirect_uri"], ["http://localhost:4545/callback"])
|
||||
self.assertEqual(params["response_type"], ["code"])
|
||||
self.assertEqual(params["scope"], ["openid profile"])
|
||||
self.assertEqual(params["state"], ["test-state"])
|
||||
self.assertEqual(params["code_challenge"], ["test-challenge"])
|
||||
self.assertEqual(params["code_challenge_method"], ["S256"])
|
||||
|
||||
def test_build_url_scopes_joined(self):
|
||||
req = self._make_req(scopes=["openid", "profile", "email"])
|
||||
url = req.build_url()
|
||||
params = parse_qs(urlparse(url).query)
|
||||
self.assertEqual(params["scope"], ["openid profile email"])
|
||||
|
||||
def test_build_url_starts_with_authorize_url(self):
|
||||
req = self._make_req(authorize_url="https://custom.example.com/auth")
|
||||
url = req.build_url()
|
||||
self.assertTrue(url.startswith("https://custom.example.com/auth?"))
|
||||
|
||||
|
||||
class TestOAuthTokenExchangeRequest(unittest.TestCase):
|
||||
def test_form_params(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="auth-code-123",
|
||||
redirect_uri="http://localhost:4545/callback",
|
||||
client_id="app-hanzo",
|
||||
code_verifier="verifier-abc",
|
||||
state="state-xyz",
|
||||
)
|
||||
params = req.form_params()
|
||||
self.assertEqual(params["grant_type"], "authorization_code")
|
||||
self.assertEqual(params["code"], "auth-code-123")
|
||||
self.assertEqual(params["redirect_uri"], "http://localhost:4545/callback")
|
||||
self.assertEqual(params["client_id"], "app-hanzo")
|
||||
self.assertEqual(params["code_verifier"], "verifier-abc")
|
||||
self.assertEqual(params["state"], "state-xyz")
|
||||
|
||||
def test_default_grant_type(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="c", redirect_uri="r", client_id="ci", code_verifier="cv", state="s"
|
||||
)
|
||||
self.assertEqual(req.grant_type, "authorization_code")
|
||||
|
||||
def test_form_params_keys(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="c", redirect_uri="r", client_id="ci", code_verifier="cv", state="s"
|
||||
)
|
||||
params = req.form_params()
|
||||
expected_keys = {"grant_type", "code", "redirect_uri", "client_id", "code_verifier", "state"}
|
||||
self.assertEqual(set(params.keys()), expected_keys)
|
||||
|
||||
|
||||
class TestOAuthTokenSet(unittest.TestCase):
|
||||
def test_required_fields(self):
|
||||
ts = OAuthTokenSet(access_token="tok123", scopes=["openid"])
|
||||
self.assertEqual(ts.access_token, "tok123")
|
||||
self.assertEqual(ts.scopes, ["openid"])
|
||||
self.assertIsNone(ts.refresh_token)
|
||||
self.assertIsNone(ts.expires_at)
|
||||
|
||||
def test_all_fields(self):
|
||||
ts = OAuthTokenSet(
|
||||
access_token="tok",
|
||||
scopes=["openid"],
|
||||
refresh_token="ref",
|
||||
expires_at=1234567890,
|
||||
)
|
||||
self.assertEqual(ts.refresh_token, "ref")
|
||||
self.assertEqual(ts.expires_at, 1234567890)
|
||||
|
||||
|
||||
class TestOAuthCredentialStore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.store = OAuthCredentialStore(config_home=self.tmpdir)
|
||||
|
||||
def test_credentials_path_default(self):
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
# Remove HANZO_CONFIG_HOME if set
|
||||
env = os.environ.copy()
|
||||
env.pop("HANZO_CONFIG_HOME", None)
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
store = OAuthCredentialStore()
|
||||
expected = Path.home() / ".hanzo" / "credentials.json"
|
||||
self.assertEqual(store.credentials_path(), expected)
|
||||
|
||||
def test_credentials_path_custom_config_home(self):
|
||||
self.assertEqual(
|
||||
self.store.credentials_path(),
|
||||
Path(self.tmpdir) / "credentials.json",
|
||||
)
|
||||
|
||||
def test_credentials_path_env_override(self):
|
||||
with patch.dict(os.environ, {"HANZO_CONFIG_HOME": "/custom/path"}):
|
||||
store = OAuthCredentialStore()
|
||||
self.assertEqual(store.credentials_path(), Path("/custom/path/credentials.json"))
|
||||
|
||||
def test_save_and_load_roundtrip(self):
|
||||
token_set = OAuthTokenSet(
|
||||
access_token="access123",
|
||||
refresh_token="refresh456",
|
||||
expires_at=9999999999,
|
||||
scopes=["openid", "profile"],
|
||||
)
|
||||
self.store.save(token_set)
|
||||
|
||||
loaded = self.store.load()
|
||||
self.assertIsNotNone(loaded)
|
||||
self.assertEqual(loaded.access_token, "access123")
|
||||
self.assertEqual(loaded.refresh_token, "refresh456")
|
||||
self.assertEqual(loaded.expires_at, 9999999999)
|
||||
self.assertEqual(loaded.scopes, ["openid", "profile"])
|
||||
|
||||
def test_save_preserves_other_keys(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "keep-me", "other": "data"}, f)
|
||||
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
with open(creds_path) as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.assertEqual(data["api_key"], "keep-me")
|
||||
self.assertEqual(data["other"], "data")
|
||||
self.assertIn("oauth", data)
|
||||
|
||||
def test_load_returns_none_when_no_file(self):
|
||||
self.assertIsNone(self.store.load())
|
||||
|
||||
def test_load_returns_none_when_no_oauth_key(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "something"}, f)
|
||||
|
||||
self.assertIsNone(self.store.load())
|
||||
|
||||
def test_clear_removes_oauth_preserves_rest(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "keep", "oauth": {"access_token": "remove"}}, f)
|
||||
|
||||
self.store.clear()
|
||||
|
||||
with open(creds_path) as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.assertEqual(data["api_key"], "keep")
|
||||
self.assertNotIn("oauth", data)
|
||||
|
||||
def test_clear_noop_when_no_file(self):
|
||||
self.store.clear() # should not raise
|
||||
|
||||
def test_save_atomic_no_tmp_left(self):
|
||||
"""After save, .tmp file should not exist."""
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
creds_path = self.store.credentials_path()
|
||||
tmp_path = creds_path.with_suffix(".json.tmp")
|
||||
self.assertTrue(creds_path.exists())
|
||||
self.assertFalse(tmp_path.exists())
|
||||
|
||||
def test_save_sets_restrictive_permissions(self):
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
creds_path = self.store.credentials_path()
|
||||
mode = os.stat(creds_path).st_mode
|
||||
self.assertEqual(stat.S_IMODE(mode), 0o600)
|
||||
|
||||
def test_save_creates_parent_dirs(self):
|
||||
nested = os.path.join(self.tmpdir, "deep", "nested")
|
||||
store = OAuthCredentialStore(config_home=nested)
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
store.save(token_set)
|
||||
self.assertTrue(store.credentials_path().exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,410 @@
|
||||
"""Tests for hanzoai.config — hierarchical config discovery and merging."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.config import (
|
||||
ConfigEntry,
|
||||
ConfigError,
|
||||
ConfigLoader,
|
||||
ConfigSource,
|
||||
McpClaudeAiProxyConfig,
|
||||
McpOAuthConfig,
|
||||
McpRemoteConfig,
|
||||
McpSdkConfig,
|
||||
McpStdioConfig,
|
||||
McpTransport,
|
||||
McpWsConfig,
|
||||
OAuthConfig,
|
||||
RuntimeConfig,
|
||||
_deep_merge,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deep merge
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeepMerge:
|
||||
def test_flat_override(self):
|
||||
assert _deep_merge({"a": 1}, {"a": 2}) == {"a": 2}
|
||||
|
||||
def test_flat_union(self):
|
||||
assert _deep_merge({"a": 1}, {"b": 2}) == {"a": 1, "b": 2}
|
||||
|
||||
def test_nested_merge(self):
|
||||
base = {"a": {"x": 1, "y": 2}}
|
||||
over = {"a": {"y": 3, "z": 4}}
|
||||
assert _deep_merge(base, over) == {"a": {"x": 1, "y": 3, "z": 4}}
|
||||
|
||||
def test_nested_replace_non_dict(self):
|
||||
assert _deep_merge({"a": {"x": 1}}, {"a": 42}) == {"a": 42}
|
||||
|
||||
def test_empty(self):
|
||||
assert _deep_merge({}, {"a": 1}) == {"a": 1}
|
||||
assert _deep_merge({"a": 1}, {}) == {"a": 1}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.discover
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDiscover:
|
||||
def test_returns_three_entries(self, tmp_path: Path):
|
||||
loader = ConfigLoader(cwd=tmp_path, config_home=tmp_path / ".hanzo")
|
||||
entries = loader.discover()
|
||||
assert len(entries) == 3
|
||||
assert entries[0].source == ConfigSource.User
|
||||
assert entries[1].source == ConfigSource.Project
|
||||
assert entries[2].source == ConfigSource.Local
|
||||
|
||||
def test_paths(self, tmp_path: Path):
|
||||
home = tmp_path / "home" / ".hanzo"
|
||||
cwd = tmp_path / "project"
|
||||
loader = ConfigLoader(cwd=cwd, config_home=home)
|
||||
entries = loader.discover()
|
||||
assert entries[0].path == home / "settings.json"
|
||||
assert entries[1].path == cwd / ".hanzo" / "settings.json"
|
||||
assert entries[2].path == cwd / ".hanzo" / "settings.local.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.default_for
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDefaultFor:
|
||||
def test_uses_env_var(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
custom = tmp_path / "custom"
|
||||
monkeypatch.setenv("HANZO_CONFIG_HOME", str(custom))
|
||||
loader = ConfigLoader.default_for(tmp_path)
|
||||
assert loader.config_home == custom
|
||||
|
||||
def test_falls_back_to_home(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.delenv("HANZO_CONFIG_HOME", raising=False)
|
||||
loader = ConfigLoader.default_for(tmp_path)
|
||||
assert loader.config_home == Path.home() / ".hanzo"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.load — file handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_json(path: Path, data: object) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
|
||||
|
||||
class TestLoad:
|
||||
def test_empty_when_no_files(self, tmp_path: Path):
|
||||
loader = ConfigLoader(cwd=tmp_path, config_home=tmp_path / "h")
|
||||
cfg = loader.load()
|
||||
assert cfg.merged == {}
|
||||
assert cfg.loaded_entries == []
|
||||
|
||||
def test_single_user_file(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"theme": "dark"})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"theme": "dark"}
|
||||
assert len(cfg.loaded_entries) == 1
|
||||
assert cfg.loaded_entries[0].source == ConfigSource.User
|
||||
|
||||
def test_project_overrides_user(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"a": 1, "b": 2})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {"b": 99})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"a": 1, "b": 99}
|
||||
assert len(cfg.loaded_entries) == 2
|
||||
|
||||
def test_local_overrides_project(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"a": 1})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {"b": 2})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.local.json", {"b": 3, "c": 4})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"a": 1, "b": 3, "c": 4}
|
||||
assert len(cfg.loaded_entries) == 3
|
||||
|
||||
def test_invalid_json_raises(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
p = home / "settings.json"
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text("{bad json", encoding="utf-8")
|
||||
with pytest.raises(ConfigError, match="invalid JSON"):
|
||||
ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
|
||||
def test_non_object_top_level_raises(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", [1, 2, 3])
|
||||
with pytest.raises(ConfigError, match="top-level value must be a JSON object"):
|
||||
ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
|
||||
def test_mcp_servers_replaced_not_merged(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {
|
||||
"mcpServers": {
|
||||
"alpha": {"command": "a", "transport": "stdio"},
|
||||
"beta": {"command": "b", "transport": "stdio"},
|
||||
}
|
||||
})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {
|
||||
"mcpServers": {
|
||||
"gamma": {"command": "g", "transport": "stdio"},
|
||||
}
|
||||
})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
servers = cfg.mcp_servers()
|
||||
assert set(servers.keys()) == {"gamma"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RuntimeConfig accessors
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRuntimeConfig:
|
||||
def test_empty(self):
|
||||
cfg = RuntimeConfig.empty()
|
||||
assert cfg.merged == {}
|
||||
assert cfg.get("anything") is None
|
||||
|
||||
def test_get(self):
|
||||
cfg = RuntimeConfig(merged={"foo": "bar"})
|
||||
assert cfg.get("foo") == "bar"
|
||||
assert cfg.get("nope") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP server parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMcpServers:
|
||||
def test_stdio(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"myserver": {
|
||||
"transport": "stdio",
|
||||
"command": "node",
|
||||
"args": ["server.js"],
|
||||
"env": {"PORT": "8080"},
|
||||
}
|
||||
}
|
||||
})
|
||||
servers = cfg.mcp_servers()
|
||||
s = servers["myserver"]
|
||||
assert isinstance(s, McpStdioConfig)
|
||||
assert s.command == "node"
|
||||
assert s.args == ["server.js"]
|
||||
assert s.env == {"PORT": "8080"}
|
||||
assert s.transport == McpTransport.Stdio
|
||||
|
||||
def test_stdio_default_transport(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"s": {"command": "echo"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["s"]
|
||||
assert isinstance(s, McpStdioConfig)
|
||||
|
||||
def test_http_remote(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"remote": {
|
||||
"transport": "http",
|
||||
"url": "https://example.com/mcp",
|
||||
"headers": {"Authorization": "Bearer tok"},
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["remote"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.url == "https://example.com/mcp"
|
||||
assert s.transport == McpTransport.Http
|
||||
|
||||
def test_sse_remote(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"sse": {"transport": "sse", "url": "https://x.com/events"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["sse"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.transport == McpTransport.Sse
|
||||
|
||||
def test_ws(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"ws": {"transport": "ws", "url": "wss://x.com/ws"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["ws"]
|
||||
assert isinstance(s, McpWsConfig)
|
||||
assert s.transport == McpTransport.Ws
|
||||
assert s.url == "wss://x.com/ws"
|
||||
|
||||
def test_sdk(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"my-sdk": {"transport": "sdk", "name": "my-sdk-server"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["my-sdk"]
|
||||
assert isinstance(s, McpSdkConfig)
|
||||
assert s.transport == McpTransport.Sdk
|
||||
assert s.name == "my-sdk-server"
|
||||
|
||||
def test_claudeai_proxy(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"proxy": {
|
||||
"transport": "claudeai-proxy",
|
||||
"url": "https://claude.ai/proxy",
|
||||
"id": "server-123",
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["proxy"]
|
||||
assert isinstance(s, McpClaudeAiProxyConfig)
|
||||
assert s.transport == McpTransport.ClaudeAiProxy
|
||||
assert s.url == "https://claude.ai/proxy"
|
||||
assert s.id == "server-123"
|
||||
|
||||
def test_remote_with_oauth(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"authed": {
|
||||
"transport": "http",
|
||||
"url": "https://example.com/mcp",
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"callbackPort": 9999,
|
||||
"authServerMetadataUrl": "https://auth.example.com/.well-known",
|
||||
"xaa": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["authed"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.oauth is not None
|
||||
assert s.oauth.client_id == "cid"
|
||||
assert s.oauth.callback_port == 9999
|
||||
assert s.oauth.xaa is True
|
||||
|
||||
def test_sdk_missing_name_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "sdk"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'name'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_claudeai_proxy_missing_id_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "claudeai-proxy", "url": "https://x.com"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'id'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_unknown_transport_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "grpc"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="unknown transport"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_stdio_missing_command_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "stdio"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'command'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_http_missing_url_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "http"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'url'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_no_mcp_servers_returns_empty(self):
|
||||
assert RuntimeConfig.empty().mcp_servers() == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestOAuth:
|
||||
def test_parses_camel_case(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "my-id",
|
||||
"authorizeUrl": "https://hanzo.id/authorize",
|
||||
"tokenUrl": "https://hanzo.id/token",
|
||||
"scopes": ["read", "write"],
|
||||
"callbackPort": 8080,
|
||||
}
|
||||
})
|
||||
oauth = cfg.oauth()
|
||||
assert isinstance(oauth, OAuthConfig)
|
||||
assert oauth.client_id == "my-id"
|
||||
assert oauth.authorize_url == "https://hanzo.id/authorize"
|
||||
assert oauth.token_url == "https://hanzo.id/token"
|
||||
assert oauth.scopes == ["read", "write"]
|
||||
assert oauth.callback_port == 8080
|
||||
|
||||
def test_parses_snake_case(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"client_id": "cid",
|
||||
"authorize_url": "https://hanzo.id/authorize",
|
||||
"token_url": "https://hanzo.id/token",
|
||||
}
|
||||
})
|
||||
oauth = cfg.oauth()
|
||||
assert oauth is not None
|
||||
assert oauth.client_id == "cid"
|
||||
assert oauth.authorize_url == "https://hanzo.id/authorize"
|
||||
|
||||
def test_returns_none_when_missing(self):
|
||||
assert RuntimeConfig.empty().oauth() is None
|
||||
|
||||
def test_returns_none_when_no_client_id(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"authorizeUrl": "https://x.com/auth",
|
||||
"tokenUrl": "https://x.com/token",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
|
||||
def test_returns_none_when_no_authorize_url(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"tokenUrl": "https://x.com/token",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
|
||||
def test_returns_none_when_no_token_url(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"authorizeUrl": "https://x.com/auth",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Tests for MCP client implementation."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.mcp import (
|
||||
JsonRpcError,
|
||||
JsonRpcRequest,
|
||||
JsonRpcResponse,
|
||||
MCPClient,
|
||||
MCPClientError,
|
||||
MCPHttpClient,
|
||||
mcp_tool_name,
|
||||
normalize_mcp_name,
|
||||
)
|
||||
|
||||
|
||||
# --- Name normalization tests ---
|
||||
|
||||
|
||||
class TestNormalizeMcpName:
|
||||
def test_passthrough(self):
|
||||
assert normalize_mcp_name("hello") == "hello"
|
||||
|
||||
def test_hyphens_kept(self):
|
||||
# Rust parity: hyphens are valid chars and kept as-is
|
||||
assert normalize_mcp_name("my-tool") == "my-tool"
|
||||
|
||||
def test_dots_to_underscores(self):
|
||||
assert normalize_mcp_name("my.tool") == "my_tool"
|
||||
|
||||
def test_slashes_to_underscores(self):
|
||||
assert normalize_mcp_name("my/tool") == "my_tool"
|
||||
|
||||
def test_mixed_keeps_hyphens(self):
|
||||
assert normalize_mcp_name("my-tool.v2/action") == "my-tool_v2_action"
|
||||
|
||||
def test_non_claude_keeps_consecutive_underscores(self):
|
||||
# Non-claude.ai names do NOT collapse underscores
|
||||
assert normalize_mcp_name("my--tool") == "my--tool"
|
||||
|
||||
def test_non_claude_keeps_leading_trailing(self):
|
||||
# Non-claude.ai names do NOT strip leading/trailing
|
||||
assert normalize_mcp_name("-tool-") == "-tool-"
|
||||
|
||||
def test_non_claude_trailing_underscore(self):
|
||||
assert normalize_mcp_name("tool name!") == "tool_name_"
|
||||
|
||||
def test_github_dot_com(self):
|
||||
assert normalize_mcp_name("github.com") == "github_com"
|
||||
|
||||
def test_claude_ai_collapses_and_strips(self):
|
||||
# claude.ai names get underscore collapsing + strip
|
||||
assert normalize_mcp_name("claude.ai Example Server!!") == "claude_ai_Example_Server"
|
||||
|
||||
def test_claude_ai_simple(self):
|
||||
assert normalize_mcp_name("claude.ai My Tool") == "claude_ai_My_Tool"
|
||||
|
||||
|
||||
class TestMcpToolName:
|
||||
def test_basic(self):
|
||||
assert mcp_tool_name("server", "tool") == "mcp__server__tool"
|
||||
|
||||
def test_normalizes_both(self):
|
||||
# Hyphens kept, so my-server stays my-server
|
||||
assert mcp_tool_name("my-server", "my-tool") == "mcp__my-server__my-tool"
|
||||
|
||||
def test_claude_ai_server(self):
|
||||
assert mcp_tool_name("claude.ai Example Server", "weather tool") == "mcp__claude_ai_Example_Server__weather_tool"
|
||||
|
||||
|
||||
# --- Protocol types tests ---
|
||||
|
||||
|
||||
class TestJsonRpcRequest:
|
||||
def test_to_json(self):
|
||||
req = JsonRpcRequest(method="initialize", params={"a": 1}, id=1)
|
||||
d = req.to_dict()
|
||||
assert d["jsonrpc"] == "2.0"
|
||||
assert d["method"] == "initialize"
|
||||
assert d["params"] == {"a": 1}
|
||||
assert d["id"] == 1
|
||||
|
||||
def test_to_line(self):
|
||||
req = JsonRpcRequest(method="test", params={}, id=42)
|
||||
line = req.to_line()
|
||||
assert line.endswith(b"\n")
|
||||
parsed = json.loads(line)
|
||||
assert parsed["id"] == 42
|
||||
|
||||
|
||||
class TestJsonRpcResponse:
|
||||
def test_from_dict_success(self):
|
||||
resp = JsonRpcResponse.from_dict({"jsonrpc": "2.0", "result": {"ok": True}, "id": 1})
|
||||
assert resp.result == {"ok": True}
|
||||
assert resp.error is None
|
||||
assert resp.id == 1
|
||||
|
||||
def test_from_dict_error(self):
|
||||
resp = JsonRpcResponse.from_dict({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {"code": -32600, "message": "Invalid Request"},
|
||||
"id": 1,
|
||||
})
|
||||
assert resp.result is None
|
||||
assert resp.error is not None
|
||||
assert resp.error.code == -32600
|
||||
assert resp.error.message == "Invalid Request"
|
||||
|
||||
|
||||
# --- Stdio MCPClient tests ---
|
||||
|
||||
# A minimal MCP server script that we spawn as a subprocess.
|
||||
# It reads JSON-RPC from stdin and responds on stdout.
|
||||
MOCK_SERVER_SCRIPT = textwrap.dedent("""\
|
||||
import json
|
||||
import sys
|
||||
|
||||
def respond(id, result):
|
||||
msg = json.dumps({"jsonrpc": "2.0", "result": result, "id": id})
|
||||
sys.stdout.write(msg + "\\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def respond_error(id, code, message):
|
||||
msg = json.dumps({"jsonrpc": "2.0", "error": {"code": code, "message": message}, "id": id})
|
||||
sys.stdout.write(msg + "\\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
req = json.loads(line)
|
||||
method = req["method"]
|
||||
rid = req["id"]
|
||||
|
||||
if method == "initialize":
|
||||
respond(rid, {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}, "resources": {}},
|
||||
"serverInfo": {"name": "mock", "version": "0.1.0"},
|
||||
})
|
||||
elif method == "tools/list":
|
||||
respond(rid, {
|
||||
"tools": [
|
||||
{"name": "echo", "description": "echoes input", "inputSchema": {"type": "object"}},
|
||||
],
|
||||
})
|
||||
elif method == "tools/call":
|
||||
tool_name = req["params"]["name"]
|
||||
args = req["params"].get("arguments", {})
|
||||
if tool_name == "echo":
|
||||
respond(rid, {"content": [{"type": "text", "text": json.dumps(args)}]})
|
||||
elif tool_name == "fail":
|
||||
respond_error(rid, -32000, "tool failed")
|
||||
else:
|
||||
respond_error(rid, -32601, f"unknown tool: {tool_name}")
|
||||
elif method == "resources/list":
|
||||
respond(rid, {"resources": [{"uri": "file:///test.txt", "name": "test.txt"}]})
|
||||
else:
|
||||
respond_error(rid, -32601, f"unknown method: {method}")
|
||||
""")
|
||||
|
||||
|
||||
class TestMCPClientStdio:
|
||||
@pytest.fixture
|
||||
async def client(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
c = MCPClient(server_command=[sys.executable, str(script)])
|
||||
await c.connect()
|
||||
yield c
|
||||
await c.disconnect()
|
||||
|
||||
async def test_connect_initializes(self, client):
|
||||
assert client.server_info is not None
|
||||
assert client.server_info["name"] == "mock"
|
||||
|
||||
async def test_list_tools(self, client):
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["name"] == "echo"
|
||||
|
||||
async def test_call_tool(self, client):
|
||||
result = await client.call_tool("echo", arguments={"msg": "hello"})
|
||||
assert result["content"][0]["text"] == '{"msg": "hello"}'
|
||||
|
||||
async def test_call_tool_error(self, client):
|
||||
with pytest.raises(MCPClientError, match="tool failed"):
|
||||
await client.call_tool("fail")
|
||||
|
||||
async def test_list_resources(self, client):
|
||||
resources = await client.list_resources()
|
||||
assert len(resources) == 1
|
||||
assert resources[0]["uri"] == "file:///test.txt"
|
||||
|
||||
async def test_context_manager(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
async with MCPClient(server_command=[sys.executable, str(script)]) as c:
|
||||
tools = await c.list_tools()
|
||||
assert len(tools) == 1
|
||||
|
||||
async def test_request_ids_increment(self, client):
|
||||
# Each call increments the request id
|
||||
id_before = client._next_id
|
||||
await client.list_tools()
|
||||
assert client._next_id == id_before + 1
|
||||
|
||||
async def test_disconnect_kills_process(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
c = MCPClient(server_command=[sys.executable, str(script)])
|
||||
await c.connect()
|
||||
proc = c._process
|
||||
await c.disconnect()
|
||||
assert c._process is None
|
||||
# Process should be terminated
|
||||
assert proc.returncode is not None
|
||||
|
||||
|
||||
class TestMCPClientNotConnected:
|
||||
async def test_list_tools_without_connect(self):
|
||||
c = MCPClient(server_command=["true"])
|
||||
with pytest.raises(MCPClientError, match="not connected"):
|
||||
await c.list_tools()
|
||||
|
||||
async def test_call_tool_without_connect(self):
|
||||
c = MCPClient(server_command=["true"])
|
||||
with pytest.raises(MCPClientError, match="not connected"):
|
||||
await c.call_tool("x")
|
||||
|
||||
|
||||
class TestMCPClientBadProcess:
|
||||
async def test_connect_to_nonexistent_binary(self):
|
||||
c = MCPClient(server_command=["__nonexistent_binary_xyz__"])
|
||||
with pytest.raises(MCPClientError):
|
||||
await c.connect()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Tests for session management and compaction."""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.session import (
|
||||
CompactionConfig,
|
||||
CompactionResult,
|
||||
ConversationMessage,
|
||||
MessageRole,
|
||||
Session,
|
||||
TextBlock,
|
||||
TokenUsage,
|
||||
ToolResultBlock,
|
||||
ToolUseBlock,
|
||||
compact_session,
|
||||
estimate_session_tokens,
|
||||
format_compact_summary,
|
||||
should_compact,
|
||||
)
|
||||
|
||||
|
||||
class TestTokenUsage:
|
||||
def test_total_tokens(self):
|
||||
u = TokenUsage(input_tokens=10, output_tokens=4, cache_creation_input_tokens=2, cache_read_input_tokens=1)
|
||||
assert u.total_tokens() == 17
|
||||
|
||||
def test_roundtrip(self):
|
||||
u = TokenUsage(input_tokens=5, output_tokens=3, cache_creation_input_tokens=1, cache_read_input_tokens=0)
|
||||
assert TokenUsage.from_dict(u.to_dict()) == u
|
||||
|
||||
|
||||
class TestContentBlocks:
|
||||
def test_text_block_roundtrip(self):
|
||||
b = TextBlock(text="hello")
|
||||
d = b.to_dict()
|
||||
assert d == {"type": "text", "text": "hello"}
|
||||
assert TextBlock.from_dict(d) == b
|
||||
|
||||
def test_tool_use_block_roundtrip(self):
|
||||
b = ToolUseBlock(id="t1", name="bash", input={"command": "ls"})
|
||||
d = b.to_dict()
|
||||
assert d["type"] == "tool_use"
|
||||
assert ToolUseBlock.from_dict(d) == b
|
||||
|
||||
def test_tool_result_block_roundtrip(self):
|
||||
b = ToolResultBlock(tool_use_id="t1", tool_name="bash", output="ok", is_error=False)
|
||||
d = b.to_dict()
|
||||
assert d["type"] == "tool_result"
|
||||
assert ToolResultBlock.from_dict(d) == b
|
||||
|
||||
|
||||
class TestConversationMessage:
|
||||
def test_user_text_factory(self):
|
||||
m = ConversationMessage.user_text("hello")
|
||||
assert m.role == MessageRole.User
|
||||
assert len(m.blocks) == 1
|
||||
assert isinstance(m.blocks[0], TextBlock)
|
||||
|
||||
def test_tool_result_factory(self):
|
||||
m = ConversationMessage.tool_result("t1", "bash", "output", is_error=True)
|
||||
assert m.role == MessageRole.Tool
|
||||
assert isinstance(m.blocks[0], ToolResultBlock)
|
||||
assert m.blocks[0].is_error
|
||||
|
||||
def test_assistant_with_usage(self):
|
||||
u = TokenUsage(input_tokens=10, output_tokens=4)
|
||||
m = ConversationMessage.assistant_with_usage([TextBlock(text="hi")], u)
|
||||
assert m.usage == u
|
||||
|
||||
def test_roundtrip(self):
|
||||
m = ConversationMessage.user_text("test")
|
||||
assert ConversationMessage.from_dict(m.to_dict()).role == m.role
|
||||
|
||||
|
||||
class TestSession:
|
||||
def test_empty_session(self):
|
||||
s = Session()
|
||||
assert s.version == 1
|
||||
assert s.messages == []
|
||||
|
||||
def test_roundtrip_dict(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="hi")]))
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "ok", False))
|
||||
|
||||
restored = Session.from_dict(s.to_dict())
|
||||
assert len(restored.messages) == 3
|
||||
assert restored.messages[0].role == MessageRole.User
|
||||
assert restored.messages[2].role == MessageRole.Tool
|
||||
|
||||
def test_save_and_load(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
s.messages.append(
|
||||
ConversationMessage.assistant_with_usage(
|
||||
[
|
||||
TextBlock(text="thinking"),
|
||||
ToolUseBlock(id="t1", name="bash", input={"cmd": "echo hi"}),
|
||||
],
|
||||
TokenUsage(input_tokens=10, output_tokens=4, cache_creation_input_tokens=1, cache_read_input_tokens=2),
|
||||
)
|
||||
)
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "hi", False))
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
|
||||
path = f.name
|
||||
|
||||
s.save(path)
|
||||
restored = Session.load(path)
|
||||
|
||||
assert restored == s
|
||||
assert restored.messages[1].usage.total_tokens() == 17
|
||||
assert restored.messages[2].role == MessageRole.Tool
|
||||
Path(path).unlink()
|
||||
|
||||
|
||||
class TestCompaction:
|
||||
def _big_session(self, n: int = 10) -> Session:
|
||||
s = Session()
|
||||
for i in range(n):
|
||||
s.messages.append(ConversationMessage.user_text("x " * 200))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="y " * 200)]))
|
||||
return s
|
||||
|
||||
def test_small_session_unchanged(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
result = compact_session(s, CompactionConfig())
|
||||
assert result.removed_message_count == 0
|
||||
assert result.compacted_session is s
|
||||
|
||||
def test_compacts_old_messages(self):
|
||||
s = self._big_session(10)
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
|
||||
assert result.removed_message_count == 18
|
||||
assert result.compacted_session.messages[0].role == MessageRole.System
|
||||
assert len(result.compacted_session.messages) == 3 # system + 2 preserved
|
||||
assert result.summary != ""
|
||||
assert result.formatted_summary != ""
|
||||
|
||||
def test_should_compact_checks_both(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("short"))
|
||||
# Few messages, low tokens - should not compact
|
||||
assert not should_compact(s, CompactionConfig(preserve_recent_messages=0, max_estimated_tokens=100))
|
||||
|
||||
def test_estimate_tokens(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("a" * 400))
|
||||
tokens = estimate_session_tokens(s)
|
||||
assert tokens == 400 // 4 + 1 # 101
|
||||
|
||||
def test_compaction_reduces_tokens(self):
|
||||
s = self._big_session(10)
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert estimate_session_tokens(result.compacted_session) < estimate_session_tokens(s)
|
||||
|
||||
def test_summary_contains_tool_names(self):
|
||||
s = Session()
|
||||
for _ in range(6):
|
||||
s.messages.append(ConversationMessage.user_text("x " * 200))
|
||||
s.messages.append(
|
||||
ConversationMessage.assistant(
|
||||
[TextBlock(text="ok"), ToolUseBlock(id="t1", name="bash", input={"cmd": "ls"})]
|
||||
)
|
||||
)
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "output " * 50, False))
|
||||
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert "bash" in result.summary
|
||||
|
||||
def test_summary_contains_file_candidates(self):
|
||||
s = Session()
|
||||
for _ in range(5):
|
||||
s.messages.append(ConversationMessage.user_text("Update src/main.rs and pkg/config.py next " + "x " * 200))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="done " * 200)]))
|
||||
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert "src/main.rs" in result.summary
|
||||
|
||||
def test_format_compact_summary_strips_analysis(self):
|
||||
summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>"
|
||||
formatted = format_compact_summary(summary)
|
||||
assert "analysis" not in formatted
|
||||
assert "Kept work" in formatted
|
||||
|
||||
def test_format_compact_summary_no_tags(self):
|
||||
formatted = format_compact_summary("plain text summary")
|
||||
assert formatted == "plain text summary"
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Tests for SSE chunking and decoding in hanzoai._streaming."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Iterator, AsyncIterator
|
||||
|
||||
import pytest
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "pkg" / "hanzoai"))
|
||||
|
||||
from hanzoai._streaming import SSEDecoder, ServerSentEvent
|
||||
|
||||
|
||||
def make_sync_iterator(chunks: list[bytes]) -> Iterator[bytes]:
|
||||
yield from chunks
|
||||
|
||||
|
||||
async def make_async_iterator(chunks: list[bytes]) -> AsyncIterator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
class TestIterChunks:
|
||||
"""Test that _iter_chunks correctly splits SSE frames."""
|
||||
|
||||
def _collect_chunks(self, raw_chunks: list[bytes]) -> list[bytes]:
|
||||
decoder = SSEDecoder()
|
||||
return list(decoder._iter_chunks(make_sync_iterator(raw_chunks)))
|
||||
|
||||
def test_simple_frame_lf(self):
|
||||
"""Single frame terminated by \\n\\n."""
|
||||
raw = [b"data: hello\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_simple_frame_crlf(self):
|
||||
"""Single frame terminated by \\r\\n\\r\\n."""
|
||||
raw = [b"data: hello\r\n\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_simple_frame_cr(self):
|
||||
"""Single frame terminated by \\r\\r."""
|
||||
raw = [b"data: hello\r\r"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\r"]
|
||||
|
||||
def test_multiple_frames(self):
|
||||
"""Multiple frames in one byte chunk."""
|
||||
raw = [b"data: one\n\ndata: two\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
def test_split_across_chunks(self):
|
||||
"""Frame split across multiple byte chunks."""
|
||||
raw = [b"data: hel", b"lo\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_separator_split_across_chunks_lf(self):
|
||||
"""\\n\\n separator split across chunk boundary."""
|
||||
raw = [b"data: hello\n", b"\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_separator_split_across_chunks_crlf(self):
|
||||
"""\\r\\n\\r\\n separator split across chunk boundary at various points."""
|
||||
# Split between \r\n and \r\n
|
||||
raw = [b"data: hello\r\n", b"\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_crlf_split_mid_separator(self):
|
||||
"""The critical bug case: \\r\\n\\r\\n split so \\r and \\n land in different chunks."""
|
||||
# This is the case splitlines(keepends=True) gets wrong:
|
||||
# chunk1 ends with \r, chunk2 starts with \n\r\n
|
||||
raw = [b"data: hello\r", b"\n\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_crlf_split_three_ways(self):
|
||||
"""\\r\\n\\r\\n split across three chunks."""
|
||||
raw = [b"data: hello\r\n", b"\r", b"\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_trailing_data_no_terminator(self):
|
||||
"""Data without a trailing separator is yielded on stream end."""
|
||||
raw = [b"data: hello"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello"]
|
||||
|
||||
def test_empty_input(self):
|
||||
"""No data yields nothing."""
|
||||
chunks = self._collect_chunks([])
|
||||
assert chunks == []
|
||||
|
||||
def test_only_separator(self):
|
||||
"""Just a separator yields it as a chunk."""
|
||||
raw = [b"\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"\n\n"]
|
||||
|
||||
def test_multiple_frames_split(self):
|
||||
"""Two frames, each split across chunks."""
|
||||
raw = [b"data: one\n", b"\ndata: tw", b"o\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
|
||||
class TestAiterChunks:
|
||||
"""Test that _aiter_chunks matches _iter_chunks behavior."""
|
||||
|
||||
async def _collect_chunks(self, raw_chunks: list[bytes]) -> list[bytes]:
|
||||
decoder = SSEDecoder()
|
||||
result = []
|
||||
async for chunk in decoder._aiter_chunks(make_async_iterator(raw_chunks)):
|
||||
result.append(chunk)
|
||||
return result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_frame(self):
|
||||
chunks = await self._collect_chunks([b"data: hello\n\n"])
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_crlf_split_mid_separator(self):
|
||||
"""The critical bug case for async path."""
|
||||
raw = [b"data: hello\r", b"\n\r\n"]
|
||||
chunks = await self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_frames_split(self):
|
||||
raw = [b"data: one\n", b"\ndata: tw", b"o\n\n"]
|
||||
chunks = await self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
|
||||
class TestDoneAndPingSentinel:
|
||||
"""Test that [DONE] and event: ping are skipped."""
|
||||
|
||||
def _collect_events(self, raw_chunks: list[bytes]) -> list[ServerSentEvent]:
|
||||
decoder = SSEDecoder()
|
||||
return list(decoder.iter_bytes(make_sync_iterator(raw_chunks)))
|
||||
|
||||
def test_done_sentinel_skipped(self):
|
||||
"""data: [DONE] should not produce an event."""
|
||||
raw = [b"data: {\"text\": \"hi\"}\n\n", b"data: [DONE]\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
|
||||
def test_ping_event_skipped(self):
|
||||
"""event: ping should not produce an event."""
|
||||
raw = [b"event: ping\ndata: \n\n", b"data: {\"text\": \"hi\"}\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
|
||||
def test_normal_events_pass_through(self):
|
||||
"""Normal events should still work."""
|
||||
raw = [b"data: {\"a\": 1}\n\n", b"data: {\"b\": 2}\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 2
|
||||
assert events[0].json() == {"a": 1}
|
||||
assert events[1].json() == {"b": 2}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_done_sentinel_skipped_async(self):
|
||||
"""data: [DONE] should not produce an event in async path."""
|
||||
decoder = SSEDecoder()
|
||||
raw = [b"data: {\"text\": \"hi\"}\n\n", b"data: [DONE]\n\n"]
|
||||
events = []
|
||||
async for sse in decoder.aiter_bytes(make_async_iterator(raw)):
|
||||
events.append(sse)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
Reference in New Issue
Block a user