Skip to content

Commit 0b743db

Browse files
feat: runtime inbound auth (Custom JWT) for agents (#657)
* test: add integ tests for evaluator and online-eval resource lifecycle * test: remove redundant levels/rating-scales section * refactor: type readProjectConfig return as AgentCoreProjectSpec * refactor(schema): extract shared auth schemas to auth.ts Move JWT authorizer schemas (GatewayAuthorizerTypeSchema, CustomJwtAuthorizerConfigSchema, etc.) from mcp.ts into a dedicated auth.ts module. Add RuntimeAuthorizerTypeSchema (AWS_IAM | CUSTOM_JWT) for Runtime resources. The resource-agnostic AuthorizerConfigSchema replaces GatewayAuthorizerConfigSchema (kept as deprecated alias). * refactor(tui): extract JWT config components to shared module Move JWT authorizer TUI components from AddGatewayScreen into src/cli/tui/components/jwt-config/: - types.ts: shared types and constants - CustomClaimForm.tsx: single claim tab-field form - CustomClaimsManager.tsx: CRUD for custom claims - JwtConfigInput.tsx: main JWT config component - useJwtConfigFlow.ts: hook for JWT wizard state management AddGatewayScreen now uses the shared useJwtConfigFlow hook and JwtConfigInput component. * refactor(primitives): extract shared auth utilities from GatewayPrimitive Extract buildAuthorizerConfigFromJwtConfig and createManagedOAuthCredential into auth-utils.ts so they can be reused by AgentPrimitive for Runtime inbound auth. GatewayPrimitive now imports these shared functions instead of defining them as private methods. * refactor(validation): extract shared JWT authorizer validation Extract JWT validation logic from validateAddGatewayOptions into a reusable validateJwtAuthorizerOptions function in auth-options.ts. This enables reuse by the agent command for Runtime inbound auth. * feat(schema): add authorizerType and authorizerConfiguration to AgentEnvSpec Add Runtime inbound auth fields to the agent schema. Supports AWS_IAM (default) and CUSTOM_JWT authorizer types with cross-field validation ensuring authorizerConfiguration is required for CUSTOM_JWT. * feat(tui): add inbound auth step to BYO agent wizard Add authorizerType selection and JWT config flow to the BYO agent creation wizard. Users can choose between AWS_IAM (default) and CUSTOM_JWT with full OIDC configuration including discovery URL, audience/client/scope constraints, and custom claims. * feat(cli): wire Runtime inbound auth to AgentPrimitive and CLI flags Add --authorizer-type, --discovery-url, --allowed-audience, --allowed-clients, --allowed-scopes, --custom-claims, --client-id, and --client-secret CLI flags to `agentcore add agent`. Wire auth options through to the BYO path for authorizer configuration building and OAuth credential auto-creation. Add validation for agent auth options reusing shared validateJwtAuthorizerOptions. * fix(tui): add missing OAuth credential creation in BYO agent path The TUI BYO path silently discarded OAuth client credentials when CUSTOM_JWT was selected. Now calls createManagedOAuthCredential to persist the credential and write secrets to .env, matching the CLI path. * test(integ): add integration tests for agent inbound auth CLI flags Tests adding BYO agents with CUSTOM_JWT authorizer via CLI, including: - Audience constraints, client/scope constraints, custom claims - OAuth credential auto-creation with .env secrets - Default AWS_IAM behavior - Validation of missing discovery URL, missing constraints, misplaced client credentials * feat(tui): add inbound auth to all agent paths (create, BYO, import) Add authorizerType and jwtConfig steps to the GenerateWizard (create/template path) and import path, matching existing BYO path support. Auth config is persisted to agentcore.json via schema-mapper. OAuth credential is auto-created for CUSTOM_JWT in all three paths. Also fix request header allowlist TextInput to allow empty submission (skip), and add E2E test for BYO custom JWT auth flow with real Cognito. * fix: wire auth config through useCreateFlow (agentcore create path) The create command's useCreateFlow.ts duplicates agent creation logic from useAddAgent.ts but was missing authorizerType/jwtConfig in the GenerateConfig, import params, and OAuth credential creation for all three sub-paths (create, import, BYO). * feat(invoke): add bearer token auth for CUSTOM_JWT agents Add thin HTTP invoke client that bypasses SigV4 when a bearer token is provided, supporting agents configured with CUSTOM_JWT authorizers. - Add --bearer-token CLI flag to invoke command - Add bearer token input to invoke TUI (auto-prompts for CUSTOM_JWT agents, press T to set/change token in chat mode) - Build invokeWithBearerToken and invokeWithBearerTokenStreaming HTTP clients with SSE parsing matching the SDK streaming pattern - Wire bearerToken through InvokeOptions, useInvokeFlow, and InvokeScreen - Expose authorizerType in invoke flow config for agent auth detection * revert: remove allowEmpty from requestHeaderAllowlist TextInputs Being handled separately by another agent. * feat(invoke): auto-fetch OAuth token for CUSTOM_JWT agents Extract shared OAuth client_credentials flow from fetchGatewayToken into oauth-token.ts (fetchOAuthToken), then add fetchRuntimeToken for agents. TUI: auto-fetches token on agent selection for CUSTOM_JWT agents. Shows fetch status in header (fetching/fetched/error). Press R to refresh, T for manual entry. Falls back to manual input on fetch failure. CLI: auto-fetches token when --bearer-token is not provided and agent has CUSTOM_JWT auth. Surfaces clear error with manual fallback hint. * fix(invoke): silently skip token auto-fetch when no OAuth credential exists Add canFetchRuntimeToken pre-check that verifies the managed OAuth credential and client secret exist before attempting the fetch. When credentials aren't configured, the TUI shows the manual token prompt (T key) without an error, and the CLI proceeds without a token. * fix(invoke): use Ctrl+T and Ctrl+R for token shortcuts Plain T/R conflict with typing those characters in chat mode. * refactor(invoke): replace key shortcuts with dedicated token screen Instead of Ctrl+T/Ctrl+R shortcuts in chat mode (which conflict with terminal input), show a dedicated bearer token screen before chat when a CUSTOM_JWT agent is selected. The screen auto-populates with a fetched token if OAuth credentials are configured, otherwise shows an empty field for the user to paste into. Enter confirms, Esc skips. * feat(fetch): extend fetch access to support agents Add --type flag to agentcore fetch access (gateway | agent). When --type agent is used, performs OAuth client_credentials flow for the agent's CUSTOM_JWT credential instead of a gateway. Usage: agentcore fetch access --type agent --name MyAgent [--json] * feat(tui): extend fetch access screen to support both gateways and agents The TUI fetch access screen previously only listed gateways. Now it loads both gateways and agents, shows the resource type in the picker, and dispatches to the correct token fetcher based on resource type. * fix(tui): handle agent fetch access gracefully without OAuth credentials For agents, use canFetchRuntimeToken pre-check before attempting token fetch. AWS_IAM agents show auth guidance without fetching. CUSTOM_JWT agents without managed credentials show a message instead of erroring. * fix(tui): gate auth screen behind advanced config and fix lint errors Auth type selection (CUSTOM_JWT vs AWS_IAM) now only appears when the user selects advanced config in the generate wizard and BYO agent paths. Skipping advanced goes directly to confirm with default AWS_IAM auth. Also fixes 4 lint errors: unused import, optional chain preferences, and unsafe argument type cast. * fix(tui): fetch agent token directly instead of pre-check in fetch access The canFetchRuntimeToken pre-check swallows all errors and returns false, which incorrectly showed "no managed OAuth credential" even when one was configured. For the fetch access screen (explicit user action), call fetchRuntimeToken directly and let real errors surface in the error phase. The pre-check remains in the invoke flow where silent skip is desired. * fix(fetch): use _CLIENT_SECRET suffix when reading OAuth secret from env The write path stores secrets as AGENTCORE_CREDENTIAL_{NAME}_CLIENT_SECRET but the read path was looking for AGENTCORE_CREDENTIAL_{NAME} without the suffix, causing token fetch to always fail with "Client secret not found". * fix(ci): add @aws-sdk/client-cognito-identity-provider dev dependency The e2e test for BYO custom JWT imports this SDK but it wasn't in package.json, causing typecheck to fail in CI. * style: fix prettier formatting across all changed files * fix: wire auth config through handleCreatePath for template agents handleCreatePath was silently dropping auth options (authorizerType, discoveryUrl, etc.) because generateConfig didn't include them. This meant `agentcore add agent --type create --authorizer-type CUSTOM_JWT` would succeed but write agentcore.json without auth fields. Fixes: pass authorizerType and jwtConfig into generateConfig so mapGenerateConfigToAgent picks them up. Also fixes misleading comment in validate.ts that said auth validation was BYO-only. --------- Co-authored-by: Harrison Weinstock <hkobew@amazon.com>
1 parent f24609f commit 0b743db

54 files changed

Lines changed: 3541 additions & 1379 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

e2e-tests/byo-custom-jwt.test.ts

Lines changed: 271 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,271 @@
1+
/**
2+
* E2E test: BYO agent with CUSTOM_JWT inbound auth (Cognito).
3+
*
4+
* Creates a Cognito user pool as the OIDC provider, deploys a BYO agent
5+
* configured with CUSTOM_JWT authorizer, and verifies that:
6+
* - Deploy succeeds with AuthorizerConfiguration in the CloudFormation template
7+
* - SigV4 invocation is rejected (auth method mismatch)
8+
* - Status reports the agent as deployed
9+
*
10+
* Unlike other E2E tests that use the globally installed CLI, this test uses
11+
* the local build (`runCLI`) because it exercises unreleased schema and CDK
12+
* changes. Set CDK_TARBALL to a path to the modified CDK package tarball.
13+
*
14+
* Requires: AWS credentials, npm, git, uv, CDK_TARBALL env var.
15+
*/
16+
import {
17+
type RunResult,
18+
hasAwsCredentials,
19+
parseJsonOutput,
20+
prereqs,
21+
runCLI,
22+
stripAnsi,
23+
} from '../src/test-utils/index.js';
24+
import { CloudFormationClient, GetTemplateCommand } from '@aws-sdk/client-cloudformation';
25+
import {
26+
CognitoIdentityProviderClient,
27+
CreateResourceServerCommand,
28+
CreateUserPoolClientCommand,
29+
CreateUserPoolCommand,
30+
CreateUserPoolDomainCommand,
31+
DeleteResourceServerCommand,
32+
DeleteUserPoolCommand,
33+
DeleteUserPoolDomainCommand,
34+
} from '@aws-sdk/client-cognito-identity-provider';
35+
import { execSync } from 'node:child_process';
36+
import { randomUUID } from 'node:crypto';
37+
import { mkdir, readFile, rm, writeFile } from 'node:fs/promises';
38+
import { tmpdir } from 'node:os';
39+
import { join } from 'node:path';
40+
import { afterAll, beforeAll, describe, expect, it } from 'vitest';
41+
42+
const hasAws = hasAwsCredentials();
43+
const hasCdkTarball = !!process.env.CDK_TARBALL;
44+
const canRun = prereqs.npm && prereqs.git && prereqs.uv && hasAws && hasCdkTarball;
45+
const region = process.env.AWS_REGION ?? 'us-east-1';
46+
47+
/**
48+
* Run the local CLI build without skipping install (needed for deploy).
49+
*/
50+
function runLocalCLI(args: string[], cwd: string): Promise<RunResult> {
51+
return runCLI(args, cwd, /* skipInstall */ false);
52+
}
53+
54+
describe.sequential('e2e: BYO agent with CUSTOM_JWT auth', () => {
55+
let testDir: string;
56+
let projectPath: string;
57+
let agentName: string;
58+
59+
// Cognito resources
60+
let userPoolId: string;
61+
let clientId: string;
62+
let domainPrefix: string;
63+
let discoveryUrl: string;
64+
65+
const cognitoClient = new CognitoIdentityProviderClient({ region });
66+
const cfnClient = new CloudFormationClient({ region });
67+
68+
beforeAll(async () => {
69+
if (!canRun) return;
70+
71+
// ── Create Cognito user pool as OIDC provider ──
72+
const suffix = randomUUID().slice(0, 8);
73+
const poolName = `agentcore-e2e-jwt-${suffix}`;
74+
domainPrefix = `agentcore-e2e-jwt-${suffix}`;
75+
76+
const poolResult = await cognitoClient.send(new CreateUserPoolCommand({ PoolName: poolName }));
77+
userPoolId = poolResult.UserPool!.Id!;
78+
79+
await cognitoClient.send(new CreateUserPoolDomainCommand({ UserPoolId: userPoolId, Domain: domainPrefix }));
80+
81+
await cognitoClient.send(
82+
new CreateResourceServerCommand({
83+
UserPoolId: userPoolId,
84+
Identifier: 'agentcore',
85+
Name: 'AgentCore API',
86+
Scopes: [{ ScopeName: 'invoke', ScopeDescription: 'Invoke the runtime' }],
87+
})
88+
);
89+
90+
const clientResult = await cognitoClient.send(
91+
new CreateUserPoolClientCommand({
92+
UserPoolId: userPoolId,
93+
ClientName: 'e2e-test-client',
94+
GenerateSecret: true,
95+
AllowedOAuthFlows: ['client_credentials'],
96+
AllowedOAuthScopes: ['agentcore/invoke'],
97+
AllowedOAuthFlowsUserPoolClient: true,
98+
ExplicitAuthFlows: ['ALLOW_REFRESH_TOKEN_AUTH'],
99+
})
100+
);
101+
clientId = clientResult.UserPoolClient!.ClientId!;
102+
103+
discoveryUrl = `https://cognito-idp.${region}.amazonaws.com/${userPoolId}/.well-known/openid-configuration`;
104+
105+
// ── Create test project using local CLI build ──
106+
testDir = join(tmpdir(), `agentcore-e2e-jwt-${randomUUID()}`);
107+
await mkdir(testDir, { recursive: true });
108+
109+
agentName = `E2eJwt${String(Date.now()).slice(-8)}`;
110+
const createResult = await runLocalCLI(
111+
[
112+
'create',
113+
'--name',
114+
agentName,
115+
'--language',
116+
'Python',
117+
'--framework',
118+
'Strands',
119+
'--model-provider',
120+
'Bedrock',
121+
'--memory',
122+
'none',
123+
'--json',
124+
],
125+
testDir
126+
);
127+
expect(createResult.exitCode, `Create failed: ${createResult.stderr}`).toBe(0);
128+
const createJson = parseJsonOutput(createResult.stdout) as { projectPath: string };
129+
projectPath = createJson.projectPath;
130+
131+
// Write AWS targets
132+
const account =
133+
process.env.AWS_ACCOUNT_ID ??
134+
execSync('aws sts get-caller-identity --query Account --output text').toString().trim();
135+
await writeFile(
136+
join(projectPath, 'agentcore', 'aws-targets.json'),
137+
JSON.stringify([{ name: 'default', account, region }])
138+
);
139+
140+
// Install modified CDK tarball (required for auth fields support)
141+
execSync(`npm install -f ${process.env.CDK_TARBALL}`, {
142+
cwd: join(projectPath, 'agentcore', 'cdk'),
143+
stdio: 'pipe',
144+
});
145+
146+
// ── Patch agent with CUSTOM_JWT auth ──
147+
const specPath = join(projectPath, 'agentcore', 'agentcore.json');
148+
const spec = JSON.parse(await readFile(specPath, 'utf8'));
149+
const agent = spec.agents[0];
150+
agent.authorizerType = 'CUSTOM_JWT';
151+
agent.authorizerConfiguration = {
152+
customJwtAuthorizer: {
153+
discoveryUrl,
154+
allowedAudience: [clientId],
155+
},
156+
};
157+
await writeFile(specPath, JSON.stringify(spec, null, 2));
158+
}, 300000);
159+
160+
afterAll(async () => {
161+
if (!canRun) return;
162+
163+
// ── Tear down deployed stack ──
164+
if (projectPath) {
165+
try {
166+
await runLocalCLI(['remove', 'all', '--json'], projectPath);
167+
await runLocalCLI(['deploy', '--yes', '--json'], projectPath);
168+
} catch {
169+
// Best-effort cleanup
170+
}
171+
}
172+
173+
// ── Delete Cognito resources ──
174+
if (userPoolId) {
175+
try {
176+
await cognitoClient.send(new DeleteResourceServerCommand({ UserPoolId: userPoolId, Identifier: 'agentcore' }));
177+
} catch {
178+
/* best-effort */
179+
}
180+
try {
181+
await cognitoClient.send(new DeleteUserPoolDomainCommand({ UserPoolId: userPoolId, Domain: domainPrefix }));
182+
} catch {
183+
/* best-effort */
184+
}
185+
try {
186+
await cognitoClient.send(new DeleteUserPoolCommand({ UserPoolId: userPoolId }));
187+
} catch {
188+
/* best-effort */
189+
}
190+
}
191+
192+
// ── Clean up temp directory ──
193+
if (testDir) {
194+
await rm(testDir, { recursive: true, force: true, maxRetries: 3, retryDelay: 1000 });
195+
}
196+
}, 600000);
197+
198+
it.skipIf(!canRun)(
199+
'deploys with CUSTOM_JWT authorizer configuration',
200+
async () => {
201+
expect(projectPath, 'Project should have been created').toBeTruthy();
202+
203+
const result = await runLocalCLI(['deploy', '--yes', '--json'], projectPath);
204+
205+
if (result.exitCode !== 0) {
206+
console.log('Deploy stdout:', result.stdout);
207+
console.log('Deploy stderr:', result.stderr);
208+
}
209+
210+
expect(result.exitCode, `Deploy failed: ${result.stderr}`).toBe(0);
211+
212+
const json = parseJsonOutput(result.stdout) as { success: boolean; stackName: string };
213+
expect(json.success, 'Deploy should report success').toBe(true);
214+
215+
// Verify CloudFormation template contains AuthorizerConfiguration
216+
const templateResult = await cfnClient.send(new GetTemplateCommand({ StackName: json.stackName }));
217+
const template = JSON.parse(templateResult.TemplateBody!) as {
218+
Resources: Record<string, { Type: string; Properties: Record<string, unknown> }>;
219+
};
220+
221+
const runtimeResource = Object.values(template.Resources).find(r => r.Type === 'AWS::BedrockAgentCore::Runtime');
222+
expect(runtimeResource, 'Template should contain a Runtime resource').toBeDefined();
223+
224+
const props = runtimeResource!.Properties;
225+
const authConfig = props.AuthorizerConfiguration as {
226+
CustomJWTAuthorizer: { DiscoveryUrl: string; AllowedAudience: string[] };
227+
};
228+
expect(authConfig, 'Runtime should have AuthorizerConfiguration').toBeDefined();
229+
expect(authConfig.CustomJWTAuthorizer.DiscoveryUrl).toBe(discoveryUrl);
230+
expect(authConfig.CustomJWTAuthorizer.AllowedAudience).toContain(clientId);
231+
},
232+
600000
233+
);
234+
235+
it.skipIf(!canRun)(
236+
'rejects SigV4 invocation (auth method mismatch)',
237+
async () => {
238+
expect(projectPath, 'Project should have been deployed').toBeTruthy();
239+
240+
// The CLI uses SigV4 by default — a CUSTOM_JWT runtime should reject it
241+
const result = await runLocalCLI(
242+
['invoke', '--prompt', 'Say hello', '--agent', agentName, '--json'],
243+
projectPath
244+
);
245+
246+
// Expect failure due to auth method mismatch
247+
const output = stripAnsi(result.stdout + result.stderr);
248+
expect(output).toMatch(/[Aa]uthoriz(ation|er).*mismatch|different.*authorization/i);
249+
},
250+
180000
251+
);
252+
253+
it.skipIf(!canRun)(
254+
'status shows the deployed agent',
255+
async () => {
256+
const result = await runLocalCLI(['status', '--json'], projectPath);
257+
expect(result.exitCode, `Status failed: ${result.stderr}`).toBe(0);
258+
259+
const json = parseJsonOutput(result.stdout) as {
260+
success: boolean;
261+
resources: { resourceType: string; name: string; deploymentState: string }[];
262+
};
263+
expect(json.success).toBe(true);
264+
265+
const agent = json.resources.find(r => r.resourceType === 'agent' && r.name === agentName);
266+
expect(agent, `Agent "${agentName}" should appear in status`).toBeDefined();
267+
expect(agent!.deploymentState).toBe('deployed');
268+
},
269+
120000
270+
);
271+
});

0 commit comments

Comments
 (0)