Update remote_uploader.py

remote_uploader.py

@@ -7,7 +7,6 @@ import json
 import base64
 import logging
 import sys
-import time
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -21,22 +20,25 @@ class WebSocketFileUploader:
         
     async def connect_and_upload(self, upload_dir):
         """连接到 WebSocket 服务器并上传文件"""
-        
-        # --- 修复点: 统一将 URL 转换为小写，以处理类似 "HTTPS" 这样的异常情况 ---
-        normalized_server_url = self.server_url.lower()
-        
-        # 使用转换后的 URL 进行协议替换
-        ws_url = normalized_server_url.replace('https://', 'wss://').replace('http://', 'ws://').rstrip('/')
+        # 将 HTTP URL 转换为 WebSocket URL
+        ws_url = self.server_url.replace('http://', 'ws://').replace('https://', 'wss://').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, open_timeout=10, ping_interval=20, ping_timeout=60) as websocket:
+            async with websockets.connect(ws_url, ping_interval=20, ping_timeout=60) as websocket:
                 logger.info("WebSocket 连接成功")
                 
-                # 等待服务器的认证成功响应
+                # 发送认证信息
+                auth_msg = {
+                    "type": "auth",
+                    "api_key": self.api_key,
+                    "space_id": self.space_id
+                }
+                await websocket.send(json.dumps(auth_msg))
+                
+                # 等待认证响应
                 response = await websocket.recv()
                 auth_response = json.loads(response)
                 
@@ -45,19 +47,21 @@ class WebSocketFileUploader:
                     return await self._upload_files(websocket, upload_dir)
                 else:
                     logger.error(f"认证失败: {auth_response.get('message', '未知错误')}")
-                    logger.info("回退到 HTTP API 上传...")
-                    return self._fallback_http_upload(upload_dir)
+                    return False
                     
         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:
@@ -83,6 +87,7 @@ class WebSocketFileUploader:
                     failed_count += 1
                     logger.error(f"❌ 文件 '{os.path.basename(file_path)}' 上传失败")
                 
+                # 稍微延迟一下
                 await asyncio.sleep(0.5)
                 
             except Exception as e:
@@ -93,14 +98,17 @@ 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,
@@ -108,8 +116,10 @@ class WebSocketFileUploader:
                 "space_id": self.space_id
             }
             
+            # 发送文件
             await websocket.send(json.dumps(upload_msg))
             
+            # 等待响应
             response = await websocket.recv()
             result = json.loads(response)
             
@@ -124,17 +134,25 @@ 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:
@@ -157,14 +175,18 @@ class WebSocketFileUploader:
                 logger.info(f"正在上传文件: {filename}")
                 
                 with open(file_path, 'rb') as f:
-                    files = {'file': (filename, f)}
-                    response = requests.post(upload_url, headers=headers, data=data, files=files, timeout=60)
+                    files = {
+                        'file': (filename, f)
+                    }
+                    
+                    # 发送 POST 请求
+                    response = requests.post(upload_url, headers=headers, data=data, files=files, timeout=30)
                     
                     if response.status_code == 200:
                         logger.info(f"✅ 文件 '{filename}' 上传成功!")
                         success_count += 1
                     else:
-                        logger.error(f"❌ 上传失败: HTTP {response.status_code} - {response.text}")
+                        logger.error(f"❌ 上传失败: HTTP {response.status_code}")
                         failed_count += 1
                 
                 time.sleep(0.5)
@@ -177,6 +199,7 @@ 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}")
@@ -196,9 +219,10 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="WebSocket 文件上传器 - 扫描并上传指定文件夹中的所有文件")
     parser.add_argument("api_key", help="您的 API 密钥")
     parser.add_argument("space_id", help="Space ID")
-    parser.add_argument("--server", default="http://127.0.0.1:5001", help="服务器的 URL 地址 (例如: http://127.0.0.1:5001)")
+    parser.add_argument("--server", default="https://gkbtyo-rqvays-5001.preview.cloudstudio.work", help="服务器的 URL 地址")
     parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
     
     args = parser.parse_args()
     
-    asyncio.run(upload_directory_websocket(args.upload_dir, args.server, args.api_key, args.space_id))
+    # 运行异步上传
+    asyncio.run(upload_directory_websocket(args.upload_dir, args.server, args.api_key, args.space_id))