From 3d60f9b899da2e0a0df93c5dd909baf9bee85eef Mon Sep 17 00:00:00 2001 From: Tejas Kashinath Date: Wed, 25 Feb 2026 14:17:45 -0500 Subject: [PATCH] feat: skip MCP client generation for VPC agents VPC agents may not have public internet access, so Exa AI MCP server should not be included in generated code. Adds isVpc flag to template rendering and conditionally excludes mcp_client/ directory and MCP imports in main.py for all frameworks. --- .../assets.snapshot.test.ts.snap | 57 +++++++++++++++++++ src/assets/python/autogen/base/main.py | 8 +++ src/assets/python/googleadk/base/main.py | 8 +++ .../python/langchain_langgraph/base/main.py | 8 +++ src/assets/python/openaiagents/base/main.py | 21 +++++++ src/assets/python/strands/base/main.py | 12 ++++ .../agent/generate/schema-mapper.ts | 1 + src/cli/templates/BaseRenderer.ts | 5 +- .../templates/__tests__/BaseRenderer.test.ts | 3 +- src/cli/templates/render.ts | 9 ++- src/cli/templates/types.ts | 2 + 11 files changed, 130 insertions(+), 4 deletions(-) diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index 4212766d..73589d8f 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -904,7 +904,9 @@ from autogen_agentchat.agents import AssistantAgent from autogen_core.tools import FunctionTool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_tools +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -928,14 +930,20 @@ tools = [add_numbers_tool] async def invoke(payload, context): log.info("Invoking Agent.....") +{{#unless isVpc}} # Get MCP Tools mcp_tools = await get_streamable_http_mcp_tools() +{{/unless}} # Define an AssistantAgent with the model and tools agent = AssistantAgent( name="{{ name }}", model_client=load_model(), +{{#unless isVpc}} tools=tools + mcp_tools, +{{else}} + tools=tools, +{{/unless}} system_message="You are a helpful assistant. Use tools when appropriate.", ) @@ -1584,7 +1592,9 @@ from google.adk.sessions import InMemorySessionService from google.genai import types from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -1601,8 +1611,10 @@ def add_numbers(a: int, b: int) -> int: return a + b +{{#unless isVpc}} # Get MCP Toolset mcp_toolset = [get_streamable_http_mcp_client()] +{{/unless}} _credentials_loaded = False @@ -1619,7 +1631,11 @@ agent = Agent( name="{{ name }}", description="Agent to answer questions", instruction="I can answer your questions using the knowledge I have!", +{{#unless isVpc}} tools=mcp_toolset + [add_numbers], +{{else}} + tools=[add_numbers], +{{/unless}} ) @@ -1864,7 +1880,9 @@ from langgraph.prebuilt import create_react_agent from langchain.tools import tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -1893,14 +1911,20 @@ tools = [add_numbers] async def invoke(payload, context): log.info("Invoking Agent.....") +{{#unless isVpc}} # Get MCP Client mcp_client = get_streamable_http_mcp_client() # Load MCP Tools mcp_tools = await mcp_client.get_tools() +{{/unless}} # Define the agent using create_react_agent +{{#unless isVpc}} graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools) +{{else}} + graph = create_react_agent(get_or_create_model(), tools=tools) +{{/unless}} # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") @@ -2210,13 +2234,17 @@ exports[`Assets Directory Snapshots > Python framework assets > python/python/op from agents import Agent, Runner, function_tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger +{{#unless isVpc}} # Get MCP Server mcp_server = get_streamable_http_mcp_client() +{{/unless}} _credentials_loaded = False @@ -2234,6 +2262,7 @@ def add_numbers(a: int, b: int) -> int: return a + b +{{#unless isVpc}} # Define the agent execution async def main(query): ensure_credentials_loaded() @@ -2251,6 +2280,22 @@ async def main(query): except Exception as e: log.error(f"Error during agent execution: {e}", exc_info=True) raise e +{{else}} +# Define the agent execution +async def main(query): + ensure_credentials_loaded() + try: + agent = Agent( + name="{{ name }}", + model="gpt-4.1", + tools=[add_numbers] + ) + result = await Runner.run(agent, query) + return result + except Exception as e: + log.error(f"Error during agent execution: {e}", exc_info=True) + raise e +{{/unless}} @app.entrypoint @@ -2455,7 +2500,9 @@ exports[`Assets Directory Snapshots > Python framework assets > python/python/st "from strands import Agent, tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} {{#if hasMemory}} from memory.session import get_memory_session_manager {{/if}} @@ -2463,8 +2510,10 @@ from memory.session import get_memory_session_manager app = BedrockAgentCoreApp() log = app.logger +{{#unless isVpc}} # Define a Streamable HTTP MCP Client mcp_client = get_streamable_http_mcp_client() +{{/unless}} # Define a collection of tools used by the model tools = [] @@ -2490,7 +2539,11 @@ def agent_factory(): system_prompt=""" You are a helpful assistant. Use tools when appropriate. """, +{{#unless isVpc}} tools=tools+[mcp_client] +{{else}} + tools=tools +{{/unless}} ) return cache[key] return get_or_create_agent @@ -2506,7 +2559,11 @@ def get_or_create_agent(): system_prompt=""" You are a helpful assistant. Use tools when appropriate. """, +{{#unless isVpc}} tools=tools+[mcp_client] +{{else}} + tools=tools +{{/unless}} ) return _agent {{/if}} diff --git a/src/assets/python/autogen/base/main.py b/src/assets/python/autogen/base/main.py index 43789e63..bf162464 100644 --- a/src/assets/python/autogen/base/main.py +++ b/src/assets/python/autogen/base/main.py @@ -3,7 +3,9 @@ from autogen_core.tools import FunctionTool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_tools +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -27,14 +29,20 @@ def add_numbers(a: int, b: int) -> int: async def invoke(payload, context): log.info("Invoking Agent.....") +{{#unless isVpc}} # Get MCP Tools mcp_tools = await get_streamable_http_mcp_tools() +{{/unless}} # Define an AssistantAgent with the model and tools agent = AssistantAgent( name="{{ name }}", model_client=load_model(), +{{#unless isVpc}} tools=tools + mcp_tools, +{{else}} + tools=tools, +{{/unless}} system_message="You are a helpful assistant. Use tools when appropriate.", ) diff --git a/src/assets/python/googleadk/base/main.py b/src/assets/python/googleadk/base/main.py index 2e89f01a..b5fc0882 100644 --- a/src/assets/python/googleadk/base/main.py +++ b/src/assets/python/googleadk/base/main.py @@ -5,7 +5,9 @@ from google.genai import types from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -22,8 +24,10 @@ def add_numbers(a: int, b: int) -> int: return a + b +{{#unless isVpc}} # Get MCP Toolset mcp_toolset = [get_streamable_http_mcp_client()] +{{/unless}} _credentials_loaded = False @@ -40,7 +44,11 @@ def ensure_credentials_loaded(): name="{{ name }}", description="Agent to answer questions", instruction="I can answer your questions using the knowledge I have!", +{{#unless isVpc}} tools=mcp_toolset + [add_numbers], +{{else}} + tools=[add_numbers], +{{/unless}} ) diff --git a/src/assets/python/langchain_langgraph/base/main.py b/src/assets/python/langchain_langgraph/base/main.py index 88bfe2d8..54764a6c 100644 --- a/src/assets/python/langchain_langgraph/base/main.py +++ b/src/assets/python/langchain_langgraph/base/main.py @@ -4,7 +4,9 @@ from langchain.tools import tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger @@ -33,14 +35,20 @@ def add_numbers(a: int, b: int) -> int: async def invoke(payload, context): log.info("Invoking Agent.....") +{{#unless isVpc}} # Get MCP Client mcp_client = get_streamable_http_mcp_client() # Load MCP Tools mcp_tools = await mcp_client.get_tools() +{{/unless}} # Define the agent using create_react_agent +{{#unless isVpc}} graph = create_react_agent(get_or_create_model(), tools=mcp_tools + tools) +{{else}} + graph = create_react_agent(get_or_create_model(), tools=tools) +{{/unless}} # Process the user prompt prompt = payload.get("prompt", "What can you help me with?") diff --git a/src/assets/python/openaiagents/base/main.py b/src/assets/python/openaiagents/base/main.py index d75f6002..1a2565df 100644 --- a/src/assets/python/openaiagents/base/main.py +++ b/src/assets/python/openaiagents/base/main.py @@ -2,13 +2,17 @@ from agents import Agent, Runner, function_tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} app = BedrockAgentCoreApp() log = app.logger +{{#unless isVpc}} # Get MCP Server mcp_server = get_streamable_http_mcp_client() +{{/unless}} _credentials_loaded = False @@ -26,6 +30,7 @@ def add_numbers(a: int, b: int) -> int: return a + b +{{#unless isVpc}} # Define the agent execution async def main(query): ensure_credentials_loaded() @@ -43,6 +48,22 @@ async def main(query): except Exception as e: log.error(f"Error during agent execution: {e}", exc_info=True) raise e +{{else}} +# Define the agent execution +async def main(query): + ensure_credentials_loaded() + try: + agent = Agent( + name="{{ name }}", + model="gpt-4.1", + tools=[add_numbers] + ) + result = await Runner.run(agent, query) + return result + except Exception as e: + log.error(f"Error during agent execution: {e}", exc_info=True) + raise e +{{/unless}} @app.entrypoint diff --git a/src/assets/python/strands/base/main.py b/src/assets/python/strands/base/main.py index a5557405..a89a6e2e 100644 --- a/src/assets/python/strands/base/main.py +++ b/src/assets/python/strands/base/main.py @@ -1,7 +1,9 @@ from strands import Agent, tool from bedrock_agentcore.runtime import BedrockAgentCoreApp from model.load import load_model +{{#unless isVpc}} from mcp_client.client import get_streamable_http_mcp_client +{{/unless}} {{#if hasMemory}} from memory.session import get_memory_session_manager {{/if}} @@ -9,8 +11,10 @@ app = BedrockAgentCoreApp() log = app.logger +{{#unless isVpc}} # Define a Streamable HTTP MCP Client mcp_client = get_streamable_http_mcp_client() +{{/unless}} # Define a collection of tools used by the model tools = [] @@ -36,7 +40,11 @@ def get_or_create_agent(session_id, user_id): system_prompt=""" You are a helpful assistant. Use tools when appropriate. """, +{{#unless isVpc}} tools=tools+[mcp_client] +{{else}} + tools=tools +{{/unless}} ) return cache[key] return get_or_create_agent @@ -52,7 +60,11 @@ def get_or_create_agent(): system_prompt=""" You are a helpful assistant. Use tools when appropriate. """, +{{#unless isVpc}} tools=tools+[mcp_client] +{{else}} + tools=tools +{{/unless}} ) return _agent {{/if}} diff --git a/src/cli/operations/agent/generate/schema-mapper.ts b/src/cli/operations/agent/generate/schema-mapper.ts index 7d860942..37b19cdc 100644 --- a/src/cli/operations/agent/generate/schema-mapper.ts +++ b/src/cli/operations/agent/generate/schema-mapper.ts @@ -202,6 +202,7 @@ export function mapGenerateConfigToRenderConfig( modelProvider: config.modelProvider, hasMemory: config.memory !== 'none', hasIdentity: identityProviders.length > 0, + isVpc: config.networkMode === 'VPC', buildType: config.buildType, memoryProviders: mapMemoryOptionToMemoryProviders(config.memory, config.projectName), identityProviders, diff --git a/src/cli/templates/BaseRenderer.ts b/src/cli/templates/BaseRenderer.ts index cda526a5..fb7889b2 100644 --- a/src/cli/templates/BaseRenderer.ts +++ b/src/cli/templates/BaseRenderer.ts @@ -45,9 +45,10 @@ export abstract class BaseRenderer { hasMcp: false, // MCP is configured separately }; - // Always render base template + // Always render base template (skip mcp_client for VPC - no public internet for external MCP servers) const baseDir = path.join(templateDir, 'base'); - await copyAndRenderDir(baseDir, projectDir, templateData); + const skipDirs = this.config.isVpc ? new Set(['mcp_client']) : undefined; + await copyAndRenderDir(baseDir, projectDir, templateData, skipDirs); // Render capability templates based on config // Only render if the capability directory exists (not all SDKs have all capabilities) diff --git a/src/cli/templates/__tests__/BaseRenderer.test.ts b/src/cli/templates/__tests__/BaseRenderer.test.ts index a017ba55..6ed84688 100644 --- a/src/cli/templates/__tests__/BaseRenderer.test.ts +++ b/src/cli/templates/__tests__/BaseRenderer.test.ts @@ -56,7 +56,8 @@ describe('BaseRenderer', () => { expect(mockCopyAndRenderDir).toHaveBeenCalledWith( '/templates/python/strands/base', '/output/app/MyAgent', - expect.objectContaining({ projectName: 'MyAgent', Name: 'MyAgent', hasMcp: false }) + expect.objectContaining({ projectName: 'MyAgent', Name: 'MyAgent', hasMcp: false }), + undefined ); }); diff --git a/src/cli/templates/render.ts b/src/cli/templates/render.ts index 166c90a3..d072becb 100644 --- a/src/cli/templates/render.ts +++ b/src/cli/templates/render.ts @@ -43,8 +43,14 @@ export async function copyDir(src: string, dest: string): Promise { /** * Recursively copies a directory, rendering Handlebars templates. + * @param skipDirs - Optional set of directory names to skip (top-level only) */ -export async function copyAndRenderDir(src: string, dest: string, data: T): Promise { +export async function copyAndRenderDir( + src: string, + dest: string, + data: T, + skipDirs?: Set +): Promise { await fs.mkdir(dest, { recursive: true }); const entries = await fs.readdir(src, { withFileTypes: true }); @@ -54,6 +60,7 @@ export async function copyAndRenderDir(src: string, dest: stri const destPath = path.join(dest, destName); if (entry.isDirectory()) { + if (skipDirs?.has(entry.name)) continue; await copyAndRenderDir(srcPath, destPath, data); } else { const content = await fs.readFile(srcPath, 'utf-8'); diff --git a/src/cli/templates/types.ts b/src/cli/templates/types.ts index 37dded4e..1dd014b2 100644 --- a/src/cli/templates/types.ts +++ b/src/cli/templates/types.ts @@ -29,6 +29,8 @@ export interface AgentRenderConfig { modelProvider: ModelProvider; hasMemory: boolean; hasIdentity: boolean; + /** Whether the agent uses VPC network mode (disables public MCP endpoints) */ + isVpc: boolean; /** Build type: CodeZip (default) or Container */ buildType?: BuildType; /** Memory providers for template rendering */