service-forge 0.1.18__py3-none-any.whl → 0.1.39__py3-none-any.whl

service_forge/api/deprecated_websocket_api.py

@@ -1,37 +1,95 @@
 from __future__ import annotations
-import uuid
+
 import asyncio
 import json
-from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Request, Query
+import uuid
+from dataclasses import replace
+from typing import Any, Dict, Optional
+
+from fastapi import APIRouter, Depends, Query, Request, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse
 from loguru import logger
-from typing import Dict, Any, Optional
+from opentelemetry import context as otel_context_api
+from opentelemetry import propagate, trace
+
+from service_forge.execution_context import (
+    ExecutionContext,
+    reset_current_context,
+    set_current_context,
+)
+
 from .websocket_manager import websocket_manager
 
 router = APIRouter(prefix="/ws", tags=["websocket"])
+tracer = trace.get_tracer("service_forge.api.websocket")
+
 
 @router.websocket("/connect")
-async def websocket_endpoint(websocket: WebSocket, client_id: Optional[str] = Query(None)):
+async def websocket_endpoint(
+    websocket: WebSocket, client_id: Optional[str] = Query(None)
+):
     """WebSocket连接端点，支持指定客户端ID"""
-    client_id = await websocket_manager.connect(websocket, client_id)
     try:
-        while True:
-            # 接收客户端消息
-            data = await websocket.receive_text()
-            try:
-                message = json.loads(data)
-                await handle_client_message(client_id, message)
-            except json.JSONDecodeError:
-                logger.error(f"从客户端 {client_id} 收到无效JSON消息: {data}")
-                await websocket_manager.send_personal_message(
-                    json.dumps({"error": "Invalid JSON format"}),
-                    client_id
-                )
-    except WebSocketDisconnect:
-        websocket_manager.disconnect(client_id)
-    except Exception as e:
-        logger.error(f"WebSocket连接处理异常: {e}")
-        websocket_manager.disconnect(client_id)
+        parent_context = propagate.extract(dict(websocket.headers))
+    except Exception:
+        parent_context = otel_context_api.get_current()
+    if parent_context is None:
+        parent_context = otel_context_api.get_current()
+
+    span_name = f"WS {websocket.url.path}"
+    with tracer.start_as_current_span(
+        span_name,
+        context=parent_context,
+        kind=trace.SpanKind.SERVER,
+    ) as span:
+        if websocket.client:
+            span.set_attribute("net.peer.ip", websocket.client.host)
+        span.set_attribute("net.protocol.name", "websocket")
+        span.set_attribute("net.host.port", websocket.url.port or 80)
+
+        base_context = ExecutionContext(
+            trace_context=otel_context_api.get_current(),
+            span=span,
+            logger=logger.bind(entrypoint="websocket"),
+            metadata={
+                "entrypoint": "websocket",
+                "path": websocket.url.path,
+            },
+        )
+        token = set_current_context(base_context)
+
+        client_id = await websocket_manager.connect(websocket, client_id)
+        span.set_attribute("client.id", client_id)
+
+        execution_context = replace(
+            base_context,
+            metadata={**base_context.metadata, "client_id": client_id},
+        )
+        websocket_manager.set_context(client_id, execution_context)
+
+        reset_current_context(token)
+        token = set_current_context(execution_context)
+
+        try:
+            while True:
+                data = await websocket.receive_text()
+                try:
+                    message = json.loads(data)
+                    await handle_client_message(client_id, message)
+                except json.JSONDecodeError:
+                    logger.error(f"从客户端 {client_id} 收到无效JSON消息: {data}")
+                    await websocket_manager.send_personal_message(
+                        json.dumps({"error": "Invalid JSON format"}), client_id
+                    )
+        except WebSocketDisconnect:
+            websocket_manager.disconnect(client_id)
+        except Exception as e:
+            logger.error(f"WebSocket连接处理异常: {e}")
+            websocket_manager.disconnect(client_id)
+        finally:
+            reset_current_context(token)
+            websocket_manager.clear_context(client_id)
+
 
 async def handle_client_message(client_id: str, message: Dict[str, Any]):
     """处理来自客户端的消息"""
@@ -42,8 +100,7 @@ async def handle_client_message(client_id: str, message: Dict[str, Any]):
         task_id_str = message.get("task_id")
         if not task_id_str:
             await websocket_manager.send_personal_message(
-                json.dumps({"error": "Missing task_id in subscribe message"}),
-                client_id
+                json.dumps({"error": "Missing task_id in subscribe message"}), client_id
             )
             return
 
@@ -51,11 +108,12 @@ async def handle_client_message(client_id: str, message: Dict[str, Any]):
             task_id = uuid.UUID(task_id_str)
             success = await websocket_manager.subscribe_to_task(client_id, task_id)
             response = {"success": success}
-            await websocket_manager.send_personal_message(json.dumps(response), client_id)
+            await websocket_manager.send_personal_message(
+                json.dumps(response), client_id
+            )
         except ValueError:
             await websocket_manager.send_personal_message(
-                json.dumps({"error": "Invalid task_id format"}),
-                client_id
+                json.dumps({"error": "Invalid task_id format"}), client_id
             )
 
     elif message_type == "unsubscribe":
@@ -64,7 +122,7 @@ async def handle_client_message(client_id: str, message: Dict[str, Any]):
         if not task_id_str:
             await websocket_manager.send_personal_message(
                 json.dumps({"error": "Missing task_id in unsubscribe message"}),
-                client_id
+                client_id,
             )
             return
 
@@ -72,15 +130,15 @@ async def handle_client_message(client_id: str, message: Dict[str, Any]):
             task_id = uuid.UUID(task_id_str)
             success = await websocket_manager.unsubscribe_from_task(client_id, task_id)
             response = {"success": success}
-            await websocket_manager.send_personal_message(json.dumps(response), client_id)
+            await websocket_manager.send_personal_message(
+                json.dumps(response), client_id
+            )
         except ValueError:
             await websocket_manager.send_personal_message(
-                json.dumps({"error": "Invalid task_id format"}),
-                client_id
+                json.dumps({"error": "Invalid task_id format"}), client_id
             )
 
     else:
         await websocket_manager.send_personal_message(
-            json.dumps({"error": f"Unknown message type: {message_type}"}),
-            client_id
+            json.dumps({"error": f"Unknown message type: {message_type}"}), client_id
         )

service_forge/api/deprecated_websocket_manager.py

@@ -1,22 +1,29 @@
 from __future__ import annotations
+
 import asyncio
-import uuid
 import json
-from typing import Dict, List, Set, Any
+import uuid
+from typing import Any
+
 from fastapi import WebSocket, WebSocketDisconnect
 from loguru import logger
+
+from ..execution_context import ExecutionContext
 from .task_manager import TaskManager
 
+
 class WebSocketManager:
     def __init__(self):
         # 存储活动连接: {client_id: websocket}
-        self.active_connections: Dict[str, WebSocket] = {}
+        self.active_connections: dict[str, WebSocket] = {}
+        # 存储链接的 ExecutionContext
+        self.connection_contexts: dict[str, ExecutionContext] = {}
         # 存储任务与客户端的映射: {task_id: client_id}
-        self.task_client_mapping: Dict[uuid.UUID, str] = {}
+        self.task_client_mapping: dict[uuid.UUID, str] = {}
         # 存储客户端订阅的任务: {client_id: set(task_id)}
-        self.client_task_subscriptions: Dict[str, Set[uuid.UUID]] = {}
+        self.client_task_subscriptions: dict[str, set[uuid.UUID]] = {}
         # 存储客户端历史记录，用于重连时恢复订阅: {client_id: last_active_time}
-        self.client_history: Dict[str, float] = {}
+        self.client_history: dict[str, float] = {}
         # 设置客户端记录过期时间（秒），默认0.5小时
         self.client_history_expiry = 0.5 * 60 * 60
         # 初始化任务管理器
@@ -51,11 +58,14 @@ class WebSocketManager:
             "type": "connection established",
             "client_id": client_id,
             "timestamp": str(asyncio.get_event_loop().time()),
-            "restored_subscriptions": []
+            "restored_subscriptions": [],
         }
 
         # 如果有历史订阅，恢复它们
-        if client_id in self.client_task_subscriptions and self.client_task_subscriptions[client_id]:
+        if (
+            client_id in self.client_task_subscriptions
+            and self.client_task_subscriptions[client_id]
+        ):
             restored_tasks = []
             for task_id in self.client_task_subscriptions[client_id]:
                 restored_tasks.append(str(task_id))
@@ -74,6 +84,7 @@ class WebSocketManager:
             # 更新客户端的最后活动时间
             self.client_history[client_id] = asyncio.get_event_loop().time()
             logger.info(f"客户端 {client_id} 已断开WebSocket连接，保留订阅信息")
+        self.clear_context(client_id)
 
     async def subscribe_to_task(self, client_id: str, task_id: uuid.UUID) -> bool:
         """客户端订阅任务"""
@@ -98,7 +109,22 @@ class WebSocketManager:
 
         return False
 
-    def create_task_with_client(self, task_id: uuid.UUID, client_id: str, workflow_name: str = "Unknown", steps: int = 1) -> bool:
+    def set_context(self, client_id: str, context: ExecutionContext) -> None:
+        self.connection_contexts[client_id] = context
+
+    def get_context(self, client_id: str) -> ExecutionContext | None:
+        return self.connection_contexts.get(client_id)
+
+    def clear_context(self, client_id: str) -> None:
+        self.connection_contexts.pop(client_id, None)
+
+    def create_task_with_client(
+        self,
+        task_id: uuid.UUID,
+        client_id: str,
+        workflow_name: str = "Unknown",
+        steps: int = 1,
+    ) -> bool:
         """创建任务与客户端的映射，并添加到任务管理器"""
         # 建立任务与客户端的映射
         self.task_client_mapping[task_id] = client_id
@@ -107,7 +133,7 @@ class WebSocketManager:
         if client_id not in self.client_task_subscriptions:
             self.client_task_subscriptions[client_id] = set()
         self.client_task_subscriptions[client_id].add(task_id)
-        
+
         # 添加任务到任务管理器
         self.task_manager.add_task(task_id, client_id, workflow_name, steps)
 
@@ -132,11 +158,11 @@ class WebSocketManager:
             return  # 没有关联的客户端
 
         client_id = self.task_client_mapping[task_id]
-        
+
         # 确保task_id是字符串，避免JSON序列化问题
         if "task_id" in message and isinstance(message["task_id"], uuid.UUID):
             message["task_id"] = str(message["task_id"])
-            
+
         # 递归处理嵌套字典中的UUID
         def convert_uuids(obj):
             if isinstance(obj, dict):
@@ -147,18 +173,21 @@ class WebSocketManager:
                 return str(obj)
             else:
                 return obj
-                
+
         message = convert_uuids(message)
         message_str = json.dumps(message)
         await self.send_personal_message(message_str, client_id)
 
-    async def send_task_status(self, task_id: uuid.UUID, status: str, node: str = None, progress: float = None, error: str = None):
+    async def send_task_status(
+        self,
+        task_id: uuid.UUID,
+        status: str,
+        node: str = None,
+        progress: float = None,
+        error: str = None,
+    ):
         """发送任务状态更新"""
-        message = {
-            "task_id": str(task_id),
-            "type": "status",
-            "status": status
-        }
+        message = {"task_id": str(task_id), "type": "status", "status": status}
 
         if node is not None:
             message["node"] = node
@@ -175,7 +204,7 @@ class WebSocketManager:
         """发送任务开始执行消息"""
         # 获取客户端ID
         client_id = self.task_client_mapping.get(task_id)
-        
+
         # 更新任务状态为运行中
         self.task_manager.start_task(task_id)
 
@@ -186,19 +215,16 @@ class WebSocketManager:
 
         # 获取全局任务队列信息
         global_queue_info = self.task_manager.get_global_queue_info()
-        
+
         # 获取当前任务在队列中的位置
         queue_position = self.task_manager.get_queue_position(task_id)
 
         message = {
             "task_id": str(task_id),
             "type": "execution start",
-            "client_tasks": {
-                "total": len(client_tasks),
-                "tasks": client_tasks
-            },
+            "client_tasks": {"total": len(client_tasks), "tasks": client_tasks},
             "global_queue": global_queue_info,
-            "queue_position": queue_position
+            "queue_position": queue_position,
         }
         await self.send_to_task_client(task_id, message)
 
@@ -214,7 +240,7 @@ class WebSocketManager:
 
         # 获取全局任务队列信息
         global_queue_info = self.task_manager.get_global_queue_info()
-        
+
         # 获取当前任务在队列中的位置
         queue_position = self.task_manager.get_queue_position(task_id)
 
@@ -222,12 +248,9 @@ class WebSocketManager:
             "task_id": str(task_id),
             "type": "executing",
             "node": node,
-            "client_tasks": {
-                "total": len(client_tasks),
-                "tasks": client_tasks
-            },
+            "client_tasks": {"total": len(client_tasks), "tasks": client_tasks},
             "global_queue": global_queue_info,
-            "queue_position": queue_position
+            "queue_position": queue_position,
         }
         await self.send_to_task_client(task_id, message)
 
@@ -243,7 +266,7 @@ class WebSocketManager:
 
         # 获取全局任务队列信息
         global_queue_info = self.task_manager.get_global_queue_info()
-        
+
         # 获取当前任务在队列中的位置
         queue_position = self.task_manager.get_queue_position(task_id)
 
@@ -252,12 +275,9 @@ class WebSocketManager:
             "type": "progress",
             "node": node,
             "progress": progress,
-            "client_tasks": {
-                "total": len(client_tasks),
-                "tasks": client_tasks
-            },
+            "client_tasks": {"total": len(client_tasks), "tasks": client_tasks},
             "global_queue": global_queue_info,
-            "queue_position": queue_position
+            "queue_position": queue_position,
         }
         await self.send_to_task_client(task_id, message)
 
@@ -273,7 +293,7 @@ class WebSocketManager:
 
         # 获取全局任务队列信息
         global_queue_info = self.task_manager.get_global_queue_info()
-        
+
         # 获取当前任务在队列中的位置
         queue_position = self.task_manager.get_queue_position(task_id)
 
@@ -281,17 +301,15 @@ class WebSocketManager:
             "task_id": str(task_id),
             "type": "executed",
             "node": node,
-            "client_tasks": {
-                "total": len(client_tasks),
-                "tasks": client_tasks
-            },
+            "client_tasks": {"total": len(client_tasks), "tasks": client_tasks},
             "global_queue": global_queue_info,
-            "queue_position": queue_position
+            "queue_position": queue_position,
         }
 
         if result is not None:
             # 检查是否为协程对象
             import asyncio
+
             if asyncio.iscoroutine(result):
                 message["result"] = "<coroutine object>"
             else:
@@ -307,7 +325,7 @@ class WebSocketManager:
         """发送执行错误消息"""
         # 获取客户端ID
         client_id = self.task_client_mapping.get(task_id)
-        
+
         # 更新任务状态为失败
         self.task_manager.fail_task(task_id, error)
 
@@ -318,7 +336,7 @@ class WebSocketManager:
 
         # 获取全局任务队列信息
         global_queue_info = self.task_manager.get_global_queue_info()
-        
+
         # 获取当前任务在队列中的位置
         queue_position = self.task_manager.get_queue_position(task_id)
 
@@ -327,12 +345,9 @@ class WebSocketManager:
             "type": "execution error",
             "node": node,
             "error": error,
-            "client_tasks": {
-                "total": len(client_tasks),
-                "tasks": client_tasks
-            },
+            "client_tasks": {"total": len(client_tasks), "tasks": client_tasks},
             "global_queue": global_queue_info,
-            "queue_position": queue_position
+            "queue_position": queue_position,
         }
         await self.send_to_task_client(task_id, message)
 
@@ -360,8 +375,10 @@ class WebSocketManager:
                 # 查找过期的客户端记录
                 for client_id, last_active_time in self.client_history.items():
                     # 如果客户端不在活动连接中且超过过期时间，则标记为过期
-                    if (client_id not in self.active_connections and 
-                        current_time - last_active_time > self.client_history_expiry):
+                    if (
+                        client_id not in self.active_connections
+                        and current_time - last_active_time > self.client_history_expiry
+                    ):
                         expired_clients.append(client_id)
 
                 # 清理过期客户端的订阅记录