feat: improve model rate limit tracking

apps/api/internal/runner/limits.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
 )
 
@@ -108,3 +109,21 @@ func estimateRequestTokens(body map[string]any) int {
 	}
 	return len([]rune(text))/4 + 1
 }
+
+func tokenUsageAmounts(usage clients.Usage) map[string]float64 {
+	out := map[string]float64{}
+	if usage.InputTokens > 0 {
+		out["tpm_input"] = float64(usage.InputTokens)
+	}
+	if usage.OutputTokens > 0 {
+		out["tpm_output"] = float64(usage.OutputTokens)
+	}
+	total := usage.TotalTokens
+	if total <= 0 {
+		total = usage.InputTokens + usage.OutputTokens
+	}
+	if total > 0 {
+		out["tpm_total"] = float64(total)
+	}
+	return out
+}

apps/api/internal/runner/limits_test.go

@@ -0,0 +1,29 @@
+package runner
+
+import (
+	"testing"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
+)
+
+func TestTokenUsageAmountsUsesActualUsageForTPM(t *testing.T) {
+	got := tokenUsageAmounts(clients.Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 21})
+
+	if got["tpm_input"] != 12 {
+		t.Fatalf("expected input token amount 12, got %v", got["tpm_input"])
+	}
+	if got["tpm_output"] != 8 {
+		t.Fatalf("expected output token amount 8, got %v", got["tpm_output"])
+	}
+	if got["tpm_total"] != 21 {
+		t.Fatalf("expected total token amount 21, got %v", got["tpm_total"])
+	}
+}
+
+func TestTokenUsageAmountsFallsBackToInputOutputTotal(t *testing.T) {
+	got := tokenUsageAmounts(clients.Usage{InputTokens: 3, OutputTokens: 5})
+
+	if got["tpm_total"] != 8 {
+		t.Fatalf("expected total token fallback 8, got %v", got["tpm_total"])
+	}
+}

apps/api/internal/runner/runtime_policy.go

@@ -29,8 +29,8 @@ func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID stri
 	degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
 	if failurePolicyMatches(degradePolicy, code, message) {
 		cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
-		if err := s.store.CooldownCandidatePlatform(ctx, candidate.PlatformID, cooldownSeconds); err == nil {
-			_ = s.emit(ctx, taskID, "task.policy.degraded", "running", "degrade", 0.5, "candidate platform cooled down by failure policy", map[string]any{
+		if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, cooldownSeconds); err == nil {
+			_ = s.emit(ctx, taskID, "task.policy.degraded", "running", "degrade", 0.5, "candidate model cooled down by failure policy", map[string]any{
 				"platformId":      candidate.PlatformID,
 				"platformModelId": candidate.PlatformModelID,
 				"cooldownSeconds": cooldownSeconds,
@@ -52,8 +52,8 @@ func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candid
 			}, decision.Match, decision.Info), simulated)
 		}
 	case "cooldown_and_next":
-		if err := s.store.CooldownCandidatePlatform(ctx, candidate.PlatformID, decision.CooldownSeconds); err == nil {
-			_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate platform cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
+		if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds); err == nil {
+			_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate model cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
 				"platformId":      candidate.PlatformID,
 				"platformModelId": candidate.PlatformModelID,
 				"cooldownSeconds": decision.CooldownSeconds,

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",