251 lines
8.4 KiB
Python
251 lines
8.4 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Enhanced JSON-RPC Communication Module
|
||
增强版 JSON-RPC 通信模块
|
||
|
||
Uses json-rpc library for robust JSON-RPC 2.0 communication.
|
||
"""
|
||
|
||
import json
|
||
import sys
|
||
import time
|
||
import uuid
|
||
import asyncio
|
||
from typing import Any, Dict, Optional, Union, Callable, List
|
||
from dataclasses import dataclass, asdict
|
||
from enum import Enum
|
||
|
||
# json-rpc library imports
|
||
from jsonrpc import JSONRPCResponseManager, dispatcher
|
||
from jsonrpc.exceptions import JSONRPCError, JSONRPCInvalidRequest, JSONRPCMethodNotFound
|
||
from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
|
||
|
||
class ProgressLevel(str, Enum):
|
||
"""进度级别"""
|
||
INFO = "info"
|
||
SUCCESS = "success"
|
||
WARNING = "warning"
|
||
ERROR = "error"
|
||
DEBUG = "debug"
|
||
|
||
|
||
@dataclass
|
||
class ProgressUpdate:
|
||
"""进度更新数据结构"""
|
||
step: str
|
||
type: str
|
||
progress: int
|
||
message: str
|
||
timestamp: float
|
||
level: ProgressLevel = ProgressLevel.INFO
|
||
data: Optional[Dict[str, Any]] = None
|
||
|
||
|
||
class EnhancedJSONRPCResponse:
|
||
"""增强版 JSON-RPC 2.0 Response handler"""
|
||
|
||
def __init__(self, request_id: Optional[Union[str, int]] = None):
|
||
# 如果没有提供request_id，生成一个UUID
|
||
if request_id is None:
|
||
self.request_id = str(uuid.uuid4())
|
||
else:
|
||
self.request_id = request_id
|
||
|
||
def success(self, result: Any) -> None:
|
||
"""发送成功响应"""
|
||
response = JSONRPC20Response(result=result, _id=self.request_id)
|
||
self._send_response(response.data)
|
||
|
||
def error(self, code: int, message: str, data: Any = None) -> None:
|
||
"""发送错误响应"""
|
||
error_obj = {
|
||
"code": code,
|
||
"message": message
|
||
}
|
||
if data is not None:
|
||
error_obj["data"] = data
|
||
response = JSONRPC20Response(error=error_obj, _id=self.request_id)
|
||
self._send_response(response.data)
|
||
|
||
def progress(self, step: str, progress: int, message: str,
|
||
level: ProgressLevel = ProgressLevel.INFO, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""发送进度通知"""
|
||
progress_data = ProgressUpdate(
|
||
step=step,
|
||
type="progress",
|
||
progress=progress,
|
||
message=message,
|
||
timestamp=time.time(),
|
||
level=level,
|
||
data=data
|
||
)
|
||
|
||
self.notification(asdict(progress_data))
|
||
|
||
def notification(self, params: Any = None) -> None:
|
||
"""发送通知（无需响应）"""
|
||
response = JSONRPC20Response(result=params, _id=self.request_id)
|
||
self._send_response(response.data)
|
||
|
||
def _send_response(self, response: Dict[str, Any]) -> None:
|
||
"""发送响应到标准输出"""
|
||
json_str = json.dumps(response, ensure_ascii=False, separators=(',', ':'))
|
||
print(f"JSONRPC:{json_str}", file=sys.stdout, flush=True)
|
||
|
||
class EnhancedProgressReporter:
|
||
"""增强版进度报告器"""
|
||
|
||
def __init__(self, total: int = 0):
|
||
self.response = EnhancedJSONRPCResponse()
|
||
self.step = 0
|
||
self.total = total
|
||
|
||
def update(self, message: str, step: Optional[int] = None, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""更新进度"""
|
||
if step is not None:
|
||
self.step = step
|
||
else:
|
||
self.step += 1
|
||
|
||
progress = int((self.step / self.total * 100)) if self.total > 0 else -1
|
||
self.response.progress("update", progress, message, ProgressLevel.INFO, data)
|
||
|
||
def info(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""信息消息"""
|
||
self.response.progress("info", -1, message, ProgressLevel.INFO, data)
|
||
|
||
def success(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""成功消息"""
|
||
self.response.progress("success", -1, message, ProgressLevel.SUCCESS, data)
|
||
|
||
def warning(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""警告消息"""
|
||
self.response.progress("warning", -1, message, ProgressLevel.WARNING, data)
|
||
|
||
def error(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""错误消息"""
|
||
self.response.progress("error", -1, message, ProgressLevel.ERROR, data)
|
||
|
||
def debug(self, message: str, data: Optional[Dict[str, Any]] = None) -> None:
|
||
"""调试消息"""
|
||
self.response.progress("debug", -1, message, ProgressLevel.DEBUG, data)
|
||
|
||
|
||
class JSONRPCMethodRegistry:
|
||
"""JSON-RPC 方法注册器"""
|
||
|
||
def __init__(self):
|
||
self.methods: Dict[str, Callable] = {}
|
||
self.dispatcher = dispatcher
|
||
|
||
def register(self, name: Optional[str] = None):
|
||
"""注册方法装饰器"""
|
||
def decorator(func: Callable):
|
||
method_name = name or func.__name__
|
||
self.methods[method_name] = func
|
||
self.dispatcher.add_method(func, method_name)
|
||
return func
|
||
return decorator
|
||
|
||
def register_function(self, func: Callable, name: Optional[str] = None) -> None:
|
||
"""直接注册函数"""
|
||
method_name = name or func.__name__
|
||
self.methods[method_name] = func
|
||
self.dispatcher.add_method(func, method_name)
|
||
|
||
|
||
def handle_request(self, request_data: str) -> Optional[str]:
|
||
"""处理JSON-RPC请求"""
|
||
try:
|
||
# 解析请求
|
||
request_json = json.loads(request_data)
|
||
method_name = request_json.get("method")
|
||
params = request_json.get("params", {})
|
||
request_id = request_json.get("id")
|
||
|
||
# 检查方法是否存在
|
||
if method_name not in self.methods:
|
||
error_response = {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"error": {
|
||
"code": -32601,
|
||
"message": f"Method '{method_name}' not found"
|
||
}
|
||
}
|
||
return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))
|
||
|
||
try:
|
||
# 调用方法
|
||
if isinstance(params, dict):
|
||
result = self.methods[method_name](**params)
|
||
elif isinstance(params, list):
|
||
result = self.methods[method_name](*params)
|
||
else:
|
||
result = self.methods[method_name](params)
|
||
|
||
# 如果方法返回None，表示响应已经异步发送，不需要额外响应
|
||
if result is None:
|
||
return None
|
||
|
||
# 正常响应
|
||
success_response = {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"result": result
|
||
}
|
||
return json.dumps(success_response, ensure_ascii=False, separators=(',', ':'))
|
||
|
||
except Exception as e:
|
||
error_response = {
|
||
"jsonrpc": "2.0",
|
||
"id": request_id,
|
||
"error": {
|
||
"code": -32603,
|
||
"message": f"Internal error: {str(e)}"
|
||
}
|
||
}
|
||
return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))
|
||
|
||
except Exception as e:
|
||
error_response = {
|
||
"jsonrpc": "2.0",
|
||
"id": None,
|
||
"error": {
|
||
"code": -32700,
|
||
"message": f"Parse error: {str(e)}"
|
||
}
|
||
}
|
||
return json.dumps(error_response, ensure_ascii=False, separators=(',', ':'))
|
||
|
||
# 全局实例
|
||
enhanced_progress_reporter = EnhancedProgressReporter()
|
||
method_registry = JSONRPCMethodRegistry()
|
||
|
||
|
||
# 便捷函数
|
||
def create_response_handler(request_id: Optional[Union[str, int]] = None) -> EnhancedJSONRPCResponse:
|
||
"""创建响应处理器"""
|
||
return EnhancedJSONRPCResponse(request_id)
|
||
|
||
|
||
def create_progress_reporter(total: int = 0) -> EnhancedProgressReporter:
|
||
"""创建进度报告器"""
|
||
return EnhancedProgressReporter(total)
|
||
|
||
|
||
def register_method(name: Optional[str] = None):
|
||
"""注册JSON-RPC方法"""
|
||
return method_registry.register(name)
|
||
|
||
|
||
def handle_jsonrpc_request(request_data: str) -> str:
|
||
"""处理JSON-RPC请求"""
|
||
return method_registry.handle_request(request_data)
|
||
|
||
|
||
# 向后兼容的别名
|
||
JSONRPCResponse = EnhancedJSONRPCResponse
|
||
ProgressReporter = EnhancedProgressReporter
|