Add max-cache-misses guardrail for API proxy token budget enforcement (#5202)

* Initial plan

* feat: add max-cache-misses guardrail

* Fix help text: use === instead of = for cache_read_tokens comparison

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Author: Copilot

containers/api-proxy/guards/common-guard-checks.js

@@ -23,6 +23,8 @@
  * @param {Function} deps.buildEffectiveTokenLimitError
  * @param {Function} deps.getMaxRunsBlockState
  * @param {Function} deps.buildMaxRunsExceededError
+ * @param {Function} deps.getMaxCacheMissesBlockState
+ * @param {Function} deps.buildMaxCacheMissesExceededError
  * @param {Function} deps.getPermissionDeniedBlockState
  * @param {Function} deps.buildPermissionDeniedLimitError
  * @param {Function} deps.getAiCreditsBlockState
@@ -44,6 +46,8 @@ function buildCommonGuardChecks(deps, model) {
     buildEffectiveTokenLimitError,
     getMaxRunsBlockState,
     buildMaxRunsExceededError,
+    getMaxCacheMissesBlockState,
+    buildMaxCacheMissesExceededError,
     getPermissionDeniedBlockState,
     buildPermissionDeniedLimitError,
     getAiCreditsBlockState,
@@ -80,6 +84,17 @@ function buildCommonGuardChecks(deps, model) {
         max_runs: block.maxRuns,
       }),
     },
+    {
+      block: getMaxCacheMissesBlockState(),
+      isBlocked: block => block && block.maxExceeded,
+      statusCode: 429,
+      eventName: 'max_cache_misses_exceeded',
+      buildError: buildMaxCacheMissesExceededError,
+      buildLogFields: block => ({
+        consecutive_cache_misses: block.consecutiveCacheMisses,
+        max_cache_misses: block.maxCacheMisses,
+      }),
+    },
     {
       block: getPermissionDeniedBlockState(),
       isBlocked: block => block && block.maxExceeded,

containers/api-proxy/guards/max-cache-misses-guard.js

@@ -0,0 +1,104 @@
+'use strict';
+
+const { parsePositiveInteger } = require('./guard-utils');
+
+function createMaxCacheMissesState(configKey = null) {
+  return {
+    configKey,
+    consecutiveCacheMisses: 0,
+  };
+}
+
+let guardState = createMaxCacheMissesState();
+const configCache = { rawMax: undefined, parsed: null };
+
+function getMaxCacheMissesConfig() {
+  const rawMax = process.env.AWF_MAX_CACHE_MISSES;
+  if (configCache.rawMax === rawMax) return configCache.parsed;
+  configCache.rawMax = rawMax;
+  configCache.parsed = parsePositiveInteger(rawMax);
+  return configCache.parsed;
+}
+
+function getMaxCacheMissesState(maxCacheMisses) {
+  if (!maxCacheMisses) return null;
+  const configKey = String(maxCacheMisses);
+  if (guardState.configKey !== configKey) {
+    guardState = createMaxCacheMissesState(configKey);
+  }
+  return guardState;
+}
+
+function applyMaxCacheMissesUsage(normalizedUsage) {
+  const maxCacheMisses = getMaxCacheMissesConfig();
+  const state = getMaxCacheMissesState(maxCacheMisses);
+  if (!state || !normalizedUsage) return;
+
+  const inputTokens = normalizedUsage.input_tokens || 0;
+  const cacheReadTokens = normalizedUsage.cache_read_tokens || 0;
+
+  // Only runs with non-zero input tokens are considered for cache-miss streaks.
+  if (inputTokens <= 0) return;
+
+  if (cacheReadTokens > 0) {
+    state.consecutiveCacheMisses = 0;
+    return;
+  }
+
+  state.consecutiveCacheMisses += 1;
+}
+
+function getMaxCacheMissesBlockState() {
+  const maxCacheMisses = getMaxCacheMissesConfig();
+  const state = getMaxCacheMissesState(maxCacheMisses);
+  if (!state) return null;
+  return {
+    maxCacheMisses,
+    consecutiveCacheMisses: state.consecutiveCacheMisses,
+    maxExceeded: state.consecutiveCacheMisses >= maxCacheMisses,
+  };
+}
+
+function getMaxCacheMissesReflectState() {
+  const maxCacheMisses = getMaxCacheMissesConfig();
+  const state = getMaxCacheMissesState(maxCacheMisses);
+  if (!state) {
+    return {
+      enabled: false,
+      max_cache_misses: null,
+      consecutive_cache_misses: 0,
+      remaining_cache_misses: null,
+    };
+  }
+  return {
+    enabled: true,
+    max_cache_misses: maxCacheMisses,
+    consecutive_cache_misses: state.consecutiveCacheMisses,
+    remaining_cache_misses: Math.max(0, maxCacheMisses - state.consecutiveCacheMisses),
+  };
+}
+
+function resetMaxCacheMissesGuardForTests() {
+  guardState = createMaxCacheMissesState();
+  configCache.rawMax = undefined;
+  configCache.parsed = null;
+}
+
+function buildMaxCacheMissesExceededError(state) {
+  return {
+    error: {
+      type: 'max_cache_misses_exceeded',
+      message: `Maximum consecutive cache misses exceeded (${state.consecutiveCacheMisses} / ${state.maxCacheMisses}).`,
+      consecutive_cache_misses: state.consecutiveCacheMisses,
+      max_cache_misses: state.maxCacheMisses,
+    },
+  };
+}
+
+module.exports = {
+  applyMaxCacheMissesUsage,
+  getMaxCacheMissesBlockState,
+  getMaxCacheMissesReflectState,
+  resetMaxCacheMissesGuardForTests,
+  buildMaxCacheMissesExceededError,
+};

containers/api-proxy/guards/max-cache-misses-guard.test.js

@@ -0,0 +1,96 @@
+'use strict';
+
+const {
+  applyMaxCacheMissesUsage,
+  getMaxCacheMissesBlockState,
+  getMaxCacheMissesReflectState,
+  resetMaxCacheMissesGuardForTests,
+  buildMaxCacheMissesExceededError,
+} = require('./max-cache-misses-guard');
+
+describe('max-cache-misses-guard', () => {
+  beforeEach(() => {
+    delete process.env.AWF_MAX_CACHE_MISSES;
+    resetMaxCacheMissesGuardForTests();
+  });
+
+  afterEach(() => {
+    delete process.env.AWF_MAX_CACHE_MISSES;
+    resetMaxCacheMissesGuardForTests();
+  });
+
+  it('is disabled when AWF_MAX_CACHE_MISSES is not configured', () => {
+    applyMaxCacheMissesUsage({ input_tokens: 100, cache_read_tokens: 0 });
+    expect(getMaxCacheMissesBlockState()).toBeNull();
+    expect(getMaxCacheMissesReflectState()).toEqual({
+      enabled: false,
+      max_cache_misses: null,
+      consecutive_cache_misses: 0,
+      remaining_cache_misses: null,
+    });
+  });
+
+  it('tracks consecutive cache misses only for non-zero input runs', () => {
+    process.env.AWF_MAX_CACHE_MISSES = '3';
+    resetMaxCacheMissesGuardForTests();
+
+    applyMaxCacheMissesUsage({ input_tokens: 100, cache_read_tokens: 0 });
+    applyMaxCacheMissesUsage({ input_tokens: 0, cache_read_tokens: 0 });
+    applyMaxCacheMissesUsage({ input_tokens: 200, cache_read_tokens: 0 });
+
+    expect(getMaxCacheMissesBlockState()).toEqual({
+      maxCacheMisses: 3,
+      consecutiveCacheMisses: 2,
+      maxExceeded: false,
+    });
+  });
+
+  it('resets streak when cache_read_tokens is non-zero', () => {
+    process.env.AWF_MAX_CACHE_MISSES = '3';
+    resetMaxCacheMissesGuardForTests();
+
+    applyMaxCacheMissesUsage({ input_tokens: 100, cache_read_tokens: 0 });
+    applyMaxCacheMissesUsage({ input_tokens: 100, cache_read_tokens: 25 });
+
+    expect(getMaxCacheMissesBlockState()).toEqual({
+      maxCacheMisses: 3,
+      consecutiveCacheMisses: 0,
+      maxExceeded: false,
+    });
+  });
+
+  it('blocks once streak reaches the configured max', () => {
+    process.env.AWF_MAX_CACHE_MISSES = '2';
+    resetMaxCacheMissesGuardForTests();
+
+    applyMaxCacheMissesUsage({ input_tokens: 50, cache_read_tokens: 0 });
+    applyMaxCacheMissesUsage({ input_tokens: 60, cache_read_tokens: 0 });
+
+    expect(getMaxCacheMissesBlockState()).toEqual({
+      maxCacheMisses: 2,
+      consecutiveCacheMisses: 2,
+      maxExceeded: true,
+    });
+    expect(getMaxCacheMissesReflectState()).toEqual({
+      enabled: true,
+      max_cache_misses: 2,
+      consecutive_cache_misses: 2,
+      remaining_cache_misses: 0,
+    });
+  });
+
+  it('builds structured guard error payload', () => {
+    const error = buildMaxCacheMissesExceededError({
+      maxCacheMisses: 3,
+      consecutiveCacheMisses: 3,
+    });
+    expect(error).toEqual({
+      error: {
+        type: 'max_cache_misses_exceeded',
+        message: expect.stringContaining('3 / 3'),
+        consecutive_cache_misses: 3,
+        max_cache_misses: 3,
+      },
+    });
+  });
+});

containers/api-proxy/management.js

@@ -29,6 +29,7 @@ const metrics = require('./metrics');
  * @property {() => Record<string, { enabled: boolean, strategy: string, suppressed: boolean, suppression_reason?: string }>} getEffectiveModelFallback - Returns provider-effective fallback summary
  * @property {() => object}         getAiCreditsUsage     - Returns AI credits usage summary
  * @property {() => object}         getMaxRunsUsage        - Returns max-runs usage summary
+ * @property {() => object}         getMaxCacheMissesUsage - Returns max-cache-misses usage summary
  * @property {() => object}         getPermissionDeniedUsage - Returns permission-denied usage summary
  */
 
@@ -53,6 +54,7 @@ function createManagementHandlers(deps) {
     getEffectiveModelFallback,
     getAiCreditsUsage,
     getMaxRunsUsage,
+    getMaxCacheMissesUsage,
     getPermissionDeniedUsage,
   } = deps;
 
@@ -105,6 +107,7 @@ function createManagementHandlers(deps) {
       model_fallback_effective: getEffectiveModelFallback(),
       ai_credits: getAiCreditsUsage(),
       runs: getMaxRunsUsage(),
+      cache_misses: getMaxCacheMissesUsage(),
       permission_denied: getPermissionDeniedUsage(),
     };
   }

containers/api-proxy/proxy-request.js

@@ -39,6 +39,12 @@ const {
   resetMaxRunsGuardForTests,
   buildMaxRunsExceededError,
 } = require('./guards/max-runs-guard');
+const {
+  getMaxCacheMissesBlockState,
+  getMaxCacheMissesReflectState,
+  resetMaxCacheMissesGuardForTests,
+  buildMaxCacheMissesExceededError,
+} = require('./guards/max-cache-misses-guard');
 const {
   applyPermissionDenied,
   getPermissionDeniedBlockState,
@@ -216,6 +222,8 @@ const proxyWebSocket = createProxyWebSocket({
   buildEffectiveTokenLimitError,
   getMaxRunsBlockState,
   buildMaxRunsExceededError,
+  getMaxCacheMissesBlockState,
+  buildMaxCacheMissesExceededError,
   getPermissionDeniedBlockState,
   buildPermissionDeniedLimitError,
   getAiCreditsBlockState,
@@ -408,6 +416,8 @@ function enforceGuards({ body, provider, req, res, requestId, startTime, span, i
     buildEffectiveTokenLimitError,
     getMaxRunsBlockState,
     buildMaxRunsExceededError,
+    getMaxCacheMissesBlockState,
+    buildMaxCacheMissesExceededError,
     getPermissionDeniedBlockState,
     buildPermissionDeniedLimitError,
     getAiCreditsBlockState,
@@ -537,10 +547,12 @@ module.exports = {
   getEffectiveTokenReflectState,
   getAiCreditsReflectState,
   getMaxRunsReflectState,
+  getMaxCacheMissesReflectState,
   getPermissionDeniedReflectState,
   resetEffectiveTokenGuardForTests,
   resetAiCreditsGuardForTests,
   resetMaxRunsGuardForTests,
+  resetMaxCacheMissesGuardForTests,
   resetPermissionDeniedGuardForTests,
   resetMaxModelMultiplierGuardForTests,
   resetTimeoutSteeringForTests,

containers/api-proxy/server.js

@@ -50,6 +50,7 @@ const {
   getEffectiveTokenReflectState,
   getAiCreditsReflectState,
   getMaxRunsReflectState,
+  getMaxCacheMissesReflectState,
   getPermissionDeniedReflectState,
 } = require('./proxy-request');
 
@@ -128,6 +129,7 @@ const { healthResponse, reflectEndpoints, handleManagementEndpoint } = createMan
   getEffectiveTokenUsage: () => getEffectiveTokenReflectState(),
   getAiCreditsUsage: () => getAiCreditsReflectState(),
   getMaxRunsUsage: () => getMaxRunsReflectState(),
+  getMaxCacheMissesUsage: () => getMaxCacheMissesReflectState(),
   getPermissionDeniedUsage: () => getPermissionDeniedReflectState(),
 });
 

containers/api-proxy/server.network.test.js

@@ -415,6 +415,16 @@ describe('reflectEndpoints', () => {
     });
   });
 
+  it('should include cache_misses in reflect output', () => {
+    const result = reflectEndpoints();
+    expect(result.cache_misses).toEqual({
+      enabled: false,
+      max_cache_misses: null,
+      consecutive_cache_misses: 0,
+      remaining_cache_misses: null,
+    });
+  });
+
   it('should expose Copilot fallback suppression in reflect output for BYOK non-githubcopilot targets', () => {
     const prevTarget = process.env.COPILOT_API_TARGET;
     const prevProviderType = process.env.COPILOT_PROVIDER_TYPE;