@aws/nx-plugin 1.0.0-rc.7 → 1.0.0-rc.70

package/src/py/agent/a2a-connection/generator.js.map

@@ -1 +1 @@
-{"version":3,"file":"generator.js","sourceRoot":"","sources":["../../../../../../../packages/nx-plugin/src/py/agent/a2a-connection/generator.ts"],"names":[],"mappings":";;;AAAA;;;GAGG;AACH,uCAQoB;AACpB,uFAO0D;AAC1D,4CAI4B;AAC5B,kDAA6D;AAC7D,oDAAyE;AACzE,gDAAiD;AACjD,0CAK2B;AAC3B,0CAG2B;AAG3B,qDAAqD;AACrD,MAAM,EAAE,GAAG,CAAC,OAAe,EAAE,EAAE,CAAC,oBAAoB,OAAO,EAAE,CAAC;AAEjD,QAAA,sCAAsC,GACjD,IAAA,qBAAgB,EAAC,UAAU,CAAC,CAAC;AAExB,MAAM,6BAA6B,GAAG,KAAK,EAChD,IAAU,EACV,OAA4C,EAChB,EAAE;IAC9B,MAAM,aAAa,GAAG,IAAA,wCAAmC,EACvD,IAAI,EACJ,OAAO,CAAC,aAAa,CACtB,CAAC;IACF,MAAM,aAAa,GAAG,IAAA,wCAAmC,EACvD,IAAI,EACJ,OAAO,CAAC,aAAa,CACtB,CAAC;IAEF,MAAM,cAAc,GAAG,OAAO,CAAC,eAAe,CAAC;IAC/C,MAAM,oBAAoB,GAAG,OAAO,CAAC,eAAe,CAAC;IAErD,IAAI,CAAC,cAAc,IAAI,CAAC,oBAAoB,EAAE,CAAC;QAC7C,MAAM,IAAI,KAAK,CACb,2FAA2F,CAC5F,CAAC;IACJ,CAAC;IAED,IAAI,CAAC,oBAAoB,CAAC,QAAQ,IAAI,EAAE,CAAC,CAAC,WAAW,EAAE,KAAK,KAAK,EAAE,CAAC;QAClE,MAAM,IAAI,KAAK,CACb,iBAAiB,oBAAoB,CAAC,IAAI,cAAc,oBAAoB,CAAC,QAAQ,IAAI,MAAM,wDAAwD,CACxJ,CAAC;IACJ,CAAC;IAED,IACE,oBAAoB,CAAC,IAAI;QACzB,oBAAoB,CAAC,IAAI,CAAC,WAAW,EAAE,KAAK,KAAK,EACjD,CAAC;QACD,MAAM,IAAI,KAAK,CACb,yEAAyE,oBAAoB,CAAC,IAAI,WAAW,oBAAoB,CAAC,IAAI,mBAAmB,CAC1J,CAAC;IACJ,CAAC;IAED,MAAM,wBAAwB,GAAG,oBAAoB,CAAC,IAAI,IAAI,OAAO,CAAC;IACtE,MAAM,oBAAoB,GAAG,oBAAoB,CAAC,EAAY,CAAC;IAC/D,MAAM,oBAAoB,GAAG,IAAA,iBAAS,EAAC,oBAAoB,CAAC,CAAC;IAC7D,MAAM,eAAe,GAAG,oBAAoB,CAAC,IAAI,IAAI,IAAI,CAAC;IAE1D,wEAAwE;IACxE,uDAAuD;IACvD,MAAM,IAAA,qDAAkC,EAAC,IAAI,CAAC,CAAC;IAC/C,IAAA,sCAAmB,EAAC,IAAI,EAAE,MAAM,CAAC,CAAC;IAClC,IAAA,sCAAmB,EAAC,IAAI,EAAE,KAAK,CAAC,CAAC;IAEjC,MAAM,yBAAyB,GAAG,IAAA,qDAAkC,EAAC,IAAI,CAAC,CAAC;IAC3E,MAAM,yBAAyB,GAAG,IAAA,qDAAkC,EAAC,IAAI,CAAC,CAAC;IAC3E,MAAM,0BAA0B,GAAG,IAAA,sDAAmC,EAAC,IAAI,CAAC,CAAC;IAE7E,oEAAoE;IACpE,IAAA,mCAA8B,EAAC,IAAI,EAAE,yBAAyB,EAAE;QAC9D,OAAO;QACP,OAAO;QACP,qBAAqB;KACtB,CAAC,CAAC;IAEH,iFAAiF;IACjF,MAAM,MAAM,GAAG,IAAA,0BAAiB,EAC9B,yBAAyB,EACzB,yBAAyB,EACzB,KAAK,CACN,CAAC;IACF,IAAA,sBAAa,EACX,IAAI,EACJ,IAAA,0BAAiB,EAAC,SAAS,EAAE,OAAO,EAAE,kBAAkB,EAAE,KAAK,CAAC,EAChE,MAAM,EACN;QACE,oBAAoB;QACpB,oBAAoB;QACpB,eAAe;QACf,yBAAyB;KAC1B,EACD,EAAE,iBAAiB,EAAE,0BAAiB,CAAC,YAAY,EAAE,CACtD,CAAC;IAEF,gCAAgC;IAChC,MAAM,WAAW,GAAG,IAAA,0BAAiB,EAAC,MAAM,EAAE,aAAa,CAAC,CAAC;IAC7D,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,WAAW,CAAC,EAAE,CAAC;QAC9B,IAAI,CAAC,KAAK,CAAC,WAAW,EAAE,EAAE,CAAC,CAAC;IAC9B,CAAC;IAED,sCAAsC;IACtC,MAAM,cAAc,GAAG,IAAA,0BAAiB,EACtC,yBAAyB,EACzB,yBAAyB,EACzB,aAAa,CACd,CAAC;IACF,MAAM,IAAA,oCAAiB,EACrB,IAAI,EACJ,cAAc,EACd,QAAQ,oBAAoB,SAAS,EACrC,GAAG,oBAAoB,QAAQ,CAChC,CAAC;IAEF,yEAAyE;IACzE,MAAM,cAAc,GAAG,IAAA,0BAAiB,EACtC,aAAa,CAAC,IAAI,EAClB,cAAc,CAAC,IAAI,IAAI,KAAK,CAC7B,CAAC;IACF,MAAM,aAAa,GAAG,IAAA,0BAAiB,EAAC,cAAc,EAAE,UAAU,CAAC,CAAC;IAEpE,IAAI,IAAI,CAAC,MAAM,CAAC,aAAa,CAAC,EAAE,CAAC;QAC/B,MAAM,eAAe,GAAG,GAAG,oBAAoB,QAAQ,CAAC;QACxD,MAAM,aAAa,GAAG,oBAAoB,CAAC;QAC3C,MAAM,QAAQ,GAAG,OAAO,oBAAoB,EAAE,CAAC;QAE/C,MAAM,oBAAoB,CACxB,IAAI,EACJ,aAAa,EACb,yBAAyB,EACzB,eAAe,CAChB,CAAC;QACF,MAAM,cAAc,CAAC,IAAI,EAAE,aAAa,EAAE,QAAQ,CAAC,CAAC;QACpD,MAAM,uBAAuB,CAC3B,IAAI,EACJ,aAAa,EACb,eAAe,EACf,aAAa,EACb,QAAQ,EACR,oBAAoB,CACrB,CAAC;IACJ,CAAC;IAED,6EAA6E;IAC7E,IAAA,sCAAiC,EAC/B,IAAI,EACJ,aAAa,CAAC,IAAI,EAClB,0BAA0B,CAC3B,CAAC;IAEF,0EAA0E;IAC1E,MAAM,SAAS,GAAG,cAAc,CAAC,IAAI,IAAI,OAAO,CAAC;IACjD,MAAM,oBAAoB,GAAG,GAAG,SAAS,cAAc,CAAC;IACxD,MAAM,0BAA0B,GAAG,GAAG,wBAAwB,cAAc,CAAC;IAE7E,IAAI,aAAa,CAAC,OAAO,EAAE,CAAC,oBAAoB,CAAC,EAAE,CAAC;QAClD,IAAA,sCAAiC,EAAC,aAAa,EAAE,oBAAoB,EAAE;YACrE,QAAQ,EAAE,CAAC,aAAa,CAAC,IAAI,CAAC;YAC9B,MAAM,EAAE,0BAA0B;SACnC,CAAC,CAAC;QACH,IAAA,mCAA0B,EAAC,IAAI,EAAE,aAAa,CAAC,IAAI,EAAE,aAAa,CAAC,CAAC;IACtE,CAAC;IAED,MAAM,IAAA,yCAA+B,EAAC,IAAI,EAAE;QAC1C,8CAAsC;KACvC,CAAC,CAAC;IAEH,MAAM,IAAA,6BAAoB,EAAC,IAAI,CAAC,CAAC;IACjC,OAAO,GAAG,EAAE;QACV,IAAA,4BAAmB,EAAC,IAAI,CAAC,CAAC;IAC5B,CAAC,CAAC;AACJ,CAAC,CAAC;AA1JW,QAAA,6BAA6B,iCA0JxC;AAEF;;;;GAIG;AACH,MAAM,oBAAoB,GAAG,KAAK,EAChC,IAAU,EACV,QAAgB,EAChB,yBAAiC,EACjC,eAAuB,EACR,EAAE;IACjB,MAAM,IAAA,iCAA2B,EAAC,IAAI,EAAE,QAAQ,EAAE,CAAC,MAAM,CAAC,EAAE,SAAS,CAAC,CAAC;IACvE,MAAM,IAAA,iCAA2B,EAC/B,IAAI,EACJ,QAAQ,EACR,CAAC,eAAe,CAAC,EACjB,yBAAyB,CAC1B,CAAC;AACJ,CAAC,CAAC;AAEF;;;;;GAKG;AACH,MAAM,cAAc,GAAG,KAAK,EAC1B,IAAU,EACV,QAAgB,EAChB,QAAgB,EACD,EAAE;IACjB,MAAM,IAAA,iBAAW,EACf,IAAI,EACJ,QAAQ,EACR,EAAE,CAAC;;2BAEoB,QAAQ;;iBAElB,QAAQ;;iDAEwB,QAAQ;;EAEvD,CAAC,CACA,CAAC;AACJ,CAAC,CAAC;AAEF;;;;;;;GAOG;AACH,MAAM,uBAAuB,GAAG,KAAK,EACnC,IAAU,EACV,QAAgB,EAChB,eAAuB,EACvB,aAAqB,EACrB,QAAgB,EAChB,oBAA4B,EACb,EAAE;IACjB,IAAI,MAAM,IAAA,iBAAW,EAAC,IAAI,EAAE,QAAQ,EAAE,KAAK,eAAe,WAAW,CAAC,EAAE,CAAC;QACvE,OAAO;IACT,CAAC;IAED,mEAAmE;IACnE,oEAAoE;IACpE,sEAAsE;IACtE,MAAM,SAAS,GAAG,GAAG,aAAa,MAAM,eAAe;;;UAG/C,QAAQ;+CAC6B,oBAAoB;qBAC9C,aAAa;CACjC,CAAC;IAEA,kEAAkE;IAClE,yEAAyE;IACzE,yCAAyC;IACzC,MAAM,IAAA,iBAAW,EACf,IAAI,EACJ,QAAQ,EACR,EAAE,CAAC;;;4BAGqB,eAAe;;MAErC,SAAS;YACH,CAAC,CACV,CAAC;AACJ,CAAC,CAAC;AAEF,kBAAe,qCAA6B,CAAC"}
+{"version":3,"sources":["../../../../../../../packages/nx-plugin/src/py/agent/a2a-connection/generator.ts"],"sourcesContent":["/**\n * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n * SPDX-License-Identifier: Apache-2.0\n */\nimport {\n  type GeneratorCallback,\n  generateFiles,\n  joinPathFragments,\n  OverwriteStrategy,\n  type Tree,\n  updateProjectConfiguration,\n} from '@nx/devkit';\nimport { addPyDependencies } from '../../../utils/add-dependencies';\nimport {\n  AGENT_CONNECTION_PY_DEPENDENCIES,\n  type AgentFramework,\n  addPythonCoreClient,\n  addPythonReExport,\n  ensurePythonAgentConnectionProject,\n  getPythonAgentConnectionModuleName,\n  getPythonAgentConnectionProject,\n  getPythonAgentConnectionProjectDir,\n  PY_CLIENT_NAMING,\n  resolveAgentFramework,\n} from '../../../utils/agent-connection/agent-connection';\nimport {\n  addPythonDestructuredImport,\n  applyGritQL,\n  matchGritQL,\n} from '../../../utils/ast';\nimport {\n  declareDependencies,\n  ownedElsewhere,\n} from '../../../utils/declared-dependencies';\nimport { formatFilesInSubtree } from '../../../utils/format';\nimport { installDependencies } from '../../../utils/install';\nimport { addGeneratorMetricsIfApplicable } from '../../../utils/metrics';\nimport { snakeCase } from '../../../utils/names';\nimport {\n  addComponentGeneratorMetadata,\n  addDependencyToTargetIfNotPresent,\n  getGeneratorInfo,\n  type NxGeneratorInfo,\n  readProjectConfigurationUnqualified,\n} from '../../../utils/nx';\nimport { toProjectRelativePath } from '../../../utils/paths';\nimport { addWorkspaceDependencyToPyProject } from '../../../utils/py';\nimport type { PyAgentA2aConnectionGeneratorSchema } from './schema';\n\n/** Prefix a GritQL pattern with `language python` */\nconst py = (pattern: string) => `language python\\n${pattern}`;\n\n/** The module each framework imports the `tool` decorator from. */\nconst TOOL_IMPORT_MODULE: Record<AgentFramework, string> = {\n  strands: 'strands',\n  langchain: 'langchain_core.tools',\n};\n\n/** The GritQL pattern matching each framework's agent constructor. */\nconst AGENT_CONSTRUCTOR: Record<AgentFramework, string> = {\n  strands: 'yield Agent($_)',\n  langchain: 'create_agent($_)',\n};\n\n/** The metadata this generator records, which its predicates read. */\nexport interface PyAgentA2aConnectionMetadata {\n  readonly framework: AgentFramework;\n}\n\n// Each framework's Layer-2 client needs its own A2A transport: Strands wraps\n// strands' `A2AAgent` (strands-agents[a2a]); LangChain drives the a2a SDK\n// directly (a2a-sdk) and must not pull Strands in.\nexport const DEPENDENCIES = declareDependencies<PyAgentA2aConnectionMetadata>()(\n  {\n    py: [\n      { name: 'boto3' },\n      { name: 'httpx' },\n      { name: 'strands-agents[a2a]', when: (m) => m.framework === 'strands' },\n      { name: 'a2a-sdk', when: (m) => m.framework === 'langchain' },\n      ...ownedElsewhere(AGENT_CONNECTION_PY_DEPENDENCIES),\n    ],\n  },\n);\n\nexport const PY_AGENT_A2A_CONNECTION_GENERATOR_INFO: NxGeneratorInfo =\n  getGeneratorInfo(import.meta.filename);\n\nexport const pyAgentA2aConnectionGenerator = async (\n  tree: Tree,\n  options: PyAgentA2aConnectionGeneratorSchema,\n): Promise<GeneratorCallback> => {\n  const sourceProject = readProjectConfigurationUnqualified(\n    tree,\n    options.sourceProject,\n  );\n  const targetProject = readProjectConfigurationUnqualified(\n    tree,\n    options.targetProject,\n  );\n\n  const agentComponent = options.sourceComponent;\n  const targetAgentComponent = options.targetComponent;\n\n  if (!agentComponent || !targetAgentComponent) {\n    throw new Error(\n      'Both sourceComponent and targetComponent must be provided for py#agent -> a2a connections',\n    );\n  }\n\n  if ((targetAgentComponent.protocol ?? '').toLowerCase() !== 'a2a') {\n    throw new Error(\n      `Target agent '${targetAgentComponent.name}' uses the ${targetAgentComponent.protocol ?? 'http'} protocol — only A2A agents can be connected as tools.`,\n    );\n  }\n\n  if (\n    targetAgentComponent.auth &&\n    targetAgentComponent.auth.toLowerCase() !== 'iam'\n  ) {\n    throw new Error(\n      `A2A agent connection currently only supports IAM authentication, but '${targetAgentComponent.name}' uses '${targetAgentComponent.auth}' authentication.`,\n    );\n  }\n\n  const targetAgentComponentName = targetAgentComponent.name ?? 'agent';\n  const targetAgentClassName = targetAgentComponent.rc as string;\n  const targetAgentSnakeCase = snakeCase(targetAgentClassName);\n  const targetAgentPort = targetAgentComponent.port ?? 9000;\n\n  // The source agent's framework drives everything that varies: the client\n  // naming + app template (shared with the mcp/gateway generators via\n  // PY_CLIENT_NAMING), the A2A transport dep (see DEPENDENCIES), and the\n  // agent.py transform below. Each framework wraps its own A2A Layer-2 client\n  // (Strands' A2AAgent vs an a2a-sdk-only client) over the shared,\n  // framework-agnostic A2A client config. Keyed on the framework, not a\n  // boolean, so a third framework stays additive.\n  const framework = resolveAgentFramework(agentComponent.framework);\n  const naming = PY_CLIENT_NAMING[framework];\n\n  // Recorded below and read by the declaration's predicates, so the packages\n  // added here are exactly the ones the version sync will own.\n  const metadata: PyAgentA2aConnectionMetadata = { framework };\n\n  // 1. Ensure the shared Python agent-connection project exists + has the\n  //    A2A core client (framework Layer-2) and its shared SigV4 auth helper.\n  await ensurePythonAgentConnectionProject(tree, DEPENDENCIES);\n  await addPythonCoreClient(tree, 'a2a', DEPENDENCIES, framework);\n\n  const agentConnectionProjectDir = getPythonAgentConnectionProjectDir(tree);\n  const agentConnectionModuleName = getPythonAgentConnectionModuleName(tree);\n\n  // A2A client config + shared auth helper deps, plus the framework's A2A\n  // transport dep. These belong to the shared agent-connection project, which\n  // holds the client.\n  addPyDependencies(tree, DEPENDENCIES, {\n    metadata,\n    projectRoot: agentConnectionProjectDir,\n  });\n\n  // 2. Generate the per-connection client into the shared agent-connection project\n  const appDir = joinPathFragments(\n    agentConnectionProjectDir,\n    agentConnectionModuleName,\n    'app',\n  );\n  generateFiles(\n    tree,\n    joinPathFragments(import.meta.dirname, 'files', naming.appTemplateSubdir),\n    appDir,\n    {\n      targetAgentSnakeCase,\n      targetAgentClassName,\n      targetAgentPort,\n      agentConnectionModuleName,\n    },\n    { overwriteStrategy: OverwriteStrategy.KeepExisting },\n  );\n\n  // Ensure app/__init__.py exists\n  const appInitPath = joinPathFragments(appDir, '__init__.py');\n  if (!tree.exists(appInitPath)) {\n    tree.write(appInitPath, '');\n  }\n\n  // Add re-export to module __init__.py\n  const moduleInitPath = joinPathFragments(\n    agentConnectionProjectDir,\n    agentConnectionModuleName,\n    '__init__.py',\n  );\n  await addPythonReExport(\n    tree,\n    moduleInitPath,\n    `.app.${targetAgentSnakeCase}_client_${naming.clientModuleSuffix}`,\n    `${targetAgentClassName}Client${naming.clientClassSuffix}`,\n  );\n\n  // 3. Transform agent.py to add the A2A client import + wrap it as a tool\n  const agentSourceDir = joinPathFragments(\n    sourceProject.root,\n    agentComponent.path ?? 'src',\n  );\n  const agentFilePath = joinPathFragments(agentSourceDir, 'agent.py');\n\n  if (tree.exists(agentFilePath)) {\n    const clientClassName = `${targetAgentClassName}Client${naming.clientClassSuffix}`;\n    const clientVarName = targetAgentSnakeCase;\n    const toolName = `ask_${targetAgentSnakeCase}`;\n\n    await addImportToAgentFile(\n      tree,\n      agentFilePath,\n      agentConnectionModuleName,\n      clientClassName,\n      framework,\n    );\n    await addToolToAgent(tree, agentFilePath, toolName, framework);\n    await addClientToolToGetAgent(\n      tree,\n      agentFilePath,\n      clientClassName,\n      clientVarName,\n      toolName,\n      targetAgentClassName,\n      framework,\n    );\n  }\n\n  // 4. Add workspace dependency from agent project to agent-connection project\n  addWorkspaceDependencyToPyProject(\n    tree,\n    sourceProject,\n    getPythonAgentConnectionProject(tree),\n  );\n\n  // 5. Set up dev target dependencies — chain onto the target agent\n  const agentName = agentComponent.name ?? 'agent';\n  const devTargetName = `${agentName}-dev`;\n  const targetDevTargetName = `${targetAgentComponentName}-dev`;\n\n  if (sourceProject.targets?.[devTargetName]) {\n    addDependencyToTargetIfNotPresent(sourceProject, devTargetName, {\n      projects: [targetProject.name],\n      target: targetDevTargetName,\n    });\n    updateProjectConfiguration(tree, sourceProject.name, sourceProject);\n  }\n\n  // Recorded so the version sync knows this connection's dependencies are ours.\n  addComponentGeneratorMetadata(\n    tree,\n    sourceProject.name,\n    PY_AGENT_A2A_CONNECTION_GENERATOR_INFO,\n    toProjectRelativePath(sourceProject, agentFilePath),\n    targetAgentClassName,\n    // `sourcePath` names the source component this connection is made from, so\n    // the pair is identifiable rather than just the two projects.\n    { ...metadata, sourcePath: options.sourceComponent?.path },\n  );\n\n  await addGeneratorMetricsIfApplicable(tree, [\n    PY_AGENT_A2A_CONNECTION_GENERATOR_INFO,\n  ]);\n\n  await formatFilesInSubtree(tree);\n  return () =>\n    installDependencies(tree, options.preferInstallDependencies, {\n      languages: ['typescript'],\n    });\n};\n\n/**\n * Add imports to agent.py:\n *  - `tool` (from `strands`, or `langchain_core.tools` for LangChain agents)\n *  - the per-connection client class from the shared agent-connection module\n */\nconst addImportToAgentFile = async (\n  tree: Tree,\n  filePath: string,\n  agentConnectionModuleName: string,\n  clientClassName: string,\n  framework: AgentFramework,\n): Promise<void> => {\n  await addPythonDestructuredImport(\n    tree,\n    filePath,\n    ['tool'],\n    TOOL_IMPORT_MODULE[framework],\n  );\n  await addPythonDestructuredImport(\n    tree,\n    filePath,\n    [clientClassName],\n    agentConnectionModuleName,\n  );\n};\n\n/**\n * Add the A2A tool to the Agent's tools array using GritQL.\n *\n * Scoped to the framework's agent constructor so unrelated `tools=` keyword\n * arguments elsewhere in the file are not touched.\n */\nconst addToolToAgent = async (\n  tree: Tree,\n  filePath: string,\n  toolName: string,\n  framework: AgentFramework,\n): Promise<void> => {\n  const scope = AGENT_CONSTRUCTOR[framework];\n  await applyGritQL(\n    tree,\n    filePath,\n    py(`\\`tools=$old\\` where {\n  $old <: within \\`${scope}\\`,\n  $old <: not contains \\`${toolName}\\`,\n  if ($old <: \\`[]\\`) {\n    $old => \\`[${toolName}]\\`\n  } else {\n    $old <: \\`[$items]\\` where { $items += \\`, ${toolName}\\` }\n  }\n}`),\n  );\n};\n\n/**\n * Wrap the get_agent body so the remote A2A client is created and exposed\n * as a `@tool`-decorated closure before constructing the Agent.\n *\n * Unlike MCP clients, A2A clients aren't context managers — we don't need\n * a `with` block around them. Creation + tool definition happen at the\n * top of the function body, anchored on the framework's agent constructor.\n */\nconst addClientToolToGetAgent = async (\n  tree: Tree,\n  filePath: string,\n  clientClassName: string,\n  clientVarName: string,\n  toolName: string,\n  targetAgentClassName: string,\n  framework: AgentFramework,\n): Promise<void> => {\n  if (await matchGritQL(tree, filePath, py(`\\`${clientClassName}.create\\``))) {\n    return;\n  }\n\n  // Both frameworks' tools are @tool-decorated callables (the decorator's\n  // import differs, handled in addImportToAgentFile), defined inline in\n  // get_agent. The client is directly callable, returning the remote's reply.\n  const toolBlock = `${clientVarName} = ${clientClassName}.create()\n\n    @tool\n    def ${toolName}(prompt: str) -> str:\n        \"\"\"Delegate a question to the remote ${targetAgentClassName} A2A agent and return its reply.\"\"\"\n        return str(${clientVarName}(prompt))\n`;\n\n  // Prepend client creation + tool definition to the function body, anchored on\n  // the framework's agent constructor. This works whether or not a prior MCP\n  // block is present since we insert before any existing statements.\n  const anchor = AGENT_CONSTRUCTOR[framework];\n  await applyGritQL(\n    tree,\n    filePath,\n    py(`\\`def get_agent($params):\n    $body\\` where {\n  $body <: contains \\`${anchor}\\`,\n  $body <: not contains \\`${clientClassName}.create\\`\n} => \\`def get_agent($params):\n    ${toolBlock}\n    $body\\``),\n  );\n};\n\nexport default pyAgentA2aConnectionGenerator;\n"],"names":["generateFiles","joinPathFragments","OverwriteStrategy","updateProjectConfiguration","addPyDependencies","AGENT_CONNECTION_PY_DEPENDENCIES","addPythonCoreClient","addPythonReExport","ensurePythonAgentConnectionProject","getPythonAgentConnectionModuleName","getPythonAgentConnectionProject","getPythonAgentConnectionProjectDir","PY_CLIENT_NAMING","resolveAgentFramework","addPythonDestructuredImport","applyGritQL","matchGritQL","declareDependencies","ownedElsewhere","formatFilesInSubtree","installDependencies","addGeneratorMetricsIfApplicable","snakeCase","addComponentGeneratorMetadata","addDependencyToTargetIfNotPresent","getGeneratorInfo","readProjectConfigurationUnqualified","toProjectRelativePath","addWorkspaceDependencyToPyProject","py","pattern","TOOL_IMPORT_MODULE","strands","langchain","AGENT_CONSTRUCTOR","DEPENDENCIES","name","when","m","framework","PY_AGENT_A2A_CONNECTION_GENERATOR_INFO","filename","pyAgentA2aConnectionGenerator","tree","options","sourceProject","targetProject","agentComponent","sourceComponent","targetAgentComponent","targetComponent","Error","protocol","toLowerCase","auth","targetAgentComponentName","targetAgentClassName","rc","targetAgentSnakeCase","targetAgentPort","port","naming","metadata","agentConnectionProjectDir","agentConnectionModuleName","projectRoot","appDir","dirname","appTemplateSubdir","overwriteStrategy","KeepExisting","appInitPath","exists","write","moduleInitPath","clientModuleSuffix","clientClassSuffix","agentSourceDir","root","path","agentFilePath","clientClassName","clientVarName","toolName","addImportToAgentFile","addToolToAgent","addClientToolToGetAgent","agentName","devTargetName","targetDevTargetName","targets","projects","target","sourcePath","preferInstallDependencies","languages","filePath","scope","toolBlock","anchor"],"mappings":"AAAA;;;CAGC,GACD,SAEEA,aAAa,EACbC,iBAAiB,EACjBC,iBAAiB,EAEjBC,0BAA0B,QACrB,aAAa;AACpB,SAASC,iBAAiB,QAAQ,qCAAkC;AACpE,SACEC,gCAAgC,EAEhCC,mBAAmB,EACnBC,iBAAiB,EACjBC,kCAAkC,EAClCC,kCAAkC,EAClCC,+BAA+B,EAC/BC,kCAAkC,EAClCC,gBAAgB,EAChBC,qBAAqB,QAChB,sDAAmD;AAC1D,SACEC,2BAA2B,EAC3BC,WAAW,EACXC,WAAW,QACN,wBAAqB;AAC5B,SACEC,mBAAmB,EACnBC,cAAc,QACT,0CAAuC;AAC9C,SAASC,oBAAoB,QAAQ,2BAAwB;AAC7D,SAASC,mBAAmB,QAAQ,4BAAyB;AAC7D,SAASC,+BAA+B,QAAQ,4BAAyB;AACzE,SAASC,SAAS,QAAQ,0BAAuB;AACjD,SACEC,6BAA6B,EAC7BC,iCAAiC,EACjCC,gBAAgB,EAEhBC,mCAAmC,QAC9B,uBAAoB;AAC3B,SAASC,qBAAqB,QAAQ,0BAAuB;AAC7D,SAASC,iCAAiC,QAAQ,uBAAoB;AAGtE,mDAAmD,GACnD,MAAMC,KAAK,CAACC,UAAoB,CAAC,iBAAiB,EAAEA,SAAS;AAE7D,iEAAiE,GACjE,MAAMC,qBAAqD;IACzDC,SAAS;IACTC,WAAW;AACb;AAEA,oEAAoE,GACpE,MAAMC,oBAAoD;IACxDF,SAAS;IACTC,WAAW;AACb;AAOA,6EAA6E;AAC7E,0EAA0E;AAC1E,mDAAmD;AACnD,OAAO,MAAME,eAAelB,sBAC1B;IACEY,IAAI;QACF;YAAEO,MAAM;QAAQ;QAChB;YAAEA,MAAM;QAAQ;QAChB;YAAEA,MAAM;YAAuBC,MAAM,CAACC,IAAMA,EAAEC,SAAS,KAAK;QAAU;QACtE;YAAEH,MAAM;YAAWC,MAAM,CAACC,IAAMA,EAAEC,SAAS,KAAK;QAAY;WACzDrB,eAAeb;KACnB;AACH,GACA;AAEF,OAAO,MAAMmC,yCACXf,iBAAiB,YAAYgB,QAAQ,EAAE;AAEzC,OAAO,MAAMC,gCAAgC,OAC3CC,MACAC;IAEA,MAAMC,gBAAgBnB,oCACpBiB,MACAC,QAAQC,aAAa;IAEvB,MAAMC,gBAAgBpB,oCACpBiB,MACAC,QAAQE,aAAa;IAGvB,MAAMC,iBAAiBH,QAAQI,eAAe;IAC9C,MAAMC,uBAAuBL,QAAQM,eAAe;IAEpD,IAAI,CAACH,kBAAkB,CAACE,sBAAsB;QAC5C,MAAM,IAAIE,MACR;IAEJ;IAEA,IAAI,AAACF,CAAAA,qBAAqBG,QAAQ,IAAI,EAAC,EAAGC,WAAW,OAAO,OAAO;QACjE,MAAM,IAAIF,MACR,CAAC,cAAc,EAAEF,qBAAqBb,IAAI,CAAC,WAAW,EAAEa,qBAAqBG,QAAQ,IAAI,OAAO,sDAAsD,CAAC;IAE3J;IAEA,IACEH,qBAAqBK,IAAI,IACzBL,qBAAqBK,IAAI,CAACD,WAAW,OAAO,OAC5C;QACA,MAAM,IAAIF,MACR,CAAC,sEAAsE,EAAEF,qBAAqBb,IAAI,CAAC,QAAQ,EAAEa,qBAAqBK,IAAI,CAAC,iBAAiB,CAAC;IAE7J;IAEA,MAAMC,2BAA2BN,qBAAqBb,IAAI,IAAI;IAC9D,MAAMoB,uBAAuBP,qBAAqBQ,EAAE;IACpD,MAAMC,uBAAuBpC,UAAUkC;IACvC,MAAMG,kBAAkBV,qBAAqBW,IAAI,IAAI;IAErD,yEAAyE;IACzE,oEAAoE;IACpE,uEAAuE;IACvE,4EAA4E;IAC5E,iEAAiE;IACjE,sEAAsE;IACtE,gDAAgD;IAChD,MAAMrB,YAAY1B,sBAAsBkC,eAAeR,SAAS;IAChE,MAAMsB,SAASjD,gBAAgB,CAAC2B,UAAU;IAE1C,2EAA2E;IAC3E,6DAA6D;IAC7D,MAAMuB,WAAyC;QAAEvB;IAAU;IAE3D,wEAAwE;IACxE,2EAA2E;IAC3E,MAAM/B,mCAAmCmC,MAAMR;IAC/C,MAAM7B,oBAAoBqC,MAAM,OAAOR,cAAcI;IAErD,MAAMwB,4BAA4BpD,mCAAmCgC;IACrE,MAAMqB,4BAA4BvD,mCAAmCkC;IAErE,wEAAwE;IACxE,4EAA4E;IAC5E,oBAAoB;IACpBvC,kBAAkBuC,MAAMR,cAAc;QACpC2B;QACAG,aAAaF;IACf;IAEA,iFAAiF;IACjF,MAAMG,SAASjE,kBACb8D,2BACAC,2BACA;IAEFhE,cACE2C,MACA1C,kBAAkB,YAAYkE,OAAO,EAAE,SAASN,OAAOO,iBAAiB,GACxEF,QACA;QACER;QACAF;QACAG;QACAK;IACF,GACA;QAAEK,mBAAmBnE,kBAAkBoE,YAAY;IAAC;IAGtD,gCAAgC;IAChC,MAAMC,cAActE,kBAAkBiE,QAAQ;IAC9C,IAAI,CAACvB,KAAK6B,MAAM,CAACD,cAAc;QAC7B5B,KAAK8B,KAAK,CAACF,aAAa;IAC1B;IAEA,sCAAsC;IACtC,MAAMG,iBAAiBzE,kBACrB8D,2BACAC,2BACA;IAEF,MAAMzD,kBACJoC,MACA+B,gBACA,CAAC,KAAK,EAAEhB,qBAAqB,QAAQ,EAAEG,OAAOc,kBAAkB,EAAE,EAClE,GAAGnB,qBAAqB,MAAM,EAAEK,OAAOe,iBAAiB,EAAE;IAG5D,yEAAyE;IACzE,MAAMC,iBAAiB5E,kBACrB4C,cAAciC,IAAI,EAClB/B,eAAegC,IAAI,IAAI;IAEzB,MAAMC,gBAAgB/E,kBAAkB4E,gBAAgB;IAExD,IAAIlC,KAAK6B,MAAM,CAACQ,gBAAgB;QAC9B,MAAMC,kBAAkB,GAAGzB,qBAAqB,MAAM,EAAEK,OAAOe,iBAAiB,EAAE;QAClF,MAAMM,gBAAgBxB;QACtB,MAAMyB,WAAW,CAAC,IAAI,EAAEzB,sBAAsB;QAE9C,MAAM0B,qBACJzC,MACAqC,eACAhB,2BACAiB,iBACA1C;QAEF,MAAM8C,eAAe1C,MAAMqC,eAAeG,UAAU5C;QACpD,MAAM+C,wBACJ3C,MACAqC,eACAC,iBACAC,eACAC,UACA3B,sBACAjB;IAEJ;IAEA,6EAA6E;IAC7EX,kCACEe,MACAE,eACAnC,gCAAgCiC;IAGlC,kEAAkE;IAClE,MAAM4C,YAAYxC,eAAeX,IAAI,IAAI;IACzC,MAAMoD,gBAAgB,GAAGD,UAAU,IAAI,CAAC;IACxC,MAAME,sBAAsB,GAAGlC,yBAAyB,IAAI,CAAC;IAE7D,IAAIV,cAAc6C,OAAO,EAAE,CAACF,cAAc,EAAE;QAC1ChE,kCAAkCqB,eAAe2C,eAAe;YAC9DG,UAAU;gBAAC7C,cAAcV,IAAI;aAAC;YAC9BwD,QAAQH;QACV;QACAtF,2BAA2BwC,MAAME,cAAcT,IAAI,EAAES;IACvD;IAEA,8EAA8E;IAC9EtB,8BACEoB,MACAE,cAAcT,IAAI,EAClBI,wCACAb,sBAAsBkB,eAAemC,gBACrCxB,sBACA,2EAA2E;IAC3E,8DAA8D;IAC9D;QAAE,GAAGM,QAAQ;QAAE+B,YAAYjD,QAAQI,eAAe,EAAE+B;IAAK;IAG3D,MAAM1D,gCAAgCsB,MAAM;QAC1CH;KACD;IAED,MAAMrB,qBAAqBwB;IAC3B,OAAO,IACLvB,oBAAoBuB,MAAMC,QAAQkD,yBAAyB,EAAE;YAC3DC,WAAW;gBAAC;aAAa;QAC3B;AACJ,EAAE;AAEF;;;;CAIC,GACD,MAAMX,uBAAuB,OAC3BzC,MACAqD,UACAhC,2BACAiB,iBACA1C;IAEA,MAAMzB,4BACJ6B,MACAqD,UACA;QAAC;KAAO,EACRjE,kBAAkB,CAACQ,UAAU;IAE/B,MAAMzB,4BACJ6B,MACAqD,UACA;QAACf;KAAgB,EACjBjB;AAEJ;AAEA;;;;;CAKC,GACD,MAAMqB,iBAAiB,OACrB1C,MACAqD,UACAb,UACA5C;IAEA,MAAM0D,QAAQ/D,iBAAiB,CAACK,UAAU;IAC1C,MAAMxB,YACJ4B,MACAqD,UACAnE,GAAG,CAAC;mBACW,EAAEoE,MAAM;yBACF,EAAEd,SAAS;;eAErB,EAAEA,SAAS;;+CAEqB,EAAEA,SAAS;;CAEzD,CAAC;AAEF;AAEA;;;;;;;CAOC,GACD,MAAMG,0BAA0B,OAC9B3C,MACAqD,UACAf,iBACAC,eACAC,UACA3B,sBACAjB;IAEA,IAAI,MAAMvB,YAAY2B,MAAMqD,UAAUnE,GAAG,CAAC,EAAE,EAAEoD,gBAAgB,SAAS,CAAC,IAAI;QAC1E;IACF;IAEA,wEAAwE;IACxE,sEAAsE;IACtE,4EAA4E;IAC5E,MAAMiB,YAAY,GAAGhB,cAAc,GAAG,EAAED,gBAAgB;;;QAGlD,EAAEE,SAAS;6CAC0B,EAAE3B,qBAAqB;mBACjD,EAAE0B,cAAc;AACnC,CAAC;IAEC,8EAA8E;IAC9E,2EAA2E;IAC3E,mEAAmE;IACnE,MAAMiB,SAASjE,iBAAiB,CAACK,UAAU;IAC3C,MAAMxB,YACJ4B,MACAqD,UACAnE,GAAG,CAAC;;sBAEc,EAAEsE,OAAO;0BACL,EAAElB,gBAAgB;;IAExC,EAAEiB,UAAU;WACL,CAAC;AAEZ;AAEA,eAAexD,8BAA8B"}

package/src/py/agent/a2a-connection/schema.d.js

@@ -0,0 +1,8 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */ /**
+ * TypeScript types for options defined in schema.json
+ */ export { };
+
+//# sourceMappingURL=schema.d.js.map

package/src/py/agent/a2a-connection/schema.d.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["../../../../../../../packages/nx-plugin/src/py/agent/a2a-connection/schema.d.ts"],"sourcesContent":["/**\n * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n * SPDX-License-Identifier: Apache-2.0\n */\nimport { ComponentMetadata } from '../../../utils/nx';\n\n/**\n * TypeScript types for options defined in schema.json\n */\nexport interface PyAgentA2aConnectionGeneratorSchema {\n  sourceProject: string;\n  targetProject: string;\n  sourceComponent?: ComponentMetadata;\n  targetComponent?: ComponentMetadata;\n  preferInstallDependencies?: boolean;\n}\n"],"names":[],"mappings":"AAAA;;;CAGC,GAGD;;CAEC,GACD,WAMC"}

package/src/py/agent/a2a-connection/schema.d.ts

@@ -12,4 +12,5 @@ export interface PyAgentA2aConnectionGeneratorSchema {
   targetProject: string;
   sourceComponent?: ComponentMetadata;
   targetComponent?: ComponentMetadata;
+  preferInstallDependencies?: boolean;
 }

package/src/py/agent/a2a-connection/schema.json

@@ -20,6 +20,11 @@
     "targetComponent": {
       "type": "string",
       "description": "The target A2A agent component name"
+    },
+    "preferInstallDependencies": {
+      "type": "boolean",
+      "description": "Whether to prefer installing dependencies after the generator runs. Set to false to defer installing when batching multiple generators (an install still runs if needed so subsequent generators can compute the Nx project graph); install once at the end.",
+      "default": true
     }
   },
   "required": ["sourceProject", "targetProject"]

package/src/py/agent/files/deploy/Dockerfile.template

@@ -1,4 +1,11 @@
-FROM public.ecr.aws/docker/library/python:3.14-slim
+FROM <%- pythonBaseImage %>
+
+# Remove the base image's pip: the agent runs from the bundled environment and
+# never needs it at runtime, and pip's vendored packages carry known
+# HIGH/CRITICAL vulnerabilities flagged by the image scan.
+RUN python -m pip uninstall -y pip && \
+    rm -rf /usr/local/lib/python*/site-packages/pip \
+           /usr/local/lib/python*/site-packages/pip-*.dist-info
 
 WORKDIR /app
 

package/src/py/agent/files/langchain/a2a/main.py.template

@@ -0,0 +1,94 @@
+import logging
+import os
+import uuid
+from contextlib import asynccontextmanager
+
+from a2a.server.agent_execution import AgentExecutor, RequestContext
+from a2a.server.apps import A2AStarletteApplication
+from a2a.server.events import EventQueue
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
+from a2a.types import AgentCapabilities, AgentCard, Part, TextPart
+from a2a.utils import new_task
+from fastapi import FastAPI, Request
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from <%- agentConnectionModuleName %> import get_current_session_id, session_id_context
+
+from .agent import get_agent
+
+logging.basicConfig(level=logging.INFO)
+
+PORT = int(os.environ.get("PORT", "9000"))
+RUNTIME_URL = os.environ.get("AGENTCORE_RUNTIME_URL", f"http://localhost:{PORT}/")
+SESSION_ID_HEADER = "x-amzn-bedrock-agentcore-runtime-session-id"
+DEFAULT_MODES = ["text/plain"]
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.agent = get_agent()
+    yield
+
+
+class _GraphAgentExecutor(AgentExecutor):
+    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
+        if context.message is None:
+            raise ValueError("A2A request is missing a message")
+        task = context.current_task or new_task(context.message)
+        await event_queue.enqueue_event(task)
+        updater = TaskUpdater(event_queue, task.id, task.context_id)
+        await updater.start_work()
+        # The session bound by the middleware drives the LangGraph thread, keeping
+        # checkpointed conversation state separate per caller session.
+        session_id = get_current_session_id() or task.context_id
+        result = await app.state.agent.ainvoke(
+            {"messages": [{"role": "user", "content": context.get_user_input()}]},
+            {"configurable": {"thread_id": session_id}},
+        )
+        reply = result["messages"][-1].content
+        await updater.add_artifact([Part(TextPart(kind="text", text=reply))])
+        await updater.complete()
+
+    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
+        raise NotImplementedError("Cancellation is not supported.")
+
+
+_handler = DefaultRequestHandler(
+    agent_executor=_GraphAgentExecutor(),
+    task_store=InMemoryTaskStore(),
+)
+
+_agent_card = AgentCard(
+    name="<%= agentNameClassName %>",
+    description="A LangChain/LangGraph Agent exposed via the Agent-to-Agent (A2A) protocol.",
+    url=RUNTIME_URL,
+    version="1.0.0",
+    default_input_modes=DEFAULT_MODES,
+    default_output_modes=DEFAULT_MODES,
+    capabilities=AgentCapabilities(streaming=True),
+    skills=[],
+)
+
+a2a_app = A2AStarletteApplication(agent_card=_agent_card, http_handler=_handler).build()
+
+
+class _SessionIdMiddleware(BaseHTTPMiddleware):
+    """Bind the inbound session (or a fresh UUID) to async context."""
+
+    async def dispatch(self, request: Request, call_next):
+        session_id = request.headers.get(SESSION_ID_HEADER) or str(uuid.uuid4())
+        with session_id_context(session_id):
+            return await call_next(request)
+
+
+app = FastAPI(title="<%= agentNameClassName %>", lifespan=lifespan)
+app.add_middleware(_SessionIdMiddleware)
+
+
+@app.get("/ping")
+def ping() -> dict[str, str]:
+    return {"status": "Healthy"}
+
+
+app.mount("/", a2a_app)

package/src/py/agent/files/langchain/ag-ui/main.py.template

@@ -0,0 +1,87 @@
+import logging
+import time
+import uuid
+from contextlib import asynccontextmanager
+
+from ag_ui.core import EventType, RunAgentInput, RunErrorEvent
+from ag_ui.encoder import EventEncoder
+from ag_ui_langgraph import LangGraphAgent
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, StreamingResponse
+from <%- agentConnectionModuleName %> import get_current_session_id, session_id_context
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from .agent import get_agent
+
+logging.basicConfig(level=logging.INFO)
+
+SESSION_ID_HEADER = "x-amzn-bedrock-agentcore-runtime-session-id"
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    graph = get_agent()
+    app.state.agui_agent = LangGraphAgent(
+        name="<%= agentNameClassName %>",
+        graph=graph,
+        description="A LangChain/LangGraph Agent exposed via the AG-UI protocol.",
+    )
+    yield
+
+
+class _SessionIdMiddleware(BaseHTTPMiddleware):
+    """Bind the session ID for this request so downstream MCP / A2A clients forward it on outbound calls."""
+
+    async def dispatch(self, request: Request, call_next):
+        session_id = request.headers.get(SESSION_ID_HEADER) or str(uuid.uuid4())
+        with session_id_context(session_id):
+            return await call_next(request)
+
+
+app = FastAPI(title="<%= agentNameClassName %>", lifespan=lifespan)
+# Allow browser-based AG-UI clients (e.g. a connected React website) to call
+# this agent cross-origin, including the CORS preflight.
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+app.add_middleware(_SessionIdMiddleware)
+
+
+@app.post("/invocations")
+async def invocations(request: Request):
+    # Validate the body manually since AgentCore may omit Content-Type.
+    accept = request.headers.get("accept") or ""
+    encoder = EventEncoder(accept=accept)
+    raw = await request.body()
+    try:
+        input_data = RunAgentInput.model_validate_json(raw)
+    except Exception as exc:
+        message = f"Invalid RunAgentInput: {str(exc)[:200]}"
+
+        async def _bad():
+            yield encoder.encode(RunErrorEvent(type=EventType.RUN_ERROR, message=message, code="BAD_REQUEST"))
+
+        return StreamingResponse(_bad(), media_type=encoder.get_content_type())
+
+    session_id = request.headers.get(SESSION_ID_HEADER) or get_current_session_id()
+
+    async def event_generator():
+        # Re-bind the session: the streaming body runs outside the middleware.
+        with session_id_context(session_id or str(uuid.uuid4())):
+            try:
+                async for event in request.app.state.agui_agent.run(input_data):
+                    yield encoder.encode(event)
+            except Exception as e:
+                yield encoder.encode(RunErrorEvent(type=EventType.RUN_ERROR, message=str(e)[:300], code="AGENT_ERROR"))
+
+    return StreamingResponse(event_generator(), media_type=encoder.get_content_type())
+
+
+@app.get("/ping")
+async def ping():
+    return JSONResponse({"status": "Healthy", "time_of_last_update": int(time.time())})

package/src/py/agent/files/langchain/common/agent.py.template

@@ -0,0 +1,30 @@
+import os
+
+from langchain.agents import create_agent
+from langchain_aws import ChatBedrockConverse
+from langchain_core.tools import tool
+from langgraph.checkpoint.memory import InMemorySaver
+
+REGION = os.environ.get("AWS_REGION", "us-east-1")
+MODEL_ID = os.environ.get("MODEL_ID", "global.anthropic.claude-haiku-4-5-20251001-v1:0")
+
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a."""
+    return a - b
+
+
+def get_agent():
+    model = ChatBedrockConverse(model=MODEL_ID, region_name=REGION)
+    # Swap InMemorySaver for a durable checkpointer in a multi-replica deployment.
+    return create_agent(
+        model=model,
+        tools=[subtract],
+        system_prompt="""
+You are a mathematical wizard.
+Use your tools for mathematical tasks.
+Refer to tools as your 'spellbook'.
+""",
+        checkpointer=InMemorySaver(),
+    )

package/src/py/agent/files/{http → langchain/http}/init.py.template

@@ -1,14 +1,16 @@
+import logging
 from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.openapi.utils import get_openapi
 from fastapi.responses import JSONResponse, StreamingResponse
-from fastapi.routing import APIRoute
 from pydantic import BaseModel
 from starlette.middleware.exceptions import ExceptionMiddleware
 
+from .agent import get_agent
+
 
 class InternalServerErrorDetails(BaseModel):
     detail: str
@@ -61,7 +63,18 @@ class JsonStreamingResponse(StreamingResponse):
         }
 
 
-app = FastAPI(title="<%= agentNameClassName %>", responses={500: {"model": InternalServerErrorDetails}})
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.graph = get_agent()
+    yield
+
+
+app = FastAPI(
+    title="<%= agentNameClassName %>",
+    responses={500: {"model": InternalServerErrorDetails}},
+    generate_unique_id_function=lambda route: route.name,
+    lifespan=lifespan,
+)
 
 # Add cors middleware
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
@@ -72,28 +85,7 @@ app.add_middleware(ExceptionMiddleware, handlers=app.exception_handlers)
 
 @app.exception_handler(Exception)
 async def unhandled_exception_handler(request, err):
-    print(request)
-    print(err)
+    logging.exception("Unhandled exception while handling request")
     return JSONResponse(
         status_code=500, content=InternalServerErrorDetails(detail="Internal Server Error").model_dump()
     )
-
-
-def custom_openapi() -> dict[str, Any]:
-    if app.openapi_schema:
-        return app.openapi_schema
-    for route in app.routes:
-        if isinstance(route, APIRoute):
-            route.operation_id = route.name
-    openapi_schema = get_openapi(
-        title=app.title,
-        version=app.version,
-        openapi_version=app.openapi_version,
-        description=app.description,
-        routes=app.routes,
-    )
-    app.openapi_schema = openapi_schema
-    return app.openapi_schema
-
-
-app.openapi = custom_openapi  # ty: ignore[invalid-assignment]

package/src/py/agent/files/langchain/http/main.py.template

@@ -0,0 +1,81 @@
+import uuid
+
+import uvicorn
+from bedrock_agentcore.runtime.models import PingStatus
+from fastapi import Request
+from pydantic import BaseModel, Field
+from starlette.middleware.base import BaseHTTPMiddleware
+from <%- agentConnectionModuleName %> import get_current_session_id, session_id_context
+
+from .init import JsonStreamingResponse, app
+
+SESSION_ID_HEADER = "x-amzn-bedrock-agentcore-runtime-session-id"
+
+
+class InvokeInput(BaseModel):
+    prompt: str = Field(max_length=100000)
+
+
+class StreamChunk(BaseModel):
+    content: str
+
+
+def _text(content) -> str:
+    # ChatBedrockConverse streams content as a list of blocks, not a plain string.
+    if isinstance(content, str):
+        return content
+    return "".join(
+        block.get("text", "")
+        for block in content
+        if isinstance(block, dict) and block.get("type") == "text"
+    )
+
+
+async def handle_invoke(input: InvokeInput, session_id: str):
+    """Streaming handler for agent invocation"""
+    # The session id is the LangGraph thread id, so the checkpointer keeps memory per session.
+    config = {"configurable": {"thread_id": session_id}}
+    # Re-bind the session: the streaming body runs outside the middleware scope.
+    with session_id_context(session_id):
+        stream = app.state.graph.astream(
+            {"messages": [{"role": "user", "content": input.prompt}]},
+            config,
+            stream_mode="messages",
+        )
+        async for message_chunk, _metadata in stream:
+            text = _text(message_chunk.content)
+            if text:
+                yield StreamChunk(content=text)
+
+
+@app.post(
+    "/invocations",
+    response_class=JsonStreamingResponse,
+    responses={200: JsonStreamingResponse.openapi_response(StreamChunk, "Stream of agent response chunks")},
+)
+async def invoke(input: InvokeInput) -> JsonStreamingResponse:
+    """Entry point for agent invocation"""
+    session_id = get_current_session_id() or str(uuid.uuid4())
+    return JsonStreamingResponse(handle_invoke(input, session_id))
+
+
+class _SessionIdMiddleware(BaseHTTPMiddleware):
+    """Bind the inbound session (or a fresh UUID) to async context."""
+
+    async def dispatch(self, request: Request, call_next):
+        session_id = request.headers.get(SESSION_ID_HEADER) or str(uuid.uuid4())
+        with session_id_context(session_id):
+            return await call_next(request)
+
+
+app.add_middleware(_SessionIdMiddleware)
+
+
+@app.get("/ping")
+def ping() -> str:
+    # TODO: if running an async task, return PingStatus.HEALTHY_BUSY
+    return PingStatus.HEALTHY
+
+
+if __name__ == "__main__":
+    uvicorn.run("<%- moduleName %>.<%- agentNameSnakeCase %>.main:app", port=8080)

package/src/py/agent/files/{a2a → strands/a2a}/main.py.template

@@ -1,9 +1,13 @@
 import logging
 import os
 import uuid
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from typing import cast
 
 from fastapi import FastAPI, Request
 from starlette.middleware.base import BaseHTTPMiddleware
+from strands import Agent
 from strands.multiagent.a2a import A2AServer
 from <%- agentConnectionModuleName %> import session_id_context, with_session_id
 
@@ -14,23 +18,27 @@ logging.basicConfig(level=logging.INFO)
 PORT = int(os.environ.get("PORT", "9000"))
 RUNTIME_URL = os.environ.get("AGENTCORE_RUNTIME_URL", f"http://localhost:{PORT}/")
 SESSION_ID_HEADER = "x-amzn-bedrock-agentcore-runtime-session-id"
+AGENT_NAME = "<%= agentNameClassName %>"
+AGENT_DESCRIPTION = "A Strands Agent exposed via the Agent-to-Agent (A2A) protocol."
 
-_agent_ctx = with_session_id(
-    get_agent,
-    name="<%= agentNameClassName %>",
-    description="A Strands Agent exposed via the Agent-to-Agent (A2A) protocol.",
+# A2AServer.__init__ calls agent_factory once, synchronously, before the lifespan
+# below has run - this is what it reads name/description from at that point.
+# The cast is a lie about runtime shape, not behavior: only .name/.description
+# are ever read off this value.
+_card_placeholder = cast(
+    Agent, SimpleNamespace(name=AGENT_NAME, description=AGENT_DESCRIPTION)
 )
-_agent = _agent_ctx.__enter__()
 
-a2a_server = A2AServer(
-    agent=_agent,
-    port=PORT,
-    http_url=RUNTIME_URL,
-    serve_at_root=True,
-    # Skip tool-registry introspection for the agent card — the per-session
-    # agents aren't constructed yet at import time.
-    skills=[],
-)
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    with with_session_id(
+        get_agent,
+        name=AGENT_NAME,
+        description=AGENT_DESCRIPTION,
+    ) as agent:
+        app.state.agent = agent
+        yield
 
 
 class _SessionIdMiddleware(BaseHTTPMiddleware):
@@ -42,7 +50,7 @@ class _SessionIdMiddleware(BaseHTTPMiddleware):
             return await call_next(request)
 
 
-app = FastAPI()
+app = FastAPI(lifespan=lifespan)
 app.add_middleware(_SessionIdMiddleware)
 
 
@@ -51,4 +59,14 @@ def ping() -> dict[str, str]:
     return {"status": "Healthy"}
 
 
+a2a_server = A2AServer(
+    agent_factory=lambda _context_id: getattr(app.state, "agent", _card_placeholder),
+    port=PORT,
+    http_url=RUNTIME_URL,
+    serve_at_root=True,
+    # Skip tool-registry introspection for the agent card — the per-session
+    # agents aren't constructed yet at startup.
+    skills=[],
+)
+
 app.mount("/", a2a_server.to_fastapi_app())

package/src/py/agent/files/strands/ag-ui/main.py.template

@@ -0,0 +1,89 @@
+import logging
+import uuid
+from contextlib import asynccontextmanager
+
+from ag_ui.core import EventType, RunAgentInput, RunErrorEvent
+from ag_ui.encoder import EventEncoder
+from ag_ui_strands import StrandsAgent
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+from <%- agentConnectionModuleName %> import get_current_session_id, session_id_context
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from .agent import get_agent
+
+logging.basicConfig(level=logging.INFO)
+
+SESSION_ID_HEADER = "x-amzn-bedrock-agentcore-runtime-session-id"
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    with get_agent() as agent:
+        app.state.agui_agent = StrandsAgent(
+            agent=agent,
+            name="<%= agentNameClassName %>",
+            description="A Strands Agent exposed via the AG-UI protocol.",
+        )
+        yield
+
+
+class _SessionIdMiddleware(BaseHTTPMiddleware):
+    """Bind the session ID for this request so downstream MCP / A2A clients forward it on outbound calls."""
+
+    async def dispatch(self, request: Request, call_next):
+        session_id = request.headers.get(SESSION_ID_HEADER) or str(uuid.uuid4())
+        with session_id_context(session_id):
+            return await call_next(request)
+
+
+app = FastAPI(title="AWS Strands - <%= agentNameClassName %>", lifespan=lifespan)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+app.add_middleware(_SessionIdMiddleware)
+
+
+@app.post("/invocations")
+async def invocations(request: Request):
+    # Validate the body manually since AgentCore may omit Content-Type.
+    encoder = EventEncoder(accept=request.headers.get("accept") or "")
+    raw = await request.body()
+    try:
+        input_data = RunAgentInput.model_validate_json(raw)
+    except Exception as exc:
+        message = f"Invalid RunAgentInput: {str(exc)[:200]}"
+
+        async def _bad():
+            yield encoder.encode(RunErrorEvent(type=EventType.RUN_ERROR, message=message, code="BAD_REQUEST"))
+
+        return StreamingResponse(_bad(), media_type=encoder.get_content_type())
+
+    session_id = request.headers.get(SESSION_ID_HEADER) or get_current_session_id()
+
+    async def event_generator():
+        # Re-bind the session: the streaming body runs outside the middleware.
+        with session_id_context(session_id or str(uuid.uuid4())):
+            async for event in request.app.state.agui_agent.run(input_data):
+                try:
+                    yield encoder.encode(event)
+                except Exception as e:
+                    error_event = RunErrorEvent(
+                        type=EventType.RUN_ERROR,
+                        message=f"Encoding error: {e}",
+                        code="ENCODING_ERROR",
+                    )
+                    yield encoder.encode(error_event)
+                    break
+
+    return StreamingResponse(event_generator(), media_type=encoder.get_content_type())
+
+
+@app.get("/ping")
+async def ping():
+    return {"status": "healthy"}

package/src/py/agent/files/{common → strands/common}/agent.py.template

@@ -2,6 +2,7 @@ from contextlib import contextmanager
 
 from strands import Agent, tool
 from strands_tools import current_time
+from <%- agentConnectionModuleName %> import log_model_errors
 
 
 @tool
@@ -20,4 +21,5 @@ Use your tools for mathematical tasks.
 Refer to tools as your 'spellbook'.
 """,
         tools=[subtract, current_time],
+        hooks=[log_model_errors],
     )