Upload proxy_handler.py

proxy_handler.py

@@ -1,15 +1,13 @@
 """
 Proxy handler for Z.AI API requests
 """
-import json, logging, re, time, uuid
+import json, logging, re, time, uuid, base64, hashlib, hmac, urllib.parse
 from typing import AsyncGenerator, Dict, Any, Tuple, List
+from datetime import datetime, timezone, timedelta
+
 import httpx
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
-import hashlib
-import hmac
-import urllib.parse
-from datetime import datetime, timezone
 
 from config import settings
 from cookie_manager import cookie_manager
@@ -32,22 +30,26 @@ class ProxyHandler:
         if not self.client.is_closed:
             await self.client.aclose()
             
-    # --- 新增方法：获取毫秒级时间戳 ---
     def _get_timestamp_millis(self) -> int:
         return int(time.time() * 1000)
 
-    # --- 新增方法：解析JWT以获取用户ID ---
     def _parse_jwt_token(self, token: str) -> Dict[str, str]:
+        """A simple JWT payload decoder to get user ID."""
         try:
             parts = token.split('.')
             if len(parts) != 3:
                 return {"userId": ""}
-            payload = json.loads(base64.urlsafe_b64decode(parts[1] + '==').decode('utf-8'))
+            # Add padding if necessary
+            payload_b64 = parts[1]
+            payload_b64 += '=' * (-len(payload_b64) % 4)
+            payload_json = base64.urlsafe_b64decode(payload_b64).decode('utf-8')
+            payload = json.loads(payload_json)
+            # The JS code checks multiple keys, 'sub' is the most standard one.
             return {"userId": payload.get("sub", "")}
-        except Exception:
+        except Exception as e:
+            logger.warning(f"Failed to parse JWT token to get user ID: {e}")
             return {"userId": ""}
 
-    # --- 新增方法：生成签名 ---
     def _generate_signature(self, token: str, payload_str: str, mt: str) -> Tuple[str, int]:
         timestamp_ms = self._get_timestamp_millis()
         minute_bucket = str(timestamp_ms // 60000)
@@ -63,55 +65,25 @@ class ProxyHandler:
         return signature, timestamp_ms
 
     def _clean_thinking_content(self, text: str) -> str:
-        """
-        Aggressively cleans raw thinking content strings based on observed patterns
-        from the Z.AI API.
-        """
-        if not text:
-            return ""
-
+        if not text: return ""
         cleaned_text = text
-
-        # 1. Remove specific unwanted blocks like tool calls and summaries.
         cleaned_text = re.sub(r'<summary>.*?</summary>', '', cleaned_text, flags=re.DOTALL)
         cleaned_text = re.sub(r'<glm_block.*?</glm_block>', '', cleaned_text, flags=re.DOTALL)
-        
-        # 2. **FIX**: Remove tag-like metadata containing `duration` attribute.
-        # This handles the reported issue: `true" duration="0" ... >`
         cleaned_text = re.sub(r'<[^>]*duration="[^"]*"[^>]*>', '', cleaned_text)
-
-        # 3. Remove specific structural tags, but keep the content between them.
-        cleaned_text = cleaned_text.replace("</thinking>", "")
-        cleaned_text = cleaned_text.replace("<Full>", "")
-        cleaned_text = cleaned_text.replace("</Full>", "")
-        # This regex handles <details>, <details open>, and </details>
+        cleaned_text = cleaned_text.replace("</thinking>", "").replace("<Full>", "").replace("</Full>", "")
         cleaned_text = re.sub(r'</?details[^>]*>', '', cleaned_text)
-
-        # 4. Handle markdown blockquotes, preserving multi-level ones.
         cleaned_text = re.sub(r'^\s*>\s*(?!>)', '', cleaned_text, flags=re.MULTILINE)
-
-        # 5. Remove other known text artifacts.
         cleaned_text = cleaned_text.replace("Thinking…", "")
-        
-        # 6. Final strip to clean up residual whitespace.
         return cleaned_text.strip()
 
     def _clean_answer_content(self, text: str) -> str:
-        """
-        Cleans unwanted tags from answer content.
-        Does NOT strip whitespace to preserve markdown in streams.
-        """
-        if not text:
-            return ""
-        # Remove tool call blocks
+        if not text: return ""
         cleaned_text = re.sub(r'<glm_block.*?</glm_block>', '', text, flags=re.DOTALL)
-        # Remove any residual details/summary blocks that might leak into the answer
         cleaned_text = re.sub(r'<details[^>]*>.*?</details>', '', cleaned_text, flags=re.DOTALL)
         cleaned_text = re.sub(r'<summary>.*?</summary>', '', cleaned_text, flags=re.DOTALL)
         return cleaned_text
 
     def _serialize_msgs(self, msgs) -> list:
-        """Converts message objects to a list of dictionaries."""
         out = []
         for m in msgs:
             if hasattr(m, "dict"): out.append(m.dict())
@@ -120,7 +92,6 @@ class ProxyHandler:
             else: out.append({"role": getattr(m, "role", "user"), "content": getattr(m, "content", str(m))})
         return out
 
-    # --- 重构 _prep_upstream 方法以加入签名逻辑 ---
     async def _prep_upstream(self, req: ChatCompletionRequest) -> Tuple[Dict[str, Any], Dict[str, str], str, str]:
         """Prepares the request body, headers, cookie, and URL for the upstream API."""
         ck = await cookie_manager.get_next_cookie()
@@ -134,22 +105,30 @@ class ProxyHandler:
         
         body = { "stream": True, "model": model, "messages": self._serialize_msgs(req.messages), "background_tasks": {"title_generation": True, "tags_generation": True}, "chat_id": chat_id, "features": {"image_generation": False, "code_interpreter": False, "web_search": False, "auto_web_search": False, "enable_thinking": True,}, "id": request_id, "mcp_servers": ["deep-web-search"], "model_item": {"id": model, "name": "GLM-4.6", "owned_by": "openai"}, "params": {}, "tool_servers": [], "variables": {"{{USER_NAME}}": "User", "{{USER_LOCATION}}": "Unknown", "{{CURRENT_DATETIME}}": time.strftime("%Y-%m-%d %H:%M:%S"),},}
         
-        # 构造用于签名的负载
-        timestamp = self._get_timestamp_millis()
-        now = datetime.now(timezone.utc)
+        # --- FIX: Timezone-aware payload generation ---
+        # The JS code uses the local timezone. We'll simulate a common one (Asia/Shanghai, UTC+8).
+        # This is the most likely cause of the signature mismatch.
+        target_tz_offset_hours = 8
+        target_tz = timezone(timedelta(hours=target_tz_offset_hours))
+        now_local = datetime.now(target_tz)
+        now_utc = datetime.now(timezone.utc)
         
+        # JS getTimezoneOffset returns the difference in minutes from UTC, and the sign is inverted.
+        # e.g., for UTC+8, it returns -480.
+        timezone_offset_minutes = -target_tz_offset_hours * 60
+
         payload_data = {
-            'timestamp': str(timestamp),
+            'timestamp': str(self._get_timestamp_millis()),
             'requestId': request_id,
             'user_id': user_id,
             'token': ck,
-            'user_agent': "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36",
+            'user_agent': "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36",
             'current_url': f"https://chat.z.ai/c/{chat_id}",
             'pathname': f"/c/{chat_id}",
-            'timezone': 'UTC', # 简化时区处理
-            'timezone_offset': '0',
-            'local_time': now.isoformat(),
-            'utc_time': now.strftime('%a, %d %b %Y %H:%M:%S GMT'),
+            'timezone': 'Asia/Shanghai',  # Use a specific timezone name
+            'timezone_offset': str(timezone_offset_minutes),
+            'local_time': now_local.isoformat(timespec='milliseconds'),
+            'utc_time': now_utc.strftime('%a, %d %b %Y %H:%M:%S GMT'),
             'version': '0.0.1',
             'platform': 'web',
             'language': 'zh-CN',
@@ -181,14 +160,17 @@ class ProxyHandler:
         sorted_payload = ",".join([f"{k},{payload_data[k]}" for k in keys])
         url_params = urllib.parse.urlencode(payload_data)
 
-        # 获取最后一条消息作为 mt
-        last_message = req.messages[-1].content if req.messages else ""
+        last_message_content = ""
+        if req.messages:
+            last_message = req.messages[-1]
+            if isinstance(last_message.content, str):
+                last_message_content = last_message.content
         
-        signature, sig_timestamp = self._generate_signature(ck, sorted_payload, last_message)
+        signature, sig_timestamp = self._generate_signature(ck, sorted_payload, last_message_content)
         
         final_url = f"{settings.UPSTREAM_URL}?{url_params}&signature_timestamp={sig_timestamp}"
 
-        headers = { "Content-Type": "application/json", "Authorization": f"Bearer {ck}", "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36"), "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN", "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.79", "X-Signature": signature, "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",}
+        headers = { "Content-Type": "application/json", "Authorization": f"Bearer {ck}", "User-Agent": ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/141.0.0.0 Safari/537.36"), "Accept": "application/json, text/event-stream", "Accept-Language": "zh-CN", "sec-ch-ua": '"Not)A;Brand";v="8", "Chromium";v="141", "Google Chrome";v="141"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"macOS"', "x-fe-version": "prod-fe-1.0.79", "X-Signature": signature, "Origin": "https://chat.z.ai", "Referer": "https://chat.z.ai/",}
         
         return body, headers, ck, final_url
         
@@ -200,7 +182,6 @@ class ProxyHandler:
             think_open = False
             yielded_think_buffer = "" 
             current_raw_thinking = ""
-            # **FIX**: State to handle the transition from thinking to answer
             is_first_answer_chunk = True
 
             async def yield_delta(content_type: str, text: str):
@@ -209,19 +190,15 @@ class ProxyHandler:
                     if not think_open:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '<think>'}, 'finish_reason': None}]})}\n\n"
                         think_open = True
-                    
                     cleaned_full_text = self._clean_thinking_content(text)
-                    delta_to_send = cleaned_full_text[len(yielded_think_buffer):] if cleaned_full_text.startswith(yielded_think_buffer) else cleaned_full_text
-                    
+                    delta_to_send = cleaned_full_text[len(yielded_think_buffer):]
                     if delta_to_send:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': delta_to_send}, 'finish_reason': None}]})}\n\n"
                     yielded_think_buffer = cleaned_full_text
-
                 elif content_type == "answer":
                     if think_open:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': '</think>'}, 'finish_reason': None}]})}\n\n"
                         think_open = False
-                    
                     cleaned_text = self._clean_answer_content(text)
                     if cleaned_text:
                         yield f"data: {json.dumps({'id': comp_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': req.model, 'choices': [{'index': 0, 'delta': {'content': cleaned_text}, 'finish_reason': None}]})}\n\n"
@@ -230,6 +207,7 @@ class ProxyHandler:
                 if resp.status_code != 200:
                     await cookie_manager.mark_cookie_failed(ck); err_body = await resp.aread()
                     err_msg = f"Error: {resp.status_code} - {err_body.decode(errors='ignore')}"
+                    logger.error(f"Upstream error: {err_msg}")
                     err = {"id": comp_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, "choices": [{"index": 0, "delta": {"content": err_msg}, "finish_reason": "stop"}],}
                     yield f"data: {json.dumps(err)}\n\n"; yield "data: [DONE]\n\n"; return
                 await cookie_manager.mark_cookie_success(ck)
@@ -238,7 +216,6 @@ class ProxyHandler:
                     for line in raw.strip().split('\n'):
                         line = line.strip()
                         if not line.startswith('data: '): continue
-                        
                         payload_str = line[6:]
                         if payload_str == '[DONE]':
                             if think_open: 
@@ -248,34 +225,23 @@ class ProxyHandler:
                             return
                         try:
                             dat = json.loads(payload_str).get("data", {})
-                        except (json.JSONDecodeError, AttributeError): 
-                            continue
+                        except (json.JSONDecodeError, AttributeError): continue
                         
                         phase = dat.get("phase")
                         content_chunk = dat.get("delta_content") or dat.get("edit_content")
-
-                        if not content_chunk:
-                            continue
+                        if not content_chunk: continue
 
                         if phase == "thinking":
-                            if dat.get("edit_content") is not None:
-                                current_raw_thinking = content_chunk
-                            else:
-                                current_raw_thinking += content_chunk
+                            current_raw_thinking = content_chunk if dat.get("edit_content") is not None else current_raw_thinking + content_chunk
                             async for item in yield_delta("thinking", current_raw_thinking):
                                 yield item
-                        
                         elif phase == "answer":
                             content_to_process = content_chunk
-                            # **FIX**: Special handling for the first answer chunk
                             if is_first_answer_chunk:
-                                # The first answer chunk often contains leftover thinking content.
-                                # We split by '</details>' and only use the part after it.
                                 if '</details>' in content_to_process:
                                     parts = content_to_process.split('</details>', 1)
                                     content_to_process = parts[1] if len(parts) > 1 else ""
                                 is_first_answer_chunk = False
-                            
                             if content_to_process:
                                 async for item in yield_delta("answer", content_to_process):
                                     yield item
@@ -291,6 +257,7 @@ class ProxyHandler:
             async with self.client.stream("POST", url, json=body, headers=headers) as resp:
                 if resp.status_code != 200:
                     await cookie_manager.mark_cookie_failed(ck); error_detail = await resp.text()
+                    logger.error(f"Upstream error: {resp.status_code} - {error_detail}")
                     raise HTTPException(resp.status_code, f"Upstream error: {error_detail}")
                 await cookie_manager.mark_cookie_success(ck)
                 
@@ -309,26 +276,18 @@ class ProxyHandler:
                         
                         phase = dat.get("phase")
                         content_chunk = dat.get("delta_content") or dat.get("edit_content")
-
-                        if not content_chunk:
-                            continue
+                        if not content_chunk: continue
 
                         if phase == "thinking":
-                            if dat.get("edit_content") is not None:
-                                current_raw_thinking = content_chunk
-                            else:
-                                current_raw_thinking += content_chunk
+                            current_raw_thinking = content_chunk if dat.get("edit_content") is not None else current_raw_thinking + content_chunk
                             last_thinking_content = current_raw_thinking
-                        
                         elif phase == "answer":
                             content_to_process = content_chunk
-                            # **FIX**: Apply same logic to non-stream mode
                             if is_first_answer_chunk:
                                 if '</details>' in content_to_process:
                                     parts = content_to_process.split('</details>', 1)
                                     content_to_process = parts[1] if len(parts) > 1 else ""
                                 is_first_answer_chunk = False
-                            
                             if content_to_process:
                                 raw_answer_parts.append(content_to_process)
                     else: 
@@ -336,7 +295,6 @@ class ProxyHandler:
                     break
 
             full_answer = ''.join(raw_answer_parts)
-            # The final cleaning is still useful for any other residual tags
             cleaned_ans_text = self._clean_answer_content(full_answer).strip()
             final_content = cleaned_ans_text