fix(export): export agentcore_gateway tools from harnesses

src/assets/templates/export-harness-python/main.py

@@ -38,14 +38,13 @@
 {{/if}}
 from bedrock_agentcore.runtime import BedrockAgentCoreApp
 from model.load import load_model
-{{#if remoteMcpTools}}
-from mcp_client.client import get_all_remote_mcp_clients
-{{/if}}
-{{#unless remoteMcpTools}}
+{{#if (or gatewayTools remoteMcpTools)}}
+from mcp_client.client import get_all_mcp_clients
+{{else}}
 {{#unless isExportHarness}}
 from mcp_client.client import get_streamable_http_mcp_client
 {{/unless}}
-{{/unless}}
+{{/if}}
 {{#if hasMemory}}
 from memory.session import get_memory_session_manager
 {{/if}}
@@ -58,12 +57,9 @@
 app = BedrockAgentCoreApp()
 log = app.logger
 
-{{#if remoteMcpTools}}
+{{#if (or gatewayTools remoteMcpTools)}}
 # Define MCP clients for all configured MCP servers (gateways and/or remote MCP)
-mcp_clients = []
-{{#if remoteMcpTools}}
-mcp_clients += get_all_remote_mcp_clients()
-{{/if}}
+mcp_clients = get_all_mcp_clients()
 {{else}}
 {{#unless isExportHarness}}
 # Define a Streamable HTTP MCP Client
@@ -279,7 +275,7 @@ def list_files(path: str) -> str:
 tools.extend([file_read, file_write, list_files])
 {{/unless}}{{/if}}
 
-{{#if remoteMcpTools}}
+{{#if (or gatewayTools remoteMcpTools)}}
 # Add MCP clients to tools
 for mcp_client in mcp_clients:
     if mcp_client:

src/assets/templates/export-harness-python/mcp_client/client.py

@@ -5,8 +5,7 @@
 
 logger = logging.getLogger(__name__)
 
-{{#if remoteMcpTools}}
-{{#if (some remoteMcpTools "toolPatterns")}}
+{{#if (or (some remoteMcpTools "toolPatterns") (some gatewayTools "toolPatterns"))}}
 import fnmatch
 
 
@@ -18,6 +17,92 @@ def allowed(tool, **_):
     return {"allowed": [allowed]}
 
 {{/if}}
+{{#if gatewayTools}}
+import httpx
+
+{{/if}}
+{{#if (includes gatewayAuthTypes "AWS_IAM")}}
+from botocore.auth import SigV4Auth
+from botocore.awsrequest import AWSRequest
+from botocore.session import get_session
+
+
+class _SigV4Auth(httpx.Auth):
+    """Signs each gateway request with SigV4 using the runtime's AWS credentials."""
+
+    requires_request_body = True
+
+    def __init__(self):
+        region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION"))
+        self._signer = SigV4Auth(get_session().get_credentials(), "bedrock-agentcore", region)
+
+    def auth_flow(self, request):
+        # The gateway does not sign the connection header, so leave it out of the signature.
+        headers = {k: v for k, v in request.headers.items() if k.lower() != "connection"}
+        signed = AWSRequest(method=request.method, url=str(request.url), data=request.content, headers=headers)
+        self._signer.add_auth(signed)
+        request.headers.update(dict(signed.headers))
+        yield request
+
+{{/if}}
+{{#if (includes gatewayAuthTypes "OAUTH")}}
+from bedrock_agentcore.identity.auth import requires_access_token
+
+
+class _BearerAuth(httpx.Auth):
+    """Sends an AgentCore Identity access token, fetching a new one when the gateway rejects it."""
+
+    requires_request_body = True
+
+    def __init__(self, get_token):
+        self._get_token = get_token
+        self._token = None
+
+    async def async_auth_flow(self, request):
+        if self._token is None:
+            self._token = await self._get_token()
+        request.headers["Authorization"] = f"Bearer {self._token}"
+        response = yield request
+        if response.status_code == 401:
+            self._token = await self._get_token()
+            request.headers["Authorization"] = f"Bearer {self._token}"
+            yield request
+
+{{/if}}
+{{#each gatewayTools}}
+{{#if oauth}}
+@requires_access_token(
+    provider_name={{safeJson oauth.providerName}},
+    scopes={{safeJson oauth.scopes}},
+    auth_flow="M2M",
+{{#if oauth.customParameters}}
+    custom_parameters={{safeJson oauth.customParameters}},
+{{/if}}
+)
+async def _get_{{pythonName}}_token(*, access_token: str) -> str:
+    """Fetch an OAuth access token for the {{name}} gateway from AgentCore Identity."""
+    return access_token
+
+{{/if}}
+def get_{{pythonName}}_mcp_client() -> MCPClient | None:
+    """Returns an MCP Client for the {{name}} AgentCore Gateway."""
+    url = os.environ.get("{{urlEnvVarName}}")
+    if not url:
+        logger.warning("{{urlEnvVarName}} is not set; {{name}} gateway tools are unavailable")
+        return None
+    {{#if (eq authType "AWS_IAM")}}
+    auth = _SigV4Auth()
+    {{/if}}
+    {{#if oauth}}
+    auth = _BearerAuth(_get_{{pythonName}}_token)
+    {{/if}}
+    {{#if (eq authType "NONE")}}
+    auth = None
+    {{/if}}
+    return MCPClient(lambda: streamablehttp_client(url, auth=auth){{#if toolPatterns}}, tool_filters=_allowed_tools({{safeJson name}}, {{#each toolPatterns}}{{safeJson this}}{{#unless @last}}, {{/unless}}{{/each}}){{/if}})
+
+{{/each}}
+{{#if remoteMcpTools}}
 {{#if (some remoteMcpTools "headerCredentials")}}
 from bedrock_agentcore.identity.auth import requires_api_key
 {{/if}}
@@ -48,12 +133,21 @@ def transport():
     {{/if}}
 
 {{/each}}
-def get_all_remote_mcp_clients() -> list[MCPClient]:
-    """Returns all configured remote MCP clients."""
-    clients = [{{#each remoteMcpTools}}get_{{pythonName}}_mcp_client(){{#unless @last}}, {{/unless}}{{/each}}]
+{{/if}}
+{{#if (or gatewayTools remoteMcpTools)}}
+def get_all_mcp_clients() -> list[MCPClient]:
+    """Returns all configured gateway and remote MCP clients."""
+    clients = [
+{{#each gatewayTools}}
+        get_{{pythonName}}_mcp_client(),
+{{/each}}
+{{#each remoteMcpTools}}
+        get_{{pythonName}}_mcp_client(),
+{{/each}}
+    ]
     return [c for c in clients if c is not None]
 {{/if}}
-{{#unless remoteMcpTools}}
+{{#unless (or gatewayTools remoteMcpTools)}}
 {{#if isVpc}}
 # VPC mode: external MCP endpoints are not reachable without a NAT gateway.
 # Add an AgentCore Gateway with `agentcore add gateway`, or configure your own endpoint below.

src/core/project/manager.export.test.ts

@@ -117,6 +117,62 @@ describe("FsProjectManager.exportHarness rendered tree", () => {
     expect(client).toContain('tool_filters=_allowed_tools("exa", "web_*")');
   });
 
+  test("connects gateway tools through their URL variable and records the connections", async () => {
+    const { manager: subject } = manager();
+    const gatewayArn = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/example-abc123";
+    const project = await projectWithHarness(subject, {
+      allowedTools: ["@tools/*", "@partner"],
+      tools: [
+        {
+          type: "agentcore_gateway",
+          name: "tools",
+          config: { agentCoreGateway: { gatewayArn, outboundAuth: { awsIam: {} } } },
+        },
+        {
+          type: "agentcore_gateway",
+          name: "partner",
+          config: {
+            agentCoreGateway: {
+              gatewayArn,
+              outboundAuth: {
+                oauth: {
+                  providerArn:
+                    "arn:aws:bedrock-agentcore:us-east-1:123456789012:token-vault/default/oauth2credentialprovider/partner-oauth",
+                  scopes: ["partner/invoke"],
+                },
+              },
+            },
+          },
+        },
+      ],
+    });
+
+    const result = await drain(subject.exportHarness(project, exportInput()));
+
+    const client = await Bun.file(join(result.agentPath, "mcp_client", "client.py")).text();
+    expect(client).toContain('url = os.environ.get("GATEWAY_TOOLS_URL")');
+    expect(client).toMatch(/GATEWAY_TOOLS_URL"\)[\s\S]*?auth = _SigV4Auth\(\)/);
+    expect(client).toContain('provider_name="partner-oauth"');
+    expect(client).toMatch(/auth = _BearerAuth\(_get_partner_[a-f0-9]{10}_token\)/);
+    const main = await Bun.file(join(result.agentPath, "main.py")).text();
+    expect(main).toContain("mcp_clients = get_all_mcp_clients()");
+    const pyproject = await Bun.file(join(result.agentPath, "pyproject.toml")).text();
+    expect(pyproject).not.toMatch(/^ {5,}"/m);
+    const spec = await Bun.file(join(project.rootPath, "agentcore", "agentcore.json")).json();
+    expect(
+      spec.runtimes.find((runtime: { name: string }) => runtime.name === "assistantAgent")
+        .connections,
+    ).toEqual([
+      { id: "tools", to: { type: "gateway", arn: gatewayArn, outboundAuth: { awsIam: {} } } },
+      expect.objectContaining({ id: "partner" }),
+    ]);
+    expect(
+      (await readEnvFile(join(project.rootPath, "agentcore", ".env.local"))).GATEWAY_TOOLS_URL,
+    ).toBe("https://example-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp");
+    const notes = await Bun.file(join(result.agentPath, "EXPORT_NOTES.md")).text();
+    expect(notes).not.toContain("Gateway tool not exported");
+  });
+
   test("merges service model parameters under the explicit settings", async () => {
     const { manager: subject } = manager();
     const project = await projectWithHarness(subject);