feat(account): 支持账号级别拦截预热请求

- 新增 intercept_warmup_requests 配置项，存储在 credentials 字段 - 启用后，标题生成、Warmup 等预热请求返回 mock 响应，不消耗上游 token - 前端支持所有账号类型（OAuth、Setup Token、API Key）的开关配置 - 修复 OAuth 凭证刷新时丢失非 token 配置的问题
2025-12-19 16:39:25 +08:00
parent ee86dbca9d
commit d1f0902ec0
7 changed files with 220 additions and 10 deletions
--- a/backend/internal/handler/admin/account_handler.go
+++ b/backend/internal/handler/admin/account_handler.go
@@ -241,16 +241,20 @@ func (h *AccountHandler) Refresh(c *gin.Context) {
 		return
 	}

-	// Update account credentials
-	newCredentials := map[string]interface{}{
-		"access_token":  tokenInfo.AccessToken,
-		"token_type":    tokenInfo.TokenType,
-		"expires_in":    tokenInfo.ExpiresIn,
-		"expires_at":    tokenInfo.ExpiresAt,
-		"refresh_token": tokenInfo.RefreshToken,
-		"scope":         tokenInfo.Scope,
+	// Copy existing credentials to preserve non-token settings (e.g., intercept_warmup_requests)
+	newCredentials := make(map[string]interface{})
+	for k, v := range account.Credentials {
+		newCredentials[k] = v
 	}

+	// Update token-related fields
+	newCredentials["access_token"] = tokenInfo.AccessToken
+	newCredentials["token_type"] = tokenInfo.TokenType
+	newCredentials["expires_in"] = tokenInfo.ExpiresIn
+	newCredentials["expires_at"] = tokenInfo.ExpiresAt
+	newCredentials["refresh_token"] = tokenInfo.RefreshToken
+	newCredentials["scope"] = tokenInfo.Scope
+
 	updatedAccount, err := h.adminService.UpdateAccount(c.Request.Context(), accountID, &service.UpdateAccountInput{
 		Credentials: newCredentials,
 	})
--- a/backend/internal/handler/gateway_handler.go
+++ b/backend/internal/handler/gateway_handler.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"log"
 	"net/http"
+	"strings"
 	"time"

 	"sub2api/internal/middleware"
@@ -127,6 +128,16 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
 		return
 	}

+	// 检查预热请求拦截（在账号选择后、转发前检查）
+	if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
+		if req.Stream {
+			sendMockWarmupStream(c, req.Model)
+		} else {
+			sendMockWarmupResponse(c, req.Model)
+		}
+		return
+	}
+
 	// 3. 获取账号并发槽位
 	accountReleaseFunc, err := h.acquireAccountSlotWithWait(c, account, req.Stream, &streamStarted)
 	if err != nil {
@@ -489,3 +500,89 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
 		return
 	}
 }
+
+// isWarmupRequest 检测是否为预热请求（标题生成、Warmup等）
+func isWarmupRequest(body []byte) bool {
+	// 快速检查：如果body不包含关键字，直接返回false
+	bodyStr := string(body)
+	if !strings.Contains(bodyStr, "title") && !strings.Contains(bodyStr, "Warmup") {
+		return false
+	}
+
+	// 解析完整请求
+	var req struct {
+		Messages []struct {
+			Content []struct {
+				Type string `json:"type"`
+				Text string `json:"text"`
+			} `json:"content"`
+		} `json:"messages"`
+		System []struct {
+			Text string `json:"text"`
+		} `json:"system"`
+	}
+	if err := json.Unmarshal(body, &req); err != nil {
+		return false
+	}
+
+	// 检查 messages 中的标题提示模式
+	for _, msg := range req.Messages {
+		for _, content := range msg.Content {
+			if content.Type == "text" {
+				if strings.Contains(content.Text, "Please write a 5-10 word title for the following conversation:") ||
+					content.Text == "Warmup" {
+					return true
+				}
+			}
+		}
+	}
+
+	// 检查 system 中的标题提取模式
+	for _, system := range req.System {
+		if strings.Contains(system.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
+			return true
+		}
+	}
+
+	return false
+}
+
+// sendMockWarmupStream 发送流式 mock 响应（用于预热请求拦截）
+func sendMockWarmupStream(c *gin.Context, model string) {
+	c.Header("Content-Type", "text/event-stream")
+	c.Header("Cache-Control", "no-cache")
+	c.Header("Connection", "keep-alive")
+	c.Header("X-Accel-Buffering", "no")
+
+	events := []string{
+		`event: message_start` + "\n" + `data: {"message":{"content":[],"id":"msg_mock_warmup","model":"` + model + `","role":"assistant","stop_reason":null,"stop_sequence":null,"type":"message","usage":{"input_tokens":10,"output_tokens":0}},"type":"message_start"}`,
+		`event: content_block_start` + "\n" + `data: {"content_block":{"text":"","type":"text"},"index":0,"type":"content_block_start"}`,
+		`event: content_block_delta` + "\n" + `data: {"delta":{"text":"New","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
+		`event: content_block_delta` + "\n" + `data: {"delta":{"text":" Conversation","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
+		`event: content_block_stop` + "\n" + `data: {"index":0,"type":"content_block_stop"}`,
+		`event: message_delta` + "\n" + `data: {"delta":{"stop_reason":"end_turn","stop_sequence":null},"type":"message_delta","usage":{"input_tokens":10,"output_tokens":2}}`,
+		`event: message_stop` + "\n" + `data: {"type":"message_stop"}`,
+	}
+
+	for _, event := range events {
+		_, _ = c.Writer.WriteString(event + "\n\n")
+		c.Writer.Flush()
+		time.Sleep(20 * time.Millisecond)
+	}
+}
+
+// sendMockWarmupResponse 发送非流式 mock 响应（用于预热请求拦截）
+func sendMockWarmupResponse(c *gin.Context, model string) {
+	c.JSON(http.StatusOK, gin.H{
+		"id":      "msg_mock_warmup",
+		"type":    "message",
+		"role":    "assistant",
+		"model":   model,
+		"content": []gin.H{{"type": "text", "text": "New Conversation"}},
+		"stop_reason": "end_turn",
+		"usage": gin.H{
+			"input_tokens":  10,
+			"output_tokens": 2,
+		},
+	})
+}
--- a/backend/internal/model/account.go
+++ b/backend/internal/model/account.go
@@ -263,3 +263,17 @@ func (a *Account) ShouldHandleErrorCode(statusCode int) bool {
 	}
 	return false
 }
+
+// IsInterceptWarmupEnabled 检查是否启用预热请求拦截
+// 启用后，标题生成、Warmup等预热请求将返回mock响应，不消耗上游token
+func (a *Account) IsInterceptWarmupEnabled() bool {
+	if a.Credentials == nil {
+		return false
+	}
+	if v, ok := a.Credentials["intercept_warmup_requests"]; ok {
+		if enabled, ok := v.(bool); ok {
+			return enabled
+		}
+	}
+	return false
+}
--- a/frontend/src/components/account/CreateAccountModal.vue
+++ b/frontend/src/components/account/CreateAccountModal.vue
@@ -418,6 +418,31 @@
        </div>
      </div>

+      <!-- Intercept Warmup Requests (all account types) -->
+      <div class="border-t border-gray-200 dark:border-dark-600 pt-4">
+        <div class="flex items-center justify-between">
+          <div>
+            <label class="input-label mb-0">{{ t('admin.accounts.interceptWarmupRequests') }}</label>
+            <p class="text-xs text-gray-500 dark:text-gray-400 mt-1">{{ t('admin.accounts.interceptWarmupRequestsDesc') }}</p>
+          </div>
+          <button
+            type="button"
+            @click="interceptWarmupRequests = !interceptWarmupRequests"
+            :class="[
+              'relative inline-flex h-6 w-11 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-primary-500 focus:ring-offset-2',
+              interceptWarmupRequests ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
+            ]"
+          >
+            <span
+              :class="[
+                'pointer-events-none inline-block h-5 w-5 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
+                interceptWarmupRequests ? 'translate-x-5' : 'translate-x-0'
+              ]"
+            />
+          </button>
+        </div>
+      </div>
+
      <div>
        <label class="input-label">{{ t('admin.accounts.proxy') }}</label>
        <ProxySelector
@@ -590,6 +615,7 @@ const allowedModels = ref<string[]>([])
 const customErrorCodesEnabled = ref(false)
 const selectedErrorCodes = ref<number[]>([])
 const customErrorCodeInput = ref<number | null>(null)
+const interceptWarmupRequests = ref(false)

 // Common models for whitelist
 const commonModels = [
@@ -758,6 +784,7 @@ const resetForm = () => {
  customErrorCodesEnabled.value = false
  selectedErrorCodes.value = []
  customErrorCodeInput.value = null
+  interceptWarmupRequests.value = false
  oauth.resetState()
  oauthFlowRef.value?.reset()
 }
@@ -801,6 +828,11 @@ const handleSubmit = async () => {
    credentials.custom_error_codes = [...selectedErrorCodes.value]
  }

+  // Add intercept warmup requests setting
+  if (interceptWarmupRequests.value) {
+    credentials.intercept_warmup_requests = true
+  }
+
  form.credentials = credentials

  submitting.value = true
@@ -847,11 +879,17 @@ const handleExchangeCode = async () => {

    const extra = oauth.buildExtraInfo(tokenInfo)

+    // Merge interceptWarmupRequests into credentials
+    const credentials = {
+      ...tokenInfo,
+      ...(interceptWarmupRequests.value ? { intercept_warmup_requests: true } : {})
+    }
+
    await adminAPI.accounts.create({
      name: form.name,
      platform: form.platform,
      type: addMethod.value, // Use addMethod as type: 'oauth' or 'setup-token'
-      credentials: tokenInfo,
+      credentials,
      extra,
      proxy_id: form.proxy_id,
      concurrency: form.concurrency,
@@ -901,11 +939,17 @@ const handleCookieAuth = async (sessionKey: string) => {
        const extra = oauth.buildExtraInfo(tokenInfo)
        const accountName = keys.length > 1 ? `${form.name} #${i + 1}` : form.name

+        // Merge interceptWarmupRequests into credentials
+        const credentials = {
+          ...tokenInfo,
+          ...(interceptWarmupRequests.value ? { intercept_warmup_requests: true } : {})
+        }
+
        await adminAPI.accounts.create({
          name: accountName,
          platform: form.platform,
          type: addMethod.value, // Use addMethod as type: 'oauth' or 'setup-token'
-          credentials: tokenInfo,
+          credentials,
          extra,
          proxy_id: form.proxy_id,
          concurrency: form.concurrency,
--- a/frontend/src/components/account/EditAccountModal.vue
+++ b/frontend/src/components/account/EditAccountModal.vue
@@ -286,6 +286,31 @@
        </div>
      </div>

+      <!-- Intercept Warmup Requests (all account types) -->
+      <div class="border-t border-gray-200 dark:border-dark-600 pt-4">
+        <div class="flex items-center justify-between">
+          <div>
+            <label class="input-label mb-0">{{ t('admin.accounts.interceptWarmupRequests') }}</label>
+            <p class="text-xs text-gray-500 dark:text-gray-400 mt-1">{{ t('admin.accounts.interceptWarmupRequestsDesc') }}</p>
+          </div>
+          <button
+            type="button"
+            @click="interceptWarmupRequests = !interceptWarmupRequests"
+            :class="[
+              'relative inline-flex h-6 w-11 flex-shrink-0 cursor-pointer rounded-full border-2 border-transparent transition-colors duration-200 ease-in-out focus:outline-none focus:ring-2 focus:ring-primary-500 focus:ring-offset-2',
+              interceptWarmupRequests ? 'bg-primary-600' : 'bg-gray-200 dark:bg-dark-600'
+            ]"
+          >
+            <span
+              :class="[
+                'pointer-events-none inline-block h-5 w-5 transform rounded-full bg-white shadow ring-0 transition duration-200 ease-in-out',
+                interceptWarmupRequests ? 'translate-x-5' : 'translate-x-0'
+              ]"
+            />
+          </button>
+        </div>
+      </div>
+
      <div>
        <label class="input-label">{{ t('admin.accounts.proxy') }}</label>
        <ProxySelector
@@ -401,6 +426,7 @@ const allowedModels = ref<string[]>([])
 const customErrorCodesEnabled = ref(false)
 const selectedErrorCodes = ref<number[]>([])
 const customErrorCodeInput = ref<number | null>(null)
+const interceptWarmupRequests = ref(false)

 // Common models for whitelist
 const commonModels = [
@@ -459,6 +485,10 @@ watch(() => props.account, (newAccount) => {
    form.status = newAccount.status as 'active' | 'inactive'
    form.group_ids = newAccount.group_ids || []

+    // Load intercept warmup requests setting (applies to all account types)
+    const credentials = newAccount.credentials as Record<string, unknown> | undefined
+    interceptWarmupRequests.value = credentials?.intercept_warmup_requests === true
+
    // Initialize API Key fields for apikey type
    if (newAccount.type === 'apikey' && newAccount.credentials) {
      const credentials = newAccount.credentials as Record<string, unknown>
@@ -630,6 +660,23 @@ const handleSubmit = async () => {
        newCredentials.custom_error_codes = [...selectedErrorCodes.value]
      }

+      // Add intercept warmup requests setting
+      if (interceptWarmupRequests.value) {
+        newCredentials.intercept_warmup_requests = true
+      }
+
+      updatePayload.credentials = newCredentials
+    } else {
+      // For oauth/setup-token types, only update intercept_warmup_requests if changed
+      const currentCredentials = props.account.credentials as Record<string, unknown> || {}
+      const newCredentials: Record<string, unknown> = { ...currentCredentials }
+
+      if (interceptWarmupRequests.value) {
+        newCredentials.intercept_warmup_requests = true
+      } else {
+        delete newCredentials.intercept_warmup_requests
+      }
+
      updatePayload.credentials = newCredentials
    }

--- a/frontend/src/i18n/locales/en.ts
+++ b/frontend/src/i18n/locales/en.ts
@@ -696,6 +696,8 @@ export default {
      enterErrorCode: 'Enter error code (100-599)',
      invalidErrorCode: 'Please enter a valid HTTP error code (100-599)',
      errorCodeExists: 'This error code is already selected',
+      interceptWarmupRequests: 'Intercept Warmup Requests',
+      interceptWarmupRequestsDesc: 'When enabled, warmup requests like title generation will return mock responses without consuming upstream tokens',
      proxy: 'Proxy',
      noProxy: 'No Proxy',
      concurrency: 'Concurrency',
--- a/frontend/src/i18n/locales/zh.ts
+++ b/frontend/src/i18n/locales/zh.ts
@@ -786,6 +786,8 @@ export default {
      enterErrorCode: '输入错误码 (100-599)',
      invalidErrorCode: '请输入有效的 HTTP 错误码 (100-599)',
      errorCodeExists: '该错误码已被选中',
+      interceptWarmupRequests: '拦截预热请求',
+      interceptWarmupRequestsDesc: '启用后，标题生成等预热请求将返回 mock 响应，不消耗上游 token',
      proxy: '代理',
      noProxy: '无代理',
      concurrency: '并发数',