Update remote_client.py

remote_client.py

@@ -1,9 +1,16 @@
+#!/usr/bin/env python3
+"""
+WebSocket远程客户端 - 在inferless环境中运行
+用于连接到主服务器并执行远程命令和文件传输
+"""
+
 import asyncio
-import socketio
+import websockets
 import subprocess
 import json
 import logging
 import sys
+import urllib.parse
 import base64
 import os
 import threading
@@ -15,217 +22,237 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 logger = logging.getLogger(__name__)
 
 class RemoteClient:
-    def __init__(self, space_id, user_token, server_url):
-        self.space_id = space_id
-        self.user_token = user_token
+    def __init__(self, server_url, user_token, space_id):
         self.server_url = server_url
-        self.sio = socketio.AsyncClient(ping_timeout=60, ping_interval=25)
-        self.connected = False
+        self.user_token = user_token
+        self.space_id = space_id
+        self.websocket = None
+        self.is_connected = False
         
-        # 绑定事件
-        self.sio.on('connect', self.on_connect)
-        self.sio.on('disconnect', self.on_disconnect)
-        self.sio.on('registered', self.on_registered)
-        self.sio.on('execute_command', self.on_execute_command)
-        self.sio.on('error', self.on_error)
-        self.sio.on('upload_success', self.on_upload_success)
-
-    async def on_connect(self):
-        logger.info("Connected to server")
-        # 注册客户端
-        await self.sio.emit('register_client', {
-            'space_id': self.space_id,
-            'token': self.user_token
-        })
-
-    async def on_disconnect(self):
-        logger.info("Disconnected from server")
-        self.connected = False
-
-    async def on_registered(self, data):
-        logger.info(f"Successfully registered: {data}")
-        self.connected = True
+    async def connect(self):
+        """连接到WebSocket服务器"""
+        # 构建WebSocket URL
+        ws_url = self.server_url.replace('http', 'ws').replace('https', 'wss')
+        uri = f"{ws_url}/ws/client/{self.space_id}?token={urllib.parse.quote(self.user_token)}"
+        
+        logger.info(f"正在连接到服务器: {uri}")
+        
+        try:
+            async with websockets.connect(uri, ping_interval=20, ping_timeout=60) as websocket:
+                self.websocket = websocket
+                self.is_connected = True
+                logger.info("✅ 已连接到WebSocket服务器")
+                
+                # 发送客户端注册信息
+                await self.register_client()
+                
+                # 启动消息监听循环
+                await self.message_loop()
+                
+        except Exception as e:
+            logger.error(f"❌ 连接失败: {e}")
+            self.is_connected = False
 
-    async def on_error(self, data):
-        logger.error(f"Server error: {data}")
+    async def register_client(self):
+        """注册客户端到服务器"""
+        registration = {
+            "type": "register",
+            "space_id": self.space_id,
+            "client_info": {
+                "environment": "inferless",
+                "python_version": sys.version,
+                "working_directory": os.getcwd()
+            }
+        }
+        await self.websocket.send(json.dumps(registration))
+        logger.info("📝 已发送客户端注册信息")
 
-    async def on_upload_success(self, data):
-        logger.info(f"File upload successful: {data}")
+    async def message_loop(self):
+        """主消息处理循环"""
+        while self.is_connected:
+            try:
+                message = await self.websocket.recv()
+                data = json.loads(message)
+                await self.handle_message(data)
+                
+            except websockets.exceptions.ConnectionClosed:
+                logger.error("🔌 WebSocket连接已关闭")
+                self.is_connected = False
+                break
+            except Exception as e:
+                logger.error(f"❌ 处理消息时出错: {e}")
 
-    async def on_execute_command(self, data):
-        """处理执行命令请求"""
-        task_id = data.get('task_id')
-        command = data.get('command')
-        token = data.get('token')
-        
-        logger.info(f"Received command for task {task_id}: {command}")
+    async def handle_message(self, data):
+        """处理收到的消息"""
+        message_type = data.get("type")
         
-        # 发送开始执行的状态
-        await self.sio.emit('command_result', {
-            'task_id': task_id,
-            'status': 'running',
-            'output': f'开始执行命令: {command}'
-        })
-        
-        # 在后台执行命令
-        asyncio.create_task(self.execute_command_async(task_id, command, token))
+        if message_type == "command":
+            await self.execute_command(data)
+        elif message_type == "upload_files":
+            await self.upload_files(data)
+        elif message_type == "ping":
+            await self.send_pong()
+        else:
+            logger.warning(f"⚠️ 未知消息类型: {message_type}")
 
-    async def execute_command_async(self, task_id, command, token):
-        """异步执行命令"""
+    async def execute_command(self, data):
+        """执行远程命令"""
+        command = data.get("command", "")
+        logger.info(f"🚀 执行命令: {command}")
+        
         try:
-            # 创建 output 目录
-            os.makedirs('./output', exist_ok=True)
-            
-            logger.info(f"Executing: {command}")
-            
             # 执行命令
-            process = await asyncio.create_subprocess_shell(
+            process = subprocess.run(
                 command,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-                text=True
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=300  # 5分钟超时
             )
             
-            # 实时读取输出
-            output_buffer = ""
-            while True:
-                line = await process.stdout.readline()
-                if not line:
-                    break
-                line_str = line.decode() if isinstance(line, bytes) else line
-                output_buffer += line_str
-                logger.info(f"Command output: {line_str.strip()}")
-                
-                # 发送实时输出
-                await self.sio.emit('command_result', {
-                    'task_id': task_id,
-                    'status': 'running',
-                    'output': line_str
-                })
+            # 发送执行结果
+            result = {
+                "type": "command_result",
+                "command": command,
+                "returncode": process.returncode,
+                "stdout": process.stdout,
+                "stderr": process.stderr,
+                "success": process.returncode == 0
+            }
             
-            # 等待进程完成
-            await process.wait()
+            await self.websocket.send(json.dumps(result))
             
             if process.returncode == 0:
-                logger.info("Command executed successfully")
-                
-                # 扫描 output 目录并上传文件
-                await self.upload_output_files(token)
-                
-                await self.sio.emit('command_result', {
-                    'task_id': task_id,
-                    'status': 'completed',
-                    'output': '命令执行完成，文件已上传到网盘'
-                })
+                logger.info("✅ 命令执行成功")
             else:
-                logger.error(f"Command failed with return code: {process.returncode}")
-                await self.sio.emit('command_result', {
-                    'task_id': task_id,
-                    'status': 'failed',
-                    'output': f'命令执行失败，退出码: {process.returncode}'
-                })
+                logger.error(f"❌ 命令执行失败，返回码: {process.returncode}")
                 
+        except subprocess.TimeoutExpired:
+            error_result = {
+                "type": "command_result",
+                "command": command,
+                "error": "命令执行超时",
+                "success": False
+            }
+            await self.websocket.send(json.dumps(error_result))
+            logger.error("⏰ 命令执行超时")
+            
         except Exception as e:
-            logger.error(f"Error executing command: {e}")
-            await self.sio.emit('command_result', {
-                'task_id': task_id,
-                'status': 'failed',
-                'output': f'执行命令时出错: {str(e)}'
-            })
+            error_result = {
+                "type": "command_result",
+                "command": command,
+                "error": str(e),
+                "success": False
+            }
+            await self.websocket.send(json.dumps(error_result))
+            logger.error(f"❌ 命令执行异常: {e}")
 
-    async def upload_output_files(self, token):
-        """上传 output 目录中的文件"""
-        output_dir = Path('./output')
-        if not output_dir.exists():
-            logger.warning("Output directory does not exist")
-            return
+    async def upload_files(self, data):
+        """上传指定目录中的所有文件"""
+        upload_dir = data.get("directory", "output")
         
-        uploaded_count = 0
-        failed_count = 0
-        
-        for file_path in output_dir.iterdir():
-            if file_path.is_file():
-                try:
-                    logger.info(f"Uploading file: {file_path.name}")
-                    
-                    # 读取文件内容
-                    with open(file_path, 'rb') as f:
-                        file_content = f.read()
-                    
-                    # 编码为 base64
-                    file_b64 = base64.b64encode(file_content).decode('utf-8')
-                    
-                    # 发送文件上传请求
-                    await self.sio.emit('file_upload', {
-                        'filename': file_path.name,
-                        'content': file_b64,
-                        'token': token
-                    })
-                    
-                    uploaded_count += 1
-                    logger.info(f"Successfully uploaded: {file_path.name}")
-                    
-                    # 删除已上传的文件
-                    file_path.unlink()
-                    
-                    # 避免发送过大的数据包，稍作延迟
-                    await asyncio.sleep(0.5)
-                    
-                except Exception as e:
-                    logger.error(f"Failed to upload {file_path.name}: {e}")
-                    failed_count += 1
+        logger.info(f"📁 开始上传目录: {upload_dir}")
         
-        logger.info(f"Upload completed. Success: {uploaded_count}, Failed: {failed_count}")
+        if not os.path.exists(upload_dir):
+            error_msg = f"目录不存在: {upload_dir}"
+            logger.error(error_msg)
+            await self.send_upload_result(False, error_msg)
+            return
 
-    async def connect(self):
-        """连接到服务器"""
         try:
-            # 构建 WebSocket URL
-            ws_url = self.server_url.replace('https://', 'wss://').replace('http://', 'ws://')
-            if not ws_url.endswith('/'):
-                ws_url += '/'
-            ws_url += 'socket.io/'
+            # 扫描目录中的所有文件
+            files_to_upload = []
+            for root, dirs, files in os.walk(upload_dir):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    if os.path.isfile(file_path):
+                        files_to_upload.append(file_path)
             
-            logger.info(f"Connecting to {ws_url}")
-            await self.sio.connect(ws_url, transports=['websocket'])
+            if not files_to_upload:
+                msg = f"目录 {upload_dir} 中没有找到文件"
+                logger.info(msg)
+                await self.send_upload_result(True, msg)
+                return
             
-            # 保持连接
-            while self.connected:
-                await asyncio.sleep(1)
+            logger.info(f"📦 找到 {len(files_to_upload)} 个文件，开始上传...")
+            
+            success_count = 0
+            failed_count = 0
+            
+            for file_path in files_to_upload:
+                try:
+                    await self.upload_single_file(file_path)
+                    success_count += 1
+                    logger.info(f"✅ 上传成功: {os.path.basename(file_path)}")
+                except Exception as e:
+                    failed_count += 1
+                    logger.error(f"❌ 上传失败: {file_path} - {e}")
                 
+                # 稍作延迟避免服务器压力
+                await asyncio.sleep(0.1)
+            
+            result_msg = f"上传完成! 成功: {success_count}, 失败: {failed_count}"
+            logger.info(f"📊 {result_msg}")
+            await self.send_upload_result(True, result_msg)
+            
         except Exception as e:
-            logger.error(f"Connection failed: {e}")
-            raise
+            error_msg = f"上传过程中出错: {e}"
+            logger.error(error_msg)
+            await self.send_upload_result(False, error_msg)
+
+    async def upload_single_file(self, file_path):
+        """上传单个文件"""
+        filename = os.path.basename(file_path)
+        
+        with open(file_path, 'rb') as f:
+            file_content = f.read()
+            file_b64 = base64.b64encode(file_content).decode('utf-8')
+        
+        file_message = {
+            "type": "file_upload",
+            "filename": filename,
+            "content": file_b64,
+            "file_size": len(file_content)
+        }
+        
+        await self.websocket.send(json.dumps(file_message))
+
+    async def send_upload_result(self, success, message):
+        """发送上传结果"""
+        result = {
+            "type": "upload_result",
+            "success": success,
+            "message": message
+        }
+        await self.websocket.send(json.dumps(result))
 
-    async def disconnect(self):
-        """断开连接"""
-        if self.sio.connected:
-            await self.sio.disconnect()
+    async def send_pong(self):
+        """响应ping消息"""
+        pong = {"type": "pong"}
+        await self.websocket.send(json.dumps(pong))
 
 async def main():
-    if len(sys.argv) != 4:
-        print("Usage: python3 remote_client.py <space_id> <user_token> <server_url>")
-        print("Example: python3 remote_client.py 'abc123' 'token456' 'https://gkbtyo-rqvays-5001.preview.cloudstudio.work'")
+    if len(sys.argv) < 4:
+        print("使用方法: python remote_client.py <server_url> <user_token> <space_id>")
+        print("示例: python remote_client.py https://gkbtyo-rqvays-5001.preview.cloudstudio.work abc123 space456")
         sys.exit(1)
     
-    space_id = sys.argv[1]
+    server_url = sys.argv[1]
     user_token = sys.argv[2]
-    server_url = sys.argv[3]
+    space_id = sys.argv[3]
     
-    logger.info(f"Starting remote client for Space: {space_id}")
-    logger.info(f"Server: {server_url}")
+    client = RemoteClient(server_url, user_token, space_id)
     
-    client = RemoteClient(space_id, user_token, server_url)
-    
-    try:
-        await client.connect()
-    except KeyboardInterrupt:
-        logger.info("Received interrupt signal")
-    except Exception as e:
-        logger.error(f"Client error: {e}")
-    finally:
-        await client.disconnect()
-        logger.info("Client stopped")
+    while True:
+        try:
+            await client.connect()
+        except KeyboardInterrupt:
+            logger.info("👋 收到中断信号，正在退出...")
+            break
+        except Exception as e:
+            logger.error(f"❌ 连接异常: {e}")
+            logger.info("🔄 5秒后重新尝试连接...")
+            await asyncio.sleep(5)
 
 if __name__ == "__main__":
     asyncio.run(main())