diff --git a/core/auth.py b/core/auth.py index 20c5762..f665890 100644 --- a/core/auth.py +++ b/core/auth.py @@ -58,17 +58,35 @@ def get_queue_status() -> dict: # ---------------------------------------------------------------------- # 账号选择与释放 - 轮询(Round-Robin)策略 # ---------------------------------------------------------------------- -def choose_new_account(exclude_ids=None): +def choose_new_account(exclude_ids=None, target_id=None): """轮询选择策略: 1. 使用线程锁保证并发安全 - 2. 优先选择队首的有 token 账号 - 3. 从队列头部取出账号(FIFO) - 4. 请求完成后调用 release_account 将账号放回队尾 + 2. 如果指定了 target_id,优先尝试获取该账号 + 3. 优先选择队首的有 token 账号 + 4. 从队列头部取出账号(FIFO) + 5. 请求完成后调用 release_account 将账号放回队尾 """ if exclude_ids is None: exclude_ids = [] with _queue_lock: + # 0. 如果指定了目标账号,优先尝试获取 + if target_id: + for i in range(len(account_queue)): + acc = account_queue[i] + acc_id = get_account_identifier(acc) + if acc_id == target_id: + selected = account_queue.pop(i) + in_use_accounts[acc_id] = selected + logger.info(f"[choose_new_account] 指定选择: {acc_id} | 队列剩余: {len(account_queue)}") + return selected + # 如果队列中没找到,且不在 in_use 中,说明账号不存在 + if target_id not in in_use_accounts: + logger.warning(f"[choose_new_account] 指定账号不存在: {target_id}") + else: + logger.warning(f"[choose_new_account] 指定账号正忙: {target_id}") + return None + # 第一轮:优先选择已有 token 的账号 for i in range(len(account_queue)): acc = account_queue[i] @@ -145,11 +163,17 @@ def determine_mode_and_token(request: Request): if caller_key in config_keys: request.state.use_config_token = True request.state.tried_accounts = [] # 初始化已尝试账号 - selected_account = choose_new_account() + + target_account = request.headers.get("X-Ds2-Target-Account") + selected_account = choose_new_account(target_id=target_account) + if not selected_account: + detail_msg = "No accounts configured or all accounts are busy." + if target_account: + detail_msg = f"Target account {target_account} is busy or not found." raise HTTPException( status_code=429, - detail="No accounts configured or all accounts are busy.", + detail=detail_msg, ) if not selected_account.get("token", "").strip(): try: diff --git a/webui/src/components/ApiTester.jsx b/webui/src/components/ApiTester.jsx index 017e7f6..b464aec 100644 --- a/webui/src/components/ApiTester.jsx +++ b/webui/src/components/ApiTester.jsx @@ -14,7 +14,9 @@ import { ChevronDown, ShieldCheck, Terminal, - Zap + Zap, + ToggleLeft, + ToggleRight } from 'lucide-react' import clsx from 'clsx' import { useI18n } from '../i18n' @@ -31,6 +33,7 @@ export default function ApiTester({ config, onMessage, authFetch }) { const [streamingContent, setStreamingContent] = useState('') const [streamingThinking, setStreamingThinking] = useState('') const [isStreaming, setIsStreaming] = useState(false) + const [streamingMode, setStreamingMode] = useState(true) const abortControllerRef = useRef(null) const defaultMessageRef = useRef(defaultMessage) @@ -55,7 +58,7 @@ export default function ApiTester({ config, onMessage, authFetch }) { setIsStreaming(false) } - const directTest = async () => { + const runTest = async () => { if (loading) return setLoading(true) @@ -75,69 +78,78 @@ export default function ApiTester({ config, onMessage, authFetch }) { return } + const headers = { + 'Content-Type': 'application/json', + 'Authorization': `Bearer ${key}`, + } + if (selectedAccount) { + headers['X-Ds2-Target-Account'] = selectedAccount + } + const res = await fetch('/v1/chat/completions', { method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'Authorization': `Bearer ${key}`, - }, + headers, body: JSON.stringify({ model, messages: [{ role: 'user', content: message }], - stream: true, + stream: streamingMode, }), signal: abortControllerRef.current.signal, }) if (!res.ok) { - const data = await res.json() - setResponse({ success: false, error: data.error?.message || t('apiTester.requestFailed') }) - onMessage('error', data.error?.message || t('apiTester.requestFailed')) + const data = await res.json().catch(() => ({})) + const errorMsg = data.error?.message || t('apiTester.requestFailed') + setResponse({ success: false, error: errorMsg }) + onMessage('error', errorMsg) setLoading(false) setIsStreaming(false) return } - setResponse({ success: true, status_code: res.status }) + if (streamingMode) { + setResponse({ success: true, status_code: res.status }) - const reader = res.body.getReader() - const decoder = new TextDecoder() - let buffer = '' + const reader = res.body.getReader() + const decoder = new TextDecoder() + let buffer = '' - while (true) { - const { done, value } = await reader.read() - if (done) break + while (true) { + const { done, value } = await reader.read() + if (done) break - buffer += decoder.decode(value, { stream: true }) - const lines = buffer.split('\n') - buffer = lines.pop() || '' + buffer += decoder.decode(value, { stream: true }) + const lines = buffer.split('\n') + buffer = lines.pop() || '' - for (const line of lines) { - const trimmed = line.trim() - if (!trimmed || !trimmed.startsWith('data: ')) continue + for (const line of lines) { + const trimmed = line.trim() + if (!trimmed || !trimmed.startsWith('data: ')) continue - const dataStr = trimmed.slice(6) - if (dataStr === '[DONE]') continue + const dataStr = trimmed.slice(6) + if (dataStr === '[DONE]') continue - try { - const json = JSON.parse(dataStr) - console.log('[ApiTester] Parsed JSON:', json) - const choice = json.choices?.[0] - if (choice?.delta) { - const delta = choice.delta - console.log('[ApiTester] Delta:', delta) - if (delta.reasoning_content) { - setStreamingThinking(prev => prev + delta.reasoning_content) - } - if (delta.content) { - console.log('[ApiTester] Content:', delta.content) - setStreamingContent(prev => prev + delta.content) + try { + const json = JSON.parse(dataStr) + const choice = json.choices?.[0] + if (choice?.delta) { + const delta = choice.delta + if (delta.reasoning_content) { + setStreamingThinking(prev => prev + delta.reasoning_content) + } + if (delta.content) { + setStreamingContent(prev => prev + delta.content) + } } + } catch (e) { + console.error('Invalid JSON hunk:', dataStr, e) } - } catch (e) { - console.error('Invalid JSON hunk:', dataStr, e) } } + } else { + const data = await res.json() + setResponse({ success: true, status_code: res.status, ...data }) + onMessage('success', t('apiTester.testSuccess', { account: selectedAccount || 'Auto', time: 'N/A' })) } } catch (e) { if (e.name === 'AbortError') { @@ -153,260 +165,235 @@ export default function ApiTester({ config, onMessage, authFetch }) { } } - const sendTest = async () => { - // 清除上次的流式/思考内容,防止残留 - setStreamingContent('') - setStreamingThinking('') +useEffect(() => { + setMessage((prev) => (prev === defaultMessageRef.current ? defaultMessage : prev)) + defaultMessageRef.current = defaultMessage +}, [defaultMessage]) - if (selectedAccount) { - setLoading(true) - setResponse(null) - try { - const res = await apiFetch('/admin/accounts/test', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - identifier: selectedAccount, - model, - message, - }), - }) - const data = await res.json() - setResponse({ - success: data.success, - status_code: res.status, - response: data, - account: selectedAccount, - }) - if (data.success) { - onMessage('success', t('apiTester.testSuccess', { account: selectedAccount, time: data.response_time })) - } else { - onMessage('error', `${selectedAccount}: ${data.message}`) - } - } catch (e) { - onMessage('error', t('apiTester.networkError', { error: e.message })) - setResponse({ error: e.message }) - } finally { - setLoading(false) - } - return - } - - directTest() - } - - useEffect(() => { - setMessage((prev) => (prev === defaultMessageRef.current ? defaultMessage : prev)) - defaultMessageRef.current = defaultMessage - }, [defaultMessage]) - - return ( -