mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 15:38:22 +08:00
fix(ollama): 避免刷新候选饥饿
ListDue 在 LIMIT 前用与 service 纯函数一致的 debounce/max-wait/backoff 规则筛真正 due 组,防止有活动但未到期的组占满每轮 20 名额。
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user