From 85ff0073fb90296db78bc33f7a627c8bd6794970 Mon Sep 17 00:00:00 2001 From: "agermel@foxmail.com" Date: Thu, 27 Aug 2026 02:33:23 +0800 Subject: [PATCH] =?UTF-8?q?feat(account):=20=E6=96=B0=E5=A2=9E=E8=B4=A6?= =?UTF-8?q?=E5=8F=B7=E6=B5=8B=E8=AF=95=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=B5=8B?= =?UTF-8?q?=E8=AF=95=E8=AF=A5=E8=B4=A6=E5=8F=B7=E7=9A=84=E6=96=87=E5=AD=97?= =?UTF-8?q?/=E5=9B=BE=E7=89=87=E6=A8=A1=E5=9E=8B=E6=98=AF=E5=90=A6?= =?UTF-8?q?=E5=8F=AF=E7=94=A8=EF=BC=8C=E4=B8=8E=E5=87=AD=E8=AF=81=E5=8F=AF?= =?UTF-8?q?=E7=94=A8=E6=80=A7=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增「测试账号」入口：对账号发一次真实上游请求，验证其使用文字/图片模型的凭据是否可用，通过 SSE / Tauri 事件把进度与结果实时回传到弹窗。 - 图片测试走 image_generation 工具（顶层文字模型 + gpt-image-2 工具），而不是把 gpt-image-2 当顶层主模型直连——后者会被上游以「ChatGPT 账号 不支持该模型」拒绝。 - 依据所选模型的能力（supports_image_generation）自动修正测试类型，并把kind 字段端到端打通（修复 Tauri 命令静默丢弃 kind 的问题）。 - 依据测试结果更新账号状态：401/403 → 不可用，429 → 限流，成功 → 恢复active（手动 disabled/inactive 不自动恢复）。 --- apps/src-tauri/src/commands/account/remote.rs | 53 + apps/src-tauri/src/commands/registry.rs | 2 + apps/src-tauri/src/lib.rs | 7 + apps/src/app/accounts/accounts-page-view.tsx | 20 + apps/src/app/accounts/page.tsx | 34 + .../components/modals/account-test-modal.tsx | 501 +++++++ apps/src/components/ui/select.tsx | 4 +- apps/src/hooks/useAccounts.ts | 4 + apps/src/lib/api/account-client.ts | 28 + apps/src/lib/api/account-maintenance.ts | 14 + apps/src/lib/api/account-test-events.ts | 93 ++ .../lib/api/transport-web-commands/account.ts | 2 + crates/service/src/account/account_status.rs | 76 ++ crates/service/src/account/account_test.rs | 1152 +++++++++++++++++ crates/service/src/account/account_warmup.rs | 12 +- crates/service/src/account/mod.rs | 2 + .../src/gateway/core/runtime_config.rs | 64 + crates/service/src/gateway/mod.rs | 4 + .../service/src/http/account_test_events.rs | 184 +++ crates/service/src/http/backend_router.rs | 7 + crates/service/src/http/mod.rs | 1 + crates/service/src/http/proxy_runtime.rs | 4 + crates/service/src/lib.rs | 2 + crates/service/src/models_v2/mod.rs | 9 + crates/service/src/rpc_dispatch/account.rs | 16 +- crates/service/src/rpc_dispatch/mod.rs | 2 + 26 files changed, 2288 insertions(+), 9 deletions(-) create mode 100644 apps/src/components/modals/account-test-modal.tsx create mode 100644 apps/src/lib/api/account-test-events.ts create mode 100644 crates/service/src/account/account_test.rs create mode 100644 crates/service/src/http/account_test_events.rs diff --git a/apps/src-tauri/src/commands/account/remote.rs b/apps/src-tauri/src/commands/account/remote.rs index ffbee7657..8a05f7ff8 100644 --- a/apps/src-tauri/src/commands/account/remote.rs +++ b/apps/src-tauri/src/commands/account/remote.rs @@ -276,6 +276,59 @@ pub async fn service_account_warmup( rpc_call_in_background("account/warmup", addr, Some(params)).await } +/// 函数 `service_account_test_start` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - addr: 参数 addr +/// - account_id: 参数 account_id +/// - model: 参数 model +/// - prompt: 参数 prompt +/// - kind: 参数 kind +/// +/// # 返回 +/// 返回函数执行结果 +#[tauri::command] +pub async fn service_account_test_start( + addr: Option, + account_id: String, + model: Option, + prompt: Option, + kind: Option, +) -> Result { + let params = serde_json::json!({ + "accountId": account_id, + "model": model, + "prompt": prompt, + "kind": kind, + }); + rpc_call_in_background("account/test", addr, Some(params)).await +} + +/// 函数 `service_account_test_cancel` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - addr: 参数 addr +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回函数执行结果 +#[tauri::command] +pub async fn service_account_test_cancel( + addr: Option, + account_id: String, +) -> Result { + let params = serde_json::json!({ "accountId": account_id }); + rpc_call_in_background("account/test/cancel", addr, Some(params)).await +} + #[tauri::command] pub async fn service_account_proxy_get( addr: Option, diff --git a/apps/src-tauri/src/commands/registry.rs b/apps/src-tauri/src/commands/registry.rs index 2613b9b6d..04fe976d3 100644 --- a/apps/src-tauri/src/commands/registry.rs +++ b/apps/src-tauri/src/commands/registry.rs @@ -128,6 +128,8 @@ macro_rules! invoke_handler { crate::commands::account::remote::service_account_update, crate::commands::account::remote::service_account_update_sorts, crate::commands::account::remote::service_account_warmup, + crate::commands::account::remote::service_account_test_start, + crate::commands::account::remote::service_account_test_cancel, crate::commands::account::remote::service_account_proxy_get, crate::commands::account::remote::service_account_proxy_set, crate::commands::account::remote::service_account_proxy_clear, diff --git a/apps/src-tauri/src/lib.rs b/apps/src-tauri/src/lib.rs index 23913eecc..fbfc92cfc 100644 --- a/apps/src-tauri/src/lib.rs +++ b/apps/src-tauri/src/lib.rs @@ -19,6 +19,7 @@ use app_shell::{ }; const USAGE_REFRESH_COMPLETED_EVENT: &str = "usage-refresh-completed"; +const ACCOUNT_TEST_EVENT: &str = "account-test-event"; #[cfg(target_os = "linux")] const AYATANA_APPINDICATOR_LOG_DOMAIN: &str = "libayatana-appindicator"; #[cfg(target_os = "linux")] @@ -209,6 +210,12 @@ pub fn run() { log::warn!("emit usage refresh completed event failed: {}", err); } }); + let account_test_event_app = app.handle().clone(); + codexmanager_service::set_account_test_event_handler(move |event| { + if let Err(err) = account_test_event_app.emit(ACCOUNT_TEST_EVENT, &event) { + log::warn!("emit account test event failed: {}", err); + } + }); if let Err(err) = setup_tray(app.handle()) { TRAY_AVAILABLE.store(false, std::sync::atomic::Ordering::Relaxed); CLOSE_TO_TRAY_ON_CLOSE.store(false, std::sync::atomic::Ordering::Relaxed); diff --git a/apps/src/app/accounts/accounts-page-view.tsx b/apps/src/app/accounts/accounts-page-view.tsx index 7dd2d5192..0881b2d56 100644 --- a/apps/src/app/accounts/accounts-page-view.tsx +++ b/apps/src/app/accounts/accounts-page-view.tsx @@ -29,6 +29,7 @@ import { import { AddAccountModal } from "@/components/modals/add-account-modal"; import { AccountResetCreditControl } from "@/components/account-reset-credit-control"; import { ConfirmDialog } from "@/components/modals/confirm-dialog"; +import { AccountTestModal } from "@/components/modals/account-test-modal"; import UsageModal from "@/components/modals/usage-modal"; import { Button, buttonVariants } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; @@ -160,6 +161,10 @@ export interface AccountsPageViewProps { proxySourceDraft: AccountProxySource; proxyProfileIdDraft: string; proxyUrlDraft: string; + accountTestAccount: Account | null; + openAccountTest: (account: Account) => void; + handleAccountTestOpenChange: (open: boolean) => void; + onAccountTestFinished: (accountId: string) => void; selectedAccount: Account | null; accountEditorState: AccountEditorState | null; deleteDialogState: DeleteDialogState; @@ -375,6 +380,7 @@ export function AccountsPageView(props: AccountsPageViewProps) { importByFile, importByDirectory, refreshAccount, + onAccountTestFinished, clearPreferredAccount, setPreferredAccount, toggleAccountStatus, @@ -548,6 +554,14 @@ export function AccountsPageView(props: AccountsPageViewProps) { {t("账号代理")} + props.openAccountTest(account)} + > + + {t("测试账号")} + + { diff --git a/apps/src/app/accounts/page.tsx b/apps/src/app/accounts/page.tsx index a7d0d4609..f7abdad4d 100644 --- a/apps/src/app/accounts/page.tsx +++ b/apps/src/app/accounts/page.tsx @@ -75,6 +75,7 @@ export default function AccountsPage() { refreshAllAccountRt, refreshAllAccounts, refreshAccountList, + refreshAccountsSilently, deleteAccount, deleteManyAccounts, cleanupAccountsByStatuses, @@ -141,6 +142,19 @@ export default function AccountsPage() { useState("custom"); const [proxyProfileIdDraft, setProxyProfileIdDraft] = useState(""); const [proxyUrlDraft, setProxyUrlDraft] = useState(""); + const [accountTestAccountId, setAccountTestAccountId] = useState( + null, + ); + const [accountTestAccountSnapshot, setAccountTestAccountSnapshot] = + useState(null); + // 从最新账号列表派生弹窗里的账号，测试结束后状态徽章可自动刷新； + // 列表短暂重取时回退到快照，避免弹窗闪烁关闭。 + const accountTestAccount = useMemo( + () => + accounts.find((account) => account.id === accountTestAccountId) ?? + accountTestAccountSnapshot, + [accounts, accountTestAccountId, accountTestAccountSnapshot], + ); const [accountEditorState, setAccountEditorState] = useState(null); @@ -556,6 +570,22 @@ const toggleCleanupStatus = (rawStatus: string) => { setProxyUrlDraft(""); }; + const openAccountTest = (account: Account) => { + setAccountTestAccountId(account.id); + setAccountTestAccountSnapshot(account); + }; + + const handleAccountTestOpenChange = (open: boolean) => { + if (open) return; + setAccountTestAccountId(null); + setAccountTestAccountSnapshot(null); + }; + + // 测试结束后静默刷新账号状态（不弹「账号用量已刷新」），让弹窗徽章与列表同步。 + const handleAccountTestFinished = () => { + void refreshAccountsSilently(); + }; + const handleTestProxySettings = async () => { if (!proxyDialogAccount) return; try { @@ -877,6 +907,7 @@ const toggleCleanupStatus = (rawStatus: string) => { proxyDialogAccount={proxyDialogAccount} proxySettings={proxySettings} proxyProfiles={proxyProfiles} + accountTestAccount={accountTestAccount} isProxySettingsLoading={isProxySettingsLoading} proxyEnabledDraft={proxyEnabledDraft} proxySourceDraft={proxySourceDraft} @@ -949,6 +980,9 @@ const toggleCleanupStatus = (rawStatus: string) => { handleDeleteSingle={handleDeleteSingle} openProxyDialog={openProxyDialog} handleProxyDialogOpenChange={handleProxyDialogOpenChange} + openAccountTest={openAccountTest} + handleAccountTestOpenChange={handleAccountTestOpenChange} + onAccountTestFinished={handleAccountTestFinished} handleSaveProxySettings={handleSaveProxySettings} handleClearProxySettings={handleClearProxySettings} handleTestProxySettings={handleTestProxySettings} diff --git a/apps/src/components/modals/account-test-modal.tsx b/apps/src/components/modals/account-test-modal.tsx new file mode 100644 index 000000000..60eba67a0 --- /dev/null +++ b/apps/src/components/modals/account-test-modal.tsx @@ -0,0 +1,501 @@ +"use client"; + +import { useCallback, useEffect, useRef, useState } from "react"; +import { CheckCircle2, Loader2, XCircle } from "lucide-react"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Button } from "@/components/ui/button"; +import { + Select, + SelectContent, + SelectGroup, + SelectItem, + SelectLabel, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; +import { accountClient } from "@/lib/api/account-client"; +import { managedModelsV2Client } from "@/lib/api/managed-models-v2"; +import { + listenAccountTestEvent, + type AccountTestEventPayload, +} from "@/lib/api/account-test-events"; +import { AccountStatusCell } from "@/app/accounts/accounts-page-helpers"; +import type { ManagedModelV2 } from "@/types/model-v2"; +import type { Account } from "@/types"; + +interface AccountTestModalProps { + account: Account | null; + open: boolean; + onOpenChange: (open: boolean) => void; + onFinished?: (accountId: string) => void; +} + +interface TestImage { + url: string; + mimeType: string; +} + +interface Accumulated { + text: string; + images: TestImage[]; + model?: string; + status?: string; + success?: boolean; + error?: string; +} + +type Phase = "idle" | "running" | "done"; + +function isImageModel(model: ManagedModelV2): boolean { + const caps = (model.capabilities ?? {}) as Record; + return ( + caps.supports_image_generation === true || + caps.supportsImageGeneration === true + ); +} + +function isManuallyDisabled(account: Account | null): boolean { + const status = String(account?.status ?? "").trim().toLowerCase(); + return status === "disabled" || status === "inactive"; +} + +function modelLabel(model: ManagedModelV2): string { + const name = model.displayName?.trim(); + return name || model.slug; +} + +export function AccountTestModal({ + account, + open, + onOpenChange, + onFinished, +}: AccountTestModalProps) { + const [phase, setPhase] = useState("idle"); + const [state, setState] = useState({ text: "", images: [] }); + const [models, setModels] = useState([]); + const [selectedModel, setSelectedModel] = useState(null); + const [testKind, setTestKind] = useState<"text" | "image">("text"); + const [canceled, setCanceled] = useState(false); + + // 测试类型只决定发哪种请求（文字直连 / 图片工具），不干预模型列表的选择。 + const handleTestKindChange = (value: string | null) => { + const nextKind = value === "image" ? "image" : "text"; + setTestKind(nextKind); + // 类型切换后保持模型一致：当前模型不符合新类型时，自动选一个匹配的模型。 + const current = models.find((item) => item.slug === selectedModel); + if (!current || isImageModel(current) !== (nextKind === "image")) { + const match = models.find( + (item) => isImageModel(item) === (nextKind === "image"), + ); + setSelectedModel(match?.slug ?? null); + } + }; + + const testIdRef = useRef(null); + const finishedRef = useRef(false); + const unlistenRef = useRef<(() => void) | null>(null); + const phaseRef = useRef("idle"); + const accountIdRef = useRef(account?.id ?? null); + const onFinishedRef = useRef(onFinished); + const terminalRef = useRef(null); + + const accountId = account?.id ?? null; + accountIdRef.current = accountId; + onFinishedRef.current = onFinished; + + useEffect(() => { + phaseRef.current = phase; + }, [phase]); + + useEffect(() => { + const el = terminalRef.current; + if (el) { + el.scrollTop = el.scrollHeight; + } + }, [state.text, state.status, phase, state.images.length]); + + const handleEvent = useCallback((payload: AccountTestEventPayload) => { + const currentId = testIdRef.current; + if (currentId && payload.testId && payload.testId !== currentId) { + return; + } + if (finishedRef.current) { + return; + } + switch (payload.type) { + case "test_start": + setState((prev) => ({ ...prev, model: payload.model ?? prev.model })); + break; + case "content": + setState((prev) => ({ ...prev, text: prev.text + (payload.text ?? "") })); + break; + case "image": { + const imageUrl = payload.imageUrl; + if (imageUrl) { + setState((prev) => ({ + ...prev, + images: [ + ...prev.images, + { url: imageUrl, mimeType: payload.mimeType ?? "image/png" }, + ], + })); + } + break; + } + case "status": + setState((prev) => ({ ...prev, status: payload.status ?? prev.status })); + break; + case "test_complete": + setState((prev) => ({ ...prev, success: payload.success ?? true })); + setPhase("done"); + finishedRef.current = true; + if (accountIdRef.current) { + onFinishedRef.current?.(accountIdRef.current); + } + break; + case "error": + setState((prev) => ({ + ...prev, + error: payload.error ?? "测试失败", + })); + setPhase("done"); + finishedRef.current = true; + if (accountIdRef.current) { + onFinishedRef.current?.(accountIdRef.current); + } + break; + } + }, []); + + const startTest = useCallback(async () => { + const id = accountIdRef.current; + if (!id || phaseRef.current === "running") { + return; + } + unlistenRef.current?.(); + unlistenRef.current = null; + testIdRef.current = null; + finishedRef.current = false; + setState({ text: "", images: [] }); + setCanceled(false); + setPhase("running"); + + try { + const unlisten = await listenAccountTestEvent(handleEvent); + unlistenRef.current = unlisten; + const result = await accountClient.testAccount({ + accountId: id, + model: selectedModel ?? undefined, + kind: testKind, + }); + testIdRef.current = result.testId ?? null; + setState((prev) => ({ ...prev, model: result.model ?? prev.model })); + } catch (err) { + unlistenRef.current?.(); + unlistenRef.current = null; + setState((prev) => ({ + ...prev, + error: err instanceof Error ? err.message : "启动测试失败", + })); + finishedRef.current = true; + setPhase("done"); + if (accountIdRef.current) { + onFinishedRef.current?.(accountIdRef.current); + } + } + }, [handleEvent, selectedModel, testKind]); + + const cancelTest = useCallback(() => { + const id = accountIdRef.current; + if (id) { + void accountClient.cancelAccountTest(id).catch(() => {}); + } + unlistenRef.current?.(); + unlistenRef.current = null; + testIdRef.current = null; + finishedRef.current = true; + setState({ text: "", images: [] }); + setCanceled(true); + setPhase("idle"); + }, []); + + const handleOpenChange = useCallback( + (nextOpen: boolean) => { + if (!nextOpen && phaseRef.current === "running") { + const id = accountIdRef.current; + if (id) { + void accountClient.cancelAccountTest(id).catch(() => {}); + } + } + onOpenChange(nextOpen); + }, + [onOpenChange], + ); + + useEffect(() => { + if (!open || !accountId) { + return; + } + + setPhase("idle"); + setState({ text: "", images: [] }); + setCanceled(false); + testIdRef.current = null; + finishedRef.current = false; + unlistenRef.current?.(); + unlistenRef.current = null; + setModels([]); + setSelectedModel(null); + setTestKind("text"); + + let disposed = false; + (async () => { + try { + const result = await managedModelsV2Client.list(true); + if (disposed) { + return; + } + const enabled = result.items.filter((model) => model.enabled); + setModels(enabled); + const textModel = enabled.find((model) => !isImageModel(model)); + setSelectedModel((textModel ?? enabled[0])?.slug ?? null); + } catch { + // 模型列表加载失败不阻塞测试，后端会用默认文字模型兜底。 + } + })(); + + return () => { + disposed = true; + unlistenRef.current?.(); + unlistenRef.current = null; + }; + }, [open, accountId]); + + const { text, images, status, success, error } = state; + + // 按来源分组展示：官方内置模型与自定义模型分开，方便识别哪些是官方目录、哪些可增删。 + const builtinModels = models.filter((model) => model.origin === "builtin"); + const customModels = models.filter((model) => model.origin !== "builtin"); + + return ( + + + + 测试账号 + + {account?.name || account?.label || accountId} + + + +
+ {account ? ( +
+ +
+ ) : null} + +
+ + +
+ +
+ + + {models.length === 0 ? ( + + 未加载到可用模型，测试将使用后端默认模型。 + + ) : null} +
+ +
+ {phase === "idle" ? ( +
+ + {canceled + ? "已取消测试，可再次点击「开始测试」。" + : "准备就绪，点击「开始测试」发起一次真实请求。"} + +
+ ) : ( + <> + {state.model ? ( +
模型：{state.model}
+ ) : null} + {status ? ( +
+ {phase === "running" ? ( + + ) : null} + {status} +
+ ) : null} + {text ? ( +
+ {text} + {phase === "running" ? ( + _ + ) : null} +
+ ) : null} + {phase === "done" ? ( + <> +
+ {success ? ( + + ) : ( + + )} + {success ? "测试成功" : error || "测试失败"} +
+ {success && isManuallyDisabled(account) ? ( +
+ 该账号为手动禁用，测试虽成功但不会被自动恢复为「可用」。 +
+ ) : null} + + ) : null} + + )} +
+ + {images.length > 0 ? ( +
+ + 图片预览 + +
+ {images.map((image, index) => ( + // eslint-disable-next-line @next/next/no-img-element + {`test-result-${index + ))} +
+
+ ) : null} +
+ + + {phase === "running" ? ( + + ) : null} + {phase === "done" ? ( + + ) : null} + {phase === "idle" ? ( + + ) : null} + + + + + ); +} diff --git a/apps/src/components/ui/select.tsx b/apps/src/components/ui/select.tsx index 2523b2c78..24b457107 100644 --- a/apps/src/components/ui/select.tsx +++ b/apps/src/components/ui/select.tsx @@ -115,7 +115,9 @@ function SelectContent({ sideOffset = 4, align = "center", alignOffset = 0, - alignItemWithTrigger = true, + // 默认向下展开（锚定触发器），而不是让选中项与触发器对齐：后者在长列表里会把整个 + // 弹出层往上顶、盖住触发器上方的控件。 + alignItemWithTrigger = false, ...props }: SelectPrimitive.Popup.Props & Pick< diff --git a/apps/src/hooks/useAccounts.ts b/apps/src/hooks/useAccounts.ts index 682dbe961..a65c95410 100644 --- a/apps/src/hooks/useAccounts.ts +++ b/apps/src/hooks/useAccounts.ts @@ -1178,6 +1178,10 @@ export function useAccounts() { await invalidateAccountData(); toast.success(t("账号列表已刷新")); }, + // 静默刷新账号数据（不弹 toast）：测试结束后只回读最新账号状态，避免误触「用量刷新」提示。 + refreshAccountsSilently: async () => { + await invalidateAccountData(); + }, deleteAccount: (accountId: string) => { if (!ensureServiceReady("删除账号")) return; deleteMutation.mutate(accountId); diff --git a/apps/src/lib/api/account-client.ts b/apps/src/lib/api/account-client.ts index a47476bff..256152f19 100644 --- a/apps/src/lib/api/account-client.ts +++ b/apps/src/lib/api/account-client.ts @@ -46,11 +46,13 @@ import { import { AccountExportResult, AccountImportResult, + AccountTestStartResult, AccountWarmupResult, DeleteAccountsByStatusesResult, DeleteUnavailableFreeResult, readAccountExportResult, readAccountImportResult, + readAccountTestStartResult, readAccountWarmupResult, readDeleteAccountsByStatusesResult, readApiKeySecret, @@ -91,6 +93,13 @@ export interface AccountWarmupPayload { message?: string; } +export interface AccountTestPayload { + accountId: string; + model?: string; + prompt?: string; + kind?: "text" | "image"; +} + export interface AccountProxyLatencyTestPayload { accountId: string; } @@ -527,6 +536,25 @@ export const accountClient = { }), ), ), + testAccount: async (params: AccountTestPayload): Promise => + readAccountTestStartResult( + await invoke( + "service_account_test_start", + withAddr({ + accountId: params.accountId, + model: params.model ?? null, + prompt: params.prompt ?? null, + kind: params.kind ?? "text", + }), + ), + ), + cancelAccountTest: async (accountId: string): Promise => + Boolean( + await invoke( + "service_account_test_cancel", + withAddr({ accountId }), + ), + ), getProxySettings: async ( accountId: string, diff --git a/apps/src/lib/api/account-maintenance.ts b/apps/src/lib/api/account-maintenance.ts index 17df2ef1d..ea68fa8b7 100644 --- a/apps/src/lib/api/account-maintenance.ts +++ b/apps/src/lib/api/account-maintenance.ts @@ -92,6 +92,12 @@ export interface AccountWarmupResult { results?: AccountWarmupItemResult[]; } +export interface AccountTestStartResult { + testId?: string; + started?: boolean; + model?: string; +} + export function readAccountImportResult(payload: unknown): AccountImportResult { const source = asRecord(payload); const hasUsageRefreshAccountIds = @@ -189,6 +195,14 @@ export function readAccountWarmupResult(payload: unknown): AccountWarmupResult { }; } +export function readAccountTestStartResult(payload: unknown): AccountTestStartResult { + return { + testId: readStringField(payload, "testId"), + started: readBooleanField(payload, "started"), + model: readStringField(payload, "model"), + }; +} + export function readApiKeySecret(payload: unknown): string { return readStringField(payload, "key"); } diff --git a/apps/src/lib/api/account-test-events.ts b/apps/src/lib/api/account-test-events.ts new file mode 100644 index 000000000..41dcc5264 --- /dev/null +++ b/apps/src/lib/api/account-test-events.ts @@ -0,0 +1,93 @@ +import { isTauriRuntime } from "./transport"; + +export const ACCOUNT_TEST_EVENT = "account-test-event"; + +export interface AccountTestEventPayload { + testId?: string; + type?: string; + text?: string; + model?: string; + status?: string; + imageUrl?: string; + mimeType?: string; + success?: boolean; + error?: string; +} + +export type AccountTestEventHandler = (payload: AccountTestEventPayload) => void; + +type Unlisten = () => void; + +function readAccountTestEventPayload(event: Event): AccountTestEventPayload { + if (event instanceof CustomEvent && typeof event.detail === "object" && event.detail) { + return event.detail as AccountTestEventPayload; + } + return {}; +} + +function readAccountTestMessagePayload(event: MessageEvent): AccountTestEventPayload { + if (typeof event.data !== "string" || !event.data.trim()) { + return {}; + } + try { + const payload = JSON.parse(event.data); + return typeof payload === "object" && payload + ? (payload as AccountTestEventPayload) + : {}; + } catch { + return {}; + } +} + +export async function listenAccountTestEvent( + handler: AccountTestEventHandler +): Promise { + if (typeof window === "undefined") { + return () => {}; + } + + const handleWindowEvent = (event: Event) => { + handler(readAccountTestEventPayload(event)); + }; + window.addEventListener(ACCOUNT_TEST_EVENT, handleWindowEvent); + + let eventSource: EventSource | null = null; + let handleEventSourceEvent: ((event: MessageEvent) => void) | null = null; + if ( + !isTauriRuntime() && + typeof EventSource !== "undefined" && + window.location.protocol.startsWith("http") + ) { + eventSource = new EventSource("/api/events/account-test"); + handleEventSourceEvent = (event: MessageEvent) => { + handler(readAccountTestMessagePayload(event)); + }; + eventSource.addEventListener( + ACCOUNT_TEST_EVENT, + handleEventSourceEvent as EventListener + ); + } + + let unlistenTauri: Unlisten | null = null; + if (isTauriRuntime()) { + const { listen } = await import("@tauri-apps/api/event"); + unlistenTauri = await listen( + ACCOUNT_TEST_EVENT, + (event) => { + handler(event.payload || {}); + }, + ); + } + + return () => { + window.removeEventListener(ACCOUNT_TEST_EVENT, handleWindowEvent); + if (eventSource && handleEventSourceEvent) { + eventSource.removeEventListener( + ACCOUNT_TEST_EVENT, + handleEventSourceEvent as EventListener + ); + } + eventSource?.close(); + unlistenTauri?.(); + }; +} diff --git a/apps/src/lib/api/transport-web-commands/account.ts b/apps/src/lib/api/transport-web-commands/account.ts index 71f8a66c5..4b6b56147 100644 --- a/apps/src/lib/api/transport-web-commands/account.ts +++ b/apps/src/lib/api/transport-web-commands/account.ts @@ -18,6 +18,8 @@ export function createAccountWebCommands(postWebRpc: WebRpcCaller): Record exportAccountsViaBrowser(postWebRpc, asRecord(params), options), }, service_account_warmup: { rpcMethod: "account/warmup" }, + service_account_test_start: { rpcMethod: "account/test" }, + service_account_test_cancel: { rpcMethod: "account/test/cancel" }, service_account_proxy_get: { rpcMethod: "account/proxy/get" }, service_account_proxy_set: { rpcMethod: "account/proxy/set" }, service_account_proxy_clear: { rpcMethod: "account/proxy/clear" }, diff --git a/crates/service/src/account/account_status.rs b/crates/service/src/account/account_status.rs index d3dcb7263..50d25523d 100644 --- a/crates/service/src/account/account_status.rs +++ b/crates/service/src/account/account_status.rs @@ -473,6 +473,82 @@ pub(crate) fn mark_account_unavailable_for_refresh_token_error( } } +/// 函数 `mark_account_unavailable_for_test_auth_status` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - storage: 参数 storage +/// - account_id: 参数 account_id +/// - status_code: 参数 status_code +/// +/// # 返回 +/// 返回是否已变更账号状态 +/// +/// 测试账号在真实上游请求中遇到 401/403 时，将账号标记为不可用。 +pub(crate) fn mark_account_unavailable_for_test_auth_status( + storage: &Storage, + account_id: &str, + status_code: u16, +) -> bool { + set_account_unavailable_with_reason(storage, account_id, &format!("test_http_{status_code}")) +} + +/// 函数 `mark_account_limited_for_test_rate_limit` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - storage: 参数 storage +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回是否已变更账号状态 +/// +/// 测试账号在真实上游请求中遇到 429 时，将账号标记为限流。 +pub(crate) fn mark_account_limited_for_test_rate_limit( + storage: &Storage, + account_id: &str, +) -> bool { + set_account_limited_with_reason(storage, account_id, "test_rate_limited") +} + +/// 函数 `restore_account_active_after_test` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - storage: 参数 storage +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回是否已变更账号状态 +/// +/// 测试成功后，仅当账号当前处于自动失败状态（unavailable/limited/banned）时恢复为 active； +/// 手动 disabled/inactive 账号不会被自动恢复。 +pub(crate) fn restore_account_active_after_test(storage: &Storage, account_id: &str) -> bool { + if should_preserve_manual_account_status(storage, account_id) { + return false; + } + let current = storage + .find_account_status_by_id(account_id) + .ok() + .flatten() + .unwrap_or_default(); + let normalized = current.trim().to_ascii_lowercase(); + if !matches!(normalized.as_str(), "unavailable" | "limited" | "banned") { + return false; + } + set_account_status(storage, account_id, "active", "test_ok"); + true +} + #[cfg(test)] #[path = "account_status_tests.rs"] mod tests; diff --git a/crates/service/src/account/account_test.rs b/crates/service/src/account/account_test.rs new file mode 100644 index 000000000..c327891a1 --- /dev/null +++ b/crates/service/src/account/account_test.rs @@ -0,0 +1,1152 @@ +use codexmanager_core::storage::Storage; +use crossbeam_channel::{bounded, Receiver, Sender, TrySendError}; +use reqwest::blocking::Client; +use reqwest::header::HeaderMap; +use serde::Serialize; +use serde_json::json; +use std::collections::{HashMap, HashSet}; +use std::io::{BufRead, BufReader}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use crate::account_status::{ + mark_account_limited_for_test_rate_limit, mark_account_unavailable_for_test_auth_status, + restore_account_active_after_test, +}; +use crate::account_warmup::{ + build_warmup_headers, resolve_warmup_authorization, summarize_warmup_error, WARMUP_UPSTREAM_URL, +}; +use crate::storage_helpers::open_storage; + +const DEFAULT_TEXT_TEST_PROMPT: &str = "hi"; +const DEFAULT_IMAGE_TEST_PROMPT: &str = + "Generate a cute orange cat astronaut sticker on a clean pastel background."; +const DEFAULT_TEXT_TEST_MODEL: &str = "gpt-5.3-codex"; +const DEFAULT_IMAGE_TEST_MODEL: &str = "gpt-image-2"; +const ACCOUNT_TEST_OVERALL_TIMEOUT: Duration = Duration::from_secs(120); + +/// 测试类型：文字模型直连，或图片模型走 image_generation 工具。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TestKind { + Text, + Image, +} + +impl TestKind { + fn parse(value: Option<&str>) -> TestKind { + match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() { + Some("image") => TestKind::Image, + _ => TestKind::Text, + } + } + + fn is_image(self) -> bool { + self == TestKind::Image + } +} + +/// 账号测试事件，序列化为 PRD 约定的 SSE 事件结构。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountTestEvent { + pub test_id: String, + #[serde(rename = "type")] + pub event_type: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub status: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub image_url: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub mime_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub success: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl AccountTestEvent { + fn new(test_id: &str, event_type: &str) -> Self { + Self { + test_id: test_id.to_string(), + event_type: event_type.to_string(), + text: None, + model: None, + status: None, + image_url: None, + mime_type: None, + success: None, + error: None, + } + } + + fn with_model(mut self, model: impl Into) -> Self { + self.model = Some(model.into()); + self + } + + fn with_text(mut self, text: impl Into) -> Self { + self.text = Some(text.into()); + self + } + + fn with_status(mut self, status: impl Into) -> Self { + self.status = Some(status.into()); + self + } + + fn with_image(mut self, image_url: impl Into, mime_type: impl Into) -> Self { + self.image_url = Some(image_url.into()); + self.mime_type = Some(mime_type.into()); + self + } + + fn with_success(mut self, success: bool) -> Self { + self.success = Some(success); + self + } + + fn with_error(mut self, error: impl Into) -> Self { + self.error = Some(error.into()); + self + } +} + +/// 账号测试启动结果，由 `account/test` RPC 直接返回给前端。 +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AccountTestStartResult { + pub test_id: String, + pub started: bool, + pub model: String, +} + +enum AccountTestOutcome { + Success, + AuthError(u16), + RateLimited, + Failed(String), + Canceled, +} + +type AccountTestEventHandler = Arc; + +static ACCOUNT_TEST_EVENT_HANDLER: OnceLock>> = + OnceLock::new(); +static ACCOUNT_TEST_EVENT_SUBSCRIBERS: OnceLock>>> = + OnceLock::new(); +static ACTIVE_ACCOUNT_TESTS: OnceLock>>> = OnceLock::new(); +static ACCOUNT_TEST_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// 函数 `set_account_test_event_handler` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - handler: 参数 handler +/// +/// # 返回 +/// 无 +/// +/// 桌面端通过该回调将测试事件转发到前端。 +pub fn set_account_test_event_handler(handler: F) +where + F: Fn(AccountTestEvent) + Send + Sync + 'static, +{ + let slot = ACCOUNT_TEST_EVENT_HANDLER.get_or_init(|| Mutex::new(None)); + let mut guard = crate::lock_utils::lock_recover(slot, "account_test_event_handler"); + *guard = Some(Arc::new(handler)); +} + +/// 函数 `subscribe_account_test_events` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 返回 +/// 返回账号测试事件订阅通道 +pub(crate) fn subscribe_account_test_events() -> Receiver { + let (sender, receiver) = bounded(64); + let subscribers = ACCOUNT_TEST_EVENT_SUBSCRIBERS.get_or_init(|| Mutex::new(Vec::new())); + let mut guard = crate::lock_utils::lock_recover(subscribers, "account_test_event_subscribers"); + guard.push(sender); + receiver +} + +/// 函数 `notify_account_test_event` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - event: 参数 event +/// +/// # 返回 +/// 无 +pub(crate) fn notify_account_test_event(event: AccountTestEvent) { + let handler = ACCOUNT_TEST_EVENT_HANDLER.get().and_then(|slot| { + let guard = crate::lock_utils::lock_recover(slot, "account_test_event_handler"); + guard.clone() + }); + if let Some(handler) = handler { + handler(event.clone()); + } + if let Some(subscribers) = ACCOUNT_TEST_EVENT_SUBSCRIBERS.get() { + let mut guard = + crate::lock_utils::lock_recover(subscribers, "account_test_event_subscribers"); + guard.retain(|sender| match sender.try_send(event.clone()) { + Ok(()) | Err(TrySendError::Full(_)) => true, + Err(TrySendError::Disconnected(_)) => false, + }); + } +} + +/// 函数 `start_account_test` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - account_id: 参数 account_id +/// - model: 参数 model +/// - prompt: 参数 prompt +/// +/// # 返回 +/// 返回测试启动结果 +pub(crate) fn start_account_test( + account_id: &str, + model: Option, + prompt: Option, + kind: Option, +) -> Result { + let account_id = account_id.trim(); + if account_id.is_empty() { + return Err("缺少账号 ID".to_string()); + } + + let storage = open_storage().ok_or_else(|| "storage unavailable".to_string())?; + let account = storage + .find_account_by_id(account_id) + .map_err(|err| err.to_string())? + .ok_or_else(|| "账号不存在".to_string())?; + let token = storage + .find_token_by_account_id(account_id) + .map_err(|err| err.to_string())? + .ok_or_else(|| "账号缺少访问令牌".to_string())?; + drop(account); + drop(token); + + let test_id = format!("test-{account_id}-{}", ACCOUNT_TEST_COUNTER.fetch_add(1, Ordering::Relaxed)); + let cancel_flag = register_active_test(account_id)?; + + let test_kind = resolve_test_kind(&storage, model.as_deref(), TestKind::parse(kind.as_deref())); + let resolved_model = resolve_model_slug(&storage, model.as_deref(), test_kind); + let resolved_prompt = resolve_prompt(prompt.as_deref(), test_kind); + drop(storage); + + let thread_account_id = account_id.to_string(); + let thread_test_id = test_id.clone(); + let thread_model = resolved_model.clone(); + let thread_prompt = resolved_prompt.clone(); + let thread_kind = test_kind; + std::thread::spawn(move || { + run_account_test( + &thread_account_id, + &thread_test_id, + &thread_model, + &thread_prompt, + thread_kind, + cancel_flag, + ); + }); + + Ok(AccountTestStartResult { + test_id, + started: true, + model: resolved_model, + }) +} + +/// 函数 `cancel_account_test` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回是否取消了进行中的测试 +pub(crate) fn cancel_account_test(account_id: &str) -> Result { + let account_id = account_id.trim(); + if account_id.is_empty() { + return Err("缺少账号 ID".to_string()); + } + let registry = ACTIVE_ACCOUNT_TESTS.get_or_init(|| Mutex::new(HashMap::new())); + let guard = crate::lock_utils::lock_recover(registry, "account_test_active_tests"); + match guard.get(account_id) { + Some(flag) => { + flag.store(true, Ordering::Relaxed); + Ok(true) + } + None => Ok(false), + } +} + +fn register_active_test(account_id: &str) -> Result, String> { + let registry = ACTIVE_ACCOUNT_TESTS.get_or_init(|| Mutex::new(HashMap::new())); + let mut guard = crate::lock_utils::lock_recover(registry, "account_test_active_tests"); + if guard.contains_key(account_id) { + return Err("该账号已有进行中的测试".to_string()); + } + let flag = Arc::new(AtomicBool::new(false)); + guard.insert(account_id.to_string(), flag.clone()); + Ok(flag) +} + +fn remove_active_test(account_id: &str) { + if let Some(registry) = ACTIVE_ACCOUNT_TESTS.get() { + let mut guard = crate::lock_utils::lock_recover(registry, "account_test_active_tests"); + guard.remove(account_id); + } +} + +/// 依据所选模型的真实能力修正测试类型，避免「文字直连 + 图片专用模型」这类组合把 +/// `gpt-image-2` 当成顶层主模型直连、被上游判定为「ChatGPT 账号不支持该模型」。 +/// 仅当模型在托管模型目录里能查到能力时才自动修正；未知模型沿用调用方传入的显式类型。 +fn resolve_test_kind(storage: &Storage, requested: Option<&str>, explicit: TestKind) -> TestKind { + let Some(slug) = requested.map(str::trim).filter(|value| !value.is_empty()) else { + return explicit; + }; + let Ok(Some(model)) = storage.get_managed_model_v2(slug) else { + return explicit; + }; + let supports_image = crate::models_v2::supports_image_generation(&model); + let supports_text = crate::models_v2::supports_text_generation(&model); + match (supports_image, supports_text) { + (true, false) => TestKind::Image, + (false, true) => TestKind::Text, + _ => explicit, + } +} + +fn resolve_model_slug(storage: &Storage, requested: Option<&str>, kind: TestKind) -> String { + if let Some(slug) = requested.map(str::trim).filter(|value| !value.is_empty()) { + return slug.to_string(); + } + let predicate = if kind.is_image() { + crate::models_v2::supports_image_generation + } else { + crate::models_v2::supports_text_generation + }; + storage + .list_api_models_v2() + .ok() + .and_then(|models| models.into_iter().find(predicate).map(|model| model.slug)) + .filter(|slug| !slug.trim().is_empty()) + .unwrap_or_else(|| { + if kind.is_image() { + DEFAULT_IMAGE_TEST_MODEL.to_string() + } else { + DEFAULT_TEXT_TEST_MODEL.to_string() + } + }) +} + +fn resolve_prompt(requested: Option<&str>, kind: TestKind) -> String { + if let Some(prompt) = requested.map(str::trim).filter(|value| !value.is_empty()) { + return prompt.to_string(); + } + if kind.is_image() { + DEFAULT_IMAGE_TEST_PROMPT.to_string() + } else { + DEFAULT_TEXT_TEST_PROMPT.to_string() + } +} + +fn run_account_test( + account_id: &str, + test_id: &str, + model: &str, + prompt: &str, + kind: TestKind, + cancel_flag: Arc, +) { + let outcome = execute_account_test(account_id, test_id, model, prompt, kind, &cancel_flag); + + if let Some(storage) = open_storage() { + match &outcome { + AccountTestOutcome::Success => { + let _ = restore_account_active_after_test(&storage, account_id); + } + AccountTestOutcome::AuthError(status_code) => { + let _ = mark_account_unavailable_for_test_auth_status( + &storage, + account_id, + *status_code, + ); + } + AccountTestOutcome::RateLimited => { + let _ = mark_account_limited_for_test_rate_limit(&storage, account_id); + } + AccountTestOutcome::Failed(_) | AccountTestOutcome::Canceled => {} + } + } + + remove_active_test(account_id); +} + +fn execute_account_test( + account_id: &str, + test_id: &str, + model: &str, + prompt: &str, + kind: TestKind, + cancel_flag: &Arc, +) -> AccountTestOutcome { + notify_account_test_event( + AccountTestEvent::new(test_id, "test_start").with_model(model), + ); + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("正在连接上游…"), + ); + + let client = match build_test_client(account_id) { + Ok(client) => client, + Err(err) => { + emit_redacted_error(test_id, &err, &[]); + return AccountTestOutcome::Failed(err); + } + }; + + let storage = match open_storage() { + Some(storage) => storage, + None => { + let message = "storage unavailable".to_string(); + emit_redacted_error(test_id, &message, &[]); + return AccountTestOutcome::Failed(message); + } + }; + + let account = match storage.find_account_by_id(account_id) { + Ok(Some(account)) => account, + Ok(None) => { + let message = "账号不存在".to_string(); + emit_redacted_error(test_id, &message, &[]); + return AccountTestOutcome::Failed(message); + } + Err(err) => { + let message = err.to_string(); + emit_redacted_error(test_id, &message, &[]); + return AccountTestOutcome::Failed(message); + } + }; + let token = match storage.find_token_by_account_id(account_id) { + Ok(Some(token)) => token, + Ok(None) => { + let message = "账号缺少访问令牌".to_string(); + emit_redacted_error(test_id, &message, &[]); + return AccountTestOutcome::Failed(message); + } + Err(err) => { + let message = err.to_string(); + emit_redacted_error(test_id, &message, &[]); + return AccountTestOutcome::Failed(message); + } + }; + + let secrets = vec![ + token.access_token.clone(), + token.refresh_token.clone(), + token.id_token.clone(), + ]; + let authorization = + match resolve_warmup_authorization(&storage, &client, &account, &token) { + Ok(authorization) => authorization, + Err(err) => { + emit_redacted_error(test_id, &err, &secrets); + return AccountTestOutcome::Failed(redact(&err, &secrets)); + } + }; + let headers = match build_warmup_headers(&account, &authorization) { + Ok(headers) => headers, + Err(err) => { + emit_redacted_error(test_id, &err, &secrets); + return AccountTestOutcome::Failed(redact(&err, &secrets)); + } + }; + // 诊断：账号测试与「裸 Bearer curl」不一致时，靠这行定位差异来源。 + // 只记录布尔/非敏感字段，绝不落 token 或 chatgpt-account-id 的值。 + log::info!( + "event=account_test_request_shape account_id={} uses_agent_identity={} has_chatgpt_account_id_header={} user_agent={}", + account_id, + authorization.uses_agent_identity, + headers.contains_key("chatgpt-account-id"), + headers + .get(reqwest::header::USER_AGENT) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default(), + ); + // 测试不再需要数据库连接，尽早归还到连接池，避免长时间占用。 + drop(storage); + + if kind.is_image() { + execute_image_test(&client, &headers, test_id, model, prompt, cancel_flag, &secrets) + } else { + execute_text_test(&client, &headers, test_id, model, prompt, cancel_flag, &secrets) + } +} + +fn build_test_client(account_id: &str) -> Result { + let proxy_url = crate::gateway::account_test_proxy_url_for_account(account_id)?; + // 只记录是否套代理，不落代理地址（地址可能含账号密码）。 + log::info!( + "event=account_test_proxy account_id={} has_proxy={}", + account_id, + proxy_url.is_some() + ); + crate::gateway::build_account_test_client_with_timeouts( + proxy_url.as_deref(), + ACCOUNT_TEST_OVERALL_TIMEOUT, + ) +} + +fn execute_text_test( + client: &Client, + headers: &HeaderMap, + test_id: &str, + model: &str, + prompt: &str, + cancel_flag: &Arc, + secrets: &[String], +) -> AccountTestOutcome { + let body = json!({ + "model": model, + "instructions": "", + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": prompt + }] + }], + "stream": true, + "store": false + }); + + let response = match client + .post(WARMUP_UPSTREAM_URL) + .headers(headers.clone()) + .json(&body) + .send() + { + Ok(response) => response, + Err(err) => { + let message = redact(&format!("测试请求发送失败: {err}"), secrets); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + }; + + let status = response.status(); + if !status.is_success() { + let body_text = response.text().unwrap_or_default(); + let message = redact( + &summarize_warmup_error(status.as_u16(), headers, &body_text), + secrets, + ); + emit_redacted_error(test_id, &message, secrets); + return classify_http_outcome(status.as_u16(), &message); + } + + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("已连接上游"), + ); + + let mut reader = BufReader::new(response); + let mut line = String::new(); + let mut event_name: Option = None; + let mut data_lines: Vec = Vec::new(); + + loop { + if cancel_flag.load(Ordering::Relaxed) { + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("已取消测试"), + ); + notify_account_test_event( + AccountTestEvent::new(test_id, "test_complete").with_success(false), + ); + return AccountTestOutcome::Canceled; + } + + line.clear(); + let bytes = match reader.read_line(&mut line) { + Ok(bytes) => bytes, + Err(err) => { + let message = redact(&format!("读取测试流失败: {err}"), secrets); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + }; + if bytes == 0 { + let message = "连接中断".to_string(); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + + let trimmed = line.trim_end_matches(['\r', '\n']); + if trimmed.is_empty() { + if let Some(outcome) = + process_text_sse_event(test_id, event_name.as_deref(), &data_lines) + { + return finish_text_test(test_id, outcome, secrets); + } + event_name = None; + data_lines.clear(); + continue; + } + if let Some(value) = trimmed.strip_prefix("event:") { + event_name = Some(value.trim().to_string()); + continue; + } + if let Some(value) = trimmed.strip_prefix("data:") { + data_lines.push(value.trim().to_string()); + } + } +} + +fn process_text_sse_event( + test_id: &str, + event_name: Option<&str>, + data_lines: &[String], +) -> Option { + let name = event_name.map(str::trim).filter(|value| !value.is_empty()); + + if data_lines.is_empty() { + if let Some(name) = name { + if is_terminal_event(name) { + return Some(AccountTestOutcome::Success); + } + if is_error_event(name) { + return Some(AccountTestOutcome::Failed(format!("测试失败: {name}"))); + } + } + return None; + } + + let data = data_lines.join("\n"); + let trimmed = data.trim(); + if trimmed == "[DONE]" { + return Some(AccountTestOutcome::Success); + } + + let Ok(value) = serde_json::from_str::(trimmed) else { + return None; + }; + let event_type = value + .get("type") + .and_then(serde_json::Value::as_str) + .or(name); + + match event_type { + Some("response.output_text.delta") => { + if let Some(delta) = value.get("delta").and_then(serde_json::Value::as_str) { + notify_account_test_event( + AccountTestEvent::new(test_id, "content").with_text(delta), + ); + } + None + } + Some("response.completed") | Some("response.done") => { + Some(AccountTestOutcome::Success) + } + Some("error") => { + let message = extract_stream_error_message(&value); + Some(AccountTestOutcome::Failed(message)) + } + Some("response.failed") | Some("response.incomplete") => { + let message = extract_stream_error_message(&value); + Some(AccountTestOutcome::Failed(message)) + } + _ => None, + } +} + +fn finish_text_test( + test_id: &str, + outcome: AccountTestOutcome, + secrets: &[String], +) -> AccountTestOutcome { + match &outcome { + AccountTestOutcome::Success => { + notify_account_test_event( + AccountTestEvent::new(test_id, "test_complete").with_success(true), + ); + } + AccountTestOutcome::Failed(message) => { + emit_redacted_error(test_id, message, secrets); + } + _ => {} + } + outcome +} + +fn execute_image_test( + client: &Client, + headers: &HeaderMap, + test_id: &str, + model: &str, + prompt: &str, + cancel_flag: &Arc, + secrets: &[String], +) -> AccountTestOutcome { + if cancel_flag.load(Ordering::Relaxed) { + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("已取消测试"), + ); + notify_account_test_event( + AccountTestEvent::new(test_id, "test_complete").with_success(false), + ); + return AccountTestOutcome::Canceled; + } + + let image_model = model.trim(); + // 图片测试与网关转发走同一条上游（chatgpt.com/backend-api/codex/responses），工具字段与 + // 网关 local_validation/request.rs 的 build_images_tool_from_request 保持一致（不带 + // `action`、带 `output_format:"png"`）。注意：这条直连上游不认 `metadata` 字段，带上会直接 + // 400「Unsupported parameter: metadata」，所以这里不能像网关内部那样塞 metadata。 + // 图片经 SSE 的 `response.output_item.done` / `response.completed` 事件回传（result 为 base64）。 + let image_headers = headers.clone(); + let body = json!({ + "model": crate::gateway::current_codex_image_main_model(), + "instructions": "", + "input": [{ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": prompt + }] + }], + "tools": [{ + "type": "image_generation", + "model": image_model, + "output_format": "png" + }], + "tool_choice": { + "type": "image_generation" + }, + "stream": true, + "store": false, + "reasoning": { + "effort": "medium", + "summary": "auto" + }, + "parallel_tool_calls": true, + "include": ["reasoning.encrypted_content"] + }); + + let response = match client + .post(WARMUP_UPSTREAM_URL) + .headers(image_headers.clone()) + .json(&body) + .send() + { + Ok(response) => response, + Err(err) => { + let message = redact(&format!("测试请求发送失败: {err}"), secrets); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + }; + + let status = response.status(); + if !status.is_success() { + let body_text = response.text().unwrap_or_default(); + let message = redact( + &summarize_warmup_error(status.as_u16(), &image_headers, &body_text), + secrets, + ); + emit_redacted_error(test_id, &message, secrets); + return classify_http_outcome(status.as_u16(), &message); + } + + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("已连接上游"), + ); + + let mut reader = BufReader::new(response); + let mut line = String::new(); + let mut event_name: Option = None; + let mut data_lines: Vec = Vec::new(); + let mut seen_images = HashSet::new(); + + loop { + if cancel_flag.load(Ordering::Relaxed) { + notify_account_test_event( + AccountTestEvent::new(test_id, "status").with_status("已取消测试"), + ); + notify_account_test_event( + AccountTestEvent::new(test_id, "test_complete").with_success(false), + ); + return AccountTestOutcome::Canceled; + } + + line.clear(); + let bytes = match reader.read_line(&mut line) { + Ok(bytes) => bytes, + Err(err) => { + let message = redact(&format!("读取图片流失败: {err}"), secrets); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + }; + if bytes == 0 { + let message = "连接中断".to_string(); + emit_redacted_error(test_id, &message, secrets); + return AccountTestOutcome::Failed(message); + } + + let trimmed = line.trim_end_matches(['\r', '\n']); + if trimmed.is_empty() { + if let Some(outcome) = process_image_sse_event( + test_id, + event_name.as_deref(), + &data_lines, + &mut seen_images, + ) { + return finish_image_test(test_id, outcome, &seen_images, secrets); + } + event_name = None; + data_lines.clear(); + continue; + } + if let Some(value) = trimmed.strip_prefix("event:") { + event_name = Some(value.trim().to_string()); + continue; + } + if let Some(value) = trimmed.strip_prefix("data:") { + data_lines.push(value.trim().to_string()); + } + } +} + +fn process_image_sse_event( + test_id: &str, + event_name: Option<&str>, + data_lines: &[String], + seen_images: &mut HashSet, +) -> Option { + let name = event_name.map(str::trim).filter(|value| !value.is_empty()); + + if data_lines.is_empty() { + if let Some(name) = name { + if is_terminal_event(name) { + return Some(AccountTestOutcome::Success); + } + if is_error_event(name) { + return Some(AccountTestOutcome::Failed(format!("测试失败: {name}"))); + } + } + return None; + } + + let data = data_lines.join("\n"); + let trimmed = data.trim(); + if trimmed == "[DONE]" { + return Some(AccountTestOutcome::Success); + } + + let Ok(value) = serde_json::from_str::(trimmed) else { + return None; + }; + let event_type = value + .get("type") + .and_then(serde_json::Value::as_str) + .or(name); + + match event_type { + Some("response.output_item.done") => { + if let Some(item) = value.get("item") { + emit_image_item(test_id, item, seen_images); + } + None + } + Some("response.completed") | Some("response.done") => { + // 兜底：`response.completed` 可能携带完整的 `response.output[]`（未走增量事件）。 + if let Some(output) = value + .get("response") + .and_then(|response| response.get("output")) + .and_then(serde_json::Value::as_array) + { + for item in output { + emit_image_item(test_id, item, seen_images); + } + } + Some(AccountTestOutcome::Success) + } + Some("error") => { + let message = extract_stream_error_message(&value); + Some(AccountTestOutcome::Failed(message)) + } + Some("response.failed") | Some("response.incomplete") => { + let message = extract_stream_error_message(&value); + Some(AccountTestOutcome::Failed(message)) + } + _ => None, + } +} + +fn finish_image_test( + test_id: &str, + outcome: AccountTestOutcome, + seen_images: &HashSet, + secrets: &[String], +) -> AccountTestOutcome { + match outcome { + AccountTestOutcome::Success => { + if seen_images.is_empty() { + let message = "未收到图片结果".to_string(); + emit_redacted_error(test_id, &message, secrets); + AccountTestOutcome::Failed(message) + } else { + notify_account_test_event( + AccountTestEvent::new(test_id, "test_complete").with_success(true), + ); + AccountTestOutcome::Success + } + } + AccountTestOutcome::Failed(message) => { + emit_redacted_error(test_id, &message, secrets); + AccountTestOutcome::Failed(message) + } + _ => outcome, + } +} + +fn image_item_to_data_uri(item: &serde_json::Value) -> Option<(String, String)> { + if item.get("type").and_then(serde_json::Value::as_str) != Some("image_generation_call") { + return None; + } + let base64_data = item.get("result").and_then(serde_json::Value::as_str)?.trim(); + if base64_data.is_empty() { + return None; + } + let format = item + .get("output_format") + .and_then(serde_json::Value::as_str) + .unwrap_or("png"); + let mime_type = image_mime_type(format); + let image_url = format!("data:{mime_type};base64,{base64_data}"); + Some((image_url, mime_type.to_string())) +} + +fn emit_image_item( + test_id: &str, + item: &serde_json::Value, + seen_images: &mut HashSet, +) -> bool { + let Some((image_url, mime_type)) = image_item_to_data_uri(item) else { + return false; + }; + // 以 item.id（缺失时退化为 data URI）去重，避免 `output_item.done` 与 `response.completed` 重复。 + let dedup_key = item + .get("id") + .and_then(serde_json::Value::as_str) + .map(str::to_string) + .unwrap_or_else(|| image_url.clone()); + if !seen_images.insert(dedup_key) { + return false; + } + notify_account_test_event( + AccountTestEvent::new(test_id, "image").with_image(image_url, mime_type), + ); + true +} + +fn image_mime_type(format: &str) -> &'static str { + match format.trim().to_ascii_lowercase().as_str() { + "webp" => "image/webp", + "jpeg" | "jpg" => "image/jpeg", + "gif" => "image/gif", + _ => "image/png", + } +} + +fn classify_http_outcome(status: u16, message: &str) -> AccountTestOutcome { + match status { + 401 | 403 => AccountTestOutcome::AuthError(status), + 429 => AccountTestOutcome::RateLimited, + _ => AccountTestOutcome::Failed(message.to_string()), + } +} + +fn extract_stream_error_message(value: &serde_json::Value) -> String { + value + .get("error") + .and_then(|error| { + error + .get("message") + .and_then(serde_json::Value::as_str) + .or_else(|| error.as_str()) + }) + .or_else(|| { + value + .get("response") + .and_then(|response| response.get("error")) + .and_then(|error| { + error + .get("message") + .and_then(serde_json::Value::as_str) + .or_else(|| error.as_str()) + }) + }) + .or_else(|| value.get("message").and_then(serde_json::Value::as_str)) + .map(str::trim) + .filter(|message| !message.is_empty()) + .unwrap_or("unknown stream error") + .to_string() +} + +fn is_terminal_event(value: &str) -> bool { + matches!(value.trim(), "response.completed" | "response.done") +} + +fn is_error_event(value: &str) -> bool { + matches!( + value.trim(), + "error" | "response.failed" | "response.incomplete" + ) +} + +fn redact(message: &str, secrets: &[String]) -> String { + let mut out = message.to_string(); + for secret in secrets { + let secret = secret.trim(); + if secret.len() < 4 { + continue; + } + out = out.replace(secret, "***"); + } + out +} + +fn emit_redacted_error(test_id: &str, message: &str, secrets: &[String]) { + notify_account_test_event( + AccountTestEvent::new(test_id, "error").with_error(redact(message, secrets)), + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_kind_parse() { + assert_eq!(TestKind::parse(Some("image")), TestKind::Image); + assert_eq!(TestKind::parse(Some("IMAGE")), TestKind::Image); + assert_eq!(TestKind::parse(Some("text")), TestKind::Text); + assert_eq!(TestKind::parse(Some("")), TestKind::Text); + assert_eq!(TestKind::parse(None), TestKind::Text); + } + + #[test] + fn resolve_test_kind_auto_switches_image_only_model() { + use codexmanager_core::storage::{ManagedModelV2, ManagedModelV2Upsert, ModelPriceV2}; + let storage = Storage::open_in_memory().expect("open storage"); + storage.init().expect("init storage"); + storage + .upsert_managed_model_v2(&ManagedModelV2Upsert { + model: ManagedModelV2 { + slug: "custom-image-model".to_string(), + display_name: "Custom Image Model".to_string(), + origin: "custom".to_string(), + enabled: true, + supported_in_api: true, + visibility: "list".to_string(), + instructions_mode: "passthrough".to_string(), + capabilities: serde_json::json!({ + "supports_image_generation": true, + "supports_text_generation": false + }), + price: ModelPriceV2 { + price_status: "missing".to_string(), + ..Default::default() + }, + ..ManagedModelV2::default() + }, + ..ManagedModelV2Upsert::default() + }) + .expect("save image model"); + + // 图片专用模型：即使前端传了默认的 text 类型，也要自动改成图片测试。 + assert_eq!( + resolve_test_kind(&storage, Some("custom-image-model"), TestKind::Text), + TestKind::Image + ); + assert_eq!( + resolve_test_kind(&storage, Some("custom-image-model"), TestKind::Image), + TestKind::Image + ); + // 未知模型 / 未指定模型：沿用显式类型，避免误判。 + assert_eq!( + resolve_test_kind(&storage, Some("external-model"), TestKind::Text), + TestKind::Text + ); + assert_eq!(resolve_test_kind(&storage, None, TestKind::Text), TestKind::Text); + } + + #[test] + fn image_mime_type_mapping() { + assert_eq!(image_mime_type("png"), "image/png"); + assert_eq!(image_mime_type("webp"), "image/webp"); + assert_eq!(image_mime_type("jpeg"), "image/jpeg"); + assert_eq!(image_mime_type("gif"), "image/gif"); + assert_eq!(image_mime_type("unknown"), "image/png"); + } + + #[test] + fn redaction_masks_secrets() { + let message = "auth error token=secret-token-value here"; + let redacted = redact(message, &["secret-token-value".to_string()]); + assert!(!redacted.contains("secret-token-value")); + assert!(redacted.contains("***")); + } + + #[test] + fn image_item_data_uri_extraction() { + let image = serde_json::json!({ + "type": "image_generation_call", + "id": "ig_1", + "result": "aGVsbG8=", + "output_format": "png" + }); + let (url, mime) = image_item_to_data_uri(&image).expect("extract image"); + assert_eq!(url, "data:image/png;base64,aGVsbG8="); + assert_eq!(mime, "image/png"); + + let text = serde_json::json!({"type": "message", "role": "assistant"}); + assert!(image_item_to_data_uri(&text).is_none()); + + let empty = serde_json::json!({"type": "image_generation_call", "result": ""}); + assert!(image_item_to_data_uri(&empty).is_none()); + + let webp = serde_json::json!({ + "type": "image_generation_call", + "result": "aGVsbG8=", + "output_format": "webp" + }); + assert_eq!(image_item_to_data_uri(&webp).unwrap().1, "image/webp"); + } +} diff --git a/crates/service/src/account/account_warmup.rs b/crates/service/src/account/account_warmup.rs index 3614b1058..f20b45461 100644 --- a/crates/service/src/account/account_warmup.rs +++ b/crates/service/src/account/account_warmup.rs @@ -13,7 +13,7 @@ use crate::usage_token_refresh::{refresh_and_persist_access_token, token_refresh const DEFAULT_WARMUP_MESSAGE: &str = "hi"; const FALLBACK_WARMUP_MESSAGE: &str = "你好"; -const WARMUP_UPSTREAM_URL: &str = "https://chatgpt.com/backend-api/codex/responses"; +pub(crate) const WARMUP_UPSTREAM_URL: &str = "https://chatgpt.com/backend-api/codex/responses"; const DEFAULT_WARMUP_MODEL: &str = "gpt-5.3-codex"; const X_OPENAI_FEDRAMP_HEADER_NAME: &str = "x-openai-fedramp"; @@ -40,11 +40,11 @@ struct AccountWarmupTarget { token: Token, } -struct WarmupAuthorization { +pub(crate) struct WarmupAuthorization { value: String, task_id: Option, is_fedramp: bool, - uses_agent_identity: bool, + pub(crate) uses_agent_identity: bool, account_scope_id: Option, } @@ -323,7 +323,7 @@ fn resolve_warmup_model_slug(storage: &Storage) -> String { .unwrap_or_else(|| DEFAULT_WARMUP_MODEL.to_string()) } -fn resolve_warmup_authorization( +pub(crate) fn resolve_warmup_authorization( storage: &Storage, client: &Client, account: &Account, @@ -606,7 +606,7 @@ fn summarize_warmup_stream_error(value: &serde_json::Value) -> String { #[path = "account_warmup_tests.rs"] mod tests; -fn build_warmup_headers( +pub(crate) fn build_warmup_headers( account: &Account, authorization: &WarmupAuthorization, ) -> Result { @@ -664,7 +664,7 @@ fn header_value(value: &str) -> Result { HeaderValue::from_str(value).map_err(|err| format!("invalid header value: {err}")) } -fn summarize_warmup_error(status: u16, headers: &HeaderMap, body: &str) -> String { +pub(crate) fn summarize_warmup_error(status: u16, headers: &HeaderMap, body: &str) -> String { let invalid_agent_task = crate::agent_identity::is_agent_identity_task_invalid_response(status, body.as_bytes()); let body_hint = if invalid_agent_task { diff --git a/crates/service/src/account/mod.rs b/crates/service/src/account/mod.rs index 452d74414..294fc940d 100644 --- a/crates/service/src/account/mod.rs +++ b/crates/service/src/account/mod.rs @@ -28,3 +28,5 @@ pub(crate) mod status; pub(crate) mod update; #[path = "account_warmup.rs"] pub(crate) mod warmup; +#[path = "account_test.rs"] +pub(crate) mod test; diff --git a/crates/service/src/gateway/core/runtime_config.rs b/crates/service/src/gateway/core/runtime_config.rs index 7df244740..6c80c6e3a 100644 --- a/crates/service/src/gateway/core/runtime_config.rs +++ b/crates/service/src/gateway/core/runtime_config.rs @@ -368,6 +368,70 @@ pub(crate) fn fresh_async_upstream_client_for_account( } } +/// 函数 `account_test_proxy_url_for_account` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回账号测试请求应使用的代理地址（显式账号代理 → 代理池 → 全局代理）， +/// 显式代理配置无效时 fail-closed。 +pub(crate) fn account_test_proxy_url_for_account( + account_id: &str, +) -> Result, String> { + ensure_runtime_config_loaded(); + match account_proxy_client_cache_entry(account_id) { + AccountProxyClientCacheEntry::Ready { proxy_url, .. } => return Ok(Some(proxy_url)), + AccountProxyClientCacheEntry::Invalid { proxy_url: _, error } => { + // 不回显 proxy_url：显式代理地址可能内嵌账号密码（http://user:pass@host）。 + return Err(format!( + "account explicit proxy for {account_id} is invalid and fail-closed. {error}" + )); + } + AccountProxyClientCacheEntry::NotConfigured => {} + } + let pool = crate::lock_utils::read_recover(upstream_client_pool_lock(), "upstream_client_pool"); + if let Some(proxy_url) = pool.proxy_for_account(account_id) { + return Ok(Some(proxy_url.to_string())); + } + Ok(current_upstream_proxy_url()) +} + +/// 函数 `build_account_test_client_with_timeouts` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - proxy_url: 参数 proxy_url +/// - overall_timeout: 参数 overall_timeout +/// +/// # 返回 +/// 返回带整体超时的阻塞式上游客户端，用于有界生命周期的账号测试请求。 +pub(crate) fn build_account_test_client_with_timeouts( + proxy_url: Option<&str>, + overall_timeout: Duration, +) -> Result { + let mut builder = Client::builder() + .timeout(overall_timeout) + .connect_timeout(upstream_connect_timeout_cached()) + .pool_max_idle_per_host(32) + .pool_idle_timeout(Some(Duration::from_secs(90))) + .tcp_keepalive(Some(Duration::from_secs(30))); + if let Some(proxy_url) = proxy_url.map(str::trim).filter(|value| !value.is_empty()) { + let proxy = Proxy::all(proxy_url).map_err(|err| format!("invalid proxy url: {err}"))?; + builder = builder.proxy(proxy); + } + builder + .build() + .map_err(|err| format!("build account test client failed: {err}")) +} + #[cfg(test)] pub(crate) fn upstream_proxy_url_for_account(account_id: &str) -> Option { ensure_runtime_config_loaded(); diff --git a/crates/service/src/gateway/mod.rs b/crates/service/src/gateway/mod.rs index 4300c505c..ca54ac7d2 100644 --- a/crates/service/src/gateway/mod.rs +++ b/crates/service/src/gateway/mod.rs @@ -421,6 +421,10 @@ pub(crate) use runtime_config::{ fresh_upstream_client_for_account, prepare_upstream_client_for_account, upstream_client_for_account, }; +pub(crate) use runtime_config::{ + account_test_proxy_url_for_account, build_account_test_client_with_timeouts, + current_codex_image_main_model, +}; pub(crate) use runtime_config::{front_proxy_max_body_bytes, front_proxy_zstd_max_body_bytes}; use runtime_config::{ prepare_upstream_client_for_aggregate_api_candidate, request_gate_wait_timeout, diff --git a/crates/service/src/http/account_test_events.rs b/crates/service/src/http/account_test_events.rs new file mode 100644 index 000000000..cb745bc0e --- /dev/null +++ b/crates/service/src/http/account_test_events.rs @@ -0,0 +1,184 @@ +use std::convert::Infallible; +use std::io::{self, Read}; +use std::time::Duration; + +use axum::body::{Body, Bytes}; +use axum::http::{ + HeaderMap as AxumHeaderMap, HeaderValue as AxumHeaderValue, StatusCode as AxumStatusCode, +}; +use axum::response::{IntoResponse, Response as AxumResponse}; +use crossbeam_channel::{Receiver, RecvTimeoutError}; +use futures_util::stream; +use tiny_http::{Header, Request, Response, StatusCode}; + +const EVENT_NAME: &str = "account-test-event"; +const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15); + +fn request_header_value<'a>(request: &'a Request, name: &str) -> Option<&'a str> { + request + .headers() + .iter() + .find(|header| header.field.as_str().as_str().eq_ignore_ascii_case(name)) + .map(|header| header.value.as_str().trim()) + .filter(|value| !value.is_empty()) +} + +fn rpc_token_valid(request: &Request) -> bool { + request_header_value(request, "X-CodexManager-Rpc-Token") + .is_some_and(crate::rpc_auth_token_matches) +} + +fn axum_rpc_token_valid(headers: &AxumHeaderMap) -> bool { + headers + .get("X-CodexManager-Rpc-Token") + .and_then(|value| value.to_str().ok()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .is_some_and(crate::rpc_auth_token_matches) +} + +fn response_header(name: &'static str, value: &'static str) -> Header { + Header::from_bytes(name.as_bytes(), value.as_bytes()).expect("valid static header") +} + +fn account_test_event_data(event: &crate::account_test::AccountTestEvent) -> String { + serde_json::to_string(event).unwrap_or_else(|_| "{}".to_string()) +} + +fn account_test_sse_frame(event: &crate::account_test::AccountTestEvent) -> Vec { + format!( + "event: {EVENT_NAME}\ndata: {}\n\n", + account_test_event_data(event) + ) + .into_bytes() +} + +fn next_account_test_event_chunk( + receiver: Receiver, +) -> Option<( + Receiver, + Vec, +)> { + let chunk = match receiver.recv_timeout(KEEPALIVE_INTERVAL) { + Ok(event) => account_test_sse_frame(&event), + Err(RecvTimeoutError::Timeout) => b": keep-alive\n\n".to_vec(), + Err(RecvTimeoutError::Disconnected) => return None, + }; + Some((receiver, chunk)) +} + +struct AccountTestEventStream { + receiver: Receiver, + pending: Vec, + pending_offset: usize, + opened: bool, +} + +impl AccountTestEventStream { + fn new(receiver: Receiver) -> Self { + Self { + receiver, + pending: Vec::new(), + pending_offset: 0, + opened: false, + } + } + + fn refill(&mut self) -> io::Result { + if !self.opened { + self.opened = true; + self.pending = b": connected\n\n".to_vec(); + self.pending_offset = 0; + return Ok(true); + } + + self.pending = match self.receiver.recv_timeout(KEEPALIVE_INTERVAL) { + Ok(event) => account_test_sse_frame(&event), + Err(RecvTimeoutError::Timeout) => b": keep-alive\n\n".to_vec(), + Err(RecvTimeoutError::Disconnected) => return Ok(false), + }; + self.pending_offset = 0; + Ok(true) + } +} + +impl Read for AccountTestEventStream { + fn read(&mut self, out: &mut [u8]) -> io::Result { + if out.is_empty() { + return Ok(0); + } + + if self.pending_offset >= self.pending.len() && !self.refill()? { + return Ok(0); + } + + let remaining = &self.pending[self.pending_offset..]; + let count = remaining.len().min(out.len()); + out[..count].copy_from_slice(&remaining[..count]); + self.pending_offset += count; + Ok(count) + } +} + +pub(crate) fn handle_account_test_events(request: Request) { + if request.method().as_str() != "GET" { + let _ = request.respond(Response::from_string("{}").with_status_code(405)); + return; + } + if !rpc_token_valid(&request) { + let _ = request.respond(Response::from_string("{}").with_status_code(401)); + return; + } + + let receiver = crate::account_test::subscribe_account_test_events(); + let headers = vec![ + response_header("Content-Type", "text/event-stream"), + response_header("Cache-Control", "no-cache"), + response_header("Connection", "keep-alive"), + response_header("X-Accel-Buffering", "no"), + ]; + let response = Response::new( + StatusCode(200), + headers, + AccountTestEventStream::new(receiver), + None, + None, + ); + let _ = request.respond(response); +} + +pub(crate) async fn handle_account_test_events_http(headers: AxumHeaderMap) -> AxumResponse { + if !axum_rpc_token_valid(&headers) { + return (AxumStatusCode::UNAUTHORIZED, "{}").into_response(); + } + + let receiver = crate::account_test::subscribe_account_test_events(); + let event_stream = stream::unfold((receiver, false), |(receiver, opened)| async move { + if !opened { + return Some(( + Ok::(Bytes::from_static(b": connected\n\n")), + (receiver, true), + )); + } + + let next = tokio::task::spawn_blocking(move || next_account_test_event_chunk(receiver)) + .await + .ok() + .flatten()?; + Some((Ok(Bytes::from(next.1)), (next.0, true))) + }); + + let mut response = AxumResponse::new(Body::from_stream(event_stream)); + *response.status_mut() = AxumStatusCode::OK; + response.headers_mut().insert( + "content-type", + AxumHeaderValue::from_static("text/event-stream"), + ); + response + .headers_mut() + .insert("cache-control", AxumHeaderValue::from_static("no-cache")); + response + .headers_mut() + .insert("x-accel-buffering", AxumHeaderValue::from_static("no")); + response +} diff --git a/crates/service/src/http/backend_router.rs b/crates/service/src/http/backend_router.rs index 637d111a6..e0d1586d6 100644 --- a/crates/service/src/http/backend_router.rs +++ b/crates/service/src/http/backend_router.rs @@ -5,6 +5,7 @@ pub(crate) enum BackendRoute { Rpc, AuthCallback, UsageRefreshEvents, + AccountTestEvents, Metrics, Gateway, } @@ -30,6 +31,9 @@ pub(crate) fn resolve_backend_route(method: &str, path: &str) -> BackendRoute { if method == "GET" && path == "/events/usage-refresh" { return BackendRoute::UsageRefreshEvents; } + if method == "GET" && path == "/events/account-test" { + return BackendRoute::AccountTestEvents; + } if method == "GET" && path == "/metrics" { return BackendRoute::Metrics; } @@ -55,6 +59,9 @@ pub(crate) fn handle_backend_request(request: Request) { BackendRoute::UsageRefreshEvents => { crate::http::usage_events::handle_usage_refresh_events(request) } + BackendRoute::AccountTestEvents => { + crate::http::account_test_events::handle_account_test_events(request) + } BackendRoute::Metrics => crate::http::gateway_endpoint::handle_metrics(request), BackendRoute::Gateway => crate::http::gateway_endpoint::handle_gateway(request), } diff --git a/crates/service/src/http/mod.rs b/crates/service/src/http/mod.rs index 8739eeb3e..1d3544b9c 100644 --- a/crates/service/src/http/mod.rs +++ b/crates/service/src/http/mod.rs @@ -2,6 +2,7 @@ pub mod callback_endpoint; pub mod gateway_endpoint; pub mod rpc_endpoint; pub mod server; +pub(crate) mod account_test_events; pub(crate) mod usage_events; pub(crate) mod backend_router; diff --git a/crates/service/src/http/proxy_runtime.rs b/crates/service/src/http/proxy_runtime.rs index 91b7e4ef9..0f88e690f 100644 --- a/crates/service/src/http/proxy_runtime.rs +++ b/crates/service/src/http/proxy_runtime.rs @@ -437,6 +437,10 @@ fn build_front_proxy_app(state: ProxyState) -> Router { "/events/usage-refresh", get(crate::http::usage_events::handle_usage_refresh_events_http), ) + .route( + "/events/account-test", + get(crate::http::account_test_events::handle_account_test_events_http), + ) .route("/v1/responses", any(responses_handler)) .route("/proxy-test-upload", post(proxy_test_upload)) .fallback(any(proxy_handler)) diff --git a/crates/service/src/lib.rs b/crates/service/src/lib.rs index f99ebe911..eda393fc8 100644 --- a/crates/service/src/lib.rs +++ b/crates/service/src/lib.rs @@ -47,6 +47,7 @@ pub(crate) use account::plan as account_plan; pub(crate) use account::proxy as account_proxy; pub(crate) use account::proxy_testing::presets::proxy_test_presets; pub(crate) use account::status as account_status; +pub(crate) use account::test as account_test; pub(crate) use account::update as account_update; pub(crate) use account::warmup as account_warmup; pub(crate) use aggregate_api::{ @@ -175,6 +176,7 @@ pub use logging::init_logging; pub use rpc_actor::{RpcActor, ROLE_ADMIN, ROLE_MEMBER, ROLE_SYSTEM_ADMIN}; pub use usage::tray_summary::{read_tray_usage_reset_summary, TrayUsageResetSummary}; pub use usage_refresh::{set_usage_refresh_completed_handler, UsageRefreshCompletedEvent}; +pub use account_test::{set_account_test_event_handler, AccountTestEvent}; /// 函数 `test_env_guard` /// diff --git a/crates/service/src/models_v2/mod.rs b/crates/service/src/models_v2/mod.rs index 6adadbbc5..7e9cd3082 100644 --- a/crates/service/src/models_v2/mod.rs +++ b/crates/service/src/models_v2/mod.rs @@ -116,6 +116,15 @@ pub(crate) fn supports_text_generation(model: &ManagedModelV2) -> bool { .unwrap_or(true) } +pub(crate) fn supports_image_generation(model: &ManagedModelV2) -> bool { + capability( + model, + &["supports_image_generation", "supportsImageGeneration"], + ) + .and_then(Value::as_bool) + .unwrap_or(false) +} + pub(crate) fn ensure_text_generation_model( storage: &codexmanager_core::storage::Storage, slug: Option<&str>, diff --git a/crates/service/src/rpc_dispatch/account.rs b/crates/service/src/rpc_dispatch/account.rs index d277627da..ee9bbcb39 100644 --- a/crates/service/src/rpc_dispatch/account.rs +++ b/crates/service/src/rpc_dispatch/account.rs @@ -3,8 +3,8 @@ use codexmanager_core::rpc::types::{JsonRpcRequest, JsonRpcResponse}; use crate::RpcActor; use crate::{ account_cleanup, account_delete, account_delete_many, account_export, account_import, - account_list, account_proxy, account_update, account_warmup, auth_account, auth_login, - auth_tokens, + account_list, account_proxy, account_test, account_update, account_warmup, auth_account, + auth_login, auth_tokens, }; /// 函数 `try_handle` @@ -120,6 +120,18 @@ pub(super) fn try_handle(req: &JsonRpcRequest, actor: &RpcActor) -> Option { + let account_id = first_str_param(req, &["accountId", "account_id"]).unwrap_or(""); + let model = first_str_param(req, &["model", "modelSlug", "model_slug"]) + .map(str::to_string); + let prompt = first_str_param(req, &["prompt", "message"]).map(str::to_string); + let kind = first_str_param(req, &["kind", "testType", "test_type"]).map(str::to_string); + super::value_or_error(account_test::start_account_test(account_id, model, prompt, kind)) + } + "account/test/cancel" => { + let account_id = first_str_param(req, &["accountId", "account_id"]).unwrap_or(""); + super::value_or_error(account_test::cancel_account_test(account_id)) + } "account/proxy/get" => { let account_id = first_str_param(req, &["accountId", "account_id"]).unwrap_or(""); super::value_or_error(account_proxy::get_account_proxy_settings(account_id)) diff --git a/crates/service/src/rpc_dispatch/mod.rs b/crates/service/src/rpc_dispatch/mod.rs index df636bd00..880fcf147 100644 --- a/crates/service/src/rpc_dispatch/mod.rs +++ b/crates/service/src/rpc_dispatch/mod.rs @@ -208,6 +208,8 @@ const MEMBER_METHOD_ALLOWLIST: &[&str] = &[ "account/usage/resetCredits", "account/usage/refresh", "account/warmup", + "account/test", + "account/test/cancel", "accountManager/password/change", "accountManager/profile/update", "accountManager/session/current",