From 83048e72e911863af381698facad972df0d643b1 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:12:24 -0500 Subject: [PATCH 1/9] feat: extend JWT wizard with allowedScopes and agent OAuth credential inputs --- src/cli/tui/screens/mcp/AddGatewayScreen.tsx | 91 ++++++++++++++++--- src/cli/tui/screens/mcp/types.ts | 3 + .../tui/screens/mcp/useAddGatewayWizard.ts | 9 +- 3 files changed, 89 insertions(+), 14 deletions(-) diff --git a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx index 13269eef..d10c6202 100644 --- a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx +++ b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx @@ -4,6 +4,7 @@ import { ConfirmReview, Panel, Screen, + SecretInput, StepIndicator, TextInput, WizardMultiSelect, @@ -29,10 +30,13 @@ interface AddGatewayScreenProps { export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassignedTargets }: AddGatewayScreenProps) { const wizard = useAddGatewayWizard(unassignedTargets.length); - // JWT config sub-step tracking (0 = discoveryUrl, 1 = audience, 2 = clients) + // JWT config sub-step tracking (0=discoveryUrl, 1=audience, 2=clients, 3=scopes, 4=agentClientId, 5=agentClientSecret) const [jwtSubStep, setJwtSubStep] = useState(0); const [jwtDiscoveryUrl, setJwtDiscoveryUrl] = useState(''); const [jwtAudience, setJwtAudience] = useState(''); + const [jwtClients, setJwtClients] = useState(''); + const [jwtScopes, setJwtScopes] = useState(''); + const [jwtAgentClientId, setJwtAgentClientId] = useState(''); const unassignedTargetItems: SelectableItem[] = useMemo( () => unassignedTargets.map(name => ({ id: name, title: name })), @@ -85,23 +89,33 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig }; const handleJwtClients = (clients: string) => { - // Parse comma-separated values - const audienceList = jwtAudience - .split(',') - .map(s => s.trim()) - .filter(Boolean); - const clientsList = clients - .split(',') - .map(s => s.trim()) - .filter(Boolean); + setJwtClients(clients); + setJwtSubStep(3); + }; + + const handleJwtScopes = (scopes: string) => { + setJwtScopes(scopes); + setJwtSubStep(4); + }; + + const handleJwtAgentClientId = (clientId: string) => { + setJwtAgentClientId(clientId); + setJwtSubStep(5); + }; + + const handleJwtAgentClientSecret = (clientSecret: string) => { + const audienceList = jwtAudience.split(',').map(s => s.trim()).filter(Boolean); + const clientsList = jwtClients.split(',').map(s => s.trim()).filter(Boolean); + const scopesList = jwtScopes.split(',').map(s => s.trim()).filter(Boolean); wizard.setJwtConfig({ discoveryUrl: jwtDiscoveryUrl, allowedAudience: audienceList, allowedClients: clientsList, + ...(scopesList.length > 0 ? { allowedScopes: scopesList } : {}), + ...(jwtAgentClientId ? { agentClientId: jwtAgentClientId, agentClientSecret: clientSecret } : {}), }); - // Reset sub-step counter only - preserve values for potential back navigation setJwtSubStep(0); }; @@ -160,6 +174,9 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig onDiscoveryUrl={handleJwtDiscoveryUrl} onAudience={handleJwtAudience} onClients={handleJwtClients} + onScopes={handleJwtScopes} + onAgentClientId={handleJwtAgentClientId} + onAgentClientSecret={handleJwtAgentClientSecret} onCancel={handleJwtCancel} /> )} @@ -187,6 +204,12 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig { label: 'Discovery URL', value: wizard.config.jwtConfig.discoveryUrl }, { label: 'Allowed Audience', value: wizard.config.jwtConfig.allowedAudience.join(', ') }, { label: 'Allowed Clients', value: wizard.config.jwtConfig.allowedClients.join(', ') }, + ...(wizard.config.jwtConfig.allowedScopes?.length + ? [{ label: 'Allowed Scopes', value: wizard.config.jwtConfig.allowedScopes.join(', ') }] + : []), + ...(wizard.config.jwtConfig.agentClientId + ? [{ label: 'Agent Credential', value: `${wizard.config.name}-agent-oauth` }] + : []), ] : []), { @@ -209,6 +232,9 @@ interface JwtConfigInputProps { onDiscoveryUrl: (url: string) => void; onAudience: (audience: string) => void; onClients: (clients: string) => void; + onScopes: (scopes: string) => void; + onAgentClientId: (clientId: string) => void; + onAgentClientSecret: (clientSecret: string) => void; onCancel: () => void; } @@ -227,11 +253,21 @@ function validateCommaSeparatedList(value: string, fieldName: string): true | st return true; } -function JwtConfigInput({ subStep, onDiscoveryUrl, onAudience, onClients, onCancel }: JwtConfigInputProps) { +function JwtConfigInput({ + subStep, + onDiscoveryUrl, + onAudience, + onClients, + onScopes, + onAgentClientId, + onAgentClientSecret, + onCancel, +}: JwtConfigInputProps) { + const totalSteps = 6; return ( Configure Custom JWT Authorizer - Step {subStep + 1} of 3 + Step {subStep + 1} of {totalSteps} {subStep === 0 && ( validateCommaSeparatedList(value, 'client')} /> )} + {subStep === 3 && ( + + )} + {subStep === 4 && ( + + )} + {subStep === 5 && ( + value.trim().length > 0 || 'Client secret is required'} + revealChars={4} + /> + )} ); diff --git a/src/cli/tui/screens/mcp/types.ts b/src/cli/tui/screens/mcp/types.ts index fcf7d593..f24aeed5 100644 --- a/src/cli/tui/screens/mcp/types.ts +++ b/src/cli/tui/screens/mcp/types.ts @@ -16,6 +16,9 @@ export interface AddGatewayConfig { discoveryUrl: string; allowedAudience: string[]; allowedClients: string[]; + allowedScopes?: string[]; + agentClientId?: string; + agentClientSecret?: string; }; /** Selected unassigned targets to include in this gateway */ selectedTargets?: string[]; diff --git a/src/cli/tui/screens/mcp/useAddGatewayWizard.ts b/src/cli/tui/screens/mcp/useAddGatewayWizard.ts index 2bd24b75..90265bca 100644 --- a/src/cli/tui/screens/mcp/useAddGatewayWizard.ts +++ b/src/cli/tui/screens/mcp/useAddGatewayWizard.ts @@ -68,7 +68,14 @@ export function useAddGatewayWizard(unassignedTargetsCount = 0) { }, []); const setJwtConfig = useCallback( - (jwtConfig: { discoveryUrl: string; allowedAudience: string[]; allowedClients: string[] }) => { + (jwtConfig: { + discoveryUrl: string; + allowedAudience: string[]; + allowedClients: string[]; + allowedScopes?: string[]; + agentClientId?: string; + agentClientSecret?: string; + }) => { setConfig(c => ({ ...c, jwtConfig, From 5bf8feb0938d3994fec80e8fd0b3019aeaab71bc Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:18:03 -0500 Subject: [PATCH 2/9] feat: auto-create managed OAuth credential for CUSTOM_JWT gateway --- src/cli/operations/mcp/create-mcp.ts | 26 ++++++++++++++++++++++++- src/schema/schemas/agentcore-project.ts | 2 ++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/cli/operations/mcp/create-mcp.ts b/src/cli/operations/mcp/create-mcp.ts index 1f554642..6e366508 100644 --- a/src/cli/operations/mcp/create-mcp.ts +++ b/src/cli/operations/mcp/create-mcp.ts @@ -1,4 +1,4 @@ -import { ConfigIO, requireConfigRoot } from '../../../lib'; +import { ConfigIO, requireConfigRoot, setEnvVar } from '../../../lib'; import type { AgentCoreCliMcpDefs, AgentCoreGateway, @@ -11,6 +11,7 @@ import { AgentCoreCliMcpDefsSchema, ToolDefinitionSchema } from '../../../schema import { getTemplateToolDefinitions, renderGatewayTargetTemplate } from '../../templates/GatewayTargetRenderer'; import type { AddGatewayConfig, AddGatewayTargetConfig } from '../../tui/screens/mcp/types'; import { DEFAULT_HANDLER, DEFAULT_NODE_VERSION, DEFAULT_PYTHON_VERSION } from '../../tui/screens/mcp/types'; +import { computeDefaultCredentialEnvVarName } from '../identity/create-identity'; import { existsSync } from 'fs'; import { mkdir, readFile, writeFile } from 'fs/promises'; import { dirname, join } from 'path'; @@ -71,6 +72,7 @@ function buildAuthorizerConfiguration(config: AddGatewayConfig): AgentCoreGatewa discoveryUrl: config.jwtConfig.discoveryUrl, allowedAudience: config.jwtConfig.allowedAudience, allowedClients: config.jwtConfig.allowedClients, + ...(config.jwtConfig.allowedScopes?.length && { allowedScopes: config.jwtConfig.allowedScopes }), }, }; } @@ -201,6 +203,28 @@ export async function createGatewayFromWizard(config: AddGatewayConfig): Promise await configIO.writeMcpSpec(mcpSpec); + // Auto-create managed credential if agent OAuth credentials provided + if (config.jwtConfig?.agentClientId && config.jwtConfig?.agentClientSecret) { + const credName = `${config.name}-agent-oauth`; + const project = await configIO.readProjectSpec(); + + const credential = { + type: 'OAuthCredentialProvider' as const, + name: credName, + discoveryUrl: config.jwtConfig.discoveryUrl, + vendor: 'CustomOauth2', + managed: true, + usage: 'inbound' as const, + }; + + project.credentials.push(credential); + await configIO.writeProjectSpec(project); + + const envBase = computeDefaultCredentialEnvVarName(credName); + await setEnvVar(`${envBase}_CLIENT_ID`, config.jwtConfig.agentClientId); + await setEnvVar(`${envBase}_CLIENT_SECRET`, config.jwtConfig.agentClientSecret); + } + return { name: config.name }; } diff --git a/src/schema/schemas/agentcore-project.ts b/src/schema/schemas/agentcore-project.ts index 13f8241f..fda34160 100644 --- a/src/schema/schemas/agentcore-project.ts +++ b/src/schema/schemas/agentcore-project.ts @@ -101,6 +101,8 @@ export const OAuthCredentialSchema = z.object({ vendor: z.string().default('CustomOauth2'), /** Whether this credential was auto-created by the CLI (e.g., for CUSTOM_JWT inbound auth) */ managed: z.boolean().optional(), + /** Whether this credential is used for inbound or outbound auth */ + usage: z.enum(['inbound', 'outbound']).optional(), }); export type OAuthCredential = z.infer; From 9d6a94da2c860a09470963d4b267055a6418731c Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:21:13 -0500 Subject: [PATCH 3/9] feat: add CLI flags for CUSTOM_JWT agent OAuth credentials --- src/cli/commands/add/actions.ts | 11 +++++++++++ src/cli/commands/add/command.tsx | 6 ++++++ src/cli/commands/add/types.ts | 3 +++ src/cli/commands/add/validate.ts | 11 +++++++++++ 4 files changed, 31 insertions(+) diff --git a/src/cli/commands/add/actions.ts b/src/cli/commands/add/actions.ts index 675c52e8..7232f7c7 100644 --- a/src/cli/commands/add/actions.ts +++ b/src/cli/commands/add/actions.ts @@ -64,6 +64,9 @@ export interface ValidatedAddGatewayOptions { discoveryUrl?: string; allowedAudience?: string; allowedClients?: string; + allowedScopes?: string; + agentClientId?: string; + agentClientSecret?: string; agents?: string; } @@ -267,6 +270,14 @@ function buildGatewayConfig(options: ValidatedAddGatewayOptions): AddGatewayConf .allowedClients!.split(',') .map(s => s.trim()) .filter(Boolean), + allowedScopes: options.allowedScopes + ? options.allowedScopes + .split(',') + .map(s => s.trim()) + .filter(Boolean) + : undefined, + agentClientId: options.agentClientId, + agentClientSecret: options.agentClientSecret, }; } diff --git a/src/cli/commands/add/command.tsx b/src/cli/commands/add/command.tsx index 6a9370bb..22e89dc5 100644 --- a/src/cli/commands/add/command.tsx +++ b/src/cli/commands/add/command.tsx @@ -82,6 +82,9 @@ async function handleAddGatewayCLI(options: AddGatewayOptions): Promise { discoveryUrl: options.discoveryUrl, allowedAudience: options.allowedAudience, allowedClients: options.allowedClients, + allowedScopes: options.allowedScopes, + agentClientId: options.agentClientId, + agentClientSecret: options.agentClientSecret, agents: options.agents, }); @@ -272,6 +275,9 @@ export function registerAdd(program: Command) { .option('--discovery-url ', 'OIDC discovery URL (required for CUSTOM_JWT)') .option('--allowed-audience ', 'Comma-separated allowed audience values (required for CUSTOM_JWT)') .option('--allowed-clients ', 'Comma-separated allowed client IDs (required for CUSTOM_JWT)') + .option('--allowed-scopes ', 'Comma-separated allowed scopes (optional for CUSTOM_JWT)') + .option('--agent-client-id ', 'Agent OAuth client ID for Bearer token auth (CUSTOM_JWT)') + .option('--agent-client-secret ', 'Agent OAuth client secret (CUSTOM_JWT)') .option('--json', 'Output as JSON') .action(async options => { requireProject(); diff --git a/src/cli/commands/add/types.ts b/src/cli/commands/add/types.ts index c83db76d..46757121 100644 --- a/src/cli/commands/add/types.ts +++ b/src/cli/commands/add/types.ts @@ -31,6 +31,9 @@ export interface AddGatewayOptions { discoveryUrl?: string; allowedAudience?: string; allowedClients?: string; + allowedScopes?: string; + agentClientId?: string; + agentClientSecret?: string; agents?: string; json?: boolean; } diff --git a/src/cli/commands/add/validate.ts b/src/cli/commands/add/validate.ts index 9a4bc4df..0aac0a21 100644 --- a/src/cli/commands/add/validate.ts +++ b/src/cli/commands/add/validate.ts @@ -181,6 +181,17 @@ export function validateAddGatewayOptions(options: AddGatewayOptions): Validatio } } + // Validate agent OAuth credentials + if (options.agentClientId && !options.agentClientSecret) { + return { valid: false, error: 'Both --agent-client-id and --agent-client-secret must be provided together' }; + } + if (options.agentClientSecret && !options.agentClientId) { + return { valid: false, error: 'Both --agent-client-id and --agent-client-secret must be provided together' }; + } + if (options.agentClientId && options.authorizerType !== 'CUSTOM_JWT') { + return { valid: false, error: 'Agent OAuth credentials are only valid with CUSTOM_JWT authorizer' }; + } + return { valid: true }; } From fb1f7877be255ab54eb20803495973601f88a21a Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:40:38 -0500 Subject: [PATCH 4/9] feat: add CUSTOM_JWT Bearer token auth to agent templates (Strands, LangChain, OpenAI, Google ADK) --- .../assets.snapshot.test.ts.snap | 155 ++++++++++++++++++ .../googleadk/base/mcp_client/client.py | 39 +++++ .../base/mcp_client/client.py | 39 +++++ .../openaiagents/base/mcp_client/client.py | 39 +++++ .../python/strands/base/mcp_client/client.py | 38 +++++ .../agent/generate/schema-mapper.ts | 31 +++- src/cli/templates/types.ts | 6 + src/cli/tui/screens/mcp/AddGatewayScreen.tsx | 2 - 8 files changed, 342 insertions(+), 7 deletions(-) diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index d58f06b0..434821b4 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -1726,6 +1726,41 @@ logger = logging.getLogger(__name__) import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_toolsets() -> list[MCPToolset]: """Returns MCP Toolsets for all configured gateways.""" @@ -1740,6 +1775,10 @@ def get_all_gateway_mcp_toolsets() -> list[MCPToolset]: url=url, httpx_client_factory=lambda **kwargs: httpx.AsyncClient(auth=auth, **kwargs) ))) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else None + toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url, headers=headers))) {{else}} toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url))) {{/if}} @@ -2012,6 +2051,41 @@ logger = logging.getLogger(__name__) {{#if (includes gatewayAuthTypes "AWS_IAM")}} from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_client() -> MultiServerMCPClient | None: """Returns an MCP Client connected to all configured gateways.""" @@ -2023,6 +2097,10 @@ def get_all_gateway_mcp_client() -> MultiServerMCPClient | None: session = create_aws_session() auth = SigV4HTTPXAuth(session.get_credentials(), "bedrock-agentcore", session.region_name) servers["{{name}}"] = {"transport": "streamable_http", "url": url, "auth": auth} + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else None + servers["{{name}}"] = {"transport": "streamable_http", "url": url, "headers": headers} {{else}} servers["{{name}}"] = {"transport": "streamable_http", "url": url} {{/if}} @@ -2438,6 +2516,41 @@ logger = logging.getLogger(__name__) import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]: """Returns MCP servers for all configured gateways.""" @@ -2452,6 +2565,10 @@ def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]: name="{{name}}", params={"url": url, "httpx_client_factory": lambda **kwargs: httpx.AsyncClient(auth=auth, **kwargs)} )) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else {} + servers.append(MCPServerStreamableHttp(name="{{name}}", params={"url": url, "headers": headers})) {{else}} servers.append(MCPServerStreamableHttp(name="{{name}}", params={"url": url})) {{/if}} @@ -2749,7 +2866,41 @@ logger = logging.getLogger(__name__) {{#if (includes gatewayAuthTypes "AWS_IAM")}} from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] +{{/if}} +{{/each}} {{#each gatewayProviders}} def get_{{snakeCase name}}_mcp_client() -> MCPClient | None: """Returns an MCP Client connected to the {{name}} gateway.""" @@ -2759,6 +2910,10 @@ def get_{{snakeCase name}}_mcp_client() -> MCPClient | None: return None {{#if (eq authType "AWS_IAM")}} return MCPClient(lambda: aws_iam_streamablehttp_client(url, aws_service="bedrock-agentcore", aws_region=os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")))) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else {} + return MCPClient(lambda: streamablehttp_client(url, headers=headers)) {{else}} return MCPClient(lambda: streamablehttp_client(url)) {{/if}} diff --git a/src/assets/python/googleadk/base/mcp_client/client.py b/src/assets/python/googleadk/base/mcp_client/client.py index f2c1a39c..21ffb53c 100644 --- a/src/assets/python/googleadk/base/mcp_client/client.py +++ b/src/assets/python/googleadk/base/mcp_client/client.py @@ -10,6 +10,41 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_toolsets() -> list[MCPToolset]: """Returns MCP Toolsets for all configured gateways.""" @@ -24,6 +59,10 @@ def get_all_gateway_mcp_toolsets() -> list[MCPToolset]: url=url, httpx_client_factory=lambda **kwargs: httpx.AsyncClient(auth=auth, **kwargs) ))) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else None + toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url, headers=headers))) {{else}} toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url))) {{/if}} diff --git a/src/assets/python/langchain_langgraph/base/mcp_client/client.py b/src/assets/python/langchain_langgraph/base/mcp_client/client.py index adcb478a..bbd0288d 100644 --- a/src/assets/python/langchain_langgraph/base/mcp_client/client.py +++ b/src/assets/python/langchain_langgraph/base/mcp_client/client.py @@ -8,6 +8,41 @@ {{#if (includes gatewayAuthTypes "AWS_IAM")}} from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_client() -> MultiServerMCPClient | None: """Returns an MCP Client connected to all configured gateways.""" @@ -19,6 +54,10 @@ def get_all_gateway_mcp_client() -> MultiServerMCPClient | None: session = create_aws_session() auth = SigV4HTTPXAuth(session.get_credentials(), "bedrock-agentcore", session.region_name) servers["{{name}}"] = {"transport": "streamable_http", "url": url, "auth": auth} + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else None + servers["{{name}}"] = {"transport": "streamable_http", "url": url, "headers": headers} {{else}} servers["{{name}}"] = {"transport": "streamable_http", "url": url} {{/if}} diff --git a/src/assets/python/openaiagents/base/mcp_client/client.py b/src/assets/python/openaiagents/base/mcp_client/client.py index 39612c38..b7dc18cc 100644 --- a/src/assets/python/openaiagents/base/mcp_client/client.py +++ b/src/assets/python/openaiagents/base/mcp_client/client.py @@ -9,6 +9,41 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} + +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]: """Returns MCP servers for all configured gateways.""" @@ -23,6 +58,10 @@ def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]: name="{{name}}", params={"url": url, "httpx_client_factory": lambda **kwargs: httpx.AsyncClient(auth=auth, **kwargs)} )) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else {} + servers.append(MCPServerStreamableHttp(name="{{name}}", params={"url": url, "headers": headers})) {{else}} servers.append(MCPServerStreamableHttp(name="{{name}}", params={"url": url})) {{/if}} diff --git a/src/assets/python/strands/base/mcp_client/client.py b/src/assets/python/strands/base/mcp_client/client.py index 3b77cdac..77147b13 100644 --- a/src/assets/python/strands/base/mcp_client/client.py +++ b/src/assets/python/strands/base/mcp_client/client.py @@ -9,7 +9,41 @@ {{#if (includes gatewayAuthTypes "AWS_IAM")}} from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} +{{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} +import httpx as _httpx +import time as _time +{{/if}} + +{{#each gatewayProviders}} +{{#if (eq authType "CUSTOM_JWT")}} +_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} +def _get_bearer_token_{{snakeCase name}}(): + """Obtain OAuth access token via client_credentials grant for {{name}}.""" + cache = _token_cache_{{snakeCase name}} + if cache["token"] and _time.time() < cache["expires_at"]: + return cache["token"] + client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") + client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") + if not client_id or not client_secret: + logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") + return None + with _httpx.Client() as c: + disc = c.get("{{discoveryUrl}}") + token_ep = disc.json()["token_endpoint"] + resp = c.post(token_ep, data={ + "grant_type": "client_credentials", + "client_id": client_id, + "client_secret": client_secret, + {{#if scopes}}"scope": "{{scopes}}",{{/if}} + }) + data = resp.json() + cache["token"] = data["access_token"] + cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 + return cache["token"] + +{{/if}} +{{/each}} {{#each gatewayProviders}} def get_{{snakeCase name}}_mcp_client() -> MCPClient | None: """Returns an MCP Client connected to the {{name}} gateway.""" @@ -19,6 +53,10 @@ def get_{{snakeCase name}}_mcp_client() -> MCPClient | None: return None {{#if (eq authType "AWS_IAM")}} return MCPClient(lambda: aws_iam_streamablehttp_client(url, aws_service="bedrock-agentcore", aws_region=os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")))) + {{else if (eq authType "CUSTOM_JWT")}} + token = _get_bearer_token_{{snakeCase name}}() + headers = {"Authorization": f"Bearer {token}"} if token else {} + return MCPClient(lambda: streamablehttp_client(url, headers=headers)) {{else}} return MCPClient(lambda: streamablehttp_client(url)) {{/if}} diff --git a/src/cli/operations/agent/generate/schema-mapper.ts b/src/cli/operations/agent/generate/schema-mapper.ts index 996fab70..0f32afeb 100644 --- a/src/cli/operations/agent/generate/schema-mapper.ts +++ b/src/cli/operations/agent/generate/schema-mapper.ts @@ -188,11 +188,32 @@ async function mapMcpGatewaysToGatewayProviders(): Promise ({ - name: gateway.name, - envVarName: computeDefaultGatewayEnvVarName(gateway.name), - authType: gateway.authorizerType, - })); + const project = await configIO.readProjectSpec(); + + return mcpSpec.agentCoreGateways.map(gateway => { + const config: GatewayProviderRenderConfig = { + name: gateway.name, + envVarName: computeDefaultGatewayEnvVarName(gateway.name), + authType: gateway.authorizerType, + }; + + if (gateway.authorizerType === 'CUSTOM_JWT' && gateway.authorizerConfiguration?.customJwtAuthorizer) { + const jwtConfig = gateway.authorizerConfiguration.customJwtAuthorizer; + const credName = `${gateway.name}-agent-oauth`; + const credential = project.credentials.find(c => c.name === credName); + + if (credential) { + config.credentialEnvVarBase = computeDefaultCredentialEnvVarName(credName); + config.discoveryUrl = jwtConfig.discoveryUrl; + const scopes = 'allowedScopes' in jwtConfig ? (jwtConfig as { allowedScopes?: string[] }).allowedScopes : undefined; + if (scopes?.length) { + config.scopes = scopes.join(' '); + } + } + } + + return config; + }); } catch { return []; } diff --git a/src/cli/templates/types.ts b/src/cli/templates/types.ts index dbde6651..098acf07 100644 --- a/src/cli/templates/types.ts +++ b/src/cli/templates/types.ts @@ -25,6 +25,12 @@ export interface GatewayProviderRenderConfig { name: string; envVarName: string; authType: string; // AWS_IAM, CUSTOM_JWT, NONE + /** Env var prefix for managed credential (CUSTOM_JWT only) */ + credentialEnvVarBase?: string; + /** OIDC discovery URL for token endpoint lookup (CUSTOM_JWT only) */ + discoveryUrl?: string; + /** Space-separated scopes for token request (CUSTOM_JWT only) */ + scopes?: string; } /** diff --git a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx index d10c6202..fc93dea3 100644 --- a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx +++ b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx @@ -320,11 +320,9 @@ function JwtConfigInput({ {subStep === 4 && ( )} {subStep === 5 && ( From 7e002e0844555fa5ae5c0d1309edeafb9fe4b995 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:46:26 -0500 Subject: [PATCH 5/9] feat: protect managed credentials from accidental deletion --- src/cli/commands/remove/actions.ts | 2 +- src/cli/operations/remove/remove-identity.ts | 18 +++++++++++++++++- src/cli/tui/hooks/useRemove.ts | 2 +- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/cli/commands/remove/actions.ts b/src/cli/commands/remove/actions.ts index 35681c69..74604ea2 100644 --- a/src/cli/commands/remove/actions.ts +++ b/src/cli/commands/remove/actions.ts @@ -72,7 +72,7 @@ export async function handleRemove(options: ValidatedRemoveOptions): Promise { +export async function removeCredential(credentialName: string, options?: { force?: boolean }): Promise { try { const configIO = new ConfigIO(); const project = await configIO.readProjectSpec(); @@ -95,6 +101,16 @@ export async function removeCredential(credentialName: string): Promise => { setState({ isLoading: true, result: null }); - const result = await removeIdentity(identityName); + const result = await removeIdentity(identityName, { force: true }); setState({ isLoading: false, result }); let logPath: string | undefined; From d0cc543a6a3d16ad74a5770cf22c2c5cc2498bff Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 13:52:41 -0500 Subject: [PATCH 6/9] test: add tests for CUSTOM_JWT CLI validation and managed credential protection --- .../commands/add/__tests__/validate.test.ts | 41 +++++++++++++ .../remove/__tests__/remove-identity.test.ts | 57 ++++++++++++++++++- 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/src/cli/commands/add/__tests__/validate.test.ts b/src/cli/commands/add/__tests__/validate.test.ts index 3b319c9e..0d4f7961 100644 --- a/src/cli/commands/add/__tests__/validate.test.ts +++ b/src/cli/commands/add/__tests__/validate.test.ts @@ -240,6 +240,47 @@ describe('validate', () => { expect(validateAddGatewayOptions(validGatewayOptionsNone)).toEqual({ valid: true }); expect(validateAddGatewayOptions(validGatewayOptionsJwt)).toEqual({ valid: true }); }); + + // AC15: agentClientId and agentClientSecret must be provided together + it('returns error when agentClientId provided without agentClientSecret', () => { + const result = validateAddGatewayOptions({ + ...validGatewayOptionsJwt, + agentClientId: 'my-client-id', + }); + expect(result.valid).toBe(false); + expect(result.error).toBe('Both --agent-client-id and --agent-client-secret must be provided together'); + }); + + it('returns error when agentClientSecret provided without agentClientId', () => { + const result = validateAddGatewayOptions({ + ...validGatewayOptionsJwt, + agentClientSecret: 'my-secret', + }); + expect(result.valid).toBe(false); + expect(result.error).toBe('Both --agent-client-id and --agent-client-secret must be provided together'); + }); + + // AC16: agent credentials only valid with CUSTOM_JWT + it('returns error when agent credentials used with non-CUSTOM_JWT authorizer', () => { + const result = validateAddGatewayOptions({ + ...validGatewayOptionsNone, + agentClientId: 'my-client-id', + agentClientSecret: 'my-secret', + }); + expect(result.valid).toBe(false); + expect(result.error).toBe('Agent OAuth credentials are only valid with CUSTOM_JWT authorizer'); + }); + + // AC17: valid CUSTOM_JWT with agent credentials passes + it('passes for CUSTOM_JWT with agent credentials', () => { + const result = validateAddGatewayOptions({ + ...validGatewayOptionsJwt, + agentClientId: 'my-client-id', + agentClientSecret: 'my-secret', + allowedScopes: 'scope1,scope2', + }); + expect(result.valid).toBe(true); + }); }); describe('validateAddGatewayTargetOptions', () => { diff --git a/src/cli/operations/remove/__tests__/remove-identity.test.ts b/src/cli/operations/remove/__tests__/remove-identity.test.ts index b6172a33..d5f97e90 100644 --- a/src/cli/operations/remove/__tests__/remove-identity.test.ts +++ b/src/cli/operations/remove/__tests__/remove-identity.test.ts @@ -1,8 +1,9 @@ -import { previewRemoveCredential } from '../remove-identity.js'; +import { previewRemoveCredential, removeCredential } from '../remove-identity.js'; import { describe, expect, it, vi } from 'vitest'; -const { mockReadProjectSpec, mockConfigExists, mockReadMcpSpec } = vi.hoisted(() => ({ +const { mockReadProjectSpec, mockWriteProjectSpec, mockConfigExists, mockReadMcpSpec } = vi.hoisted(() => ({ mockReadProjectSpec: vi.fn(), + mockWriteProjectSpec: vi.fn(), mockConfigExists: vi.fn(), mockReadMcpSpec: vi.fn(), })); @@ -10,6 +11,7 @@ const { mockReadProjectSpec, mockConfigExists, mockReadMcpSpec } = vi.hoisted(() vi.mock('../../../../lib/index.js', () => ({ ConfigIO: class { readProjectSpec = mockReadProjectSpec; + writeProjectSpec = mockWriteProjectSpec; configExists = mockConfigExists; readMcpSpec = mockReadMcpSpec; }, @@ -118,4 +120,55 @@ describe('previewRemoveCredential', () => { 'Warning: Credential "test-cred" is referenced by gateway targets: gateway2/target2. Removing it may break these targets.' ); }); + + it('shows managed credential warning in preview', async () => { + mockReadProjectSpec.mockResolvedValue({ + credentials: [{ name: 'gw-agent-oauth', type: 'OAuthCredentialProvider', managed: true, usage: 'inbound' }], + }); + mockConfigExists.mockReturnValue(false); + + const result = await previewRemoveCredential('gw-agent-oauth'); + + const warning = result.summary.find(s => s.includes('auto-created')); + expect(warning).toBeTruthy(); + }); +}); + +describe('removeCredential', () => { + it('blocks removal of managed credential without force', async () => { + mockReadProjectSpec.mockResolvedValue({ + credentials: [{ name: 'gw-agent-oauth', type: 'OAuthCredentialProvider', managed: true, usage: 'inbound' }], + }); + mockConfigExists.mockReturnValue(false); + + const result = await removeCredential('gw-agent-oauth'); + + expect(result.ok).toBe(false); + expect(result.error).toContain('auto-created'); + expect(result.error).toContain('--force'); + }); + + it('allows removal of managed credential with force', async () => { + mockReadProjectSpec.mockResolvedValue({ + credentials: [{ name: 'gw-agent-oauth', type: 'OAuthCredentialProvider', managed: true, usage: 'inbound' }], + }); + mockConfigExists.mockReturnValue(false); + mockWriteProjectSpec.mockResolvedValue(undefined); + + const result = await removeCredential('gw-agent-oauth', { force: true }); + + expect(result.ok).toBe(true); + }); + + it('allows removal of non-managed credential without force', async () => { + mockReadProjectSpec.mockResolvedValue({ + credentials: [{ name: 'regular-cred', type: 'OAuthCredentialProvider' }], + }); + mockConfigExists.mockReturnValue(false); + mockWriteProjectSpec.mockResolvedValue(undefined); + + const result = await removeCredential('regular-cred'); + + expect(result.ok).toBe(true); + }); }); From a246fed9a1bb1e2de9dfe4ab186dc3a158558429 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 26 Feb 2026 14:05:43 -0500 Subject: [PATCH 7/9] fix: resolve httpx import collision between AWS_IAM and CUSTOM_JWT templates --- .../assets.snapshot.test.ts.snap | 24 +++++++++---------- .../googleadk/base/mcp_client/client.py | 6 ++--- .../base/mcp_client/client.py | 6 ++--- .../openaiagents/base/mcp_client/client.py | 6 ++--- .../python/strands/base/mcp_client/client.py | 6 ++--- 5 files changed, 24 insertions(+), 24 deletions(-) diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index 434821b4..c4273fa2 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -1727,8 +1727,8 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -1745,7 +1745,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ @@ -2052,8 +2052,8 @@ logger = logging.getLogger(__name__) from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -2070,7 +2070,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ @@ -2517,8 +2517,8 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -2535,7 +2535,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ @@ -2867,8 +2867,8 @@ logger = logging.getLogger(__name__) from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -2885,7 +2885,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ diff --git a/src/assets/python/googleadk/base/mcp_client/client.py b/src/assets/python/googleadk/base/mcp_client/client.py index 21ffb53c..d3cc6cab 100644 --- a/src/assets/python/googleadk/base/mcp_client/client.py +++ b/src/assets/python/googleadk/base/mcp_client/client.py @@ -11,8 +11,8 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -29,7 +29,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ diff --git a/src/assets/python/langchain_langgraph/base/mcp_client/client.py b/src/assets/python/langchain_langgraph/base/mcp_client/client.py index bbd0288d..bb78656d 100644 --- a/src/assets/python/langchain_langgraph/base/mcp_client/client.py +++ b/src/assets/python/langchain_langgraph/base/mcp_client/client.py @@ -9,8 +9,8 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -27,7 +27,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ diff --git a/src/assets/python/openaiagents/base/mcp_client/client.py b/src/assets/python/openaiagents/base/mcp_client/client.py index b7dc18cc..762bd11a 100644 --- a/src/assets/python/openaiagents/base/mcp_client/client.py +++ b/src/assets/python/openaiagents/base/mcp_client/client.py @@ -10,8 +10,8 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -28,7 +28,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ diff --git a/src/assets/python/strands/base/mcp_client/client.py b/src/assets/python/strands/base/mcp_client/client.py index 77147b13..8e71d5a0 100644 --- a/src/assets/python/strands/base/mcp_client/client.py +++ b/src/assets/python/strands/base/mcp_client/client.py @@ -10,8 +10,8 @@ from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -import httpx as _httpx -import time as _time +{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx +{{/unless}}import time as _time {{/if}} {{#each gatewayProviders}} @@ -28,7 +28,7 @@ def _get_bearer_token_{{snakeCase name}}(): if not client_id or not client_secret: logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") return None - with _httpx.Client() as c: + with httpx.Client() as c: disc = c.get("{{discoveryUrl}}") token_ep = disc.json()["token_endpoint"] resp = c.post(token_ep, data={ From 84f8aa4a1b29c39c53092a193940b5f65d416078 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Fri, 27 Feb 2026 12:08:57 -0500 Subject: [PATCH 8/9] fix: use placeholder instead of initialValue for gateway discovery URL --- src/cli/tui/screens/mcp/AddGatewayScreen.tsx | 23 +++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx index fc93dea3..dca25086 100644 --- a/src/cli/tui/screens/mcp/AddGatewayScreen.tsx +++ b/src/cli/tui/screens/mcp/AddGatewayScreen.tsx @@ -104,9 +104,18 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig }; const handleJwtAgentClientSecret = (clientSecret: string) => { - const audienceList = jwtAudience.split(',').map(s => s.trim()).filter(Boolean); - const clientsList = jwtClients.split(',').map(s => s.trim()).filter(Boolean); - const scopesList = jwtScopes.split(',').map(s => s.trim()).filter(Boolean); + const audienceList = jwtAudience + .split(',') + .map(s => s.trim()) + .filter(Boolean); + const clientsList = jwtClients + .split(',') + .map(s => s.trim()) + .filter(Boolean); + const scopesList = jwtScopes + .split(',') + .map(s => s.trim()) + .filter(Boolean); wizard.setJwtConfig({ discoveryUrl: jwtDiscoveryUrl, @@ -267,12 +276,14 @@ function JwtConfigInput({ return ( Configure Custom JWT Authorizer - Step {subStep + 1} of {totalSteps} + + Step {subStep + 1} of {totalSteps} + {subStep === 0 && ( { From 8d5d92f3f6198b6045d8225f94e676448462828d Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Fri, 27 Feb 2026 15:42:24 -0500 Subject: [PATCH 9/9] feat: wire CUSTOM_JWT inbound auth through AgentCore identity system --- .../assets.snapshot.test.ts.snap | 144 +++++------------- .../googleadk/base/mcp_client/client.py | 36 ++--- .../base/mcp_client/client.py | 36 ++--- .../openaiagents/base/mcp_client/client.py | 36 ++--- .../python/strands/base/mcp_client/client.py | 36 ++--- .../agent/generate/schema-mapper.ts | 2 +- .../operations/identity/create-identity.ts | 2 + src/cli/operations/mcp/create-mcp.ts | 22 +-- .../remove/__tests__/remove-identity.test.ts | 6 +- src/cli/templates/types.ts | 4 +- 10 files changed, 88 insertions(+), 236 deletions(-) diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index c4273fa2..7de977a1 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -1727,37 +1727,19 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} @@ -2052,37 +2034,19 @@ logger = logging.getLogger(__name__) from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} @@ -2517,37 +2481,19 @@ import httpx from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} @@ -2867,37 +2813,19 @@ logger = logging.getLogger(__name__) from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} diff --git a/src/assets/python/googleadk/base/mcp_client/client.py b/src/assets/python/googleadk/base/mcp_client/client.py index d3cc6cab..e6dddd62 100644 --- a/src/assets/python/googleadk/base/mcp_client/client.py +++ b/src/assets/python/googleadk/base/mcp_client/client.py @@ -11,37 +11,19 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} diff --git a/src/assets/python/langchain_langgraph/base/mcp_client/client.py b/src/assets/python/langchain_langgraph/base/mcp_client/client.py index bb78656d..71b336d2 100644 --- a/src/assets/python/langchain_langgraph/base/mcp_client/client.py +++ b/src/assets/python/langchain_langgraph/base/mcp_client/client.py @@ -9,37 +9,19 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} diff --git a/src/assets/python/openaiagents/base/mcp_client/client.py b/src/assets/python/openaiagents/base/mcp_client/client.py index 762bd11a..2fe91136 100644 --- a/src/assets/python/openaiagents/base/mcp_client/client.py +++ b/src/assets/python/openaiagents/base/mcp_client/client.py @@ -10,37 +10,19 @@ from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} diff --git a/src/assets/python/strands/base/mcp_client/client.py b/src/assets/python/strands/base/mcp_client/client.py index 8e71d5a0..01457de2 100644 --- a/src/assets/python/strands/base/mcp_client/client.py +++ b/src/assets/python/strands/base/mcp_client/client.py @@ -10,37 +10,19 @@ from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client {{/if}} {{#if (includes gatewayAuthTypes "CUSTOM_JWT")}} -{{#unless (includes gatewayAuthTypes "AWS_IAM")}}import httpx -{{/unless}}import time as _time +from bedrock_agentcore.identity import requires_access_token {{/if}} {{#each gatewayProviders}} {{#if (eq authType "CUSTOM_JWT")}} -_token_cache_{{snakeCase name}} = {"token": None, "expires_at": 0} - -def _get_bearer_token_{{snakeCase name}}(): - """Obtain OAuth access token via client_credentials grant for {{name}}.""" - cache = _token_cache_{{snakeCase name}} - if cache["token"] and _time.time() < cache["expires_at"]: - return cache["token"] - client_id = os.environ.get("{{credentialEnvVarBase}}_CLIENT_ID") - client_secret = os.environ.get("{{credentialEnvVarBase}}_CLIENT_SECRET") - if not client_id or not client_secret: - logger.warning("Agent OAuth credentials not set — {{name}} CUSTOM_JWT auth unavailable") - return None - with httpx.Client() as c: - disc = c.get("{{discoveryUrl}}") - token_ep = disc.json()["token_endpoint"] - resp = c.post(token_ep, data={ - "grant_type": "client_credentials", - "client_id": client_id, - "client_secret": client_secret, - {{#if scopes}}"scope": "{{scopes}}",{{/if}} - }) - data = resp.json() - cache["token"] = data["access_token"] - cache["expires_at"] = _time.time() + data.get("expires_in", 3600) - 60 - return cache["token"] +@requires_access_token( + provider_name="{{credentialProviderName}}", + scopes=[{{#if scopes}}"{{scopes}}"{{/if}}], + auth_flow="M2M", +) +def _get_bearer_token_{{snakeCase name}}(*, access_token: str): + """Obtain OAuth access token via AgentCore Identity for {{name}}.""" + return access_token {{/if}} {{/each}} diff --git a/src/cli/operations/agent/generate/schema-mapper.ts b/src/cli/operations/agent/generate/schema-mapper.ts index 0f32afeb..b90212c6 100644 --- a/src/cli/operations/agent/generate/schema-mapper.ts +++ b/src/cli/operations/agent/generate/schema-mapper.ts @@ -203,7 +203,7 @@ async function mapMcpGatewaysToGatewayProviders(): Promise c.name === credName); if (credential) { - config.credentialEnvVarBase = computeDefaultCredentialEnvVarName(credName); + config.credentialProviderName = credName; config.discoveryUrl = jwtConfig.discoveryUrl; const scopes = 'allowedScopes' in jwtConfig ? (jwtConfig as { allowedScopes?: string[] }).allowedScopes : undefined; if (scopes?.length) { diff --git a/src/cli/operations/identity/create-identity.ts b/src/cli/operations/identity/create-identity.ts index f42bee61..26a0c672 100644 --- a/src/cli/operations/identity/create-identity.ts +++ b/src/cli/operations/identity/create-identity.ts @@ -14,6 +14,7 @@ export type CreateCredentialConfig = clientSecret: string; scopes?: string[]; vendor?: string; + managed?: boolean; }; /** @@ -143,6 +144,7 @@ export async function createCredential(config: CreateCredentialConfig): Promise< discoveryUrl: config.discoveryUrl, vendor: config.vendor ?? 'CustomOauth2', ...(config.scopes && config.scopes.length > 0 ? { scopes: config.scopes } : {}), + ...(config.managed ? { managed: true } : {}), }; project.credentials.push(credential); await configIO.writeProjectSpec(project); diff --git a/src/cli/operations/mcp/create-mcp.ts b/src/cli/operations/mcp/create-mcp.ts index 6e366508..f8bb6e63 100644 --- a/src/cli/operations/mcp/create-mcp.ts +++ b/src/cli/operations/mcp/create-mcp.ts @@ -1,4 +1,4 @@ -import { ConfigIO, requireConfigRoot, setEnvVar } from '../../../lib'; +import { ConfigIO, requireConfigRoot } from '../../../lib'; import type { AgentCoreCliMcpDefs, AgentCoreGateway, @@ -11,7 +11,7 @@ import { AgentCoreCliMcpDefsSchema, ToolDefinitionSchema } from '../../../schema import { getTemplateToolDefinitions, renderGatewayTargetTemplate } from '../../templates/GatewayTargetRenderer'; import type { AddGatewayConfig, AddGatewayTargetConfig } from '../../tui/screens/mcp/types'; import { DEFAULT_HANDLER, DEFAULT_NODE_VERSION, DEFAULT_PYTHON_VERSION } from '../../tui/screens/mcp/types'; -import { computeDefaultCredentialEnvVarName } from '../identity/create-identity'; +import { createCredential } from '../identity/create-identity'; import { existsSync } from 'fs'; import { mkdir, readFile, writeFile } from 'fs/promises'; import { dirname, join } from 'path'; @@ -206,23 +206,15 @@ export async function createGatewayFromWizard(config: AddGatewayConfig): Promise // Auto-create managed credential if agent OAuth credentials provided if (config.jwtConfig?.agentClientId && config.jwtConfig?.agentClientSecret) { const credName = `${config.name}-agent-oauth`; - const project = await configIO.readProjectSpec(); - - const credential = { - type: 'OAuthCredentialProvider' as const, + await createCredential({ + type: 'OAuthCredentialProvider', name: credName, discoveryUrl: config.jwtConfig.discoveryUrl, + clientId: config.jwtConfig.agentClientId, + clientSecret: config.jwtConfig.agentClientSecret, vendor: 'CustomOauth2', managed: true, - usage: 'inbound' as const, - }; - - project.credentials.push(credential); - await configIO.writeProjectSpec(project); - - const envBase = computeDefaultCredentialEnvVarName(credName); - await setEnvVar(`${envBase}_CLIENT_ID`, config.jwtConfig.agentClientId); - await setEnvVar(`${envBase}_CLIENT_SECRET`, config.jwtConfig.agentClientSecret); + }); } return { name: config.name }; diff --git a/src/cli/operations/remove/__tests__/remove-identity.test.ts b/src/cli/operations/remove/__tests__/remove-identity.test.ts index d5f97e90..2426b345 100644 --- a/src/cli/operations/remove/__tests__/remove-identity.test.ts +++ b/src/cli/operations/remove/__tests__/remove-identity.test.ts @@ -144,8 +144,10 @@ describe('removeCredential', () => { const result = await removeCredential('gw-agent-oauth'); expect(result.ok).toBe(false); - expect(result.error).toContain('auto-created'); - expect(result.error).toContain('--force'); + if (!result.ok) { + expect(result.error).toContain('auto-created'); + expect(result.error).toContain('--force'); + } }); it('allows removal of managed credential with force', async () => { diff --git a/src/cli/templates/types.ts b/src/cli/templates/types.ts index 098acf07..c9366def 100644 --- a/src/cli/templates/types.ts +++ b/src/cli/templates/types.ts @@ -25,8 +25,8 @@ export interface GatewayProviderRenderConfig { name: string; envVarName: string; authType: string; // AWS_IAM, CUSTOM_JWT, NONE - /** Env var prefix for managed credential (CUSTOM_JWT only) */ - credentialEnvVarBase?: string; + /** Credential provider name for @requires_access_token (CUSTOM_JWT only) */ + credentialProviderName?: string; /** OIDC discovery URL for token endpoint lookup (CUSTOM_JWT only) */ discoveryUrl?: string; /** Space-separated scopes for token request (CUSTOM_JWT only) */