feat: improve model rate limit tracking

apps/api/internal/runner/service.go

@@ -64,7 +64,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	}
 	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user)
 	if err != nil {
-		failed, finishErr := s.failTask(ctx, task.ID, "no_model_candidate", err.Error(), task.RunMode == "simulation", err)
+		failed, finishErr := s.failTask(ctx, task.ID, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
 		if finishErr != nil {
 			return Result{}, finishErr
 		}
@@ -279,7 +279,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 		return clients.Response{}, err
 	}
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
-	limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, reservations)
+	limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, attemptID, reservations)
 	if err != nil {
 		clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: false}
 		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
@@ -292,6 +292,12 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 		})
 		return clients.Response{}, clientErr
 	}
+	rateReservationsFinalized := false
+	defer func() {
+		if !rateReservationsFinalized {
+			_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
+		}
+	}()
 	defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
 
 	if err := s.store.RecordClientAssignment(ctx, candidate); err != nil {
@@ -397,6 +403,10 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 			return clients.Response{}, err
 		}
 	}
+	if err := s.store.CommitRateLimitReservations(ctx, limitResult.Reservations, tokenUsageAmounts(response.Usage)); err != nil {
+		return clients.Response{}, err
+	}
+	rateReservationsFinalized = true
 	if err := s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
 		AttemptID:          attemptID,
 		Status:             "succeeded",