feat: Major fixes for Bedrock AgentCore Gateway support

- Expose BedrockAgentCoreGatewayTargetHandler in top-level mcp_lambda
- Fix both Py and TS method names
- Fix TS content lookup for tool name
- Translate gateway tool name to MCP server tool name

Author: clareliguori

README.md

@@ -321,9 +321,9 @@ const requestHandler = new BedrockAgentCoreGatewayTargetHandler(
 );
 
 export const handler: Handler = async (
-  event: event: Record<string, unknown>,
+  event: Record<string, unknown>,
   context: Context
-): Promise<event: Record<string, unknown>> => {
+): Promise<Record<string, unknown>> => {
   return requestHandler.handle(event, context);
 };
 ```

src/python/src/mcp_lambda/__init__.py

@@ -2,6 +2,7 @@
 from .handlers import (
     APIGatewayProxyEventHandler,
     APIGatewayProxyEventV2Handler,
+    BedrockAgentCoreGatewayTargetHandler,
     LambdaFunctionURLEventHandler,
     RequestHandler,
 )
@@ -18,5 +19,6 @@
     "RequestHandler",
     "APIGatewayProxyEventHandler",
     "APIGatewayProxyEventV2Handler",
+    "BedrockAgentCoreGatewayTargetHandler",
     "LambdaFunctionURLEventHandler",
 ]

src/python/src/mcp_lambda/handlers/bedrock_agent_core_gateway_target_handler.py

@@ -18,15 +18,24 @@ class BedrockAgentCoreGatewayTargetHandler:
     def __init__(self, request_handler: RequestHandler):
         self.request_handler = request_handler
 
-    def handle_event(self, event: Dict[str, Any], context: LambdaContext) -> Any:
+    def handle(self, event: Dict[str, Any], context: LambdaContext) -> Any:
         """Handle Lambda invocation from Bedrock AgentCore Gateway"""
         # Extract tool metadata from context
-        tool_name = None
-        if context.client_context and hasattr(context.client_context, "custom"):
-            tool_name = context.client_context.custom.get("bedrockagentcoreToolName")
+        gateway_tool_name = None
 
-        if not tool_name:
-            raise ValueError("Missing bedrockagentcoreToolName in context")
+        if context.client_context and hasattr(context.client_context, "custom"):
+            gateway_tool_name = context.client_context.custom.get(
+                "bedrockAgentCoreToolName"
+            )
+
+        if not gateway_tool_name:
+            raise ValueError("Missing bedrockAgentCoreToolName in context")
+
+        # Gateway names the tools like <target name>___<tool name>
+        parts = gateway_tool_name.split("___", 1)
+        if len(parts) != 2:
+            raise ValueError(f"Invalid tool name format: {gateway_tool_name}")
+        tool_name = parts[1]
 
         # Create JSON-RPC request from gateway event
         jsonrpc_request = JSONRPCRequest(

src/python/tests/test_handlers.py

@@ -527,13 +527,13 @@ def test_handle_valid_tool_invocation(self):
 
         handler = BedrockAgentCoreGatewayTargetHandler(mock_handler)
 
-        # Mock context with tool name
+        # Mock context with gateway tool name
         context = Mock(spec=LambdaContext)
         context.client_context = Mock()
-        context.client_context.custom = {"bedrockagentcoreToolName": "test_tool"}
+        context.client_context.custom = {"bedrockAgentCoreToolName": "target___test_tool"}
 
         event = {"param1": "value1", "param2": "value2"}
-        result = handler.handle_event(event, context)
+        result = handler.handle(event, context)
 
         assert result == {"message": "Tool executed successfully"}
 
@@ -556,9 +556,52 @@ def test_missing_tool_name_raises_error(self):
         event = {"param1": "value1"}
 
         with pytest.raises(
-            ValueError, match="Missing bedrockagentcoreToolName in context"
+            ValueError, match="Missing bedrockAgentCoreToolName in context"
         ):
-            handler.handle_event(event, context)
+            handler.handle(event, context)
+
+    def test_invalid_tool_name_format_raises_error(self):
+        """Test that invalid tool name format raises ValueError."""
+        handler = BedrockAgentCoreGatewayTargetHandler(Mock(spec=RequestHandler))
+
+        # Mock context with invalid tool name format
+        context = Mock(spec=LambdaContext)
+        context.client_context = Mock()
+        context.client_context.custom = {"bedrockAgentCoreToolName": "invalid_format"}
+
+        event = {"param1": "value1"}
+
+        with pytest.raises(
+            ValueError, match="Invalid tool name format: invalid_format"
+        ):
+            handler.handle(event, context)
+
+    def test_multiple_delimiters_in_tool_name(self):
+        """Test that tool name with multiple delimiters works correctly."""
+        # Create a mock request handler that handles tools/call
+        mock_handler = Mock(spec=RequestHandler)
+        mock_handler.handle_request.return_value = JSONRPCResponse(
+            jsonrpc="2.0",
+            result={"message": "Tool executed successfully"},
+            id=1,
+        )
+
+        handler = BedrockAgentCoreGatewayTargetHandler(mock_handler)
+
+        # Mock context with gateway tool name containing multiple delimiters
+        context = Mock(spec=LambdaContext)
+        context.client_context = Mock()
+        context.client_context.custom = {"bedrockAgentCoreToolName": "target___test___tool"}
+
+        event = {"param1": "value1"}
+        result = handler.handle(event, context)
+
+        assert result == {"message": "Tool executed successfully"}
+
+        # Verify the extracted tool name is everything after the first delimiter
+        call_args = mock_handler.handle_request.call_args[0]
+        request = call_args[0]
+        assert request.params["name"] == "test___tool"
 
     def test_request_handler_error_raises_exception(self):
         """Test that request handler errors are raised as exceptions."""
@@ -572,12 +615,12 @@ def test_request_handler_error_raises_exception(self):
 
         handler = BedrockAgentCoreGatewayTargetHandler(mock_handler)
 
-        # Mock context with tool name
+        # Mock context with gateway tool name
         context = Mock(spec=LambdaContext)
         context.client_context = Mock()
-        context.client_context.custom = {"bedrockagentcoreToolName": "unknown_tool"}
+        context.client_context.custom = {"bedrockAgentCoreToolName": "target___unknown_tool"}
 
         event = {"param1": "value1"}
 
         with pytest.raises(Exception, match="Tool not found"):
-            handler.handle_event(event, context)
+            handler.handle(event, context)

src/typescript/src/handlers/bedrockAgentCoreGatewayTargetHandler.ts

@@ -18,18 +18,29 @@ export class BedrockAgentCoreGatewayTargetHandler {
   /**
    * Handle Lambda invocation from Bedrock AgentCore Gateway
    */
-  async handleEvent(
+  async handle(
     event: Record<string, unknown>,
     context: Context
   ): Promise<unknown> {
     // Extract tool metadata from context
-    const toolName =
-      context.clientContext?.Custom?.["bedrockagentcoreToolName"];
+    const clientContext = context.clientContext as unknown as Record<string, unknown> | undefined;
+    const custom = clientContext?.["custom"] as Record<string, unknown> | undefined;
+    const gatewayToolName = custom?.["bedrockAgentCoreToolName"] as string | undefined;
 
-    if (!toolName) {
-      throw new Error("Missing bedrockagentcoreToolName in context");
+    if (!gatewayToolName) {
+      throw new Error("Missing bedrockAgentCoreToolName in context");
     }
 
+    // Gateway names the tools like <target name>___<tool name>
+    const delimiter = "___";
+    const delimiterIndex = gatewayToolName.indexOf(delimiter);
+    if (delimiterIndex === -1) {
+      throw new Error(`Invalid gateway tool name format: ${gatewayToolName}`);
+    }
+    const toolName = gatewayToolName.substring(
+      delimiterIndex + delimiter.length
+    );
+
     // Create JSON-RPC request from gateway event
     const jsonRpcRequest: JSONRPCRequest = {
       jsonrpc: "2.0",

src/typescript/src/handlers/handlers.test.ts

@@ -23,6 +23,7 @@ import { RequestHandler } from "./requestHandler.js";
 class MockRequestHandler implements RequestHandler {
   private responses: Map<string, JSONRPCResponse | JSONRPCError> = new Map();
   private shouldThrow = false;
+  private lastCall: JSONRPCRequest | null = null;
 
   setResponse(method: string, response: JSONRPCResponse | JSONRPCError) {
     this.responses.set(method, response);
@@ -32,10 +33,16 @@ class MockRequestHandler implements RequestHandler {
     this.shouldThrow = shouldThrow;
   }
 
+  getLastCall(): JSONRPCRequest | null {
+    return this.lastCall;
+  }
+
   async handleRequest(
     request: JSONRPCRequest,
     _context: Context
   ): Promise<JSONRPCResponse | JSONRPCError> {
+    this.lastCall = request;
+
     if (this.shouldThrow) {
       throw new Error("Mock handler error");
     }
@@ -910,11 +917,11 @@ describe("BedrockAgentCoreGatewayTargetHandler", () => {
     handler = new BedrockAgentCoreGatewayTargetHandler(mockRequestHandler);
     mockContext = {
       clientContext: {
-        Custom: {
-          bedrockagentcoreToolName: "test_tool",
+        custom: {
+          bedrockAgentCoreToolName: "target___test_tool",
         },
       },
-    } as Context;
+    } as unknown as Context;
   });
 
   it("should handle valid tool invocation", async () => {
@@ -926,18 +933,63 @@ describe("BedrockAgentCoreGatewayTargetHandler", () => {
     mockRequestHandler.setResponse("tools/call", expectedResponse);
 
     const event = { param1: "value1", param2: "value2" };
-    const result = await handler.handleEvent(event, mockContext);
+    const result = await handler.handle(event, mockContext);
 
     expect(result).toEqual({ message: "Tool executed successfully" });
+
+    // Verify the extracted tool name is everything after the first delimiter
+    const lastCall = mockRequestHandler.getLastCall();
+    expect(lastCall?.params?.name).toBe("test_tool");
   });
 
   it("should throw error when tool name is missing", async () => {
-    const contextWithoutTool = { clientContext: { Custom: {} } } as Context;
+    const contextWithoutTool = { clientContext: { custom: {} } } as unknown as Context;
+    const event = { param1: "value1" };
+
+    await expect(handler.handle(event, contextWithoutTool)).rejects.toThrow(
+      "Missing bedrockAgentCoreToolName in context"
+    );
+  });
+
+  it("should throw error when tool name format is invalid", async () => {
+    const contextWithInvalidFormat = {
+      clientContext: {
+        custom: {
+          bedrockAgentCoreToolName: "invalid_format",
+        },
+      },
+    } as unknown as Context;
     const event = { param1: "value1" };
 
     await expect(
-      handler.handleEvent(event, contextWithoutTool)
-    ).rejects.toThrow("Missing bedrockagentcoreToolName in context");
+      handler.handle(event, contextWithInvalidFormat)
+    ).rejects.toThrow("Invalid gateway tool name format: invalid_format");
+  });
+
+  it("should handle tool name with multiple delimiters", async () => {
+    const contextWithMultipleDelimiters = {
+      clientContext: {
+        custom: {
+          bedrockAgentCoreToolName: "target___test___tool",
+        },
+      },
+    } as unknown as Context;
+
+    const expectedResponse: JSONRPCResponse = {
+      jsonrpc: "2.0",
+      result: { message: "Tool executed successfully" },
+      id: 1,
+    };
+    mockRequestHandler.setResponse("tools/call", expectedResponse);
+
+    const event = { param1: "value1" };
+    const result = await handler.handle(event, contextWithMultipleDelimiters);
+
+    expect(result).toEqual({ message: "Tool executed successfully" });
+
+    // Verify the extracted tool name is everything after the first delimiter
+    const lastCall = mockRequestHandler.getLastCall();
+    expect(lastCall?.params?.name).toBe("test___tool");
   });
 
   it("should throw error when request handler returns error", async () => {
@@ -949,7 +1001,7 @@ describe("BedrockAgentCoreGatewayTargetHandler", () => {
     mockRequestHandler.setResponse("tools/call", errorResponse);
 
     const event = { param1: "value1" };
-    await expect(handler.handleEvent(event, mockContext)).rejects.toThrow(
+    await expect(handler.handle(event, mockContext)).rejects.toThrow(
       "Method not found"
     );
   });