diff --git a/containers/api-proxy/Dockerfile b/containers/api-proxy/Dockerfile index ca32a6a89..bbac404e4 100644 --- a/containers/api-proxy/Dockerfile +++ b/containers/api-proxy/Dockerfile @@ -24,6 +24,7 @@ COPY server.js logging.js metrics.js rate-limiter.js \ oidc-token-provider-base.js \ github-oidc.js aws-oidc-token-provider.js gcp-oidc-token-provider.js \ anthropic-oidc-token-provider.js \ + ai-credits-pricing.js \ oidc-refresh-utils.js body-transform.js body-utils.js rate-limit.js websocket-proxy.js \ deprecated-header-tracker.js billing-headers.js upstream-response.js \ anthropic-cache.js otel.js ./ diff --git a/containers/api-proxy/ai-credits-pricing.js b/containers/api-proxy/ai-credits-pricing.js new file mode 100644 index 000000000..6e1d53fa4 --- /dev/null +++ b/containers/api-proxy/ai-credits-pricing.js @@ -0,0 +1,28 @@ +'use strict'; + +// Per-model pricing in dollars per 1M tokens. +// Source: https://docs.github.com/en/copilot/reference/copilot-billing/models-and-pricing +module.exports = Object.freeze({ + 'gpt-5-mini': { input: 0.25, cachedInput: 0.025, cacheWrite: null, output: 2.00 }, + 'gpt-5.2': { input: 1.75, cachedInput: 0.175, cacheWrite: null, output: 14.00 }, + 'gpt-5.2-codex': { input: 1.75, cachedInput: 0.175, cacheWrite: null, output: 14.00 }, + 'gpt-5.3-codex': { input: 1.75, cachedInput: 0.175, cacheWrite: null, output: 14.00 }, + 'gpt-5.4': { input: 2.50, cachedInput: 0.25, cacheWrite: null, output: 15.00 }, + 'gpt-5.4-mini': { input: 0.75, cachedInput: 0.075, cacheWrite: null, output: 4.50 }, + 'gpt-5.4-nano': { input: 0.20, cachedInput: 0.02, cacheWrite: null, output: 1.25 }, + 'gpt-5.5': { input: 5.00, cachedInput: 0.50, cacheWrite: null, output: 30.00 }, + 'claude-haiku-4-5': { input: 1.00, cachedInput: 0.10, cacheWrite: 1.25, output: 5.00 }, + 'claude-sonnet-4': { input: 3.00, cachedInput: 0.30, cacheWrite: 3.75, output: 15.00 }, + 'claude-sonnet-4-5': { input: 3.00, cachedInput: 0.30, cacheWrite: 3.75, output: 15.00 }, + 'claude-sonnet-4-6': { input: 3.00, cachedInput: 0.30, cacheWrite: 3.75, output: 15.00 }, + 'claude-opus-4-5': { input: 5.00, cachedInput: 0.50, cacheWrite: 6.25, output: 25.00 }, + 'claude-opus-4-6': { input: 5.00, cachedInput: 0.50, cacheWrite: 6.25, output: 25.00 }, + 'claude-opus-4-7': { input: 5.00, cachedInput: 0.50, cacheWrite: 6.25, output: 25.00 }, + 'claude-opus-4-8': { input: 5.00, cachedInput: 0.50, cacheWrite: 6.25, output: 25.00 }, + 'gemini-2.5-pro': { input: 1.25, cachedInput: 0.125, cacheWrite: null, output: 10.00 }, + 'gemini-3-flash': { input: 0.50, cachedInput: 0.05, cacheWrite: null, output: 3.00 }, + 'gemini-3.1-pro': { input: 2.00, cachedInput: 0.20, cacheWrite: null, output: 12.00 }, + 'gemini-3.5-flash': { input: 1.50, cachedInput: 0.15, cacheWrite: null, output: 9.00 }, + 'mai-code-1-flash': { input: 0.75, cachedInput: 0.075, cacheWrite: null, output: 4.50 }, + 'raptor-mini': { input: 0.25, cachedInput: 0.025, cacheWrite: null, output: 2.00 }, +}); diff --git a/containers/api-proxy/guards/ai-credits-guard.js b/containers/api-proxy/guards/ai-credits-guard.js new file mode 100644 index 000000000..2164c2436 --- /dev/null +++ b/containers/api-proxy/guards/ai-credits-guard.js @@ -0,0 +1,173 @@ +'use strict'; + +const { logRequest, sanitizeForLog } = require('../logging'); +const pricingByModel = require('../ai-credits-pricing'); +const { parsePositiveNumber } = require('./guard-utils'); + +const TOKENS_PER_MILLION = 1_000_000; +const DOLLARS_PER_CREDIT = 0.01; +const CREDIT_DENOMINATOR = TOKENS_PER_MILLION * DOLLARS_PER_CREDIT; + +function roundCredits(value) { + return Math.round((value + Number.EPSILON) * 1_000_000) / 1_000_000; +} + +function createAiCreditsState() { + return { + totalAiCredits: 0, + byModel: {}, + warnedUnknownModels: new Set(), + }; +} + +let aiCreditsState = createAiCreditsState(); + +const aiCreditsConfigCache = { + rawMax: undefined, + parsed: { max: null }, +}; + +function getAiCreditsConfig() { + const rawMax = process.env.AWF_MAX_AI_CREDITS; + if (aiCreditsConfigCache.rawMax === rawMax) { + return aiCreditsConfigCache.parsed; + } + aiCreditsConfigCache.rawMax = rawMax; + aiCreditsConfigCache.parsed = { + max: parsePositiveNumber(rawMax), + }; + return aiCreditsConfigCache.parsed; +} + +function resolveModelPricing(model, state = aiCreditsState) { + if (Object.hasOwn(pricingByModel, model)) return pricingByModel[model]; + + let prefixMatch = null; + for (const [configuredModel, pricing] of Object.entries(pricingByModel)) { + if (model.startsWith(`${configuredModel}-`)) { + if (!prefixMatch || configuredModel.length > prefixMatch.model.length) { + prefixMatch = { model: configuredModel, pricing }; + } + } + } + if (prefixMatch) return prefixMatch.pricing; + + if (!state.warnedUnknownModels.has(model)) { + logRequest('warn', 'unknown_model_ai_credits_pricing', { + model: sanitizeForLog(model), + }); + state.warnedUnknownModels.add(model); + } + return null; +} + +function calculateAiCredits(normalizedUsage, model, state = aiCreditsState) { + const pricing = resolveModelPricing(model, state); + if (!pricing) return null; + + const inputCredits = ((normalizedUsage.input_tokens || 0) * pricing.input) / CREDIT_DENOMINATOR; + const cachedInputCredits = ((normalizedUsage.cache_read_tokens || 0) * pricing.cachedInput) / CREDIT_DENOMINATOR; + const cacheWriteCredits = pricing.cacheWrite + ? ((normalizedUsage.cache_write_tokens || 0) * pricing.cacheWrite) / CREDIT_DENOMINATOR + : 0; + const outputCredits = ((normalizedUsage.output_tokens || 0) * pricing.output) / CREDIT_DENOMINATOR; + const totalCredits = inputCredits + cachedInputCredits + cacheWriteCredits + outputCredits; + + return { + inputCredits, + cachedInputCredits, + cacheWriteCredits, + outputCredits, + totalCredits, + }; +} + +function applyAiCreditsUsage(normalizedUsage, model) { + if (!normalizedUsage) return null; + const safeModel = model || 'unknown'; + const calc = calculateAiCredits(normalizedUsage, safeModel); + if (!calc) return null; + + if (!Object.hasOwn(aiCreditsState.byModel, safeModel)) { + aiCreditsState.byModel[safeModel] = { + inputCredits: 0, + cachedInputCredits: 0, + cacheWriteCredits: 0, + outputCredits: 0, + totalCredits: 0, + }; + } + + const modelBucket = aiCreditsState.byModel[safeModel]; + modelBucket.inputCredits += calc.inputCredits; + modelBucket.cachedInputCredits += calc.cachedInputCredits; + modelBucket.cacheWriteCredits += calc.cacheWriteCredits; + modelBucket.outputCredits += calc.outputCredits; + modelBucket.totalCredits += calc.totalCredits; + aiCreditsState.totalAiCredits += calc.totalCredits; + + process.env.AWF_AI_CREDITS_USED = String(roundCredits(aiCreditsState.totalAiCredits)); + + return { + aiCreditsThisResponse: roundCredits(calc.totalCredits), + inputCreditsThisResponse: roundCredits(calc.inputCredits), + cachedInputCreditsThisResponse: roundCredits(calc.cachedInputCredits), + cacheWriteCreditsThisResponse: roundCredits(calc.cacheWriteCredits), + outputCreditsThisResponse: roundCredits(calc.outputCredits), + totalAiCredits: roundCredits(aiCreditsState.totalAiCredits), + }; +} + +function getAiCreditsReflectState() { + const byModel = {}; + for (const [model, usage] of Object.entries(aiCreditsState.byModel)) { + byModel[model] = { + input_credits: roundCredits(usage.inputCredits), + cached_input_credits: roundCredits(usage.cachedInputCredits), + cache_write_credits: roundCredits(usage.cacheWriteCredits), + output_credits: roundCredits(usage.outputCredits), + total: roundCredits(usage.totalCredits), + }; + } + return { + total: roundCredits(aiCreditsState.totalAiCredits), + by_model: byModel, + }; +} + +function getAiCreditsBlockState() { + const config = getAiCreditsConfig(); + if (!config.max) return null; + const roundedTotalAiCredits = roundCredits(aiCreditsState.totalAiCredits); + return { + maxAiCredits: config.max, + totalAiCredits: roundedTotalAiCredits, + maxExceeded: roundedTotalAiCredits >= config.max, + }; +} + +function buildAiCreditsLimitError(aiCreditsBlockState) { + return { + error: { + type: 'ai_credits_limit_exceeded', + message: `Maximum AI credits exceeded (${aiCreditsBlockState.totalAiCredits.toFixed(6)} / ${aiCreditsBlockState.maxAiCredits}).`, + total_ai_credits: aiCreditsBlockState.totalAiCredits, + max_ai_credits: aiCreditsBlockState.maxAiCredits, + }, + }; +} + +function resetAiCreditsGuardForTests() { + aiCreditsState = createAiCreditsState(); + aiCreditsConfigCache.rawMax = undefined; + aiCreditsConfigCache.parsed = { max: null }; + delete process.env.AWF_AI_CREDITS_USED; +} + +module.exports = { + applyAiCreditsUsage, + getAiCreditsReflectState, + getAiCreditsBlockState, + buildAiCreditsLimitError, + resetAiCreditsGuardForTests, +}; diff --git a/containers/api-proxy/guards/ai-credits-guard.test.js b/containers/api-proxy/guards/ai-credits-guard.test.js new file mode 100644 index 000000000..c88e6de45 --- /dev/null +++ b/containers/api-proxy/guards/ai-credits-guard.test.js @@ -0,0 +1,107 @@ +const { + applyAiCreditsUsage, + getAiCreditsReflectState, + getAiCreditsBlockState, + buildAiCreditsLimitError, + resetAiCreditsGuardForTests, +} = require('./ai-credits-guard'); +const { collectLogOutput } = require('../test-helpers/log-test-helpers'); + +describe('ai-credits-guard', () => { + let originalMaxAiCredits; + + beforeEach(() => { + originalMaxAiCredits = process.env.AWF_MAX_AI_CREDITS; + delete process.env.AWF_MAX_AI_CREDITS; + resetAiCreditsGuardForTests(); + }); + + afterEach(() => { + resetAiCreditsGuardForTests(); + if (originalMaxAiCredits === undefined) { + delete process.env.AWF_MAX_AI_CREDITS; + } else { + process.env.AWF_MAX_AI_CREDITS = originalMaxAiCredits; + } + jest.restoreAllMocks(); + }); + + it('calculates and accumulates AI credits by model', () => { + const usage = applyAiCreditsUsage({ + input_tokens: 1000, + cache_read_tokens: 100, + output_tokens: 500, + }, 'gpt-5-mini'); + + expect(usage).toMatchObject({ + aiCreditsThisResponse: 0.12525, + totalAiCredits: 0.12525, + }); + expect(process.env.AWF_AI_CREDITS_USED).toBe('0.12525'); + expect(getAiCreditsReflectState()).toEqual({ + total: 0.12525, + by_model: { + 'gpt-5-mini': { + input_credits: 0.025, + cached_input_credits: 0.00025, + cache_write_credits: 0, + output_credits: 0.1, + total: 0.12525, + }, + }, + }); + }); + + it('matches pricing table entries by model prefix', () => { + const usage = applyAiCreditsUsage({ + input_tokens: 2000, + cache_read_tokens: 1000, + cache_write_tokens: 500, + output_tokens: 100, + }, 'claude-sonnet-4-6-20260601'); + + expect(usage.aiCreditsThisResponse).toBeCloseTo(0.9675, 10); + expect(getAiCreditsReflectState().by_model['claude-sonnet-4-6-20260601'].total).toBeCloseTo(0.9675, 10); + }); + + it('warns and skips usage for unknown models', () => { + const { lines } = collectLogOutput(); + const usage = applyAiCreditsUsage({ input_tokens: 100 }, 'unknown-model'); + + expect(usage).toBeNull(); + expect(getAiCreditsReflectState()).toEqual({ total: 0, by_model: {} }); + expect(lines).toContainEqual(expect.objectContaining({ + event: 'unknown_model_ai_credits_pricing', + level: 'warn', + model: 'unknown-model', + })); + }); + + it('reports block state when max ai credits is configured and exceeded', () => { + process.env.AWF_MAX_AI_CREDITS = '0.1'; + applyAiCreditsUsage({ + input_tokens: 1000, + output_tokens: 500, + }, 'gpt-5-mini'); + + expect(getAiCreditsBlockState()).toEqual({ + maxAiCredits: 0.1, + totalAiCredits: 0.125, + maxExceeded: true, + }); + }); + + it('builds a structured max ai credits limit error payload', () => { + expect(buildAiCreditsLimitError({ + totalAiCredits: 0.125, + maxAiCredits: 0.1, + })).toEqual({ + error: { + type: 'ai_credits_limit_exceeded', + message: 'Maximum AI credits exceeded (0.125000 / 0.1).', + total_ai_credits: 0.125, + max_ai_credits: 0.1, + }, + }); + }); +}); diff --git a/containers/api-proxy/management.js b/containers/api-proxy/management.js index 7078c6f13..d1264d0ad 100644 --- a/containers/api-proxy/management.js +++ b/containers/api-proxy/management.js @@ -28,6 +28,7 @@ const metrics = require('./metrics'); * @property {() => { enabled: boolean, strategy: string }} getModelFallback - Returns fallback config * @property {() => Record} getEffectiveModelFallback - Returns provider-effective fallback summary * @property {() => object} getEffectiveTokenUsage - Returns effective token usage summary + * @property {() => object} getAiCreditsUsage - Returns AI credits usage summary * @property {() => object} getMaxRunsUsage - Returns max-runs usage summary * @property {() => object} getPermissionDeniedUsage - Returns permission-denied usage summary */ @@ -52,6 +53,7 @@ function createManagementHandlers(deps) { getModelFallback, getEffectiveModelFallback, getEffectiveTokenUsage, + getAiCreditsUsage, getMaxRunsUsage, getPermissionDeniedUsage, } = deps; @@ -104,6 +106,7 @@ function createManagementHandlers(deps) { model_fallback: getModelFallback(), model_fallback_effective: getEffectiveModelFallback(), effective_tokens: getEffectiveTokenUsage(), + ai_credits: getAiCreditsUsage(), runs: getMaxRunsUsage(), permission_denied: getPermissionDeniedUsage(), }; diff --git a/containers/api-proxy/proxy-request.js b/containers/api-proxy/proxy-request.js index efd433973..bda005f75 100644 --- a/containers/api-proxy/proxy-request.js +++ b/containers/api-proxy/proxy-request.js @@ -55,6 +55,13 @@ const { buildModelMultiplierCapError, resetMaxModelMultiplierGuardForTests, } = require('./guards/max-model-multiplier-guard'); +const { + applyAiCreditsUsage, + getAiCreditsReflectState, + getAiCreditsBlockState, + buildAiCreditsLimitError, + resetAiCreditsGuardForTests, +} = require('./guards/ai-credits-guard'); // ── Optional token tracker (graceful degradation when not bundled) ──────────── let trackTokenUsage; @@ -214,8 +221,11 @@ const proxyWebSocket = createProxyWebSocket({ buildMaxRunsExceededError, getPermissionDeniedBlockState, buildPermissionDeniedLimitError, + getAiCreditsBlockState, + buildAiCreditsLimitError, trackWebSocketTokenUsage, applyEffectiveTokenUsage, + applyAiCreditsUsage, }); // ── Proxy helpers ───────────────────────────────────────────────────────────── @@ -281,6 +291,7 @@ const { handleUpstreamResponse } = createUpstreamResponseHandlers({ handleRequestError, trackTokenUsage, applyEffectiveTokenUsage, + applyAiCreditsUsage, applyMaxRunsInvocation, applyPermissionDenied, extractBillingHeaders, @@ -558,6 +569,24 @@ function proxyRequest(req, res, targetHost, injectHeaders, provider, basePath = return; } + const aiCreditsBlock = getAiCreditsBlockState(); + if (aiCreditsBlock && aiCreditsBlock.maxExceeded) { + const duration = Date.now() - startTime; + metrics.gaugeDec('active_requests', { provider }); + metrics.increment('requests_total', { provider, method: req.method, status_class: '4xx' }); + metrics.observe('request_duration_ms', duration, { provider }); + logRequest('warn', 'ai_credits_limit_exceeded', { + request_id: requestId, + provider, + total_ai_credits: aiCreditsBlock.totalAiCredits, + max_ai_credits: aiCreditsBlock.maxAiCredits, + }); + otel.endSpan(span, 429); + res.writeHead(429, { 'Content-Type': 'application/json', 'X-Request-ID': requestId }); + res.end(JSON.stringify(buildAiCreditsLimitError(aiCreditsBlock))); + return; + } + if (req.method === 'POST' || req.method === 'PUT' || req.method === 'PATCH') { const bodyModel = extractModelFromBody(body); const mmBlock = getModelMultiplierCapBlockState(bodyModel); @@ -596,9 +625,11 @@ module.exports = { proxyAgent, HTTPS_PROXY, getEffectiveTokenReflectState, + getAiCreditsReflectState, getMaxRunsReflectState, getPermissionDeniedReflectState, resetEffectiveTokenGuardForTests, + resetAiCreditsGuardForTests, resetMaxRunsGuardForTests, resetPermissionDeniedGuardForTests, resetMaxModelMultiplierGuardForTests, diff --git a/containers/api-proxy/server.js b/containers/api-proxy/server.js index 850db624f..f01b1e1e0 100644 --- a/containers/api-proxy/server.js +++ b/containers/api-proxy/server.js @@ -48,6 +48,7 @@ const { HTTPS_PROXY, extractBillingHeaders, getEffectiveTokenReflectState, + getAiCreditsReflectState, getMaxRunsReflectState, getPermissionDeniedReflectState, } = require('./proxy-request'); @@ -125,6 +126,7 @@ const { healthResponse, reflectEndpoints, handleManagementEndpoint } = createMan getModelFallback: () => MODEL_FALLBACK, getEffectiveModelFallback: () => getEffectiveModelFallbackForReflect(registeredAdapters), getEffectiveTokenUsage: () => getEffectiveTokenReflectState(), + getAiCreditsUsage: () => getAiCreditsReflectState(), getMaxRunsUsage: () => getMaxRunsReflectState(), getPermissionDeniedUsage: () => getPermissionDeniedReflectState(), }); diff --git a/containers/api-proxy/server.network.test.js b/containers/api-proxy/server.network.test.js index c7827e2c8..4df927e6e 100644 --- a/containers/api-proxy/server.network.test.js +++ b/containers/api-proxy/server.network.test.js @@ -407,6 +407,14 @@ describe('reflectEndpoints', () => { }); }); + it('should include ai_credits in reflect output', () => { + const result = reflectEndpoints(); + expect(result.ai_credits).toEqual({ + total: 0, + by_model: {}, + }); + }); + 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; diff --git a/containers/api-proxy/server.token-guards.test.js b/containers/api-proxy/server.token-guards.test.js index 1297351ad..b65890c4e 100644 --- a/containers/api-proxy/server.token-guards.test.js +++ b/containers/api-proxy/server.token-guards.test.js @@ -10,6 +10,7 @@ const { EventEmitter } = require('events'); const { makeReq: makeReqFactory, makeRes, + getStructuredLogs, setupServerTestEnv, } = require('./test-helpers/server-mock-factories'); @@ -18,6 +19,7 @@ let resetEffectiveTokenGuardForTests; let resetMaxRunsGuardForTests; let resetPermissionDeniedGuardForTests; let resetMaxModelMultiplierGuardForTests; +let resetAiCreditsGuardForTests; setupServerTestEnv(() => { ({ proxyRequest } = require('./server')); @@ -26,8 +28,9 @@ setupServerTestEnv(() => { resetMaxRunsGuardForTests, resetPermissionDeniedGuardForTests, resetMaxModelMultiplierGuardForTests, + resetAiCreditsGuardForTests, } = require('./proxy-request')); - return { proxyRequest, resetEffectiveTokenGuardForTests, resetMaxRunsGuardForTests, resetPermissionDeniedGuardForTests, resetMaxModelMultiplierGuardForTests }; + return { proxyRequest, resetEffectiveTokenGuardForTests, resetMaxRunsGuardForTests, resetPermissionDeniedGuardForTests, resetMaxModelMultiplierGuardForTests, resetAiCreditsGuardForTests }; }); describe('proxyRequest effective token guard', () => { @@ -41,12 +44,14 @@ describe('proxyRequest effective token guard', () => { process.env.AWF_MAX_EFFECTIVE_TOKENS = '10'; delete process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS; resetEffectiveTokenGuardForTests(); + resetAiCreditsGuardForTests(); }); afterEach(() => { delete process.env.AWF_MAX_EFFECTIVE_TOKENS; delete process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS; resetEffectiveTokenGuardForTests(); + resetAiCreditsGuardForTests(); jest.restoreAllMocks(); }); @@ -94,6 +99,47 @@ describe('proxyRequest effective token guard', () => { expect(payload.error.max_effective_tokens).toBe(10); expect(payload.error.total_effective_tokens).toBeGreaterThanOrEqual(10); }); + + it('logs ai credits and effective tokens for each response usage update', () => { + const writeSpy = jest.spyOn(process.stdout, 'write').mockImplementation(() => true); + let responseHandler; + const upstreamRequest = new EventEmitter(); + upstreamRequest.end = jest.fn(); + upstreamRequest.write = jest.fn(); + upstreamRequest.destroy = jest.fn(); + jest.spyOn(https, 'request').mockImplementation((options, cb) => { + responseHandler = cb; + return upstreamRequest; + }); + + const req = makeReq(); + const res = makeRes(); + proxyRequest(req, res, 'api.openai.com', { Authorization: '******' }, 'openai'); + req.emit('end'); + + const proxyRes = new EventEmitter(); + proxyRes.statusCode = 200; + proxyRes.headers = { 'content-type': 'application/json' }; + proxyRes.pipe = jest.fn(); + responseHandler(proxyRes); + proxyRes.emit('data', Buffer.from(JSON.stringify({ + model: 'gpt-5-mini', + usage: { prompt_tokens: 1000, completion_tokens: 500 }, + }))); + proxyRes.emit('end'); + + const budgetLogs = getStructuredLogs(writeSpy, 'token_budget_usage'); + expect(budgetLogs).toEqual(expect.arrayContaining([ + expect.objectContaining({ + effective_tokens_this_response: 3000, + ai_credits_this_response: 0.125, + ai_credits_total: 0.125, + }), + ])); + expect(process.env.AWF_AI_CREDITS_USED).toBe('0.125'); + + writeSpy.mockRestore(); + }); }); describe('proxyRequest max-runs guard', () => { @@ -175,6 +221,70 @@ describe('proxyRequest max-runs guard', () => { }); }); +describe('proxyRequest max-ai-credits guard', () => { + function makeReq(headers = {}) { + return makeReqFactory('/v1/chat/completions', headers); + } + + beforeEach(() => { + process.env.AWF_MAX_AI_CREDITS = '0.1'; + delete process.env.AWF_MAX_EFFECTIVE_TOKENS; + resetEffectiveTokenGuardForTests(); + resetAiCreditsGuardForTests(); + }); + + afterEach(() => { + delete process.env.AWF_MAX_AI_CREDITS; + resetEffectiveTokenGuardForTests(); + resetAiCreditsGuardForTests(); + jest.restoreAllMocks(); + }); + + it('returns 429 with structured payload when ai credits limit is reached', () => { + let responseHandler; + const upstreamRequest = new EventEmitter(); + upstreamRequest.end = jest.fn(); + upstreamRequest.write = jest.fn(); + upstreamRequest.destroy = jest.fn(); + + const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => { + responseHandler = cb; + return upstreamRequest; + }); + + const req1 = makeReq(); + const res1 = makeRes(); + proxyRequest(req1, res1, 'api.openai.com', { Authorization: '******' }, 'openai'); + req1.emit('end'); + + const proxyRes = new EventEmitter(); + proxyRes.statusCode = 200; + proxyRes.headers = { 'content-type': 'application/json' }; + proxyRes.pipe = jest.fn(); + + responseHandler(proxyRes); + proxyRes.emit('data', Buffer.from(JSON.stringify({ + model: 'gpt-5-mini', + usage: { prompt_tokens: 1000, completion_tokens: 500 }, + }))); + proxyRes.emit('end'); + + const req2 = makeReq(); + const res2 = makeRes(); + proxyRequest(req2, res2, 'api.openai.com', { Authorization: '******' }, 'openai'); + req2.emit('end'); + + expect(httpsRequestSpy).toHaveBeenCalledTimes(1); + expect(res2.writeHead).toHaveBeenCalledWith(429, expect.objectContaining({ + 'Content-Type': 'application/json', + })); + const payload = JSON.parse(res2.end.mock.calls[0][0]); + expect(payload.error.type).toBe('ai_credits_limit_exceeded'); + expect(payload.error.max_ai_credits).toBe(0.1); + expect(payload.error.total_ai_credits).toBeGreaterThanOrEqual(0.1); + }); +}); + describe('proxyRequest permission-denied guard', () => { function makeReq(headers = {}) { return makeReqFactory('/v1/chat/completions', headers); diff --git a/containers/api-proxy/upstream-response.js b/containers/api-proxy/upstream-response.js index b03e84e18..fe22c8937 100644 --- a/containers/api-proxy/upstream-response.js +++ b/containers/api-proxy/upstream-response.js @@ -29,6 +29,7 @@ function createUpstreamResponseHandlers({ handleRequestError, trackTokenUsage, applyEffectiveTokenUsage, + applyAiCreditsUsage, applyMaxRunsInvocation, applyPermissionDenied, extractBillingHeaders, @@ -198,7 +199,18 @@ function createUpstreamResponseHandlers({ requestId, provider, path: sanitizeForLog(req.url), startTime, metrics, billingInfo, initiatorSent, onUsage: (normalizedUsage, model) => { otel.setTokenAttributes(span, { provider, model, normalizedUsage, streaming: isStreaming }); - applyEffectiveTokenUsage(normalizedUsage, model); + const effectiveTokenUsage = applyEffectiveTokenUsage(normalizedUsage, model); + const aiCreditsUsage = applyAiCreditsUsage(normalizedUsage, model); + if (effectiveTokenUsage || aiCreditsUsage) { + logRequest('info', 'token_budget_usage', { + request_id: requestId, + provider, + model: model || 'unknown', + effective_tokens_this_response: effectiveTokenUsage?.effectiveTokensThisResponse ?? null, + ai_credits_this_response: aiCreditsUsage?.aiCreditsThisResponse ?? null, + ai_credits_total: aiCreditsUsage?.totalAiCredits ?? null, + }); + } }, onSpanEnd: (statusCode) => { otel.endSpan(span, statusCode); diff --git a/containers/api-proxy/websocket-proxy.js b/containers/api-proxy/websocket-proxy.js index 5aad74a14..a6c420105 100644 --- a/containers/api-proxy/websocket-proxy.js +++ b/containers/api-proxy/websocket-proxy.js @@ -20,8 +20,11 @@ function createProxyWebSocket({ buildMaxRunsExceededError, getPermissionDeniedBlockState, buildPermissionDeniedLimitError, + getAiCreditsBlockState, + buildAiCreditsLimitError, trackWebSocketTokenUsage, applyEffectiveTokenUsage, + applyAiCreditsUsage, }) { /** * Handle a WebSocket upgrade request by tunnelling through the Squid proxy. @@ -105,6 +108,20 @@ function createProxyWebSocket({ return; } + const aiCreditsBlock = getAiCreditsBlockState(); + if (aiCreditsBlock && aiCreditsBlock.maxExceeded) { + logRequest('warn', 'ai_credits_limit_exceeded', { + request_id: requestId, + provider, + total_ai_credits: aiCreditsBlock.totalAiCredits, + max_ai_credits: aiCreditsBlock.maxAiCredits, + }); + socket.write('HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n'); + socket.write(JSON.stringify(buildAiCreditsLimitError(aiCreditsBlock))); + socket.destroy(); + return; + } + const rateCheck = limiter.check(provider, 0); if (!rateCheck.allowed) { metrics.increment('rate_limit_rejected_total', { provider, limit_type: rateCheck.limitType }); @@ -215,7 +232,18 @@ function createProxyWebSocket({ startTime, metrics, onUsage: (normalizedUsage, model) => { - applyEffectiveTokenUsage(normalizedUsage, model); + const effectiveTokenUsage = applyEffectiveTokenUsage(normalizedUsage, model); + const aiCreditsUsage = applyAiCreditsUsage(normalizedUsage, model); + if (effectiveTokenUsage || aiCreditsUsage) { + logRequest('info', 'token_budget_usage', { + request_id: requestId, + provider, + model: model || 'unknown', + effective_tokens_this_response: effectiveTokenUsage?.effectiveTokensThisResponse ?? null, + ai_credits_this_response: aiCreditsUsage?.aiCreditsThisResponse ?? null, + ai_credits_total: aiCreditsUsage?.totalAiCredits ?? null, + }); + } }, }); diff --git a/docs/awf-config-spec.md b/docs/awf-config-spec.md index f6cd5da90..2a1e5c443 100644 --- a/docs/awf-config-spec.md +++ b/docs/awf-config-spec.md @@ -709,6 +709,18 @@ current effective-token state: When `maxEffectiveTokens` is not configured, the `enabled` field MUST be `false` and numeric fields MUST be `0` or `null`. +### 10.7 Max AI Credits Configuration + +`maxAiCredits` is a positive number. It is supplied via the AWF config file +(including stdin config via `--config -`) and maps to the +`AWF_MAX_AI_CREDITS` environment variable injected into the api-proxy +container. + +When configured, the proxy MUST enforce this budget in addition to any +configured `maxEffectiveTokens` budget. Once cumulative AI credits reach or +exceed `maxAiCredits`, subsequent requests MUST be rejected with HTTP `429` +and error type `ai_credits_limit_exceeded`. + ## 11. Max-Runs Enforcement *This section is normative.* diff --git a/docs/awf-config.schema.json b/docs/awf-config.schema.json index 7c2b9749d..6e3cbe2d1 100644 --- a/docs/awf-config.schema.json +++ b/docs/awf-config.schema.json @@ -72,6 +72,11 @@ "minimum": 1, "description": "Maximum cumulative effective tokens allowed for a run. When reached, the API proxy rejects subsequent requests with HTTP 429 and error type 'effective_tokens_limit_exceeded'. Tokens are weighted: input ×1, cache-read ×0.1, output ×4, reasoning ×4. See spec §10." }, + "maxAiCredits": { + "type": "number", + "exclusiveMinimum": 0, + "description": "Maximum cumulative AI credits allowed for a run. When reached, the API proxy rejects subsequent requests with HTTP 429 and error type 'ai_credits_limit_exceeded'." + }, "modelMultipliers": { "type": "object", "description": "Per-model multipliers for effective token accounting. Each model's weighted tokens are multiplied by this value before accumulation. Unlisted models use defaultModelMultiplier when set, otherwise the highest configured multiplier. See spec §10.2.", diff --git a/src/awf-config-schema.json b/src/awf-config-schema.json index 7c2b9749d..6e3cbe2d1 100644 --- a/src/awf-config-schema.json +++ b/src/awf-config-schema.json @@ -72,6 +72,11 @@ "minimum": 1, "description": "Maximum cumulative effective tokens allowed for a run. When reached, the API proxy rejects subsequent requests with HTTP 429 and error type 'effective_tokens_limit_exceeded'. Tokens are weighted: input ×1, cache-read ×0.1, output ×4, reasoning ×4. See spec §10." }, + "maxAiCredits": { + "type": "number", + "exclusiveMinimum": 0, + "description": "Maximum cumulative AI credits allowed for a run. When reached, the API proxy rejects subsequent requests with HTTP 429 and error type 'ai_credits_limit_exceeded'." + }, "modelMultipliers": { "type": "object", "description": "Per-model multipliers for effective token accounting. Each model's weighted tokens are multiplied by this value before accumulation. Unlisted models use defaultModelMultiplier when set, otherwise the highest configured multiplier. See spec §10.2.", diff --git a/src/commands/build-config.test.ts b/src/commands/build-config.test.ts index 39438ab1c..5fd5e9799 100644 --- a/src/commands/build-config.test.ts +++ b/src/commands/build-config.test.ts @@ -35,6 +35,7 @@ function makeInputs(overrides: Partial[0]> = {}): agentImage: undefined, modelAliases: undefined, maxEffectiveTokens: undefined, + maxAiCredits: undefined, effectiveTokenModelMultipliers: undefined, effectiveTokenDefaultModelMultiplier: undefined, maxRuns: undefined, diff --git a/src/commands/build-config.ts b/src/commands/build-config.ts index ac752657a..9bd29c07d 100644 --- a/src/commands/build-config.ts +++ b/src/commands/build-config.ts @@ -22,6 +22,7 @@ interface BuildConfigInputs { agentImage: string | undefined; modelAliases: Record | undefined; maxEffectiveTokens: number | undefined; + maxAiCredits: number | undefined; effectiveTokenModelMultipliers: Record | undefined; effectiveTokenDefaultModelMultiplier: number | undefined; maxModelMultiplierCap?: number; @@ -57,6 +58,7 @@ export function buildConfig(inputs: BuildConfigInputs): WrapperConfig { agentImage, modelAliases, maxEffectiveTokens, + maxAiCredits, effectiveTokenModelMultipliers, effectiveTokenDefaultModelMultiplier, maxModelMultiplierCap, @@ -112,6 +114,7 @@ export function buildConfig(inputs: BuildConfigInputs): WrapperConfig { anthropicCacheTailTtl: options.anthropicCacheTailTtl as '5m' | '1h' | undefined, modelAliases, maxEffectiveTokens, + maxAiCredits, effectiveTokenModelMultipliers, effectiveTokenDefaultModelMultiplier, maxModelMultiplierCap, diff --git a/src/commands/validate-options.test.ts b/src/commands/validate-options.test.ts index eb1cb8d75..f2e713dbe 100644 --- a/src/commands/validate-options.test.ts +++ b/src/commands/validate-options.test.ts @@ -69,6 +69,7 @@ const STUB_CONFIG = { anthropicCacheTailTtl: undefined, modelAliases: undefined, maxEffectiveTokens: undefined, + maxAiCredits: undefined, effectiveTokenModelMultipliers: undefined, effectiveTokenDefaultModelMultiplier: undefined, maxRuns: undefined, @@ -214,6 +215,25 @@ describe('validateOptions', () => { consoleSpy.mockRestore(); }); + describe('maxAiCredits validation', () => { + it('exits when maxAiCredits is not a positive number', () => { + const consoleSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined); + expect(() => + validateOptions({ logLevel: 'info', maxAiCredits: 'abc' }, 'echo hi'), + ).toThrow('process.exit called'); + expect(consoleSpy).toHaveBeenCalledWith(expect.stringContaining('Invalid maxAiCredits')); + consoleSpy.mockRestore(); + }); + + it('exits when maxAiCredits is zero', () => { + const consoleSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined); + expect(() => + validateOptions({ logLevel: 'info', maxAiCredits: 0 }, 'echo hi'), + ).toThrow('process.exit called'); + consoleSpy.mockRestore(); + }); + }); + it('exits when maxEffectiveTokens is zero', () => { const consoleSpy = jest.spyOn(console, 'error').mockImplementation(() => undefined); expect(() => diff --git a/src/commands/validators/config-assembly.test.ts b/src/commands/validators/config-assembly.test.ts index 2fe319ec5..6e3ca3b30 100644 --- a/src/commands/validators/config-assembly.test.ts +++ b/src/commands/validators/config-assembly.test.ts @@ -87,6 +87,7 @@ describe('config-assembly', () => { agentImage: undefined, modelAliases: {}, maxEffectiveTokens: undefined, + maxAiCredits: undefined, effectiveTokenModelMultipliers: {}, effectiveTokenDefaultModelMultiplier: undefined, maxRuns: undefined, diff --git a/src/commands/validators/config-assembly.ts b/src/commands/validators/config-assembly.ts index b9a55aeee..9c15699b4 100644 --- a/src/commands/validators/config-assembly.ts +++ b/src/commands/validators/config-assembly.ts @@ -88,6 +88,7 @@ export function assembleAndValidateConfig( agentImage: logAndLimits.agentImage, modelAliases: logAndLimits.modelAliases, maxEffectiveTokens: logAndLimits.maxEffectiveTokens, + maxAiCredits: logAndLimits.maxAiCredits, effectiveTokenModelMultipliers: logAndLimits.effectiveTokenModelMultipliers, effectiveTokenDefaultModelMultiplier: logAndLimits.effectiveTokenDefaultModelMultiplier, maxModelMultiplierCap: logAndLimits.maxModelMultiplierCap, diff --git a/src/commands/validators/log-and-limits.ts b/src/commands/validators/log-and-limits.ts index 7db8e861a..03ed13e73 100644 --- a/src/commands/validators/log-and-limits.ts +++ b/src/commands/validators/log-and-limits.ts @@ -16,6 +16,7 @@ export interface LogAndLimitsResult { logLevel: LogLevel; modelAliases: Record | undefined; maxEffectiveTokens: number | undefined; + maxAiCredits: number | undefined; effectiveTokenModelMultipliers: Record | undefined; effectiveTokenDefaultModelMultiplier: number | undefined; maxModelMultiplierCap?: number; @@ -60,6 +61,10 @@ export function validateLogAndLimits(options: Record): LogAndLi | string | number | undefined; + const maxAiCreditsOption = (options as Record).maxAiCredits as + | string + | number + | undefined; const effectiveTokenDefaultModelMultiplierOption = (options as Record) .effectiveTokenDefaultModelMultiplier as string | number | undefined; // Config-file multipliers (already a Record) @@ -85,6 +90,8 @@ export function validateLogAndLimits(options: Record): LogAndLi : undefined; const maxEffectiveTokens = maxEffectiveTokensOption !== undefined ? Number(maxEffectiveTokensOption) : undefined; + const maxAiCredits = + maxAiCreditsOption !== undefined ? Number(maxAiCreditsOption) : undefined; const effectiveTokenDefaultModelMultiplier = effectiveTokenDefaultModelMultiplierOption !== undefined ? Number(effectiveTokenDefaultModelMultiplierOption) @@ -98,6 +105,14 @@ export function validateLogAndLimits(options: Record): LogAndLi process.exit(1); } + if ( + maxAiCredits !== undefined && + (!Number.isFinite(maxAiCredits) || maxAiCredits <= 0) + ) { + console.error('Error: Invalid maxAiCredits value (must be > 0)'); + process.exit(1); + } + if ( effectiveTokenDefaultModelMultiplier !== undefined && (!Number.isFinite(effectiveTokenDefaultModelMultiplier) || effectiveTokenDefaultModelMultiplier <= 0) @@ -172,6 +187,7 @@ export function validateLogAndLimits(options: Record): LogAndLi logLevel, modelAliases, maxEffectiveTokens, + maxAiCredits, effectiveTokenModelMultipliers, effectiveTokenDefaultModelMultiplier, maxModelMultiplierCap, diff --git a/src/config-file-mapping.test.ts b/src/config-file-mapping.test.ts index c70575852..1caddf97a 100644 --- a/src/config-file-mapping.test.ts +++ b/src/config-file-mapping.test.ts @@ -142,6 +142,7 @@ describe('mapAwfFileConfigToCliOptions', () => { const result = mapAwfFileConfigToCliOptions({ apiProxy: { maxEffectiveTokens: 6000, + maxAiCredits: 1.2, modelMultipliers: { 'gpt-4o': 2, 'claude-sonnet-4': 1.5, @@ -151,6 +152,7 @@ describe('mapAwfFileConfigToCliOptions', () => { }, }); expect(result.maxEffectiveTokens).toBe(6000); + expect(result.maxAiCredits).toBe(1.2); expect(result.effectiveTokenModelMultipliers).toEqual({ 'gpt-4o': 2, 'claude-sonnet-4': 1.5, diff --git a/src/config-file-validation.test.ts b/src/config-file-validation.test.ts index 7a1281ab5..b7565b478 100644 --- a/src/config-file-validation.test.ts +++ b/src/config-file-validation.test.ts @@ -110,6 +110,7 @@ describe('validateAwfFileConfig', () => { expect(validateAwfFileConfig({ apiProxy: { maxEffectiveTokens: 5000, + maxAiCredits: 1.5, modelMultipliers: { 'gpt-4o': 2, 'claude-sonnet-4': 1.5 }, defaultModelMultiplier: 27, maxModelMultiplierCap: 5, @@ -118,6 +119,8 @@ describe('validateAwfFileConfig', () => { expect(validateAwfFileConfig({ apiProxy: { maxEffectiveTokens: 0 } })) .toContain('config.apiProxy.maxEffectiveTokens must be a positive integer'); + expect(validateAwfFileConfig({ apiProxy: { maxAiCredits: 0 } })) + .toContain('config.apiProxy.maxAiCredits must be > 0'); expect(validateAwfFileConfig({ apiProxy: { modelMultipliers: { 'gpt-4o': 0 } } })) .toContain('config.apiProxy.modelMultipliers.gpt-4o must be > 0'); expect(validateAwfFileConfig({ apiProxy: { defaultModelMultiplier: 0 } })) diff --git a/src/config-file.ts b/src/config-file.ts index 6342a289f..1b99fb6b7 100644 --- a/src/config-file.ts +++ b/src/config-file.ts @@ -18,6 +18,7 @@ interface AwfFileConfig { anthropicAutoCache?: boolean; anthropicCacheTailTtl?: string; maxEffectiveTokens?: number; + maxAiCredits?: number; modelMultipliers?: Record; defaultModelMultiplier?: number; maxModelMultiplierCap?: number; @@ -198,6 +199,7 @@ export function mapAwfFileConfigToCliOptions(config: AwfFileConfig): Record { anthropicAutoCache: true, anthropicCacheTailTtl: '5m', maxEffectiveTokens: 100000, + maxAiCredits: 5.5, modelMultipliers: { 'gpt-4o': 2, 'claude-sonnet-4': 1.5, @@ -149,6 +150,8 @@ describe('awf-config.schema.json', () => { it('validates effective-token guard apiProxy fields', () => { expect(validate({ apiProxy: { maxEffectiveTokens: 1000 } })).toBe(true); expect(validate({ apiProxy: { maxEffectiveTokens: 0 } })).toBe(false); + expect(validate({ apiProxy: { maxAiCredits: 1.2 } })).toBe(true); + expect(validate({ apiProxy: { maxAiCredits: 0 } })).toBe(false); expect(validate({ apiProxy: { modelMultipliers: { 'gpt-4o': 2, 'claude': 1.5 } } })).toBe(true); expect(validate({ apiProxy: { modelMultipliers: { 'gpt-4o': 0 } } })).toBe(false); expect(validate({ apiProxy: { defaultModelMultiplier: 27 } })).toBe(true); diff --git a/src/services/api-proxy-service-config.ts b/src/services/api-proxy-service-config.ts index b38f28e6b..d217da3c5 100644 --- a/src/services/api-proxy-service-config.ts +++ b/src/services/api-proxy-service-config.ts @@ -146,6 +146,9 @@ export function buildApiProxyServiceConfig(params: ApiProxyServiceConfigParams): ...(config.maxEffectiveTokens !== undefined && { AWF_MAX_EFFECTIVE_TOKENS: String(config.maxEffectiveTokens), }), + ...(config.maxAiCredits !== undefined && { + AWF_MAX_AI_CREDITS: String(config.maxAiCredits), + }), ...(config.effectiveTokenModelMultipliers && { AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS: JSON.stringify(config.effectiveTokenModelMultipliers), }), diff --git a/src/services/api-proxy-service-rate-limit.test.ts b/src/services/api-proxy-service-rate-limit.test.ts index 9239ae7b8..500201e59 100644 --- a/src/services/api-proxy-service-rate-limit.test.ts +++ b/src/services/api-proxy-service-rate-limit.test.ts @@ -64,6 +64,7 @@ describe('API proxy sidecar: rate limiting and token guard', () => { enableApiProxy: true, openaiApiKey: 'sk-test-key', maxEffectiveTokens: 5000, + maxAiCredits: 1.25, effectiveTokenModelMultipliers: { 'gpt-4o': 2, 'claude-sonnet-4': 1.5, @@ -75,11 +76,19 @@ describe('API proxy sidecar: rate limiting and token guard', () => { const proxy = result.services['api-proxy']; const env = proxy.environment as Record; expect(env.AWF_MAX_EFFECTIVE_TOKENS).toBe('5000'); + expect(env.AWF_MAX_AI_CREDITS).toBe('1.25'); expect(env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS).toBe('{"gpt-4o":2,"claude-sonnet-4":1.5}'); expect(env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER).toBe('27'); expect(env.AWF_MAX_MODEL_MULTIPLIER).toBe('4'); }); + it('should not set AWF_MAX_AI_CREDITS in api-proxy when maxAiCredits is not configured', () => { + const result = generateDockerCompose({ ...mockConfig, enableApiProxy: true, openaiApiKey: 'sk-test-key' }, mockNetworkConfigWithProxy); + const proxy = result.services['api-proxy']; + const env = proxy.environment as Record; + expect(env.AWF_MAX_AI_CREDITS).toBeUndefined(); + }); + it('should set AWF_MAX_MODEL_MULTIPLIER when maxModelMultiplierCap is configured', () => { const configWithCap = { ...mockConfig, diff --git a/src/types/rate-limit-options.ts b/src/types/rate-limit-options.ts index 36c2fd0bd..0155e5587 100644 --- a/src/types/rate-limit-options.ts +++ b/src/types/rate-limit-options.ts @@ -23,6 +23,15 @@ export interface RateLimitOptions { */ maxEffectiveTokens?: number; + /** + * Maximum total AI credits allowed for the current AWF run. + * + * When set, the API proxy tracks AI credits across responses using + * model-specific pricing and rejects additional requests once this limit is + * reached. + */ + maxAiCredits?: number; + /** * Model-specific multipliers used by effective token accounting. *