From b403f88f513da8f81f5a3cb451f22b4493bbd917 Mon Sep 17 00:00:00 2001 From: alfadb Date: Sat, 25 Jul 2026 11:28:49 +0800 Subject: [PATCH] =?UTF-8?q?fix(ollama):=20=E9=81=BF=E5=85=8D=E5=88=B7?= =?UTF-8?q?=E6=96=B0=E5=80=99=E9=80=89=E9=A5=A5=E9=A5=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ListDue 在 LIMIT 前用与 service 纯函数一致的 debounce/max-wait/backoff 规则筛真正 due 组，防止有活动但未到期的组占满每轮 20 名额。 --- .../admin/account_ollama_cloud_usage_test.go | 2 +- .../account_repo_ollama_cloud_usage.go | 86 ++++++++++---- ...epo_ollama_cloud_usage_integration_test.go | 112 +++++++++++++++++- .../account_repo_ollama_cloud_usage_test.go | 13 +- ...teway_anthropic_apikey_passthrough_test.go | 111 +++++++++++++++++ .../internal/service/ollama_cloud_usage.go | 4 +- .../service/ollama_cloud_usage_test.go | 2 +- ...ai_upstream_transport_error_handle_test.go | 65 ++++++++++ 8 files changed, 362 insertions(+), 33 deletions(-) diff --git a/backend/internal/handler/admin/account_ollama_cloud_usage_test.go b/backend/internal/handler/admin/account_ollama_cloud_usage_test.go index 38db30a2c1..7dd7b712f8 100644 --- a/backend/internal/handler/admin/account_ollama_cloud_usage_test.go +++ b/backend/internal/handler/admin/account_ollama_cloud_usage_test.go @@ -61,7 +61,7 @@ func (r *ollamaCloudUsageHandlerTestRepo) UpdateOllamaCloudUsageSnapshot(context func (r *ollamaCloudUsageHandlerTestRepo) DisableOllamaCloudUsageAutoRefresh(context.Context, *service.Account) error { return nil } -func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]service.Account, error) { +func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]service.Account, error) { return nil, nil } diff --git a/backend/internal/repository/account_repo_ollama_cloud_usage.go b/backend/internal/repository/account_repo_ollama_cloud_usage.go index 9d06c44ef1..45ed06fe41 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage.go @@ -372,22 +372,37 @@ func ollamaCloudUsageParseRFC3339SQL(expression string) string { END` } -// ListDueOllamaCloudUsageAccounts returns at most one activity-driven candidate -// per exact API key. It stamps Account.LastUsedAt with the group MAX(last_used_at) -// so the service due pure function can decide without hydrating the whole table. +// ListDueOllamaCloudUsageAccounts returns at most one truly-due activity-driven +// candidate per exact API key. Due timing (debounce, max-wait, failure backoff) +// is evaluated in SQL before LIMIT so non-due active groups cannot starve due ones. +// Account.LastUsedAt is stamped with the group MAX(last_used_at) for a service +// pure-function recheck against races between list and refresh. // -// SQL only prefilters potential candidates (missing/invalid snapshot, or group -// activity after the last successful fetch / failed attempt). Final due timing -// (debounce, max-wait, failure backoff) is evaluated in service code. -func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) { +// Rules mirror service.ollamaCloudUsageAutoRefreshDueAt: +// - missing/invalid snapshot or times → fail-open first due +// - success: activity after fetched_at; due_at = LEAST(last_used+debounce, fetched+maxWait) +// - failed/unauthorized: activity after last_attempt; activity_due = LEAST(...); +// final due_at is not earlier than a valid next_refresh_at (invalid/missing fail-open) +func (r *accountRepository) ListDueOllamaCloudUsageAccounts( + ctx context.Context, + now time.Time, + debounce, maxWait time.Duration, + limit int, +) ([]service.Account, error) { if limit <= 0 { return []service.Account{}, nil } if r == nil || r.sql == nil { return nil, errors.New("account repository SQL executor not configured") } - // now is retained in the signature for callers; activity due is evaluated in service. - _ = now + if debounce <= 0 { + debounce = time.Minute + } + if maxWait <= 0 { + maxWait = time.Hour + } + debounceSeconds := debounce.Seconds() + maxWaitSeconds := maxWait.Seconds() rows, err := r.sql.QueryContext(ctx, ` WITH eligible AS ( SELECT id, @@ -412,14 +427,41 @@ func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, SELECT e.id, e.api_key, e.snapshot, g.group_last_used_at, e.snapshot #>> '{status}' AS status, e.snapshot #>> '{fetched_at}' AS fetched_at, - e.snapshot #>> '{last_attempt_at}' AS last_attempt_at + e.snapshot #>> '{last_attempt_at}' AS last_attempt_at, + e.snapshot #>> '{next_refresh_at}' AS next_refresh_at FROM eligible e JOIN group_activity g ON g.api_key = e.api_key ), parsed AS MATERIALIZED ( SELECT id, api_key, snapshot, group_last_used_at, status, `+ollamaCloudUsageParseRFC3339SQL("fetched_at")+` AS parsed_fetched_at, - `+ollamaCloudUsageParseRFC3339SQL("last_attempt_at")+` AS parsed_last_attempt_at + `+ollamaCloudUsageParseRFC3339SQL("last_attempt_at")+` AS parsed_last_attempt_at, + `+ollamaCloudUsageParseRFC3339SQL("next_refresh_at")+` AS parsed_next_refresh_at FROM joined + ), timed AS ( + SELECT *, + CASE + WHEN status = 'ok' + AND parsed_fetched_at IS NOT NULL + AND group_last_used_at IS NOT NULL + AND group_last_used_at > parsed_fetched_at::timestamptz + THEN LEAST( + group_last_used_at + make_interval(secs => $2::double precision), + parsed_fetched_at::timestamptz + make_interval(secs => $3::double precision) + ) + WHEN status IN ('failed', 'unauthorized') + AND parsed_last_attempt_at IS NOT NULL + AND group_last_used_at IS NOT NULL + AND group_last_used_at > parsed_last_attempt_at::timestamptz + THEN GREATEST( + LEAST( + group_last_used_at + make_interval(secs => $2::double precision), + parsed_last_attempt_at::timestamptz + make_interval(secs => $3::double precision) + ), + COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz) + ) + ELSE NULL + END AS activity_due_at + FROM parsed ), candidates AS ( SELECT *, CASE @@ -427,21 +469,17 @@ func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, OR status NOT IN ('ok', 'failed', 'unauthorized') THEN 0 WHEN status = 'ok' AND parsed_fetched_at IS NULL THEN 0 WHEN status IN ('failed', 'unauthorized') AND parsed_last_attempt_at IS NULL THEN 0 - WHEN status = 'ok' - AND group_last_used_at IS NOT NULL - AND group_last_used_at > parsed_fetched_at::timestamptz THEN 1 - WHEN status IN ('failed', 'unauthorized') - AND group_last_used_at IS NOT NULL - AND group_last_used_at > parsed_last_attempt_at::timestamptz THEN 1 + WHEN activity_due_at IS NOT NULL AND $1 >= activity_due_at THEN 1 ELSE NULL - END AS due_class - FROM parsed + END AS due_class, + activity_due_at AS due_at + FROM timed ), ranked AS ( - SELECT id, api_key, group_last_used_at, due_class, + SELECT id, api_key, group_last_used_at, due_class, due_at, row_number() OVER ( PARTITION BY api_key ORDER BY due_class, - group_last_used_at NULLS FIRST, + due_at NULLS FIRST, id ) AS group_rank FROM candidates @@ -450,9 +488,9 @@ func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, SELECT id, group_last_used_at FROM ranked WHERE group_rank = 1 - ORDER BY due_class, group_last_used_at NULLS FIRST, id - LIMIT $1 - `, limit) + ORDER BY due_class, due_at NULLS FIRST, id + LIMIT $4 + `, now.UTC(), debounceSeconds, maxWaitSeconds, limit) if err != nil { return nil, err } 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 49bb79127f..47de81095d 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 @@ -58,7 +58,7 @@ func TestListDueOllamaCloudUsageAccountsOrderingLimitAndProxyHydration(t *testin }, nil) _ = createAccount("ollama-ineligible", "https://ollama.com.evil.test", nil, nil, nil) - accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 2) + accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 2) require.NoError(t, err) require.Len(t, accounts, 2) @@ -127,7 +127,7 @@ func TestListDueOllamaCloudUsageAccountsUsesGroupMaxLastUsedAndFailsOpen(t *test }, }) - accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 10) + accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10) require.NoError(t, err, "invalid stored values must not abort the query") ids := accountIDs(accounts) @@ -456,3 +456,111 @@ func TestUpdateCredentialsUnchangedCredentialsPreserveManagedExtra(t *testing.T) require.NotContains(t, probeLoaded.Extra, service.UpstreamBillingProbeExtraKey, "changed credentials must keep clearing the probe snapshot") } + +// TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService proves the SQL +// candidate layer applies debounce / max-wait / failure-backoff before LIMIT, +// matching service.ollamaCloudUsageIsAutoRefreshDue, and that >20 active-but- +// not-yet-due groups cannot starve a truly due max-wait group. +func TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService(t *testing.T) { + ctx := context.Background() + tx := testEntTx(t) + repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil) + now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC) + debounce := time.Minute + maxWait := time.Hour + + createOK := func(name string, fetched, lastUsed time.Time) *service.Account { + t.Helper() + return mustCreateAccount(t, tx.Client(), &service.Account{ + Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"}, + Extra: map[string]any{ + service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture", + service.OllamaCloudUsageAutoRefreshExtraKey: true, + service.OllamaCloudUsageSnapshotExtraKey: map[string]any{ + "status": service.OllamaCloudUsageStatusOK, + "fetched_at": fetched.UTC().Format(time.RFC3339Nano), + "last_attempt_at": fetched.UTC().Format(time.RFC3339Nano), + "next_refresh_at": fetched.Add(maxWait).UTC().Format(time.RFC3339Nano), + }, + }, + LastUsedAt: &lastUsed, + }) + } + createFailed := func(name string, lastAttempt, lastUsed, nextRefresh time.Time, nextRefreshRaw string) *service.Account { + t.Helper() + snapshot := map[string]any{ + "status": service.OllamaCloudUsageStatusFailed, + "last_attempt_at": lastAttempt.UTC().Format(time.RFC3339Nano), + "failure_count": 1, + } + if nextRefreshRaw != "" { + snapshot["next_refresh_at"] = nextRefreshRaw + } else { + snapshot["next_refresh_at"] = nextRefresh.UTC().Format(time.RFC3339Nano) + } + return mustCreateAccount(t, tx.Client(), &service.Account{ + Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, + Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"}, + Extra: map[string]any{ + service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture", + service.OllamaCloudUsageAutoRefreshExtraKey: true, + service.OllamaCloudUsageSnapshotExtraKey: snapshot, + }, + LastUsedAt: &lastUsed, + }) + } + + // 21 groups with activity after fetch but debounce not elapsed — previously + // these alone could fill LIMIT 20 every minute and starve true due groups. + notDueIDs := make(map[int64]struct{}, 21) + for i := 0; i < 21; i++ { + // fetched 10m ago, last used 10s ago → due_at = lastUsed+debounce = now+50s (not due) + acc := createOK(fmt.Sprintf("ollama-not-due-debounce-%02d", i), now.Add(-10*time.Minute), now.Add(-10*time.Second)) + notDueIDs[acc.ID] = struct{}{} + } + + // Truly due via max-wait: fetched 2h ago, continuous activity 10s ago. + // due_at = min(now-10s+1m, now-2h+1h) = now-1h → due. + maxWaitDue := createOK("ollama-due-maxwait", now.Add(-2*time.Hour), now.Add(-10*time.Second)) + + // Success debounce elapsed: last used 2m ago with debounce 1m → due. + debounceDue := createOK("ollama-due-debounce", now.Add(-30*time.Minute), now.Add(-2*time.Minute)) + + // Success still within debounce → not due. + _ = createOK("ollama-not-due-fresh", now.Add(-30*time.Minute), now.Add(-20*time.Second)) + + // Failure blocked by next_refresh_at backoff even with new activity. + _ = createFailed("ollama-fail-backoff", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(10*time.Minute), "") + + // Failure after backoff with new request → due. + failDue := createFailed("ollama-fail-due", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(-time.Minute), "") + + // Invalid next_refresh_at must fail open (not abort query / not block activity due). + failInvalidNext := createFailed("ollama-fail-invalid-next", now.Add(-30*time.Minute), now.Add(-2*time.Minute), time.Time{}, "not-a-timestamp") + + accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, 20) + require.NoError(t, err) + + ids := accountIDs(accounts) + require.Contains(t, ids, maxWaitDue.ID, "max-wait due group must not be starved by not-yet-due activity groups") + require.Contains(t, ids, debounceDue.ID, "success debounce elapsed must be due in SQL") + require.Contains(t, ids, failDue.ID, "failure after backoff with new activity must be due in SQL") + require.Contains(t, ids, failInvalidNext.ID, "invalid next_refresh_at must fail open to activity due") + require.LessOrEqual(t, len(accounts), 20) + + // Fixtures below match service.ollamaCloudUsageIsAutoRefreshDue semantics; + // none of the not-yet-due groups may appear even when they outnumber the limit. + for _, id := range ids { + _, isNotDue := notDueIDs[id] + require.False(t, isNotDue, "not-yet-due debounce group %d must not be returned by SQL LIMIT layer", id) + } + require.NotContains(t, ids, int64(0)) + + // Explicit not-due names must stay out: fresh success and failure still in backoff. + for _, account := range accounts { + require.NotContains(t, account.Name, "not-due") + require.NotEqual(t, "ollama-fail-backoff", account.Name) + require.NotEqual(t, "ollama-not-due-fresh", account.Name) + } +} 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 6adb819144..64d715e2d0 100644 --- a/backend/internal/repository/account_repo_ollama_cloud_usage_test.go +++ b/backend/internal/repository/account_repo_ollama_cloud_usage_test.go @@ -190,13 +190,15 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { _ = db.Close() }) now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC) + debounce := time.Minute + maxWait := time.Hour var capturedSQL string mock.ExpectQuery("WITH eligible AS"). - WithArgs(20). + WithArgs(now.UTC(), debounce.Seconds(), maxWait.Seconds(), 20). WillReturnRows(sqlmock.NewRows([]string{"id", "group_last_used_at"})) repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil) - accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, 20) + accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, debounce, maxWait, 20) require.NoError(t, err) require.Empty(t, accounts) @@ -212,9 +214,14 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) { "MAX(last_used_at) AS group_last_used_at", "PARTITION BY api_key", "WHERE group_rank = 1", - "LIMIT $1", + "LIMIT $4", + "make_interval(secs => $2::double precision)", + "make_interval(secs => $3::double precision)", "group_last_used_at > parsed_fetched_at::timestamptz", "group_last_used_at > parsed_last_attempt_at::timestamptz", + "$1 >= activity_due_at", + "COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)", + "ORDER BY due_class, due_at NULLS FIRST, id", } { require.Contains(t, normalized, clause) } diff --git a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go index b687696cfe..3722ec5189 100644 --- a/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go +++ b/backend/internal/service/gateway_anthropic_apikey_passthrough_test.go @@ -1623,3 +1623,114 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_StreamingUpstreamReadErrorAft require.True(t, result.clientDisconnect) require.Equal(t, 8, result.usage.InputTokens) } + +func TestGatewayService_AnthropicAPIKeyPassthrough_TransportErrorRecordsOllamaActivity(t *testing.T) { + gin.SetMode(gin.TestMode) + deferred := NewDeferredService(nil, nil, time.Second) + upstream := &anthropicHTTPUpstreamRecorder{err: errors.New("dial tcp timeout")} + svc := &GatewayService{ + cfg: &config.Config{ + Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }, + }, + httpUpstream: upstream, + deferredService: deferred, + } + + ollama := &Account{ + ID: 601, Name: "ollama-anthropic", Platform: PlatformAnthropic, Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + Extra: map[string]any{"anthropic_passthrough": true}, + Status: StatusActive, Schedulable: true, + } + other := newAnthropicAPIKeyAccountForTest() + other.ID = 602 + + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + _, err := svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now()) + require.Error(t, err) + + rec2 := httptest.NewRecorder() + c2, _ := gin.CreateTestContext(rec2) + c2.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + _, err = svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c2, other, []byte(`{"model":"x"}`), "x", "x", false, time.Now()) + require.Error(t, err) + + _, ok := deferred.lastUsedUpdates.Load(int64(601)) + require.True(t, ok, "Anthropic passthrough transport error on Ollama account must record activity") + _, ok = deferred.lastUsedUpdates.Load(int64(602)) + require.False(t, ok, "non-Ollama Anthropic passthrough transport error must not record Ollama activity") +} + +func TestGatewayService_AnthropicAPIKeyPassthrough_ContextCanceledSkipsOllamaActivity(t *testing.T) { + gin.SetMode(gin.TestMode) + deferred := NewDeferredService(nil, nil, time.Second) + upstream := &anthropicHTTPUpstreamRecorder{err: context.Canceled} + svc := &GatewayService{ + cfg: &config.Config{ + Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }, + }, + httpUpstream: upstream, + deferredService: deferred, + } + ollama := &Account{ + ID: 603, Name: "ollama-canceled", Platform: PlatformAnthropic, Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + Extra: map[string]any{"anthropic_passthrough": true}, + Status: StatusActive, Schedulable: true, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + _, err := svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now()) + + require.Error(t, err) + _, ok := deferred.lastUsedUpdates.Load(int64(603)) + require.False(t, ok, "context.Canceled on Anthropic passthrough must not count as Ollama activity") +} + +func TestGatewayService_AnthropicAPIKeyPassthrough_Non2xxRecordsOllamaActivity(t *testing.T) { + gin.SetMode(gin.TestMode) + deferred := NewDeferredService(nil, nil, time.Second) + // 400 is non-retryable / non-failover for default API-key accounts, so it reaches handleErrorResponse. + upstream := &anthropicHTTPUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"invalid_request_error","message":"bad"}}`)), + }, + } + svc := &GatewayService{ + cfg: &config.Config{ + Security: config.SecurityConfig{ + URLAllowlist: config.URLAllowlistConfig{Enabled: false}, + }, + }, + httpUpstream: upstream, + deferredService: deferred, + rateLimitService: &RateLimitService{}, + } + ollama := &Account{ + ID: 604, Name: "ollama-400", Platform: PlatformAnthropic, Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + Extra: map[string]any{"anthropic_passthrough": true}, + Status: StatusActive, Schedulable: true, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + _, _ = svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now()) + + _, ok := deferred.lastUsedUpdates.Load(int64(604)) + require.True(t, ok, "Anthropic passthrough non-2xx on Ollama account must record activity via handleErrorResponse") +} diff --git a/backend/internal/service/ollama_cloud_usage.go b/backend/internal/service/ollama_cloud_usage.go index d489ec1721..d0e7077b7b 100644 --- a/backend/internal/service/ollama_cloud_usage.go +++ b/backend/internal/service/ollama_cloud_usage.go @@ -156,7 +156,7 @@ type ollamaCloudUsageRepository interface { SetOllamaCloudUsageAutoRefresh(context.Context, *Account, bool) error UpdateOllamaCloudUsageSnapshot(context.Context, *Account, *OllamaCloudUsageSnapshot) error DisableOllamaCloudUsageAutoRefresh(context.Context, *Account) error - ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]Account, error) + ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]Account, error) } // GetOllamaCloudUsageSettings returns fail-safe defaults when the setting is absent. @@ -703,7 +703,7 @@ func (s *OllamaCloudUsageService) RunDue(ctx context.Context) error { } now := s.currentTime() debounce, maxWait := ollamaCloudUsageDurations(settings) - accounts, err := writer.ListDueOllamaCloudUsageAccounts(ctx, now, ollamaCloudUsageMaxPerCycle) + accounts, err := writer.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, ollamaCloudUsageMaxPerCycle) if err != nil { return fmt.Errorf("list due Ollama Cloud usage accounts: %w", err) } diff --git a/backend/internal/service/ollama_cloud_usage_test.go b/backend/internal/service/ollama_cloud_usage_test.go index 9269e4f75f..cfe1864ad0 100644 --- a/backend/internal/service/ollama_cloud_usage_test.go +++ b/backend/internal/service/ollama_cloud_usage_test.go @@ -185,7 +185,7 @@ func applyOllamaUsageTestManagedExtra(account, source *Account) { } } -func (r *ollamaUsageTestRepo) ListDueOllamaCloudUsageAccounts(_ context.Context, _ time.Time, limit int) ([]Account, error) { +func (r *ollamaUsageTestRepo) ListDueOllamaCloudUsageAccounts(_ context.Context, _ time.Time, _, _ time.Duration, limit int) ([]Account, error) { r.mu.Lock() defer r.mu.Unlock() if len(r.due) > 0 { diff --git a/backend/internal/service/openai_upstream_transport_error_handle_test.go b/backend/internal/service/openai_upstream_transport_error_handle_test.go index 3ea8b43181..38e392893e 100644 --- a/backend/internal/service/openai_upstream_transport_error_handle_test.go +++ b/backend/internal/service/openai_upstream_transport_error_handle_test.go @@ -210,3 +210,68 @@ func TestForwardAsRawChatCompletions_TransportErrorFailsOver(t *testing.T) { require.Empty(t, repo.tempUnschedCalls, "plain EOF is transient: fail over but do not evict") require.Equal(t, 0, rec.Body.Len(), "service must not write a hard 502 before handler can fail over") } + +func TestHandleOpenAIUpstreamTransportError_RecordsOllamaActivityOnly(t *testing.T) { + deferred := NewDeferredService(nil, nil, time.Second) + svc := &OpenAIGatewayService{ + accountRepo: &openaiTransportAccountRepoStub{}, + deferredService: deferred, + } + ollama := &Account{ + ID: 501, Name: "ollama-cloud", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + } + other := &Account{ + ID: 502, Name: "openai-official", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k-openai", "base_url": "https://api.openai.com"}, + } + c, _ := newOpenAITransportErrTestContext() + + _ = svc.handleOpenAIUpstreamTransportError(context.Background(), c, ollama, errors.New("connection reset"), false) + _ = svc.handleOpenAIUpstreamTransportError(context.Background(), c, other, errors.New("connection reset"), false) + + _, ok := deferred.lastUsedUpdates.Load(int64(501)) + require.True(t, ok, "Ollama Cloud transport error must schedule last_used activity") + _, ok = deferred.lastUsedUpdates.Load(int64(502)) + require.False(t, ok, "non-Ollama transport error must not schedule Ollama activity") +} + +func TestHandleOpenAIUpstreamTransportError_ContextCanceledSkipsOllamaActivity(t *testing.T) { + deferred := NewDeferredService(nil, nil, time.Second) + svc := &OpenAIGatewayService{ + accountRepo: &openaiTransportAccountRepoStub{}, + deferredService: deferred, + } + ollama := &Account{ + ID: 503, Name: "ollama-canceled", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + } + c, _ := newOpenAITransportErrTestContext() + + err := svc.handleOpenAIUpstreamTransportError(context.Background(), c, ollama, context.Canceled, false) + + require.ErrorIs(t, err, context.Canceled) + _, ok := deferred.lastUsedUpdates.Load(int64(503)) + require.False(t, ok, "context.Canceled is client disconnect before a fault; do not count as Ollama activity") +} + +func TestHandleOpenAIAccountUpstreamError_RecordsOllamaActivityOnly(t *testing.T) { + deferred := NewDeferredService(nil, nil, time.Second) + svc := &OpenAIGatewayService{deferredService: deferred} + ollama := &Account{ + ID: 504, Name: "ollama-429", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"}, + } + other := &Account{ + ID: 505, Name: "openai-429", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, + Credentials: map[string]any{"api_key": "k-openai", "base_url": "https://api.openai.com"}, + } + + _ = svc.handleOpenAIAccountUpstreamError(context.Background(), ollama, http.StatusTooManyRequests, http.Header{}, []byte(`{"error":{"message":"rate"}}`), "gpt-test") + _ = svc.handleOpenAIAccountUpstreamError(context.Background(), other, http.StatusTooManyRequests, http.Header{}, []byte(`{"error":{"message":"rate"}}`), "gpt-test") + + _, ok := deferred.lastUsedUpdates.Load(int64(504)) + require.True(t, ok, "Ollama Cloud non-2xx must schedule last_used activity") + _, ok = deferred.lastUsedUpdates.Load(int64(505)) + require.False(t, ok, "non-Ollama non-2xx must not schedule Ollama activity") +}