Skip to content

Commit d4b28bf

Browse files
authored
feat: max-model-multiplier-cap guardrail (#4215)
* Initial plan * feat: add max-model-multiplier guardrail Add a new guard that rejects API requests whose model's resolved cost multiplier exceeds the operator-configured cap (AWF_MAX_MODEL_MULTIPLIER). Changes: - containers/api-proxy/guards/max-model-multiplier-guard.js: new guard module - containers/api-proxy/proxy-request.js: extractModelFromBody(), guard integration - src/types/rate-limit-options.ts: maxModelMultiplier?: number field - src/config-file.ts: config-file mapping for apiProxy.maxModelMultiplier - src/services/api-proxy-service.ts: AWF_MAX_MODEL_MULTIPLIER env var passthrough - src/commands/build-config.ts: BuildConfigInputs.maxModelMultiplier? - src/commands/validators/log-and-limits.ts: parsing/validation - src/commands/validators/config-assembly.ts: passthrough to buildConfig - docs/awf-config.schema.json + src/awf-config-schema.json: schema field - Tests: guard unit tests, server integration tests, TS config/service tests * fix: rename cap option to maxModelMultiplierCap; fix guard JSDoc --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
1 parent ab2c781 commit d4b28bf

17 files changed

Lines changed: 1771 additions & 1108 deletions
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
'use strict';
2+
3+
const { sanitizeForLog } = require('../logging');
4+
5+
const maxModelMultiplierConfigCache = {
6+
rawCap: undefined,
7+
rawMultipliers: undefined,
8+
rawDefaultMultiplier: undefined,
9+
parsed: { cap: null, multipliers: {}, defaultMultiplier: 1 },
10+
};
11+
12+
function parseModelMultipliers(raw) {
13+
if (!raw || String(raw).trim() === '') return {};
14+
try {
15+
const parsed = JSON.parse(raw);
16+
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return {};
17+
const result = {};
18+
for (const [model, value] of Object.entries(parsed)) {
19+
const num = Number(value);
20+
if (Number.isFinite(num) && num > 0) {
21+
result[model] = num;
22+
}
23+
}
24+
return result;
25+
} catch {
26+
return {};
27+
}
28+
}
29+
30+
function parsePositiveNumber(raw) {
31+
const value = Number(raw);
32+
return Number.isFinite(value) && value > 0 ? value : null;
33+
}
34+
35+
function getMaxModelMultiplierConfig() {
36+
const rawCap = process.env.AWF_MAX_MODEL_MULTIPLIER;
37+
const rawMultipliers = process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS;
38+
const rawDefaultMultiplier = process.env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER;
39+
40+
if (
41+
maxModelMultiplierConfigCache.rawCap === rawCap &&
42+
maxModelMultiplierConfigCache.rawMultipliers === rawMultipliers &&
43+
maxModelMultiplierConfigCache.rawDefaultMultiplier === rawDefaultMultiplier
44+
) {
45+
return maxModelMultiplierConfigCache.parsed;
46+
}
47+
48+
maxModelMultiplierConfigCache.rawCap = rawCap;
49+
maxModelMultiplierConfigCache.rawMultipliers = rawMultipliers;
50+
maxModelMultiplierConfigCache.rawDefaultMultiplier = rawDefaultMultiplier;
51+
52+
const parsedMultipliers = Object.freeze(parseModelMultipliers(rawMultipliers));
53+
const configuredDefaultMultiplier = parsePositiveNumber(rawDefaultMultiplier);
54+
const cap = parsePositiveNumber(rawCap);
55+
56+
maxModelMultiplierConfigCache.parsed = Object.freeze({
57+
cap,
58+
multipliers: parsedMultipliers,
59+
defaultMultiplier: configuredDefaultMultiplier ?? 1,
60+
});
61+
return maxModelMultiplierConfigCache.parsed;
62+
}
63+
64+
function resolveMultiplierForModel(model, config) {
65+
if (Object.hasOwn(config.multipliers, model)) {
66+
return config.multipliers[model];
67+
}
68+
69+
let bestMatch = null;
70+
for (const [configuredModel, multiplier] of Object.entries(config.multipliers)) {
71+
if (model.startsWith(`${configuredModel}-`)) {
72+
if (!bestMatch || configuredModel.length > bestMatch.key.length) {
73+
bestMatch = { key: configuredModel, multiplier };
74+
}
75+
}
76+
}
77+
78+
if (bestMatch) return bestMatch.multiplier;
79+
return config.defaultMultiplier;
80+
}
81+
82+
/**
83+
* Returns a block state object when the given model's resolved multiplier
84+
* exceeds the configured cap (AWF_MAX_MODEL_MULTIPLIER), or null when no cap
85+
* is configured, the model is absent, or the multiplier is within the cap.
86+
*
87+
* @param {string|null} model - The model name from the request body (may be null)
88+
* @returns {{ model: string, multiplier: number, maxModelMultiplier: number } | null}
89+
*/
90+
function getModelMultiplierCapBlockState(model) {
91+
const config = getMaxModelMultiplierConfig();
92+
if (!config.cap || !model) return null;
93+
94+
const multiplier = resolveMultiplierForModel(model, config);
95+
if (multiplier <= config.cap) return null;
96+
97+
return {
98+
model: sanitizeForLog(model),
99+
multiplier,
100+
maxModelMultiplier: config.cap,
101+
};
102+
}
103+
104+
/**
105+
* Builds the structured error response body for a model-multiplier cap rejection.
106+
*
107+
* @param {{ model: string, multiplier: number, maxModelMultiplier: number }} state
108+
* @returns {{ error: object }}
109+
*/
110+
function buildModelMultiplierCapError(state) {
111+
return {
112+
error: {
113+
type: 'model_multiplier_cap_exceeded',
114+
message: `Model multiplier cap exceeded: model "${state.model}" has multiplier ${state.multiplier} which exceeds the configured maximum of ${state.maxModelMultiplier}.`,
115+
model: state.model,
116+
model_multiplier: state.multiplier,
117+
max_model_multiplier: state.maxModelMultiplier,
118+
},
119+
};
120+
}
121+
122+
/** @internal Test-only: reset cached config state between test cases. */
123+
function resetMaxModelMultiplierGuardForTests() {
124+
maxModelMultiplierConfigCache.rawCap = undefined;
125+
maxModelMultiplierConfigCache.rawMultipliers = undefined;
126+
maxModelMultiplierConfigCache.rawDefaultMultiplier = undefined;
127+
maxModelMultiplierConfigCache.parsed = { cap: null, multipliers: {}, defaultMultiplier: 1 };
128+
}
129+
130+
module.exports = {
131+
getModelMultiplierCapBlockState,
132+
buildModelMultiplierCapError,
133+
resetMaxModelMultiplierGuardForTests,
134+
};
Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
'use strict';
2+
3+
const {
4+
getModelMultiplierCapBlockState,
5+
buildModelMultiplierCapError,
6+
resetMaxModelMultiplierGuardForTests,
7+
} = require('./max-model-multiplier-guard');
8+
9+
describe('max-model-multiplier-guard', () => {
10+
beforeEach(() => {
11+
delete process.env.AWF_MAX_MODEL_MULTIPLIER;
12+
delete process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS;
13+
delete process.env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER;
14+
resetMaxModelMultiplierGuardForTests();
15+
});
16+
17+
afterEach(() => {
18+
delete process.env.AWF_MAX_MODEL_MULTIPLIER;
19+
delete process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS;
20+
delete process.env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER;
21+
resetMaxModelMultiplierGuardForTests();
22+
});
23+
24+
describe('getModelMultiplierCapBlockState', () => {
25+
it('returns null when AWF_MAX_MODEL_MULTIPLIER is not set', () => {
26+
expect(getModelMultiplierCapBlockState('claude-opus-4.7')).toBeNull();
27+
});
28+
29+
it('returns null when model is null', () => {
30+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
31+
expect(getModelMultiplierCapBlockState(null)).toBeNull();
32+
});
33+
34+
it('returns null when model is empty string', () => {
35+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
36+
expect(getModelMultiplierCapBlockState('')).toBeNull();
37+
});
38+
39+
it('returns null when model multiplier equals the cap', () => {
40+
process.env.AWF_MAX_MODEL_MULTIPLIER = '4';
41+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 4 });
42+
43+
expect(getModelMultiplierCapBlockState('gpt-4o')).toBeNull();
44+
});
45+
46+
it('returns block state when model multiplier exceeds the cap', () => {
47+
process.env.AWF_MAX_MODEL_MULTIPLIER = '4';
48+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({
49+
'claude-opus-4.7': 27,
50+
'gpt-4o': 2,
51+
});
52+
53+
const state = getModelMultiplierCapBlockState('claude-opus-4.7');
54+
expect(state).not.toBeNull();
55+
expect(state.multiplier).toBe(27);
56+
expect(state.maxModelMultiplier).toBe(4);
57+
expect(state.model).toBe('claude-opus-4.7');
58+
});
59+
60+
it('returns null when model multiplier is below the cap', () => {
61+
process.env.AWF_MAX_MODEL_MULTIPLIER = '4';
62+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({
63+
'claude-opus-4.7': 27,
64+
'gpt-4o': 2,
65+
});
66+
67+
expect(getModelMultiplierCapBlockState('gpt-4o')).toBeNull();
68+
});
69+
70+
it('resolves multiplier via prefix match and blocks when exceeds cap', () => {
71+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
72+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({
73+
'claude-opus-4.7': 27,
74+
});
75+
76+
const state = getModelMultiplierCapBlockState('claude-opus-4.7-20260501');
77+
expect(state).not.toBeNull();
78+
expect(state.multiplier).toBe(27);
79+
});
80+
81+
it('returns null when model is unknown and default multiplier (1) is within cap', () => {
82+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
83+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 2 });
84+
85+
expect(getModelMultiplierCapBlockState('unknown-model')).toBeNull();
86+
});
87+
88+
it('blocks when configured default multiplier for unknown model exceeds cap', () => {
89+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
90+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 2 });
91+
process.env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER = '10';
92+
93+
const state = getModelMultiplierCapBlockState('unknown-model');
94+
expect(state).not.toBeNull();
95+
expect(state.multiplier).toBe(10);
96+
});
97+
98+
it('blocks unknown models when default multiplier exceeds cap', () => {
99+
process.env.AWF_MAX_MODEL_MULTIPLIER = '3';
100+
process.env.AWF_EFFECTIVE_TOKEN_DEFAULT_MODEL_MULTIPLIER = '27';
101+
102+
const state = getModelMultiplierCapBlockState('any-model');
103+
expect(state).not.toBeNull();
104+
expect(state.multiplier).toBe(27);
105+
});
106+
107+
it('caches config across calls with same env vars', () => {
108+
process.env.AWF_MAX_MODEL_MULTIPLIER = '1';
109+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 2 });
110+
111+
const state1 = getModelMultiplierCapBlockState('gpt-4o');
112+
const state2 = getModelMultiplierCapBlockState('gpt-4o');
113+
expect(state1).not.toBeNull();
114+
expect(state1.multiplier).toBe(state2.multiplier);
115+
});
116+
117+
it('invalidates cache when env vars change', () => {
118+
process.env.AWF_MAX_MODEL_MULTIPLIER = '1';
119+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 2 });
120+
const state1 = getModelMultiplierCapBlockState('gpt-4o');
121+
expect(state1).not.toBeNull();
122+
expect(state1.multiplier).toBe(2);
123+
124+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 0.5 });
125+
resetMaxModelMultiplierGuardForTests();
126+
expect(getModelMultiplierCapBlockState('gpt-4o')).toBeNull(); // 0.5 <= 1 cap
127+
});
128+
129+
it('uses longest prefix match when multiple prefixes match', () => {
130+
process.env.AWF_MAX_MODEL_MULTIPLIER = '10';
131+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({
132+
'claude-opus': 20,
133+
'claude-opus-4.7': 27,
134+
});
135+
136+
const state = getModelMultiplierCapBlockState('claude-opus-4.7-20260501');
137+
expect(state).not.toBeNull();
138+
expect(state.multiplier).toBe(27); // longer match wins
139+
});
140+
});
141+
142+
describe('buildModelMultiplierCapError', () => {
143+
it('returns a structured error object', () => {
144+
const state = { model: 'claude-opus-4.7', multiplier: 27, maxModelMultiplier: 5 };
145+
const err = buildModelMultiplierCapError(state);
146+
147+
expect(err.error.type).toBe('model_multiplier_cap_exceeded');
148+
expect(err.error.model).toBe('claude-opus-4.7');
149+
expect(err.error.model_multiplier).toBe(27);
150+
expect(err.error.max_model_multiplier).toBe(5);
151+
expect(typeof err.error.message).toBe('string');
152+
expect(err.error.message).toContain('claude-opus-4.7');
153+
expect(err.error.message).toContain('27');
154+
expect(err.error.message).toContain('5');
155+
});
156+
});
157+
158+
describe('resetMaxModelMultiplierGuardForTests', () => {
159+
it('clears cached config so new env vars take effect', () => {
160+
process.env.AWF_MAX_MODEL_MULTIPLIER = '5';
161+
process.env.AWF_EFFECTIVE_TOKEN_MODEL_MULTIPLIERS = JSON.stringify({ 'gpt-4o': 10 });
162+
getModelMultiplierCapBlockState('gpt-4o'); // populate cache
163+
164+
// Reset and change env
165+
resetMaxModelMultiplierGuardForTests();
166+
delete process.env.AWF_MAX_MODEL_MULTIPLIER;
167+
168+
expect(getModelMultiplierCapBlockState('gpt-4o')).toBeNull();
169+
});
170+
});
171+
});

containers/api-proxy/proxy-request.js

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,11 @@ const {
4343
getAndClearPendingTimeoutSteeringMessage,
4444
resetTimeoutSteeringForTests,
4545
} = require('./guards/timeout-steering');
46+
const {
47+
getModelMultiplierCapBlockState,
48+
buildModelMultiplierCapError,
49+
resetMaxModelMultiplierGuardForTests,
50+
} = require('./guards/max-model-multiplier-guard');
4651

4752
// ── Optional token tracker (graceful degradation when not bundled) ────────────
4853
let trackTokenUsage;
@@ -128,6 +133,27 @@ function isValidRequestId(id) {
128133
return typeof id === 'string' && id.length <= 128 && /^[\w\-\.]+$/.test(id);
129134
}
130135

136+
/**
137+
* Attempt to extract the `model` field from a JSON request body.
138+
* Returns null for non-JSON bodies, bodies without a string `model` field,
139+
* or any parse failures.
140+
*
141+
* @param {Buffer} body
142+
* @returns {string|null}
143+
*/
144+
function extractModelFromBody(body) {
145+
if (!body || body.length === 0) return null;
146+
try {
147+
const parsed = JSON.parse(body.toString('utf8'));
148+
if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
149+
return typeof parsed.model === 'string' ? parsed.model : null;
150+
}
151+
return null;
152+
} catch {
153+
return null;
154+
}
155+
}
156+
131157
function handleRequestError(err, {
132158
res,
133159
requestId,
@@ -504,6 +530,28 @@ function proxyRequest(req, res, targetHost, injectHeaders, provider, basePath =
504530
return;
505531
}
506532

533+
if (req.method === 'POST' || req.method === 'PUT' || req.method === 'PATCH') {
534+
const bodyModel = extractModelFromBody(body);
535+
const mmBlock = getModelMultiplierCapBlockState(bodyModel);
536+
if (mmBlock) {
537+
const duration = Date.now() - startTime;
538+
metrics.gaugeDec('active_requests', { provider });
539+
metrics.increment('requests_total', { provider, method: req.method, status_class: '4xx' });
540+
metrics.observe('request_duration_ms', duration, { provider });
541+
logRequest('warn', 'model_multiplier_cap_exceeded', {
542+
request_id: requestId,
543+
provider,
544+
model: mmBlock.model,
545+
model_multiplier: mmBlock.multiplier,
546+
max_model_multiplier: mmBlock.maxModelMultiplier,
547+
});
548+
otel.endSpan(span, 400);
549+
res.writeHead(400, { 'Content-Type': 'application/json', 'X-Request-ID': requestId });
550+
res.end(JSON.stringify(buildModelMultiplierCapError(mmBlock)));
551+
return;
552+
}
553+
}
554+
507555
sendUpstreamRequest(headers, {
508556
body, targetHost, upstreamPath, req, res, provider, requestId, startTime, span, requestBytes,
509557
});
@@ -523,6 +571,7 @@ module.exports = {
523571
getMaxRunsReflectState,
524572
resetEffectiveTokenGuardForTests,
525573
resetMaxRunsGuardForTests,
574+
resetMaxModelMultiplierGuardForTests,
526575
resetTimeoutSteeringForTests,
527576
resetAnthropicDeprecatedBetaHeadersForTests: resetDeprecatedHeaderValuesForTests,
528577
getAndClearPendingSteeringMessage,

0 commit comments

Comments
 (0)