Update remote_uploader.py

remote_uploader.py

@@ -7,8 +7,7 @@ import json
 import base64
 import logging
 import sys
-import time  # 修复点 1: 导入 time 模块以使用 time.sleep()
-from urllib.parse import urlparse # 修复点 2: 导入 urlparse 以便解析 URL
+import time
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -23,36 +22,19 @@ class WebSocketFileUploader:
     async def connect_and_upload(self, upload_dir):
         """连接到 WebSocket 服务器并上传文件"""
         
-        # --- 修复点 3: 智能地构建 WebSocket URL ---
-        # 目标: 将主应用传来的 HTTP URL (如 http://...:5001) 转换为正确的 WebSocket URL (ws://...:8765)
-        try:
-            parsed_http_url = urlparse(self.server_url)
-            # 根据 HTTP 协议确定 WebSocket 协议 (http -> ws, https -> wss)
-            ws_scheme = 'wss' if parsed_http_url.scheme == 'https' else 'ws'
-            
-            # 使用从 HTTP URL 获取的主机名，但硬编码正确的 WebSocket 端口 8765
-            ws_netloc = f"{parsed_http_url.hostname}:8765"
-            
-            # 重新构建完整的 WebSocket URL
-            ws_url = f"{ws_scheme}://{ws_netloc}/ws/upload/{self.space_id}?token={self.api_key}"
-        except Exception as e:
-            logger.error(f"解析服务器 URL '{self.server_url}' 时出错: {e}")
-            logger.info("回退到 HTTP API 上传...")
-            return self._fallback_http_upload(upload_dir)
-        # --- 修复结束 ---
+        # 修复点: 简化 WebSocket URL 的构建。
+        # 服务器在同一端口处理 HTTP 和 WS，只需替换协议。
+        ws_url = self.server_url.replace('https://', 'wss://').replace('http://', 'ws://').rstrip('/')
+        ws_url += f'/ws/upload/{self.space_id}?token={self.api_key}'
 
         logger.info(f"正在连接到 WebSocket: {ws_url}")
         
         try:
-            async with websockets.connect(ws_url, ping_interval=20, ping_timeout=60) as websocket:
+            # 增加连接超时以更快地失败
+            async with websockets.connect(ws_url, open_timeout=10, ping_interval=20, ping_timeout=60) as websocket:
                 logger.info("WebSocket 连接成功")
                 
-                # 注意：原始代码中的认证流程是先连接再发认证消息，这很好。
-                # 我们将遵循这个流程。客户端不需要发送 auth 消息，
-                # 因为服务器端的 handle_websocket_upload 函数会直接从连接路径和参数中验证。
-                # 我们只需要等待服务器的 "auth_success" 响应。
-                
-                # 等待认证响应
+                # 等待服务器的认证成功响应
                 response = await websocket.recv()
                 auth_response = json.loads(response)
                 
@@ -61,23 +43,19 @@ class WebSocketFileUploader:
                     return await self._upload_files(websocket, upload_dir)
                 else:
                     logger.error(f"认证失败: {auth_response.get('message', '未知错误')}")
-                    # 认证失败也回退到 HTTP
                     logger.info("回退到 HTTP API 上传...")
                     return self._fallback_http_upload(upload_dir)
                     
         except Exception as e:
             logger.error(f"WebSocket 连接失败: {e}")
-            # 如果 WebSocket 失败，回退到 HTTP API
             logger.info("回退到 HTTP API 上传...")
             return self._fallback_http_upload(upload_dir)
     
     async def _upload_files(self, websocket, upload_dir):
-        """通过 WebSocket 上传目录中的所有文件"""
         if not os.path.exists(upload_dir):
             logger.error(f"目录不存在: {upload_dir}")
             return False
         
-        # 获取所有文件
         all_files = []
         for root, dirs, files in os.walk(upload_dir):
             for file in files:
@@ -103,7 +81,6 @@ class WebSocketFileUploader:
                     failed_count += 1
                     logger.error(f"❌ 文件 '{os.path.basename(file_path)}' 上传失败")
                 
-                # 稍微延迟一下
                 await asyncio.sleep(0.5)
                 
             except Exception as e:
@@ -114,17 +91,14 @@ class WebSocketFileUploader:
         return success_count > 0
     
     async def _upload_single_file(self, websocket, file_path):
-        """上传单个文件通过 WebSocket"""
         try:
             filename = os.path.basename(file_path)
             logger.info(f"正在上传文件: {filename}")
             
-            # 读取文件内容并编码为 base64
             with open(file_path, 'rb') as f:
                 file_content = f.read()
                 file_b64 = base64.b64encode(file_content).decode('utf-8')
             
-            # 构建上传消息
             upload_msg = {
                 "type": "file_upload",
                 "filename": filename,
@@ -132,10 +106,8 @@ class WebSocketFileUploader:
                 "space_id": self.space_id
             }
             
-            # 发送文件
             await websocket.send(json.dumps(upload_msg))
             
-            # 等待响应
             response = await websocket.recv()
             result = json.loads(response)
             
@@ -150,25 +122,17 @@ class WebSocketFileUploader:
             return False
     
     def _fallback_http_upload(self, upload_dir):
-        """HTTP API 回退方案"""
         logger.info("使用 HTTP API 回退方案上传文件...")
         
         if not os.path.exists(upload_dir):
             logger.error(f"目录不存在: {upload_dir}")
             return False
         
-        # 构建完整的 API 端点 URL
         upload_url = f"{self.server_url.rstrip('/')}/api/remote_upload"
         
-        # 准备请求头
-        headers = {
-            'X-API-Key': self.api_key
-        }
-        data = {
-            'space_id': self.space_id
-        }
+        headers = {'X-API-Key': self.api_key}
+        data = {'space_id': self.space_id}
         
-        # 获取所有文件
         all_files = []
         for root, dirs, files in os.walk(upload_dir):
             for file in files:
@@ -191,11 +155,7 @@ class WebSocketFileUploader:
                 logger.info(f"正在上传文件: {filename}")
                 
                 with open(file_path, 'rb') as f:
-                    files = {
-                        'file': (filename, f)
-                    }
-                    
-                    # 发送 POST 请求
+                    files = {'file': (filename, f)}
                     response = requests.post(upload_url, headers=headers, data=data, files=files, timeout=60)
                     
                     if response.status_code == 200:
@@ -205,7 +165,6 @@ class WebSocketFileUploader:
                         logger.error(f"❌ 上传失败: HTTP {response.status_code} - {response.text}")
                         failed_count += 1
                 
-                # 使用 time.sleep()
                 time.sleep(0.5)
                         
             except Exception as e:
@@ -216,7 +175,6 @@ class WebSocketFileUploader:
         return success_count > 0
 
 async def upload_directory_websocket(upload_dir, server_url, api_key, space_id):
-    """使用 WebSocket 方式上传目录中的所有文件"""
     logger.info(f"🔍 开始扫描目录: {upload_dir}")
     logger.info(f"📡 服务器地址: {server_url}")
     logger.info(f"🔑 Space ID: {space_id}")
@@ -241,5 +199,4 @@ if __name__ == "__main__":
     
     args = parser.parse_args()
     
-    # 运行异步上传
     asyncio.run(upload_directory_websocket(args.upload_dir, args.server, args.api_key, args.space_id))