From d129bcccbd8ce45f699ba661cd7af15ed4fda56b Mon Sep 17 00:00:00 2001 From: wangbo Date: Mon, 3 Aug 2026 21:06:23 +0800 Subject: [PATCH] =?UTF-8?q?fix(runner):=20=E4=BF=9D=E7=95=99=E5=90=8C?= =?UTF-8?q?=E5=B9=B3=E5=8F=B0=E9=87=8D=E8=AF=95=E7=9A=84=E5=B9=B6=E5=8F=91?= =?UTF-8?q?=E7=A7=9F=E7=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 同平台 retry_same 之前会在首个 attempt 失败时释放任务级 admission 租约，下一次 attempt 继续续租旧租约并进入 upstream_submission_unknown。 现在明确区分 attempt 自有租约与任务级 admission 租约，只在 attempt 结束时释放前者；补充回归测试覆盖同平台重试继续续租 admission 租约。 验证：env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1；go vet ./...；gofmt；git diff --check。 --- .../internal/runner/execution_lease_test.go | 26 +++++++++++++++++++ apps/api/internal/runner/service.go | 22 ++++++++++------ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/apps/api/internal/runner/execution_lease_test.go b/apps/api/internal/runner/execution_lease_test.go index 6a304d9..a6664aa 100644 --- a/apps/api/internal/runner/execution_lease_test.go +++ b/apps/api/internal/runner/execution_lease_test.go @@ -31,3 +31,29 @@ func TestLeaseRenewalFatalityDistinguishesTransientDatabaseErrors(t *testing.T) t.Fatal("concurrency lease renewal error at expiry must be fatal") } } + +func TestBindAttemptConcurrencyLeasesPreservesAdmissionLeaseAcrossRetry(t *testing.T) { + attemptLease := store.ConcurrencyLease{ID: "attempt-lease", TTL: time.Minute} + admissionLeases := []store.ConcurrencyLease{ + {ID: "admission-user-group", TTL: 2 * time.Minute}, + {ID: "admission-worker-capacity", TTL: 2 * time.Minute}, + } + + active, release := bindAttemptConcurrencyLeases([]store.ConcurrencyLease{attemptLease}, admissionLeases) + if len(active) != 3 || active[0].ID != attemptLease.ID || active[1].ID != admissionLeases[0].ID || active[2].ID != admissionLeases[1].ID { + t.Fatalf("active leases must renew attempt and admission ownership together: %+v", active) + } + if len(release) != 1 || release[0].ID != attemptLease.ID { + t.Fatalf("attempt cleanup must not release task-scoped admission leases: %+v", release) + } + + // A same-platform retry receives the same admission lease IDs. The first + // attempt cleanup must leave them renewable by the next attempt. + retryActive, retryRelease := bindAttemptConcurrencyLeases(nil, admissionLeases) + if len(retryActive) != 2 || retryActive[0].ID != admissionLeases[0].ID || retryActive[1].ID != admissionLeases[1].ID { + t.Fatalf("retry must keep using the live admission leases: %+v", retryActive) + } + if len(retryRelease) != 0 { + t.Fatalf("retry without attempt-owned leases must not release admission leases: %+v", retryRelease) + } +} diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index da9685a..5fcb50a 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -1357,20 +1357,18 @@ func (s *Service) runCandidate( clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable} return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)} } + attemptOwnedLeases := append([]store.ConcurrencyLease(nil), limitResult.Leases...) if admittedPlatformModelID == candidate.PlatformModelID { - limitResult.Leases = append(limitResult.Leases, admittedLeases...) + limitResult.Leases, attemptOwnedLeases = bindAttemptConcurrencyLeases(limitResult.Leases, admittedLeases) } rateReservationsFinalized := false - retainAdmittedLeases := false defer func() { if !rateReservationsFinalized { _ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed") } }() defer func() { - if !retainAdmittedLeases { - _ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases) - } + _ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), attemptOwnedLeases) }() attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{ @@ -1798,12 +1796,20 @@ func (s *Service) runCandidate( }); err != nil { s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID) } - if admittedPlatformModelID == candidate.PlatformModelID { - retainAdmittedLeases = true - } return response, nil } +// bindAttemptConcurrencyLeases combines task-scoped admission leases with +// leases acquired by the current attempt. Only the latter are returned for +// attempt cleanup: admission owns its leases across retry_same and releases +// them when the task changes binding or reaches a terminal state. +func bindAttemptConcurrencyLeases(attemptLeases []store.ConcurrencyLease, admissionLeases []store.ConcurrencyLease) ([]store.ConcurrencyLease, []store.ConcurrencyLease) { + attemptOwned := append([]store.ConcurrencyLease(nil), attemptLeases...) + active := append([]store.ConcurrencyLease(nil), attemptLeases...) + active = append(active, admissionLeases...) + return active, attemptOwned +} + func (s *Service) observeProviderQuotaWait(metric string) { observer, ok := s.billingMetrics.(interface { ObserveProviderQuotaWait(string)