#!/usr/bin/env python3
"""
GitLab Duo Chat → OpenAI Compatible API Proxy (v2)
====================================================
基于真实逆向分析的 GitLab Duo Chat 协议构建。
协议发现 (2026-06-17 通过浏览器网络拦截验证):
==================================================
1. GitLab Duo Chat 使用 GraphQL 端点: POST /api/graphql
2. 聊天基于 Duo Workflow 系统 (Ai::DuoWorkflows::Workflow)
3. 消息通过 GraphQL mutation 发送到工作流
4. 响应通过 getWorkflowLatestCheckpoint 查询轮询
5. 工作流状态: INPUT_REQUIRED → processing → complete
6. 消息类型: user(用户) / agent(AI助手)
7. 认证: Cookie (_gitlab_session) + CSRF Token
功能：
- /v1/chat/completions — 完全兼容 OpenAI SDK
- 支持流式响应 (SSE)
- 多模型切换
- Cookie/Token 切换账号
- 对话历史管理 (conversation_id = workflow_id)
使用方式：
pip install -r requirements.txt
python server.py
"""
import asyncio
import json
import os
import re
import sys
import time
import uuid
import logging
import secrets
from dataclasses import dataclass, field, asdict
from typing import Optional, AsyncGenerator, Dict, List, Any, Callable, Awaitable
from pathlib import Path
try:
from fastapi import FastAPI, Request, HTTPException, Header, Body, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse, StreamingResponse, HTMLResponse, FileResponse, RedirectResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
import httpx
import yaml
except ImportError as e:
print(f"Missing dependency: {e}")
print("Run: pip install -r requirements.txt")
sys.exit(1)
# Account pool + browser login (local modules)
sys.path.insert(0, str(Path(__file__).parent))
from account_pool import AccountPool, Account, SCHEDULE_STRATEGIES # noqa: E402
from browser_login import BrowserLoginManager, BrowserLoginSession # noqa: E402
from chat_driver import get_driver, close_driver # noqa: E402
from api_keys import ApiKeyManager # noqa: E402
# ============================================================
# Configuration
# ============================================================
CONFIG_PATH = Path(__file__).parent / "config.yaml"
DEFAULT_CONFIG = {
"server": {
"host": "0.0.0.0",
"port": 7860,
"debug": False,
},
"gitlab": {
"base_url": "https://gitlab.com",
"auth_type": "cookie", # cookie | token | session | oauth
"auth_value": "",
"graphql_endpoint": "/api/graphql",
"timeout": 120,
"default_model": "claude-opus-4.8",
# CSRF token (auto-fetched or manually set)
"csrf_token": "",
# Polling interval for workflow checkpoint (seconds)
"poll_interval": 1.0,
# Max polling rounds before timeout
"max_poll_rounds": 180,
},
"models": {
"claude-opus-4.8": {"id": "anthropic/claude-opus-4.8", "provider": "anthropic"},
"claude-sonnet-4": {"id": "anthropic/claude-sonnet-4", "provider": "anthropic"},
"claude-haiku-3.5": {"id": "anthropic/claude-haiku-3.5", "provider": "anthropic"},
"gpt-5.5": {"id": "openai/gpt-5.5", "provider": "openai"},
"gitlab-duo": {"id": "gitlab_duo", "provider": "gitlab"},
"duo-chat": {"id": "duo_chat", "provider": "gitlab"},
},
}
@dataclass
class AppConfig:
host: str = "0.0.0.0"
port: int = 8080
debug: bool = False
gitlab_base_url: str = "https://gitlab.com"
auth_type: str = "cookie"
auth_value: str = ""
graphql_endpoint: str = "/api/graphql"
timeout: int = 120
default_model: str = "claude-opus-4.8"
csrf_token: str = ""
poll_interval: float = 1.0
max_poll_rounds: int = 180
models: Dict[str, Dict[str, str]] = field(default_factory=dict)
# Account pool
pool_enabled: bool = True
pool_strategy: str = "round_robin"
pool_cooldown_seconds: int = 60
pool_max_failures: int = 3
pool_retry_count: int = 3
pool_invalid_on_auth_error: bool = True
# WebUI access token (auto-generated if empty; protects management UI)
webui_token: str = ""
# Allow requests without auth header to use pool (vs require Authorization)
allow_anonymous_pool: bool = True
def load_config(path: Path = CONFIG_PATH) -> AppConfig:
cfg_dict = DEFAULT_CONFIG.copy()
if path.exists():
with open(path, 'r', encoding='utf-8') as f:
user_cfg = yaml.safe_load(f) or {}
for section, values in user_cfg.items():
if section in cfg_dict and isinstance(cfg_dict[section], dict):
cfg_dict[section].update(values)
else:
cfg_dict[section] = values
env_map = {
("server", "host"): "GITLAB_PROXY_HOST",
("server", "port"): "GITLAB_PROXY_PORT",
("gitlab", "auth_type"): "GITLAB_AUTH_TYPE",
("gitlab", "auth_value"): "GITLAB_AUTH_VALUE",
("gitlab", "base_url"): "GITLAB_BASE_URL",
("gitlab", "default_model"): "GITLAB_DEFAULT_MODEL",
("gitlab", "csrf_token"): "GITLAB_CSRF_TOKEN",
}
for (section, key), env_var in env_map.items():
val = os.environ.get(env_var)
if val is not None:
if key == "port":
val = int(val)
cfg_dict[section][key] = val
sc = cfg_dict["server"]
gc = cfg_dict["gitlab"]
pool = cfg_dict.get("pool", {})
return AppConfig(
host=sc.get("host", "0.0.0.0"),
port=sc.get("port", 8080),
debug=sc.get("debug", False),
gitlab_base_url=gc.get("base_url", "https://gitlab.com"),
auth_type=gc.get("auth_type", "cookie"),
auth_value=gc.get("auth_value", ""),
graphql_endpoint=gc.get("graphql_endpoint", "/api/graphql"),
timeout=gc.get("timeout", 120),
default_model=gc.get("default_model", "claude-opus-4.8"),
csrf_token=gc.get("csrf_token", ""),
poll_interval=gc.get("poll_interval", 1.0),
max_poll_rounds=gc.get("max_poll_rounds", 180),
models=cfg_dict.get("models", DEFAULT_CONFIG["models"]),
pool_enabled=pool.get("enabled", True),
pool_strategy=pool.get("strategy", "round_robin"),
pool_cooldown_seconds=pool.get("cooldown_seconds", 60),
pool_max_failures=pool.get("max_failures", 3),
pool_retry_count=pool.get("retry_count", 3),
pool_invalid_on_auth_error=pool.get("invalid_on_auth_error", True),
webui_token = os.getenv("WEBUI_TOKEN") or pool.get("webui_token", "") or cfg_dict.get("webui_token", ""),
allow_anonymous_pool=pool.get("allow_anonymous", True),
)
# ============================================================
# Pydantic Models (OpenAI-compatible)
# ============================================================
class ChatMessage(BaseModel):
role: str
content: Optional[str] = None
name: Optional[str] = None
tool_calls: Optional[List[Dict]] = None
tool_call_id: Optional[str] = None
class ChatCompletionRequest(BaseModel):
model: str = ""
messages: List[ChatMessage] = Field(..., description="Chat messages")
stream: bool = False
temperature: Optional[float] = Field(default=None, ge=0, le=2)
max_tokens: Optional[int] = Field(default=None, ge=1)
top_p: Optional[float] = Field(default=None, ge=0, le=1)
stop: Optional[List[str]] = None
presence_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
frequency_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
user: Optional[str] = None
conversation_id: Optional[str] = Field(
default=None,
description="Existing GitLab Workflow ID (gid://gitlab/Ai::DuoWorkflows::Workflow/xxx) to continue conversation"
)
resource: Optional[str] = Field(default=None, description="GitLab resource context")
class UsageInfo(BaseModel):
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
class ChoiceMessage(BaseModel):
role: str = "assistant"
content: Optional[str] = None
tool_calls: Optional[List[Dict]] = None
class ChatCompletionChoice(BaseModel):
index: int = 0
message: Optional[ChoiceMessage] = None
delta: Optional[Dict[str, Any]] = None
finish_reason: Optional[str] = None
class ChatCompletionResponse(BaseModel):
id: str = ""
object: str = "chat.completion"
created: int = 0
model: str = ""
choices: List[ChatCompletionChoice] = []
usage: Optional[UsageInfo] = None
class ModelInfo(BaseModel):
id: str
object: str = "model"
created: int = 1700000000
owned_by: str = "gitlab-duo"
class ModelListResponse(BaseModel):
object: str = "list"
data: List[ModelInfo] = []
# ============================================================
# GitLab Duo Chat Protocol Client (v2 - Workflow Based)
# ============================================================
class GitLabDuoClientV2:
"""
GitLab Duo Chat API 客户端 v2
基于真实逆向分析的协议：
=== 协议流程 ===
1. 发送消息: GraphQL mutation → 创建/更新 Duo Workflow
2. 轮询响应: getWorkflowLatestCheckpoint 查询 → 获取消息列表
3. 流式输出: 将轮询到的增量内容转换为 SSE 格式
=== 关键发现 (2026-06-17) ===
- 端点: POST /api/graphql
- 系统: Ai::DuoWorkflows::Workflow
- 查询: getWorkflowLatestCheckpoint($workflowId)
- 消息结构: latestCheckpoint.duoMessages[]
- 请求头: x-csrf-token, x-gitlab-feature-category=duo_agent_platform
"""
def __init__(self, config: AppConfig):
self.config = config
self.base_url = config.gitlab_base_url.rstrip("/")
self.graphql_url = f"{self.base_url}{config.graphql_endpoint}"
def _build_headers(self, override_auth: Optional[str] = None) -> Dict[str, str]:
"""构建请求头（包含认证和GitLab特定头）"""
headers = {
"Content-Type": "application/json",
"Accept": "application/json, text/event-stream",
"User-Agent": "GitLab-Duo-Proxy/2.0",
"Origin": self.base_url,
"Referer": f"{self.base_url}/dashboard/home",
"X-Gitlab-Feature-Category": "duo_agent_platform",
"X-Gitlab-Version": "19.1.0-pre",
}
# CSRF token
if self.config.csrf_token:
headers["X-Csrf-Token"] = self.config.csrf_token
# Auth
auth_value = override_auth or self.config.auth_value
if self.config.auth_type == "cookie":
headers["Cookie"] = auth_value
elif self.config.auth_type == "token":
headers["PRIVATE-TOKEN"] = auth_value
headers["Authorization"] = f"Bearer {auth_value}"
elif self.config.auth_type == "session":
headers["Cookie"] = f"_gitlab_session={auth_value}"
elif self.config.auth_type == "oauth":
headers["Authorization"] = f"Bearer {auth_value}"
return headers
async def _fetch_csrf_token(self) -> str:
"""从 GitLab 页面获取 CSRF token"""
async with httpx.AsyncClient(follow_redirects=True, timeout=30) as client:
resp = await client.get(f"{self.base_url}/dashboard/home")
match = re.search(r'name="csrf-token" content="([^"]+)"', resp.text)
if match:
return match.group(1)
# Try meta tag pattern
match = re.search(r'csrf-token.*?content="([^"]+)"', resp.text)
if match:
return match.group(1)
return ""
async def _graphql_request(
self,
operation_name: str,
query: str,
variables: Dict[str, Any],
override_auth: Optional[str] = None,
) -> Dict[str, Any]:
"""执行 GraphQL 请求"""
payload = {
"operationName": operation_name,
"query": query.strip(),
"variables": variables,
}
headers = self._build_headers(override_auth)
async with httpx.AsyncClient(timeout=self.config.timeout, follow_redirects=True) as client:
resp = await client.post(self.graphql_url, json=payload, headers=headers)
result = resp.json()
if "errors" in result:
errors = [e.get("message", "Unknown") for e in result["errors"]]
raise Exception(f"GraphQL Error ({resp.status_code}): {'; '.join(errors)}")
return result
# ---- GraphQL Operations (based on reverse-engineered schema) ----
QUERY_GET_WORKFLOW_CHECKPOINT = """
query getWorkflowLatestCheckpoint($workflowId: AiDuoWorkflowsWorkflowID!) {
duoWorkflowWorkflows(workflowId: $workflowId) {
nodes {
id
status
aiCatalogItemVersionId
workflowDefinition
archived
stalled
latestCheckpoint {
workflowGoal
workflowStatus
errors
duoMessages {
content
messageType
messageSubType
status
toolInfo
timestamp
correlationId
messageId
role
additionalContext {
category
id
content
metadata
__typename
}
__typename
}
__typename
}
__typename
}
__typename
}
}
"""
MUTATION_SEND_CHAT_MESSAGE = """
mutation sendChatMessage($input: AiDuoWorkflowsSendMessageInput!) {
sendDuoChatMessage(input: $input) {
errors
workflow {
id
status
latestCheckpoint {
workflowStatus
duoMessages {
messageId
content
messageType
__typename
}
__typename
}
__typename
}
__typename
}
}
"""
MUTATION_CREATE_WORKFLOW = """
mutation createDuoWorkflow($input: CreateDuoWorkflowInput!) {
createDuoWorkflow(input: $input) {
errors
workflow {
id
status
__typename
}
__typename
}
}
"""
# Fallback: direct aiAction mutation (older/simpler API path)
MUTATION_AI_ACTION = """
mutation aiAction($question: String!, $modelId: ModelID!, $conversationId: ConversationID, $resource: AiAgentResourceInput) {
aiAction(input: { question: $question, modelId: $modelId, conversationId: $conversationId, resource: $resource }) {
errors
messageId
requestId
chatId
}
}
"""
SUBSCRIPTION_AI_RESPONSE = """
subscription aiMessageResponse($chatId: ID!, $requestId: String!) {
aiMessageResponse(chatId: $chatId, requestId: $requestId) {
... on AiMessageType {
id
role
content
timestamp
chunkId
}
... on AiErrorMessage {
message
errorCode
}
... on AiCompleteMessage {
completionReason
}
}
}
"""
async def send_message_to_workflow(
self,
prompt: str,
model_id: str,
conversation_id: Optional[str] = None,
**kwargs,
) -> Dict[str, Any]:
"""
发送聊天消息到 GitLab Duo Workflow
尝试多种 mutation 方式以兼容不同版本的 GitLab
"""
# Strategy 1: Try sendDuoChatMessage mutation (preferred for newer GitLab)
try:
result = await self._graphql_request(
operation_name="sendChatMessage",
query=self.MUTATION_SEND_CHAT_MESSAGE,
variables={
"input": {
"prompt": prompt,
"modelId": model_id,
"conversationId": conversation_id,
}
},
)
data = result.get("data", {}).get("sendDuoChatMessage", {})
if data.get("workflow"):
wf = data["workflow"]
return {
"workflow_id": wf["id"],
"status": wf["status"],
"method": "sendDuoChatMessage",
}
except Exception as e:
logging.debug(f"sendDuoChatMessage failed: {e}")
# Strategy 2: Try aiAction mutation (fallback)
try:
result = await self._graphql_request(
operation_name="aiAction",
query=self.MUTATION_AI_ACTION,
variables={
"question": prompt,
"modelId": model_id,
"conversationId": conversation_id,
},
)
data = result.get("data", {}).get("aiAction", {})
if data.get("requestId"):
return {
"request_id": data["requestId"],
"chat_id": data.get("chatId") or conversation_id,
"message_id": data.get("messageId"),
"method": "aiAction",
}
except Exception as e:
logging.debug(f"aiAction failed: {e}")
# Strategy 3: Create new workflow then poll
try:
result = await self._graphql_request(
operation_name="createDuoWorkflow",
query=self.MUTATION_CREATE_WORKFLOW,
variables={
"input": {
"goal": prompt,
"definition": "chat",
"modelId": model_id,
}
},
)
data = result.get("data", {}).get("createDuoWorkflow", {})
if data.get("workflow"):
wf = data["workflow"]
return {
"workflow_id": wf["id"],
"status": wf["status"],
"method": "createDuoWorkflow",
}
except Exception as e:
logging.debug(f"createDuoWorkflow failed: {e}")
raise Exception("All message sending strategies failed")
async def poll_workflow_response(
self,
workflow_id: str,
last_message_count: int = 0,
) -> AsyncGenerator[Dict[str, Any], None]:
"""
轮询工作流检查点，yield 新增的消息
基于 getWorkflowLatestCheckpoint 查询
"""
round_num = 0
seen_message_ids = set()
while round_num < self.config.max_poll_rounds:
round_num += 1
await asyncio.sleep(self.config.poll_interval)
try:
result = await self._graphql_request(
operation_name="getWorkflowLatestCheckpoint",
query=self.QUERY_GET_WORKFLOW_CHECKPOINT,
variables={"workflowId": workflow_id},
)
nodes = (
result
.get("data", {})
.get("duoWorkflowWorkflows", {})
.get("nodes", [])
)
if not nodes:
continue
node = nodes[0]
checkpoint = node.get("latestCheckpoint")
if not checkpoint:
continue
messages = checkpoint.get("duoMessages", [])
status = checkpoint.get("workflowStatus", "")
errors = checkpoint.get("errors", [])
# Yield new messages
new_messages = []
for msg in messages:
msg_id = msg.get("messageId", "")
if msg_id and msg_id not in seen_message_ids:
seen_message_ids.add(msg_id)
new_messages.append(msg)
for msg in new_messages:
yield {
"type": "message",
"message": msg,
"workflow_status": status,
}
# Check terminal states
if status in ("COMPLETED", "FINISHED", "FAILED", "ERROR"):
if errors:
yield {
"type": "error",
"errors": errors,
"workflow_status": status,
}
yield {
"type": "done",
"workflow_status": status,
"total_messages": len(messages),
}
return
# Also check node-level status
node_status = node.get("status", "")
if node_status in ("COMPLETED", "FINISHED", "FAILED"):
yield {
"type": "done",
"workflow_status": node_status,
"total_messages": len(messages),
}
return
except Exception as e:
yield {
"type": "poll_error",
"error": str(e),
"round": round_num,
}
# Timeout
yield {
"type": "timeout",
"rounds": round_num,
}
async def stream_chat(
self,
messages: List[ChatMessage],
model_name: str,
conversation_id: Optional[str] = None,
override_auth: Optional[str] = None,
on_started: Optional[Callable[[str], Awaitable[None]]] = None,
on_send_error: Optional[Callable[[str], Awaitable[None]]] = None,
raise_on_send_error: bool = False,
emit_initial_role: bool = True,
**kwargs,
) -> AsyncGenerator[str, None]:
"""
完整的流式聊天流程:
1. 构建 prompt
2. 发送消息到工作流
3. 轮询响应并转换为 OpenAI SSE 格式
钩子 (供账号池使用):
- on_started(workflow_id): 发送成功、开始轮询前调用
- on_send_error(error_msg): 发送阶段失败时调用
- raise_on_send_error=True 时, 发送失败直接抛出异常 (不 yield error chunk),
便于外层捕获后切换账号重试。此时 emit_initial_role 自动置为 False。
"""
if raise_on_send_error:
emit_initial_role = False
model_info = self._resolve_model(model_name)
model_id = model_info["id"]
prompt = self._build_prompt(messages)
completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
created_ts = int(time.time())
# Initial role chunk (deferred until after send succeeds when raise_on_send_error)
initial_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {"role": "assistant", "content": ""},
"finish_reason": None,
}],
}
if emit_initial_role:
yield f"data: {json.dumps(initial_chunk)}\n\n"
_send_succeeded = False
try:
# Step 1: Send message
send_result = await self.send_message_to_workflow(
prompt=prompt,
model_id=model_id,
conversation_id=conversation_id,
**kwargs,
)
workflow_id = send_result.get("workflow_id") or send_result.get("chat_id")
# If we don't have a workflow_id, we can't poll
if not workflow_id:
# For aiAction method, try subscription-style response
request_id = send_result.get("request_id")
if request_id and send_result.get("chat_id"):
if on_started:
await on_started(send_result["chat_id"])
_send_succeeded = True
if not emit_initial_role:
yield f"data: {json.dumps(initial_chunk)}\n\n"
async for chunk in self._stream_ai_action_response(
request_id, send_result["chat_id"], completion_id, created_ts, model_name
):
yield chunk
return
raise Exception(f"No workflow/chat ID in response: {send_result}")
# Send succeeded
_send_succeeded = True
if on_started:
await on_started(workflow_id)
if not emit_initial_role:
yield f"data: {json.dumps(initial_chunk)}\n\n"
# Step 2: Poll for response
full_content_parts = []
async for event in self.poll_workflow_response(workflow_id):
etype = event.get("type")
if etype == "message":
msg = event["message"]
content = msg.get("content", "")
mtype = msg.get("messageType", "")
# Only forward agent/assistant messages as content chunks
if mtype == "agent" and content:
full_content_parts.append(content)
chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": content},
"finish_reason": None,
}],
}
yield f"data: {json.dumps(chunk)}\n\n"
elif etype == "error":
err_content = "\n\n[Error] " + "; ".join(event.get("errors", []))
err_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": err_content},
"finish_reason": "error",
}],
}
yield f"data: {json.dumps(err_chunk)}\n\n"
elif etype == "done":
# Final done chunk
done_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}],
}
yield f"data: {json.dumps(done_chunk)}\n\n"
yield "data: [DONE]\n\n"
return
elif etype == "timeout":
timeout_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": "\n\n[Timeout waiting for response]"},
"finish_reason": "length",
}],
}
yield f"data: {json.dumps(timeout_chunk)}\n\n"
yield "data: [DONE]\n\n"
return
except Exception as e:
# Send-phase failure: notify pool and optionally re-raise for retry
if not _send_succeeded:
if on_send_error:
try:
await on_send_error(str(e))
except Exception:
pass
if raise_on_send_error:
raise
error_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{
"index": 0,
"delta": {"content": f"\n\n[Proxy Error] {str(e)}"},
"finish_reason": "error",
}],
}
yield f"data: {json.dumps(error_chunk)}\n\n"
yield "data: [DONE]\n\n"
async def _stream_ai_action_response(
self, request_id: str, chat_id: str,
completion_id: str, created_ts: int, model_name: str,
) -> AsyncGenerator[str, None]:
"""Fallback: use subscription-style polling for aiAction responses"""
# Poll using the subscription query as a regular query
for i in range(self.config.max_poll_rounds):
await asyncio.sleep(self.config.poll_interval)
try:
result = await self._graphql_request(
operation_name="aiMessageResponse",
query=self.SUBSCRIPTION_AI_RESPONSE,
variables={"chatId": chat_id, "requestId": request_id},
)
# Subscription via regular POST won't work well, but let's try
data = result.get("data", {}).get("aiMessageResponse")
if data:
content = data.get("content", "")
if content:
chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}],
}
yield f"data: {json.dumps(chunk)}\n\n"
reason = data.get("completionReason") or data.get("errorCode")
if reason:
done_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model_name,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
}
yield f"data: {json.dumps(done_chunk)}\n\n"
yield "data: [DONE]\n\n"
return
except Exception:
pass
yield "data: [DONE]\n\n"
def _resolve_model(self, model_name: str) -> Dict[str, str]:
if not model_name:
model_name = self.config.default_model
if model_name in self.config.models:
return self.config.models[model_name]
lower = model_name.lower()
for k, v in self.config.models.items():
if lower == k.lower():
return v
return {"id": model_name, "provider": "unknown"}
def _build_prompt(self, messages: List[ChatMessage]) -> str:
parts = []
for msg in messages:
role = msg.role.upper()
content = msg.content or ""
if role == "SYSTEM":
parts.append(f"[System Instructions]\n{content}")
elif role == "USER":
parts.append(content)
elif role == "ASSISTANT":
parts.append(f"[Previous Assistant Response]\n{content}")
elif role == "TOOL":
parts.append(f"[Tool Result]\n{content}")
return "\n\n".join(parts)
# ============================================================
# FastAPI Application
# ============================================================
app = FastAPI(
title="GitLab Duo Chat → OpenAI API Proxy v2",
version="2.0.0",
description="Convert GitLab Duo Chat (Duo Workflow) to OpenAI-compatible API. "
"Based on real protocol reverse-engineering of gitlab.com.",
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
config: Optional[AppConfig] = None
client: Optional[GitLabDuoClientV2] = None
pool: Optional[AccountPool] = None
login_mgr: Optional[BrowserLoginManager] = None
api_key_mgr: Optional[ApiKeyManager] = None
# Storage for the account pool
POOL_STORAGE_PATH = Path(__file__).parent / "accounts.json"
API_KEYS_STORAGE_PATH = Path(__file__).parent / "api_keys.json"
WEB_DIR = Path(__file__).parent / "web"
@app.on_event("startup")
async def startup():
global config, client, pool, login_mgr, api_key_mgr
config = load_config()
# Auto-fetch CSRF token if not provided
if not config.csrf_token and config.auth_type in ("cookie", "session"):
logging.info("Auto-fetching CSRF token...")
temp_client = GitLabDuoClientV2(config)
config.csrf_token = await temp_client._fetch_csrf_token()
if config.csrf_token:
logging.info(f"CSRF token fetched: {config.csrf_token[:16]}...")
else:
logging.warning("Could not auto-fetch CSRF token. Set it manually in config.yaml.")
client = GitLabDuoClientV2(config)
# Initialize account pool
pool = AccountPool(
storage_path=POOL_STORAGE_PATH,
strategy=config.pool_strategy,
cooldown_seconds=config.pool_cooldown_seconds,
max_consecutive_failures=config.pool_max_failures,
invalid_on_auth_error=config.pool_invalid_on_auth_error,
)
await pool.load()
# Initialize browser login manager
login_mgr = BrowserLoginManager(max_sessions=5, session_ttl=600)
# Initialize API key manager
api_key_mgr = ApiKeyManager(API_KEYS_STORAGE_PATH)
await api_key_mgr.load()
# Auto-generate a WebUI access token if none set
if not config.webui_token:
config.webui_token = secrets.token_urlsafe(16)
logging.info(f"[WebUI] Auto-generated access token: {config.webui_token}")
logging.basicConfig(level=logging.DEBUG if config.debug else logging.INFO)
logging.info("=" * 60)
logging.info(" GitLab Duo Proxy v2 (Workflow-Based) + Account Pool")
logging.info(f" Listening: http://{config.host}:{config.port}")
logging.info(f" WebUI: http://{config.host}:{config.port}/web")
logging.info(f" WebUI Token: {config.webui_token}")
logging.info(f" Auth Type: {config.auth_type}")
logging.info(f" Base URL: {config.gitlab_base_url}")
logging.info(f" Models: {', '.join(config.models.keys())}")
pool_cfg = await pool.get_config()
logging.info(f" Pool: enabled={config.pool_enabled} strategy={pool_cfg['strategy']} "
f"accounts={pool_cfg['total_accounts']} active={pool_cfg['active_accounts']}")
logging.info("=" * 60)
@app.get("/health")
async def health():
return {"status": "ok", "service": "gitlab-duo-proxy-v2", "version": "2.0.0"}
@app.get("/v1/models", response_model=ModelListResponse)
async def list_models():
models = [
ModelInfo(id=k, owned_by=v.get("provider", "gitlab"))
for k, v in config.models.items()
]
return ModelListResponse(data=models)
@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest, authorization: Optional[str] = Header(None)):
if not config or not client:
raise HTTPException(status_code=503, detail="Service not initialized")
req_auth = None
is_api_key = False
api_key_raw = None
if authorization:
token = authorization.removeprefix("Bearer ") if authorization.startswith("Bearer ") else authorization
if token.startswith("sk-"):
is_api_key = True
api_key_raw = token
else:
req_auth = token
# API key auth: verify key, then use pool (key tracks usage)
if is_api_key and api_key_mgr:
key_obj = await api_key_mgr.verify(api_key_raw)
if not key_obj or not key_obj.enabled:
raise HTTPException(status_code=401, detail="Invalid or revoked API key")
await api_key_mgr.report_usage(api_key_raw)
# Fall through to pool — API key just authenticates, pool handles actual sending
model = req.model or config.default_model
completion_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
created_ts = int(time.time())
# Decide auth source:
# 1. Per-request Authorization (non-API-key) → use it directly (backward compatible)
# 2. API key or anonymous → pool scheduling with retry
# 3. Fallback → config.auth_value (original behavior)
use_pool = (
config.pool_enabled
and pool is not None
and not req_auth
and (config.allow_anonymous_pool or is_api_key)
)
if use_pool:
pool_summary = await pool.get_summary()
active_in_pool = (pool_summary.get("by_status", {}) or {}).get("active", 0)
if active_in_pool == 0:
# No active account in pool → fall back to config auth if available
if config.auth_value and config.auth_value.strip() not in ("", "_gitlab_session=YOUR_SESSION_HERE; _gitlab_session_random=..."):
use_pool = False
else:
raise HTTPException(
status_code=503,
detail="账号池中没有可用账号。请先在 WebUI 通过「浏览器登录」或手动添加一个 GitLab 账号。",
)
if use_pool:
async def pool_stream():
tried: List[str] = []
last_error = "No available account"
# 把多轮 messages 合并成单个 prompt（Duo Chat UI 单消息发送）
prompt_parts = []
for m in req.messages:
role = m.role.upper()
c = m.content or ""
if role == "SYSTEM":
prompt_parts.append(f"[System]\n{c}")
elif role == "USER":
prompt_parts.append(c)
elif role == "ASSISTANT":
prompt_parts.append(f"[Assistant]\n{c}")
prompt = "\n\n".join(prompt_parts)
for attempt in range(max(1, config.pool_retry_count)):
account = await pool.acquire(exclude=tried)
if account is None:
break
tried.append(account.id)
# 优先复用 pinned 会话，否则用 cookie 创建临时浏览器会话
sess = login_mgr.get_pinned(account.id) if login_mgr else None
is_tmp = False
if sess is None:
# 创建临时会话：跳过初始登录页导航，直接设 cookie 后跳 dashboard
tmp_sid = uuid.uuid4().hex[:10]
try:
sess = await login_mgr.create(tmp_sid, base_url=config.gitlab_base_url, skip_nav=True)
is_tmp = True
except Exception as e:
last_error = f"创建浏览器会话失败: {e}"
logging.warning("[Pool] temp session create failed: %s", e)
await pool.report_failure(account.id, last_error)
continue
# 用账号 cookie 覆盖
from urllib.parse import urlparse as _up
host = _up(config.gitlab_base_url).hostname
domain = "." + ".".join(host.split(".")[-2:])
cookies_to_set = []
for pair in account.auth_value.split(";"):
pair = pair.strip()
if "=" not in pair: continue
n, _, v = pair.partition("=")
cookies_to_set.append({"name": n.strip(), "value": v.strip(),
"domain": domain, "path": "/", "httpOnly": False, "secure": True, "sameSite": "Lax"})
try:
await sess._context.add_cookies(cookies_to_set)
await sess.page.goto(config.gitlab_base_url + "/dashboard/home",
wait_until="domcontentloaded", timeout=30000)
await asyncio.sleep(2)
except Exception as e:
logging.warning("[Pool] cookie set/nav failed: %s", e)
streamed_any = False
try:
async for chunk in sess.chat_stream(
prompt=prompt, model_name=model,
):
streamed_any = True
yield chunk
await pool.report_success(account.id)
return
except Exception as e:
last_error = str(e)
logging.warning(
"[Pool] account '%s' chat failed (attempt %d/%d): %s",
account.name, attempt + 1, config.pool_retry_count, e,
)
await pool.report_failure(account.id, last_error)
if streamed_any:
yield "data: [DONE]\n\n"
return
finally:
if is_tmp:
try: await sess.close()
except Exception: pass
continue
# All retries exhausted
err_chunk = {
"id": completion_id,
"object": "chat.completion.chunk",
"created": created_ts,
"model": model,
"choices": [{
"index": 0,
"delta": {"content": f"\n\n[Proxy Error] All pool accounts failed. Last error: {last_error}"},
"finish_reason": "error",
}],
}
yield f"data: {json.dumps(err_chunk)}\n\n"
yield "data: [DONE]\n\n"
if req.stream:
return StreamingResponse(
pool_stream(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
else:
full_content = []
async for chunk_str in pool_stream():
if chunk_str.startswith("data: ") and "[DONE]" not in chunk_str:
try:
data = json.loads(chunk_str[6:])
delta = data.get("choices", [{}])[0].get("delta", {})
c = delta.get("content", "")
if c:
full_content.append(c)
except (json.JSONDecodeError, IndexError):
pass
content = "".join(full_content)
response = ChatCompletionResponse(
id=completion_id, object="chat.completion", created=created_ts, model=model,
choices=[ChatCompletionChoice(index=0,
message=ChoiceMessage(role="assistant", content=content),
finish_reason="stop" if content else "error")],
usage=UsageInfo(
prompt_tokens=sum(len(m.content or "") // 4 for m in req.messages),
completion_tokens=len(content) // 4,
total_tokens=sum(len(m.content or "") // 4 for m in req.messages) + len(content) // 4,
),
)
return JSONResponse(content=json.loads(response.model_dump_json()))
# ---- Non-pool path (per-request auth or config fallback) ----
active_client = client
if req_auth:
active_client = GitLabDuoClientV2(AppConfig(**asdict(config), auth_value=req_auth))
if req.stream:
async def generate():
async for chunk in active_client.stream_chat(
messages=req.messages,
model_name=model,
conversation_id=req.conversation_id,
override_auth=req_auth,
temperature=req.temperature,
max_tokens=req.max_tokens,
):
yield chunk
return StreamingResponse(
generate(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
else:
# Non-streaming: collect all chunks
full_content = []
async for chunk_str in active_client.stream_chat(
messages=req.messages,
model_name=model,
conversation_id=req.conversation_id,
override_auth=req_auth,
):
if chunk_str.startswith("data: ") and "[DONE]" not in chunk_str:
try:
data = json.loads(chunk_str[6:])
delta = data.get("choices", [{}])[0].get("delta", {})
c = delta.get("content", "")
if c:
full_content.append(c)
except (json.JSONDecodeError, IndexError):
pass
content = "".join(full_content)
response = ChatCompletionResponse(
id=completion_id,
object="chat.completion",
created=created_ts,
model=model,
choices=[ChatCompletionChoice(
index=0,
message=ChoiceMessage(role="assistant", content=content),
finish_reason="stop",
)],
usage=UsageInfo(
prompt_tokens=sum(len(m.content or "") // 4 for m in req.messages),
completion_tokens=len(content) // 4,
total_tokens=sum(len(m.content or "") // 4 for m in req.messages) + len(content) // 4,
),
)
return JSONResponse(content=json.loads(response.model_dump_json()))
@app.post("/v1/accounts/switch")
async def switch_account(request: Request):
global config, client
body = await request.json()
auth_type = body.get("auth_type", "cookie")
auth_value = body.get("auth_value", "")
if not auth_value:
raise HTTPException(status_code=400, detail="auth_value required")
config.auth_type = auth_type
config.auth_value = auth_value
client = GitLabDuoClientV2(config)
return {
"status": "ok",
"message": f"Account switched (auth_type={auth_type})",
"auth_preview": auth_value[:20] + "..." if len(auth_value) > 20 else auth_value,
}
@app.get("/v1/accounts/info")
async def account_info():
if not config:
raise HTTPException(status_code=503, detail="Service not initialized")
val = config.auth_value
preview = val[:8] + "..." + val[-4:] if len(val) > 12 else "***"
return {
"auth_type": config.auth_type,
"auth_value_preview": preview,
"base_url": config.gitlab_base_url,
"default_model": config.default_model,
"csrf_token_set": bool(config.csrf_token),
"available_models": list(config.models.keys()),
"protocol_version": "v2-workflow",
}
# ============================================================
# Account Pool Management API
# ============================================================
def _check_webui_token(request: Request) -> bool:
"""Verify WebUI management token from header or query."""
if not config:
return False
token = (
request.headers.get("x-webui-token")
or request.query_params.get("token")
or ""
)
return secrets.compare_digest(token, config.webui_token)
async def _require_webui(request: Request):
if not _check_webui_token(request):
raise HTTPException(status_code=401, detail="Invalid or missing WebUI token")
return True
@app.get("/v1/accounts/pool")
async def pool_list(request: Request):
"""List all accounts in the pool."""
await _require_webui(request)
return {"accounts": await pool.list_all(mask=True)}
@app.get("/v1/accounts/pool/summary")
async def pool_summary(request: Request):
"""Pool-level statistics."""
await _require_webui(request)
return await pool.get_summary()
@app.get("/v1/accounts/pool/config")
async def pool_get_config(request: Request):
await _require_webui(request)
return await pool.get_config()
@app.put("/v1/accounts/pool/config")
async def pool_put_config(request: Request):
await _require_webui(request)
body = await request.json()
await pool.set_config(
cooldown_seconds=body.get("cooldown_seconds"),
max_consecutive_failures=body.get("max_consecutive_failures"),
invalid_on_auth_error=body.get("invalid_on_auth_error"),
)
if body.get("strategy"):
await pool.set_strategy(body["strategy"])
return {"status": "ok", "config": await pool.get_config()}
@app.post("/v1/accounts/pool")
async def pool_add(request: Request):
"""Add a new account to the pool."""
await _require_webui(request)
body = await request.json()
name = (body.get("name") or "").strip()
auth_type = (body.get("auth_type") or "cookie").strip()
auth_value = (body.get("auth_value") or "").strip()
if not name or not auth_value:
raise HTTPException(status_code=400, detail="name and auth_value are required")
if auth_type not in ("cookie", "token", "session", "oauth"):
raise HTTPException(status_code=400, detail="invalid auth_type")
acc = await pool.add(
name=name, auth_type=auth_type, auth_value=auth_value,
note=body.get("note", ""), enabled=body.get("enabled", True),
)
return {"status": "ok", "account": acc.to_dict(mask=True)}
@app.get("/v1/accounts/pool/{account_id}")
async def pool_get(account_id: str, request: Request):
await _require_webui(request)
acc = await pool.get(account_id)
if not acc:
raise HTTPException(status_code=404, detail="account not found")
return {"account": acc.to_dict(mask=True)}
@app.put("/v1/accounts/pool/{account_id}")
async def pool_update(account_id: str, request: Request):
await _require_webui(request)
body = await request.json()
acc = await pool.update(account_id, **body)
if not acc:
raise HTTPException(status_code=404, detail="account not found")
return {"status": "ok", "account": acc.to_dict(mask=True)}
@app.delete("/v1/accounts/pool/{account_id}")
async def pool_delete(account_id: str, request: Request):
await _require_webui(request)
ok = await pool.delete(account_id)
if not ok:
raise HTTPException(status_code=404, detail="account not found")
return {"status": "ok", "deleted": account_id}
@app.post("/v1/accounts/pool/{account_id}/reset")
async def pool_reset(account_id: str, request: Request):
"""Reset an account to active status (clear cooldown/invalid)."""
await _require_webui(request)
acc = await pool.reset_status(account_id)
if not acc:
raise HTTPException(status_code=404, detail="account not found")
return {"status": "ok", "account": acc.to_dict(mask=True)}
@app.post("/v1/accounts/pool/{account_id}/test")
async def pool_test(account_id: str, request: Request):
"""
Test an account by calling GitLab /api/v4/user.
Returns 200 with user info on success, marks account invalid on auth failure.
"""
await _require_webui(request)
acc = await pool.get(account_id)
if not acc:
raise HTTPException(status_code=404, detail="account not found")
base = config.gitlab_base_url.rstrip("/")
headers = {"User-Agent": "GitLab-Duo-Proxy/2.0"}
if acc.auth_type == "cookie":
headers["Cookie"] = acc.auth_value
elif acc.auth_type == "token":
headers["PRIVATE-TOKEN"] = acc.auth_value
elif acc.auth_type == "session":
headers["Cookie"] = f"_gitlab_session={acc.auth_value}"
elif acc.auth_type == "oauth":
headers["Authorization"] = f"Bearer {acc.auth_value}"
try:
async with httpx.AsyncClient(timeout=30, follow_redirects=True) as http:
resp = await http.get(f"{base}/api/v4/user", headers=headers)
if resp.status_code == 200:
user = resp.json()
await pool.report_success(acc.id)
return {
"status": "ok",
"account_id": acc.id,
"user": {
"id": user.get("id"),
"username": user.get("username"),
"name": user.get("name"),
"email": user.get("email"),
"state": user.get("state"),
},
}
else:
err = f"HTTP {resp.status_code}"
await pool.report_failure(acc.id, err)
return JSONResponse(
status_code=200,
content={"status": "fail", "account_id": acc.id,
"http_code": resp.status_code,
"detail": resp.text[:300]},
)
except Exception as e:
await pool.report_failure(acc.id, str(e))
return JSONResponse(
status_code=200,
content={"status": "error", "account_id": acc.id, "detail": str(e)},
)
@app.get("/v1/accounts/pool/token")
async def pool_get_token(request: Request):
"""Return the current WebUI token (used by the UI to bootstrap)."""
await _require_webui(request)
return {"token": config.webui_token, "pool_enabled": config.pool_enabled}
# ============================================================
# Browser Login Assistant (Playwright + WebSocket 串流)
# ============================================================
@app.post("/v1/accounts/pool/assist/create")
async def assist_create(request: Request):
"""启动一个新的浏览器登录会话，返回 sid。"""
await _require_webui(request)
if login_mgr is None:
raise HTTPException(status_code=503, detail="login manager not initialized")
sid = uuid.uuid4().hex[:12]
try:
# 会话在 WebSocket 连接时真正启动；这里仅预注册 sid 并检查 playwright 可用性
import playwright # noqa: F401
except ImportError:
raise HTTPException(
status_code=503,
detail="playwright not installed on server. Run: pip install playwright && playwright install chromium",
)
return {"sid": sid, "base_url": config.gitlab_base_url}
@app.websocket("/v1/accounts/pool/assist/ws")
async def assist_ws(ws: WebSocket, token: str = ""):
"""浏览器登录串流 WebSocket。
前端消息:
{type: "start"} 启动浏览器
{type: "click", x, y} 点击
{type: "type", text} 输入文本
{type: "key", key} 按键 (Enter/Tab/Backspace/...)
{type: "scroll", dx, dy} 滚动
{type: "goto", url} 导航
{type: "reload"} 刷新
{type: "close"} 关闭
后端消息:
{type: "ready", sid, viewport}
{type: "frame", data, url, title, logged_in, status}
{type: "logged_in", cookie_preview}
{type: "error", message}
"""
# 鉴权 (query param)
if not config or not secrets.compare_digest(token, config.webui_token):
await ws.close(code=4401)
return
await ws.accept()
sid = uuid.uuid4().hex[:12]
sess: Optional[BrowserLoginSession] = None
push_task: Optional[asyncio.Task] = None
logged_in_notified = False
async def on_logged_in(cookie_str: str) -> None:
nonlocal logged_in_notified
if not logged_in_notified:
logged_in_notified = True
preview = cookie_str[:32] + "..." if len(cookie_str) > 32 else cookie_str
try:
await ws.send_json({"type": "logged_in", "cookie_preview": preview})
except Exception:
pass
try:
# 等待前端 start 指令
first = await ws.receive_json()
if first.get("type") != "start":
await ws.send_json({"type": "error", "message": "expected start first"})
await ws.close()
return
sess = await login_mgr.create(sid, base_url=config.gitlab_base_url, on_logged_in=on_logged_in)
await ws.send_json({
"type": "ready",
"sid": sid,
"viewport": list(sess.viewport),
"base_url": config.gitlab_base_url,
"status": sess.status,
"error": sess.error,
})
async def push_frames():
while not sess._closed:
frame = await sess.get_frame_b64()
if frame:
await ws.send_json({
"type": "frame",
"data": frame,
"url": sess.current_url,
"title": sess.title,
"logged_in": sess.logged_in,
"status": sess.status,
})
await asyncio.sleep(0.3)
push_task = asyncio.create_task(push_frames())
while True:
msg = await ws.receive_json()
mtype = msg.get("type")
if sess._closed:
break
if mtype == "click":
await sess.click(int(msg.get("x", 0)), int(msg.get("y", 0)))
elif mtype == "type":
await sess.type_text(msg.get("text", ""))
elif mtype == "key":
await sess.press_key(msg.get("key", ""))
elif mtype == "scroll":
await sess.scroll(int(msg.get("dx", 0)), int(msg.get("dy", 0)))
elif mtype == "goto":
await sess.goto(msg.get("url", ""))
elif mtype == "reload":
await sess.reload()
elif mtype == "login":
# 用 httpx 直接 POST 登录（绕过 CF），成功后 cookie 注入浏览器
result = await sess.login_via_httpx(
username=msg.get("username", ""),
password=msg.get("password", ""),
)
await ws.send_json({"type": "login_result", **result})
elif mtype == "close":
break
except WebSocketDisconnect:
pass
except Exception as e:
logging.exception("assist ws error")
try:
await ws.send_json({"type": "error", "message": str(e)})
except Exception:
pass
finally:
if push_task:
push_task.cancel()
if sess and not sess.pinned_account_id:
await login_mgr.close(sid)
try:
await ws.close()
except Exception:
pass
@app.post("/v1/accounts/pool/assist/{sid}/save")
async def assist_save(sid: str, request: Request):
"""从指定登录会话抓取 Cookie 并保存为新账号。"""
await _require_webui(request)
sess = login_mgr.get(sid) if login_mgr else None
if not sess:
raise HTTPException(status_code=404, detail="session not found or expired")
if not sess.logged_in:
# 兜底：再检查一次
await sess._check_login()
if not sess.logged_in:
raise HTTPException(status_code=400, detail="not logged in yet")
body = await request.json()
name = (body.get("name") or "").strip()
if not name:
raise HTTPException(status_code=400, detail="name required")
cookie_str = await sess.get_cookies_str()
if not cookie_str or "_gitlab_session" not in cookie_str:
raise HTTPException(status_code=400, detail="no valid gitlab session cookie found")
note = body.get("note", "").strip() or f"browser login · {sess.current_url}"
acc = await pool.add(
name=name, auth_type="cookie", auth_value=cookie_str, note=note,
)
# 把已登录会话钉住给该账号聊天用（Cloudflare 已过，复用同一浏览器上下文）
login_mgr.pin_for_account(acc.id, sess)
logging.info("[assist] session %s pinned for account %s (%s)", sid, acc.id, name)
return {"status": "ok", "account": acc.to_dict(mask=True), "pinned": True}
@app.on_event("shutdown")
async def shutdown():
if login_mgr:
await login_mgr.close_all()
await close_driver()
# ============================================================
# API Key Management
# ============================================================
@app.get("/v1/api-keys")
async def api_keys_list(request: Request):
"""列出所有 API 密钥。"""
await _require_webui(request)
return {"keys": await api_key_mgr.list_all_full()}
@app.post("/v1/api-keys")
async def api_keys_create(request: Request):
"""生成新的 API 密钥。返回原始密钥（仅此一次可见）。"""
await _require_webui(request)
body = await request.json()
name = (body.get("name") or "").strip()
if not name:
raise HTTPException(status_code=400, detail="name required")
raw_key = await api_key_mgr.create(name=name, note=body.get("note", ""))
return {"status": "ok", "key": raw_key, "message": "请立即复制保存，关闭后无法再次查看完整密钥"}
@app.delete("/v1/api-keys/{key_id}")
async def api_keys_revoke(key_id: str, request: Request):
"""吊销（禁用）API 密钥。"""
await _require_webui(request)
ok = await api_key_mgr.revoke(key_id)
if not ok:
raise HTTPException(status_code=404, detail="key not found")
return {"status": "ok", "revoked": key_id}
@app.put("/v1/api-keys/{key_id}")
async def api_keys_rename(key_id: str, request: Request):
await _require_webui(request)
body = await request.json()
name = (body.get("name") or "").strip()
if not name:
raise HTTPException(status_code=400, detail="name required")
ok = await api_key_mgr.rename(key_id, name)
if not ok:
raise HTTPException(status_code=404, detail="key not found")
return {"status": "ok", "renamed": key_id}
# ============================================================
# ============================================================
# 真实浏览器登录代理 (反向代理 gitlab.com 捕获 Cookie)
# ============================================================
# 待捕获的 session (sid -> captured cookies)
_auth_sessions: Dict[str, Dict] = {}
_auth_lock = asyncio.Lock()
AUTH_HTML = """
GitLab 登录授权
正在连接 GitLab，请登录…
"""
async def _proxy_fetch(url: str, method: str, headers: Dict, body: bytes = b"") -> httpx.Response:
proxy_headers = {k: v for k, v in headers.items()
if k.lower() not in ("host", "accept-encoding")}
proxy_headers["accept-encoding"] = "identity"
async with httpx.AsyncClient(timeout=30, follow_redirects=False, verify=False) as cl:
return await cl.request(method, url, headers=proxy_headers, content=body or None)
@app.api_route("/auth/proxy/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def auth_proxy(request: Request, path: str):
"""反向代理 gitlab.com。用 cookie 里的 auth_sid 鉴权。"""
sid = request.cookies.get("auth_sid", "")
if not sid:
raise HTTPException(status_code=401)
async with _auth_lock:
if sid not in _auth_sessions:
raise HTTPException(status_code=401, detail="session not found")
qs = f"?{request.url.query}" if request.url.query else ""
target = f"{config.gitlab_base_url.rstrip('/')}/{path}{qs}"
body = await request.body()
resp = await _proxy_fetch(target, request.method, dict(request.headers), body)
set_cookies = resp.headers.get_list("set-cookie")
if set_cookies:
async with _auth_lock:
session = _auth_sessions.setdefault(sid, {"cookies": {}, "logged": False, "url": ""})
for sc in set_cookies:
for part in sc.split(","):
part = part.strip()
if "=" in part:
n, _, v = part.partition("=")
n = n.strip(); v = v.split(";")[0].strip()
session["cookies"][n] = v
if resp.status_code in (302, 303) and "location" in resp.headers:
loc = resp.headers["location"]
if "/dashboard" in loc and "/sign_in" not in loc:
session["logged"] = True; session["url"] = loc
content_type = resp.headers.get("content-type", "")
if "text/html" in content_type and resp.status_code == 200:
html = resp.text
proxy_base = f"{request.url.scheme}://{request.url.netloc}"
base_tag = f''
if "" in html:
html = html.replace("", f"\n{base_tag}", 1)
elif "\n{base_tag}', 1)
return HTMLResponse(html, status_code=resp.status_code, headers=dict(resp.headers))
return HTMLResponse(resp.content, status_code=resp.status_code, headers=dict(resp.headers))
@app.get("/auth/start")
async def auth_start(token: str = "", request: Request = None):
if not secrets.compare_digest(token, config.webui_token):
raise HTTPException(status_code=401)
sid = uuid.uuid4().hex[:12]
# 创建空 session 并设 cookie 方便后续代理请求鉴权
async with _auth_lock:
_auth_sessions[sid] = {"cookies": {}, "logged": False, "url": ""}
# 使用请求里的实际 host（而非 config 里的 0.0.0.0）
proxy_base = f"{request.url.scheme}://{request.url.netloc}" if request else f"http://{config.host}:{config.port}"
html = AUTH_HTML.replace("__SID__", sid).replace("__TOKEN__", token).replace("__PROXY_BASE__", proxy_base)
resp = HTMLResponse(html)
resp.set_cookie("auth_sid", sid, path="/", samesite="lax")
return resp
@app.get("/auth/check")
async def auth_check(sid: str, token: str = ""):
if not secrets.compare_digest(token, config.webui_token):
return {"logged_in": False}
async with _auth_lock:
s = _auth_sessions.get(sid, {})
return {"logged_in": s.get("logged", False)}
@app.post("/auth/save")
async def auth_save(sid: str, token: str = "", request: Request = None):
if not secrets.compare_digest(token, config.webui_token):
raise HTTPException(401)
body = await request.json()
name = (body.get("name") or "").strip()
if not name:
raise HTTPException(400, detail="name required")
async with _auth_lock:
s = _auth_sessions.pop(sid, {})
cookies = s.get("cookies", {})
cookie_str = "; ".join(f"{n}={v}" for n, v in cookies.items())
if "_gitlab_session" not in cookie_str:
return {"status": "error", "error": "未检测到 _gitlab_session cookie"}
acc = await pool.add(name=name, auth_type="cookie", auth_value=cookie_str,
note=f"proxy auth - {s.get('url','')}")
return {"status": "ok", "account": acc.to_dict(mask=True)}
# WebUI (Claude-style)
# ============================================================
@app.get("/web", response_class=HTMLResponse)
@app.get("/web/", response_class=HTMLResponse)
async def webui_index():
index = WEB_DIR / "index.html"
if not index.exists():
return HTMLResponse("WebUI not built
web/index.html missing
", status_code=404)
return HTMLResponse(index.read_text(encoding="utf-8"))
if WEB_DIR.exists():
app.mount("/web/static", StaticFiles(directory=str(WEB_DIR)), name="web-static")
if __name__ == "__main__":
import uvicorn
cfg = load_config()
uvicorn.run(app, host=cfg.host, port=cfg.port)