fix(gateway): add provider key concurrency and cache affinity modes

2026-10-08 02:17:46 +08:00 · 2026-09-01 08:06:58 +08:00
parent 36daba7a34
commit 9631b229b3
26 changed files with 1176 additions and 154 deletions
@@ -20,6 +20,7 @@
      >
        {{ accountQuotaText }}
      </div>
+      <ResetCredits />
    </div>
    <div
      v-else-if="accountQuotaText || fallbackText"
@@ -47,6 +48,7 @@
      >
        {{ accountQuotaText }}
      </div>
+      <ResetCredits />
    </div>
    <span
      v-else-if="accountQuotaText || fallbackText"
@@ -72,23 +74,61 @@ export interface PoolQuotaProgressDisplayItem {
  meterText: string
  barClass: string
  meterClass: string
+  numericOnly?: boolean
 }

-withDefaults(defineProps<{
+const props = withDefaults(defineProps<{
  items: PoolQuotaProgressDisplayItem[]
  accountQuotaText?: string | null
  fallbackText?: string | null
  textClass?: string
  variant?: 'desktop' | 'mobile'
+  resetCreditText?: string | null
+  resetCreditItems?: string[]
+  canConsumeResetCredit?: boolean
+  consumingResetCredit?: boolean
 }>(), {
  accountQuotaText: null,
  fallbackText: null,
  textClass: '',
  variant: 'desktop',
+  resetCreditText: null,
+  resetCreditItems: () => [],
+  canConsumeResetCredit: false,
+  consumingResetCredit: false,
 })

+const emit = defineEmits<{
+  'consume-reset-credit': []
+}>()
+
 const { legacyT } = useI18n()

+const ResetCredits = defineComponent({
+  name: 'PoolQuotaResetCredits',
+  setup() {
+    return () => props.resetCreditText ? h('div', {
+      'data-testid': 'pool-quota-reset-credits',
+      class: 'mt-2 border-t border-border/50 pt-1.5 text-[10px] leading-4 text-muted-foreground',
+    }, [
+      h('div', { class: 'flex flex-wrap items-center gap-x-1' }, [
+        props.canConsumeResetCredit
+          ? h('button', {
+            type: 'button',
+            disabled: props.consumingResetCredit,
+            class: 'font-medium text-primary hover:underline disabled:pointer-events-none disabled:opacity-60',
+            onClick: () => emit('consume-reset-credit'),
+          }, props.consumingResetCredit ? legacyT('重置中...') : legacyT('点击以进行重置'))
+          : null,
+        h('span', props.resetCreditText),
+      ]),
+      props.resetCreditItems.length
+        ? h('div', { class: 'truncate tabular-nums', title: props.resetCreditItems.join(' · ') }, props.resetCreditItems.join(' · '))
+        : null,
+    ]) : null
+  },
+})
+
 const QuotaProgressRows = defineComponent({
  name: 'QuotaProgressRows',
  props: {
@@ -122,15 +162,24 @@ const QuotaProgressRows = defineComponent({
          : null,
      ]),
      h('div', { class: 'flex items-center gap-1.5' }, [
-        h('div', { class: 'relative flex-1 h-1.5 rounded-full bg-border overflow-hidden' }, [
-          h('div', {
-            class: ['absolute left-0 top-0 h-full rounded-full transition-all duration-300', item.barClass],
-            style: { width: `${item.remainingPercent}%` },
-          }),
-        ]),
+        item.numericOnly
+          ? null
+          : h('div', {
+            'data-testid': 'pool-quota-progress-track',
+            class: 'relative flex-1 h-1.5 rounded-full bg-border overflow-hidden',
+          }, [
+            h('div', {
+              class: ['absolute left-0 top-0 h-full rounded-full transition-all duration-300', item.barClass],
+              style: { width: `${item.remainingPercent}%` },
+            }),
+          ]),
        h('span', {
          'data-testid': 'pool-quota-meter-text',
-          class: ['shrink-0 text-[10px] font-medium tabular-nums leading-none', item.meterClass],
+          class: [
+            'shrink-0 text-[10px] font-medium tabular-nums leading-none',
+            item.numericOnly ? 'ml-auto' : '',
+            item.meterClass,
+          ],
        }, item.meterText),
      ]),
    ]))
@@ -48,6 +48,27 @@
          >
            {{ activeDistributionDesc }}
          </p>
+          <div
+            v-if="activeDistributionItem?.modeOptions.length"
+            data-testid="pool-cache-affinity-secondary-mode"
+            class="mt-2 flex w-fit flex-wrap gap-1 rounded-lg bg-muted/50 p-1"
+          >
+            <button
+              v-for="modeOpt in activeDistributionItem.modeOptions"
+              :key="modeOpt.value"
+              type="button"
+              :data-mode="modeOpt.value"
+              class="rounded-md px-2.5 py-1 text-xs font-medium transition-all"
+              :class="[
+                activeDistributionItem.mode === modeOpt.value
+                  ? 'bg-primary text-primary-foreground shadow-sm'
+                  : 'text-muted-foreground hover:bg-background/70 hover:text-foreground'
+              ]"
+              @click="setPresetModeByPreset(activeDistributionItem.preset, modeOpt.value)"
+            >
+              {{ modeOpt.label }}
+            </button>
+          </div>
        </div>
      </div>

@@ -267,13 +288,16 @@ const FALLBACK_PRESET_DEFS: PoolPresetMeta[] = [
  {
    name: 'cache_affinity',
    label: '缓存亲和',
-    description: '优先复用最近使用过的 Key，利用 Prompt Caching',
+    description: '同一用户持续复用 Key，首次分配可集中或轮转',
    mutex_group: DISTRIBUTION_GROUP,
-    evidence_hint: '依据 LRU 时间戳（最近使用优先，与 LRU 轮转相反）',
+    evidence_hint: '先复用用户粘性 Key，未命中时按所选二级模式分配',
    providers: [],
    default_enabled: true,
-    modes: null,
-    default_mode: null,
+    modes: [
+      { value: 'single_account', label: '单号优先' },
+      { value: 'lru', label: 'LRU 轮号' },
+    ],
+    default_mode: 'single_account',
  },
  {
    name: 'lru',
@@ -733,14 +757,17 @@ const activeDistributionPreset = computed(() => {
  return found?.item.preset ?? null
 })

-const activeDistributionDesc = computed(() => {
+const activeDistributionItem = computed(() => {
  const found = distributionItems.value.find(({ item }) => item.enabled && item.applicable)
-  return found?.item.desc ?? null
+  return found?.item ?? null
+})
+
+const activeDistributionDesc = computed(() => {
+  return activeDistributionItem.value?.desc ?? null
 })

 const activeDistributionLabel = computed(() => {
-  const found = distributionItems.value.find(({ item }) => item.enabled && item.applicable)
-  return found?.item.label ?? null
+  return activeDistributionItem.value?.label ?? null
 })

 const strategyItems = computed(() => {
@@ -90,6 +90,30 @@ describe('pool key display panels', () => {
    root.remove()
  })

+  it('renders Antigravity quota as numeric values without progress tracks', () => {
+    const root = document.createElement('div')
+    document.body.appendChild(root)
+    const app = createApp(PoolKeyQuotaPanel, {
+      items: [{
+        label: 'Gemini 3.1 Pro (High)',
+        remainingPercent: 42,
+        resetText: '1h 后重置',
+        meterText: '42',
+        barClass: 'bg-amber-500',
+        meterClass: 'text-amber-600',
+        numericOnly: true,
+      }],
+    })
+    app.use(createI18n())
+    app.mount(root)
+
+    expect(root.querySelector('[data-testid="pool-quota-meter-text"]')?.textContent).toBe('42')
+    expect(root.querySelector('[data-testid="pool-quota-progress-track"]')).toBeNull()
+
+    app.unmount()
+    root.remove()
+  })
+
  it('renders single-cycle stats as plain text', () => {
    const root = document.createElement('div')
    document.body.appendChild(root)
@@ -0,0 +1,125 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { createApp, defineComponent, h, nextTick, ref, type App } from 'vue'
+
+import PoolSchedulingDialog from '../PoolSchedulingDialog.vue'
+
+const endpointMocks = vi.hoisted(() => ({
+  getPoolSchedulingPresets: vi.fn(),
+  getProvider: vi.fn(),
+  updateProvider: vi.fn(),
+}))
+
+vi.mock('@/api/endpoints', () => ({
+  getProvider: endpointMocks.getProvider,
+  updateProvider: endpointMocks.updateProvider,
+}))
+
+vi.mock('@/api/endpoints/pool', () => ({
+  getPoolSchedulingPresets: endpointMocks.getPoolSchedulingPresets,
+}))
+
+vi.mock('@/composables/useToast', () => ({
+  useToast: () => ({
+    success: vi.fn(),
+    error: vi.fn(),
+  }),
+}))
+
+const mountedApps: Array<{ app: App, root: HTMLElement }> = []
+const cacheAffinityConfig = {
+  scheduling_presets: [{
+    preset: 'cache_affinity',
+    enabled: true,
+    mode: 'single_account',
+  }],
+}
+
+async function settle(): Promise<void> {
+  for (let index = 0; index < 4; index += 1) {
+    await Promise.resolve()
+    await nextTick()
+  }
+}
+
+function mountDialog(): void {
+  const root = document.createElement('div')
+  document.body.appendChild(root)
+  const TestHost = defineComponent({
+    setup() {
+      const open = ref(false)
+      void nextTick(() => { open.value = true })
+      return () => h(PoolSchedulingDialog, {
+        modelValue: open.value,
+        providerId: 'provider-1',
+        providerType: 'openai',
+        currentConfig: cacheAffinityConfig,
+        'onUpdate:modelValue': (value: boolean) => { open.value = value },
+      })
+    },
+  })
+  const app = createApp(TestHost)
+  app.mount(root)
+  mountedApps.push({ app, root })
+}
+
+beforeEach(() => {
+  endpointMocks.getPoolSchedulingPresets.mockReset()
+  endpointMocks.getProvider.mockReset()
+  endpointMocks.updateProvider.mockReset()
+  endpointMocks.getPoolSchedulingPresets.mockResolvedValue([{
+    name: 'cache_affinity',
+    label: '缓存亲和',
+    description: '同一用户持续复用 Key，首次分配可集中或轮转',
+    providers: [],
+    default_enabled: true,
+    modes: [
+      { value: 'single_account', label: '单号优先' },
+      { value: 'lru', label: 'LRU 轮号' },
+    ],
+    default_mode: 'single_account',
+    mutex_group: 'distribution_mode',
+  }])
+  endpointMocks.getProvider.mockResolvedValue({
+    id: 'provider-1',
+    pool_advanced: cacheAffinityConfig,
+  })
+  endpointMocks.updateProvider.mockResolvedValue({ id: 'provider-1' })
+})
+
+afterEach(() => {
+  for (const { app, root } of mountedApps.splice(0)) {
+    app.unmount()
+    root.remove()
+  }
+  document.body.innerHTML = ''
+})
+
+describe('PoolSchedulingDialog cache affinity modes', () => {
+  it('shows and saves the LRU secondary mode', async () => {
+    mountDialog()
+    await settle()
+
+    const modeControl = document.body.querySelector(
+      '[data-testid="pool-cache-affinity-secondary-mode"]',
+    )
+    expect(modeControl?.textContent).toContain('单号优先')
+    expect(modeControl?.textContent).toContain('LRU 轮号')
+
+    modeControl?.querySelector<HTMLButtonElement>('[data-mode="lru"]')?.click()
+    await nextTick()
+    const saveButton = [...document.body.querySelectorAll<HTMLButtonElement>('button')]
+      .find(button => button.textContent?.trim() === '保存')
+    saveButton?.click()
+    await settle()
+
+    expect(endpointMocks.updateProvider).toHaveBeenCalledWith('provider-1', {
+      pool_advanced: {
+        scheduling_presets: [{
+          preset: 'cache_affinity',
+          enabled: true,
+          mode: 'lru',
+        }],
+      },
+    })
+  })
+})
@@ -35,6 +35,18 @@ describe('pool key batch settings', () => {
    })
  })

+  it('serializes a positive concurrency limit as a number', () => {
+    const selection = createPoolKeyBatchSettingSelection()
+    const draft = createPoolKeyBatchSettingsDraft()
+    selection.concurrent_limit = true
+    draft.concurrent_limit = 6
+
+    expect(validatePoolKeyBatchSettings(selection, draft)).toEqual([])
+    expect(buildPoolKeySettingsPatch(selection, draft)).toEqual({
+      concurrent_limit: 6,
+    })
+  })
+
  it('requires a selected field and a proxy node for set mode', () => {
    const selection = createPoolKeyBatchSettingSelection()
    const draft = createPoolKeyBatchSettingsDraft()
@@ -5,24 +5,34 @@ export function mergePoolKeyQuotaSnapshots(
  keys: PoolKeyDetail[],
  results: RefreshQuotaResult['results'],
 ): PoolKeyDetail[] {
-  const quotaByKeyId = new Map<string, NonNullable<RefreshQuotaResult['results'][number]['quota_snapshot']>>()
+  const resultByKeyId = new Map<string, RefreshQuotaResult['results'][number]>()
  for (const result of results) {
-    if (result.quota_snapshot) {
-      quotaByKeyId.set(result.key_id, result.quota_snapshot)
+    if (result.quota_snapshot || result.metadata) {
+      resultByKeyId.set(result.key_id, result)
    }
  }
-  if (quotaByKeyId.size === 0) return keys
+  if (resultByKeyId.size === 0) return keys

  return keys.map((key) => {
-    const quotaSnapshot = quotaByKeyId.get(key.key_id)
-    if (!quotaSnapshot) return key
+    const result = resultByKeyId.get(key.key_id)
+    if (!result) return key
+    const quotaSnapshot = result.quota_snapshot
+    const providerType = String(quotaSnapshot?.provider_type || key.provider_type || '').trim().toLowerCase()
    return {
      ...key,
-      quota_updated_at: quotaSnapshot.updated_at ?? quotaSnapshot.observed_at ?? key.quota_updated_at ?? null,
-      status_snapshot: {
-        ...(key.status_snapshot ?? {}),
-        quota: quotaSnapshot,
-      },
+      ...(quotaSnapshot ? {
+        quota_updated_at: quotaSnapshot.updated_at ?? quotaSnapshot.observed_at ?? key.quota_updated_at ?? null,
+        status_snapshot: {
+          ...(key.status_snapshot ?? {}),
+          quota: quotaSnapshot,
+        },
+      } : {}),
+      ...(result.metadata && providerType ? {
+        upstream_metadata: {
+          ...(key.upstream_metadata ?? {}),
+          [providerType]: result.metadata,
+        },
+      } : {}),
    }
  })
 }