Skip to content

Commit b1c8a50

Browse files
aidandaly24aidandaly24
authored
feat(gateway): add custom claims validation and TUI wizard for JWT auth (#599)
* feat(gateway): add custom claims validation and TUI wizard for JWT auth Add custom JWT claims validation support and a full TUI wizard flow for configuring Custom JWT gateway authorization. Schema: - Add ClaimMatchOperator, ClaimMatchValue, InboundTokenClaimValueType, and CustomClaimValidation schemas with strict validation - Add customClaims to CustomJwtAuthorizerConfigSchema and deployed-state - Add --custom-claims CLI flag with JSON parsing and validation TUI Wizard: - Expand JWT config flow with custom claims manager (add/edit/done) - Add claim name, operator, value, and value type sub-steps - Show human-readable claim summary in confirm review - Make client credentials optional (skip with empty Enter) Testing: - Add AddGatewayJwtConfig.test.tsx — full TUI component tests - Add finishJwtConfig.test.ts — unit tests for config assembly - Extend useAddGatewayWizard.test.tsx with JWT + custom claims flows - Add GatewayPrimitive.test.ts for custom claims round-trip - Extend validate.test.ts with custom claims validation cases - Add TUI integration test (add-gateway-jwt.test.ts) Constraint: Stacked on fix/inbound-auth-hardening (#598) Confidence: high Scope-risk: moderate * fix(gateway): improve custom claim form navigation UX Enter now advances to the next field instead of immediately submitting, and up/down arrow keys navigate between fields for a more intuitive form experience. * fix(gateway): position cursor before placeholder in custom claim form When a text field is empty, the cursor now appears before the placeholder hint instead of after it, matching expected input behavior. * test(gateway): update claim form test for Enter-advances-fields behavior The test expected Enter to immediately submit and show a validation error, but Enter now advances to the next field. Updated the test to press Enter through all fields before expecting the submission validation error. * fix: restore CLIENT_ID env var and move inline import to top-level Restore writing both CLIENT_ID and CLIENT_SECRET to .env in createManagedOAuthCredential, matching main branch behavior. Move dynamic import of policyEnginePrimitive to a static top-level import per AGENTS.md conventions. * style: run prettier and fix test prop Run prettier on 3 files and add missing existingPolicyEngines prop to AddGatewayJwtConfig test defaults. * ci: retrigger checks
1 parent 8758f9b commit b1c8a50

18 files changed

Lines changed: 2370 additions & 160 deletions

‎integ-tests/tui/add-gateway-jwt.test.ts‎

Lines changed: 415 additions & 0 deletions
Large diffs are not rendered by default.

‎integ-tests/tui/setup.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
/**
2+
* TUI integration test setup.
3+
*
4+
* This file is referenced by vitest.config.ts as a setupFile for the 'tui' project.
5+
* It runs before each test file in integ-tests/tui/.
6+
*/

‎src/cli/commands/add/__tests__/validate.test.ts‎

Lines changed: 91 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -218,40 +218,29 @@ describe('validate', () => {
218218
expect(result.error?.includes('Invalid authorizer type')).toBeTruthy();
219219
});
220220

221-
// AC11: CUSTOM_JWT requires discoveryUrl
222-
it('returns error for CUSTOM_JWT missing discoveryUrl', () => {
223-
const opts = { ...validGatewayOptionsJwt, discoveryUrl: undefined };
224-
const result = validateAddGatewayOptions(opts);
221+
// AC11: CUSTOM_JWT requires discoveryUrl; at least one of allowedAudience/allowedClients/allowedScopes
222+
it('returns error for CUSTOM_JWT missing required fields', () => {
223+
// discoveryUrl is always required
224+
const result = validateAddGatewayOptions({ ...validGatewayOptionsJwt, discoveryUrl: undefined });
225225
expect(result.valid).toBe(false);
226226
expect(result.error).toBe('--discovery-url is required for CUSTOM_JWT authorizer');
227-
});
228227

229-
// AC11b: at least one of audience/clients/scopes required
230-
it('returns error when all of audience, clients, and scopes are missing', () => {
231-
const opts = {
228+
// All three optional fields absent fails
229+
const noneResult = validateAddGatewayOptions({
232230
...validGatewayOptionsJwt,
233231
allowedAudience: undefined,
234232
allowedClients: undefined,
235233
allowedScopes: undefined,
236-
};
237-
const result = validateAddGatewayOptions(opts);
238-
expect(result.valid).toBe(false);
239-
expect(result.error).toContain('At least one of');
240-
});
241-
242-
it('allows CUSTOM_JWT with only allowedScopes', () => {
243-
const opts = {
244-
...validGatewayOptionsJwt,
245-
allowedAudience: undefined,
246-
allowedClients: undefined,
247-
allowedScopes: 'scope1',
248-
};
249-
const result = validateAddGatewayOptions(opts);
250-
expect(result.valid).toBe(true);
234+
});
235+
expect(noneResult.valid).toBe(false);
236+
expect(noneResult.error).toBe(
237+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer'
238+
);
251239
});
252240

253-
it('allows CUSTOM_JWT with only allowedAudience', () => {
254-
const opts = { ...validGatewayOptionsJwt, allowedClients: undefined, allowedScopes: undefined };
241+
// AC11b: allowedAudience is optional
242+
it('allows CUSTOM_JWT without allowedAudience', () => {
243+
const opts = { ...validGatewayOptionsJwt, allowedAudience: undefined };
255244
const result = validateAddGatewayOptions(opts);
256245
expect(result.valid).toBe(true);
257246
});
@@ -269,21 +258,88 @@ describe('validate', () => {
269258
expect(result.error?.includes('.well-known/openid-configuration')).toBeTruthy();
270259
});
271260

272-
it('returns error for HTTP discoveryUrl (HTTPS required)', () => {
261+
// AC13: At least one of audience/clients/scopes must be non-empty
262+
it('returns error when all of audience, clients, and scopes are empty', () => {
273263
const result = validateAddGatewayOptions({
274264
...validGatewayOptionsJwt,
275-
discoveryUrl: 'http://example.com/.well-known/openid-configuration',
265+
allowedAudience: ' ',
266+
allowedClients: undefined,
267+
allowedScopes: undefined,
276268
});
277269
expect(result.valid).toBe(false);
278-
expect(result.error).toBe('Discovery URL must use HTTPS');
270+
expect(result.error).toBe(
271+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer'
272+
);
279273
});
280274

281-
it('allows CUSTOM_JWT with only allowedClients', () => {
282-
const opts = { ...validGatewayOptionsJwt, allowedAudience: undefined, allowedScopes: undefined };
283-
const result = validateAddGatewayOptions(opts);
275+
// AC-claims1: --custom-claims with valid JSON passes validation
276+
it('accepts valid --custom-claims JSON', () => {
277+
const result = validateAddGatewayOptions({
278+
...validGatewayOptionsJwt,
279+
customClaims: JSON.stringify([
280+
{
281+
inboundTokenClaimName: 'dept',
282+
inboundTokenClaimValueType: 'STRING',
283+
authorizingClaimMatchValue: {
284+
claimMatchOperator: 'EQUALS',
285+
claimMatchValue: { matchValueString: 'engineering' },
286+
},
287+
},
288+
]),
289+
});
284290
expect(result.valid).toBe(true);
285291
});
286292

293+
// AC-claims2: --custom-claims alone satisfies the "at least one constraint" check
294+
it('allows CUSTOM_JWT with only --custom-claims (no audience/clients/scopes)', () => {
295+
const result = validateAddGatewayOptions({
296+
name: 'test-gw',
297+
authorizerType: 'CUSTOM_JWT',
298+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
299+
customClaims: JSON.stringify([
300+
{
301+
inboundTokenClaimName: 'role',
302+
inboundTokenClaimValueType: 'STRING_ARRAY',
303+
authorizingClaimMatchValue: {
304+
claimMatchOperator: 'CONTAINS_ANY',
305+
claimMatchValue: { matchValueStringList: ['admin'] },
306+
},
307+
},
308+
]),
309+
});
310+
expect(result.valid).toBe(true);
311+
});
312+
313+
// AC-claims3: --custom-claims with invalid JSON fails
314+
it('returns error for --custom-claims with invalid JSON', () => {
315+
const result = validateAddGatewayOptions({
316+
...validGatewayOptionsJwt,
317+
customClaims: 'not json',
318+
});
319+
expect(result.valid).toBe(false);
320+
expect(result.error).toBe('--custom-claims must be valid JSON');
321+
});
322+
323+
// AC-claims4: --custom-claims with empty array fails
324+
it('returns error for --custom-claims with empty array', () => {
325+
const result = validateAddGatewayOptions({
326+
...validGatewayOptionsJwt,
327+
customClaims: '[]',
328+
});
329+
expect(result.valid).toBe(false);
330+
expect(result.error).toBe('--custom-claims must be a non-empty JSON array');
331+
});
332+
333+
// AC-claims5: --custom-claims with invalid claim structure fails
334+
it('returns error for --custom-claims with invalid claim structure', () => {
335+
const result = validateAddGatewayOptions({
336+
...validGatewayOptionsJwt,
337+
customClaims: JSON.stringify([{ badField: 'value' }]),
338+
});
339+
expect(result.valid).toBe(false);
340+
expect(result.error).toContain('Invalid custom claim at index 0');
341+
});
342+
287343
// AC14: Valid options pass
288344
it('passes for valid options', () => {
289345
expect(validateAddGatewayOptions(validGatewayOptionsNone)).toEqual({ valid: true });
@@ -309,8 +365,8 @@ describe('validate', () => {
309365
expect(result.error).toBe('Both --client-id and --client-secret must be provided together');
310366
});
311367

312-
// AC16: OAuth credentials only valid with CUSTOM_JWT
313-
it('returns error when OAuth credentials used with non-CUSTOM_JWT authorizer', () => {
368+
// AC16: OAuth client credentials only valid with CUSTOM_JWT
369+
it('returns error when OAuth client credentials used with non-CUSTOM_JWT authorizer', () => {
314370
const result = validateAddGatewayOptions({
315371
...validGatewayOptionsNone,
316372
clientId: 'my-client-id',
@@ -320,8 +376,8 @@ describe('validate', () => {
320376
expect(result.error).toBe('OAuth client credentials are only valid with CUSTOM_JWT authorizer');
321377
});
322378

323-
// AC17: valid CUSTOM_JWT with OAuth credentials passes
324-
it('passes for CUSTOM_JWT with OAuth credentials', () => {
379+
// AC17: valid CUSTOM_JWT with OAuth client credentials passes
380+
it('passes for CUSTOM_JWT with OAuth client credentials', () => {
325381
const result = validateAddGatewayOptions({
326382
...validGatewayOptionsJwt,
327383
clientId: 'my-client-id',

‎src/cli/commands/add/types.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export interface AddGatewayOptions {
3737
allowedAudience?: string;
3838
allowedClients?: string;
3939
allowedScopes?: string;
40+
customClaims?: string;
4041
clientId?: string;
4142
clientSecret?: string;
4243
agents?: string;

‎src/cli/commands/add/validate.ts‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { ConfigIO, findConfigRoot } from '../../../lib';
22
import {
33
AgentNameSchema,
44
BuildTypeSchema,
5+
CustomClaimValidationSchema,
56
GatewayExceptionLevelSchema,
67
GatewayNameSchema,
78
ModelProviderSchema,
@@ -275,16 +276,36 @@ export function validateAddGatewayOptions(options: AddGatewayOptions): Validatio
275276
return { valid: false, error: `Discovery URL must end with ${OIDC_WELL_KNOWN_SUFFIX}` };
276277
}
277278

278-
// allowedAudience, allowedClients, allowedScopes are all optional individually,
279+
// Validate custom claims JSON if provided
280+
if (options.customClaims) {
281+
let parsed: unknown;
282+
try {
283+
parsed = JSON.parse(options.customClaims);
284+
} catch {
285+
return { valid: false, error: '--custom-claims must be valid JSON' };
286+
}
287+
if (!Array.isArray(parsed) || parsed.length === 0) {
288+
return { valid: false, error: '--custom-claims must be a non-empty JSON array' };
289+
}
290+
for (const [i, entry] of parsed.entries()) {
291+
const result = CustomClaimValidationSchema.safeParse(entry);
292+
if (!result.success) {
293+
return { valid: false, error: `Invalid custom claim at index ${i}: ${result.error.issues[0]?.message}` };
294+
}
295+
}
296+
}
297+
298+
// allowedAudience, allowedClients, allowedScopes, customClaims are all optional individually,
279299
// but at least one must be provided
280300
const hasAudience = !!options.allowedAudience?.trim();
281301
const hasClients = !!options.allowedClients?.trim();
282302
const hasScopes = !!options.allowedScopes?.trim();
283-
if (!hasAudience && !hasClients && !hasScopes) {
303+
const hasClaims = !!options.customClaims?.trim();
304+
if (!hasAudience && !hasClients && !hasScopes && !hasClaims) {
284305
return {
285306
valid: false,
286307
error:
287-
'At least one of --allowed-audience, --allowed-clients, or --allowed-scopes must be provided for CUSTOM_JWT authorizer',
308+
'At least one of --allowed-audience, --allowed-clients, --allowed-scopes, or --custom-claims must be provided for CUSTOM_JWT authorizer',
288309
};
289310
}
290311
}

‎src/cli/primitives/GatewayPrimitive.ts‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import type {
44
AgentCoreGatewayTarget,
55
AgentCoreMcpSpec,
66
AgentCoreProjectSpec,
7+
CustomClaimValidation,
78
GatewayAuthorizerType,
89
} from '../../schema';
910
import { AgentCoreGatewaySchema, PolicyEngineModeSchema } from '../../schema';
@@ -29,6 +30,7 @@ export interface AddGatewayOptions {
2930
allowedAudience?: string;
3031
allowedClients?: string;
3132
allowedScopes?: string;
33+
customClaims?: CustomClaimValidation[];
3234
clientId?: string;
3335
clientSecret?: string;
3436
agents?: string;
@@ -164,6 +166,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
164166
.option('--allowed-audience <audience>', 'Comma-separated allowed audiences (for CUSTOM_JWT)')
165167
.option('--allowed-clients <clients>', 'Comma-separated allowed client IDs (for CUSTOM_JWT)')
166168
.option('--allowed-scopes <scopes>', 'Comma-separated allowed scopes (for CUSTOM_JWT)')
169+
.option('--custom-claims <json>', 'Custom claim validations as JSON array (for CUSTOM_JWT)')
167170
.option('--client-id <id>', 'OAuth client ID for gateway bearer token')
168171
.option('--client-secret <secret>', 'OAuth client secret')
169172
.option('--agents <agents>', 'Comma-separated agent names')
@@ -190,6 +193,11 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
190193
process.exit(1);
191194
}
192195

196+
// Parse custom claims JSON if provided (already validated)
197+
const parsedCustomClaims = cliOptions.customClaims
198+
? (JSON.parse(cliOptions.customClaims) as CustomClaimValidation[])
199+
: undefined;
200+
193201
const result = await this.add({
194202
name: cliOptions.name!,
195203
description: cliOptions.description,
@@ -198,6 +206,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
198206
allowedAudience: cliOptions.allowedAudience,
199207
allowedClients: cliOptions.allowedClients,
200208
allowedScopes: cliOptions.allowedScopes,
209+
customClaims: parsedCustomClaims,
201210
clientId: cliOptions.clientId,
202211
clientSecret: cliOptions.clientSecret,
203212
agents: cliOptions.agents,
@@ -334,6 +343,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
334343
...(allowedAudience?.length ? { allowedAudience } : {}),
335344
...(allowedClients?.length ? { allowedClients } : {}),
336345
...(allowedScopes?.length ? { allowedScopes } : {}),
346+
...(options.customClaims?.length ? { customClaims: options.customClaims } : {}),
337347
...(options.clientId ? { clientId: options.clientId } : {}),
338348
...(options.clientSecret ? { clientSecret: options.clientSecret } : {}),
339349
};
@@ -415,9 +425,10 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
415425
});
416426
await this.writeProjectSpec(project);
417427

418-
// Write client secret to .env
419-
const envVarName = computeDefaultCredentialEnvVarName(credentialName);
420-
await setEnvVar(envVarName, jwtConfig.clientSecret!);
428+
// Write client ID and client secret to .env
429+
const envVarPrefix = computeDefaultCredentialEnvVarName(credentialName);
430+
await setEnvVar(`${envVarPrefix}_CLIENT_ID`, jwtConfig.clientId!);
431+
await setEnvVar(`${envVarPrefix}_CLIENT_SECRET`, jwtConfig.clientSecret!);
421432
}
422433

423434
/**
@@ -434,6 +445,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
434445
...(config.jwtConfig.allowedAudience?.length ? { allowedAudience: config.jwtConfig.allowedAudience } : {}),
435446
...(config.jwtConfig.allowedClients?.length ? { allowedClients: config.jwtConfig.allowedClients } : {}),
436447
...(config.jwtConfig.allowedScopes?.length ? { allowedScopes: config.jwtConfig.allowedScopes } : {}),
448+
...(config.jwtConfig.customClaims?.length ? { customClaims: config.jwtConfig.customClaims } : {}),
437449
},
438450
};
439451
}

‎src/cli/primitives/__tests__/GatewayPrimitive.test.ts‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,78 @@ describe('GatewayPrimitive', () => {
5151
primitive = new GatewayPrimitive();
5252
});
5353

54+
describe('customClaims pipeline', () => {
55+
const SAMPLE_CLAIMS = [
56+
{
57+
inboundTokenClaimName: 'department',
58+
inboundTokenClaimValueType: 'STRING_ARRAY' as const,
59+
authorizingClaimMatchValue: {
60+
claimMatchOperator: 'CONTAINS_ANY' as const,
61+
claimMatchValue: { matchValueStringList: ['engineering', 'sales'] },
62+
},
63+
},
64+
];
65+
66+
it('custom claims from TUI flow are written to authorizerConfiguration', async () => {
67+
await primitive.add({
68+
name: 'jwt-gw',
69+
authorizerType: 'CUSTOM_JWT',
70+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
71+
allowedAudience: 'aud1',
72+
customClaims: SAMPLE_CLAIMS,
73+
});
74+
75+
const gw = getWrittenGateway();
76+
expect(gw.authorizerConfiguration?.customJwtAuthorizer).toBeDefined();
77+
expect(gw.authorizerConfiguration!.customJwtAuthorizer!.customClaims).toEqual(SAMPLE_CLAIMS);
78+
});
79+
80+
it('custom claims are preserved alongside audience and clients', async () => {
81+
await primitive.add({
82+
name: 'jwt-gw',
83+
authorizerType: 'CUSTOM_JWT',
84+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
85+
allowedAudience: 'aud1,aud2',
86+
allowedClients: 'client1',
87+
customClaims: SAMPLE_CLAIMS,
88+
});
89+
90+
const gw = getWrittenGateway();
91+
const jwtConfig = gw.authorizerConfiguration!.customJwtAuthorizer!;
92+
expect(jwtConfig.allowedAudience).toEqual(['aud1', 'aud2']);
93+
expect(jwtConfig.allowedClients).toEqual(['client1']);
94+
expect(jwtConfig.customClaims).toEqual(SAMPLE_CLAIMS);
95+
});
96+
97+
it('omits customClaims from authorizerConfiguration when not provided', async () => {
98+
await primitive.add({
99+
name: 'jwt-gw',
100+
authorizerType: 'CUSTOM_JWT',
101+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
102+
allowedAudience: 'aud1',
103+
});
104+
105+
const gw = getWrittenGateway();
106+
expect(gw.authorizerConfiguration!.customJwtAuthorizer!.customClaims).toBeUndefined();
107+
});
108+
109+
it('custom claims only (no audience/clients/scopes) produces valid config', async () => {
110+
await primitive.add({
111+
name: 'jwt-gw',
112+
authorizerType: 'CUSTOM_JWT',
113+
discoveryUrl: 'https://example.com/.well-known/openid-configuration',
114+
customClaims: SAMPLE_CLAIMS,
115+
});
116+
117+
const gw = getWrittenGateway();
118+
const jwtConfig = gw.authorizerConfiguration!.customJwtAuthorizer!;
119+
expect(jwtConfig.allowedAudience).toBeUndefined();
120+
expect(jwtConfig.allowedClients).toBeUndefined();
121+
expect(jwtConfig.allowedScopes).toBeUndefined();
122+
expect(jwtConfig.customClaims).toEqual(SAMPLE_CLAIMS);
123+
});
124+
});
125+
54126
describe('exceptionLevel', () => {
55127
it('defaults to exceptionLevel NONE', async () => {
56128
await primitive.add({ name: 'test-gw', authorizerType: 'NONE' });

‎src/cli/tui/hooks/useCreateMcp.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ export function useCreateGateway() {
3131
allowedAudience: config.jwtConfig?.allowedAudience?.join(','),
3232
allowedClients: config.jwtConfig?.allowedClients?.join(','),
3333
allowedScopes: config.jwtConfig?.allowedScopes?.join(','),
34+
customClaims: config.jwtConfig?.customClaims,
3435
clientId: config.jwtConfig?.clientId,
3536
clientSecret: config.jwtConfig?.clientSecret,
3637
enableSemanticSearch: config.enableSemanticSearch,

0 commit comments

Comments
 (0)