fix: 性能、系统工具不显示、rate limit不正常

Mag1cFall · Mag1cFall · commit 058628a4573a · 2025-12-09T23:40:44.000+08:00
diff --git a/data/excluded_models.txt b/data/excluded_models.txt
@@ -1,7 +1,2 @@
-veo-2.0-generate-001
-imagen-4.0-fast-generate-001
-imagen-4.0-ultra-generate-001
-imagen-4.0-generate-001
 gemini-2.5-flash-native-audio-preview-09-2025
-gemini-2.5-flash-image
 gemini-3-pro-image-preview
diff --git a/src/gateway.py b/src/gateway.py
@@ -2,6 +2,7 @@
 import json
 import asyncio
 import logging
+import time
 from typing import Optional, AsyncGenerator
 from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import StreamingResponse, Response
@@ -14,134 +15,159 @@
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
 PROJECT_ROOT = os.path.dirname(SCRIPT_DIR)
 DATA_DIR = os.path.join(PROJECT_ROOT, 'data')
-WORKERS_CONFIG_PATH = os.path.join(DATA_DIR, 'workers.json')
+
+MANAGER_URL = "http://127.0.0.1:9000"
+RATE_LIMIT_KEYWORDS = [b"exceeded quota", b"out of free generations", b"rate limit"]
 
 app = FastAPI(title="AIStudio2API Gateway")
 
-workers = []
-current_index = 0
+_session: Optional[aiohttp.ClientSession] = None
+_worker_cache = {"workers": [], "last_update": 0, "index": 0}
+CACHE_TTL = 5
 
-def load_workers():
-    global workers
-    if os.path.exists(WORKERS_CONFIG_PATH):
-        try:
-            with open(WORKERS_CONFIG_PATH, 'r', encoding='utf-8') as f:
-                config = json.load(f)
-            workers = [w['port'] for w in config.get('workers', [])]
-            logger.info(f"Loaded {len(workers)} workers: {workers}")
-        except Exception as e:
-            logger.error(f"Load workers failed: {e}")
+async def get_session() -> aiohttp.ClientSession:
+    global _session
+    if _session is None or _session.closed:
+        connector = aiohttp.TCPConnector(limit=100, limit_per_host=20, keepalive_timeout=30)
+        _session = aiohttp.ClientSession(connector=connector)
+    return _session
 
-def get_next_worker() -> Optional[int]:
-    global current_index, workers
-    if not workers:
+async def refresh_workers():
+    cache = _worker_cache
+    if time.time() - cache["last_update"] < CACHE_TTL and cache["workers"]:
+        return
+    try:
+        session = await get_session()
+        async with session.get(f"{MANAGER_URL}/api/workers", timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            workers = await resp.json()
+            cache["workers"] = [w for w in workers if w.get("status") == "running"]
+            cache["last_update"] = time.time()
+    except Exception as e:
+        logger.warning(f"Refresh workers failed: {e}")
+
+def get_next_worker(model: str = "") -> Optional[dict]:
+    cache = _worker_cache
+    available = cache["workers"]
+    if not available:
         return None
-    port = workers[current_index % len(workers)]
-    current_index += 1
-    return port
+    worker = available[cache["index"] % len(available)]
+    cache["index"] += 1
+    return worker
+
+async def report_rate_limit(worker_id: str, model: str):
+    try:
+        session = await get_session()
+        await session.post(f"{MANAGER_URL}/api/workers/{worker_id}/rate-limit", json={"model": model}, timeout=aiohttp.ClientTimeout(total=2))
+    except:
+        pass
+
+def check_rate_limit_in_response(content: bytes) -> bool:
+    content_lower = content.lower()
+    return any(kw in content_lower for kw in RATE_LIMIT_KEYWORDS)
 
 @app.on_event("startup")
 async def startup():
-    load_workers()
-    logger.info(f"Gateway started with {len(workers)} workers")
+    await refresh_workers()
+    logger.info(f"Gateway started")
+
+@app.on_event("shutdown")
+async def shutdown():
+    global _session
+    if _session and not _session.closed:
+        await _session.close()
 
 @app.get("/")
 async def root():
-    return {"status": "ok", "mode": "gateway", "workers": len(workers)}
+    return {"status": "ok", "mode": "gateway", "workers": len(_worker_cache["workers"])}
 
 @app.get("/v1/models")
 async def models():
-    port = get_next_worker()
-    if not port:
+    await refresh_workers()
+    worker = get_next_worker()
+    if not worker:
         raise HTTPException(status_code=503, detail="No workers available")
     
+    port = worker["port"]
     url = f"http://127.0.0.1:{port}/v1/models"
-    logger.info(f"GET /v1/models -> worker:{port}")
     
-    timeout = aiohttp.ClientTimeout(total=60)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        try:
-            async with session.get(url) as resp:
-                content = await resp.read()
-                return Response(content=content, status_code=resp.status, media_type=resp.content_type)
-        except Exception as e:
-            logger.error(f"Forward /v1/models failed: {e}")
-            raise HTTPException(status_code=502, detail=str(e))
+    session = await get_session()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp:
+            content = await resp.read()
+            return Response(content=content, status_code=resp.status, media_type=resp.content_type)
+    except Exception as e:
+        logger.error(f"Forward /v1/models failed: {e}")
+        raise HTTPException(status_code=502, detail=str(e))
 
 @app.post("/v1/chat/completions")
 async def chat_completions(request: Request):
+    await refresh_workers()
     body = await request.body()
     body_json = json.loads(body)
     is_stream = body_json.get("stream", False)
+    model_id = body_json.get("model", "")
     
-    port = get_next_worker()
-    if not port:
+    worker = get_next_worker(model_id)
+    if not worker:
         raise HTTPException(status_code=503, detail="No workers available")
     
+    port = worker["port"]
+    worker_id = worker.get("id", "")
     url = f"http://127.0.0.1:{port}/v1/chat/completions"
-    req_id = f"gw-{current_index}"
+    req_id = f"gw-{worker_id}"
     logger.info(f"[{req_id}] POST -> worker:{port} (stream={is_stream})")
     
     forward_headers = {'Content-Type': 'application/json'}
     for k, v in request.headers.items():
-        k_lower = k.lower()
-        if k_lower not in ('host', 'content-length', 'transfer-encoding', 'content-type'):
+        if k.lower() not in ('host', 'content-length', 'transfer-encoding', 'content-type'):
             forward_headers[k] = v
     
+    session = await get_session()
+    
     if is_stream:
         async def stream_proxy() -> AsyncGenerator[bytes, None]:
-            timeout = aiohttp.ClientTimeout(total=600, sock_read=300)
-            connector = aiohttp.TCPConnector(force_close=True)
-            async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
-                try:
-                    async with session.post(url, data=body, headers=forward_headers) as resp:
-                        logger.info(f"[{req_id}] Stream started, status={resp.status}")
-                        chunk_count = 0
-                        async for chunk in resp.content.iter_chunks():
-                            data, end_of_chunk = chunk
-                            if data:
-                                chunk_count += 1
-                                yield data
-                        logger.info(f"[{req_id}] Stream completed, chunks={chunk_count}")
-                except asyncio.CancelledError:
-                    logger.warning(f"[{req_id}] Stream cancelled")
-                except Exception as e:
-                    logger.error(f"[{req_id}] Stream error: {e}")
-        
-        return StreamingResponse(
-            stream_proxy(),
-            media_type="text/event-stream",
-            headers={
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "X-Accel-Buffering": "no",
-                "Transfer-Encoding": "chunked"
-            }
-        )
-    else:
-        timeout = aiohttp.ClientTimeout(total=300)
-        async with aiohttp.ClientSession(timeout=timeout) as session:
+            rate_limited = False
+            check_count = 0
             try:
-                async with session.post(url, data=body, headers=forward_headers) as resp:
-                    content = await resp.read()
-                    logger.info(f"[{req_id}] Non-stream response, status={resp.status}, len={len(content)}")
-                    return Response(content=content, status_code=resp.status, media_type=resp.content_type)
+                async with session.post(url, data=body, headers=forward_headers, timeout=aiohttp.ClientTimeout(total=600, sock_read=300)) as resp:
+                    async for chunk in resp.content.iter_chunks():
+                        data, _ = chunk
+                        if data:
+                            check_count += 1
+                            if check_count <= 5 and not rate_limited:
+                                if check_rate_limit_in_response(data):
+                                    rate_limited = True
+                            yield data
+                    if rate_limited and worker_id and model_id:
+                        asyncio.create_task(report_rate_limit(worker_id, model_id))
+            except asyncio.CancelledError:
+                pass
             except Exception as e:
-                logger.error(f"[{req_id}] Forward failed: {e}")
-                raise HTTPException(status_code=502, detail=str(e))
+                logger.error(f"[{req_id}] Stream error: {e}")
+        
+        return StreamingResponse(stream_proxy(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
+    else:
+        try:
+            async with session.post(url, data=body, headers=forward_headers, timeout=aiohttp.ClientTimeout(total=300)) as resp:
+                content = await resp.read()
+                if check_rate_limit_in_response(content) and worker_id and model_id:
+                    asyncio.create_task(report_rate_limit(worker_id, model_id))
+                return Response(content=content, status_code=resp.status, media_type=resp.content_type)
+        except Exception as e:
+            logger.error(f"[{req_id}] Forward failed: {e}")
+            raise HTTPException(status_code=502, detail=str(e))
 
 @app.get("/health")
 async def health():
-    return {"status": "ok", "workers": workers}
+    return {"status": "ok", "workers": len(_worker_cache["workers"])}
 
 def main():
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument('--port', type=int, default=2048)
     args = parser.parse_args()
-    
-    logger.info(f"Starting Gateway on port {args.port}")
-    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="info")
+    uvicorn.run(app, host="0.0.0.0", port=args.port, log_level="warning")
 
 if __name__ == "__main__":
     main()
+
diff --git a/src/manager.py b/src/manager.py
@@ -527,6 +527,11 @@ async def check_all_ports():
     
     if config.get('stream_port_enabled'):
         ports_to_check.append({"label": "流式代理", "port": config.get('stream_port', 3120)})
+    
+    if config.get('worker_mode_enabled') and WORKER_POOL_AVAILABLE:
+        for w in worker_pool.workers.values():
+            ports_to_check.append({"label": f"Worker-{w.id} API", "port": w.port})
+            ports_to_check.append({"label": f"Worker-{w.id} Camoufox", "port": w.camoufox_port})
         
     results = []
     for item in ports_to_check:
@@ -710,6 +715,26 @@ async def save_workers_config():
     except Exception as e:
         return {"success": False, "error": str(e)}
 
+@app.get("/api/workers/next")
+async def get_next_available_worker(model: str = ""):
+    if not WORKER_POOL_AVAILABLE:
+        raise HTTPException(status_code=503, detail="Worker pool not available")
+    worker = worker_pool.get_worker_for_model(model)
+    if worker:
+        worker.request_count += 1
+        return {"port": worker.port, "worker_id": worker.id}
+    all_limited = all(w.is_model_limited(model) for w in worker_pool.workers.values() if w.status == "running")
+    if all_limited:
+        return {"error": "all_rate_limited", "message": f"All workers rate limited for model {model}"}
+    return {"error": "no_workers", "message": "No available workers"}
+
+@app.post("/api/workers/{worker_id}/rate-limit")
+async def mark_worker_rate_limited(worker_id: str, model: str = Body(..., embed=True)):
+    if not WORKER_POOL_AVAILABLE:
+        raise HTTPException(status_code=500, detail="Worker pool not available")
+    worker_pool.mark_rate_limited(worker_id, model)
+    return {"success": True}
+
 @app.post("/api/workers/{worker_id}/start")
 async def start_worker_api(worker_id: str):
     if not WORKER_POOL_AVAILABLE:
