From a921300a535717d0d6970126ffa508704134770c Mon Sep 17 00:00:00 2001
From: diegosouzapw <diegosouza.pw@gmail.com>
Date: Sat, 30 May 2026 21:00:36 -0300
Subject: [PATCH] feat(quota): force quota-exclusive keys onto pool connection
 in account selection (Phase A4)

---
 src/lib/quota/quotaKey.ts                     | 21 ++++++++++
 src/sse/handlers/chat.ts                      | 27 ++++++++++++-
 .../unit/quota-constrain-connections.test.ts  | 38 +++++++++++++++++++
 3 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 tests/unit/quota-constrain-connections.test.ts

diff --git a/src/lib/quota/quotaKey.ts b/src/lib/quota/quotaKey.ts
index 27c0a1311e..5ec87bffac 100644
--- a/src/lib/quota/quotaKey.ts
+++ b/src/lib/quota/quotaKey.ts
@@ -20,6 +20,27 @@ export interface QuotaKeyScope {
   providers: string[];
 }
 
+/**
+ * Constrain an existing connection allow-list to the connections belonging to a
+ * quota key's pool scope.
+ *
+ * Semantics mirror `intersectAllowedConnectionIds` in chat.ts:
+ *  - Empty `quotaConnectionIds` (non-quota key)  → return `existing` unchanged.
+ *  - Empty / null `existing` (no prior constraint) → return `quotaConnectionIds`.
+ *  - Both non-empty                               → intersection.
+ *  - Disjoint sets                               → empty array (no eligible connection).
+ *
+ * This is a pure, synchronous function — easy to unit-test without DB setup.
+ */
+export function constrainConnectionsToQuota(
+  existing: string[],
+  quotaConnectionIds: string[]
+): string[] {
+  if (quotaConnectionIds.length === 0) return existing;
+  if (existing.length === 0) return quotaConnectionIds;
+  return existing.filter((id) => quotaConnectionIds.includes(id));
+}
+
 /**
  * Given the `allowedQuotas` field of an API key (array of quota-pool IDs),
  * returns the set of connection IDs and provider slugs that the key is
diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts
index 9449e7700b..04aea3f5bb 100644
--- a/src/sse/handlers/chat.ts
+++ b/src/sse/handlers/chat.ts
@@ -96,6 +96,10 @@ import {
   resolveCooldownAwareRetrySettings,
   waitForCooldownAwareRetry,
 } from "../services/cooldownAwareRetry";
+import {
+  constrainConnectionsToQuota,
+  resolveQuotaKeyScope,
+} from "../../lib/quota/quotaKey";
 
 registerCodexQuotaFetcher();
 
@@ -460,11 +464,20 @@ export async function handleChat(request: any, clientRawRequest: any = null) {
       const resolvedModel = modelInfo.model || modelString;
       const hasForcedConnection =
         typeof target?.connectionId === "string" && target.connectionId.trim().length > 0;
-      const allowedConnections = intersectAllowedConnectionIds(
+      let allowedConnections = intersectAllowedConnectionIds(
         apiKeyInfo?.allowedConnections ?? null,
         target?.allowedConnectionIds ?? null
       );
 
+      // A4: quota-exclusive keys must only use the pool's connection(s).
+      if (apiKeyInfo?.allowedQuotas && apiKeyInfo.allowedQuotas.length > 0) {
+        const quotaScope = await resolveQuotaKeyScope(apiKeyInfo.allowedQuotas);
+        allowedConnections = constrainConnectionsToQuota(
+          allowedConnections ?? [],
+          quotaScope.connectionIds
+        );
+      }
+
       if (Array.isArray(allowedConnections) && allowedConnections.length === 0) {
         return false;
       }
@@ -731,10 +744,20 @@ async function handleSingleModelChat(
   const hasForcedConnection =
     typeof runtimeOptions.forcedConnectionId === "string" &&
     runtimeOptions.forcedConnectionId.trim().length > 0;
-  const effectiveAllowedConnections = intersectAllowedConnectionIds(
+  let effectiveAllowedConnections = intersectAllowedConnectionIds(
     apiKeyInfo?.allowedConnections ?? null,
     runtimeOptions.allowedConnectionIds ?? null
   );
+
+  // A4: quota-exclusive keys must only use the pool's connection(s).
+  if (apiKeyInfo?.allowedQuotas && apiKeyInfo.allowedQuotas.length > 0) {
+    const quotaScope = await resolveQuotaKeyScope(apiKeyInfo.allowedQuotas);
+    effectiveAllowedConnections = constrainConnectionsToQuota(
+      effectiveAllowedConnections ?? [],
+      quotaScope.connectionIds
+    );
+  }
+
   const bypassReason = forceLiveComboTest
     ? "combo live test"
     : hasForcedConnection
diff --git a/tests/unit/quota-constrain-connections.test.ts b/tests/unit/quota-constrain-connections.test.ts
new file mode 100644
index 0000000000..6fc8de4d1a
--- /dev/null
+++ b/tests/unit/quota-constrain-connections.test.ts
@@ -0,0 +1,38 @@
+/**
+ * tests/unit/quota-constrain-connections.test.ts
+ *
+ * TDD coverage for src/lib/quota/quotaKey.ts::constrainConnectionsToQuota.
+ * Pure-function tests — no DB setup needed.
+ *
+ * Cases:
+ *  1. Non-quota key (empty quotaConnectionIds) → existing returned unchanged.
+ *  2. No prior constraint (empty existing) + quota [c1,c2] → [c1,c2].
+ *  3. Existing [c1,c2,c3] ∩ quota [c2] → [c2].
+ *  4. Disjoint sets → [].
+ */
+
+import { test } from "node:test";
+import assert from "node:assert/strict";
+
+import { constrainConnectionsToQuota } from "../../src/lib/quota/quotaKey.ts";
+
+test("non-quota key: empty quotaConnectionIds returns existing unchanged", () => {
+  const existing = ["c1", "c2", "c3"];
+  const result = constrainConnectionsToQuota(existing, []);
+  assert.deepStrictEqual(result, existing);
+});
+
+test("no prior constraint: empty existing with quota [c1,c2] returns [c1,c2]", () => {
+  const result = constrainConnectionsToQuota([], ["c1", "c2"]);
+  assert.deepStrictEqual(result, ["c1", "c2"]);
+});
+
+test("intersection: existing [c1,c2,c3] ∩ quota [c2] returns [c2]", () => {
+  const result = constrainConnectionsToQuota(["c1", "c2", "c3"], ["c2"]);
+  assert.deepStrictEqual(result, ["c2"]);
+});
+
+test("disjoint sets: no common elements returns empty array", () => {
+  const result = constrainConnectionsToQuota(["c1", "c2"], ["c3", "c4"]);
+  assert.deepStrictEqual(result, []);
+});