update: 更新聊天服务、客户端连接池和辅助工具

app/server/chat.py

@@ -52,6 +52,7 @@ from ..utils.helper import (
     estimate_tokens,
     extract_image_dimensions,
     extract_tool_calls,
+    retry_with_backoff,
     strip_system_hints,
     text_from_message,
 )
@@ -774,14 +775,31 @@ async def _send_with_split(
     files: list[Path | str | io.BytesIO] | None = None,
     stream: bool = False,
 ) -> AsyncGenerator[ModelOutput, None] | ModelOutput:
-    """Send text to Gemini, splitting or converting to attachment if too long."""
+    """
+    Send text to Gemini, splitting or converting to attachment if too long.
+    Includes retry with exponential backoff for transient failures.
+    """
+    
+    @retry_with_backoff(
+        max_retries=3,
+        base_delay=1.0,
+        max_delay=30.0,
+        exponential_base=2.0,
+        retryable_exceptions=(ConnectionError, TimeoutError, OSError),
+    )
+    async def _send_with_retry(
+        content: str, file_list: list | None, is_stream: bool
+    ) -> AsyncGenerator[ModelOutput, None] | ModelOutput:
+        """Internal function with retry logic."""
+        if is_stream:
+            return session.send_message_stream(content, files=file_list)
+        return await session.send_message(content, files=file_list)
+    
     if len(text) <= MAX_CHARS_PER_REQUEST:
         try:
-            if stream:
-                return session.send_message_stream(text, files=files)
-            return await session.send_message(text, files=files)
+            return await _send_with_retry(text, files, stream)
         except Exception as e:
-            logger.exception(f"Error sending message to Gemini: {e}")
+            logger.exception(f"Error sending message to Gemini after retries: {e}")
             raise
 
     logger.info(
@@ -799,11 +817,9 @@ async def _send_with_split(
             "2. Treat that content as the **primary** user prompt for this turn.\n"
             "3. Execute the instructions or answer the questions found *inside* that file immediately.\n"
         )
-        if stream:
-            return session.send_message_stream(instruction, files=final_files)
-        return await session.send_message(instruction, files=final_files)
+        return await _send_with_retry(instruction, final_files, stream)
     except Exception as e:
-        logger.exception(f"Error sending large text as file to Gemini: {e}")
+        logger.exception(f"Error sending large text as file to Gemini after retries: {e}")
         raise
 
 
@@ -820,6 +836,14 @@ class StreamingOutputFilter:
     )
     # Citation markers (e.g., [1] or 【1†source】)
     CITATION_RE = re.compile(r"【\d+†source】|\[\d+\]")
+    # Match patterns for JSON-like content
+    JSON_START_RE = re.compile(r"\s*\{")
+    # Match whitespace for faster stripping
+    LEADING_WS_RE = re.compile(r"^\s+")
+    # Pre-compiled patterns for streaming performance
+    TRAILING_FENCE_RE = re.compile(r'\n?```$')
+    URL_DELIMITER_RE = re.compile(r"[\s\]\)]")
+    BRACKET_NUM_RE = re.compile(r"\[\d+\]")
 
     def __init__(self):
         self.buffer = ""
@@ -938,7 +962,7 @@ class StreamingOutputFilter:
                     self.state = "IN_POTENTIAL_SOURCE"
 
             elif self.state == "IN_POTENTIAL_URL":
-                match = re.search(r"[\s\]\)]", self.buffer)
+                match = self.URL_DELIMITER_RE.search(self.buffer)
                 if match:
                     end_idx = match.start()
                     url_candidate = self.buffer[:end_idx]
@@ -975,7 +999,7 @@ class StreamingOutputFilter:
                 end_idx = self.buffer.find("]")
                 if end_idx != -1:
                     full_match = self.buffer[:end_idx + 1]
-                    if re.match(r"\[\d+\]", full_match):
+                    if self.BRACKET_NUM_RE.match(full_match):
                         # It's a citation, skip it
                         self.buffer = self.buffer[end_idx + 1:]
                         self.state = "NORMAL"
@@ -1096,10 +1120,16 @@ class StreamingOutputFilter:
                     args_part = self.buffer[:stop_idx]
                     if not self.args_started:
                         stripped_args = args_part.lstrip()
-                        if stripped_args and not stripped_args.startswith("{"):
-                            args_part = "" 
-                        elif stripped_args:
+                        # More robust handling: accept any non-empty content
+                        # Even if it doesn't start with '{', we still emit it
+                        # The JSON parsing will handle errors downstream
+                        if stripped_args:
                             self.args_started = True
+                            # If doesn't start with '{', try to fix common issues
+                            if not stripped_args.startswith("{"):
+                                # Check if it might be a valid JSON that's missing the opening brace
+                                # or if it's just plain text arguments
+                                args_part = "{" + args_part.lstrip()
 
                     if args_part:
                         events.append({"type": "tool_delta", "index": self.tool_call_index, "content": args_part})
@@ -1121,19 +1151,27 @@ class StreamingOutputFilter:
                             if stripped.startswith("{"):
                                 self.args_started = True
                             else:
-                                events.append({"type": "tool_end", "index": self.tool_call_index})
-                                self.tool_call_index += 1
-                                self.current_call_id = None
-                                self.args_started = False
-                                self.state = "IN_TOOL_BLOCK"
-                                continue
-
-                    keep_len = max(len(self.CALL_END), len(self.TOOL_END_PREFIX)) - 1
-                    if len(self.buffer) > keep_len:
-                        chunk_to_send = self.buffer[:-keep_len]
-                        if chunk_to_send:
-                            events.append({"type": "tool_delta", "index": self.tool_call_index, "content": chunk_to_send})
-                        self.buffer = self.buffer[-keep_len:]
+                                # Be more tolerant: accept arguments even without opening brace
+                                # Common issue: LLM outputs `key: value` instead of `{"key": "value"}`
+                                self.args_started = True
+                    
+                    if self.args_started:
+                        keep_len = max(len(self.CALL_END), len(self.TOOL_END_PREFIX)) - 1
+                        if len(self.buffer) > keep_len:
+                            chunk_to_send = self.buffer[:-keep_len]
+                            if chunk_to_send:
+                                events.append({"type": "tool_delta", "index": self.tool_call_index, "content": chunk_to_send})
+                            self.buffer = self.buffer[-keep_len:]
+                    else:
+                        # Buffer has content but we haven't found JSON start yet
+                        # Wait a bit more for potential JSON start
+                        if len(self.buffer) > 100:
+                            # Timeout waiting for JSON, emit what we have
+                            events.append({"type": "tool_end", "index": self.tool_call_index})
+                            self.tool_call_index += 1
+                            self.current_call_id = None
+                            self.args_started = False
+                            self.state = "IN_TOOL_BLOCK"
                     break
 
             elif self.state == "IN_TAG":
@@ -1256,7 +1294,7 @@ def _create_real_streaming_response(
                 new_text = chunk.text or ""
                 # Strip transient triple-backtick fence which is often added by the web-scraping backend
                 # during streaming snapshots. We will restore the real ending after the loop.
-                display_text = re.sub(r'\n?```$', '', new_text)
+                display_text = StreamingOutputFilter.TRAILING_FENCE_RE.sub('', new_text)
                 
                 if len(display_text) > len(full_text):
                     text_delta = display_text[len(full_text) :]
@@ -1462,7 +1500,7 @@ def _create_responses_real_streaming_response(
                 new_text = chunk.text or ""
                 # Strip transient triple-backtick fence which is often added by the web-scraping backend
                 # during streaming snapshots. We will restore the real ending after the loop.
-                display_text = re.sub(r'\n?```$', '', new_text)
+                display_text = StreamingOutputFilter.TRAILING_FENCE_RE.sub('', new_text)
 
                 if len(display_text) > len(full_text):
                     text_delta = display_text[len(full_text) :]
@@ -1759,7 +1797,7 @@ async def _create_anthropic_streaming_response(
             async for chunk in generator:
                 last_chunk = chunk
                 if chunk.text:
-                    display_text = re.sub(r'\n?```$', '', chunk.text)
+                    display_text = StreamingOutputFilter.TRAILING_FENCE_RE.sub('', chunk.text)
                     if len(display_text) > len(full_text):
                         delta_text = display_text[len(full_text):]
                         full_text = display_text

app/server/health.py

@@ -9,19 +9,20 @@ router = APIRouter()
 
 @router.get("/health", response_model=HealthCheckResponse)
 async def health_check():
+    """
+    Health check endpoint.
+    Only checks current status without re-initializing clients.
+    Re-initialization is handled automatically by the pool when needed.
+    """
     pool = GeminiClientPool()
     db = LMDBConversationStore()
 
-    try:
-        await pool.init()
-    except Exception as e:
-        logger.error(f"Failed to initialize Gemini clients: {e}")
-        return HealthCheckResponse(ok=False, error=str(e))
-
+    # Only check status, don't re-initialize
+    # The pool handles client restarts automatically in acquire()
     client_status = pool.status()
 
     if not all(client_status.values()):
-        logger.warning("One or more Gemini clients not running")
+        logger.warning(f"One or more Gemini clients not running: {client_status}")
 
     stat = db.stats()
     if not stat:

app/services/client.py

@@ -61,9 +61,14 @@ def clean_citations(text: str) -> str:
     return "".join(parts)
 
 
+# Pre-compiled citation pattern for performance
+CITATION_PATTERN = re.compile(r"(?<!/)(?<![a-zA-Z0-9])\[\d+\]")
+SOURCE_CITATION_PATTERN = re.compile(r"【\d+†source】")
+
+
 def _clean_text_chunk(text: str) -> str:
     # A. Always remove source-style citations
-    text = re.sub(r"【\d+†source】", "", text)
+    text = SOURCE_CITATION_PATTERN.sub("", text)
     
     # B. Remove internal technical artifacts and links
     # These are raw internal URLs or IDs that Gemini Web sometimes leaks.
@@ -73,7 +78,6 @@ def _clean_text_chunk(text: str) -> str:
     # C. Remove standard [n] citations
     # Use careful regex to avoid matching list items (start of line) or URLs
     # Negative lookbehind to ensure we aren't part of a URL or Markdown link path
-    citation_pattern = re.compile(r"(?<!/)(?<![a-zA-Z0-9])\[\d+\]")
     
     def repl(m):
         try:
@@ -91,7 +95,7 @@ def _clean_text_chunk(text: str) -> str:
             logger.error(f"Error in _clean_text_chunk.repl: {e}")
             return m.group(0)
 
-    return citation_pattern.sub(repl, text)
+    return CITATION_PATTERN.sub(repl, text)
 
 
 def _resolve(value: Any, fallback: Any):

app/services/pool.py

@@ -17,6 +17,7 @@ class GeminiClientPool(metaclass=Singleton):
         self._id_map: Dict[str, GeminiClientWrapper] = {}
         self._round_robin: deque[GeminiClientWrapper] = deque()
         self._restart_locks: Dict[str, asyncio.Lock] = {}
+        self._round_robin_lock = asyncio.Lock()  # Lock for thread-safe round-robin
 
         if len(g_config.gemini.clients) == 0:
             raise ValueError("No Gemini clients configured")
@@ -69,9 +70,19 @@ class GeminiClientPool(metaclass=Singleton):
                 f"Gemini client {client_id} is not running and could not be restarted"
             )
 
-        for _ in range(len(self._round_robin)):
-            client = self._round_robin[0]
-            self._round_robin.rotate(-1)
+        # Thread-safe round-robin: try each client once
+        tried_clients = set()
+        while len(tried_clients) < len(self._round_robin):
+            # Atomically get next client
+            async with self._round_robin_lock:
+                client = self._round_robin[0]
+                self._round_robin.rotate(-1)
+                if client in tried_clients:
+                    # Already tried all clients
+                    break
+                tried_clients.add(client)
+            
+            # Check readiness outside lock to avoid blocking other requests
             if await self._ensure_client_ready(client):
                 return client
 

app/utils/helper.py

@@ -1,12 +1,15 @@
 import base64
+import functools
 import hashlib
 import mimetypes
+import random
 import re
 import reprlib
 import struct
 import tempfile
+import json
 from pathlib import Path
-from urllib.parse import urlparse
+from typing import Any, Callable, TypeVar
 
 import httpx
 import orjson
@@ -14,6 +17,88 @@ from loguru import logger
 
 from ..models import FunctionCall, Message, ToolCall
 
+T = TypeVar("T")
+
+
+def retry_with_backoff(
+    max_retries: int = 3,
+    base_delay: float = 1.0,
+    max_delay: float = 30.0,
+    exponential_base: float = 2.0,
+    retryable_exceptions: tuple = (Exception,),
+):
+    """
+    Decorator that retries a function with exponential backoff.
+    
+    Args:
+        max_retries: Maximum number of retry attempts
+        base_delay: Initial delay in seconds
+        max_delay: Maximum delay cap in seconds
+        exponential_base: Base for exponential backoff calculation
+        retryable_exceptions: Tuple of exception types to retry on
+    """
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs) -> T:
+            last_exception = None
+            for attempt in range(max_retries + 1):
+                try:
+                    return await func(*args, **kwargs)
+                except retryable_exceptions as e:
+                    last_exception = e
+                    if attempt == max_retries:
+                        logger.error(f"All {max_retries + 1} attempts failed for {func.__name__}")
+                        raise
+                    
+                    # Calculate delay with exponential backoff and jitter
+                    delay = min(base_delay * (exponential_base ** attempt), max_delay)
+                    jitter = random.uniform(0.1, 0.3) * delay
+                    actual_delay = delay + jitter
+                    
+                    logger.warning(
+                        f"Attempt {attempt + 1}/{max_retries + 1} failed for {func.__name__}: {e}. "
+                        f"Retrying in {actual_delay:.1f}s..."
+                    )
+                    
+                    import asyncio
+                    await asyncio.sleep(actual_delay)
+            
+            raise last_exception
+        
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs) -> T:
+            last_exception = None
+            for attempt in range(max_retries + 1):
+                try:
+                    return func(*args, **kwargs)
+                except retryable_exceptions as e:
+                    last_exception = e
+                    if attempt == max_retries:
+                        logger.error(f"All {max_retries + 1} attempts failed for {func.__name__}")
+                        raise
+                    
+                    delay = min(base_delay * (exponential_base ** attempt), max_delay)
+                    jitter = random.uniform(0.1, 0.3) * delay
+                    actual_delay = delay + jitter
+                    
+                    logger.warning(
+                        f"Attempt {attempt + 1}/{max_retries + 1} failed for {func.__name__}: {e}. "
+                        f"Retrying in {actual_delay:.1f}s..."
+                    )
+                    
+                    import time
+                    time.sleep(actual_delay)
+            
+            raise last_exception
+        
+        # Return appropriate wrapper based on whether function is async
+        import asyncio
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        return sync_wrapper
+    
+    return decorator
+
 VALID_TAG_ROLES = {"user", "assistant", "system", "tool"}
 TOOL_WRAP_HINT = (
     "\nYou MUST wrap every tool call response inside a single [function_calls] block exactly like:\n"
@@ -34,6 +119,98 @@ TOOL_HINT_LINE_START = _hint_lines[0] if _hint_lines else ""
 TOOL_HINT_LINE_END = _hint_lines[-1] if _hint_lines else ""
 
 
+def safe_json_loads(s: str) -> Any:
+    """
+    Attempt to parse JSON with multiple fallback strategies for LLM-generated content.
+    Handles common issues like:
+    - Markdown code blocks
+    - Trailing commas
+    - Missing opening/closing braces
+    - Unescaped control characters
+    - Single quotes instead of double quotes
+    - Python-style literals (None, True, False)
+    """
+    if not s:
+        return None
+
+    # 1. Try direct parse with orjson
+    try:
+        return orjson.loads(s)
+    except orjson.JSONDecodeError:
+        pass
+
+    # 2. Clean markdown and extract JSON block
+    cleaned = s.strip()
+    # Remove markdown code blocks
+    cleaned = re.sub(r"```(?:json)?\s*(.*?)\s*```", r"\1", cleaned, flags=re.DOTALL)
+    
+    # Extract the first { and last }
+    match = re.search(r"(\{.*\})", cleaned, re.DOTALL)
+    if match:
+        cleaned = match.group(1)
+
+    # 3. Fix common syntax errors
+    # Fix trailing commas
+    cleaned = re.sub(r",\s*([\}\]])", r"\1", cleaned)
+    
+    # Fix missing opening brace (common in LLM output)
+    if cleaned.strip() and not cleaned.strip().startswith("{"):
+        # Check if it looks like JSON content
+        if ":" in cleaned or "[" in cleaned:
+            cleaned = "{" + cleaned
+    
+    # Fix missing closing brace
+    open_braces = cleaned.count("{") - cleaned.count("}")
+    if open_braces > 0:
+        cleaned = cleaned + "}" * open_braces
+    
+    # Fix unescaped newlines in strings
+    cleaned = re.sub(r'(?<!\\)\n(?=[^"]*"[^"]*$)', r"\\n", cleaned)
+    
+    # 4. Try orjson again
+    try:
+        return orjson.loads(cleaned)
+    except orjson.JSONDecodeError:
+        pass
+    
+    # 5. Try to fix single quotes to double quotes
+    try:
+        # Simple heuristic: if it has single quotes around keys/values
+        fixed = re.sub(r"'([^']+)'", r'"\1"', cleaned)
+        return orjson.loads(fixed)
+    except orjson.JSONDecodeError:
+        pass
+
+    # 6. Fallback to standard json module with strict=False
+    # strict=False allows control characters like unescaped newlines inside strings
+    try:
+        return json.loads(cleaned, strict=False)
+    except Exception:
+        pass
+    
+    # 7. Last resort: try to extract key-value pairs manually
+    try:
+        result = {}
+        # Match patterns like "key": "value" or "key": value
+        pattern = r'"([^"]+)"\s*:\s*(?:"([^"]*)"|([\w\d]+)|(\{[^}]*\})|(\[[^\]]*\]))'
+        for m in re.finditer(pattern, cleaned):
+            key = m.group(1)
+            value = m.group(2) or m.group(3) or m.group(4) or m.group(5)
+            if value:
+                # Try to parse nested structures
+                try:
+                    value = orjson.loads(value)
+                except:
+                    pass
+                result[key] = value
+        if result:
+            return result
+    except Exception:
+        pass
+
+    raise ValueError(f"Failed to parse JSON: {reprlib.repr(s)}")
+
+
 def add_tag(role: str, content: str, unclose: bool = False) -> str:
     """Surround content with role tags"""
     if role not in VALID_TAG_ROLES:
@@ -198,25 +375,12 @@ def _process_tools_internal(text: str, extract: bool = True) -> tuple[str, list[
 
         arguments = raw_args
         try:
-            parsed_args = orjson.loads(raw_args)
+            parsed_args = safe_json_loads(raw_args)
             arguments = orjson.dumps(parsed_args, option=orjson.OPT_SORT_KEYS).decode("utf-8")
-        except orjson.JSONDecodeError:
-            json_match = re.search(r"({.*})", raw_args, re.DOTALL)
-            if json_match:
-                potential_json = json_match.group(1)
-                try:
-                    parsed_args = orjson.loads(potential_json)
-                    arguments = orjson.dumps(parsed_args, option=orjson.OPT_SORT_KEYS).decode(
-                        "utf-8"
-                    )
-                except orjson.JSONDecodeError:
-                    logger.warning(
-                        f"Failed to parse extracted JSON arguments for '{name}': {reprlib.repr(potential_json)}"
-                    )
-            else:
-                logger.warning(
-                    f"Failed to parse tool call arguments for '{name}'. Passing raw string: {reprlib.repr(raw_args)}"
-                )
+        except Exception as e:
+            logger.warning(
+                f"Failed to parse tool call arguments for '{name}': {e}. Passing raw string: {reprlib.repr(raw_args)}"
+            )
 
         index = len(tool_calls)
         seed = f"{name}:{arguments}:{index}".encode("utf-8")

config/config.yaml

@@ -19,8 +19,8 @@ cors:
 gemini:
   clients:
     - id: "example-id-1"   # Arbitrary client ID
-      secure_1psid: "g.a0006gg0IXONYyliA8BQ9p7rvNa-tvX8KvU_yuYYLzLsV3Zsk0LIxs7vB6PIUDbHm9H-9n9dmwACgYKAV0SARYSFQHGX2Mi_xRuEWuRzLS671F_NBGIUBoVAUF8yKrBOirFIiZtZ0ivTkmPCCcd0076"
-      secure_1psidts: "sidts-CjEB7I_69FYuREW6Ijjqid7kS3bws_QplihBxYrc-uMS7QOuyYTDIK_wsh_lJ3ebHrwKEAA"
+      secure_1psid: "g.a0006ghLMnrH-xrCpx9MPvW3muBiYT-ndrFy-0jmTZqWQYn6pP3ofLfBrtYDuTRlVmiMppg2LwACgYKATISARcSFQHGX2MisJCoQlKW5k8gkHYqbDf6TBoVAUF8yKqSTWG_iCHAyw2Yhwy6-gF40076"
+      secure_1psidts: "sidts-CjEB7I_69CNJfUZU4ue_qIqhCNx91dOCLnIy_DSoOAonz924_M99NJe-qnDvDXjPUe5uEAA"
       proxy: null          # Optional proxy URL (null/empty means direct connection)
   timeout: 120             # Init timeout in seconds
   auto_refresh: true       # Auto-refresh session cookies