diff --git a/apps/src-tauri/src/commands/account/remote.rs b/apps/src-tauri/src/commands/account/remote.rs index ffbee7657..4da89f0e2 100644 --- a/apps/src-tauri/src/commands/account/remote.rs +++ b/apps/src-tauri/src/commands/account/remote.rs @@ -276,6 +276,65 @@ 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, + test_id: Option, +) -> Result { + let params = serde_json::json!({ + "accountId": account_id, + "model": model, + "prompt": prompt, + "kind": kind, + "testId": test_id, + }); + 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, + test_id: String, +) -> Result { + let params = serde_json::json!({ + "accountId": account_id, + "testId": test_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..15d1c4287 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"; @@ -155,11 +156,16 @@ export interface AccountsPageViewProps { proxyDialogAccount: Account | null; proxySettings: AccountProxySettings | null; proxyProfiles: ProxyProfile[]; + canTestAccounts: boolean; isProxySettingsLoading: boolean; proxyEnabledDraft: boolean; 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 +381,7 @@ export function AccountsPageView(props: AccountsPageViewProps) { importByFile, importByDirectory, refreshAccount, + onAccountTestFinished, clearPreferredAccount, setPreferredAccount, toggleAccountStatus, @@ -548,6 +555,16 @@ export function AccountsPageView(props: AccountsPageViewProps) { {t("账号代理")} + {props.canTestAccounts ? ( + props.openAccountTest(account)} + > + + {t("测试账号")} + + ) : null} + {props.canTestAccounts ? ( + + ) : null} { diff --git a/apps/src/app/accounts/page.tsx b/apps/src/app/accounts/page.tsx index a7d0d4609..6d11cd3a7 100644 --- a/apps/src/app/accounts/page.tsx +++ b/apps/src/app/accounts/page.tsx @@ -3,6 +3,11 @@ import { useMemo, useState } from "react"; import { toast } from "sonner"; import { useAccounts } from "@/hooks/useAccounts"; +import { + isAdminRole, + resolveSessionRole, + useAppSession, +} from "@/hooks/useAppSession"; import { useDesktopPageActive } from "@/hooks/useDesktopPageActive"; import { usePageTransitionReady } from "@/hooks/usePageTransitionReady"; import { useRuntimeCapabilities } from "@/hooks/useRuntimeCapabilities"; @@ -65,6 +70,10 @@ export default function AccountsPage() { const { t } = useI18n(); const { isDesktopRuntime, canUseBrowserDownloadExport } = useRuntimeCapabilities(); + const { data: session, isLoading: isSessionLoading } = useAppSession(); + const role = resolveSessionRole(session, isSessionLoading, isDesktopRuntime); + const canTestAccounts = + isDesktopRuntime || (!isSessionLoading && isAdminRole(role)); const { accounts, planTypes, @@ -75,6 +84,7 @@ export default function AccountsPage() { refreshAllAccountRt, refreshAllAccounts, refreshAccountList, + refreshAccountsSilently, deleteAccount, deleteManyAccounts, cleanupAccountsByStatuses, @@ -141,6 +151,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 +579,23 @@ const toggleCleanupStatus = (rawStatus: string) => { setProxyUrlDraft(""); }; + const openAccountTest = (account: Account) => { + if (!canTestAccounts) return; + 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 +917,8 @@ const toggleCleanupStatus = (rawStatus: string) => { proxyDialogAccount={proxyDialogAccount} proxySettings={proxySettings} proxyProfiles={proxyProfiles} + canTestAccounts={canTestAccounts} + accountTestAccount={accountTestAccount} isProxySettingsLoading={isProxySettingsLoading} proxyEnabledDraft={proxyEnabledDraft} proxySourceDraft={proxySourceDraft} @@ -949,6 +991,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..1c1b3379d --- /dev/null +++ b/apps/src/components/modals/account-test-modal.tsx @@ -0,0 +1,542 @@ +"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 { useI18n } from "@/lib/i18n/provider"; +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; +} + +function newTestId(): string { + const cryptoApi = globalThis.crypto; + if (typeof cryptoApi?.randomUUID === "function") { + return cryptoApi.randomUUID(); + } + if (typeof cryptoApi?.getRandomValues !== "function") { + throw new Error("Secure random values are unavailable"); + } + const bytes = cryptoApi.getRandomValues(new Uint8Array(16)); + bytes[6] = (bytes[6] & 0x0f) | 0x40; + bytes[8] = (bytes[8] & 0x3f) | 0x80; + const hex = Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")); + return [ + hex.slice(0, 4).join(""), + hex.slice(4, 6).join(""), + hex.slice(6, 8).join(""), + hex.slice(8, 10).join(""), + hex.slice(10).join(""), + ].join("-"); +} + +export function AccountTestModal({ + account, + open, + onOpenChange, + onFinished, +}: AccountTestModalProps) { + const { t } = useI18n(); + 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 ?? t("测试失败"), + })); + setPhase("done"); + finishedRef.current = true; + if (accountIdRef.current) { + onFinishedRef.current?.(accountIdRef.current); + } + break; + } + }, [t]); + + const startTest = useCallback(async () => { + const id = accountIdRef.current; + if (!id || phaseRef.current === "running") { + return; + } + unlistenRef.current?.(); + unlistenRef.current = null; + // 订阅前先持有本次测试的 testId，事件到达时即可按 testId 隔离，避免并发测试串流。 + let testId: string; + try { + testId = newTestId(); + } catch { + setState({ text: "", images: [], error: t("启动测试失败") }); + setCanceled(false); + finishedRef.current = true; + setPhase("done"); + return; + } + testIdRef.current = testId; + finishedRef.current = false; + setState({ text: "", images: [] }); + setCanceled(false); + setPhase("running"); + + try { + const unlisten = await listenAccountTestEvent(testId, handleEvent); + unlistenRef.current = unlisten; + const result = await accountClient.testAccount({ + accountId: id, + model: selectedModel ?? undefined, + kind: testKind, + testId, + }); + setState((prev) => ({ ...prev, model: result.model ?? prev.model })); + } catch (err) { + unlistenRef.current?.(); + unlistenRef.current = null; + setState((prev) => ({ + ...prev, + error: + err instanceof Error && err.message.trim() + ? err.message + : t("启动测试失败"), + })); + finishedRef.current = true; + setPhase("done"); + if (accountIdRef.current) { + onFinishedRef.current?.(accountIdRef.current); + } + } + }, [handleEvent, selectedModel, t, testKind]); + + const cancelTest = useCallback(() => { + const id = accountIdRef.current; + const testId = testIdRef.current; + if (id && testId) { + void accountClient.cancelAccountTest(id, testId).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; + const testId = testIdRef.current; + if (id && testId) { + void accountClient.cancelAccountTest(id, testId).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 ( + + + + {t("测试账号")} + + {account?.name || account?.label || accountId} + + + +
+ {account ? ( +
+ +
+ ) : null} + +
+ + +
+ +
+ + + {models.length === 0 ? ( + + {t("未加载到可用模型，测试将使用后端默认模型。")} + + ) : null} +
+ +
+ {phase === "idle" ? ( +
+ + {canceled + ? t("已取消测试，可再次点击「开始测试」。") + : t("准备就绪，点击「开始测试」发起一次真实请求。")} + +
+ ) : ( + <> + {state.model ? ( +
+ {t("模型：")}{state.model} +
+ ) : null} + {status ? ( +
+ {phase === "running" ? ( + + ) : null} + {t(status)} +
+ ) : null} + {text ? ( +
+ {text} + {phase === "running" ? ( + _ + ) : null} +
+ ) : null} + {phase === "done" ? ( + <> +
+ {success ? ( + + ) : ( + + )} + {success ? t("测试成功") : error || t("测试失败")} +
+ {success && isManuallyDisabled(account) ? ( +
+ {t("该账号为手动禁用，测试虽成功但不会被自动恢复为「可用」。")} +
+ ) : null} + + ) : null} + + )} +
+ + {images.length > 0 ? ( +
+ + {t("图片预览")} + +
+ {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/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..c33696f41 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,14 @@ export interface AccountWarmupPayload { message?: string; } +export interface AccountTestPayload { + accountId: string; + model?: string; + prompt?: string; + kind?: "text" | "image"; + testId?: string; +} + export interface AccountProxyLatencyTestPayload { accountId: string; } @@ -527,6 +537,26 @@ 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", + testId: params.testId ?? null, + }), + ), + ), + cancelAccountTest: async (accountId: string, testId: string): Promise => + Boolean( + await invoke( + "service_account_test_cancel", + withAddr({ accountId, testId }), + ), + ), 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..63d3c2722 --- /dev/null +++ b/apps/src/lib/api/account-test-events.ts @@ -0,0 +1,140 @@ +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; + +const ACCOUNT_TEST_EVENT_OPEN_TIMEOUT_MS = 5_000; + +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( + testId: string, + 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; + let unlistenTauri: Unlisten | null = null; + const cleanup = () => { + window.removeEventListener(ACCOUNT_TEST_EVENT, handleWindowEvent); + if (eventSource && handleEventSourceEvent) { + eventSource.removeEventListener( + ACCOUNT_TEST_EVENT, + handleEventSourceEvent as EventListener + ); + } + eventSource?.close(); + unlistenTauri?.(); + }; + + try { + if ( + !isTauriRuntime() && + typeof EventSource !== "undefined" && + window.location.protocol.startsWith("http") + ) { + const normalizedTestId = testId.trim(); + if (!normalizedTestId) { + throw new Error("Missing account test ID"); + } + eventSource = new EventSource( + `/api/events/account-test?testId=${encodeURIComponent(normalizedTestId)}`, + ); + handleEventSourceEvent = (event: MessageEvent) => { + handler(readAccountTestMessagePayload(event)); + }; + eventSource.addEventListener( + ACCOUNT_TEST_EVENT, + handleEventSourceEvent as EventListener, + ); + + await new Promise((resolve, reject) => { + let settled = false; + const source = eventSource; + let timeoutId: number | undefined; + const finish = (error?: Error) => { + if (settled) return; + settled = true; + if (timeoutId !== undefined) { + window.clearTimeout(timeoutId); + } + source?.removeEventListener("open", handleOpen); + source?.removeEventListener("error", handleInitialError); + if (error) reject(error); + else resolve(); + }; + const handleOpen = () => finish(); + const handleInitialError = () => + finish(new Error("Failed to connect to account test events")); + source?.addEventListener("open", handleOpen); + source?.addEventListener("error", handleInitialError); + timeoutId = window.setTimeout( + () => finish(new Error("Timed out connecting to account test events")), + ACCOUNT_TEST_EVENT_OPEN_TIMEOUT_MS, + ); + // The connection may have opened between construction and listener registration. + // Re-check after listeners are attached so the RPC never starts before the SSE channel. + if (source?.readyState === 1) { + finish(); + } + }); + } + + if (isTauriRuntime()) { + const { listen } = await import("@tauri-apps/api/event"); + unlistenTauri = await listen( + ACCOUNT_TEST_EVENT, + (event) => { + handler(event.payload || {}); + }, + ); + } + + return cleanup; + } catch (error) { + cleanup(); + throw error; + } +} 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/apps/src/lib/i18n/messages/sections/en-accounts.ts b/apps/src/lib/i18n/messages/sections/en-accounts.ts index e359ad004..cfda3028d 100644 --- a/apps/src/lib/i18n/messages/sections/en-accounts.ts +++ b/apps/src/lib/i18n/messages/sections/en-accounts.ts @@ -286,5 +286,27 @@ export const EN_ACCOUNTS_MESSAGES: MessageCatalog = { "批量{action}完成：成功{success}个": "Bulk {action} complete: {success} succeeded", "批量{action}失败: {error}": "Bulk {action} failed: {error}", + "测试账号": "Test account", + "测试类型": "Test type", + "选择测试类型": "Select test type", + "文字模型": "Text model", + "图片模型": "Image model", + "测试模型": "Test model", + "选择模型": "Select model", + "官方模型": "Official models", + "图片": "Image", + "未加载到可用模型，测试将使用后端默认模型。": + "No available models were loaded. The backend default model will be used.", + "已取消测试，可再次点击「开始测试」。": + "Test canceled. Click Start test to run it again.", + "准备就绪，点击「开始测试」发起一次真实请求。": + "Ready. Click Start test to send a real request.", + "模型：": "Model: ", + "测试成功": "Test succeeded", + "该账号为手动禁用，测试虽成功但不会被自动恢复为「可用」。": + "This account is manually disabled. A successful test will not automatically mark it as available.", + "图片预览": "Image preview", + "开始测试": "Start test", + "启动测试失败": "Failed to start test", "预计删除": "Estimated delete", }; diff --git a/apps/src/lib/i18n/messages/sections/ko-accounts.ts b/apps/src/lib/i18n/messages/sections/ko-accounts.ts index 69fb9149c..500470046 100644 --- a/apps/src/lib/i18n/messages/sections/ko-accounts.ts +++ b/apps/src/lib/i18n/messages/sections/ko-accounts.ts @@ -273,5 +273,27 @@ export const KO_ACCOUNTS_MESSAGES: MessageCatalog = { "批量{action}完成：成功{success}个": "일괄 {action} 완료: 성공 {success}개", "批量{action}失败: {error}": "일괄 {action} 실패: {error}", + "测试账号": "계정 테스트", + "测试类型": "테스트 유형", + "选择测试类型": "테스트 유형 선택", + "文字模型": "텍스트 모델", + "图片模型": "이미지 모델", + "测试模型": "테스트 모델", + "选择模型": "모델 선택", + "官方模型": "공식 모델", + "图片": "이미지", + "未加载到可用模型，测试将使用后端默认模型。": + "사용 가능한 모델을 불러오지 못했습니다. 백엔드 기본 모델을 사용합니다.", + "已取消测试，可再次点击「开始测试」。": + "테스트가 취소되었습니다. 테스트 시작을 눌러 다시 실행할 수 있습니다.", + "准备就绪，点击「开始测试」发起一次真实请求。": + "준비되었습니다. 테스트 시작을 눌러 실제 요청을 보내세요.", + "模型：": "모델: ", + "测试成功": "테스트 성공", + "该账号为手动禁用，测试虽成功但不会被自动恢复为「可用」。": + "이 계정은 수동으로 비활성화되어 있어 테스트가 성공해도 자동으로 사용 가능 상태로 복원되지 않습니다.", + "图片预览": "이미지 미리보기", + "开始测试": "테스트 시작", + "启动测试失败": "테스트를 시작하지 못했습니다", "预计删除": "예상 삭제", }; diff --git a/apps/src/lib/i18n/messages/sections/ru-accounts.ts b/apps/src/lib/i18n/messages/sections/ru-accounts.ts index d80faa8be..ed092e485 100644 --- a/apps/src/lib/i18n/messages/sections/ru-accounts.ts +++ b/apps/src/lib/i18n/messages/sections/ru-accounts.ts @@ -286,5 +286,27 @@ export const RU_ACCOUNTS_MESSAGES: MessageCatalog = { "批量{action}完成：成功{success}个": "Массовое {action} завершено: успешно {success}", "批量{action}失败: {error}": "Массовое {action} не удалось: {error}", + "测试账号": "Тест аккаунта", + "测试类型": "Тип теста", + "选择测试类型": "Выберите тип теста", + "文字模型": "Текстовая модель", + "图片模型": "Модель изображений", + "测试模型": "Тестовая модель", + "选择模型": "Выберите модель", + "官方模型": "Официальные модели", + "图片": "Изображение", + "未加载到可用模型，测试将使用后端默认模型。": + "Доступные модели не загружены. Будет использована модель backend по умолчанию.", + "已取消测试，可再次点击「开始测试」。": + "Тест отменен. Нажмите «Начать тест», чтобы запустить его снова.", + "准备就绪，点击「开始测试」发起一次真实请求。": + "Готово. Нажмите «Начать тест», чтобы отправить реальный запрос.", + "模型：": "Модель: ", + "测试成功": "Тест успешно выполнен", + "该账号为手动禁用，测试虽成功但不会被自动恢复为「可用」。": + "Аккаунт отключен вручную, поэтому успешный тест не вернет его автоматически в доступное состояние.", + "图片预览": "Предпросмотр изображения", + "开始测试": "Начать тест", + "启动测试失败": "Не удалось запустить тест", "预计删除": "Ожидается удалить", }; diff --git a/apps/tests/account-test-admin-gate.test.mjs b/apps/tests/account-test-admin-gate.test.mjs new file mode 100644 index 000000000..6bc2a5b84 --- /dev/null +++ b/apps/tests/account-test-admin-gate.test.mjs @@ -0,0 +1,51 @@ +import assert from "node:assert/strict"; +import fs from "node:fs/promises"; +import path from "node:path"; +import test from "node:test"; + +const appsRoot = path.resolve(import.meta.dirname, ".."); +const pageSource = await fs.readFile( + path.join(appsRoot, "src", "app", "accounts", "page.tsx"), + "utf8", +); +const viewSource = await fs.readFile( + path.join(appsRoot, "src", "app", "accounts", "accounts-page-view.tsx"), + "utf8", +); +const modalSource = await fs.readFile( + path.join(appsRoot, "src", "components", "modals", "account-test-modal.tsx"), + "utf8", +); + +test("account test UI is derived from the current admin session", () => { + assert.match(pageSource, /useAppSession\(\)/); + assert.match( + pageSource, + /resolveSessionRole\(session, isSessionLoading, isDesktopRuntime\)/, + ); + assert.match( + pageSource, + /const canTestAccounts\s*=\s*isDesktopRuntime \|\|\s*\(!isSessionLoading && isAdminRole\(role\)\)/, + ); + assert.match( + pageSource, + /const openAccountTest = \(account: Account\) => \{\s*if \(!canTestAccounts\) return;/, + ); + assert.match(pageSource, /canTestAccounts=\{canTestAccounts\}/); +}); + +test("account test menu and modal are both hidden from non-admin views", () => { + assert.match( + viewSource, + /\{props\.canTestAccounts \? \(\s* { + assert.match(modalSource, /cryptoApi\.getRandomValues\(new Uint8Array\(16\)\)/); + assert.doesNotMatch(modalSource, /Math\.random|Date\.now/); +}); diff --git a/crates/core/src/storage/accounts.rs b/crates/core/src/storage/accounts.rs index e7a548c7a..357c073aa 100644 --- a/crates/core/src/storage/accounts.rs +++ b/crates/core/src/storage/accounts.rs @@ -1280,6 +1280,30 @@ impl Storage { Ok((self.account_exists(account_id)?, false)) } + /// Updates an account status only when the row still matches the state observed by the + /// caller. Long-running operations use this guard so a stale completion cannot overwrite a + /// newer manual or gateway-driven status transition. + pub fn update_account_status_if_context_matches( + &self, + account_id: &str, + expected_status: &str, + expected_updated_at: i64, + status: &str, + ) -> Result { + let updated_at = now_ts().max(expected_updated_at.saturating_add(1)); + let updated = self.conn.execute( + update_account_status_if_context_matches_sql(), + ( + status, + updated_at, + account_id, + expected_status, + expected_updated_at, + ), + )?; + Ok(updated > 0) + } + /// 函数 `delete_account` /// /// 作者: gaohongshun diff --git a/crates/core/src/storage/accounts_sql.rs b/crates/core/src/storage/accounts_sql.rs index f7942545e..635673172 100644 --- a/crates/core/src/storage/accounts_sql.rs +++ b/crates/core/src/storage/accounts_sql.rs @@ -97,6 +97,12 @@ pub(super) fn update_account_status_if_changed_sql() -> &'static str { "UPDATE accounts SET status = ?1, updated_at = ?2 WHERE id = ?3 AND status != ?1" } +pub(super) fn update_account_status_if_context_matches_sql() -> &'static str { + "UPDATE accounts + SET status = ?1, updated_at = ?2 + WHERE id = ?3 AND status = ?4 AND updated_at = ?5" +} + pub(super) fn delete_account_by_id_sql() -> &'static str { "DELETE FROM accounts WHERE id = ?1" } diff --git a/crates/core/tests/storage.rs b/crates/core/tests/storage.rs index 5a8185d28..269880de0 100644 --- a/crates/core/tests/storage.rs +++ b/crates/core/tests/storage.rs @@ -1361,6 +1361,66 @@ fn storage_updates_account_status_only_when_changed() { assert_eq!(loaded.status, "inactive"); } +#[test] +fn storage_updates_account_status_only_when_observed_context_still_matches() { + let storage = Storage::open_in_memory().expect("open in memory"); + storage.init().expect("init schema"); + let observed_updated_at = 1_700_000_000; + storage + .insert_account(&Account { + id: "acc-context-cas".to_string(), + label: "context cas".to_string(), + issuer: "https://auth.openai.com".to_string(), + chatgpt_account_id: Some("acct_context_cas".to_string()), + workspace_id: None, + group_name: None, + sort: 0, + status: "active".to_string(), + created_at: observed_updated_at, + updated_at: observed_updated_at, + }) + .expect("insert account"); + + assert!(storage + .update_account_status_if_context_matches( + "acc-context-cas", + "active", + observed_updated_at, + "inactive", + ) + .expect("matching context update")); + let changed = storage + .find_account_by_id("acc-context-cas") + .expect("find account") + .expect("account exists"); + assert_eq!(changed.status, "inactive"); + assert!(changed.updated_at > observed_updated_at); + + assert!(!storage + .update_account_status_if_context_matches( + "acc-context-cas", + "active", + observed_updated_at, + "banned", + ) + .expect("stale status update")); + assert!(!storage + .update_account_status_if_context_matches( + "acc-context-cas", + "inactive", + observed_updated_at, + "banned", + ) + .expect("stale timestamp update")); + assert_eq!( + storage + .find_account_status_by_id("acc-context-cas") + .expect("read final status") + .as_deref(), + Some("inactive") + ); +} + /// 函数 `storage_gateway_candidates_exclude_unavailable_or_missing_token_accounts` /// /// 作者: gaohongshun diff --git a/crates/service/src/account/account_status.rs b/crates/service/src/account/account_status.rs index d3dcb7263..b0c2bd0ca 100644 --- a/crates/service/src/account/account_status.rs +++ b/crates/service/src/account/account_status.rs @@ -25,10 +25,11 @@ pub(crate) struct GatewayErrorFollowUp { pub should_mark_default_cooldown: bool, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct AccountStatusContext { pub status: String, pub reason: Option, + pub updated_at: Option, } /// 函数 `latest_status_reason` @@ -54,13 +55,14 @@ pub(crate) fn load_account_status_context( storage: &Storage, account_id: &str, ) -> AccountStatusContext { + let account = storage.find_account_by_id(account_id).ok().flatten(); AccountStatusContext { - status: storage - .find_account_status_by_id(account_id) - .ok() - .flatten() + status: account + .as_ref() + .map(|account| account.status.clone()) .unwrap_or_default(), reason: latest_status_reason(storage, account_id), + updated_at: account.map(|account| account.updated_at), } } @@ -473,6 +475,131 @@ 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, + context: &AccountStatusContext, +) -> bool { + set_account_status_after_test_if_context_matches( + storage, + account_id, + "unavailable", + &format!("test_http_{status_code}"), + context, + ) +} + +/// 函数 `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, + context: &AccountStatusContext, +) -> bool { + set_account_status_after_test_if_context_matches( + storage, + account_id, + "limited", + "test_rate_limited", + context, + ) +} + +fn set_account_status_after_test_if_context_matches( + storage: &Storage, + account_id: &str, + status: &str, + reason: &str, + context: &AccountStatusContext, +) -> bool { + let normalized = context.status.trim().to_ascii_lowercase(); + if matches!(normalized.as_str(), "disabled" | "inactive" | "banned") { + return false; + } + if load_account_status_context(storage, account_id) != *context { + return false; + } + let Some(expected_updated_at) = context.updated_at else { + return false; + }; + let changed = storage + .update_account_status_if_context_matches( + account_id, + &context.status, + expected_updated_at, + status, + ) + .unwrap_or(false); + if !changed { + return false; + } + crate::gateway::invalidate_candidate_cache(); + let _ = storage.insert_event(&Event { + account_id: Some(account_id.to_string()), + event_type: "account_status_update".to_string(), + message: format!("status={status} reason={reason}"), + created_at: now_ts(), + }); + true +} + +/// 函数 `restore_account_active_after_test` +/// +/// 作者: gaohongshun +/// +/// 时间: 2026-08-26 +/// +/// # 参数 +/// - storage: 参数 storage +/// - account_id: 参数 account_id +/// +/// # 返回 +/// 返回是否已变更账号状态 +/// +/// 测试成功后，仅当账号从测试开始至今仍处于同一个自动失败状态 +///（unavailable/limited）时恢复为 active；banned 与手动停用状态始终保留。 +pub(crate) fn restore_account_active_after_test( + storage: &Storage, + account_id: &str, + context: &AccountStatusContext, +) -> bool { + let normalized = context.status.trim().to_ascii_lowercase(); + if !matches!(normalized.as_str(), "unavailable" | "limited") { + return false; + } + set_account_status_after_test_if_context_matches( + storage, account_id, "active", "test_ok", context, + ) +} + #[cfg(test)] #[path = "account_status_tests.rs"] mod tests; diff --git a/crates/service/src/account/account_status_tests.rs b/crates/service/src/account/account_status_tests.rs index fbcf8134d..25be62815 100644 --- a/crates/service/src/account/account_status_tests.rs +++ b/crates/service/src/account/account_status_tests.rs @@ -1,6 +1,8 @@ use super::{ - analyze_gateway_error, classify_account_availability_signal, - mark_account_unavailable_for_gateway_error, AccountAvailabilitySignal, GatewayErrorKind, + analyze_gateway_error, classify_account_availability_signal, load_account_status_context, + mark_account_limited_for_test_rate_limit, mark_account_unavailable_for_gateway_error, + mark_account_unavailable_for_test_auth_status, restore_account_active_after_test, + AccountAvailabilitySignal, GatewayErrorKind, }; use codexmanager_core::storage::{now_ts, Account, Storage, UsageSnapshotRecord}; @@ -208,3 +210,90 @@ fn gateway_usage_limit_error_marks_account_limited_when_snapshot_exhausted() { Some("usage_limit_exhausted") ); } + +#[test] +fn stale_account_test_outcomes_do_not_overwrite_newer_banned_status() { + let _guard = crate::test_env_guard(); + let storage = Storage::open_in_memory().expect("open storage"); + storage.init().expect("init storage"); + let now = now_ts(); + let account_id = "acc-stale-account-test"; + storage + .insert_account(&Account { + id: account_id.to_string(), + label: "stale-account-test".to_string(), + issuer: "issuer".to_string(), + chatgpt_account_id: None, + workspace_id: None, + group_name: None, + sort: 0, + status: "unavailable".to_string(), + created_at: now, + updated_at: now, + }) + .expect("insert account"); + let start_context = load_account_status_context(&storage, account_id); + + storage + .update_account_status(account_id, "banned") + .expect("ban account while test is running"); + + assert!(!restore_account_active_after_test( + &storage, + account_id, + &start_context + )); + assert!(!mark_account_unavailable_for_test_auth_status( + &storage, + account_id, + 401, + &start_context + )); + assert!(!mark_account_limited_for_test_rate_limit( + &storage, + account_id, + &start_context + )); + assert_eq!( + storage + .find_account_status_by_id(account_id) + .expect("read account status") + .as_deref(), + Some("banned") + ); +} + +#[test] +fn account_test_success_never_restores_an_existing_banned_status() { + let _guard = crate::test_env_guard(); + let storage = Storage::open_in_memory().expect("open storage"); + storage.init().expect("init storage"); + let now = now_ts(); + let account_id = "acc-banned-before-test"; + storage + .insert_account(&Account { + id: account_id.to_string(), + label: "banned-before-test".to_string(), + issuer: "issuer".to_string(), + chatgpt_account_id: None, + workspace_id: None, + group_name: None, + sort: 0, + status: "banned".to_string(), + created_at: now, + updated_at: now, + }) + .expect("insert account"); + let context = load_account_status_context(&storage, account_id); + + assert!(!restore_account_active_after_test( + &storage, account_id, &context + )); + assert_eq!( + storage + .find_account_status_by_id(account_id) + .expect("read account status") + .as_deref(), + Some("banned") + ); +} diff --git a/crates/service/src/account/account_test.rs b/crates/service/src/account/account_test.rs new file mode 100644 index 000000000..50ae75bcc --- /dev/null +++ b/crates/service/src/account/account_test.rs @@ -0,0 +1,1336 @@ +use codexmanager_core::storage::Storage; +use crossbeam_channel::{bounded, Receiver, Sender, TrySendError}; +use rand::RngCore; +use reqwest::blocking::Client; +use reqwest::header::HeaderMap; +use serde::Serialize; +use serde_json::json; +use std::collections::{HashMap, HashSet}; +use std::fmt::Write as _; +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::{ + load_account_status_context, mark_account_limited_for_test_rate_limit, + mark_account_unavailable_for_test_auth_status, restore_account_active_after_test, + AccountStatusContext, +}; +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; + +#[derive(Clone)] +struct ActiveAccountTest { + test_id: String, + cancel_flag: Arc, +} + +struct AccountTestSubscriber { + id: u64, + sender: Sender, +} + +pub(crate) struct AccountTestEventSubscription { + test_id: String, + subscriber_id: u64, + receiver: Receiver, +} + +impl AccountTestEventSubscription { + pub(crate) fn recv_timeout( + &self, + timeout: Duration, + ) -> Result { + self.receiver.recv_timeout(timeout) + } +} + +impl Drop for AccountTestEventSubscription { + fn drop(&mut self) { + if let Some(subscribers) = ACCOUNT_TEST_EVENT_SUBSCRIBERS.get() { + let mut guard = + crate::lock_utils::lock_recover(subscribers, "account_test_event_subscribers"); + let remove_test_id = if let Some(entries) = guard.get_mut(&self.test_id) { + entries.retain(|entry| entry.id != self.subscriber_id); + entries.is_empty() + } else { + false + }; + if remove_test_id { + guard.remove(&self.test_id); + } + } + } +} + +static ACCOUNT_TEST_EVENT_HANDLER: OnceLock>> = + OnceLock::new(); +static ACCOUNT_TEST_EVENT_SUBSCRIBERS: OnceLock< + Mutex>>, +> = OnceLock::new(); +static ACTIVE_ACCOUNT_TESTS: OnceLock>> = OnceLock::new(); +static ACCOUNT_TEST_SUBSCRIBER_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(test_id: &str) -> AccountTestEventSubscription { + let (sender, receiver) = bounded(64); + let subscriber_id = ACCOUNT_TEST_SUBSCRIBER_COUNTER.fetch_add(1, Ordering::Relaxed); + let subscribers = ACCOUNT_TEST_EVENT_SUBSCRIBERS.get_or_init(|| Mutex::new(HashMap::new())); + let mut guard = crate::lock_utils::lock_recover(subscribers, "account_test_event_subscribers"); + guard + .entry(test_id.to_string()) + .or_default() + .push(AccountTestSubscriber { + id: subscriber_id, + sender, + }); + AccountTestEventSubscription { + test_id: test_id.to_string(), + subscriber_id, + 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"); + let remove_test_id = if let Some(entries) = guard.get_mut(&event.test_id) { + entries.retain(|entry| match entry.sender.try_send(event.clone()) { + Ok(()) | Err(TrySendError::Full(_)) => true, + Err(TrySendError::Disconnected(_)) => false, + }); + entries.is_empty() + } else { + false + }; + if remove_test_id { + guard.remove(&event.test_id); + } + } +} + +pub(crate) fn normalize_account_test_id(value: &str) -> Option { + let value = value.trim(); + if value.is_empty() + || value.len() > 128 + || !value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + { + return None; + } + Some(value.to_string()) +} + +fn generate_account_test_id() -> String { + let mut bytes = [0_u8; 24]; + rand::rngs::OsRng.fill_bytes(&mut bytes); + let mut test_id = String::with_capacity(5 + bytes.len() * 2); + test_id.push_str("test-"); + for byte in bytes { + write!(&mut test_id, "{byte:02x}").expect("writing to a String cannot fail"); + } + test_id +} + +/// 函数 `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, + test_id: 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())?; + let status_context = load_account_status_context(&storage, account_id); + drop(account); + drop(token); + + // 前端可自行生成 testId 并在订阅前持有它，从而在事件到达时就能按 testId 隔离，避免 + // 「订阅到拿到 testId」之间的竞态串流。未提供时回退到服务端自增 ID 保持向后兼容。 + let test_id = match test_id { + Some(value) => { + normalize_account_test_id(&value).ok_or_else(|| "无效测试 ID".to_string())? + } + None => generate_account_test_id(), + }; + let cancel_flag = register_active_test(account_id, &test_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, + status_context, + ); + }); + + 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, test_id: &str) -> Result { + let account_id = account_id.trim(); + if account_id.is_empty() { + return Err("缺少账号 ID".to_string()); + } + let test_id = normalize_account_test_id(test_id).ok_or_else(|| "无效测试 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(active) if active.test_id == test_id => { + active.cancel_flag.store(true, Ordering::Relaxed); + Ok(true) + } + Some(_) | None => Ok(false), + } +} + +fn register_active_test(account_id: &str, test_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(), + ActiveAccountTest { + test_id: test_id.to_string(), + cancel_flag: flag.clone(), + }, + ); + Ok(flag) +} + +fn remove_active_test(account_id: &str, test_id: &str) { + if let Some(registry) = ACTIVE_ACCOUNT_TESTS.get() { + let mut guard = crate::lock_utils::lock_recover(registry, "account_test_active_tests"); + if guard + .get(account_id) + .is_some_and(|active| active.test_id == test_id) + { + 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, + status_context: AccountStatusContext, +) { + 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, &status_context); + } + AccountTestOutcome::AuthError(status_code) => { + let _ = mark_account_unavailable_for_test_auth_status( + &storage, + account_id, + *status_code, + &status_context, + ); + } + AccountTestOutcome::RateLimited => { + let _ = + mark_account_limited_for_test_rate_limit(&storage, account_id, &status_context); + } + AccountTestOutcome::Failed(_) | AccountTestOutcome::Canceled => {} + } + } + + remove_active_test(account_id, test_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"); + } + + #[test] + fn account_test_id_validation_accepts_opaque_ascii_ids_only() { + assert_eq!( + normalize_account_test_id(" 550e8400-e29b-41d4-a716-446655440000 ").as_deref(), + Some("550e8400-e29b-41d4-a716-446655440000") + ); + assert!(normalize_account_test_id("").is_none()); + assert!(normalize_account_test_id("bad/id").is_none()); + assert!(normalize_account_test_id(&"a".repeat(129)).is_none()); + } + + #[test] + fn server_generated_account_test_ids_are_opaque_and_unique() { + let first = generate_account_test_id(); + let second = generate_account_test_id(); + assert_eq!(first.len(), 53); + assert!(first.starts_with("test-")); + assert!(normalize_account_test_id(&first).is_some()); + assert_ne!(first, second); + } + + #[test] + fn account_test_subscriptions_are_filtered_by_exact_test_id() { + let first_id = "subscription-filter-first"; + let second_id = "subscription-filter-second"; + let first = subscribe_account_test_events(first_id); + let second = subscribe_account_test_events(second_id); + + notify_account_test_event(AccountTestEvent::new(first_id, "status")); + + assert_eq!( + first + .receiver + .try_recv() + .expect("matching subscriber receives event") + .test_id, + first_id + ); + assert!(second.receiver.try_recv().is_err()); + } + + #[test] + fn cancel_account_test_requires_matching_test_id() { + let account_id = "cancel-owner-account"; + let test_id = "cancel-owner-test"; + let cancel_flag = register_active_test(account_id, test_id).expect("register active test"); + + assert!(!cancel_account_test(account_id, "stale-test").expect("stale cancel result")); + assert!(!cancel_flag.load(Ordering::Relaxed)); + assert!(cancel_account_test(account_id, test_id).expect("matching cancel result")); + assert!(cancel_flag.load(Ordering::Relaxed)); + + remove_active_test(account_id, test_id); + } +} 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..4373b3372 100644 --- a/crates/service/src/account/mod.rs +++ b/crates/service/src/account/mod.rs @@ -24,6 +24,8 @@ pub(crate) mod proxy_health; pub(crate) mod proxy_testing; #[path = "account_status.rs"] pub(crate) mod status; +#[path = "account_test.rs"] +pub(crate) mod test; #[path = "account_update.rs"] pub(crate) mod update; #[path = "account_warmup.rs"] diff --git a/crates/service/src/gateway/core/runtime_config.rs b/crates/service/src/gateway/core/runtime_config.rs index 7df244740..ee4bff49c 100644 --- a/crates/service/src/gateway/core/runtime_config.rs +++ b/crates/service/src/gateway/core/runtime_config.rs @@ -368,6 +368,73 @@ 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..1f7eb8c82 100644 --- a/crates/service/src/gateway/mod.rs +++ b/crates/service/src/gateway/mod.rs @@ -416,6 +416,10 @@ use route_quality::record_route_quality; pub(crate) use runtime_config::invalidate_account_proxy_client_cache as invalidate_account_proxy_cache; pub(crate) use runtime_config::upstream_client; pub(crate) use runtime_config::{account_max_inflight_limit, set_account_max_inflight_limit}; +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::{ async_upstream_client_for_account, fresh_async_upstream_client_for_account, fresh_upstream_client_for_account, prepare_upstream_client_for_account, 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..22f9df1d5 --- /dev/null +++ b/crates/service/src/http/account_test_events.rs @@ -0,0 +1,227 @@ +use std::convert::Infallible; +use std::io::{self, Read}; +use std::time::Duration; + +use axum::body::{Body, Bytes}; +use axum::extract::RawQuery; +use axum::http::{ + HeaderMap as AxumHeaderMap, HeaderValue as AxumHeaderValue, StatusCode as AxumStatusCode, +}; +use axum::response::{IntoResponse, Response as AxumResponse}; +use crossbeam_channel::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 account_test_id_from_query(query: Option<&str>) -> Option { + let mut values = query + .into_iter() + .flat_map(|query| url::form_urlencoded::parse(query.as_bytes())) + .filter(|(key, _)| key == "testId") + .map(|(_, value)| value.into_owned()); + let value = values.next()?; + if values.next().is_some() { + return None; + } + crate::account_test::normalize_account_test_id(&value) +} + +fn next_account_test_event_chunk( + receiver: crate::account_test::AccountTestEventSubscription, +) -> Option<(crate::account_test::AccountTestEventSubscription, 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: crate::account_test::AccountTestEventSubscription, + pending: Vec, + pending_offset: usize, + opened: bool, +} + +impl AccountTestEventStream { + fn new(receiver: crate::account_test::AccountTestEventSubscription) -> 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 test_id = account_test_id_from_query(request.url().split_once('?').map(|(_, query)| query)); + let Some(test_id) = test_id else { + let _ = request.respond(Response::from_string("{}").with_status_code(400)); + return; + }; + let receiver = crate::account_test::subscribe_account_test_events(&test_id); + 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, + RawQuery(query): RawQuery, +) -> AxumResponse { + if !axum_rpc_token_valid(&headers) { + return (AxumStatusCode::UNAUTHORIZED, "{}").into_response(); + } + + let Some(test_id) = account_test_id_from_query(query.as_deref()) else { + return (AxumStatusCode::BAD_REQUEST, "{}").into_response(); + }; + let receiver = crate::account_test::subscribe_account_test_events(&test_id); + 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 +} + +#[cfg(test)] +mod tests { + use super::account_test_id_from_query; + + #[test] + fn account_test_event_query_requires_one_valid_test_id() { + assert_eq!( + account_test_id_from_query(Some( + "other=value&testId=550e8400-e29b-41d4-a716-446655440000" + )) + .as_deref(), + Some("550e8400-e29b-41d4-a716-446655440000") + ); + assert!(account_test_id_from_query(None).is_none()); + assert!(account_test_id_from_query(Some("other=value")).is_none()); + assert!(account_test_id_from_query(Some("testId=%20%20")).is_none()); + assert!(account_test_id_from_query(Some("testId=bad%2Fid")).is_none()); + assert!(account_test_id_from_query(Some("testId=one&testId=two")).is_none()); + } +} 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..138898ad7 100644 --- a/crates/service/src/http/mod.rs +++ b/crates/service/src/http/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod account_test_events; pub mod callback_endpoint; pub mod gateway_endpoint; pub mod rpc_endpoint; 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..1237a4173 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::{ @@ -101,6 +102,7 @@ pub(crate) use usage::scheduler as usage_scheduler; pub(crate) use usage::snapshot_store as usage_snapshot_store; pub(crate) use usage::token_refresh as usage_token_refresh; +pub use account_test::{set_account_test_event_handler, AccountTestEvent}; pub use app_settings::{ app_settings_get, app_settings_get_with_overrides, app_settings_set, author_content_get, bind_all_interfaces_enabled, bind_all_interfaces_enabled_for_mode, 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..c6d4b57d7 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,33 @@ pub(super) fn try_handle(req: &JsonRpcRequest, actor: &RpcActor) -> Option { + if !actor.is_admin() { + super::value_or_error::(Err( + super::permission_denied("account/test"), + )) + } else { + 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); + let test_id = first_str_param(req, &["testId", "test_id"]).map(str::to_string); + super::value_or_error(account_test::start_account_test( + account_id, model, prompt, kind, test_id, + )) + } + } + "account/test/cancel" => { + if !actor.is_admin() { + super::value_or_error::(Err(super::permission_denied("account/test/cancel"))) + } else { + let account_id = first_str_param(req, &["accountId", "account_id"]).unwrap_or(""); + let test_id = first_str_param(req, &["testId", "test_id"]).unwrap_or(""); + super::value_or_error(account_test::cancel_account_test(account_id, test_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)) @@ -560,4 +587,23 @@ mod tests { "permission_denied: account/update groupName" ); } + + #[test] + fn member_cannot_start_or_cancel_account_tests() { + let actor = RpcActor::from_parts(Some(crate::ROLE_MEMBER), Some("member-a")); + for method in ["account/test", "account/test/cancel"] { + let response = try_handle( + &rpc_request( + method, + serde_json::json!({ "accountId": "acc-a", "testId": "test-a" }), + ), + &actor, + ) + .expect("response"); + assert_eq!( + error_message(&response), + format!("permission_denied: {method}") + ); + } + } } diff --git a/crates/web/src/main.rs b/crates/web/src/main.rs index f63dbfdfb..21920c61b 100644 --- a/crates/web/src/main.rs +++ b/crates/web/src/main.rs @@ -524,6 +524,10 @@ async fn async_main() { "/api/events/usage-refresh", get(service_gateway::usage_refresh_events), ) + .route( + "/api/events/account-test", + get(service_gateway::account_test_events), + ) .route("/__quit", get(service_gateway::quit)); let disk_ok = ensure_index_file(&index); diff --git a/crates/web/src/service_gateway.rs b/crates/web/src/service_gateway.rs index 2437d09e8..ccdffe65b 100644 --- a/crates/web/src/service_gateway.rs +++ b/crates/web/src/service_gateway.rs @@ -439,6 +439,71 @@ pub(super) async fn usage_refresh_events(State(state): State>) -> out } +fn account_test_events_target_url(service_addr: &str, uri: &axum::http::Uri) -> String { + let mut target_url = format!("http://{}/events/account-test", service_addr.trim()); + if let Some(query) = uri.query().filter(|query| !query.is_empty()) { + target_url.push('?'); + target_url.push_str(query); + } + target_url +} + +fn account_test_events_role_allowed(web_auth_mode: &str, role: Option<&str>) -> bool { + web_auth_mode != "accounts" + || matches!( + role, + Some(codexmanager_service::ROLE_ADMIN | codexmanager_service::ROLE_SYSTEM_ADMIN) + ) +} + +pub(super) async fn account_test_events( + State(state): State>, + headers: HeaderMap, + uri: axum::http::Uri, +) -> Response { + let web_auth_mode = codexmanager_service::current_web_auth_mode(); + let session = auth::current_app_session_from_headers(&headers); + if !account_test_events_role_allowed( + &web_auth_mode, + session.as_ref().map(|session| session.user.role.as_str()), + ) { + return (StatusCode::FORBIDDEN, "{}").into_response(); + } + + let target_url = account_test_events_target_url(&state.service_addr, &uri); + let resp = state + .client + .get(&target_url) + .header("accept", "text/event-stream") + .header("x-codexmanager-rpc-token", &state.rpc_token) + .send() + .await; + let resp = match resp { + Ok(value) => value, + Err(err) => { + let msg = format_upstream_error_message(state.service_addr.as_str(), &err); + return (StatusCode::BAD_GATEWAY, msg).into_response(); + } + }; + + let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); + let mut out = Response::new(Body::from_stream(resp.bytes_stream())); + *out.status_mut() = status; + out.headers_mut().insert( + "content-type", + axum::http::HeaderValue::from_static("text/event-stream"), + ); + out.headers_mut().insert( + "cache-control", + axum::http::HeaderValue::from_static("no-cache"), + ); + out.headers_mut().insert( + "x-accel-buffering", + axum::http::HeaderValue::from_static("no"), + ); + out +} + const DEFAULT_GATEWAY_PROXY_MAX_BODY_BYTES: usize = 0; const ENV_GATEWAY_PROXY_MAX_BODY_BYTES: &str = "CODEXMANAGER_GATEWAY_PROXY_MAX_BODY_BYTES"; diff --git a/crates/web/src/service_gateway_tests.rs b/crates/web/src/service_gateway_tests.rs index f1d2a42bf..dd4aebcfe 100644 --- a/crates/web/src/service_gateway_tests.rs +++ b/crates/web/src/service_gateway_tests.rs @@ -1,4 +1,5 @@ use super::{ + account_test_events_role_allowed, account_test_events_target_url, format_upstream_error_message, gateway_proxy_max_body_bytes, gateway_proxy_target_url, service_probe_client, should_skip_gateway_request_header, should_skip_gateway_response_header, tcp_probe, ENV_GATEWAY_PROXY_MAX_BODY_BYTES, @@ -140,3 +141,101 @@ async fn rpc_proxy_rejects_body_over_the_bounded_upload_limit() { assert_eq!(response.status(), axum::http::StatusCode::PAYLOAD_TOO_LARGE); } + +/// 在独立线程里起一个最小 HTTP 服务，返回固定 SSE 响应后关闭连接，供代理测试当上游用。 +fn spawn_mock_sse_upstream(body: &'static str) -> String { + let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind mock upstream"); + let addr = listener.local_addr().expect("mock upstream addr"); + std::thread::spawn(move || { + use std::io::{Read, Write}; + let (mut socket, _) = listener.accept().expect("accept mock request"); + let mut buf = [0u8; 4096]; + let mut total = 0; + loop { + let n = socket.read(&mut buf[total..]).expect("read mock request"); + if n == 0 { + break; + } + total += n; + if buf[..total].windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n{body}" + ); + let _ = socket.write_all(response.as_bytes()); + let _ = socket.shutdown(std::net::Shutdown::Both); + }); + addr.to_string() +} + +#[tokio::test] +async fn account_test_events_proxies_sse_stream_from_service() { + let upstream = + spawn_mock_sse_upstream("event: account-test-event\ndata: {\"testId\":\"t1\"}\n\n"); + let (shutdown_tx, _shutdown_rx) = tokio::sync::watch::channel(false); + let state = std::sync::Arc::new(crate::AppState { + client: reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"), + service_rpc_url: "http://127.0.0.1:1/rpc".to_string(), + service_addr: upstream, + rpc_token: "test-token".to_string(), + web_auth_session_key: "test-session".to_string(), + shutdown_tx, + spawned_service: std::sync::Arc::new(tokio::sync::Mutex::new(false)), + missing_ui_html: std::sync::Arc::new(String::new()), + }); + + let uri: Uri = "/api/events/account-test?testId=t1" + .parse() + .expect("valid account test URI"); + let response = super::account_test_events(State(state), HeaderMap::new(), uri).await; + + assert_eq!(response.status(), axum::http::StatusCode::OK); + assert_eq!( + response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()), + Some("text/event-stream") + ); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("read body"); + let text = String::from_utf8_lossy(&body); + assert!( + text.contains("account-test-event"), + "unexpected body: {text}" + ); + assert!(text.contains("t1"), "unexpected body: {text}"); +} + +#[test] +fn account_test_events_require_admin_in_accounts_mode() { + assert!(account_test_events_role_allowed("none", None)); + assert!(account_test_events_role_allowed("password", None)); + assert!(account_test_events_role_allowed("accounts", Some("admin"))); + assert!(account_test_events_role_allowed( + "accounts", + Some("system_admin") + )); + assert!(!account_test_events_role_allowed( + "accounts", + Some("member") + )); + assert!(!account_test_events_role_allowed("accounts", None)); +} + +#[test] +fn account_test_events_target_preserves_test_id_query() { + let uri: Uri = "/api/events/account-test?testId=550e8400-e29b-41d4-a716-446655440000" + .parse() + .expect("valid URI"); + assert_eq!( + account_test_events_target_url("127.0.0.1:48760", &uri), + "http://127.0.0.1:48760/events/account-test?testId=550e8400-e29b-41d4-a716-446655440000" + ); +}