fix: per‑key traffic fields

Author: riturajFi

packages/core/src/agent/agent.ts

@@ -270,6 +270,22 @@ export interface BaseGenerationOptions extends Partial<CallSettings> {
   userId?: string;
   conversationId?: string;
   tenantId?: string;
+  /**
+   * Optional key metadata for per-key rate limits.
+   */
+  apiKeyId?: string;
+  /**
+   * Optional region metadata for per-region rate limits.
+   */
+  region?: string;
+  /**
+   * Optional endpoint metadata for per-endpoint rate limits.
+   */
+  endpoint?: string;
+  /**
+   * Optional tenant tier metadata for per-tier rate limits.
+   */
+  tenantTier?: string;
   context?: ContextInput;
   elicitation?: (request: unknown) => Promise<unknown>;
   /**
@@ -2336,6 +2352,10 @@ export class Agent {
     const startTimeDate = new Date();
     const priority = this.resolveTrafficPriority(options);
     const tenantId = this.resolveTenantId(options);
+    const apiKeyId = options?.apiKeyId ?? options?.parentOperationContext?.apiKeyId;
+    const region = options?.region ?? options?.parentOperationContext?.region;
+    const endpoint = options?.endpoint ?? options?.parentOperationContext?.endpoint;
+    const tenantTier = options?.tenantTier ?? options?.parentOperationContext?.tenantTier;
 
     // Prefer reusing an existing context instance to preserve reference across calls/subagents
     const runtimeContext = toContextMap(options?.context);
@@ -2445,6 +2465,10 @@ export class Agent {
       userId: options?.userId,
       conversationId: options?.conversationId,
       tenantId,
+      apiKeyId,
+      region,
+      endpoint,
+      tenantTier,
       parentAgentId: options?.parentAgentId,
       traceContext,
       startTime: startTimeDate,
@@ -4152,6 +4176,10 @@ export class Agent {
       this.resolveProvider(this.model) ??
       undefined;
     const priority = this.resolveTrafficPriority(options);
+    const apiKeyId = options?.apiKeyId ?? options?.parentOperationContext?.apiKeyId;
+    const region = options?.region ?? options?.parentOperationContext?.region;
+    const endpoint = options?.endpoint ?? options?.parentOperationContext?.endpoint;
+    const tenantTier = options?.tenantTier ?? options?.parentOperationContext?.tenantTier;
 
     return {
       agentId: this.id, // Identify which agent issued the request
@@ -4160,6 +4188,10 @@ export class Agent {
       provider, // Allows per-provider throttling later
       priority,
       tenantId: this.resolveTenantId(options),
+      apiKeyId,
+      region,
+      endpoint,
+      tenantTier,
       taskType: options?.taskType,
       fallbackPolicyId: options?.fallbackPolicyId,
     };

packages/core/src/agent/types.ts

@@ -900,6 +900,18 @@ export type OperationContext = {
   /** Optional tenant identifier propagated across nested operations */
   tenantId?: string;
 
+  /** Optional key identifier for per-key traffic limits */
+  apiKeyId?: string;
+
+  /** Optional region identifier for per-region traffic limits */
+  region?: string;
+
+  /** Optional endpoint identifier for per-endpoint traffic limits */
+  endpoint?: string;
+
+  /** Optional tenant tier identifier for per-tier traffic limits */
+  tenantTier?: string;
+
   /** User-managed context map for this operation */
   readonly context: Map<string | symbol, unknown>;
 

packages/core/src/traffic/traffic-circuit-breaker.ts

@@ -359,12 +359,19 @@ export class TrafficCircuitBreaker {
   ): void {
     next.request = fallbackRequest;
     next.attempt = 1;
+    next.estimatedTokens = fallbackRequest.estimatedTokens;
+    next.reservedTokens = undefined;
     next.tenantConcurrencyKey = undefined;
     next.providerModelConcurrencyKey = undefined;
     next.rateLimitKey = undefined;
     next.etaMs = undefined;
     next.circuitKey = undefined;
     next.circuitStatus = undefined;
+    next.extractUsage = fallbackRequest.extractUsage;
+    if (context?.reason === "queue-timeout") {
+      next.enqueuedAt = Date.now();
+      next.dispatchedAt = undefined;
+    }
     logger?.debug?.("Switched to fallback request", {
       previousCircuitKey: context?.previousCircuitKey,
       fallbackModel: fallback,

packages/core/src/traffic/traffic-controller.spec.ts

@@ -488,3 +488,84 @@ describe("TrafficController stream reporting", () => {
     expect(order).toEqual(["fallback"]);
   });
 });
+
+describe("TrafficController queue timeouts", () => {
+  it("lets fallback requests wait after queue timeout without rejecting", async () => {
+    vi.useFakeTimers();
+
+    try {
+      vi.setSystemTime(new Date(0));
+      const controller = new TrafficController({
+        maxConcurrent: 1,
+        fallbackChains: {
+          "p::m": ["m-fallback"],
+        },
+      });
+      const order: string[] = [];
+      let releaseFirst!: () => void;
+      const firstGate = new Promise<void>((resolve) => {
+        releaseFirst = resolve;
+      });
+
+      const first = controller.handleText({
+        tenantId: "tenant-a",
+        metadata: { provider: "p", model: "m", priority: "P1" },
+        execute: async () => {
+          order.push("first");
+          await firstGate;
+          return "first";
+        },
+      });
+
+      const second = controller.handleText({
+        tenantId: "tenant-a",
+        metadata: { provider: "p", model: "m", priority: "P1" },
+        maxQueueWaitMs: 1,
+        execute: async () => {
+          order.push("primary");
+          return "primary";
+        },
+        createFallbackRequest: (target) => ({
+          tenantId: "tenant-a",
+          metadata: {
+            provider: "p",
+            model: typeof target === "string" ? target : target.model,
+            priority: "P1",
+          },
+          maxQueueWaitMs: 1,
+          execute: async () => {
+            order.push("fallback");
+            return "fallback";
+          },
+        }),
+      });
+
+      await Promise.resolve();
+      expect(order).toEqual(["first"]);
+
+      await vi.advanceTimersByTimeAsync(2);
+
+      const third = controller.handleText({
+        tenantId: "tenant-a",
+        metadata: { provider: "p", model: "other", priority: "P1" },
+        execute: async () => {
+          order.push("third");
+          return "third";
+        },
+      });
+
+      await Promise.resolve();
+      expect(order).toEqual(["first"]);
+
+      releaseFirst();
+      await vi.runAllTimersAsync();
+
+      await expect(second).resolves.toBe("fallback");
+      await Promise.all([first, third]);
+
+      expect(order).toEqual(["first", "fallback", "third"]);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+});

packages/core/src/traffic/traffic-controller.ts

@@ -857,7 +857,7 @@ export class TrafficController {
       this.trafficLogger,
     );
     if (fallbackApplied) {
-      return "expired";
+      return "none";
     }
 
     const timeoutError = this.createQueueTimeoutError(next, now);