hanzoai/python-sdk
2
0
Fork 0
mirror of https://github.com/hanzoai/python-sdk.git synced 2026-10-23 20:23:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add hanzo-kms package, MCP proxy tools, and agent reflexion

Browse Source 
- Add hanzo-kms v1.0.0: Secret management SDK compatible with Infisical API
  - Sync and async clients with multiple auth methods (Universal, AWS IAM, Azure, GCP, K8s)
  - Full CRUD operations for secrets with environment injection

- Add MCP proxy system for lazy-loading external MCP servers
  - mcp_proxy.py: Dynamic MCP server connection and tool discovery
  - proxy_tool.py: Tool for managing external MCP server proxies
  - platform_auth.py: Hanzo Platform authentication utilities
  - Built-in servers: platform, github, cloudflare, postgres, sqlite, docker, k8s

- Add agent reflexion module for self-correction and rule management
  - ReflexionEngine for reflection and rule-based behavior updates
  - Integration with hanzo-memory for persistent rules

- Add self_learning_agent example demonstrating reflexion capabilities

- Update ps_tool with improved process management and log handling
- Update browser tools with enhanced CDP bridge and automation
- Bump dependency versions across all tool packages

All tests pass (3156 passed)
This commit is contained in:
zeekay
2026-02-04 10:55:08 -08:00
parent 2940310a5d
commit 73fbba8141
31 changed files with 3344 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/self_learning_agent.py
+37
View File
@@ -0,0 +1,37 @@
from agents import Agent, ReflexionEngine, Runner
from hanzo_memory.client import SQLiteMemoryClient
from agents.reflexion import Rule
async def main():
# 1. Initialize Memory Client (SQLite)
# Ensure the DB exists or is created
db_client = SQLiteMemoryClient(user_id="test_user")
# 2. Initialize Reflexion Engine
reflexion = ReflexionEngine(memory_client=db_client)
# 3. Add some initial rules (simulating past learning)
await reflexion.add_rule("Always double check file paths before writing.", context="coding")
# 4. Create an Agent with Reflexion
agent = Agent(
name="SelfLearningBot",
instructions="You are a helpful assistant that learns from mistakes.",
reflexion=reflexion
)
# 5. Simulate a run (in a real scenario, the runner would use these)
print(f"Agent {agent.name} initialized.")
rules = await agent.reflexion.load_rules(context="coding")
print(f"Loaded {len(rules)} rules for context 'coding':")
for r in rules:
print(f" - {r.content}")
# 6. Simulate a learning event
print("\nSimulating reflection...")
new_rule = await agent.reflexion.add_rule("Use pathlib for path manipulations.", context="coding")
print(f"Added new rule: {new_rule.content}")
if __name__ == "__main__":
asyncio.run(main())
pkg/hanzo-agent/pyproject.toml
+3 -4
View File
@@ -3,12 +3,14 @@ name = "hanzo-agent"
version = "0.0.4"
description = "Hanzo AI SDK"
readme = "README.md"
requires-python = ">=3.9"
requires-python = ">=3.12"
license = "MIT"
authors = [
{ name = "Hanzo AI", email = "[email protected]" },
]
dependencies = [
"hanzo-memory>=1.0.0",
"hanzo-async>=0.1.0",
"openai>=1.66.2",
"pydantic>=2.10, <3",
"griffe>=1.5.6, <2",
@@ -21,9 +23,6 @@ classifiers = [
"Typing :: Typed",
"Intended Audience :: Developers",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Intended Audience :: Developers",
"Operating System :: OS Independent",
pkg/hanzo-agent/src/agents/__init__.py
+7
View File
@@ -3,6 +3,10 @@ import sys
from typing import Literal
from openai import AsyncOpenAI
import hanzo_async
# Configure unified event loop (uvloop)
hanzo_async.configure_loop()
from . import _config
from .agent import Agent
@@ -140,6 +144,7 @@ from .orchestration import (
StreamUpdate,
UpdateType,
)
from .reflexion import ReflexionEngine, Rule
def set_default_openai_key(key: str, use_for_tracing: bool = True) -> None:
@@ -320,4 +325,6 @@ __all__ = [
"UIStreamer",
"StreamUpdate",
"UpdateType",
"ReflexionEngine",
"Rule",
]
pkg/hanzo-agent/src/agents/agent.py
+10 -1
View File
@@ -14,6 +14,7 @@ from .items import ItemHelpers
from .logger import logger
from .model_settings import ModelSettings
from .models.interface import Model
from .reflexion import ReflexionEngine
from .run_context import RunContextWrapper, TContext
from .tool import Tool, function_tool
@@ -75,7 +76,15 @@ class Agent(Generic[TContext]):
"""Configures model-specific tuning parameters (e.g. temperature, top_p).
"""
tools: list[Tool] = field(default_factory=list)
handoffs: list[str | Agent[TContext] | Handoff] = field(default_factory=list)
"""A list of agents that this agent can handoff to."""
reflexion: ReflexionEngine | None = None
"""The reflexion engine for self-correction."""
tools: list[Tool | Callable[..., Any] | Callable[..., Awaitable[Any]]] = field(
default_factory=list
)
"""A list of tools that the agent can use."""
input_guardrails: list[InputGuardrail[TContext]] = field(default_factory=list)
pkg/hanzo-agent/src/agents/reflexion.py
+71
View File
@@ -0,0 +1,71 @@
"""Reflexion module for agent self-correction and rule management.
This module provides capabilities for agents to:
1. Reflect on past actions and outcomes.
2. Maintain a set of dynamic rules/guidelines.
3. Update their own behavior based on these rules.
"""
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field
import structlog
from .memory import Memory, MemoryType
logger = structlog.get_logger()
class Rule(BaseModel):
"""A behavioral rule for the agent."""
id: str
content: str
context: str = "general"
confidence: float = 1.0
created_at: float
updated_at: float
class ReflexionEngine:
"""Engine for managing agent reflection and rules."""
def __init__(self, memory_client: Memory):
self.memory = memory_client
self.rules: Dict[str, List[Rule]] = {}
async def load_rules(self, context: str = "general") -> List[Rule]:
"""Load rules for a specific context."""
# TODO: Implement retrieval from hanzo-memory
# functionality will use memory.recall(query=f"rules for {context}")
return self.rules.get(context, [])
async def add_rule(self, content: str, context: str = "general") -> Rule:
"""Add a new rule based on reflection."""
# TODO: Implement storage in hanzo-memory
# functionality will use memory.remember(content=f"Rule: {content}", metadata={"type": "rule", "context": context})
import time
import uuid
rule = Rule(
id=str(uuid.uuid4()),
content=content,
context=context,
created_at=time.time(),
updated_at=time.time()
)
if context not in self.rules:
self.rules[context] = []
self.rules[context].append(rule)
return rule
async def reflect(self, task: str, outcome: str, success: bool) -> List[str]:
"""Analyze a task outcome and generate strictures/improvements."""
# Simple heuristic for now: if failed, suggest checking the error.
reflections = []
if not success:
reflections.append(f"Reflection: Task '{task}' failed. Consider input validation.")
reflections.append(f"Reflection: Analyze error message: {outcome}")
else:
reflections.append(f"Reflection: Task '{task}' succeeded. Reinforce this path.")
# Todo: In production, call an LLM here to analyze the trace:
# response = await self.llm.generate(f"Analyze this execution: {task} -> {outcome}")
return reflections
pkg/hanzo-agents/pyproject.toml
+1 -3
View File
@@ -4,7 +4,7 @@ version = "0.1.2"
description = "Multi-agent orchestration CLI for Hanzo AI"
readme = "README.md"
license = { text = "Apache-2.0" }
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [{ name = "Hanzo AI", email = "[email protected]" }]
keywords = ["hanzo", "ai", "agents", "orchestration", "cli"]
classifiers = [
@@ -12,8 +12,6 @@ classifiers = [
"Intended Audience :: Developers",
"License :: OSI Approved :: Apache Software License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
dependencies = [
pkg/hanzo-kms/README.md
+152
View File
@@ -0,0 +1,152 @@
# Hanzo KMS - Python SDK
Official Python SDK for [Hanzo KMS](https://kms.hanzo.ai) - Secret management for your applications.
## Installation
```bash
pip install hanzo-kms
```
Or with uv:
```bash
uv add hanzo-kms
```
## Quick Start
```python
from hanzo_kms import KMSClient, ClientSettings, AuthenticationOptions, UniversalAuthMethod
# Initialize client
client = KMSClient(ClientSettings(
site_url="https://kms.hanzo.ai",
auth=AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id="your-client-id",
client_secret="your-client-secret",
)
)
))
# List all secrets
secrets = client.list_secrets(
project_id="my-project",
environment="production"
)
for secret in secrets:
print(f"{secret.secret_key}: {secret.secret_value}")
# Get a specific secret
db_url = client.get_value(
project_id="my-project",
environment="production",
secret_name="DATABASE_URL"
)
# Inject all secrets into environment
client.inject_env(
project_id="my-project",
environment="production"
)
```
## Environment Variables
The client can be configured via environment variables:
```bash
export HANZO_KMS_URL="https://kms.hanzo.ai"
export HANZO_KMS_CLIENT_ID="your-client-id"
export HANZO_KMS_CLIENT_SECRET="your-client-secret"
```
Then simply:
```python
from hanzo_kms import KMSClient
client = KMSClient() # Uses environment variables
secrets = client.list_secrets("my-project", "production")
```
## Authentication Methods
### Universal Auth (Recommended)
```python
from hanzo_kms import KMSClient, ClientSettings, AuthenticationOptions, UniversalAuthMethod
client = KMSClient(ClientSettings(
auth=AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id="...",
client_secret="...",
)
)
))
```
### Kubernetes Auth
For workloads running in Kubernetes:
```python
from hanzo_kms import KMSClient, ClientSettings, AuthenticationOptions, KubernetesAuthMethod
client = KMSClient(ClientSettings(
auth=AuthenticationOptions(
kubernetes=KubernetesAuthMethod(
identity_id="your-identity-id",
# Uses default service account token path
)
)
))
```
### AWS IAM Auth
```python
from hanzo_kms import KMSClient, ClientSettings, AuthenticationOptions, AWSIamAuthMethod
client = KMSClient(ClientSettings(
auth=AuthenticationOptions(
aws_iam=AWSIamAuthMethod(
identity_id="your-identity-id",
)
)
))
```
## API Reference
### KMSClient
| Method | Description |
|--------|-------------|
| `list_secrets(project_id, environment, path="/")` | List all secrets |
| `get_secret(project_id, environment, secret_name)` | Get a single secret |
| `get_value(project_id, environment, secret_name, default=None)` | Get just the value |
| `create_secret(project_id, environment, secret_name, secret_value)` | Create a secret |
| `update_secret(project_id, environment, secret_name, secret_value)` | Update a secret |
| `delete_secret(project_id, environment, secret_name)` | Delete a secret |
| `inject_env(project_id, environment, overwrite=False)` | Inject into os.environ |
## Compatibility
This SDK is compatible with:
- Hanzo KMS (https://kms.hanzo.ai)
- Lux KMS (https://kms.lux.network)
- Infisical (https://infisical.com)
The `InfisicalClient` alias is provided for drop-in compatibility:
```python
from hanzo_kms import InfisicalClient # Same as KMSClient
```
## License
MIT License - see [LICENSE](LICENSE) for details.
pkg/hanzo-kms/hanzo_kms/__init__.py
+83
View File
@@ -0,0 +1,83 @@
"""
Hanzo KMS - Secret Management SDK for Python
A pure Python SDK for Hanzo/Lux KMS (compatible with Infisical API).
Usage:
from hanzo_kms import KMSClient, ClientSettings, UniversalAuthMethod, AuthenticationOptions
client = KMSClient(ClientSettings(
site_url="https://kms.hanzo.ai",
auth=AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id="your-client-id",
client_secret="your-client-secret",
)
)
))
# List all secrets
secrets = client.list_secrets(project_id="my-project", environment="production")
# Get a specific secret
secret = client.get_secret(
project_id="my-project",
environment="production",
secret_name="DATABASE_URL"
)
print(secret.secret_value)
# Inject secrets into environment
client.inject_env(project_id="my-project", environment="production")
"""
__version__ = "1.0.0"
from .client import KMSClient
from .async_client import AsyncKMSClient
from .models import (
AuthenticationOptions,
AWSIamAuthMethod,
AzureAuthMethod,
ClientSettings,
CreateSecretOptions,
DeleteSecretOptions,
GCPIamAuthMethod,
GCPIDTokenAuthMethod,
GetSecretOptions,
KubernetesAuthMethod,
ListSecretsOptions,
SecretElement,
UniversalAuthMethod,
UpdateSecretOptions,
)
# Aliases for compatibility with infisical-python
InfisicalClient = KMSClient
AsyncInfisicalClient = AsyncKMSClient
__all__ = [
# Main clients
"KMSClient",
"AsyncKMSClient",
"InfisicalClient", # Alias for compatibility
"AsyncInfisicalClient", # Async alias
# Settings
"ClientSettings",
"AuthenticationOptions",
# Auth methods
"UniversalAuthMethod",
"AWSIamAuthMethod",
"AzureAuthMethod",
"GCPIamAuthMethod",
"GCPIDTokenAuthMethod",
"KubernetesAuthMethod",
# Options
"GetSecretOptions",
"ListSecretsOptions",
"CreateSecretOptions",
"UpdateSecretOptions",
"DeleteSecretOptions",
# Response types
"SecretElement",
]
pkg/hanzo-kms/hanzo_kms/async_client.py
+345
View File
@@ -0,0 +1,345 @@
"""
Hanzo KMS Async Client - Async Python implementation
An async-first KMS client compatible with Infisical API.
"""
import os
import time
from typing import Optional
import httpx
from .models import (
AuthenticationOptions,
ClientSettings,
CreateSecretOptions,
DeleteSecretOptions,
GetSecretOptions,
ListSecretsOptions,
SecretElement,
SecretsResponse,
TokenResponse,
UpdateSecretOptions,
)
class AsyncKMSClient:
"""
Async Hanzo KMS Client for secret management.
Example:
async with AsyncKMSClient(settings) as client:
secrets = await client.list_secrets("myproject", "production")
"""
def __init__(
self,
settings: Optional[ClientSettings] = None,
debug: bool = False,
):
self.settings = settings or self._settings_from_env()
self.debug = debug
self._access_token: Optional[str] = None
self._token_expires_at: float = 0
self._http_client: Optional[httpx.AsyncClient] = None
def _settings_from_env(self) -> ClientSettings:
"""Create settings from environment variables."""
from .models import UniversalAuthMethod
site_url = os.getenv("HANZO_KMS_URL", os.getenv("INFISICAL_SITE_URL", "https://kms.hanzo.ai"))
organization = os.getenv("HANZO_KMS_ORG", "hanzo")
client_id = os.getenv("HANZO_KMS_CLIENT_ID", os.getenv("INFISICAL_CLIENT_ID", ""))
client_secret = os.getenv("HANZO_KMS_CLIENT_SECRET", os.getenv("INFISICAL_CLIENT_SECRET", ""))
auth = None
if client_id and client_secret:
auth = AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id=client_id,
client_secret=client_secret,
)
)
return ClientSettings(site_url=site_url, organization=organization, auth=auth)
@property
def http(self) -> httpx.AsyncClient:
"""Get or create async HTTP client."""
if self._http_client is None:
self._http_client = httpx.AsyncClient(
base_url=self.settings.site_url.rstrip("/"),
timeout=30.0,
headers={
"User-Agent": self.settings.user_agent,
"Content-Type": "application/json",
},
)
return self._http_client
async def _get_access_token(self) -> str:
"""Get valid access token, refreshing if needed."""
if self._access_token and time.time() < self._token_expires_at - 60:
return self._access_token
auth = self.settings.auth
if not auth:
raise ValueError("No authentication configured")
# Universal Auth
if auth.universal_auth:
response = await self.http.post(
"/api/v1/auth/universal-auth/login",
json={
"clientId": auth.universal_auth.client_id,
"clientSecret": auth.universal_auth.client_secret,
},
)
response.raise_for_status()
data = response.json()
token_data = TokenResponse.model_validate(data)
self._access_token = token_data.access_token
self._token_expires_at = time.time() + token_data.expires_in
return self._access_token
# Kubernetes Auth
if auth.kubernetes:
token_path = auth.kubernetes.service_account_token_path
if os.path.exists(token_path):
with open(token_path) as f:
k8s_token = f.read().strip()
response = await self.http.post(
"/api/v1/auth/kubernetes-auth/login",
json={
"identityId": auth.kubernetes.identity_id,
"jwt": k8s_token,
},
)
response.raise_for_status()
data = response.json()
token_data = TokenResponse.model_validate(data)
self._access_token = token_data.access_token
self._token_expires_at = time.time() + token_data.expires_in
return self._access_token
raise ValueError("No valid authentication method configured")
async def _auth_headers(self) -> dict[str, str]:
"""Get authorization headers including organization context."""
token = await self._get_access_token()
return {
"Authorization": f"Bearer {token}",
"X-Org-Name": self.settings.organization,
}
async def get_secret(
self,
project_id: str,
environment: str,
secret_name: str,
path: str = "/",
**kwargs,
) -> SecretElement:
"""Get a single secret by name."""
options = GetSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
path=path,
**kwargs,
)
response = await self.http.get(
f"/api/v3/secrets/raw/{options.secret_name}",
params={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"type": options.type,
},
headers=await self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
async def list_secrets(
self,
project_id: str,
environment: str,
path: str = "/",
attach_to_process_env: bool = False,
**kwargs,
) -> list[SecretElement]:
"""List all secrets in a project/environment."""
options = ListSecretsOptions(
project_id=project_id,
environment=environment,
path=path,
attach_to_process_env=attach_to_process_env,
**kwargs,
)
response = await self.http.get(
"/api/v3/secrets/raw",
params={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"include_imports": str(options.include_imports).lower(),
"recursive": str(options.recursive).lower(),
"expandSecretReferences": str(options.expand_secret_references).lower(),
},
headers=await self._auth_headers(),
)
response.raise_for_status()
data = response.json()
secrets_data = SecretsResponse.model_validate(data)
if options.attach_to_process_env:
for secret in secrets_data.secrets:
if secret.secret_key not in os.environ:
os.environ[secret.secret_key] = secret.secret_value
return secrets_data.secrets
async def create_secret(
self,
project_id: str,
environment: str,
secret_name: str,
secret_value: str,
**kwargs,
) -> SecretElement:
"""Create a new secret."""
options = CreateSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
secret_value=secret_value,
**kwargs,
)
response = await self.http.post(
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"secretValue": options.secret_value,
"secretComment": options.secret_comment,
"type": options.type,
},
headers=await self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
async def update_secret(
self,
project_id: str,
environment: str,
secret_name: str,
secret_value: str,
**kwargs,
) -> SecretElement:
"""Update an existing secret."""
options = UpdateSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
secret_value=secret_value,
**kwargs,
)
response = await self.http.patch(
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"secretValue": options.secret_value,
},
headers=await self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
async def delete_secret(
self,
project_id: str,
environment: str,
secret_name: str,
**kwargs,
) -> SecretElement:
"""Delete a secret."""
options = DeleteSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
**kwargs,
)
response = await self.http.request(
"DELETE",
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
},
headers=await self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
async def inject_env(
self,
project_id: str,
environment: str,
path: str = "/",
overwrite: bool = False,
) -> int:
"""Inject all secrets into environment variables."""
secrets = await self.list_secrets(project_id, environment, path)
count = 0
for secret in secrets:
if overwrite or secret.secret_key not in os.environ:
os.environ[secret.secret_key] = secret.secret_value
count += 1
return count
async def get_value(
self,
project_id: str,
environment: str,
secret_name: str,
default: Optional[str] = None,
) -> Optional[str]:
"""Get just the value of a secret."""
try:
secret = await self.get_secret(project_id, environment, secret_name)
return secret.secret_value
except httpx.HTTPStatusError:
return default
async def close(self) -> None:
"""Close the HTTP client."""
if self._http_client:
await self._http_client.aclose()
self._http_client = None
async def __aenter__(self) -> "AsyncKMSClient":
return self
async def __aexit__(self, *args) -> None:
await self.close()
# Alias for compatibility
AsyncInfisicalClient = AsyncKMSClient
pkg/hanzo-kms/hanzo_kms/client.py
+494
View File
@@ -0,0 +1,494 @@
"""
Hanzo KMS Client - Pure Python implementation
A full-featured KMS client compatible with Infisical API.
"""
import os
import time
from typing import Optional
import httpx
from .models import (
AuthenticationOptions,
ClientSettings,
CreateSecretOptions,
DeleteSecretOptions,
GetSecretOptions,
ListSecretsOptions,
SecretElement,
SecretsResponse,
TokenResponse,
UpdateSecretOptions,
)
class KMSClient:
"""
Hanzo KMS Client for secret management.
Supports multiple authentication methods:
- Universal Auth (client_id/client_secret)
- AWS IAM
- Azure AD
- GCP IAM
- Kubernetes Service Account
Example:
client = KMSClient(ClientSettings(
site_url="https://kms.hanzo.ai",
auth=AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id="your-client-id",
client_secret="your-client-secret",
)
)
))
secrets = client.list_secrets(project_id="myproject", environment="production")
"""
def __init__(
self,
settings: Optional[ClientSettings] = None,
debug: bool = False,
):
"""Initialize KMS client.
Args:
settings: Client configuration (defaults from environment if not provided)
debug: Enable debug logging
"""
self.settings = settings or self._settings_from_env()
self.debug = debug
self._access_token: Optional[str] = None
self._token_expires_at: float = 0
self._http_client: Optional[httpx.Client] = None
def _settings_from_env(self) -> ClientSettings:
"""Create settings from environment variables."""
from .models import UniversalAuthMethod
site_url = os.getenv("HANZO_KMS_URL", os.getenv("INFISICAL_SITE_URL", "https://kms.hanzo.ai"))
organization = os.getenv("HANZO_KMS_ORG", "hanzo")
client_id = os.getenv("HANZO_KMS_CLIENT_ID", os.getenv("INFISICAL_CLIENT_ID", ""))
client_secret = os.getenv("HANZO_KMS_CLIENT_SECRET", os.getenv("INFISICAL_CLIENT_SECRET", ""))
auth = None
if client_id and client_secret:
auth = AuthenticationOptions(
universal_auth=UniversalAuthMethod(
client_id=client_id,
client_secret=client_secret,
)
)
return ClientSettings(site_url=site_url, organization=organization, auth=auth)
@property
def http(self) -> httpx.Client:
"""Get or create HTTP client."""
if self._http_client is None:
self._http_client = httpx.Client(
base_url=self.settings.site_url.rstrip("/"),
timeout=30.0,
headers={
"User-Agent": self.settings.user_agent,
"Content-Type": "application/json",
},
)
return self._http_client
def _get_access_token(self) -> str:
"""Get valid access token, refreshing if needed."""
if self._access_token and time.time() < self._token_expires_at - 60:
return self._access_token
auth = self.settings.auth
if not auth:
raise ValueError("No authentication configured")
# Universal Auth
if auth.universal_auth:
response = self.http.post(
"/api/v1/auth/universal-auth/login",
json={
"clientId": auth.universal_auth.client_id,
"clientSecret": auth.universal_auth.client_secret,
},
)
response.raise_for_status()
data = response.json()
token_data = TokenResponse.model_validate(data)
self._access_token = token_data.access_token
self._token_expires_at = time.time() + token_data.expires_in
return self._access_token
# Kubernetes Auth
if auth.kubernetes:
token_path = auth.kubernetes.service_account_token_path
if os.path.exists(token_path):
with open(token_path) as f:
k8s_token = f.read().strip()
response = self.http.post(
"/api/v1/auth/kubernetes-auth/login",
json={
"identityId": auth.kubernetes.identity_id,
"jwt": k8s_token,
},
)
response.raise_for_status()
data = response.json()
token_data = TokenResponse.model_validate(data)
self._access_token = token_data.access_token
self._token_expires_at = time.time() + token_data.expires_in
return self._access_token
# AWS IAM Auth
if auth.aws_iam:
# Get AWS credentials from environment/instance metadata
import json
try:
import boto3
session = boto3.Session()
credentials = session.get_credentials()
region = session.region_name or "us-east-1"
response = self.http.post(
"/api/v1/auth/aws-auth/login",
json={
"identityId": auth.aws_iam.identity_id,
"iamHttpRequestMethod": "POST",
"iamRequestBody": "",
"iamRequestHeaders": json.dumps({
"X-Amz-Date": credentials.token or "",
}),
},
)
response.raise_for_status()
data = response.json()
token_data = TokenResponse.model_validate(data)
self._access_token = token_data.access_token
self._token_expires_at = time.time() + token_data.expires_in
return self._access_token
except ImportError:
raise ValueError("boto3 required for AWS IAM auth: pip install boto3")
raise ValueError("No valid authentication method configured")
def _auth_headers(self) -> dict[str, str]:
"""Get authorization headers including organization context."""
token = self._get_access_token()
return {
"Authorization": f"Bearer {token}",
"X-Org-Name": self.settings.organization,
}
def _resolve_project_id(self, project_id: str) -> tuple[str, str]:
"""Resolve org-scoped project ID.
Supports formats:
- "project" -> uses settings.organization
- "org/project" -> uses specified org
Returns:
Tuple of (organization, project_id)
"""
if "/" in project_id:
org, proj = project_id.split("/", 1)
return (org, proj)
return (self.settings.organization, project_id)
# =========================================================================
# Secret Operations
# =========================================================================
def get_secret(
self,
project_id: str,
environment: str,
secret_name: str,
path: str = "/",
**kwargs,
) -> SecretElement:
"""Get a single secret by name.
Args:
project_id: Project ID or slug
environment: Environment slug (e.g., "production")
secret_name: Name of the secret
path: Secret path (default "/")
Returns:
SecretElement with the secret data
"""
options = GetSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
path=path,
**kwargs,
)
response = self.http.get(
f"/api/v3/secrets/raw/{options.secret_name}",
params={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"type": options.type,
"include_imports": str(options.include_imports).lower(),
},
headers=self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
def list_secrets(
self,
project_id: str,
environment: str,
path: str = "/",
attach_to_process_env: bool = False,
**kwargs,
) -> list[SecretElement]:
"""List all secrets in a project/environment.
Args:
project_id: Project ID or slug
environment: Environment slug
path: Secret path (default "/")
attach_to_process_env: If True, set secrets as environment variables
Returns:
List of SecretElement
"""
options = ListSecretsOptions(
project_id=project_id,
environment=environment,
path=path,
attach_to_process_env=attach_to_process_env,
**kwargs,
)
response = self.http.get(
"/api/v3/secrets/raw",
params={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"include_imports": str(options.include_imports).lower(),
"recursive": str(options.recursive).lower(),
"expandSecretReferences": str(options.expand_secret_references).lower(),
},
headers=self._auth_headers(),
)
response.raise_for_status()
data = response.json()
secrets_data = SecretsResponse.model_validate(data)
# Optionally inject into environment
if options.attach_to_process_env:
for secret in secrets_data.secrets:
if secret.secret_key not in os.environ:
os.environ[secret.secret_key] = secret.secret_value
return secrets_data.secrets
def create_secret(
self,
project_id: str,
environment: str,
secret_name: str,
secret_value: str,
**kwargs,
) -> SecretElement:
"""Create a new secret.
Args:
project_id: Project ID or slug
environment: Environment slug
secret_name: Name of the secret
secret_value: Value of the secret
Returns:
Created SecretElement
"""
options = CreateSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
secret_value=secret_value,
**kwargs,
)
response = self.http.post(
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"secretValue": options.secret_value,
"secretComment": options.secret_comment,
"type": options.type,
},
headers=self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
def update_secret(
self,
project_id: str,
environment: str,
secret_name: str,
secret_value: str,
**kwargs,
) -> SecretElement:
"""Update an existing secret.
Args:
project_id: Project ID or slug
environment: Environment slug
secret_name: Name of the secret
secret_value: New value
Returns:
Updated SecretElement
"""
options = UpdateSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
secret_value=secret_value,
**kwargs,
)
response = self.http.patch(
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"secretValue": options.secret_value,
"secretComment": options.secret_comment,
"type": options.type,
},
headers=self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
def delete_secret(
self,
project_id: str,
environment: str,
secret_name: str,
**kwargs,
) -> SecretElement:
"""Delete a secret.
Args:
project_id: Project ID or slug
environment: Environment slug
secret_name: Name of the secret
Returns:
Deleted SecretElement
"""
options = DeleteSecretOptions(
project_id=project_id,
environment=environment,
secret_name=secret_name,
**kwargs,
)
response = self.http.request(
"DELETE",
f"/api/v3/secrets/raw/{options.secret_name}",
json={
"workspaceId": options.project_id,
"environment": options.environment,
"secretPath": options.path,
"type": options.type,
},
headers=self._auth_headers(),
)
response.raise_for_status()
data = response.json()
return SecretElement.model_validate(data.get("secret", data))
# =========================================================================
# Convenience Methods
# =========================================================================
def inject_env(
self,
project_id: str,
environment: str,
path: str = "/",
overwrite: bool = False,
) -> int:
"""Inject all secrets into environment variables.
Args:
project_id: Project ID or slug
environment: Environment slug
path: Secret path
overwrite: If True, overwrite existing env vars
Returns:
Number of secrets injected
"""
secrets = self.list_secrets(project_id, environment, path)
count = 0
for secret in secrets:
if overwrite or secret.secret_key not in os.environ:
os.environ[secret.secret_key] = secret.secret_value
count += 1
return count
def get_value(
self,
project_id: str,
environment: str,
secret_name: str,
default: Optional[str] = None,
) -> Optional[str]:
"""Get just the value of a secret.
Args:
project_id: Project ID or slug
environment: Environment slug
secret_name: Name of the secret
default: Default value if secret not found
Returns:
Secret value or default
"""
try:
secret = self.get_secret(project_id, environment, secret_name)
return secret.secret_value
except httpx.HTTPStatusError:
return default
def close(self) -> None:
"""Close the HTTP client."""
if self._http_client:
self._http_client.close()
self._http_client = None
def __enter__(self) -> "KMSClient":
return self
def __exit__(self, *args) -> None:
self.close()
# Alias for compatibility
InfisicalClient = KMSClient
pkg/hanzo-kms/hanzo_kms/models.py
+199
View File
@@ -0,0 +1,199 @@
"""
Hanzo KMS - Data models for the SDK
Pydantic models for API requests and responses.
Compatible with Infisical API schema.
"""
from datetime import datetime
from typing import Any, Optional
from pydantic import BaseModel, Field
# =============================================================================
# Authentication Models
# =============================================================================
class UniversalAuthMethod(BaseModel):
"""Universal authentication using client credentials."""
client_id: str = Field(..., description="Client ID from KMS")
client_secret: str = Field(..., description="Client secret from KMS")
class AWSIamAuthMethod(BaseModel):
"""AWS IAM authentication."""
identity_id: str = Field(..., description="Identity ID in KMS")
class AzureAuthMethod(BaseModel):
"""Azure AD authentication."""
identity_id: str = Field(..., description="Identity ID in KMS")
resource: Optional[str] = Field(None, description="Azure resource")
class GCPIamAuthMethod(BaseModel):
"""GCP IAM authentication."""
identity_id: str = Field(..., description="Identity ID in KMS")
service_account_key_file_path: str = Field(..., description="Path to service account key")
class GCPIDTokenAuthMethod(BaseModel):
"""GCP ID Token authentication."""
identity_id: str = Field(..., description="Identity ID in KMS")
class KubernetesAuthMethod(BaseModel):
"""Kubernetes service account authentication."""
identity_id: str = Field(..., description="Identity ID in KMS")
service_account_token_path: str = Field(
"/var/run/secrets/kubernetes.io/serviceaccount/token",
description="Path to service account token",
)
class AuthenticationOptions(BaseModel):
"""Authentication configuration - use one method."""
universal_auth: Optional[UniversalAuthMethod] = None
aws_iam: Optional[AWSIamAuthMethod] = None
azure: Optional[AzureAuthMethod] = None
gcp_iam: Optional[GCPIamAuthMethod] = None
gcp_id_token: Optional[GCPIDTokenAuthMethod] = None
kubernetes: Optional[KubernetesAuthMethod] = None
# =============================================================================
# Client Settings
# =============================================================================
class ClientSettings(BaseModel):
"""Client configuration settings."""
site_url: str = Field("https://kms.hanzo.ai", description="KMS API URL")
organization: str = Field("hanzo", description="Organization name for multi-tenancy")
auth: Optional[AuthenticationOptions] = Field(None, description="Authentication options")
user_agent: str = Field("hanzo-kms-python", description="User agent string")
cache_ttl: int = Field(300, description="Cache TTL in seconds")
# =============================================================================
# Secret Models
# =============================================================================
class SecretElement(BaseModel):
"""A secret from KMS."""
id: str = Field(..., description="Secret ID")
secret_key: str = Field(..., alias="secretKey", description="Secret key/name")
secret_value: str = Field(..., alias="secretValue", description="Secret value")
secret_comment: Optional[str] = Field(None, alias="secretComment", description="Comment")
version: int = Field(1, description="Secret version")
type: str = Field("shared", description="Secret type")
environment: str = Field(..., description="Environment slug")
workspace: str = Field(..., description="Workspace/project ID")
created_at: Optional[datetime] = Field(None, alias="createdAt")
updated_at: Optional[datetime] = Field(None, alias="updatedAt")
class Config:
populate_by_name = True
# =============================================================================
# Request Options
# =============================================================================
class GetSecretOptions(BaseModel):
"""Options for getting a single secret."""
project_id: str = Field(..., description="Project ID or slug")
environment: str = Field(..., description="Environment slug")
secret_name: str = Field(..., description="Secret key/name")
path: str = Field("/", description="Secret path")
type: str = Field("shared", description="Secret type")
include_imports: bool = Field(True, description="Include imported secrets")
class ListSecretsOptions(BaseModel):
"""Options for listing secrets."""
project_id: str = Field(..., description="Project ID or slug")
environment: str = Field(..., description="Environment slug")
path: str = Field("/", description="Secret path")
include_imports: bool = Field(True, description="Include imported secrets")
recursive: bool = Field(False, description="Recursively fetch from subpaths")
expand_secret_references: bool = Field(True, description="Expand ${} references")
attach_to_process_env: bool = Field(False, description="Set as environment variables")
class CreateSecretOptions(BaseModel):
"""Options for creating a secret."""
project_id: str = Field(..., description="Project ID or slug")
environment: str = Field(..., description="Environment slug")
secret_name: str = Field(..., description="Secret key/name")
secret_value: str = Field(..., description="Secret value")
secret_comment: Optional[str] = Field(None, description="Comment")
path: str = Field("/", description="Secret path")
type: str = Field("shared", description="Secret type")
class UpdateSecretOptions(BaseModel):
"""Options for updating a secret."""
project_id: str = Field(..., description="Project ID or slug")
environment: str = Field(..., description="Environment slug")
secret_name: str = Field(..., description="Secret key/name")
secret_value: str = Field(..., description="New secret value")
secret_comment: Optional[str] = Field(None, description="Comment")
path: str = Field("/", description="Secret path")
type: str = Field("shared", description="Secret type")
class DeleteSecretOptions(BaseModel):
"""Options for deleting a secret."""
project_id: str = Field(..., description="Project ID or slug")
environment: str = Field(..., description="Environment slug")
secret_name: str = Field(..., description="Secret key/name")
path: str = Field("/", description="Secret path")
type: str = Field("shared", description="Secret type")
# =============================================================================
# Response Models
# =============================================================================
class TokenResponse(BaseModel):
"""Access token response from authentication."""
access_token: str = Field(..., alias="accessToken")
expires_in: int = Field(..., alias="expiresIn")
token_type: str = Field("Bearer", alias="tokenType")
class Config:
populate_by_name = True
class SecretsResponse(BaseModel):
"""Response containing list of secrets."""
secrets: list[SecretElement]
imports: Optional[list[Any]] = None
class SecretResponse(BaseModel):
"""Response containing a single secret."""
secret: SecretElement
pkg/hanzo-kms/hanzo_kms/py.typed
View File
pkg/hanzo-kms/pyproject.toml
+61
View File
@@ -0,0 +1,61 @@
[project]
name = "hanzo-kms"
version = "1.0.0"
description = "Hanzo KMS SDK - Secret management for Python"
readme = "README.md"
license = { text = "MIT" }
requires-python = ">=3.9"
authors = [{ name = "Hanzo AI", email = "[email protected]" }]
keywords = ["secrets", "kms", "hanzo", "lux", "security", "vault"]
classifiers = [
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Topic :: Security",
"Topic :: Software Development :: Libraries :: Python Modules",
"Typing :: Typed",
]
dependencies = [
"httpx>=0.25.0",
"pydantic>=2.0.0",
]
[project.optional-dependencies]
async = ["httpx[http2]>=0.25.0"]
dev = [
"pytest>=8.0.0",
"pytest-asyncio>=0.24.0",
"ruff>=0.5.0",
"mypy>=1.10.0",
]
[project.urls]
Homepage = "https://kms.hanzo.ai"
Documentation = "https://docs.hanzo.ai/kms"
Repository = "https://github.com/hanzoai/python-sdk"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["hanzo_kms"]
[tool.ruff]
target-version = "py39"
line-length = 100
[tool.ruff.lint]
select = ["E", "W", "F", "I", "B", "C4", "UP"]
ignore = ["E501"]
[tool.mypy]
python_version = "3.9"
strict = true
pkg/hanzo-repl/pyproject.toml
+1 -1
View File
@@ -23,7 +23,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
keywords = ["repl", "mcp", "ai", "llm", "hanzo", "claude", "interactive"]
requires-python = ">=3.8"
requires-python = ">=3.12"
dependencies = [
"hanzo-mcp>=0.1.0",
pkg/hanzo-tools-api/hanzo_tools/api/openapi_client.py
+6
View File
@@ -643,6 +643,12 @@ class OpenAPIClient:
Returns:
OperationListResult with matching operations
"""
# Ensure spec is parsed if we have it
if self._spec and not self._operations:
logger.debug(f"Spec loaded but operations empty, re-parsing for {self.provider}")
self._parsed = False
self._parse_spec()
ops = list(self._operations.values())
# Apply filters
pkg/hanzo-tools-browser/hanzo_tools/browser/__init__.py
+122 -6
View File
@@ -4,13 +4,17 @@ High-performance async browser automation using Playwright.
Supports shared browser instance for low latency and cross-MCP sharing via CDP.
Integration with Hanzo Browser Extension:
# Start the CDP bridge server
python -m hanzo_tools.browser.cdp_bridge_server
# The browser extension connects to ws://localhost:9223/cdp
# hanzo-mcp can then control browser tabs through the extension
The CDP bridge server starts automatically when browser tools are registered.
The browser extension connects to ws://localhost:9223/cdp automatically.
No manual setup required - just install the extension and use hanzo-mcp.
"""
import asyncio
import logging
import os
import threading
from typing import Optional
from mcp.server import FastMCP
from hanzo_tools.core import BaseTool, ToolRegistry
@@ -23,6 +27,8 @@ from hanzo_tools.browser.browser_tool import (
launch_browser_server,
)
logger = logging.getLogger(__name__)
# CDP Bridge for browser extension integration
try:
from hanzo_tools.browser.cdp_bridge_server import (
@@ -35,6 +41,11 @@ except ImportError:
CDPBridgeServer = None
CDPBridgeClient = None
# Global CDP bridge server instance (singleton)
_cdp_bridge_server: Optional["CDPBridgeServer"] = None
_cdp_bridge_thread: Optional[threading.Thread] = None
_cdp_bridge_loop: Optional[asyncio.AbstractEventLoop] = None
# Tools list for entry point discovery
TOOLS = [BrowserTool]
@@ -50,6 +61,8 @@ __all__ = [
"CDPBridgeServer",
"CDPBridgeClient",
"CDP_BRIDGE_AVAILABLE",
"start_cdp_bridge",
"stop_cdp_bridge",
# Availability check
"PLAYWRIGHT_AVAILABLE",
# Registration
@@ -59,12 +72,109 @@ __all__ = [
]
def _run_cdp_bridge_server(host: str, port: int) -> None:
"""Run CDP bridge server in a background thread."""
global _cdp_bridge_server, _cdp_bridge_loop
_cdp_bridge_loop = asyncio.new_event_loop()
asyncio.set_event_loop(_cdp_bridge_loop)
_cdp_bridge_server = CDPBridgeServer(host=host, port=port)
async def run():
await _cdp_bridge_server.start()
# Keep running until stopped
while True:
await asyncio.sleep(1)
try:
_cdp_bridge_loop.run_until_complete(run())
except Exception:
pass
def start_cdp_bridge(
host: str = "localhost",
port: int = 9223,
auto_start: bool = True,
) -> bool:
"""Start the CDP bridge server for browser extension integration.
The CDP bridge enables communication between:
- hanzo-mcp's browser tool (via HTTP API on port 9224)
- Hanzo browser extension (via WebSocket on port 9223)
Args:
host: Host to bind to (default: localhost)
port: WebSocket port (default: 9223, HTTP API on 9224)
auto_start: Whether to auto-start (can be disabled via env var)
Returns:
True if bridge started, False otherwise
"""
global _cdp_bridge_thread
# Check if disabled via environment
if os.environ.get("HANZO_CDP_BRIDGE_DISABLED", "").lower() in ("1", "true", "yes"):
logger.debug("CDP bridge disabled via environment")
return False
if not CDP_BRIDGE_AVAILABLE:
logger.debug("CDP bridge not available (websockets not installed)")
return False
if _cdp_bridge_thread is not None and _cdp_bridge_thread.is_alive():
logger.debug("CDP bridge already running")
return True
# Override from environment
host = os.environ.get("HANZO_CDP_BRIDGE_HOST", host)
port = int(os.environ.get("HANZO_CDP_BRIDGE_PORT", str(port)))
try:
_cdp_bridge_thread = threading.Thread(
target=_run_cdp_bridge_server,
args=(host, port),
daemon=True,
name="cdp-bridge-server"
)
_cdp_bridge_thread.start()
logger.info(f"CDP bridge started on ws://{host}:{port}")
return True
except Exception as e:
logger.warning(f"Failed to start CDP bridge: {e}")
return False
def stop_cdp_bridge() -> None:
"""Stop the CDP bridge server."""
global _cdp_bridge_server, _cdp_bridge_thread, _cdp_bridge_loop
if _cdp_bridge_loop is not None:
try:
if _cdp_bridge_server is not None:
asyncio.run_coroutine_threadsafe(
_cdp_bridge_server.stop(),
_cdp_bridge_loop
)
except Exception:
pass
_cdp_bridge_server = None
_cdp_bridge_thread = None
_cdp_bridge_loop = None
def register_browser_tools(mcp_server: FastMCP, **kwargs) -> list[BaseTool]:
"""Register browser tools with the MCP server.
This also auto-starts the CDP bridge server for browser extension
integration. The extension can connect to ws://localhost:9223/cdp
and control browser tabs.
Args:
mcp_server: The FastMCP server instance
**kwargs: Additional arguments (headless, cdp_endpoint)
**kwargs: Additional arguments (headless, cdp_endpoint, cdp_bridge)
Returns:
List of registered tools
@@ -72,6 +182,11 @@ def register_browser_tools(mcp_server: FastMCP, **kwargs) -> list[BaseTool]:
headless = kwargs.get("headless", True)
cdp_endpoint = kwargs.get("cdp_endpoint")
# Auto-start CDP bridge for browser extension integration
auto_cdp_bridge = kwargs.get("cdp_bridge", True)
if auto_cdp_bridge and CDP_BRIDGE_AVAILABLE:
start_cdp_bridge()
tool = create_browser_tool(headless=headless, cdp_endpoint=cdp_endpoint)
ToolRegistry.register_tool(mcp_server, tool)
return [tool]
@@ -81,5 +196,6 @@ def register_tools(mcp_server: FastMCP, **kwargs) -> list[BaseTool]:
"""Register all browser tools with the MCP server.
This is the standard entry point called by the tool discovery system.
Auto-starts CDP bridge for browser extension integration.
"""
return register_browser_tools(mcp_server, **kwargs)
pkg/hanzo-tools-browser/hanzo_tools/browser/cdp_bridge_server.py
+186 -16
View File
@@ -2,8 +2,8 @@
CDP Bridge Server for Hanzo Browser Extension Integration.
This server acts as a bridge between:
1. hanzo-mcp's browser tool (Playwright)
2. The Hanzo browser extension (Chrome/Firefox)
1. hanzo-mcp's browser tool (via HTTP API on port 9224)
2. The Hanzo browser extension (via WebSocket on port 9223)
MULTI-CLIENT SUPPORT:
- Multiple browser extensions can connect simultaneously
@@ -12,18 +12,17 @@ MULTI-CLIENT SUPPORT:
- Commands can specify client_id or target_id for routing
- Default client = most recently active
Usage:
# Start the bridge server (typically on port 9223)
python -m hanzo_tools.browser.cdp_bridge_server
ARCHITECTURE:
- WebSocket server (port 9223): Browser extensions connect here
- HTTP API server (port 9224): hanzo-mcp sends commands here
# Or programmatically
from hanzo_tools.browser.cdp_bridge_server import CDPBridgeServer
server = CDPBridgeServer(port=9223)
await server.start()
The bridge auto-starts when hanzo-mcp loads browser tools.
No manual setup required.
Environment Variables:
HANZO_CDP_BRIDGE_PORT: Port for the WebSocket server (default: 9223)
HANZO_CDP_BRIDGE_HOST: Host to bind to (default: localhost)
HANZO_CDP_HTTP_PORT: Port for the HTTP API (default: 9224)
"""
import asyncio
@@ -44,6 +43,14 @@ except ImportError:
WEBSOCKETS_AVAILABLE = False
WebSocketServerProtocol = Any
# Try to import aiohttp for HTTP API server
try:
from aiohttp import web
AIOHTTP_AVAILABLE = True
except ImportError:
AIOHTTP_AVAILABLE = False
web = None
logger = logging.getLogger(__name__)
@@ -71,18 +78,24 @@ class ExtensionClient:
class CDPBridgeServer:
"""WebSocket server that bridges hanzo-mcp and browser extensions.
"""WebSocket + HTTP server that bridges hanzo-mcp and browser extensions.
Supports multiple browser extension clients simultaneously.
Architecture:
- WebSocket (port 9223): Browser extensions connect here as CDP providers
- HTTP API (port 9224): hanzo-mcp sends commands here
"""
def __init__(
self,
host: str = "localhost",
port: int = 9223,
http_port: int = 9224,
):
self.host = host
self.port = port
self.http_port = http_port
# Multi-client registry: client_id -> ExtensionClient
self.extension_clients: dict[str, ExtensionClient] = {}
@@ -93,6 +106,8 @@ class CDPBridgeServer:
self.pending_requests: dict[int, asyncio.Future] = {}
self.request_id = 0
self._server = None
self._http_server = None
self._http_runner = None
@property
def default_client_id(self) -> Optional[str]:
@@ -112,26 +127,177 @@ class CDPBridgeServer:
return self.extension_clients.get(cid) if cid else None
async def start(self) -> None:
"""Start the WebSocket server."""
"""Start the WebSocket and HTTP servers."""
if not WEBSOCKETS_AVAILABLE:
raise ImportError(
"websockets package required for CDP bridge. "
"Install with: pip install websockets"
)
# Start WebSocket server for browser extensions
self._server = await ws_serve(
self._handle_connection,
self.host,
self.port,
)
logger.info(f"CDP Bridge Server started on ws://{self.host}:{self.port}")
logger.info(f"CDP Bridge WebSocket started on ws://{self.host}:{self.port}")
# Start HTTP API server for hanzo-mcp
if AIOHTTP_AVAILABLE:
await self._start_http_server()
else:
logger.warning("aiohttp not installed, HTTP API disabled. Install with: pip install aiohttp")
async def _start_http_server(self) -> None:
"""Start the HTTP API server."""
app = web.Application()
app.router.add_get("/status", self._http_status)
app.router.add_post("/", self._http_command)
app.router.add_options("/", self._http_cors) # CORS preflight
self._http_runner = web.AppRunner(app)
await self._http_runner.setup()
self._http_server = web.TCPSite(
self._http_runner,
self.host,
self.http_port
)
await self._http_server.start()
logger.info(f"CDP Bridge HTTP API started on http://{self.host}:{self.http_port}")
async def _http_cors(self, request: "web.Request") -> "web.Response":
"""Handle CORS preflight requests."""
return web.Response(
status=200,
headers={
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "GET, POST, OPTIONS",
"Access-Control-Allow-Headers": "Content-Type",
}
)
async def _http_status(self, request: "web.Request") -> "web.Response":
"""HTTP endpoint for status check."""
connected = len(self.extension_clients) > 0
return web.json_response({
"connected": connected,
"clients": len(self.extension_clients),
"default_client_id": self.default_client_id,
}, headers={"Access-Control-Allow-Origin": "*"})
async def _http_command(self, request: "web.Request") -> "web.Response":
"""HTTP endpoint for sending commands to browser extension."""
try:
data = await request.json()
except Exception as e:
return web.json_response(
{"error": f"Invalid JSON: {e}"},
status=400,
headers={"Access-Control-Allow-Origin": "*"}
)
# Resolve target client
target_client = self._resolve_client(
client_id=data.get("client_id"),
target_id=data.get("target_id"),
)
if not target_client:
return web.json_response({
"error": "No browser extension connected"
if not self.extension_clients
else f"Client not found: {data.get('client_id') or data.get('target_id')}"
}, status=503, headers={"Access-Control-Allow-Origin": "*"})
# Update target client's last_active
target_client.last_active = time.time()
# Map HTTP action to CDP method
action = data.get("action", "")
method = self._action_to_method(action)
params = self._build_params(data)
# Forward to extension and wait for response
request_id = self._next_request_id()
future: asyncio.Future = asyncio.get_event_loop().create_future()
self.pending_requests[request_id] = future
try:
await target_client.websocket.send(json.dumps({
"id": request_id,
"method": method,
"params": params,
}))
# Wait for response with timeout
response = await asyncio.wait_for(future, timeout=30.0)
# Return result
return web.json_response({
"success": True,
"client_id": target_client.client_id,
"result": response.get("result", {}),
}, headers={"Access-Control-Allow-Origin": "*"})
except asyncio.TimeoutError:
self.pending_requests.pop(request_id, None)
return web.json_response({
"error": "Request timeout",
"client_id": target_client.client_id,
}, status=504, headers={"Access-Control-Allow-Origin": "*"})
except Exception as e:
self.pending_requests.pop(request_id, None)
return web.json_response({
"error": str(e),
}, status=500, headers={"Access-Control-Allow-Origin": "*"})
def _action_to_method(self, action: str) -> str:
"""Map HTTP action to CDP method."""
action_map = {
"navigate": "Page.navigate",
"screenshot": "hanzo.screenshot",
"click": "hanzo.click",
"fill": "hanzo.fill",
"evaluate": "Runtime.evaluate",
"tabs": "Target.getTargets",
"status": "Browser.getVersion",
}
return action_map.get(action, action)
def _build_params(self, data: dict) -> dict:
"""Build CDP params from HTTP request data."""
params = {}
# Common params
if "url" in data:
params["url"] = data["url"]
if "selector" in data:
params["selector"] = data["selector"]
if "value" in data:
params["value"] = data["value"]
if "text" in data:
params["text"] = data["text"]
if "code" in data or "expression" in data:
params["expression"] = data.get("code") or data.get("expression")
if "full_page" in data or "fullPage" in data:
params["fullPage"] = data.get("full_page") or data.get("fullPage")
if "tabId" in data or "tab_id" in data:
params["tabId"] = data.get("tabId") or data.get("tab_id")
return params
async def stop(self) -> None:
"""Stop the WebSocket server."""
"""Stop the WebSocket and HTTP servers."""
if self._http_runner:
await self._http_runner.cleanup()
logger.info("CDP Bridge HTTP API stopped")
if self._server:
self._server.close()
await self._server.wait_closed()
logger.info("CDP Bridge Server stopped")
logger.info("CDP Bridge WebSocket stopped")
async def _handle_connection(
self,
@@ -590,11 +756,15 @@ async def main():
"""Run the CDP bridge server."""
host = os.environ.get("HANZO_CDP_BRIDGE_HOST", "localhost")
port = int(os.environ.get("HANZO_CDP_BRIDGE_PORT", "9223"))
http_port = int(os.environ.get("HANZO_CDP_HTTP_PORT", "9224"))
server = CDPBridgeServer(host=host, port=port)
server = CDPBridgeServer(host=host, port=port, http_port=http_port)
await server.start()
print(f"CDP Bridge Server running on ws://{host}:{port}")
print(f"CDP Bridge Server running:")
print(f" WebSocket: ws://{host}:{port} (browser extensions connect here)")
print(f" HTTP API: http://{host}:{http_port} (hanzo-mcp sends commands here)")
print()
print("Waiting for browser extension(s) to connect...")
print("Supports multiple browsers simultaneously")
print("Press Ctrl+C to stop")
pkg/hanzo-tools-browser/pyproject.toml
+2 -1
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "hanzo-tools-browser"
version = "0.2.2"
version = "0.3.0"
description = "Browser automation tools with extension-first routing - uses Hanzo browser extension when available, Playwright as fallback"
readme = "README.md"
requires-python = ">=3.12"
@@ -16,6 +16,7 @@ dependencies = [
"mcp>=1.25.0",
"pydantic>=2.12.5",
"aiohttp>=3.9.0",
"websockets>=12.0",
]
[project.optional-dependencies]
pkg/hanzo-tools-code/pyproject.toml
+1 -1
View File
@@ -8,7 +8,7 @@ version = "0.1.0"
description = "Code semantic tools for Hanzo AI (HIP-0300)"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [
{ name = "Hanzo AI Team", email = "[email protected]" },
]
pkg/hanzo-tools-mcp/hanzo_tools/mcp_tools/__init__.py
+39
View File
@@ -5,6 +5,7 @@ Tools:
- mcp_add: Add MCP servers
- mcp_remove: Remove MCP servers
- mcp_stats: MCP statistics
- proxy: Dynamic MCP proxy for external servers (platform, github, cloudflare, etc.)
Install:
pip install hanzo-tools-mcp
@@ -48,6 +49,35 @@ except ImportError as e:
logger.debug(f"McpStatsTool not available: {e}")
McpStatsTool = None
try:
from .proxy_tool import ProxyTool
_tools.append(ProxyTool)
except ImportError as e:
logger.debug(f"ProxyTool not available: {e}")
ProxyTool = None
# Proxy utilities for programmatic use
try:
from .mcp_proxy import (
MCPProxyRegistry,
MCPServerConfig,
MCPServerConnection,
ProxiedTool,
enable_mcp_server,
call_mcp_tool,
BUILTIN_SERVERS,
)
except ImportError as e:
logger.debug(f"MCP proxy utilities not available: {e}")
MCPProxyRegistry = None
MCPServerConfig = None
MCPServerConnection = None
ProxiedTool = None
enable_mcp_server = None
call_mcp_tool = None
BUILTIN_SERVERS = {}
TOOLS = _tools
__all__ = [
@@ -56,6 +86,15 @@ __all__ = [
"McpAddTool",
"McpRemoveTool",
"McpStatsTool",
"ProxyTool",
# Proxy utilities
"MCPProxyRegistry",
"MCPServerConfig",
"MCPServerConnection",
"ProxiedTool",
"enable_mcp_server",
"call_mcp_tool",
"BUILTIN_SERVERS",
"register_tools",
]
pkg/hanzo-tools-mcp/hanzo_tools/mcp_tools/mcp_proxy.py
+574
View File
@@ -0,0 +1,574 @@
"""Dynamic MCP Proxy for lazy-loading external MCP servers.
This module provides functionality to:
1. Connect to external MCP servers on-demand
2. Discover their tools dynamically
3. Proxy tool calls through hanzo-mcp
4. Support lazy loading - only start when needed
Example servers that can be proxied:
- platform-mcp (Hanzo Platform)
- @modelcontextprotocol/server-github (GitHub)
- @modelcontextprotocol/server-cloudflare (Cloudflare)
"""
import asyncio
import json
import logging
import os
import shutil
import subprocess
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Callable, Awaitable
logger = logging.getLogger(__name__)
@dataclass
class MCPServerConfig:
"""Configuration for an external MCP server."""
name: str
command: List[str]
env: Dict[str, str] = field(default_factory=dict)
working_dir: Optional[str] = None
description: str = ""
auto_start: bool = False
lazy_load: bool = True
auth_required: bool = False
auth_env_var: Optional[str] = None
auth_url: Optional[str] = None
@dataclass
class ProxiedTool:
"""A tool proxied from an external MCP server."""
name: str
description: str
input_schema: Dict[str, Any]
server_name: str
class MCPServerConnection:
"""Connection to an external MCP server."""
def __init__(self, config: MCPServerConfig):
self.config = config
self.process: Optional[subprocess.Popen] = None
self.tools: List[ProxiedTool] = []
self.resources: List[Dict[str, Any]] = []
self.prompts: List[Dict[str, Any]] = []
self._reader: Optional[asyncio.StreamReader] = None
self._writer: Optional[asyncio.StreamWriter] = None
self._request_id = 0
self._pending_requests: Dict[int, asyncio.Future] = {}
self._read_task: Optional[asyncio.Task] = None
async def connect(self) -> bool:
"""Connect to the MCP server."""
if self.process is not None:
return True
# Check auth if required
if self.config.auth_required and self.config.auth_env_var:
if not os.environ.get(self.config.auth_env_var):
logger.warning(
f"MCP server '{self.config.name}' requires {self.config.auth_env_var} to be set"
)
return False
# Prepare environment
env = os.environ.copy()
env.update(self.config.env)
# Process environment variable references
for key, value in list(env.items()):
if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
var_name = value[2:-1]
env[key] = os.environ.get(var_name, "")
try:
# Start the MCP server process
self.process = await asyncio.create_subprocess_exec(
*self.config.command,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
cwd=self.config.working_dir,
)
self._reader = self.process.stdout
self._writer = self.process.stdin
# Start reading responses
self._read_task = asyncio.create_task(self._read_loop())
# Initialize the connection
await self._initialize()
# Discover tools
await self._discover_tools()
logger.info(
f"Connected to MCP server '{self.config.name}' with {len(self.tools)} tools"
)
return True
except Exception as e:
logger.error(f"Failed to connect to MCP server '{self.config.name}': {e}")
await self.disconnect()
return False
async def disconnect(self):
"""Disconnect from the MCP server."""
if self._read_task:
self._read_task.cancel()
try:
await self._read_task
except asyncio.CancelledError:
pass
self._read_task = None
if self.process:
self.process.terminate()
try:
await asyncio.wait_for(self.process.wait(), timeout=5.0)
except asyncio.TimeoutError:
self.process.kill()
self.process = None
self._reader = None
self._writer = None
self.tools = []
async def _read_loop(self):
"""Read responses from the MCP server."""
try:
while True:
if not self._reader:
break
# Read content-length header
header = await self._reader.readline()
if not header:
break
# Handle JSON-RPC over stdio (newline-delimited)
line = header.decode("utf-8").strip()
if not line:
continue
try:
response = json.loads(line)
# Handle response
if "id" in response and response["id"] in self._pending_requests:
future = self._pending_requests.pop(response["id"])
if "error" in response:
future.set_exception(Exception(response["error"]))
else:
future.set_result(response.get("result"))
except json.JSONDecodeError:
continue
except asyncio.CancelledError:
pass
except Exception as e:
logger.error(f"Error reading from MCP server: {e}")
async def _send_request(self, method: str, params: Optional[Dict] = None) -> Any:
"""Send a JSON-RPC request to the MCP server."""
if not self._writer:
raise ConnectionError("Not connected to MCP server")
self._request_id += 1
request_id = self._request_id
request = {
"jsonrpc": "2.0",
"id": request_id,
"method": method,
}
if params:
request["params"] = params
# Create future for response
future: asyncio.Future = asyncio.get_event_loop().create_future()
self._pending_requests[request_id] = future
# Send request
request_json = json.dumps(request) + "\n"
self._writer.write(request_json.encode("utf-8"))
await self._writer.drain()
# Wait for response with timeout
try:
result = await asyncio.wait_for(future, timeout=30.0)
return result
except asyncio.TimeoutError:
self._pending_requests.pop(request_id, None)
raise TimeoutError(f"Request to MCP server timed out: {method}")
async def _initialize(self):
"""Initialize the MCP connection."""
result = await self._send_request("initialize", {
"protocolVersion": "2024-11-05",
"capabilities": {
"roots": {"listChanged": True},
},
"clientInfo": {
"name": "hanzo-mcp-proxy",
"version": "1.0.0",
},
})
# Send initialized notification
if self._writer:
notification = {"jsonrpc": "2.0", "method": "notifications/initialized"}
self._writer.write((json.dumps(notification) + "\n").encode("utf-8"))
await self._writer.drain()
return result
async def _discover_tools(self):
"""Discover tools from the MCP server."""
try:
result = await self._send_request("tools/list")
self.tools = []
for tool in result.get("tools", []):
self.tools.append(ProxiedTool(
name=tool["name"],
description=tool.get("description", ""),
input_schema=tool.get("inputSchema", {}),
server_name=self.config.name,
))
except Exception as e:
logger.warning(f"Failed to discover tools from '{self.config.name}': {e}")
async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Call a tool on the MCP server."""
result = await self._send_request("tools/call", {
"name": tool_name,
"arguments": arguments,
})
return result
class MCPProxyRegistry:
"""Registry for managing external MCP server proxies."""
_instance: Optional["MCPProxyRegistry"] = None
CONFIG_FILE = Path.home() / ".hanzo" / "mcp" / "proxy.json"
# Pre-configured servers
BUILTIN_SERVERS: Dict[str, MCPServerConfig] = {
"platform": MCPServerConfig(
name="platform",
command=["npx", "--yes", "@hanzo/platform-mcp"],
env={"PLATFORM_API_KEY": "${PLATFORM_API_KEY}"},
description="Hanzo Platform - deploy and manage applications",
auth_required=True,
auth_env_var="PLATFORM_API_KEY",
auth_url="https://platform.hanzo.ai/settings/api-keys",
lazy_load=True,
),
"github": MCPServerConfig(
name="github",
command=["npx", "--yes", "@modelcontextprotocol/server-github"],
env={"GITHUB_TOKEN": "${GITHUB_TOKEN}"},
description="GitHub API - repositories, issues, PRs",
auth_required=True,
auth_env_var="GITHUB_TOKEN",
lazy_load=True,
),
"cloudflare": MCPServerConfig(
name="cloudflare",
command=["npx", "--yes", "mcp-server-cloudflare"],
env={"CLOUDFLARE_API_TOKEN": "${CLOUDFLARE_API_TOKEN}"},
description="Cloudflare API - DNS, Workers, R2, KV",
auth_required=True,
auth_env_var="CLOUDFLARE_API_TOKEN",
lazy_load=True,
),
"filesystem": MCPServerConfig(
name="filesystem",
command=["npx", "--yes", "@modelcontextprotocol/server-filesystem", "/tmp"],
description="File system access",
lazy_load=True,
),
"postgres": MCPServerConfig(
name="postgres",
command=["npx", "--yes", "@modelcontextprotocol/server-postgres"],
env={"DATABASE_URL": "${DATABASE_URL}"},
description="PostgreSQL database access",
auth_required=True,
auth_env_var="DATABASE_URL",
lazy_load=True,
),
"sqlite": MCPServerConfig(
name="sqlite",
command=["uvx", "mcp-server-sqlite"],
description="SQLite database access",
lazy_load=True,
),
"docker": MCPServerConfig(
name="docker",
command=["uvx", "mcp-server-docker"],
description="Docker container management",
lazy_load=True,
),
"kubernetes": MCPServerConfig(
name="kubernetes",
command=["uvx", "mcp-server-kubernetes"],
description="Kubernetes cluster management",
lazy_load=True,
),
}
def __init__(self):
self._connections: Dict[str, MCPServerConnection] = {}
self._custom_servers: Dict[str, MCPServerConfig] = {}
self._load_config()
@classmethod
def get_instance(cls) -> "MCPProxyRegistry":
"""Get singleton instance."""
if cls._instance is None:
cls._instance = MCPProxyRegistry()
return cls._instance
def _load_config(self):
"""Load custom server configurations."""
if self.CONFIG_FILE.exists():
try:
with open(self.CONFIG_FILE, "r") as f:
data = json.load(f)
for name, config in data.get("servers", {}).items():
self._custom_servers[name] = MCPServerConfig(
name=name,
command=config.get("command", []),
env=config.get("env", {}),
working_dir=config.get("working_dir"),
description=config.get("description", ""),
auto_start=config.get("auto_start", False),
lazy_load=config.get("lazy_load", True),
auth_required=config.get("auth_required", False),
auth_env_var=config.get("auth_env_var"),
auth_url=config.get("auth_url"),
)
except Exception as e:
logger.warning(f"Failed to load proxy config: {e}")
def _save_config(self):
"""Save custom server configurations."""
self.CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
servers = {}
for name, config in self._custom_servers.items():
servers[name] = {
"command": config.command,
"env": config.env,
"working_dir": config.working_dir,
"description": config.description,
"auto_start": config.auto_start,
"lazy_load": config.lazy_load,
"auth_required": config.auth_required,
"auth_env_var": config.auth_env_var,
"auth_url": config.auth_url,
}
with open(self.CONFIG_FILE, "w") as f:
json.dump({"servers": servers}, f, indent=2)
def get_server_config(self, name: str) -> Optional[MCPServerConfig]:
"""Get server configuration by name."""
if name in self._custom_servers:
return self._custom_servers[name]
if name in self.BUILTIN_SERVERS:
return self.BUILTIN_SERVERS[name]
return None
def list_servers(self) -> List[Dict[str, Any]]:
"""List all available servers."""
servers = []
# Add builtin servers
for name, config in self.BUILTIN_SERVERS.items():
is_connected = name in self._connections
auth_configured = True
if config.auth_required and config.auth_env_var:
auth_configured = bool(os.environ.get(config.auth_env_var))
servers.append({
"name": name,
"description": config.description,
"builtin": True,
"connected": is_connected,
"auth_required": config.auth_required,
"auth_configured": auth_configured,
"auth_env_var": config.auth_env_var,
"auth_url": config.auth_url,
"tool_count": len(self._connections[name].tools) if is_connected else 0,
})
# Add custom servers
for name, config in self._custom_servers.items():
if name in self.BUILTIN_SERVERS:
continue
is_connected = name in self._connections
auth_configured = True
if config.auth_required and config.auth_env_var:
auth_configured = bool(os.environ.get(config.auth_env_var))
servers.append({
"name": name,
"description": config.description,
"builtin": False,
"connected": is_connected,
"auth_required": config.auth_required,
"auth_configured": auth_configured,
"auth_env_var": config.auth_env_var,
"auth_url": config.auth_url,
"tool_count": len(self._connections[name].tools) if is_connected else 0,
})
return servers
async def enable_server(self, name: str) -> Dict[str, Any]:
"""Enable and connect to an MCP server.
This starts the server if not running and discovers its tools.
"""
config = self.get_server_config(name)
if not config:
return {"success": False, "error": f"Server '{name}' not found"}
# Check auth
if config.auth_required and config.auth_env_var:
if not os.environ.get(config.auth_env_var):
return {
"success": False,
"error": f"Authentication required. Set {config.auth_env_var}",
"auth_url": config.auth_url,
}
# Check if already connected
if name in self._connections:
conn = self._connections[name]
return {
"success": True,
"message": f"Already connected to '{name}'",
"tools": [{"name": t.name, "description": t.description} for t in conn.tools],
}
# Create and connect
conn = MCPServerConnection(config)
if await conn.connect():
self._connections[name] = conn
return {
"success": True,
"message": f"Connected to '{name}'",
"tools": [{"name": t.name, "description": t.description} for t in conn.tools],
}
else:
return {"success": False, "error": f"Failed to connect to '{name}'"}
async def disable_server(self, name: str) -> Dict[str, Any]:
"""Disable and disconnect from an MCP server."""
if name not in self._connections:
return {"success": False, "error": f"Server '{name}' is not connected"}
conn = self._connections.pop(name)
await conn.disconnect()
return {"success": True, "message": f"Disconnected from '{name}'"}
def add_server(self, config: MCPServerConfig) -> Dict[str, Any]:
"""Add a custom MCP server configuration."""
self._custom_servers[config.name] = config
self._save_config()
return {"success": True, "message": f"Added server '{config.name}'"}
def remove_server(self, name: str) -> Dict[str, Any]:
"""Remove a custom MCP server configuration."""
if name not in self._custom_servers:
if name in self.BUILTIN_SERVERS:
return {"success": False, "error": f"Cannot remove builtin server '{name}'"}
return {"success": False, "error": f"Server '{name}' not found"}
# Disconnect if connected
if name in self._connections:
asyncio.create_task(self.disable_server(name))
del self._custom_servers[name]
self._save_config()
return {"success": True, "message": f"Removed server '{name}'"}
def get_all_proxied_tools(self) -> List[ProxiedTool]:
"""Get all tools from connected servers."""
tools = []
for conn in self._connections.values():
tools.extend(conn.tools)
return tools
async def call_proxied_tool(
self, server_name: str, tool_name: str, arguments: Dict[str, Any]
) -> Any:
"""Call a tool on a proxied server.
If lazy_load is enabled, the server will be connected on first use.
"""
# Lazy load if not connected
if server_name not in self._connections:
config = self.get_server_config(server_name)
if config and config.lazy_load:
result = await self.enable_server(server_name)
if not result.get("success"):
raise ConnectionError(result.get("error", "Failed to connect"))
else:
raise ConnectionError(f"Server '{server_name}' is not connected")
conn = self._connections.get(server_name)
if not conn:
raise ConnectionError(f"Server '{server_name}' is not connected")
return await conn.call_tool(tool_name, arguments)
# Convenience function for lazy loading
async def enable_mcp_server(name: str) -> Dict[str, Any]:
"""Enable an MCP server by name.
This is a convenience function for lazy-loading MCP servers.
Example:
result = await enable_mcp_server("platform")
if result["success"]:
print(f"Enabled with {len(result['tools'])} tools")
"""
registry = MCPProxyRegistry.get_instance()
return await registry.enable_server(name)
async def call_mcp_tool(server: str, tool: str, **kwargs) -> Any:
"""Call a tool on an MCP server.
This is a convenience function that handles lazy loading.
Example:
result = await call_mcp_tool("platform", "application-list")
"""
registry = MCPProxyRegistry.get_instance()
return await registry.call_proxied_tool(server, tool, kwargs)
# Export builtin servers for convenience
BUILTIN_SERVERS = MCPProxyRegistry.BUILTIN_SERVERS
pkg/hanzo-tools-mcp/hanzo_tools/mcp_tools/platform_auth.py
+377
View File
@@ -0,0 +1,377 @@
"""Platform authentication utilities for Hanzo Platform MCP.
This module provides utilities for authenticating with Hanzo Platform,
including:
- Getting API keys
- Getting deployment tokens
- Setting up CI/CD credentials
Usage:
# Interactive login
hanzo-platform login
# Get deploy token for CI
hanzo-platform deploy-token --app console
# Check authentication status
hanzo-platform status
"""
import asyncio
import json
import os
import sys
import webbrowser
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional
import httpx
@dataclass
class PlatformConfig:
"""Configuration for Platform API."""
base_url: str = "https://platform.hanzo.ai"
api_key: Optional[str] = None
org_id: Optional[str] = None
@classmethod
def load(cls) -> "PlatformConfig":
"""Load configuration from file and environment."""
config = cls()
# Load from file
config_file = Path.home() / ".hanzo" / "platform" / "config.json"
if config_file.exists():
try:
with open(config_file, "r") as f:
data = json.load(f)
config.base_url = data.get("base_url", config.base_url)
config.api_key = data.get("api_key")
config.org_id = data.get("org_id")
except Exception:
pass
# Environment overrides
if os.environ.get("PLATFORM_URL"):
config.base_url = os.environ["PLATFORM_URL"]
if os.environ.get("PLATFORM_API_KEY"):
config.api_key = os.environ["PLATFORM_API_KEY"]
if os.environ.get("PLATFORM_ORG_ID"):
config.org_id = os.environ["PLATFORM_ORG_ID"]
return config
def save(self):
"""Save configuration to file."""
config_file = Path.home() / ".hanzo" / "platform" / "config.json"
config_file.parent.mkdir(parents=True, exist_ok=True)
data = {
"base_url": self.base_url,
"api_key": self.api_key,
"org_id": self.org_id,
}
with open(config_file, "w") as f:
json.dump(data, f, indent=2)
# Set restrictive permissions
try:
import stat
config_file.chmod(stat.S_IRUSR | stat.S_IWUSR) # 0600
except Exception:
pass
class PlatformClient:
"""Client for Hanzo Platform API."""
def __init__(self, config: Optional[PlatformConfig] = None):
self.config = config or PlatformConfig.load()
self._client: Optional[httpx.AsyncClient] = None
async def __aenter__(self):
self._client = httpx.AsyncClient(
base_url=self.config.base_url,
headers=self._get_headers(),
timeout=30.0,
)
return self
async def __aexit__(self, *args):
if self._client:
await self._client.aclose()
def _get_headers(self) -> Dict[str, str]:
headers = {
"Content-Type": "application/json",
"Accept": "application/json",
}
if self.config.api_key:
headers["x-api-key"] = self.config.api_key
return headers
async def check_auth(self) -> Dict[str, Any]:
"""Check if authentication is valid."""
if not self._client:
raise RuntimeError("Client not initialized")
try:
response = await self._client.get("/api/v1/user/me")
if response.status_code == 200:
return {"authenticated": True, "user": response.json()}
return {"authenticated": False, "error": "Invalid API key"}
except Exception as e:
return {"authenticated": False, "error": str(e)}
async def list_projects(self) -> Dict[str, Any]:
"""List all projects."""
if not self._client:
raise RuntimeError("Client not initialized")
response = await self._client.get("/api/v1/projects")
return response.json()
async def list_applications(self, project_id: str) -> Dict[str, Any]:
"""List applications in a project."""
if not self._client:
raise RuntimeError("Client not initialized")
response = await self._client.get(f"/api/v1/projects/{project_id}/applications")
return response.json()
async def get_deploy_token(self, app_id: str) -> Optional[str]:
"""Get or create a deploy token for an application."""
if not self._client:
raise RuntimeError("Client not initialized")
# Get application details including refresh token
response = await self._client.get(f"/api/v1/applications/{app_id}")
if response.status_code != 200:
return None
app_data = response.json()
return app_data.get("application", {}).get("refreshToken")
async def create_api_key(self, name: str = "hanzo-mcp") -> Optional[str]:
"""Create a new API key."""
if not self._client:
raise RuntimeError("Client not initialized")
response = await self._client.post(
"/api/v1/api-keys",
json={"name": name, "expiresAt": None},
)
if response.status_code == 201:
data = response.json()
return data.get("apiKey", {}).get("key")
return None
async def interactive_login() -> bool:
"""Interactive login flow."""
print("=== Hanzo Platform Login ===")
print()
config = PlatformConfig.load()
# Check if already logged in
if config.api_key:
async with PlatformClient(config) as client:
result = await client.check_auth()
if result.get("authenticated"):
user = result.get("user", {})
print(f"Already logged in as: {user.get('email', 'Unknown')}")
confirm = input("Re-authenticate? [y/N]: ").strip().lower()
if confirm != "y":
return True
print(f"Opening {config.base_url}/settings/api-keys in your browser...")
print()
# Try to open browser
try:
webbrowser.open(f"{config.base_url}/settings/api-keys")
except Exception:
print(f"Could not open browser. Please visit: {config.base_url}/settings/api-keys")
print("1. Create a new API key in the Platform dashboard")
print("2. Copy the API key")
print()
api_key = input("Paste your API key here: ").strip()
if not api_key:
print("Error: No API key provided")
return False
# Validate the API key
config.api_key = api_key
async with PlatformClient(config) as client:
result = await client.check_auth()
if result.get("authenticated"):
user = result.get("user", {})
print(f"\n✅ Successfully authenticated as: {user.get('email')}")
# Save the configuration
config.save()
# Also set in environment for current session
os.environ["PLATFORM_API_KEY"] = api_key
print(f"\nConfiguration saved to: ~/.hanzo/platform/config.json")
print(f"For CI/CD, set: PLATFORM_API_KEY={api_key[:8]}...")
return True
else:
print(f"\n❌ Authentication failed: {result.get('error')}")
return False
async def get_deploy_token_cli(app_name: str) -> Optional[str]:
"""Get deploy token for an application."""
config = PlatformConfig.load()
if not config.api_key:
print("Error: Not logged in. Run 'hanzo-platform login' first.")
return None
async with PlatformClient(config) as client:
# List projects to find the app
projects = await client.list_projects()
for project in projects.get("projects", []):
apps = await client.list_applications(project["id"])
for app in apps.get("applications", []):
if app.get("name") == app_name or app.get("id") == app_name:
token = await client.get_deploy_token(app["id"])
if token:
return token
print(f"Error: Application '{app_name}' not found")
return None
async def status() -> Dict[str, Any]:
"""Check authentication status and available resources."""
config = PlatformConfig.load()
result = {
"configured": bool(config.api_key),
"base_url": config.base_url,
"authenticated": False,
"user": None,
"projects": [],
}
if not config.api_key:
return result
async with PlatformClient(config) as client:
auth_result = await client.check_auth()
result["authenticated"] = auth_result.get("authenticated", False)
result["user"] = auth_result.get("user")
if result["authenticated"]:
try:
projects = await client.list_projects()
result["projects"] = projects.get("projects", [])
except Exception:
pass
return result
def cli():
"""CLI entry point."""
import argparse
parser = argparse.ArgumentParser(
description="Hanzo Platform authentication",
prog="hanzo-platform",
)
subparsers = parser.add_subparsers(dest="command")
# Login command
subparsers.add_parser("login", help="Interactive login")
# Status command
subparsers.add_parser("status", help="Check authentication status")
# Deploy token command
deploy_parser = subparsers.add_parser("deploy-token", help="Get deploy token for app")
deploy_parser.add_argument("--app", required=True, help="Application name or ID")
deploy_parser.add_argument("--output", choices=["plain", "json"], default="plain")
# Config command
config_parser = subparsers.add_parser("config", help="Show/set configuration")
config_parser.add_argument("--set-url", help="Set Platform URL")
config_parser.add_argument("--set-key", help="Set API key")
args = parser.parse_args()
if args.command == "login":
success = asyncio.run(interactive_login())
sys.exit(0 if success else 1)
elif args.command == "status":
result = asyncio.run(status())
if result["authenticated"]:
user = result.get("user", {})
print(f"✅ Authenticated as: {user.get('email', 'Unknown')}")
print(f" Platform: {result['base_url']}")
print(f" Projects: {len(result.get('projects', []))}")
for project in result.get("projects", []):
print(f" - {project.get('name', 'Unknown')}")
else:
print("❌ Not authenticated")
print(f" Platform: {result['base_url']}")
print("\nRun 'hanzo-platform login' to authenticate")
sys.exit(0)
elif args.command == "deploy-token":
token = asyncio.run(get_deploy_token_cli(args.app))
if token:
if args.output == "json":
print(json.dumps({"token": token}))
else:
print(token)
sys.exit(0)
sys.exit(1)
elif args.command == "config":
config = PlatformConfig.load()
if args.set_url:
config.base_url = args.set_url
config.save()
print(f"Set Platform URL to: {args.set_url}")
if args.set_key:
config.api_key = args.set_key
config.save()
print(f"Set API key: {args.set_key[:8]}...")
if not args.set_url and not args.set_key:
print(f"Platform URL: {config.base_url}")
print(f"API Key: {'****' + config.api_key[-4:] if config.api_key else 'Not set'}")
print(f"Org ID: {config.org_id or 'Not set'}")
sys.exit(0)
else:
parser.print_help()
sys.exit(1)
if __name__ == "__main__":
cli()
pkg/hanzo-tools-mcp/hanzo_tools/mcp_tools/proxy_tool.py
+354
View File
@@ -0,0 +1,354 @@
"""Proxy Tool - Enable/disable external MCP servers dynamically.
This tool allows dynamically enabling external MCP servers like:
- platform (Hanzo Platform)
- github (GitHub API)
- cloudflare (Cloudflare API)
- filesystem, postgres, sqlite, docker, kubernetes, etc.
When enabled, tools from these servers become available through hanzo-mcp.
"""
import os
import json
from typing import Any, Dict, List, Optional, Unpack, Annotated, TypedDict, final, override
from pydantic import Field
from mcp.server.fastmcp import Context as MCPContext
from hanzo_tools.core import BaseTool, auto_timeout, create_tool_context
from .mcp_proxy import MCPProxyRegistry, MCPServerConfig
# Parameter types
Action = Annotated[
str,
Field(
description="Action: list, enable, disable, add, remove, call, auth",
default="list",
),
]
ServerName = Annotated[
Optional[str],
Field(
description="MCP server name (e.g., platform, github, cloudflare)",
default=None,
),
]
ToolName = Annotated[
Optional[str],
Field(
description="Tool name to call (for 'call' action)",
default=None,
),
]
ToolArgs = Annotated[
Optional[Dict[str, Any]],
Field(
description="Arguments for tool call (JSON object)",
default=None,
),
]
Command = Annotated[
Optional[str],
Field(
description="Command to run the MCP server (for 'add' action)",
default=None,
),
]
AuthToken = Annotated[
Optional[str],
Field(
description="API token for authentication (for 'auth' action)",
default=None,
),
]
class ProxyToolParams(TypedDict, total=False):
"""Parameters for the proxy tool."""
action: str
server: Optional[str]
tool: Optional[str]
args: Optional[Dict[str, Any]]
command: Optional[str]
token: Optional[str]
env: Optional[Dict[str, str]]
description: Optional[str]
@final
class ProxyTool(BaseTool):
"""Tool for managing external MCP server proxies."""
def __init__(self):
"""Initialize the proxy tool."""
self._registry = MCPProxyRegistry.get_instance()
@property
@override
def name(self) -> str:
return "proxy"
@property
@override
def description(self) -> str:
servers = self._registry.list_servers()
connected = sum(1 for s in servers if s.get("connected"))
total_tools = sum(s.get("tool_count", 0) for s in servers)
return f"""Manage external MCP server proxies. Actions: list, enable, disable, add, remove, call, auth.
Connected: {connected}/{len(servers)} servers, {total_tools} tools available.
Built-in servers: platform, github, cloudflare, filesystem, postgres, sqlite, docker, kubernetes
Usage:
proxy # List all servers
proxy --action enable --server platform # Enable Platform MCP
proxy --action disable --server github # Disable GitHub MCP
proxy --action auth --server platform --token <api-key> # Configure auth
proxy --action call --server platform --tool application-list # Call tool"""
@override
@auto_timeout("proxy")
async def call(
self,
ctx: MCPContext,
**params: Unpack[ProxyToolParams],
) -> str:
"""Execute proxy action."""
tool_ctx = None
try:
if hasattr(ctx, "client") and ctx.client:
tool_ctx = create_tool_context(ctx)
if tool_ctx:
await tool_ctx.set_tool_info(self.name)
except Exception:
pass
action = params.get("action", "list")
if action == "list":
return self._handle_list()
elif action == "enable":
return await self._handle_enable(params.get("server"))
elif action == "disable":
return await self._handle_disable(params.get("server"))
elif action == "add":
return self._handle_add(params)
elif action == "remove":
return self._handle_remove(params.get("server"))
elif action == "call":
return await self._handle_call(
params.get("server"),
params.get("tool"),
params.get("args") or {},
)
elif action == "auth":
return self._handle_auth(params.get("server"), params.get("token"))
else:
return f"Unknown action: {action}. Valid: list, enable, disable, add, remove, call, auth"
def _handle_list(self) -> str:
"""List all available MCP servers."""
servers = self._registry.list_servers()
if not servers:
return "No MCP servers available."
output = ["=== MCP Server Proxies ===", ""]
# Group by status
connected = [s for s in servers if s.get("connected")]
available = [s for s in servers if not s.get("connected") and s.get("auth_configured", True)]
needs_auth = [s for s in servers if not s.get("connected") and not s.get("auth_configured", True)]
if connected:
output.append("🟢 Connected:")
for s in connected:
output.append(f" {s['name']}: {s['description']} ({s['tool_count']} tools)")
output.append("")
if available:
output.append("⚪ Available (ready to enable):")
for s in available:
output.append(f" {s['name']}: {s['description']}")
output.append("")
if needs_auth:
output.append("🔐 Needs authentication:")
for s in needs_auth:
auth_info = f"Set {s.get('auth_env_var', 'API key')}"
if s.get("auth_url"):
auth_info += f" or visit {s['auth_url']}"
output.append(f" {s['name']}: {s['description']}")
output.append(f" → {auth_info}")
output.append("")
output.append("Commands:")
output.append(" proxy --action enable --server <name>")
output.append(" proxy --action auth --server <name> --token <key>")
return "\n".join(output)
async def _handle_enable(self, server: Optional[str]) -> str:
"""Enable an MCP server."""
if not server:
return "Error: --server is required"
result = await self._registry.enable_server(server)
if result.get("success"):
tools = result.get("tools", [])
if tools:
tool_list = "\n".join([f" - {t['name']}: {t['description']}" for t in tools[:10]])
if len(tools) > 10:
tool_list += f"\n ... and {len(tools) - 10} more"
return f"✅ Enabled '{server}' with {len(tools)} tools:\n{tool_list}"
return f"✅ {result.get('message', f'Enabled {server}')}"
else:
error = result.get("error", "Unknown error")
if result.get("auth_url"):
return f"❌ {error}\n\nGet your API key at: {result['auth_url']}\nThen run: proxy --action auth --server {server} --token <key>"
return f"❌ {error}"
async def _handle_disable(self, server: Optional[str]) -> str:
"""Disable an MCP server."""
if not server:
return "Error: --server is required"
result = await self._registry.disable_server(server)
if result.get("success"):
return f"✅ {result.get('message', f'Disabled {server}')}"
else:
return f"❌ {result.get('error', 'Unknown error')}"
def _handle_add(self, params: Dict[str, Any]) -> str:
"""Add a custom MCP server."""
server = params.get("server")
command = params.get("command")
if not server:
return "Error: --server (name) is required"
if not command:
return "Error: --command is required"
# Parse command
import shlex
cmd_parts = shlex.split(command)
config = MCPServerConfig(
name=server,
command=cmd_parts,
env=params.get("env", {}),
description=params.get("description", "Custom MCP server"),
lazy_load=True,
)
result = self._registry.add_server(config)
if result.get("success"):
return f"✅ Added server '{server}'. Use 'proxy --action enable --server {server}' to connect."
else:
return f"❌ {result.get('error', 'Unknown error')}"
def _handle_remove(self, server: Optional[str]) -> str:
"""Remove a custom MCP server."""
if not server:
return "Error: --server is required"
result = self._registry.remove_server(server)
if result.get("success"):
return f"✅ {result.get('message', f'Removed {server}')}"
else:
return f"❌ {result.get('error', 'Unknown error')}"
async def _handle_call(
self, server: Optional[str], tool: Optional[str], args: Dict[str, Any]
) -> str:
"""Call a tool on a proxied server."""
if not server:
return "Error: --server is required"
if not tool:
return "Error: --tool is required"
try:
result = await self._registry.call_proxied_tool(server, tool, args)
return json.dumps(result, indent=2) if isinstance(result, (dict, list)) else str(result)
except Exception as e:
return f"❌ Error calling {server}/{tool}: {e}"
def _handle_auth(self, server: Optional[str], token: Optional[str]) -> str:
"""Configure authentication for a server."""
if not server:
return "Error: --server is required"
config = self._registry.get_server_config(server)
if not config:
return f"Error: Server '{server}' not found"
if not config.auth_required:
return f"Server '{server}' does not require authentication"
env_var = config.auth_env_var
if not env_var:
return f"Server '{server}' has no auth configuration"
if token:
# Set the environment variable for this session
os.environ[env_var] = token
# Also save to credentials file
creds_file = self._registry.CONFIG_FILE.parent / "credentials.json"
creds_file.parent.mkdir(parents=True, exist_ok=True)
creds = {}
if creds_file.exists():
try:
with open(creds_file, "r") as f:
creds = json.load(f)
except Exception:
pass
creds[server] = {
"env_var": env_var,
"token": token, # In production, this should be encrypted
}
with open(creds_file, "w") as f:
json.dump(creds, f, indent=2)
# Set file permissions
try:
import stat
creds_file.chmod(stat.S_IRUSR | stat.S_IWUSR) # 0600
except Exception:
pass
return f"✅ Configured authentication for '{server}'.\nNow run: proxy --action enable --server {server}"
else:
# Show auth instructions
return f"""To authenticate with '{server}':
1. Get your API key:
{config.auth_url or 'Check the service documentation'}
2. Set it:
proxy --action auth --server {server} --token YOUR_API_KEY
Or set environment variable:
export {env_var}=YOUR_API_KEY"""
def register(self, mcp_server) -> None:
"""Register this tool with the MCP server."""
pass
pkg/hanzo-tools-mcp/pyproject.toml
+10 -3
View File
@@ -4,19 +4,26 @@ build-backend = "setuptools.build_meta"
[project]
name = "hanzo-tools-mcp"
version = "0.2.0"
description = "MCP management tools for Hanzo AI - server add/remove/stats"
version = "0.3.0"
description = "MCP management tools for Hanzo AI - server management, dynamic proxy for external MCPs"
requires-python = ">=3.12"
license = { text = "MIT" }
authors = [{ name = "Hanzo Industries Inc", email = "[email protected]" }]
keywords = ["hanzo", "tools", "mcp", "ai"]
keywords = ["hanzo", "tools", "mcp", "ai", "proxy", "platform"]
dependencies = [
"hanzo-tools>=0.3.0",
"mcp>=1.25.0",
"fastmcp>=2.14.1",
"pydantic>=2.12.5",
"httpx>=0.27.0",
]
[project.optional-dependencies]
platform = ["httpx>=0.27.0"]
[project.scripts]
hanzo-platform = "hanzo_tools.mcp_tools.platform_auth:cli"
[project.entry-points."hanzo.tools"]
mcp_tools = "hanzo_tools.mcp_tools:TOOLS"
pkg/hanzo-tools-net/pyproject.toml
+1 -1
View File
@@ -8,7 +8,7 @@ version = "0.1.0"
description = "Network tools for Hanzo AI (HIP-0300)"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [
{ name = "Hanzo AI Team", email = "[email protected]" },
]
pkg/hanzo-tools-plan/pyproject.toml
+1 -1
View File
@@ -8,7 +8,7 @@ version = "0.1.0"
description = "Plan orchestration tools for Hanzo AI (HIP-0300)"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [
{ name = "Hanzo AI Team", email = "[email protected]" },
]
pkg/hanzo-tools-shell/hanzo_tools/shell/ps_tool.py
+162 -129
View File
@@ -1,153 +1,170 @@
"""Process management tool (ps).
List, monitor, and control background processes.
Clean UNIX-style API.
List, monitor, and control background processes system-wide.
"""
import signal
import asyncio
from typing import Any, Dict, List, Optional, Annotated, override
import sys
import psutil
from typing import Any, Dict, List, Optional, Annotated, override, Literal
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass
from hanzo_async import read_file, path_exists
from pydantic import Field
from mcp.server import FastMCP
from mcp.server.fastmcp import Context as MCPContext
from hanzo_tools.core import BaseTool, auto_timeout
from hanzo_tools.shell.base_process import ProcessManager
@dataclass
class ProcessInfo:
"""Process information."""
id: str
pid: int
cmd: str
running: bool
exit_code: Optional[int] = None
started: Optional[datetime] = None
log_file: Optional[Path] = None
name: str
username: str
status: str
cpu_percent: float
memory_mb: float
cmdline: str
create_time: datetime
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary."""
return {
"pid": self.pid,
"name": self.name,
"username": self.username,
"status": self.status,
"cpu_percent": self.cpu_percent,
"memory_mb": self.memory_mb,
"cmdline": self.cmdline,
"create_time": self.create_time.isoformat(),
}
class PsTool(BaseTool):
"""Process management - list, kill, logs.
"""Process management - list, kill, and monitor system processes.
USAGE:
ps() # List all processes
ps(id="abc123") # Get specific process info
ps(kill="abc123") # Kill process (SIGTERM)
ps(kill="abc123", signal="KILL") # Kill with SIGKILL
ps(logs="abc123") # Get process logs
ps(logs="abc123", n=50) # Last 50 lines
ps(action="list") # List all processes (top by CPU)
ps(action="list", sort_by="memory") # List all processes (top by RAM)
ps(action="list", user="root") # List processes for user
ps(action="kill", pid=1234) # Kill process 1234 (SIGTERM)
ps(action="kill", pid=1234, sig=9) # Kill with SIGKILL
ps(action="get", pid=1234) # Get info for specific PID
"""
name = "ps"
_instance: Optional["PsTool"] = None
_process_manager: Optional[ProcessManager] = None
def __new__(cls):
"""Singleton - share process manager across instances."""
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._process_manager = ProcessManager()
return cls._instance
def __init__(self):
super().__init__()
self.process_manager = self._process_manager
@property
@override
def description(self) -> str:
return """Process management - list, kill, and read logs.
return """Process management - list, kill, and monitor system processes.
USAGE:
ps List all background processes
ps --id ID Get specific process info
ps --kill ID Kill process (SIGTERM)
ps --kill ID --sig 9 Kill with specific signal
ps --logs ID Get process stdout/stderr
ps --logs ID --n 50 Last N lines of logs"""
ps(action="list") # List top processes by CPU
ps(action="list", sort_by="memory") # List top processes by Memory
ps(action="list", limit=20) # List top 20 processes
ps(action="get", pid=1234) # Get info for process 1234
ps(action="kill", pid=1234) # Kill process 1234 (SIGTERM)
"""
def _list_processes(self) -> List[ProcessInfo]:
"""Get all tracked processes."""
processes = []
for proc_id, info in self.process_manager.list_processes().items():
processes.append(
ProcessInfo(
id=proc_id,
pid=info.get("pid", 0),
cmd=info.get("cmd", ""),
running=info.get("running", False),
exit_code=info.get("return_code"),
log_file=Path(info["log_file"]) if info.get("log_file") else None,
def _get_process_info(self, proc: psutil.Process) -> Optional[ProcessInfo]:
"""Get process info safely."""
try:
with proc.oneshot():
return ProcessInfo(
pid=proc.pid,
name=proc.name(),
username=proc.username(),
status=proc.status(),
cpu_percent=proc.cpu_percent(),
memory_mb=proc.memory_info().rss / 1024 / 1024,
cmdline=" ".join(proc.cmdline()),
create_time=datetime.fromtimestamp(proc.create_time()),
)
)
return processes
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
return None
def _get_process(self, proc_id: str) -> Optional[ProcessInfo]:
"""Get specific process info."""
for p in self._list_processes():
if p.id == proc_id or str(p.pid) == proc_id:
return p
return None
def _list_processes(
self,
sort_by: Literal["cpu", "memory", "pid", "name"] = "cpu",
limit: int = 50,
user: Optional[str] = None
) -> List[ProcessInfo]:
"""List system processes."""
processes = []
for proc in psutil.process_iter(['pid', 'name', 'username', 'status', 'cpu_percent', 'memory_info', 'cmdline', 'create_time']):
try:
# Filter by user if requested
if user and proc.info['username'] != user:
continue
async def _kill_process(self, proc_id: str, sig: int = signal.SIGTERM) -> str:
info = ProcessInfo(
pid=proc.info['pid'],
name=proc.info['name'] or "",
username=proc.info['username'] or "",
status=proc.info['status'] or "",
cpu_percent=proc.info['cpu_percent'] or 0.0,
memory_mb=(proc.info['memory_info'].rss / 1024 / 1024) if proc.info['memory_info'] else 0.0,
cmdline=" ".join(proc.info['cmdline'] or []),
create_time=datetime.fromtimestamp(proc.info['create_time']),
)
processes.append(info)
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
continue
# Sort
if sort_by == "cpu":
processes.sort(key=lambda x: x.cpu_percent, reverse=True)
elif sort_by == "memory":
processes.sort(key=lambda x: x.memory_mb, reverse=True)
elif sort_by == "pid":
processes.sort(key=lambda x: x.pid)
elif sort_by == "name":
processes.sort(key=lambda x: x.name.lower())
return processes[:limit]
def _kill_process(self, pid: int, sig: int = signal.SIGTERM) -> str:
"""Kill a process."""
process = self.process_manager.get_process(proc_id)
if not process:
return f"Process not found: {proc_id}"
try:
process = psutil.Process(pid)
process.send_signal(sig)
sig_name = signal.Signals(sig).name
return f"Sent {sig_name} to {proc_id} (PID {process.pid})"
except ProcessLookupError:
return f"Process {proc_id} already terminated"
try:
sig_name = signal.Signals(sig).name
except Exception:
sig_name = str(sig)
return f"Sent signal {sig_name} to PID {pid} ({process.name()})"
except psutil.NoSuchProcess:
return f"Process with PID {pid} not found"
except psutil.AccessDenied:
return f"Access denied to kill PID {pid}"
except Exception as e:
return f"Failed to kill {proc_id}: {e}"
async def _get_logs(self, proc_id: str, n: int = 100) -> str:
"""Get process logs."""
import asyncio
log_file = self.process_manager.get_log_file(proc_id)
if not log_file:
return f"No logs for process: {proc_id}"
# Check file exists asynchronously
if not await path_exists(log_file):
return f"No logs for process: {proc_id}"
try:
content = await read_file(log_file)
lines = content.splitlines()
if len(lines) > n:
lines = lines[-n:]
return "\n".join(lines)
except Exception as e:
return f"Error reading logs: {e}"
return f"Error killing PID {pid}: {e}"
def _format_list(self, processes: List[ProcessInfo]) -> str:
"""Format process list for display."""
if not processes:
return "No background processes"
return "No processes found"
lines = ["PID ID STATUS CMD"]
lines.append("-" * 60)
# Headers
header = f"{'PID':<8} {'USER':<15} {'%CPU':<6} {'MEM(MB)':<10} {'STATUS':<10} {'COMMAND'}"
lines = [header, "-" * len(header)]
for p in processes:
status = "running" if p.running else f"exit({p.exit_code})"
cmd = p.cmd[:30] + "..." if len(p.cmd) > 30 else p.cmd
lines.append(f"{p.pid:<7} {p.id:<15} {status:<9} {cmd}")
cmd = p.cmdline
if len(cmd) > 50:
cmd = cmd[:47] + "..."
line = f"{p.pid:<8} {p.username[:14]:<15} {p.cpu_percent:<6.1f} {p.memory_mb:<10.1f} {p.status[:9]:<10} {cmd}"
lines.append(line)
return "\n".join(lines)
@@ -156,40 +173,55 @@ USAGE:
async def call(
self,
ctx: MCPContext,
id: Optional[str] = None,
kill: Optional[str] = None,
logs: Optional[str] = None,
sig: int = 15, # SIGTERM
n: int = 100,
action: Literal["list", "kill", "get"] = "list",
pid: Optional[int] = None,
sig: int = 15,
sort_by: Literal["cpu", "memory", "pid", "name"] = "cpu",
limit: int = 50,
user: Optional[str] = None,
**kwargs,
) -> str:
"""Process management.
Args:
ctx: MCP context
id: Get specific process info
kill: Kill process by ID
logs: Get logs for process
action: Action to perform (list, kill, get)
pid: Process ID for kill/get
sig: Signal number for kill (default: 15/SIGTERM)
n: Number of log lines to show
Returns:
Process information or action result
sort_by: Sort field for list (cpu, memory, pid, name)
limit: Limit number of results for list (default: 50)
user: Filter by username
"""
if kill:
return await self._kill_process(kill, sig)
if action == "kill":
if pid is None:
return "Error: pid is required for kill action"
return self._kill_process(pid, sig)
if logs:
return await self._get_logs(logs, n)
elif action == "get":
if pid is None:
return "Error: pid is required for get action"
try:
proc = psutil.Process(pid)
info = self._get_process_info(proc)
if not info:
return f"Process {pid} not found or access denied"
return (
f"PID: {info.pid}\n"
f"Name: {info.name}\n"
f"User: {info.username}\n"
f"Status: {info.status}\n"
f"CPU: {info.cpu_percent}%\n"
f"Memory: {info.memory_mb:.1f} MB\n"
f"Created: {info.create_time}\n"
f"Command: {info.cmdline}"
)
except psutil.NoSuchProcess:
return f"Process {pid} not found"
if id:
proc = self._get_process(id)
if not proc:
return f"Process not found: {id}"
status = "running" if proc.running else f"exited ({proc.exit_code})"
return f"ID: {proc.id}\nPID: {proc.pid}\nStatus: {status}\nCommand: {proc.cmd}"
return self._format_list(self._list_processes())
else: # list
processes = self._list_processes(sort_by, limit, user)
return self._format_list(processes)
@override
def register(self, mcp_server: FastMCP) -> None:
@@ -197,15 +229,16 @@ USAGE:
tool_self = self
@mcp_server.tool(name=self.name, description=self.description)
async def ps_handler(
id: Annotated[Optional[str], Field(description="Get specific process by ID", default=None)] = None,
kill: Annotated[Optional[str], Field(description="Kill process by ID", default=None)] = None,
logs: Annotated[Optional[str], Field(description="Get logs for process ID", default=None)] = None,
async def ps(
action: Annotated[Literal["list", "kill", "get"], Field(description="Action to perform", default="list")] = "list",
pid: Annotated[Optional[int], Field(description="Process ID for kill/get", default=None)] = None,
sig: Annotated[int, Field(description="Signal for kill (default: 15/SIGTERM)", default=15)] = 15,
n: Annotated[int, Field(description="Number of log lines", default=100)] = 100,
sort_by: Annotated[Literal["cpu", "memory", "pid", "name"], Field(description="Sort field for list", default="cpu")] = "cpu",
limit: Annotated[int, Field(description="Number of processes to list", default=50)] = 50,
user: Annotated[Optional[str], Field(description="Filter by username", default=None)] = None,
ctx: MCPContext = None,
) -> str:
return await tool_self.call(ctx, id=id, kill=kill, logs=logs, sig=sig, n=n)
return await tool_self.call(ctx, action=action, pid=pid, sig=sig, sort_by=sort_by, limit=limit, user=user)
# Singleton instance
pkg/hanzo-tools-shell/pyproject.toml
+1
View File
@@ -17,6 +17,7 @@ dependencies = [
"mcp>=1.25.0",
"pydantic>=2.12.5",
"tiktoken>=0.8.0",
"psutil>=6.1.1",
]
[project.entry-points."hanzo.tools"]
pkg/hanzo-tools-shell/tests/test_ps_tool_v2.py
+42
View File
@@ -0,0 +1,42 @@
import pytest
import psutil
import os
import signal
from hanzo_tools.shell.ps_tool import ps_tool, PsTool
@pytest.mark.asyncio
async def test_ps_list():
"""Test listing processes."""
# List by CPU
result = await ps_tool.call(None, action="list", sort_by="cpu", limit=10)
assert "PID" in result
assert "COMMAND" in result
assert len(result.splitlines()) > 2 # Header + separator + at least one process
# List by Memory
result_mem = await ps_tool.call(None, action="list", sort_by="memory", limit=10)
assert "PID" in result_mem
@pytest.mark.asyncio
async def test_ps_get_self():
"""Test getting info for the current process."""
my_pid = os.getpid()
result = await ps_tool.call(None, action="get", pid=my_pid)
assert f"PID: {my_pid}" in result
assert "Status: running" in result or "Status: sleeping" in result
@pytest.mark.asyncio
async def test_ps_get_invalid():
"""Test getting info for non-existent process."""
# Max PID is usually 32768 or 4194304, so 99999999 is likely safe
result = await ps_tool.call(None, action="get", pid=99999999)
assert "not found" in result
@pytest.mark.asyncio
async def test_ps_kill_error():
"""Test killing a non-existent process."""
result = await ps_tool.call(None, action="kill", pid=99999999)
assert "not found" in result
# Note: We don't test successful kill to avoid killing random processes or the test runner itself.
pkg/hanzo-tools-test/pyproject.toml
+1 -1
View File
@@ -8,7 +8,7 @@ version = "0.1.0"
description = "Test and validation tools for Hanzo AI (HIP-0300)"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [
{ name = "Hanzo AI Team", email = "[email protected]" },
]
pkg/hanzo-tools-vcs/pyproject.toml
+1 -1
View File
@@ -8,7 +8,7 @@ version = "0.1.0"
description = "Unified version control tool for Hanzo AI (HIP-0300)"
readme = "README.md"
license = "MIT"
requires-python = ">=3.10"
requires-python = ">=3.12"
authors = [
{ name = "Hanzo AI Team", email = "[email protected]" },
]