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)