@@ -1749,6 +1749,23 @@ logger = logging.getLogger(__name__)
17491749import httpx
17501750from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session
17511751{ {/if }}
1752+ { {#if (includes gatewayAuthTypes "CUSTOM_JWT ")}}
1753+ from bedrock_agentcore.identity import requires_access_token
1754+ { {/if }}
1755+
1756+ { {#each gatewayProviders }}
1757+ { {#if (eq authType "CUSTOM_JWT ")}}
1758+ @requires_access_token(
1759+ provider_name="{ {credentialProviderName }} ",
1760+ scopes=[{ {#if scopes }} "{ {scopes }} "{ {/if }} ],
1761+ auth_flow="M2M",
1762+ )
1763+ def _get_bearer_token_{ {snakeCase name }} (*, access_token: str):
1764+ """Obtain OAuth access token via AgentCore Identity for { {name }} ."""
1765+ return access_token
1766+
1767+ { {/if }}
1768+ { {/each }}
17521769
17531770def get_all_gateway_mcp_toolsets() -> list[MCPToolset]:
17541771 """Returns MCP Toolsets for all configured gateways."""
@@ -1763,6 +1780,10 @@ def get_all_gateway_mcp_toolsets() -> list[MCPToolset]:
17631780 url=url,
17641781 httpx_client_factory=lambda **kwargs: httpx.AsyncClient(auth=auth, **kwargs)
17651782 )))
1783+ { {else if (eq authType "CUSTOM_JWT ")}}
1784+ token = _get_bearer_token_{ {snakeCase name }} ()
1785+ headers = { " Authorization" : f " Bearer {token}" } if token else None
1786+ toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url, headers=headers)))
17661787 { {else }}
17671788 toolsets.append(MCPToolset(connection_params=StreamableHTTPConnectionParams(url=url)))
17681789 { {/if }}
@@ -2035,6 +2056,23 @@ logger = logging.getLogger(__name__)
20352056{ {#if (includes gatewayAuthTypes "AWS_IAM ")}}
20362057from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session
20372058{ {/if }}
2059+ { {#if (includes gatewayAuthTypes "CUSTOM_JWT ")}}
2060+ from bedrock_agentcore.identity import requires_access_token
2061+ { {/if }}
2062+
2063+ { {#each gatewayProviders }}
2064+ { {#if (eq authType "CUSTOM_JWT ")}}
2065+ @requires_access_token(
2066+ provider_name="{ {credentialProviderName }} ",
2067+ scopes=[{ {#if scopes }} "{ {scopes }} "{ {/if }} ],
2068+ auth_flow="M2M",
2069+ )
2070+ def _get_bearer_token_{ {snakeCase name }} (*, access_token: str):
2071+ """Obtain OAuth access token via AgentCore Identity for { {name }} ."""
2072+ return access_token
2073+
2074+ { {/if }}
2075+ { {/each }}
20382076
20392077def get_all_gateway_mcp_client() -> MultiServerMCPClient | None:
20402078 """Returns an MCP Client connected to all configured gateways."""
@@ -2046,6 +2084,10 @@ def get_all_gateway_mcp_client() -> MultiServerMCPClient | None:
20462084 session = create_aws_session()
20472085 auth = SigV4HTTPXAuth(session.get_credentials(), "bedrock-agentcore", session.region_name)
20482086 servers["{ {name }} "] = { " transport" : " streamable_http" , " url" : url , " auth" : auth }
2087+ { {else if (eq authType "CUSTOM_JWT ")}}
2088+ token = _get_bearer_token_{ {snakeCase name }} ()
2089+ headers = { " Authorization" : f " Bearer {token}" } if token else None
2090+ servers["{ {name }} "] = { " transport" : " streamable_http" , " url" : url , " headers" : headers }
20492091 { {else }}
20502092 servers["{ {name }} "] = { " transport" : " streamable_http" , " url" : url }
20512093 { {/if }}
@@ -2460,6 +2502,23 @@ logger = logging.getLogger(__name__)
24602502import httpx
24612503from mcp_proxy_for_aws.sigv4_helper import SigV4HTTPXAuth, create_aws_session
24622504{ {/if }}
2505+ { {#if (includes gatewayAuthTypes "CUSTOM_JWT ")}}
2506+ from bedrock_agentcore.identity import requires_access_token
2507+ { {/if }}
2508+
2509+ { {#each gatewayProviders }}
2510+ { {#if (eq authType "CUSTOM_JWT ")}}
2511+ @requires_access_token(
2512+ provider_name="{ {credentialProviderName }} ",
2513+ scopes=[{ {#if scopes }} "{ {scopes }} "{ {/if }} ],
2514+ auth_flow="M2M",
2515+ )
2516+ def _get_bearer_token_{ {snakeCase name }} (*, access_token: str):
2517+ """Obtain OAuth access token via AgentCore Identity for { {name }} ."""
2518+ return access_token
2519+
2520+ { {/if }}
2521+ { {/each }}
24632522
24642523def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]:
24652524 """Returns MCP servers for all configured gateways."""
@@ -2474,6 +2533,10 @@ def get_all_gateway_mcp_servers() -> list[MCPServerStreamableHttp]:
24742533 name="{ {name }} ",
24752534 params={ " url" : url , " httpx_client_factory" : lambda ** kwargs : httpx .AsyncClient (auth = auth , ** kwargs )}
24762535 ))
2536+ { {else if (eq authType "CUSTOM_JWT ")}}
2537+ token = _get_bearer_token_{ {snakeCase name }} ()
2538+ headers = { " Authorization" : f " Bearer {token}" } if token else { }
2539+ servers.append(MCPServerStreamableHttp(name="{ {name }} ", params={ " url" : url , " headers" : headers } ))
24772540 { {else }}
24782541 servers.append(MCPServerStreamableHttp(name="{ {name }} ", params={ " url" : url } ))
24792542 { {/if }}
@@ -2771,7 +2834,23 @@ logger = logging.getLogger(__name__)
27712834{ {#if (includes gatewayAuthTypes "AWS_IAM ")}}
27722835from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client
27732836{ {/if }}
2837+ { {#if (includes gatewayAuthTypes "CUSTOM_JWT ")}}
2838+ from bedrock_agentcore.identity import requires_access_token
2839+ { {/if }}
2840+
2841+ { {#each gatewayProviders }}
2842+ { {#if (eq authType "CUSTOM_JWT ")}}
2843+ @requires_access_token(
2844+ provider_name="{ {credentialProviderName }} ",
2845+ scopes=[{ {#if scopes }} "{ {scopes }} "{ {/if }} ],
2846+ auth_flow="M2M",
2847+ )
2848+ def _get_bearer_token_{ {snakeCase name }} (*, access_token: str):
2849+ """Obtain OAuth access token via AgentCore Identity for { {name }} ."""
2850+ return access_token
27742851
2852+ { {/if }}
2853+ { {/each }}
27752854{ {#each gatewayProviders }}
27762855def get_{ {snakeCase name }} _mcp_client() -> MCPClient | None:
27772856 """Returns an MCP Client connected to the { {name }} gateway."""
@@ -2781,6 +2860,10 @@ def get_{{snakeCase name}}_mcp_client() -> MCPClient | None:
27812860 return None
27822861 { {#if (eq authType "AWS_IAM ")}}
27832862 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"))))
2863+ { {else if (eq authType "CUSTOM_JWT ")}}
2864+ token = _get_bearer_token_{ {snakeCase name }} ()
2865+ headers = { " Authorization" : f " Bearer {token}" } if token else { }
2866+ return MCPClient(lambda: streamablehttp_client(url, headers=headers))
27842867 { {else }}
27852868 return MCPClient(lambda: streamablehttp_client(url))
27862869 { {/if }}
0 commit comments