1 files changed, 38 insertions, 0 deletions
diff --git a/internal/limits/limits_test.go b/internal/limits/limits_test.go
index 78b346e..8e48d42 100644
--- a/internal/limits/limits_test.go
+++ b/internal/limits/limits_test.go
@@ -56,6 +56,44 @@ func TestLocalTokenAndConcurrencyLimits(t *testing.T) {
 	retry.Release()
 }
 
+func TestLocalKeyLimitDoesNotConsumeProjectQuotaWhenRejected(t *testing.T) {
+	limiter := New("", "test", 0, nil)
+	limiter.ReplacePolicies([]domain.LimitPolicy{{ProjectID: "project-1", RequestsPerMinute: 2}})
+	firstKey := domain.Principal{ProjectID: "project-1", KeyID: "key-1", RequestsPerMinute: 1}
+	lease, err := limiter.Acquire(context.Background(), firstKey, []byte(`{}`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	lease.Release()
+	if _, err := limiter.Acquire(context.Background(), firstKey, []byte(`{}`)); !errors.Is(err, ErrRequestsExceeded) {
+		t.Fatalf("second key request error = %v, want request limit", err)
+	}
+	secondKey := domain.Principal{ProjectID: "project-1", KeyID: "key-2", RequestsPerMinute: 1}
+	lease, err = limiter.Acquire(context.Background(), secondKey, []byte(`{}`))
+	if err != nil {
+		t.Fatalf("key rejection consumed project quota: %v", err)
+	}
+	lease.Release()
+	if _, err := limiter.Acquire(context.Background(), domain.Principal{ProjectID: "project-1", KeyID: "key-3"}, []byte(`{}`)); !errors.Is(err, ErrRequestsExceeded) {
+		t.Fatalf("project request limit error = %v", err)
+	}
+}
+
+func TestLocalKeyTokenLimit(t *testing.T) {
+	limiter := New("", "test", 0, nil)
+	body := []byte(`{"max_tokens":4}`)
+	estimate := EstimateTokens(body, 0)
+	principal := domain.Principal{ProjectID: "project-1", KeyID: "key-1", TokensPerMinute: estimate}
+	lease, err := limiter.Acquire(context.Background(), principal, body)
+	if err != nil {
+		t.Fatal(err)
+	}
+	lease.Release()
+	if _, err := limiter.Acquire(context.Background(), principal, body); !errors.Is(err, ErrTokensExceeded) {
+		t.Fatalf("second key token request error = %v, want token limit", err)
+	}
+}
+
 func TestEstimateTokensUsesLargestExplicitOutputLimit(t *testing.T) {
 	body := []byte(`{"max_tokens":10,"max_completion_tokens":25}`)
 	want := int64((len(body)+3)/4 + 25)