diff --git a/src/assets/templates/export-harness-python/main.py b/src/assets/templates/export-harness-python/main.py index 9424df446..2cb7bbc79 100644 --- a/src/assets/templates/export-harness-python/main.py +++ b/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: diff --git a/src/assets/templates/export-harness-python/mcp_client/client.py b/src/assets/templates/export-harness-python/mcp_client/client.py index 6b600156e..076fc5c9d 100644 --- a/src/assets/templates/export-harness-python/mcp_client/client.py +++ b/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,50 @@ def allowed(tool, **_): return {"allowed": [allowed]} {{/if}} +{{#if (includes gatewayAuthTypes "AWS_IAM")}} +from mcp_proxy_for_aws.client import aws_iam_streamablehttp_client + +{{/if}} +{{#if (includes gatewayAuthTypes "OAUTH")}} +from bedrock_agentcore.identity.auth import requires_access_token + +{{/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}} +) +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 + def transport(): + {{#if (eq authType "AWS_IAM")}} + return aws_iam_streamablehttp_client(url, aws_service="bedrock-agentcore", aws_region=os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION"))) + {{/if}} + {{#if oauth}} + return streamablehttp_client(url, headers={"Authorization": f"Bearer {_get_{{pythonName}}_token()}"}) + {{/if}} + {{#if (eq authType "NONE")}} + return streamablehttp_client(url) + {{/if}} + + return MCPClient(transport{{#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 +91,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. diff --git a/src/assets/templates/export-harness-python/pyproject.toml b/src/assets/templates/export-harness-python/pyproject.toml index a0943dfb6..be69973be 100644 --- a/src/assets/templates/export-harness-python/pyproject.toml +++ b/src/assets/templates/export-harness-python/pyproject.toml @@ -13,7 +13,8 @@ dependencies = [ "bedrock-agentcore[strands-agents] >= 1.18.1, < 2.0.0", "botocore[crt] ~= 1.43.0", "mcp >= 1.28.1, < 2.0.0", - {{#if bedrockMantle}}"aws-bedrock-token-generator >= 1.1.0, < 2.0.0", + {{#if (includes gatewayAuthTypes "AWS_IAM")}}"mcp-proxy-for-aws >= 1.1.0, < 2.0.0", + {{/if}} {{#if bedrockMantle}}"aws-bedrock-token-generator >= 1.1.0, < 2.0.0", {{/if}}"strands-agents{{#if strandsExtras}}[{{strandsExtras}}]{{/if}} >= 1.62.0, < 2.0.0", ] diff --git a/src/core/project/manager.export.test.ts b/src/core/project/manager.export.test.ts index 37995ae2c..403fba72a 100644 --- a/src/core/project/manager.export.test.ts +++ b/src/core/project/manager.export.test.ts @@ -117,6 +117,64 @@ 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).toContain('aws_iam_streamablehttp_client(url, aws_service="bedrock-agentcore"'); + expect(client).toContain('provider_name="partner-oauth"'); + expect(client).toMatch( + /headers=\{"Authorization": f"Bearer \{_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).toContain('"mcp-proxy-for-aws >= 1.1.0, < 2.0.0"'); + 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); diff --git a/src/core/project/templates/export.test.ts b/src/core/project/templates/export.test.ts index 0a58b0f50..cf8a9c8c6 100644 --- a/src/core/project/templates/export.test.ts +++ b/src/core/project/templates/export.test.ts @@ -396,15 +396,6 @@ describe("mapHarnessToExportPlan tools", () => { const result = plan({ spec: harness({ tools: [ - { - type: "agentcore_gateway", - name: "gw", - config: { - agentCoreGateway: { - gatewayArn: "arn:aws:bedrock-agentcore:us-east-1:111122223333:gateway/g-1", - }, - }, - }, { type: "agentcore_browser", name: "browser" }, { type: "agentcore_code_interpreter", name: "ci" }, ], @@ -415,14 +406,11 @@ describe("mapHarnessToExportPlan tools", () => { // branch to render them from — the notes below are the whole output. expect(result.context.hasBrowser).toBeUndefined(); expect(result.context.hasCodeInterpreter).toBeUndefined(); - expect(result.context.hasGateway).toBeUndefined(); expect(result.context.remoteMcpTools).toBeUndefined(); expect(categories(result)).toEqual([ - GATEWAY_TOOL_NOTE_CATEGORY, BROWSER_TOOL_NOTE_CATEGORY, CODE_INTERPRETER_TOOL_NOTE_CATEGORY, ]); - expect(result.notes[0]!.message).toContain("gateway/g-1"); }); test("includes the harness builtins unless allowedTools filters them out", () => { @@ -461,6 +449,148 @@ describe("mapHarnessToExportPlan tools", () => { }); }); +describe("mapHarnessToExportPlan gateway tools", () => { + const GATEWAY_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:gateway/example-abc123"; + const PROVIDER_ARN = + "arn:aws:bedrock-agentcore:us-east-1:123456789012:token-vault/default/oauth2credentialprovider/gw-oauth"; + const MEMORY_CONNECTION = { + id: "notes", + to: { + type: "memory", + arn: "arn:aws:bedrock-agentcore:us-east-1:123456789012:memory/notes-AbCdEf1234", + }, + } as const; + const gatewayTool = (name: string, outboundAuth?: Record) => ({ + type: "agentcore_gateway", + name, + config: { + agentCoreGateway: { gatewayArn: GATEWAY_ARN, ...(outboundAuth && { outboundAuth }) }, + }, + }); + const oauth = { + providerArn: PROVIDER_ARN, + scopes: ["tools/invoke"], + grantType: "CLIENT_CREDENTIALS" as const, + customParameters: { audience: "tools" }, + }; + + test.each([ + ["no outboundAuth", undefined, "AWS_IAM", undefined, { awsIam: {} }], + ["awsIam", { awsIam: {} }, "AWS_IAM", undefined, { awsIam: {} }], + ["none", { none: {} }, "NONE", undefined, { none: {} }], + [ + "oauth client credentials", + { oauth }, + "OAUTH", + { + providerName: "gw-oauth", + scopes: ["tools/invoke"], + customParameters: { audience: "tools" }, + }, + { oauth }, + ], + ])( + "exports a gateway tool with %s as a client and a runtime connection", + (_, outboundAuth, authType, oauthContext, connectionAuth) => { + const result = plan({ + spec: harness({ + tools: [gatewayTool("my-tools", outboundAuth)], + connections: [MEMORY_CONNECTION], + }), + }); + + expect(result.context.gatewayTools).toEqual([ + { + name: "my-tools", + pythonName: expect.stringMatching(/^my_tools_[a-f0-9]{10}$/), + urlEnvVarName: "GATEWAY_MY_TOOLS_URL", + authType, + oauth: oauthContext, + toolPatterns: undefined, + }, + ]); + expect(result.context.gatewayAuthTypes).toEqual([authType]); + expect(result.runtime.connections).toEqual([ + MEMORY_CONNECTION, + { id: "my-tools", to: { type: "gateway", arn: GATEWAY_ARN, outboundAuth: connectionAuth } }, + ]); + expect(result.envEntries).toEqual([ + { + key: "GATEWAY_MY_TOOLS_URL", + value: "https://example-abc123.gateway.bedrock-agentcore.us-east-1.amazonaws.com/mcp", + comment: expect.stringContaining('gateway tool "my-tools"'), + }, + ]); + expect(result.notes).toEqual([]); + }, + ); + + test("narrows a gateway to the tools @server/tool selects and drops an unselected one", () => { + const result = plan({ + spec: harness({ + tools: [gatewayTool("tools"), gatewayTool("other")], + allowedTools: ["@tools/search*"], + }), + }); + + expect(result.context.gatewayTools).toEqual([ + expect.objectContaining({ name: "tools", toolPatterns: ["search*"] }), + ]); + expect(result.runtime.connections?.map((connection) => connection.id)).toEqual(["tools"]); + }); + + test("reports a USER_FEDERATION gateway tool instead of exporting it", () => { + const result = plan({ + spec: harness({ + tools: [gatewayTool("tools", { oauth: { ...oauth, grantType: "USER_FEDERATION" } })], + }), + }); + + expect(result.context.gatewayTools).toBeUndefined(); + expect(result.runtime.connections).toBeUndefined(); + expect(result.envEntries).toEqual([]); + expect(categories(result)).toEqual([GATEWAY_TOOL_NOTE_CATEGORY]); + expect(result.notes[0]!.message).toContain(GATEWAY_ARN); + }); + + test.each([ + [ + "a name that cannot be a connection id", + [gatewayTool("1tools")], + [], + /must start with a letter/, + ], + [ + "names that collide once tokenized", + [gatewayTool("my-tools"), gatewayTool("my_tools")], + [], + /"MY_TOOLS"/, + ], + [ + "a name that collides with a harness connection", + [gatewayTool("notes")], + [MEMORY_CONNECTION], + /"NOTES"/, + ], + [ + "a gatewayArn that is not a gateway ARN", + [ + { + type: "agentcore_gateway", + name: "tools", + config: { agentCoreGateway: { gatewayArn: "example-abc123" } }, + }, + ], + [], + /not a gateway ARN/, + ], + ])("rejects a gateway tool with %s", (_, tools, connections, message) => { + const exportPlan = () => plan({ spec: harness({ tools, connections }) }); + expect(exportPlan).toThrow(InputValidationError); + expect(exportPlan).toThrow(message); + }); +}); + describe("mapHarnessToExportPlan allowedTools selection", () => { test("a bare name or glob selects builtins only", () => { const shellOnly = plan({ spec: harness({ allowedTools: ["shell"] }) }); diff --git a/src/core/project/templates/export.ts b/src/core/project/templates/export.ts index fb5751c95..4e6e613be 100644 --- a/src/core/project/templates/export.ts +++ b/src/core/project/templates/export.ts @@ -2,6 +2,7 @@ import { createHash } from "node:crypto"; import type { z } from "zod"; import type { ProjectRuntime } from "../../../projectSchemas/runtime"; import type { + HarnessGatewayOutboundAuth, HarnessMemoryRef, HarnessMemoryRetrievalConfig, HarnessSkill, @@ -14,6 +15,7 @@ import type { HarnessTruncationConfig, } from "../../../projectSchemas/harness"; import type { ProjectSpecSchema } from "../../../projectSchemas/project"; +import type { Connection } from "../../../projectSchemas/connections"; import { credentialEnvVarName, credentialProviderNameEnvVarName, @@ -23,7 +25,7 @@ import { memoryEnvVarName, type Memory } from "../../../projectSchemas/memory"; import type { EnvLocalEntry } from "../../../handlers/project/types"; import { InputValidationError } from "../../../errors/errors"; import { toPythonPackageName } from "../fsUtils"; -import { resourceNameFromArn } from "../../arn"; +import { parseArn, resourceNameFromArn, serviceIdFromArn } from "../../arn"; type ProjectSpec = z.infer; @@ -187,12 +189,12 @@ export function mapHarnessToExportPlan(input: HarnessExportInput): HarnessExport ? String(memory.retrievalConfig.relevanceScore) : undefined, actorId: memory.actorId, - // Gateways are never exported as code (see resolveTools); the template still - // needs the keys so its conditionals resolve. // Tools. Empty collections become undefined: the template's custom `or`/ // `some` helpers use JS truthiness, where [] is truthy, unlike `{{#if}}`. inlineFunctionTools: undefinedIfEmpty(tools.inlineFunctionTools), remoteMcpTools: undefinedIfEmpty(tools.remoteMcpTools), + gatewayTools: undefinedIfEmpty(tools.gatewayTools), + gatewayAuthTypes: [...new Set(tools.gatewayTools.map(({ authType }) => authType))], hasShell: tools.hasShell, hasFileOperations: tools.hasFileOperations, // Skills @@ -236,7 +238,9 @@ export function mapHarnessToExportPlan(input: HarnessExportInput): HarnessExport ...(spec.lifecycleConfig && { lifecycleConfiguration: spec.lifecycleConfig }), ...(filesystemConfigurations.length > 0 && { filesystemConfigurations }), ...(additionalPolicies.length > 0 && { additionalPolicies }), - ...(spec.connections?.length && { connections: spec.connections }), + ...((spec.connections?.length || tools.connections.length) && { + connections: [...(spec.connections ?? []), ...tools.connections], + }), ...(spec.tags && { tags: spec.tags }), // NOTE: the harness's executionRoleArn is deliberately NOT carried over. The // exported agent is a new runtime that needs its own CDK-managed role so the @@ -566,6 +570,16 @@ interface ToolsResolution { pythonName: string; }[]; }[]; + gatewayTools: { + name: string; + pythonName: string; + urlEnvVarName: string; + authType: "AWS_IAM" | "NONE" | "OAUTH"; + oauth?: { providerName: string; scopes: string[]; customParameters?: Record }; + toolPatterns?: string[]; + }[]; + /** Runtime connections that make deploy grant gateway access and set each tool's URL variable. */ + connections: Connection[]; hasShell: boolean; hasFileOperations: boolean; } @@ -581,6 +595,8 @@ function resolveTools( const result: ToolsResolution = { inlineFunctionTools: [], remoteMcpTools: [], + gatewayTools: [], + connections: [], // Builtin tools are always available in the harness runtime; include them // unless the allowedTools filter excludes them. hasShell: isBuiltinIncluded("shell", allowedPatterns), @@ -663,15 +679,57 @@ function resolveTools( break; } case "agentcore_gateway": { - const cfg = configOf(tool, "agentCoreGateway") as { gatewayArn?: string } | undefined; - notes.push({ - category: GATEWAY_TOOL_NOTE_CATEGORY, - message: - `The gateway tool "${tool.name}"${cfg?.gatewayArn ? ` (${cfg.gatewayArn})` : ""} was ` + - `not exported: gateway URL discovery, outbound auth, and IAM wiring are managed by ` + - `the harness runtime. To keep these tools, connect an MCP client to the gateway in ` + - `mcp_client/client.py and grant the runtime role bedrock-agentcore:InvokeGateway on ` + - `the gateway (or its OAuth token flow) before deploying.`, + const cfg = configOf(tool, "agentCoreGateway") as + { gatewayArn: string; outboundAuth?: HarnessGatewayOutboundAuth } | undefined; + if (!cfg) break; + const auth = cfg.outboundAuth ?? { awsIam: {} }; + if ("oauth" in auth && auth.oauth.grantType === "USER_FEDERATION") { + notes.push({ + category: GATEWAY_TOOL_NOTE_CATEGORY, + message: + `The gateway tool "${tool.name}" (${cfg.gatewayArn}) was not exported: it uses the ` + + `OAuth grant type USER_FEDERATION, which needs a user consent flow the exported ` + + `agent does not have. Connect to the gateway in mcp_client/client.py yourself, ` + + `and add a gateway connection to this runtime in agentcore.json so deploy grants ` + + `the token access.`, + }); + break; + } + if (!CONNECTION_ID_PATTERN.test(tool.name)) { + throw new InputValidationError( + `Gateway tool "${tool.name}" names the runtime connection that carries its URL, so ` + + `it must start with a letter. Rename the tool in the harness, then export again.`, + ); + } + const urlEnvVarName = `GATEWAY_${connectionEnvToken(tool.name)}_URL`; + const oauth = "oauth" in auth ? auth.oauth : undefined; + result.connections.push({ + id: tool.name, + to: { + type: "gateway", + arn: cfg.gatewayArn, + outboundAuth: + "oauth" in auth + ? { oauth: { ...auth.oauth, grantType: "CLIENT_CREDENTIALS" as const } } + : auth, + }, + }); + envEntries.push({ + key: urlEnvVarName, + value: gatewayUrl(tool.name, cfg.gatewayArn), + comment: `URL of gateway tool "${tool.name}" for agentcore dev (deploy sets it on the runtime)`, + }); + result.gatewayTools.push({ + name: tool.name, + pythonName: stablePythonIdentifier(tool.name), + urlEnvVarName, + authType: oauth ? "OAUTH" : "none" in auth ? "NONE" : "AWS_IAM", + oauth: oauth && { + providerName: resourceNameFromArn(oauth.providerArn), + scopes: oauth.scopes, + customParameters: oauth.customParameters, + }, + toolPatterns: serverToolPatterns(tool.name, allowedPatterns), }); break; } @@ -701,9 +759,41 @@ function resolveTools( } } + const tokens = [ + ...(spec.connections ?? []).flatMap((c) => (c.id ? [connectionEnvToken(c.id)] : [])), + ...result.connections.map((c) => connectionEnvToken(c.id!)), + ]; + const duplicate = tokens.find((token, index) => tokens.indexOf(token) !== index); + if (duplicate) { + throw new InputValidationError( + `Harness "${spec.name}" has gateway tools or connections whose names both become ` + + `"${duplicate}" in environment variable names. Rename one of them, then export again.`, + ); + } + return result; } +const CONNECTION_ID_PATTERN = /^[a-zA-Z][a-zA-Z0-9_-]{0,63}$/; + +/** The env-var token @aws/agentcore-cdk derives from a connection id. */ +function connectionEnvToken(id: string): string { + return id.toUpperCase().replace(/[^A-Z0-9]/g, "_"); +} + +/** The MCP endpoint of a gateway, as @aws/agentcore-cdk derives it from the ARN at deploy. */ +function gatewayUrl(toolName: string, gatewayArn: string): string { + const arn = parseArn(gatewayArn); + if (arn?.service !== "bedrock-agentcore" || !arn.resource.startsWith("gateway/")) { + throw new InputValidationError( + `Gateway tool "${toolName}" has gatewayArn "${gatewayArn}", which is not a gateway ARN.`, + ); + } + const { partition, region } = arn; + const suffix = partition === "aws-cn" ? "amazonaws.com.cn" : "amazonaws.com"; + return `https://${serviceIdFromArn(gatewayArn)}.gateway.bedrock-agentcore.${region}.${suffix}/mcp`; +} + function undefinedIfEmpty(values: T[]): T[] | undefined { return values.length > 0 ? values : undefined; }