fix: SSE keepalive comments every 5s to prevent AstrBot ReadTimeout

main.py

@@ -114,90 +114,99 @@ def verify_api_key(authorization: Optional[str] = Header(None)) -> str:
 
 
 async def _stream_chat_response(
-    browser,
-    prompt: str,
-    model: str,
-    has_tools: bool,
-    manager,
-    account,
+    browser, prompt, model, has_tools, manager, account,
 ):
-    """Shared async generator for streaming chat completions with optional tool call detection."""
     chunk_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
-    try:
-        is_tool_call = False
-        not_tool_call = False
-        content_buffer = ""
-
-        async for chunk_data in browser.stream_message(prompt, timeout=120, model=model):
-            chunk_type = chunk_data.get("type", "content")
-            chunk_text = chunk_data.get("chunk", "")
-
-            if chunk_type == "thinking":
-                delta = {"reasoning_content": chunk_text}
-            else:
-                if has_tools and not is_tool_call and not not_tool_call:
-                    content_buffer += chunk_text
-                    if len(content_buffer) < 12:
-                        if not "<tool_call>".startswith(content_buffer):
-                            not_tool_call = True
-                            delta = {"content": content_buffer}
+    q = asyncio.Queue()
+
+    async def _producer():
+        try:
+            is_tool_call = False
+            not_tool_call = False
+            content_buffer = ""
+            async for chunk_data in browser.stream_message(prompt, timeout=120, model=model):
+                chunk_type = chunk_data.get("type", "content")
+                chunk_text = chunk_data.get("chunk", "")
+                if chunk_type == "thinking":
+                    delta = {"reasoning_content": chunk_text}
+                else:
+                    if has_tools and not is_tool_call and not not_tool_call:
+                        content_buffer += chunk_text
+                        if len(content_buffer) < 12:
+                            if not "<tool_call>".startswith(content_buffer):
+                                not_tool_call = True
+                                delta = {"content": content_buffer}
+                            else:
+                                continue
                         else:
-                            continue
+                            if content_buffer.startswith("<tool_call>"):
+                                is_tool_call = True
+                                continue
+                            else:
+                                not_tool_call = True
+                                delta = {"content": content_buffer}
+                    elif has_tools and is_tool_call:
+                        content_buffer += chunk_text
+                        continue
                     else:
-                        if content_buffer.startswith("<tool_call>"):
-                            is_tool_call = True
-                            continue
-                        else:
-                            not_tool_call = True
-                            delta = {"content": content_buffer}
-                elif has_tools and is_tool_call:
-                    content_buffer += chunk_text
-                    continue
-                else:
-                    delta = {"content": chunk_text}
-
-            data = {
-                "id": chunk_id,
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": model,
-                "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
-            }
-            yield f"data: {json.dumps(data)}\n\n"
-
-        if is_tool_call:
-            m = re.search(r'<tool_call>(.*?)</tool_call>', content_buffer, re.DOTALL)
-            if m:
-                try:
-                    tcall = json.loads(m.group(1))
-                    delta = {
-                        "tool_calls": [{
-                            "index": 0,
-                            "id": f"call_{uuid.uuid4().hex[:8]}",
-                            "type": "function",
-                            "function": {
-                                "name": tcall.get("name", ""),
-                                "arguments": json.dumps(tcall.get("arguments", {}))
-                            }
-                        }]
-                    }
-                    data = {
-                        "id": chunk_id,
-                        "object": "chat.completion.chunk",
-                        "created": int(time.time()),
-                        "model": model,
-                        "choices": [{"index": 0, "delta": delta, "finish_reason": "tool_calls"}]
-                    }
-                    yield f"data: {json.dumps(data)}\n\n"
-                except Exception as e:
-                    logger.error("Failed to parse tool call: %s", e)
+                        delta = {"content": chunk_text}
+                data = {
+                    "id": chunk_id, "object": "chat.completion.chunk",
+                    "created": int(time.time()), "model": model,
+                    "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                }
+                await q.put(f"data: {json.dumps(data)}\n\n")
+            if is_tool_call:
+                m = re.search(r'<tool_call>(.*?)</tool_call>', content_buffer, re.DOTALL)
+                if m:
+                    try:
+                        tcall = json.loads(m.group(1))
+                        delta = {
+                            "tool_calls": [{
+                                "index": 0,
+                                "id": f"call_{uuid.uuid4().hex[:8]}",
+                                "type": "function",
+                                "function": {
+                                    "name": tcall.get("name", ""),
+                                    "arguments": json.dumps(tcall.get("arguments", {}))
+                                }
+                            }]
+                        }
+                        data = {
+                            "id": chunk_id, "object": "chat.completion.chunk",
+                            "created": int(time.time()), "model": model,
+                            "choices": [{"index": 0, "delta": delta, "finish_reason": "tool_calls"}]
+                        }
+                        await q.put(f"data: {json.dumps(data)}\n\n")
+                    except Exception as e:
+                        logger.error("Failed to parse tool call: %s", e)
+            await q.put(f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n")
+            await q.put("data: [DONE]\n\n")
+        except Exception as e:
+            await q.put(f"data: {json.dumps({'error': {'message': str(e)}})}\n\n")
+        finally:
+            await q.put(None)
+            await manager.release(account)
+
+    async def _keepalive():
+        while True:
+            await asyncio.sleep(5)
+            try:
+                q.put_nowait(": keepalive\n\n")
+            except asyncio.QueueFull:
+                pass
 
-        yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': int(time.time()), 'model': model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
-        yield "data: [DONE]\n\n"
-    except Exception as e:
-        yield f"data: {json.dumps({'error': {'message': str(e)}})}\n\n"
+    pt = asyncio.create_task(_producer())
+    kt = asyncio.create_task(_keepalive())
+    try:
+        while True:
+            item = await q.get()
+            if item is None:
+                break
+            yield item
     finally:
-        await manager.release(account)
+        kt.cancel()
+        pt.cancel()
 
 
 @app.get("/v1/models")