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

package/migrations.json

@@ -12,6 +12,11 @@
       "description": "Update the docs translate script to write whole translated files rather than editing them in place",
       "implementation": "./src/migrations/latest/translate-script-whole-file-writes/migration"
     },
+    "latest-py-agent-a2a-lifespan-construction": {
+      "version": "1.0.0-rc.70",
+      "description": "Move py#agent A2A agent construction (Strands and LangChain) out of module import time and into a FastAPI lifespan handler, stored on app.state; also updates with_session_id/session_id_context to use Generator instead of the deprecated Iterator return type",
+      "implementation": "./src/migrations/latest/py-agent-a2a-lifespan-construction/migration"
+    },
     "v1.0.0-rc.50-0001-modernize-function-props-cast": {
       "version": "1.0.0-rc.50",
       "description": "Replace the legacy angle-bracket FunctionProps type assertion with the modern as syntax in generated API constructs",
@@ -108,7 +113,7 @@
       "implementation": "./src/migrations/v1.0.0-rc.65/0002-rolldown-code-splitting/migration"
     },
     "sync-vended-versions": {
-      "version": "1.0.0-rc.69",
+      "version": "1.0.0-rc.70",
       "description": "Sync vended dependency versions and the tracked plugin version to those vended by this release",
       "implementation": "./src/utils/version-upgrade-migration/migration"
     }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws/nx-plugin",
-  "version": "1.0.0-rc.69",
+  "version": "1.0.0-rc.70",
   "repository": {
     "type": "git",
     "url": "https://github.com/awslabs/nx-plugin-for-aws.git",

package/src/migrations/latest/py-agent-a2a-lifespan-construction/metadata.json

@@ -0,0 +1,3 @@
+{
+  "description": "Move py#agent A2A agent construction (Strands and LangChain) out of module import time and into a FastAPI lifespan handler, stored on app.state; also updates with_session_id/session_id_context to use Generator instead of the deprecated Iterator return type"
+}

package/src/migrations/latest/py-agent-a2a-lifespan-construction/migration.d.ts

@@ -0,0 +1,6 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import { type MigrationReturnObject, type Tree } from '@nx/devkit';
+export default function migration(tree: Tree): Promise<MigrationReturnObject>;

package/src/migrations/latest/py-agent-a2a-lifespan-construction/migration.js

@@ -0,0 +1,234 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */ import { getProjects, joinPathFragments, visitNotIgnoredFiles } from "@nx/devkit";
+import { PY_AGENT_GENERATOR_INFO } from "../../../py/agent/generator.js";
+import { addPythonDestructuredImport, applyGritQL, captureGritQLVariable, matchGritQL } from "../../../utils/ast.js";
+import { formatFilesInSubtree } from "../../../utils/format.js";
+/**
+ * Move py#agent A2A agent construction (Strands and LangChain) out of module
+ * import time and into a FastAPI `lifespan` handler, stored on `app.state`.
+ *
+ * The `v1.0.0-rc.61/0001-py-agent-lifespan-construction` migration (HTTP/AG-UI)
+ * deliberately left A2A alone, reasoning that `with_session_id` only touches
+ * its agent factory lazily on first use, so the eager `_agent_ctx.__enter__()`
+ * was never actually eager in practice. That held for `with_session_id`'s own
+ * laziness, but missed a separate trigger: `strands.multiagent.a2a.executor
+ * .StrandsA2AExecutor.__init__`, in its deprecated single-`agent=` mode, does
+ * `getattr(agent, "_session_manager", None)` - which the session-routing
+ * proxy's `__getattr__` happily serves by building the real underlying Agent
+ * right then, under a bogus "default" session, since no request has bound a
+ * real session id yet. That call happens inside `A2AServer.__init__`, which
+ * the old template constructed at plain module import time - so the "lazy"
+ * proxy was forced eager anyway, just one layer further down than the
+ * previous migration checked.
+ *
+ * Strands: `agent_factory=` mode sidesteps this - `StrandsA2AExecutor` takes
+ * an entirely different branch when given a factory (no `_session_manager`
+ * probe at all), and the one placeholder call `A2AServer.__init__` makes
+ * before `lifespan` has run returns a bare `SimpleNamespace` with just
+ * `name`/`description` - no `__getattr__`, nothing to force construction.
+ *
+ * LangChain: the compiled graph from `get_agent()` has no such trap, but was
+ * still built eagerly at import time (`_graph = get_agent()`). Moved into
+ * `lifespan` for the same reason HTTP/AG-UI were: it can require env vars
+ * (model id, region) or do network setup that isn't guaranteed available at
+ * import - e.g. under test collection or OpenAPI spec generation tooling.
+ *
+ * Also folds in an unrelated but adjacent fix: Pylance (Pyright) deprecates
+ * annotating a `@contextmanager`-decorated function's return type as
+ * `Iterator[T]`, in favor of the single-argument `Generator[T]` form enabled
+ * by PEP 696 default type parameters (Python 3.13+ - generated projects
+ * target >=3.14). `with_session_id_strands.py` and `session_context.py` both
+ * use this pattern. Both files are shared, protocol-agnostic files copied
+ * once per workspace (not per py#agent component), so that part scans the
+ * whole tree by basename rather than going through component metadata.
+ */ const pyMatch = (snippet)=>`language python\n\`${snippet}\``;
+const pyDelete = (snippet)=>`${pyMatch(snippet)} => .`;
+const pyRewrite = (from, to)=>{
+    const replacement = to.includes('\n') ? `raw\`${to}\`` : `\`${to}\``;
+    return `${pyMatch(from)} => ${replacement}`;
+};
+const AGENT_ENTER = '_agent = _agent_ctx.__enter__()';
+const OLD_AGENT_CONTEXT = `_agent_ctx = with_session_id(
+    get_agent,
+    name="$name",
+    description="A Strands Agent exposed via the Agent-to-Agent (A2A) protocol.",
+)`;
+const NEW_CONSTANTS_AND_LIFESPAN = `AGENT_NAME = "$name"
+AGENT_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)
+)
+
+
+@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`;
+const A2A_SERVER_STATEMENT = 'a2a_server = A2AServer($args)';
+const OLD_MOUNT_STATEMENT = 'app.mount("/", a2a_server.to_fastapi_app())';
+const OLD_FASTAPI_APP = 'app = FastAPI()';
+const NEW_FASTAPI_APP = 'app = FastAPI(lifespan=lifespan)';
+const SERVER_MATCH = 'language python\n`A2AServer($args)` where { $args <: contains `agent=_agent` }';
+const allMatch = async (tree, filePath, patterns)=>{
+    for (const pattern of patterns){
+        if (!await matchGritQL(tree, filePath, pattern)) return false;
+    }
+    return true;
+};
+const migrateStrandsA2AAgent = async (tree, mainPath, nextSteps)=>{
+    if (!tree.exists(mainPath)) return;
+    const contents = tree.read(mainPath, 'utf-8') ?? '';
+    if (!contents.includes('_agent_ctx = with_session_id(')) return;
+    const ready = await allMatch(tree, mainPath, [
+        pyMatch('import logging'),
+        pyMatch('from $mod import session_id_context, with_session_id'),
+        pyMatch(OLD_AGENT_CONTEXT),
+        pyMatch(AGENT_ENTER),
+        pyMatch(OLD_FASTAPI_APP),
+        pyMatch(OLD_MOUNT_STATEMENT),
+        SERVER_MATCH
+    ]);
+    const diverged = ()=>{
+        nextSteps.push(`${mainPath}: diverged from the generated Strands A2A shape - left untouched. Manually rebuild the agent inside the FastAPI \`lifespan\` and have \`agent_factory\` read it off \`app.state\` (see the py#agent generator's strands/a2a template).`);
+    };
+    if (!ready) {
+        diverged();
+        return;
+    }
+    await addPythonDestructuredImport(tree, mainPath, [
+        'asynccontextmanager'
+    ], 'contextlib');
+    await addPythonDestructuredImport(tree, mainPath, [
+        'SimpleNamespace'
+    ], 'types');
+    await addPythonDestructuredImport(tree, mainPath, [
+        'cast'
+    ], 'typing');
+    await addPythonDestructuredImport(tree, mainPath, [
+        'Agent'
+    ], 'strands');
+    await applyGritQL(tree, mainPath, pyRewrite(OLD_AGENT_CONTEXT, NEW_CONSTANTS_AND_LIFESPAN));
+    await applyGritQL(tree, mainPath, pyDelete(AGENT_ENTER));
+    // Swap the kwarg while the call is still in its original spot (any extra
+    // kwargs a user added are left alone), then capture the whole arg list so
+    // it can be reinserted at the call's new home below.
+    await applyGritQL(tree, mainPath, pyRewrite('agent=_agent', 'agent_factory=lambda _context_id: getattr(app.state, "agent", _card_placeholder)'));
+    const capturedArgs = await captureGritQLVariable(tree, mainPath, `language python\n\`${A2A_SERVER_STATEMENT}\``, 'args');
+    if (capturedArgs === undefined) {
+        diverged();
+        return;
+    }
+    await applyGritQL(tree, mainPath, pyDelete(A2A_SERVER_STATEMENT));
+    await applyGritQL(tree, mainPath, pyRewrite(OLD_FASTAPI_APP, NEW_FASTAPI_APP));
+    await applyGritQL(tree, mainPath, pyRewrite(OLD_MOUNT_STATEMENT, `a2a_server = A2AServer(${capturedArgs})\n\n${OLD_MOUNT_STATEMENT}`));
+};
+// --- LangChain A2A -----------------------------------------------------------
+const LANGCHAIN_GRAPH_OLD = '_graph = get_agent()';
+const LANGCHAIN_LIFESPAN_NEW = `@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.agent = get_agent()
+    yield`;
+const LANGCHAIN_GRAPH_INVOKE_OLD = '_graph.ainvoke($args)';
+const LANGCHAIN_GRAPH_INVOKE_NEW = 'app.state.agent.ainvoke($args)';
+const LANGCHAIN_APP_OLD = 'app = FastAPI(title="$name")';
+const LANGCHAIN_APP_NEW = 'app = FastAPI(title="$name", lifespan=lifespan)';
+const migrateLangchainA2AAgent = async (tree, mainPath, nextSteps)=>{
+    if (!tree.exists(mainPath)) return;
+    const contents = tree.read(mainPath, 'utf-8') ?? '';
+    if (!contents.includes(LANGCHAIN_GRAPH_OLD)) return;
+    const ready = await allMatch(tree, mainPath, [
+        pyMatch(LANGCHAIN_GRAPH_OLD),
+        pyMatch('from .agent import get_agent'),
+        pyMatch(LANGCHAIN_GRAPH_INVOKE_OLD),
+        pyMatch(LANGCHAIN_APP_OLD)
+    ]);
+    if (!ready) {
+        nextSteps.push(`${mainPath}: diverged from the generated LangChain A2A shape - left untouched. Manually wrap \`get_agent()\` in a \`lifespan\` handler storing the graph on \`app.state.agent\` (see the py#agent generator's langchain/a2a template).`);
+        return;
+    }
+    await addPythonDestructuredImport(tree, mainPath, [
+        'asynccontextmanager'
+    ], 'contextlib');
+    await applyGritQL(tree, mainPath, pyRewrite(LANGCHAIN_GRAPH_OLD, LANGCHAIN_LIFESPAN_NEW));
+    await applyGritQL(tree, mainPath, pyRewrite(LANGCHAIN_GRAPH_INVOKE_OLD, LANGCHAIN_GRAPH_INVOKE_NEW));
+    await applyGritQL(tree, mainPath, pyRewrite(LANGCHAIN_APP_OLD, LANGCHAIN_APP_NEW));
+};
+// --- Iterator -> Generator (with_session_id_strands.py / session_context.py) ---
+const WITH_SESSION_ID_IMPORT_OLD = 'from collections.abc import Callable, Iterator';
+const WITH_SESSION_ID_IMPORT_NEW = 'from collections.abc import Callable, Generator';
+const migrateWithSessionId = async (tree, filePath, nextSteps)=>{
+    const contents = tree.read(filePath, 'utf-8') ?? '';
+    if (!contents.includes('Iterator[Any]')) return;
+    const ready = await allMatch(tree, filePath, [
+        pyMatch(WITH_SESSION_ID_IMPORT_OLD),
+        pyMatch('Iterator[Any]')
+    ]);
+    if (!ready) {
+        nextSteps.push(`${filePath}: diverged from the generated shape - left untouched. Manually change the \`Iterator\` import and \`Iterator[Any]\` return annotation on \`with_session_id\` to \`Generator\`/\`Generator[Any]\` (see the agent-connection generator's with_session_id_strands.py template).`);
+        return;
+    }
+    await applyGritQL(tree, filePath, pyRewrite(WITH_SESSION_ID_IMPORT_OLD, WITH_SESSION_ID_IMPORT_NEW));
+    await applyGritQL(tree, filePath, pyRewrite('Iterator[Any]', 'Generator[Any]'));
+};
+const SESSION_CONTEXT_IMPORT_OLD = 'from collections.abc import Iterator';
+const SESSION_CONTEXT_IMPORT_NEW = 'from collections.abc import Generator';
+const migrateSessionContext = async (tree, filePath, nextSteps)=>{
+    const contents = tree.read(filePath, 'utf-8') ?? '';
+    if (!contents.includes('Iterator[None]')) return;
+    const ready = await allMatch(tree, filePath, [
+        pyMatch(SESSION_CONTEXT_IMPORT_OLD),
+        pyMatch('Iterator[None]')
+    ]);
+    if (!ready) {
+        nextSteps.push(`${filePath}: diverged from the generated shape - left untouched. Manually change the \`Iterator\` import and \`Iterator[None]\` return annotation on \`session_id_context\` to \`Generator\`/\`Generator[None]\` (see the agent-connection generator's session_context.py template).`);
+        return;
+    }
+    await applyGritQL(tree, filePath, pyRewrite(SESSION_CONTEXT_IMPORT_OLD, SESSION_CONTEXT_IMPORT_NEW));
+    await applyGritQL(tree, filePath, pyRewrite('Iterator[None]', 'Generator[None]'));
+};
+const findAgentComponents = (components)=>(components ?? []).filter((component)=>component.generator === PY_AGENT_GENERATOR_INFO.id);
+export default async function migration(tree) {
+    const nextSteps = [];
+    for (const project of getProjects(tree).values()){
+        const components = findAgentComponents(project.metadata?.components);
+        for (const component of components){
+            if (!component.path || component.protocol !== 'a2a') {
+                continue;
+            }
+            const componentDir = component.path.endsWith('/agent.py') ? component.path.slice(0, -'/agent.py'.length) : component.path;
+            const mainPath = joinPathFragments(project.root, componentDir, 'main.py');
+            if ((component.framework ?? 'strands') === 'langchain') {
+                await migrateLangchainA2AAgent(tree, mainPath, nextSteps);
+            } else {
+                await migrateStrandsA2AAgent(tree, mainPath, nextSteps);
+            }
+        }
+    }
+    const filePaths = [];
+    visitNotIgnoredFiles(tree, '', (filePath)=>filePaths.push(filePath));
+    for (const filePath of filePaths){
+        if (filePath.endsWith('/with_session_id_strands.py')) {
+            await migrateWithSessionId(tree, filePath, nextSteps);
+        } else if (filePath.endsWith('/session_context.py')) {
+            await migrateSessionContext(tree, filePath, nextSteps);
+        }
+    }
+    await formatFilesInSubtree(tree);
+    return {
+        nextSteps
+    };
+}
+
+//# sourceMappingURL=migration.js.map

package/src/migrations/latest/py-agent-a2a-lifespan-construction/migration.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["../../../../../../../packages/nx-plugin/src/migrations/latest/py-agent-a2a-lifespan-construction/migration.ts"],"sourcesContent":["/**\n * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n * SPDX-License-Identifier: Apache-2.0\n */\nimport {\n  getProjects,\n  joinPathFragments,\n  type MigrationReturnObject,\n  type Tree,\n  visitNotIgnoredFiles,\n} from '@nx/devkit';\nimport { PY_AGENT_GENERATOR_INFO } from '../../../py/agent/generator';\nimport {\n  addPythonDestructuredImport,\n  applyGritQL,\n  captureGritQLVariable,\n  matchGritQL,\n} from '../../../utils/ast';\nimport { formatFilesInSubtree } from '../../../utils/format';\nimport type { ComponentMetadata } from '../../../utils/nx';\n\n/**\n * Move py#agent A2A agent construction (Strands and LangChain) out of module\n * import time and into a FastAPI `lifespan` handler, stored on `app.state`.\n *\n * The `v1.0.0-rc.61/0001-py-agent-lifespan-construction` migration (HTTP/AG-UI)\n * deliberately left A2A alone, reasoning that `with_session_id` only touches\n * its agent factory lazily on first use, so the eager `_agent_ctx.__enter__()`\n * was never actually eager in practice. That held for `with_session_id`'s own\n * laziness, but missed a separate trigger: `strands.multiagent.a2a.executor\n * .StrandsA2AExecutor.__init__`, in its deprecated single-`agent=` mode, does\n * `getattr(agent, \"_session_manager\", None)` - which the session-routing\n * proxy's `__getattr__` happily serves by building the real underlying Agent\n * right then, under a bogus \"default\" session, since no request has bound a\n * real session id yet. That call happens inside `A2AServer.__init__`, which\n * the old template constructed at plain module import time - so the \"lazy\"\n * proxy was forced eager anyway, just one layer further down than the\n * previous migration checked.\n *\n * Strands: `agent_factory=` mode sidesteps this - `StrandsA2AExecutor` takes\n * an entirely different branch when given a factory (no `_session_manager`\n * probe at all), and the one placeholder call `A2AServer.__init__` makes\n * before `lifespan` has run returns a bare `SimpleNamespace` with just\n * `name`/`description` - no `__getattr__`, nothing to force construction.\n *\n * LangChain: the compiled graph from `get_agent()` has no such trap, but was\n * still built eagerly at import time (`_graph = get_agent()`). Moved into\n * `lifespan` for the same reason HTTP/AG-UI were: it can require env vars\n * (model id, region) or do network setup that isn't guaranteed available at\n * import - e.g. under test collection or OpenAPI spec generation tooling.\n *\n * Also folds in an unrelated but adjacent fix: Pylance (Pyright) deprecates\n * annotating a `@contextmanager`-decorated function's return type as\n * `Iterator[T]`, in favor of the single-argument `Generator[T]` form enabled\n * by PEP 696 default type parameters (Python 3.13+ - generated projects\n * target >=3.14). `with_session_id_strands.py` and `session_context.py` both\n * use this pattern. Both files are shared, protocol-agnostic files copied\n * once per workspace (not per py#agent component), so that part scans the\n * whole tree by basename rather than going through component metadata.\n */\nconst pyMatch = (snippet: string) => `language python\\n\\`${snippet}\\``;\n\nconst pyDelete = (snippet: string) => `${pyMatch(snippet)} => .`;\n\nconst pyRewrite = (from: string, to: string): string => {\n  const replacement = to.includes('\\n') ? `raw\\`${to}\\`` : `\\`${to}\\``;\n  return `${pyMatch(from)} => ${replacement}`;\n};\n\nconst AGENT_ENTER = '_agent = _agent_ctx.__enter__()';\n\nconst OLD_AGENT_CONTEXT = `_agent_ctx = with_session_id(\n    get_agent,\n    name=\"$name\",\n    description=\"A Strands Agent exposed via the Agent-to-Agent (A2A) protocol.\",\n)`;\n\nconst NEW_CONSTANTS_AND_LIFESPAN = `AGENT_NAME = \"$name\"\nAGENT_DESCRIPTION = \"A Strands Agent exposed via the Agent-to-Agent (A2A) protocol.\"\n\n# A2AServer.__init__ calls agent_factory once, synchronously, before the lifespan\n# below has run - this is what it reads name/description from at that point.\n# The cast is a lie about runtime shape, not behavior: only .name/.description\n# are ever read off this value.\n_card_placeholder = cast(\n    Agent, SimpleNamespace(name=AGENT_NAME, description=AGENT_DESCRIPTION)\n)\n\n\n@asynccontextmanager\nasync def lifespan(app: FastAPI):\n    with with_session_id(\n        get_agent,\n        name=AGENT_NAME,\n        description=AGENT_DESCRIPTION,\n    ) as agent:\n        app.state.agent = agent\n        yield`;\n\nconst A2A_SERVER_STATEMENT = 'a2a_server = A2AServer($args)';\n\nconst OLD_MOUNT_STATEMENT = 'app.mount(\"/\", a2a_server.to_fastapi_app())';\n\nconst OLD_FASTAPI_APP = 'app = FastAPI()';\n\nconst NEW_FASTAPI_APP = 'app = FastAPI(lifespan=lifespan)';\n\nconst SERVER_MATCH =\n  'language python\\n`A2AServer($args)` where { $args <: contains `agent=_agent` }';\n\nconst allMatch = async (\n  tree: Tree,\n  filePath: string,\n  patterns: string[],\n): Promise<boolean> => {\n  for (const pattern of patterns) {\n    if (!(await matchGritQL(tree, filePath, pattern))) return false;\n  }\n  return true;\n};\n\nconst migrateStrandsA2AAgent = async (\n  tree: Tree,\n  mainPath: string,\n  nextSteps: string[],\n): Promise<void> => {\n  if (!tree.exists(mainPath)) return;\n\n  const contents = tree.read(mainPath, 'utf-8') ?? '';\n  if (!contents.includes('_agent_ctx = with_session_id(')) return;\n\n  const ready = await allMatch(tree, mainPath, [\n    pyMatch('import logging'),\n    pyMatch('from $mod import session_id_context, with_session_id'),\n    pyMatch(OLD_AGENT_CONTEXT),\n    pyMatch(AGENT_ENTER),\n    pyMatch(OLD_FASTAPI_APP),\n    pyMatch(OLD_MOUNT_STATEMENT),\n    SERVER_MATCH,\n  ]);\n\n  const diverged = () => {\n    nextSteps.push(\n      `${mainPath}: diverged from the generated Strands A2A shape - left untouched. Manually rebuild the agent inside the FastAPI \\`lifespan\\` and have \\`agent_factory\\` read it off \\`app.state\\` (see the py#agent generator's strands/a2a template).`,\n    );\n  };\n\n  if (!ready) {\n    diverged();\n    return;\n  }\n\n  await addPythonDestructuredImport(\n    tree,\n    mainPath,\n    ['asynccontextmanager'],\n    'contextlib',\n  );\n  await addPythonDestructuredImport(\n    tree,\n    mainPath,\n    ['SimpleNamespace'],\n    'types',\n  );\n  await addPythonDestructuredImport(tree, mainPath, ['cast'], 'typing');\n  await addPythonDestructuredImport(tree, mainPath, ['Agent'], 'strands');\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(OLD_AGENT_CONTEXT, NEW_CONSTANTS_AND_LIFESPAN),\n  );\n  await applyGritQL(tree, mainPath, pyDelete(AGENT_ENTER));\n\n  // Swap the kwarg while the call is still in its original spot (any extra\n  // kwargs a user added are left alone), then capture the whole arg list so\n  // it can be reinserted at the call's new home below.\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(\n      'agent=_agent',\n      'agent_factory=lambda _context_id: getattr(app.state, \"agent\", _card_placeholder)',\n    ),\n  );\n  const capturedArgs = await captureGritQLVariable(\n    tree,\n    mainPath,\n    `language python\\n\\`${A2A_SERVER_STATEMENT}\\``,\n    'args',\n  );\n  if (capturedArgs === undefined) {\n    diverged();\n    return;\n  }\n\n  await applyGritQL(tree, mainPath, pyDelete(A2A_SERVER_STATEMENT));\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(OLD_FASTAPI_APP, NEW_FASTAPI_APP),\n  );\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(\n      OLD_MOUNT_STATEMENT,\n      `a2a_server = A2AServer(${capturedArgs})\\n\\n${OLD_MOUNT_STATEMENT}`,\n    ),\n  );\n};\n\n// --- LangChain A2A -----------------------------------------------------------\n\nconst LANGCHAIN_GRAPH_OLD = '_graph = get_agent()';\n\nconst LANGCHAIN_LIFESPAN_NEW = `@asynccontextmanager\nasync def lifespan(app: FastAPI):\n    app.state.agent = get_agent()\n    yield`;\n\nconst LANGCHAIN_GRAPH_INVOKE_OLD = '_graph.ainvoke($args)';\n\nconst LANGCHAIN_GRAPH_INVOKE_NEW = 'app.state.agent.ainvoke($args)';\n\nconst LANGCHAIN_APP_OLD = 'app = FastAPI(title=\"$name\")';\n\nconst LANGCHAIN_APP_NEW = 'app = FastAPI(title=\"$name\", lifespan=lifespan)';\n\nconst migrateLangchainA2AAgent = async (\n  tree: Tree,\n  mainPath: string,\n  nextSteps: string[],\n): Promise<void> => {\n  if (!tree.exists(mainPath)) return;\n\n  const contents = tree.read(mainPath, 'utf-8') ?? '';\n  if (!contents.includes(LANGCHAIN_GRAPH_OLD)) return;\n\n  const ready = await allMatch(tree, mainPath, [\n    pyMatch(LANGCHAIN_GRAPH_OLD),\n    pyMatch('from .agent import get_agent'),\n    pyMatch(LANGCHAIN_GRAPH_INVOKE_OLD),\n    pyMatch(LANGCHAIN_APP_OLD),\n  ]);\n\n  if (!ready) {\n    nextSteps.push(\n      `${mainPath}: diverged from the generated LangChain A2A shape - left untouched. Manually wrap \\`get_agent()\\` in a \\`lifespan\\` handler storing the graph on \\`app.state.agent\\` (see the py#agent generator's langchain/a2a template).`,\n    );\n    return;\n  }\n\n  await addPythonDestructuredImport(\n    tree,\n    mainPath,\n    ['asynccontextmanager'],\n    'contextlib',\n  );\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(LANGCHAIN_GRAPH_OLD, LANGCHAIN_LIFESPAN_NEW),\n  );\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(LANGCHAIN_GRAPH_INVOKE_OLD, LANGCHAIN_GRAPH_INVOKE_NEW),\n  );\n  await applyGritQL(\n    tree,\n    mainPath,\n    pyRewrite(LANGCHAIN_APP_OLD, LANGCHAIN_APP_NEW),\n  );\n};\n\n// --- Iterator -> Generator (with_session_id_strands.py / session_context.py) ---\n\nconst WITH_SESSION_ID_IMPORT_OLD =\n  'from collections.abc import Callable, Iterator';\nconst WITH_SESSION_ID_IMPORT_NEW =\n  'from collections.abc import Callable, Generator';\n\nconst migrateWithSessionId = async (\n  tree: Tree,\n  filePath: string,\n  nextSteps: string[],\n): Promise<void> => {\n  const contents = tree.read(filePath, 'utf-8') ?? '';\n  if (!contents.includes('Iterator[Any]')) return;\n\n  const ready = await allMatch(tree, filePath, [\n    pyMatch(WITH_SESSION_ID_IMPORT_OLD),\n    pyMatch('Iterator[Any]'),\n  ]);\n\n  if (!ready) {\n    nextSteps.push(\n      `${filePath}: diverged from the generated shape - left untouched. Manually change the \\`Iterator\\` import and \\`Iterator[Any]\\` return annotation on \\`with_session_id\\` to \\`Generator\\`/\\`Generator[Any]\\` (see the agent-connection generator's with_session_id_strands.py template).`,\n    );\n    return;\n  }\n\n  await applyGritQL(\n    tree,\n    filePath,\n    pyRewrite(WITH_SESSION_ID_IMPORT_OLD, WITH_SESSION_ID_IMPORT_NEW),\n  );\n  await applyGritQL(\n    tree,\n    filePath,\n    pyRewrite('Iterator[Any]', 'Generator[Any]'),\n  );\n};\n\nconst SESSION_CONTEXT_IMPORT_OLD = 'from collections.abc import Iterator';\nconst SESSION_CONTEXT_IMPORT_NEW = 'from collections.abc import Generator';\n\nconst migrateSessionContext = async (\n  tree: Tree,\n  filePath: string,\n  nextSteps: string[],\n): Promise<void> => {\n  const contents = tree.read(filePath, 'utf-8') ?? '';\n  if (!contents.includes('Iterator[None]')) return;\n\n  const ready = await allMatch(tree, filePath, [\n    pyMatch(SESSION_CONTEXT_IMPORT_OLD),\n    pyMatch('Iterator[None]'),\n  ]);\n\n  if (!ready) {\n    nextSteps.push(\n      `${filePath}: diverged from the generated shape - left untouched. Manually change the \\`Iterator\\` import and \\`Iterator[None]\\` return annotation on \\`session_id_context\\` to \\`Generator\\`/\\`Generator[None]\\` (see the agent-connection generator's session_context.py template).`,\n    );\n    return;\n  }\n\n  await applyGritQL(\n    tree,\n    filePath,\n    pyRewrite(SESSION_CONTEXT_IMPORT_OLD, SESSION_CONTEXT_IMPORT_NEW),\n  );\n  await applyGritQL(\n    tree,\n    filePath,\n    pyRewrite('Iterator[None]', 'Generator[None]'),\n  );\n};\n\nconst findAgentComponents = (\n  components: ComponentMetadata[] | undefined,\n): ComponentMetadata[] =>\n  (components ?? []).filter(\n    (component) => component.generator === PY_AGENT_GENERATOR_INFO.id,\n  );\n\nexport default async function migration(\n  tree: Tree,\n): Promise<MigrationReturnObject> {\n  const nextSteps: string[] = [];\n\n  for (const project of getProjects(tree).values()) {\n    const components = findAgentComponents(\n      (project.metadata as { components?: ComponentMetadata[] })?.components,\n    );\n\n    for (const component of components) {\n      if (!component.path || component.protocol !== 'a2a') {\n        continue;\n      }\n\n      const componentDir = component.path.endsWith('/agent.py')\n        ? component.path.slice(0, -'/agent.py'.length)\n        : component.path;\n      const mainPath = joinPathFragments(project.root, componentDir, 'main.py');\n\n      if ((component.framework ?? 'strands') === 'langchain') {\n        await migrateLangchainA2AAgent(tree, mainPath, nextSteps);\n      } else {\n        await migrateStrandsA2AAgent(tree, mainPath, nextSteps);\n      }\n    }\n  }\n\n  const filePaths: string[] = [];\n  visitNotIgnoredFiles(tree, '', (filePath) => filePaths.push(filePath));\n\n  for (const filePath of filePaths) {\n    if (filePath.endsWith('/with_session_id_strands.py')) {\n      await migrateWithSessionId(tree, filePath, nextSteps);\n    } else if (filePath.endsWith('/session_context.py')) {\n      await migrateSessionContext(tree, filePath, nextSteps);\n    }\n  }\n\n  await formatFilesInSubtree(tree);\n\n  return { nextSteps };\n}\n"],"names":["getProjects","joinPathFragments","visitNotIgnoredFiles","PY_AGENT_GENERATOR_INFO","addPythonDestructuredImport","applyGritQL","captureGritQLVariable","matchGritQL","formatFilesInSubtree","pyMatch","snippet","pyDelete","pyRewrite","from","to","replacement","includes","AGENT_ENTER","OLD_AGENT_CONTEXT","NEW_CONSTANTS_AND_LIFESPAN","A2A_SERVER_STATEMENT","OLD_MOUNT_STATEMENT","OLD_FASTAPI_APP","NEW_FASTAPI_APP","SERVER_MATCH","allMatch","tree","filePath","patterns","pattern","migrateStrandsA2AAgent","mainPath","nextSteps","exists","contents","read","ready","diverged","push","capturedArgs","undefined","LANGCHAIN_GRAPH_OLD","LANGCHAIN_LIFESPAN_NEW","LANGCHAIN_GRAPH_INVOKE_OLD","LANGCHAIN_GRAPH_INVOKE_NEW","LANGCHAIN_APP_OLD","LANGCHAIN_APP_NEW","migrateLangchainA2AAgent","WITH_SESSION_ID_IMPORT_OLD","WITH_SESSION_ID_IMPORT_NEW","migrateWithSessionId","SESSION_CONTEXT_IMPORT_OLD","SESSION_CONTEXT_IMPORT_NEW","migrateSessionContext","findAgentComponents","components","filter","component","generator","id","migration","project","values","metadata","path","protocol","componentDir","endsWith","slice","length","root","framework","filePaths"],"mappings":"AAAA;;;CAGC,GACD,SACEA,WAAW,EACXC,iBAAiB,EAGjBC,oBAAoB,QACf,aAAa;AACpB,SAASC,uBAAuB,QAAQ,iCAA8B;AACtE,SACEC,2BAA2B,EAC3BC,WAAW,EACXC,qBAAqB,EACrBC,WAAW,QACN,wBAAqB;AAC5B,SAASC,oBAAoB,QAAQ,2BAAwB;AAG7D;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CAsCC,GACD,MAAMC,UAAU,CAACC,UAAoB,CAAC,mBAAmB,EAAEA,QAAQ,EAAE,CAAC;AAEtE,MAAMC,WAAW,CAACD,UAAoB,GAAGD,QAAQC,SAAS,KAAK,CAAC;AAEhE,MAAME,YAAY,CAACC,MAAcC;IAC/B,MAAMC,cAAcD,GAAGE,QAAQ,CAAC,QAAQ,CAAC,KAAK,EAAEF,GAAG,EAAE,CAAC,GAAG,CAAC,EAAE,EAAEA,GAAG,EAAE,CAAC;IACpE,OAAO,GAAGL,QAAQI,MAAM,IAAI,EAAEE,aAAa;AAC7C;AAEA,MAAME,cAAc;AAEpB,MAAMC,oBAAoB,CAAC;;;;CAI1B,CAAC;AAEF,MAAMC,6BAA6B,CAAC;;;;;;;;;;;;;;;;;;;;aAoBvB,CAAC;AAEd,MAAMC,uBAAuB;AAE7B,MAAMC,sBAAsB;AAE5B,MAAMC,kBAAkB;AAExB,MAAMC,kBAAkB;AAExB,MAAMC,eACJ;AAEF,MAAMC,WAAW,OACfC,MACAC,UACAC;IAEA,KAAK,MAAMC,WAAWD,SAAU;QAC9B,IAAI,CAAE,MAAMrB,YAAYmB,MAAMC,UAAUE,UAAW,OAAO;IAC5D;IACA,OAAO;AACT;AAEA,MAAMC,yBAAyB,OAC7BJ,MACAK,UACAC;IAEA,IAAI,CAACN,KAAKO,MAAM,CAACF,WAAW;IAE5B,MAAMG,WAAWR,KAAKS,IAAI,CAACJ,UAAU,YAAY;IACjD,IAAI,CAACG,SAASlB,QAAQ,CAAC,kCAAkC;IAEzD,MAAMoB,QAAQ,MAAMX,SAASC,MAAMK,UAAU;QAC3CtB,QAAQ;QACRA,QAAQ;QACRA,QAAQS;QACRT,QAAQQ;QACRR,QAAQa;QACRb,QAAQY;QACRG;KACD;IAED,MAAMa,WAAW;QACfL,UAAUM,IAAI,CACZ,GAAGP,SAAS,sOAAsO,CAAC;IAEvP;IAEA,IAAI,CAACK,OAAO;QACVC;QACA;IACF;IAEA,MAAMjC,4BACJsB,MACAK,UACA;QAAC;KAAsB,EACvB;IAEF,MAAM3B,4BACJsB,MACAK,UACA;QAAC;KAAkB,EACnB;IAEF,MAAM3B,4BAA4BsB,MAAMK,UAAU;QAAC;KAAO,EAAE;IAC5D,MAAM3B,4BAA4BsB,MAAMK,UAAU;QAAC;KAAQ,EAAE;IAC7D,MAAM1B,YACJqB,MACAK,UACAnB,UAAUM,mBAAmBC;IAE/B,MAAMd,YAAYqB,MAAMK,UAAUpB,SAASM;IAE3C,yEAAyE;IACzE,0EAA0E;IAC1E,qDAAqD;IACrD,MAAMZ,YACJqB,MACAK,UACAnB,UACE,gBACA;IAGJ,MAAM2B,eAAe,MAAMjC,sBACzBoB,MACAK,UACA,CAAC,mBAAmB,EAAEX,qBAAqB,EAAE,CAAC,EAC9C;IAEF,IAAImB,iBAAiBC,WAAW;QAC9BH;QACA;IACF;IAEA,MAAMhC,YAAYqB,MAAMK,UAAUpB,SAASS;IAC3C,MAAMf,YACJqB,MACAK,UACAnB,UAAUU,iBAAiBC;IAE7B,MAAMlB,YACJqB,MACAK,UACAnB,UACES,qBACA,CAAC,uBAAuB,EAAEkB,aAAa,KAAK,EAAElB,qBAAqB;AAGzE;AAEA,gFAAgF;AAEhF,MAAMoB,sBAAsB;AAE5B,MAAMC,yBAAyB,CAAC;;;SAGvB,CAAC;AAEV,MAAMC,6BAA6B;AAEnC,MAAMC,6BAA6B;AAEnC,MAAMC,oBAAoB;AAE1B,MAAMC,oBAAoB;AAE1B,MAAMC,2BAA2B,OAC/BrB,MACAK,UACAC;IAEA,IAAI,CAACN,KAAKO,MAAM,CAACF,WAAW;IAE5B,MAAMG,WAAWR,KAAKS,IAAI,CAACJ,UAAU,YAAY;IACjD,IAAI,CAACG,SAASlB,QAAQ,CAACyB,sBAAsB;IAE7C,MAAML,QAAQ,MAAMX,SAASC,MAAMK,UAAU;QAC3CtB,QAAQgC;QACRhC,QAAQ;QACRA,QAAQkC;QACRlC,QAAQoC;KACT;IAED,IAAI,CAACT,OAAO;QACVJ,UAAUM,IAAI,CACZ,GAAGP,SAAS,2NAA2N,CAAC;QAE1O;IACF;IAEA,MAAM3B,4BACJsB,MACAK,UACA;QAAC;KAAsB,EACvB;IAEF,MAAM1B,YACJqB,MACAK,UACAnB,UAAU6B,qBAAqBC;IAEjC,MAAMrC,YACJqB,MACAK,UACAnB,UAAU+B,4BAA4BC;IAExC,MAAMvC,YACJqB,MACAK,UACAnB,UAAUiC,mBAAmBC;AAEjC;AAEA,kFAAkF;AAElF,MAAME,6BACJ;AACF,MAAMC,6BACJ;AAEF,MAAMC,uBAAuB,OAC3BxB,MACAC,UACAK;IAEA,MAAME,WAAWR,KAAKS,IAAI,CAACR,UAAU,YAAY;IACjD,IAAI,CAACO,SAASlB,QAAQ,CAAC,kBAAkB;IAEzC,MAAMoB,QAAQ,MAAMX,SAASC,MAAMC,UAAU;QAC3ClB,QAAQuC;QACRvC,QAAQ;KACT;IAED,IAAI,CAAC2B,OAAO;QACVJ,UAAUM,IAAI,CACZ,GAAGX,SAAS,4QAA4Q,CAAC;QAE3R;IACF;IAEA,MAAMtB,YACJqB,MACAC,UACAf,UAAUoC,4BAA4BC;IAExC,MAAM5C,YACJqB,MACAC,UACAf,UAAU,iBAAiB;AAE/B;AAEA,MAAMuC,6BAA6B;AACnC,MAAMC,6BAA6B;AAEnC,MAAMC,wBAAwB,OAC5B3B,MACAC,UACAK;IAEA,MAAME,WAAWR,KAAKS,IAAI,CAACR,UAAU,YAAY;IACjD,IAAI,CAACO,SAASlB,QAAQ,CAAC,mBAAmB;IAE1C,MAAMoB,QAAQ,MAAMX,SAASC,MAAMC,UAAU;QAC3ClB,QAAQ0C;QACR1C,QAAQ;KACT;IAED,IAAI,CAAC2B,OAAO;QACVJ,UAAUM,IAAI,CACZ,GAAGX,SAAS,yQAAyQ,CAAC;QAExR;IACF;IAEA,MAAMtB,YACJqB,MACAC,UACAf,UAAUuC,4BAA4BC;IAExC,MAAM/C,YACJqB,MACAC,UACAf,UAAU,kBAAkB;AAEhC;AAEA,MAAM0C,sBAAsB,CAC1BC,aAEA,AAACA,CAAAA,cAAc,EAAE,AAAD,EAAGC,MAAM,CACvB,CAACC,YAAcA,UAAUC,SAAS,KAAKvD,wBAAwBwD,EAAE;AAGrE,eAAe,eAAeC,UAC5BlC,IAAU;IAEV,MAAMM,YAAsB,EAAE;IAE9B,KAAK,MAAM6B,WAAW7D,YAAY0B,MAAMoC,MAAM,GAAI;QAChD,MAAMP,aAAaD,oBAChBO,QAAQE,QAAQ,EAA2CR;QAG9D,KAAK,MAAME,aAAaF,WAAY;YAClC,IAAI,CAACE,UAAUO,IAAI,IAAIP,UAAUQ,QAAQ,KAAK,OAAO;gBACnD;YACF;YAEA,MAAMC,eAAeT,UAAUO,IAAI,CAACG,QAAQ,CAAC,eACzCV,UAAUO,IAAI,CAACI,KAAK,CAAC,GAAG,CAAC,YAAYC,MAAM,IAC3CZ,UAAUO,IAAI;YAClB,MAAMjC,WAAW9B,kBAAkB4D,QAAQS,IAAI,EAAEJ,cAAc;YAE/D,IAAI,AAACT,CAAAA,UAAUc,SAAS,IAAI,SAAQ,MAAO,aAAa;gBACtD,MAAMxB,yBAAyBrB,MAAMK,UAAUC;YACjD,OAAO;gBACL,MAAMF,uBAAuBJ,MAAMK,UAAUC;YAC/C;QACF;IACF;IAEA,MAAMwC,YAAsB,EAAE;IAC9BtE,qBAAqBwB,MAAM,IAAI,CAACC,WAAa6C,UAAUlC,IAAI,CAACX;IAE5D,KAAK,MAAMA,YAAY6C,UAAW;QAChC,IAAI7C,SAASwC,QAAQ,CAAC,gCAAgC;YACpD,MAAMjB,qBAAqBxB,MAAMC,UAAUK;QAC7C,OAAO,IAAIL,SAASwC,QAAQ,CAAC,wBAAwB;YACnD,MAAMd,sBAAsB3B,MAAMC,UAAUK;QAC9C;IACF;IAEA,MAAMxB,qBAAqBkB;IAE3B,OAAO;QAAEM;IAAU;AACrB"}

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

@@ -1,6 +1,7 @@
 import logging
 import os
 import uuid
+from contextlib import asynccontextmanager
 
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.apps import A2AStarletteApplication
@@ -23,7 +24,11 @@ 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"]
 
-_graph = get_agent()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.agent = get_agent()
+    yield
 
 
 class _GraphAgentExecutor(AgentExecutor):
@@ -37,7 +42,7 @@ class _GraphAgentExecutor(AgentExecutor):
         # 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 _graph.ainvoke(
+        result = await app.state.agent.ainvoke(
             {"messages": [{"role": "user", "content": context.get_user_input()}]},
             {"configurable": {"thread_id": session_id}},
         )
@@ -77,7 +82,7 @@ class _SessionIdMiddleware(BaseHTTPMiddleware):
             return await call_next(request)
 
 
-app = FastAPI(title="<%= agentNameClassName %>")
+app = FastAPI(title="<%= agentNameClassName %>", lifespan=lifespan)
 app.add_middleware(_SessionIdMiddleware)
 
 

package/src/py/agent/files/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/utils/agent-connection/files/py-core-runtime-config/session_context.py.template

@@ -1,4 +1,4 @@
-from collections.abc import Iterator
+from collections.abc import Generator
 from contextlib import contextmanager
 from contextvars import ContextVar
 
@@ -13,7 +13,7 @@ def get_current_session_id() -> str | None:
 
 
 @contextmanager
-def session_id_context(session_id: str) -> Iterator[None]:
+def session_id_context(session_id: str) -> Generator[None]:
     """Bind *session_id* as the current session for the scope of the block."""
     token = _session_id_var.set(session_id)
     try:

package/src/utils/agent-connection/files/py-core-strands/base/with_session_id_strands.py.template

@@ -1,4 +1,4 @@
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Generator
 from contextlib import AbstractContextManager, ExitStack, contextmanager
 from typing import Any
 
@@ -11,7 +11,7 @@ def with_session_id(
     *,
     name: str,
     description: str,
-) -> Iterator[Any]:
+) -> Generator[Any]:
     """Wrap an agent factory so each session gets its own cached Agent."""
     stack = ExitStack()
     agents: dict[str, Any] = {}