mirror of https://github.com/Wox-launcher/Wox
126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
import traceback
|
|
from wox_plugin import Context
|
|
import websockets
|
|
|
|
from . import logger
|
|
from .constants import PLUGIN_JSONRPC_TYPE_REQUEST, PLUGIN_JSONRPC_TYPE_RESPONSE
|
|
from .plugin_manager import waiting_for_response
|
|
from .jsonrpc import handle_request_from_wox
|
|
|
|
|
|
def _clean_for_serialization(obj):
|
|
"""Remove non-serializable properties from any object recursively"""
|
|
if obj is None:
|
|
return obj
|
|
|
|
if isinstance(obj, (str, int, float, bool)):
|
|
return obj
|
|
|
|
if isinstance(obj, (list, tuple)):
|
|
return [_clean_for_serialization(item) for item in obj]
|
|
|
|
if isinstance(obj, dict):
|
|
return {k: _clean_for_serialization(v) for k, v in obj.items()}
|
|
|
|
# Handle custom objects
|
|
if hasattr(obj, "__dict__"):
|
|
# Create a copy of the object's dict
|
|
obj_dict = obj.__dict__.copy()
|
|
|
|
# Remove callable (methods/functions) and handle nested objects
|
|
cleaned_dict = {}
|
|
for k, v in obj_dict.items():
|
|
if callable(v):
|
|
continue
|
|
cleaned_dict[k] = _clean_for_serialization(v)
|
|
|
|
return cleaned_dict
|
|
|
|
# If we can't handle it, just return None
|
|
return None
|
|
|
|
|
|
async def handle_message(ws: websockets.asyncio.server.ServerConnection, message: str):
|
|
"""Handle incoming WebSocket message"""
|
|
|
|
trace_id = str(uuid.uuid4())
|
|
try:
|
|
msg_data = json.loads(message)
|
|
if msg_data.get("TraceId"):
|
|
trace_id = msg_data.get("TraceId")
|
|
|
|
ctx = Context.new_with_value("TraceId", trace_id)
|
|
|
|
if PLUGIN_JSONRPC_TYPE_RESPONSE in message:
|
|
# Handle response from Wox
|
|
if msg_data.get("Id") in waiting_for_response:
|
|
deferred = waiting_for_response[msg_data["Id"]]
|
|
if msg_data.get("Error"):
|
|
deferred.set_exception(Exception(msg_data["Error"]))
|
|
else:
|
|
deferred.set_result(msg_data.get("Result"))
|
|
del waiting_for_response[msg_data["Id"]]
|
|
elif PLUGIN_JSONRPC_TYPE_REQUEST in message:
|
|
# Handle request from Wox
|
|
try:
|
|
result = await handle_request_from_wox(ctx, msg_data, ws)
|
|
# Clean result for serialization
|
|
cleaned_result = _clean_for_serialization(result)
|
|
|
|
response = {
|
|
"TraceId": trace_id,
|
|
"Id": msg_data["Id"],
|
|
"Method": msg_data["Method"],
|
|
"Type": PLUGIN_JSONRPC_TYPE_RESPONSE,
|
|
"Result": cleaned_result,
|
|
}
|
|
await ws.send(json.dumps(response))
|
|
except Exception as e:
|
|
error_stack = traceback.format_exc()
|
|
error_response = {
|
|
"TraceId": trace_id,
|
|
"Id": msg_data["Id"],
|
|
"Method": msg_data["Method"],
|
|
"Type": PLUGIN_JSONRPC_TYPE_RESPONSE,
|
|
"Error": str(e),
|
|
}
|
|
await logger.error(trace_id, f"handle request failed: {str(e)}\nStack trace:\n{error_stack}")
|
|
await ws.send(json.dumps(error_response))
|
|
else:
|
|
await logger.error(trace_id, f"unknown message type: {message}")
|
|
except Exception as e:
|
|
error_stack = traceback.format_exc()
|
|
await logger.error(
|
|
trace_id,
|
|
f"receive and handle msg error: {message}, err: {str(e)}\nStack trace:\n{error_stack}",
|
|
)
|
|
|
|
|
|
async def handler(websocket: websockets.asyncio.server.ServerConnection):
|
|
"""WebSocket connection handler"""
|
|
logger.update_websocket(websocket)
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
message = await websocket.recv()
|
|
asyncio.create_task(handle_message(websocket, str(message)))
|
|
except websockets.exceptions.ConnectionClosed:
|
|
await logger.info(str(uuid.uuid4()), "connection closed")
|
|
break
|
|
except Exception as e:
|
|
error_stack = traceback.format_exc()
|
|
await logger.error(str(uuid.uuid4()), f"connection error: {str(e)}\nStack trace:\n{error_stack}")
|
|
finally:
|
|
logger.update_websocket(None)
|
|
|
|
|
|
async def start_websocket(websocket_port: int):
|
|
"""Start WebSocket server"""
|
|
await logger.info(str(uuid.uuid4()), "start websocket server")
|
|
async with websockets.serve(handler, "", websocket_port):
|
|
await asyncio.Future() # run forever
|