From 30b29e51e82c5e8c0ba4fd6adf0636501ad0ffba Mon Sep 17 00:00:00 2001 From: alfadb Date: Sun, 30 Aug 2026 16:38:11 +0800 Subject: [PATCH] =?UTF-8?q?feat(ollama):=20=E7=94=A8=E9=87=8F=E7=AA=97?= =?UTF-8?q?=E5=8F=A3=E6=94=AF=E6=8C=81=E6=8C=82=E5=9C=A8=E5=9B=BD=E4=BA=A7?= =?UTF-8?q?=E4=B8=89=E5=AE=B6=E5=B9=B3=E5=8F=B0=E4=B8=8B=E7=9A=84=20Ollama?= =?UTF-8?q?=20Cloud=20=E8=B4=A6=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ollama Cloud 用量窗口（浏览器会话、5h/7d 快照、自动刷新、同 key 组共享）此前 只认 platform ∈ {openai, anthropic} 的 apikey 账号。但网关对 kimi/zhipu/deepseek 的 apikey 账号同样按 credentials.base_url 转发，用户把 Ollama Cloud 的 key 挂在 这三个平台下跑 Ollama Cloud 托管的 glm/kimi/deepseek 模型时，请求转发正常，用量 窗口却完全不可用。 - 平台白名单放开到 {openai, anthropic, kimi, zhipu, deepseek}；type='apikey' 与官方 ollama.com base_url 的限制不变 - Go 侧收敛为 isOllamaCloudUsagePlatform，SQL 侧收敛为 ollamaCloudUsagePlatformsSQL，6 处白名单只引用单一来源，并补两侧互为镜像的 漂移守卫测试 - CN 周期余额/额度探测跳过 ollama 用量账号：其探测端点由 base_url 衍生，对 ollama.com 会被出站 URL 白名单拒绝，每轮都白跑并产生噪声 - 前端账号列表 CN 分支在 eligible 时改渲染 Ollama 用量窗口，并跳过 CN 额度/ 余额子单元格（同上，渲染出来只会是一行探测报错） --- backend/internal/repository/account_repo.go | 8 +- .../account_repo_ollama_cloud_usage.go | 8 +- ...epo_ollama_cloud_usage_integration_test.go | 92 ++++++++++++++ .../account_repo_ollama_cloud_usage_test.go | 71 ++++++++++- backend/internal/repository/proxy_repo.go | 2 +- .../proxy_repo_upstream_billing_probe_test.go | 4 +- .../cn_provider_balance_check_service.go | 7 + .../cn_provider_balance_check_service_test.go | 47 +++++++ .../internal/service/ollama_cloud_usage.go | 16 ++- .../service/ollama_cloud_usage_test.go | 19 ++- .../components/account/AccountUsageCell.vue | 11 +- .../__tests__/AccountUsageCell.spec.ts | 120 ++++++++++++++++++ 12 files changed, 390 insertions(+), 15 deletions(-) diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 664a845460..ec705a0877 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -626,8 +626,8 @@ func lockAndMergeAccountProbeExtra( AND credentials = $4::jsonb AND proxy_id IS NOT DISTINCT FROM $5, COALESCE( - platform IN ('openai', 'anthropic') - AND $2 IN ('openai', 'anthropic') + platform IN (`+ollamaCloudUsagePlatformsSQL+`) + AND $2 IN (`+ollamaCloudUsagePlatformsSQL+`) AND type = 'apikey' AND $3 = 'apikey' AND credentials -> 'api_key' IS NOT DISTINCT FROM $4::jsonb -> 'api_key' @@ -814,7 +814,7 @@ func (r *accountRepository) UpdateCredentials(ctx context.Context, id int64, cre extra = CASE -- 凭证整体未变化 ⇒ Ollama 组身份必然未变化；顶层 DISTINCT 守卫防止 -- 非 Ollama 账号的无变化持久化误清探测快照或重写 NULL extra。 - WHEN platform IN ('openai', 'anthropic') + WHEN platform IN (`+ollamaCloudUsagePlatformsSQL+`) AND type = 'apikey' AND credentials IS DISTINCT FROM $1::jsonb AND ( @@ -2935,7 +2935,7 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates extraExpression = "(" + extraExpression + ") - 'ollama_cloud_usage_snapshot'" } } - eligibleAccount := "platform IN ('openai', 'anthropic') AND type = 'apikey'" + eligibleAccount := "platform IN (" + ollamaCloudUsagePlatformsSQL + ") AND type = 'apikey'" groupIdentityChanged := "" if len(ollamaGroupIdentityChanges) > 0 { groupIdentityChanged = "(" + eligibleAccount + " AND (" + joinClauses(ollamaGroupIdentityChanges, " OR ") + "))" diff --git a/backend/internal/repository/account_repo_ollama_cloud_usage.go b/backend/internal/repository/account_repo_ollama_cloud_usage.go index 4726805ae6..55634765a4 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage.go @@ -15,8 +15,12 @@ const ( ollamaCloudBaseURLRegexSQL = `^[hH][tT][tT][pP][sS]://([wW][wW][wW]\.)?[oO][lL][lL][aA][mM][aA]\.[cC][oO][mM](:443)?(/v1)?$` ollamaCloudBaseURLMatchSQLPrefix = "btrim(" ollamaCloudBaseURLMatchSQLSuffix = ") ~ '" + ollamaCloudBaseURLRegexSQL + "'" - ollamaCloudUsageEligibleSQL = ` - platform IN ('openai', 'anthropic') + // ollamaCloudUsagePlatformsSQL 是 service.isOllamaCloudUsagePlatform 的 SQL + // 镜像：Ollama Cloud key 允许挂在 openai/anthropic 与国产 OpenAI 兼容平台 + // 下复用。所有平台白名单 SQL 只允许引用本常量，不得各处重写字面量，防止漂移。 + ollamaCloudUsagePlatformsSQL = "'openai', 'anthropic', 'kimi', 'zhipu', 'deepseek'" + ollamaCloudUsageEligibleSQL = ` + platform IN (` + ollamaCloudUsagePlatformsSQL + `) AND type = 'apikey' AND ` + ollamaCloudBaseURLMatchSQLPrefix + `credentials ->> 'base_url'` + ollamaCloudBaseURLMatchSQLSuffix + ` AND jsonb_typeof(credentials -> 'api_key') = 'string' diff --git a/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go b/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go index 868901584b..71fbc0b4c7 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage_integration_test.go @@ -384,6 +384,98 @@ func accountIDs(accounts []service.Account) []int64 { return ids } +// 平台白名单放开后，官方 ollama.com key 挂在国产 OpenAI 兼容平台下同样进用量 +// 窗口：组写入跨 kimi/zhipu/deepseek 共享，白名单外平台（gemini）不得被卷入； +// 组身份守卫（lockAndMerge）与 due 列表也必须识别 CN 平台的行。 +func TestOllamaCloudUsageEligibilityExtendsToCNOpenAICompatPlatforms(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil) + now := time.Now().UTC() + activity := now.Add(-time.Minute) + create := func(name, platform, baseURL string) *service.Account { + t.Helper() + return mustCreateAccount(t, tx.Client(), &service.Account{ + Name: name, Platform: platform, Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "cn-shared-key", "base_url": baseURL}, + Extra: map[string]any{}, LastUsedAt: &activity, + }) + } + kimi := create("ollama-cn-kimi", service.PlatformKimi, "https://ollama.com") + zhipu := create("ollama-cn-zhipu", service.PlatformZhipu, "HTTPS://WWW.OLLAMA.COM:443/v1") + deepseek := create("ollama-cn-deepseek", service.PlatformDeepseek, "https://ollama.com/v1") + gemini := create("ollama-cn-gemini", service.PlatformGemini, "https://ollama.com") + + require.NoError(t, repo.SaveOllamaCloudUsageSession(ctx, kimi, "cipher:cn-shared", true)) + for _, id := range []int64{kimi.ID, zhipu.ID, deepseek.ID} { + account, err := repo.GetByID(ctx, id) + require.NoError(t, err) + require.Equal(t, "cipher:cn-shared", account.Extra[service.OllamaCloudUsageSessionExtraKey], account.Name) + require.Equal(t, true, account.Extra[service.OllamaCloudUsageAutoRefreshExtraKey], account.Name) + } + geminiLoaded, err := repo.GetByID(ctx, gemini.ID) + require.NoError(t, err) + require.NotContains(t, geminiLoaded.Extra, service.OllamaCloudUsageSessionExtraKey, + "用量窗口不随 base_url 放开到白名单外平台") + + // lockAndMerge 组身份守卫：CN 行凭证未变时必须保留 ollama 托管键。 + kimiLoaded, err := repo.GetByID(ctx, kimi.ID) + require.NoError(t, err) + merged, err := lockAndMergeAccountProbeExtra(ctx, tx.Client(), kimiLoaded, nil, nil) + require.NoError(t, err) + require.Equal(t, "cipher:cn-shared", merged[service.OllamaCloudUsageSessionExtraKey]) + require.Equal(t, true, merged[service.OllamaCloudUsageAutoRefreshExtraKey]) + + // due 列表识别 CN 行：给同 key 组挂上过期快照（fetched 2h 前、活动更新）， + // 每个 api_key 只返回一行，且必须来自 CN 组员。 + staleFetched := now.Add(-2 * time.Hour) + deepseekLoaded, err := repo.GetByID(ctx, deepseek.ID) + require.NoError(t, err) + require.NoError(t, repo.UpdateOllamaCloudUsageSnapshot(ctx, deepseekLoaded, &service.OllamaCloudUsageSnapshot{ + Status: service.OllamaCloudUsageStatusOK, + FetchedAt: &staleFetched, + LastAttemptAt: staleFetched, + })) + + due, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10) + + require.NoError(t, err) + require.Len(t, due, 1, "同 key 跨 CN 平台组按组去重后只应有一行 due") + require.Contains(t, []int64{kimi.ID, zhipu.ID, deepseek.ID}, due[0].ID, + "due 行必须来自 CN 平台的 ollama 组员") + require.NotNil(t, due[0].LastUsedAt) +} + +// 语义等价性端到端：普通（非 ollama）kimi apikey 账号改凭证落进 Ollama 分支后， +// probe 快照仍被清理，开关键与其它 extra 键不受影响。 +func TestUpdateCredentialsPlainCNAPIKeyAccountCleanupIsSemanticallyEquivalent(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil) + account := mustCreateAccount(t, tx.Client(), &service.Account{ + Name: "plain-kimi-apikey", Platform: service.PlatformKimi, Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "sk-moonshot", "base_url": "https://api.moonshot.cn"}, + Extra: map[string]any{ + service.UpstreamBillingProbeExtraKey: map[string]any{"status": "ok"}, + service.UpstreamBillingProbeEnabledExtraKey: true, + "custom_note": "keep-me", + }, + }) + + require.NoError(t, repo.UpdateCredentials(ctx, account.ID, map[string]any{ + "api_key": "sk-moonshot-rotated", "base_url": "https://api.moonshot.cn", + })) + + loaded, err := repo.GetByID(ctx, account.ID) + require.NoError(t, err) + require.NotContains(t, loaded.Extra, service.UpstreamBillingProbeExtraKey, + "probe 快照仍必须被清理") + require.Equal(t, true, loaded.Extra[service.UpstreamBillingProbeEnabledExtraKey], + "探测开关键不受清理影响") + require.Equal(t, "keep-me", loaded.Extra["custom_note"], "其它 extra 键不得误伤") + require.NotContains(t, loaded.Extra, service.OllamaCloudUsageSessionExtraKey) +} + func TestOllamaCloudUsageCredentialAndBulkUpdatesPreserveManagedStateOnlyWhenSafe(t *testing.T) { ctx := context.Background() tx := testEntTx(t) diff --git a/backend/internal/repository/account_repo_ollama_cloud_usage_test.go b/backend/internal/repository/account_repo_ollama_cloud_usage_test.go index 14512e885c..4fb3b6cce5 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage_test.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "regexp" + "strings" "testing" "time" @@ -178,7 +179,7 @@ func TestListOllamaCloudUsageGroupAccountsUsesOneStrictBatchQuery(t *testing.T) require.Empty(t, accounts) query := normalizeSQLWhitespace(capturedSQL) require.Contains(t, query, "credentials ->> 'api_key' = ANY($1)") - require.Contains(t, query, "platform IN ('openai', 'anthropic')") + require.Contains(t, query, "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek')") require.Contains(t, query, "jsonb_typeof(credentials -> 'api_key') = 'string'") require.Contains(t, query, ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")) require.NotContains(t, query, "~*") @@ -206,7 +207,7 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) { for _, clause := range []string{ "deleted_at IS NULL", "status = 'active'", - "platform IN ('openai', 'anthropic')", + "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek')", "type = 'apikey'", ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'"), "jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'", @@ -250,7 +251,7 @@ func TestBulkUpdateOllamaIdentityCleanupIsValueConditional(t *testing.T) { require.Contains(t, query, "NOT ("+ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'")) require.Contains(t, query, ollamaCloudBaseURLMatchesSQL("$1::jsonb ->> 'base_url'")) require.NotContains(t, query, "~*") - require.Contains(t, query, "platform IN ('openai', 'anthropic') AND type = 'apikey'") + require.Contains(t, query, "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek') AND type = 'apikey'") require.Contains(t, query, "- 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'") payload, ok := exec.execArgs[0][0].([]byte) require.True(t, ok) @@ -315,3 +316,67 @@ func TestUpdateCredentialsCleanupBranchRequiresChangedCredentials(t *testing.T) require.NoError(t, err) require.NoError(t, mock.ExpectationsWereMet()) } + +// SQL 平台白名单常量与 service 判定必须互为镜像：对每个已知平台，常量里的 +// 成员关系都要与 IsOllamaCloudUsageAccount（apikey + 官方 ollama.com）一致， +// 防止两侧平台列表各自漂移。 +func TestOllamaCloudUsagePlatformWhitelistMatchesServicePredicate(t *testing.T) { + matches := regexp.MustCompile(`'([^']+)'`).FindAllStringSubmatch(ollamaCloudUsagePlatformsSQL, -1) + sqlPlatforms := make(map[string]struct{}, len(matches)) + for _, match := range matches { + sqlPlatforms[match[1]] = struct{}{} + } + require.Len(t, sqlPlatforms, 5) + for _, platform := range []string{ + service.PlatformOpenAI, service.PlatformAnthropic, + service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, + service.PlatformGemini, service.PlatformGrok, service.PlatformAntigravity, + service.PlatformComposite, "kiro", + } { + account := ollamaCloudUsageRepositoryAccount() + account.Platform = platform + _, inSQL := sqlPlatforms[platform] + require.Equal(t, inSQL, service.IsOllamaCloudUsageAccount(account), platform) + } +} + +// 语义等价性：平台放开后，普通（非 ollama）kimi apikey 账号改凭证会从通用 +// apikey 分支落到 Ollama 分支——多减三个它本来就不存在的 ollama 键。分支必须 +// 仍然清掉 probe 快照，且不得减其它任何 extra 键。 +func TestUpdateCredentialsPlainCNAPIKeyAccountCleanupStaysSemanticallyEquivalent(t *testing.T) { + var capturedSQL string + matcher := sqlmock.QueryMatcherFunc(func(expectedSQL, actualSQL string) error { + if strings.Contains(actualSQL, "UPDATE accounts") { + capturedSQL = actualSQL + } + return sqlmock.QueryMatcherRegexp.Match(expectedSQL, actualSQL) + }) + db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(matcher)) + require.NoError(t, err) + t.Cleanup(func() { _ = db.Close() }) + client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db))) + t.Cleanup(func() { _ = client.Close() }) + mock.ExpectBegin() + mock.ExpectExec(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_session'.*- 'ollama_cloud_usage_auto_refresh'.*- 'ollama_cloud_usage_snapshot'`). + WithArgs(`{"api_key":"rotated-key","base_url":"https://api.moonshot.cn/v1"}`, int64(17)). + WillReturnResult(sqlmock.NewResult(0, 1)) + mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")). + WithArgs(service.SchedulerOutboxEventAccountChanged, int64(17), nil, nil, sqlmock.AnyArg()). + WillReturnResult(sqlmock.NewResult(1, 1)) + mock.ExpectCommit() + repo := newAccountRepositoryWithSQL(client, db, nil) + + err = repo.UpdateCredentials(context.Background(), 17, map[string]any{ + "api_key": "rotated-key", "base_url": "https://api.moonshot.cn/v1", + }) + + require.NoError(t, err) + query := normalizeSQLWhitespace(capturedSQL) + require.Contains(t, query, + "platform IN ('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek') AND type = 'apikey' AND credentials IS DISTINCT FROM $1::jsonb") + require.Contains(t, query, + "THEN COALESCE(extra, '{}'::jsonb) - 'upstream_billing_probe' - 'ollama_cloud_usage_session' - 'ollama_cloud_usage_auto_refresh' - 'ollama_cloud_usage_snapshot'") + require.NotContains(t, query, "- 'upstream_billing_probe_enabled'") + require.NotContains(t, query, "- 'upstream_billing_rate_sync_enabled'") + require.NoError(t, mock.ExpectationsWereMet()) +} diff --git a/backend/internal/repository/proxy_repo.go b/backend/internal/repository/proxy_repo.go index 4869721f02..679b22b26f 100644 --- a/backend/internal/repository/proxy_repo.go +++ b/backend/internal/repository/proxy_repo.go @@ -233,7 +233,7 @@ func invalidateProxyProbeSnapshots(ctx context.Context, exec sqlExecutor, proxyI AND ( (extra ? 'upstream_billing_probe' AND extra -> 'upstream_billing_probe' <> 'null'::jsonb) - OR (platform IN ('openai', 'anthropic') + OR (platform IN (`+ollamaCloudUsagePlatformsSQL+`) AND extra ? 'ollama_cloud_usage_snapshot' AND extra -> 'ollama_cloud_usage_snapshot' <> 'null'::jsonb) ) diff --git a/backend/internal/repository/proxy_repo_upstream_billing_probe_test.go b/backend/internal/repository/proxy_repo_upstream_billing_probe_test.go index 2fa420bf30..1f0ea83d34 100644 --- a/backend/internal/repository/proxy_repo_upstream_billing_probe_test.go +++ b/backend/internal/repository/proxy_repo_upstream_billing_probe_test.go @@ -33,7 +33,7 @@ func TestProxyUpdateInvalidatesBoundProbeSnapshotsAndEnqueuesOutboxAtomically(t WithArgs(int64(9)). WillReturnResult(sqlmock.NewResult(0, 0)) expectProxyUpdateReload(mock, 9, "new.example", "user", "pass") - mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`). + mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`). WithArgs(int64(9)). WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(17)).AddRow(int64(18))) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)")). @@ -76,7 +76,7 @@ func TestProxyUpdateRollsBackWhenProbeInvalidationOutboxFails(t *testing.T) { WithArgs(int64(9)). WillReturnResult(sqlmock.NewResult(0, 0)) expectProxyUpdateReload(mock, 9, "new.example", "", "") - mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`). + mock.ExpectQuery(`(?s)UPDATE accounts.*- 'upstream_billing_probe'.*- 'ollama_cloud_usage_snapshot'.*type = 'apikey'.*extra \? 'upstream_billing_probe'.*platform IN \('openai', 'anthropic', 'kimi', 'zhipu', 'deepseek'\).*extra \? 'ollama_cloud_usage_snapshot'.*RETURNING id`). WithArgs(int64(9)). WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(17))) mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload)")). diff --git a/backend/internal/service/cn_provider_balance_check_service.go b/backend/internal/service/cn_provider_balance_check_service.go index d5ddd6712e..f880605d68 100644 --- a/backend/internal/service/cn_provider_balance_check_service.go +++ b/backend/internal/service/cn_provider_balance_check_service.go @@ -113,6 +113,13 @@ func (s *CNProviderBalanceCheckService) runOnce() { if !account.IsActive() { continue } + // 挂在国产平台下、base_url 指向官方 ollama.com 的账号由 Ollama Cloud + // 用量窗口负责：CN 探测端点由 base_url 衍生，ollama.com 会被出站 + // URL 白名单拒绝（CN_BALANCE_URL_REJECTED），不跳过则每个周期都 + // 白跑并产生告警噪声。 + if IsOllamaCloudUsageAccount(account) { + continue + } // coding 账号：探测滚动窗口并落快照（不要求 Schedulable——已被 // 阈值停调的账号也需要新鲜快照决定是否续停）。 if account.IsCodingPlan() { diff --git a/backend/internal/service/cn_provider_balance_check_service_test.go b/backend/internal/service/cn_provider_balance_check_service_test.go index ddd59d5dcd..59eba6d311 100644 --- a/backend/internal/service/cn_provider_balance_check_service_test.go +++ b/backend/internal/service/cn_provider_balance_check_service_test.go @@ -2,6 +2,7 @@ package service import ( "context" + "errors" "sync" "testing" @@ -75,6 +76,52 @@ func TestCNProviderBalanceCheckRunOnceWithoutQuotaService(t *testing.T) { require.NotPanics(t, func() { svc.runOnce() }) } +// recordingCNBalanceLoadRepo 记录余额探测服务的 GetByID 调用：payg 账号进入 +// payg 检查队列必然触发 loadPayGAccount → GetByID，以此断言"未入队"这一内部 +// 状态。GetByID 直接报错，保证不发生真实外呼。 +type recordingCNBalanceLoadRepo struct { + AccountRepository + getByIDIDs []int64 +} + +func (r *recordingCNBalanceLoadRepo) GetByID(ctx context.Context, id int64) (*Account, error) { + r.getByIDIDs = append(r.getByIDIDs, id) + return nil, errors.New("not found") +} + +// 挂在国产平台下、base_url 指向官方 ollama.com 的账号由 Ollama Cloud 用量窗口 +// 负责：CN 探测端点由 base_url 衍生，ollama.com 会被出站 URL 白名单拒绝 +// （CN_BALANCE_URL_REJECTED），周期任务必须整体跳过——不进额度目标，也不进 +// payg 检查队列。 +func TestCNProviderBalanceCheckRunOnceSkipsOllamaCloudUsageAccounts(t *testing.T) { + // 对照组：普通 kimi coding 账号仍进额度探测。 + kimiCoding := Account{ID: 1, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive, + Credentials: map[string]any{"account_mode": "coding"}} + // ollama.com 挂 kimi：coding 模式不进额度目标。 + ollamaKimiCoding := Account{ID: 2, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive, + Credentials: map[string]any{"account_mode": "coding", "base_url": "https://ollama.com"}} + // ollama.com 挂 kimi：payg 模式不进 payg 检查队列。 + ollamaKimiPayg := Account{ID: 3, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, + Credentials: map[string]any{"base_url": "https://ollama.com", "api_key": "sk-ollama"}} + + loadRepo := &recordingCNBalanceLoadRepo{} + repo := &fakeCNCheckRepo{byPlatform: map[string][]Account{ + PlatformKimi: {kimiCoding, ollamaKimiCoding, ollamaKimiPayg}, + }} + prober := &fakeCNQuotaProber{} + svc := &CNProviderBalanceCheckService{ + accountRepo: repo, + balanceService: NewCNProviderBalanceService(loadRepo, nil, nil, &config.Config{}), + quotaService: prober, + cfg: &config.Config{}, + } + + svc.runOnce() + + require.Equal(t, []int64{1}, prober.probed) + require.Empty(t, loadRepo.getByIDIDs, "ollama 账号不得进入 payg 检查队列") +} + // 双币种（deepseek CNY+USD）停调判定：任一币种达标即不停调，全部低于阈值才停； // 无明细时退回主币种（兼容旧结果）。 func TestAllCNBalancesBelowThreshold(t *testing.T) { diff --git a/backend/internal/service/ollama_cloud_usage.go b/backend/internal/service/ollama_cloud_usage.go index ad0f647ae6..2956cc9e3b 100644 --- a/backend/internal/service/ollama_cloud_usage.go +++ b/backend/internal/service/ollama_cloud_usage.go @@ -1007,8 +1007,22 @@ func OllamaCloudUsageStateFromAccount(account *Account) *OllamaCloudUsageState { return state } +// isOllamaCloudUsagePlatform 收敛 Ollama Cloud 用量窗口的平台白名单。官方 +// ollama.com base_url 除官方两平台外，也允许挂在经 OpenAI 网关转发的国产 +// OpenAI 兼容平台下（用户把 Ollama Cloud key 挂在 kimi/zhipu/deepseek 分组 +// 里跑托管的 glm/kimi/deepseek 模型）。repository 侧 SQL 白名单 +// （ollamaCloudUsagePlatformsSQL）是本列表的镜像，两侧必须同步修改。 +func isOllamaCloudUsagePlatform(platform string) bool { + switch platform { + case PlatformOpenAI, PlatformAnthropic, PlatformKimi, PlatformZhipu, PlatformDeepseek: + return true + default: + return false + } +} + func IsOllamaCloudUsageAccount(account *Account) bool { - if account == nil || account.Type != AccountTypeAPIKey || (account.Platform != PlatformOpenAI && account.Platform != PlatformAnthropic) { + if account == nil || account.Type != AccountTypeAPIKey || !isOllamaCloudUsagePlatform(account.Platform) { return false } baseURL, _ := account.Credentials["base_url"].(string) diff --git a/backend/internal/service/ollama_cloud_usage_test.go b/backend/internal/service/ollama_cloud_usage_test.go index 4b4a5aadb6..5ad87d69d5 100644 --- a/backend/internal/service/ollama_cloud_usage_test.go +++ b/backend/internal/service/ollama_cloud_usage_test.go @@ -467,7 +467,14 @@ func TestIsOllamaCloudUsageAccountStrictOfficialHost(t *testing.T) { {"https://ollama.com", PlatformOpenAI, true}, {"HTTPS://OLLAMA.COM", PlatformAnthropic, true}, {"https://www.OLLAMA.com:443/v1", PlatformOpenAI, true}, - {"https://ollama.com:443", PlatformOpenAI, true}, + // 官方 ollama.com key 挂在国产 OpenAI 兼容平台下同样进用量窗口。 + {"https://ollama.com", PlatformKimi, true}, + {"https://www.ollama.com/v1", PlatformZhipu, true}, + {"https://ollama.com:443", PlatformDeepseek, true}, + // 用量窗口不随 base_url 放开到其余平台。 + {"https://ollama.com", PlatformGemini, false}, + {"https://ollama.com", PlatformGrok, false}, + {"https://ollama.com", PlatformAntigravity, false}, {"https://ollama.com/", PlatformAnthropic, false}, {"https://ollama.com/v1/", PlatformOpenAI, false}, {"http://ollama.com", PlatformOpenAI, false}, @@ -488,6 +495,16 @@ func TestIsOllamaCloudUsageAccountStrictOfficialHost(t *testing.T) { } } +// oauth 类型账号即使平台与 base_url 都命中也不进用量窗口（仅 apikey 账号）。 +func TestIsOllamaCloudUsageAccountRejectsOAuthType(t *testing.T) { + for _, platform := range []string{PlatformOpenAI, PlatformAnthropic, PlatformKimi, PlatformZhipu, PlatformDeepseek} { + account := ollamaUsageAccount(1) + account.Platform = platform + account.Type = AccountTypeOAuth + require.False(t, IsOllamaCloudUsageAccount(account), platform) + } +} + func TestNormalizeOllamaCloudUsageCookieAllowlist(t *testing.T) { normalized, err := normalizeOllamaCloudUsageCookie(" tracking=discard ; wos-session=secret ; __Secure-authjs.session-token.0=part-a ; device=discard ") require.NoError(t, err) diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 1ff84f3275..12c309cf4c 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -431,7 +431,16 @@