Fix PR 133 merge integration issues

Author: genz27

cat.jpg

@@ -4,14 +4,12 @@ services:
   flow2api:
     image: ghcr.io/thesmallhancat/flow2api:latest
     container_name: flow2api
-    network_mode: host
+    ports:
+      - "38000:8000"
     volumes:
       - ./data:/app/data
       - ./tmp:/app/tmp
       - ./config/setting.toml:/app/config/setting.toml
-      - ./src/services/proxy_manager.py:/app/src/services/proxy_manager.py
-      - ./src/services/flow_client.py:/app/src/services/flow_client.py
-      - ./src/api/admin.py:/app/src/api/admin.py
     environment:
       - PYTHONUNBUFFERED=1
     restart: unless-stopped

cat_watercolor.jpg

@@ -4,6 +4,7 @@ let heartbeatInterval = null;
 
 const DEFAULT_SETTINGS = {
     serverUrl: "ws://127.0.0.1:8000/captcha_ws",
+    apiKey: "",
     routeKey: "",
     clientLabel: ""
 };
@@ -13,6 +14,7 @@ function getSettings() {
         chrome.storage.local.get(DEFAULT_SETTINGS, (stored) => {
             resolve({
                 serverUrl: (stored.serverUrl || DEFAULT_SETTINGS.serverUrl).trim(),
+                apiKey: (stored.apiKey || "").trim(),
                 routeKey: (stored.routeKey || "").trim(),
                 clientLabel: (stored.clientLabel || "").trim()
             });
@@ -74,6 +76,9 @@ async function connectWS() {
 
     const settings = await getSettings();
     const url = new URL(settings.serverUrl || DEFAULT_SETTINGS.serverUrl);
+    if (settings.apiKey) {
+        url.searchParams.set("key", settings.apiKey);
+    }
     if (settings.routeKey) {
         url.searchParams.set("route_key", settings.routeKey);
     }
@@ -227,7 +232,7 @@ async function handleGetToken(data) {
 
 chrome.storage.onChanged.addListener((changes, areaName) => {
     if (areaName !== "local") return;
-    if (changes.routeKey || changes.serverUrl || changes.clientLabel) {
+    if (changes.routeKey || changes.serverUrl || changes.apiKey || changes.clientLabel) {
         console.log("[Flow2API] Extension settings changed, reconnecting WebSocket...");
         closeSocket();
         connectWS();

docker-compose.yml

@@ -9,7 +9,7 @@
     "scripting"
   ],
   "host_permissions": [
-    "<all_urls>"
+    "https://labs.google/*"
   ],
   "background": {
     "service_worker": "background.js"

extension/background.js

@@ -87,9 +87,12 @@ <h1>Flow2API Captcha Worker</h1>
       <label for="serverUrl">WebSocket URL</label>
       <input id="serverUrl" type="text" placeholder="ws://127.0.0.1:8000/captcha_ws">
 
+      <label for="apiKey">Flow2API API Key</label>
+      <input id="apiKey" type="password" placeholder="管理后台里的 API Key">
+
       <button id="saveBtn" type="button">保存</button>
       <div class="status" id="status"></div>
-      <div class="hint">管理台 token 里填写同样的 <code>Route Key</code>，比如 9223 对 9223。</div>
+      <div class="hint">管理台 token 里填写同样的 <code>Route Key</code>，比如 9223 对 9223。后端会使用 API Key 校验 WebSocket 连接。</div>
     </div>
   </div>
   <script src="options.js"></script>

extension/manifest 2.json

@@ -1,5 +1,6 @@
 const DEFAULT_SETTINGS = {
   serverUrl: "ws://127.0.0.1:8000/captcha_ws",
+  apiKey: "",
   routeKey: "",
   clientLabel: ""
 };
@@ -9,6 +10,7 @@ const $ = (id) => document.getElementById(id);
 function normalizeSettings(values) {
   return {
     serverUrl: (values.serverUrl || DEFAULT_SETTINGS.serverUrl).trim(),
+    apiKey: (values.apiKey || "").trim(),
     routeKey: (values.routeKey || "").trim(),
     clientLabel: (values.clientLabel || "").trim()
   };
@@ -33,6 +35,7 @@ function loadSettings() {
   chrome.storage.local.get(DEFAULT_SETTINGS, (stored) => {
     const settings = normalizeSettings(stored);
     $("serverUrl").value = settings.serverUrl;
+    $("apiKey").value = settings.apiKey;
     $("routeKey").value = settings.routeKey;
     $("clientLabel").value = settings.clientLabel;
   });
@@ -41,6 +44,7 @@ function loadSettings() {
 function saveSettings() {
   const settings = normalizeSettings({
     serverUrl: $("serverUrl").value,
+    apiKey: $("apiKey").value,
     routeKey: $("routeKey").value,
     clientLabel: $("clientLabel").value
   });
@@ -49,6 +53,10 @@ function saveSettings() {
     setStatus("WebSocket URL 必须以 ws:// 或 wss:// 开头。", true);
     return;
   }
+  if (!settings.apiKey) {
+    setStatus("请填写 Flow2API API Key。", true);
+    return;
+  }
 
   chrome.storage.local.set(settings, () => {
     if (chrome.runtime.lastError) {

extension/manifest.json

@@ -9,10 +9,10 @@
 from urllib.parse import urlparse
 
 from curl_cffi.requests import AsyncSession
-from fastapi import APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect
+from fastapi import APIRouter, Depends, HTTPException, Query, Request, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse, StreamingResponse
 
-from ..core.auth import verify_api_key_flexible
+from ..core.auth import AuthManager, verify_api_key_flexible
 from ..core.logger import debug_logger
 from ..core.model_resolver import get_base_model_aliases, resolve_model_name
 from ..core.models import (
@@ -23,7 +23,6 @@
 )
 from ..services.generation_handler import MODEL_CONFIG, GenerationHandler
 from ..services.browser_captcha_extension import ExtensionCaptchaService
-from fastapi import WebSocket, WebSocketDisconnect
 
 router = APIRouter()
 
@@ -486,6 +485,7 @@ async def _collect_non_stream_result(
     model: str,
     prompt: str,
     images: List[bytes],
+    base_url_override: Optional[str] = None,
     video_media_id: Optional[str] = None,
 ) -> str:
     handler = _ensure_generation_handler()
@@ -495,6 +495,7 @@ async def _collect_non_stream_result(
         prompt=prompt,
         images=images if images else None,
         stream=False,
+        base_url_override=base_url_override,
         video_media_id=video_media_id,
     ):
         result = chunk
@@ -723,6 +724,7 @@ async def _iterate_openai_stream(
         prompt=normalized.prompt,
         images=normalized.images if normalized.images else None,
         stream=True,
+        base_url_override=base_url_override,
         video_media_id=normalized.video_media_id,
     ):
         if chunk.startswith("data: "):
@@ -746,6 +748,7 @@ async def _iterate_gemini_stream(
         prompt=normalized.prompt,
         images=normalized.images if normalized.images else None,
         stream=True,
+        base_url_override=base_url_override,
         video_media_id=normalized.video_media_id,
     ):
         if chunk.startswith("data: "):
@@ -874,6 +877,7 @@ async def create_chat_completion(
                 normalized.model,
                 normalized.prompt,
                 normalized.images,
+                base_url_override=request_base_url,
                 video_media_id=normalized.video_media_id,
             )
         )
@@ -907,7 +911,8 @@ async def generate_content(
                     normalized.model,
                     normalized.prompt,
                     normalized.images,
-                    request_base_url,
+                    base_url_override=request_base_url,
+                    video_media_id=normalized.video_media_id,
                 )
             )
         )
@@ -970,6 +975,20 @@ async def stream_generate_content(
 @router.websocket("/captcha_ws")
 async def captcha_websocket_endpoint(websocket: WebSocket):
     from ..core.logger import debug_logger
+    api_key = (
+        websocket.query_params.get("key")
+        or websocket.query_params.get("api_key")
+        or websocket.headers.get("x-goog-api-key")
+        or ""
+    ).strip()
+    authorization = (websocket.headers.get("authorization") or "").strip()
+    if authorization.lower().startswith("bearer "):
+        api_key = authorization[7:].strip()
+
+    if not api_key or not AuthManager.verify_api_key(api_key):
+        await websocket.close(code=1008)
+        return
+
     service = await ExtensionCaptchaService.get_instance()
     await service.connect(websocket)
     try: