@aws/nx-plugin-mcp 1.0.0-rc.70 → 1.0.0-rc.72

package/docs/get_started/tutorials/dungeon-game/1.mdx

@@ -198,6 +198,7 @@ import {
   FunctionProps,
   Tracing,
 } from 'aws-cdk-lib/aws-lambda';
+import { RuntimeConfig } from '../../core/runtime-config.js';
 import {
   AuthorizationType,
   LambdaIntegration,
@@ -213,6 +214,7 @@ import {
   Grant,
 } from 'aws-cdk-lib/aws-iam';
 import {
+  ApiIntegrations,
   IntegrationBuilder,
   RestApiIntegration,
 } from '../../core/api/utils.js';
@@ -230,12 +232,18 @@ type Operations = Procedures<AppRouter>;
  * @template TIntegrations - Map of operation names to their integrations
  */
 export interface GameApiProps<
-  TIntegrations extends Record<Operations, RestApiIntegration>,
+  TIntegrations extends ApiIntegrations<Operations, RestApiIntegration>,
 > {
   /**
    * Map of operation names to their API Gateway integrations
    */
   integrations: TIntegrations;
+  /**
+   * Whether to enable AWS WAFv2 with the default managed ruleset on the API's default stage.
+   *
+   * @default true
+   */
+  enableWaf?: boolean;
 }
 
 /**
@@ -244,7 +252,7 @@ export interface GameApiProps<
  * @template TIntegrations - Map of operation names to their integrations
  */
 export class GameApi<
-  TIntegrations extends Record<Operations, RestApiIntegration>,
+  TIntegrations extends ApiIntegrations<Operations, RestApiIntegration>,
 > extends RestApi<Operations, TIntegrations> {
   private allowedOrigins: readonly string[] = ['*'];
 
@@ -256,7 +264,9 @@ export class GameApi<
    * @returns An IntegrationBuilder with default lambda integrations
    */
   public static defaultIntegrations = (scope: Construct) => {
+    const rc = RuntimeConfig.ensure(scope);
     return IntegrationBuilder.rest({
+      pattern: 'isolated',
       operations: routerToOperations(appRouter),
       defaultIntegrationOptions: {
         runtime: Runtime.NODEJS_LATEST,
@@ -274,6 +284,8 @@ export class GameApi<
       } as FunctionProps,
       buildDefaultIntegration: (op, props: FunctionProps) => {
         const handler = new Function(scope, `GameApi${op}Handler`, props);
+        handler.addEnvironment('RUNTIME_CONFIG_APP_ID', rc.appConfigApplicationId);
+        rc.grantReadAppConfig(handler);
         return {
           handler,
           integration: new LambdaIntegration(handler, {
@@ -439,6 +451,7 @@ The `py#agent` generates these files:
       - agent/
         - main.py entrypoint for your agent in Bedrock AgentCore Runtime
         - agent.py defines an example agent and tools
+        - session.py resolves a SessionManager for persisting conversation state
         - Dockerfile defines the docker image for deployment to AgentCore Runtime
   - common/constructs/
     - src
@@ -454,6 +467,7 @@ from contextlib import contextmanager
 
 from strands import Agent, tool
 from strands_tools import current_time
+from dungeon_adventure_agent_connection import log_model_errors, log_tool_errors
 
 
 @tool
@@ -465,17 +479,18 @@ def subtract(a: int, b: int) -> int:
 def get_agent():
     yield Agent(
         name="StoryAgent",
-        description="StoryAgent Agent",
+        description="StoryAgent Strands Agent",
         system_prompt="""
 You are a mathematical wizard.
 Use your tools for mathematical tasks.
 Refer to tools as your 'spellbook'.
 """,
         tools=[subtract, current_time],
+        hooks=[log_model_errors, log_tool_errors],
     )
 ```
 
-This creates an example Strands agent and defines a subtraction tool.
+This creates an example Strands agent and defines a subtraction tool. `log_model_errors` and `log_tool_errors` are hooks from the shared `dungeon_adventure_agent_connection` project that log model/tool failures instead of letting them fail silently.
 
 ```python
 # agent/main.py
@@ -485,7 +500,7 @@ 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 ag_ui_strands import StrandsAgent, StrandsAgentConfig
 from dungeon_adventure_agent_connection import get_current_session_id, session_id_context
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
@@ -493,6 +508,7 @@ from fastapi.responses import StreamingResponse
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from .agent import get_agent
+from .session import get_session_manager
 
 logging.basicConfig(level=logging.INFO)
 
@@ -506,6 +522,9 @@ async def lifespan(app: FastAPI):
             agent=agent,
             name="StoryAgent",
             description="A Strands Agent exposed via the AG-UI protocol.",
+            # A per-thread session manager, not the template Agent's own, since
+            # AG-UI caches one Strands agent per thread_id.
+            config=StrandsAgentConfig(session_manager_provider=lambda _input_data: get_session_manager()),
         )
         yield
 
@@ -570,14 +589,27 @@ async def ping():
     return {"status": "healthy"}
 ```
 
-This is the entrypoint for the agent. Because we selected `--protocol=ag-ui`, the generator wraps our Strands `Agent` with `StrandsAgent` from [`ag_ui_strands`](https://docs.copilotkit.ai/aws-strands/integration) and mounts it on a FastAPI app that speaks the [AG-UI protocol](https://docs.copilotkit.ai/aws-strands/protocol) — this is what CopilotKit will talk to from the React website. The agent is built inside a `lifespan` handler — not at import time — so container startup, not module import, owns construction, and each AgentCore session gets its own container. The `_SessionIdMiddleware` binds the inbound AgentCore runtime session ID onto a `ContextVar` so any downstream MCP/A2A client we wire up later (e.g. the Inventory MCP server in <Link path="get_started/tutorials/dungeon-game/2">Module 2</Link>) automatically forwards it on its outbound calls. In <Link path="get_started/tutorials/dungeon-game/3">Module 3</Link> we'll also add a `session_manager_provider` so each thread id gets its own `S3SessionManager` and conversation history persists across turns.
+This is the entrypoint for the agent. Because we selected `--protocol=ag-ui`, the generator wraps our Strands `Agent` with `StrandsAgent` from [`ag_ui_strands`](https://docs.copilotkit.ai/aws-strands/integration) and mounts it on a FastAPI app that speaks the [AG-UI protocol](https://docs.copilotkit.ai/aws-strands/protocol) — this is what CopilotKit will talk to from the React website. The agent is built inside a `lifespan` handler — not at import time — so container startup, not module import, owns construction, and each AgentCore session gets its own container. The `_SessionIdMiddleware` binds the inbound AgentCore runtime session ID onto a `ContextVar` so any downstream MCP/A2A client we wire up later (e.g. the Inventory MCP server in <Link path="get_started/tutorials/dungeon-game/2">Module 2</Link>) automatically forwards it on its outbound calls. Since AG-UI caches one Strands agent per `thread_id`, we plug in a `session_manager_provider` rather than the template agent's own session manager — this gives each thread its own `SessionManager` so conversation history persists across turns. That `get_session_manager()` function comes from a generated `session.py` sibling: deployed, it returns a `strands.session.S3SessionManager` backed by an S3 bucket the generator provisions automatically; under `agent-dev` (`LOCAL_DEV=true`) it always returns a `FileSessionManager` writing to a local temp directory instead, regardless of deployed configuration.
 
 ```ts
-// common/constructs/src/app/agents/story-agent.ts
-import { Lazy, Names } from 'aws-cdk-lib';
+// common/constructs/src/app/agents/story-agent/story-agent.ts
+import { Fn, Lazy, Names, RemovalPolicy, Stack } from 'aws-cdk-lib';
 import { Platform } from 'aws-cdk-lib/aws-ecr-assets';
+import { Connections, IConnectable } from 'aws-cdk-lib/aws-ec2';
+import {
+  BlockPublicAccess,
+  Bucket,
+  BucketEncryption,
+} from 'aws-cdk-lib/aws-s3';
+import { Key } from 'aws-cdk-lib/aws-kms';
+import {
+  CfnDelivery,
+  CfnDeliveryDestination,
+  CfnDeliverySource,
+  LogGroup,
+  RetentionDays,
+} from 'aws-cdk-lib/aws-logs';
 import { Construct } from 'constructs';
-import { execSync } from 'child_process';
 import * as path from 'path';
 import * as url from 'url';
 import {
@@ -585,31 +617,177 @@ import {
   ProtocolType,
   Runtime,
   RuntimeProps,
+  RuntimeAuthorizerConfiguration,
 } from 'aws-cdk-lib/aws-bedrockagentcore';
-import { IGrantable, IPrincipal } from 'aws-cdk-lib/aws-iam';
+import {
+  PolicyStatement,
+  Effect,
+  ServicePrincipal,
+  IGrantable,
+  IPrincipal,
+} from 'aws-cdk-lib/aws-iam';
+import { IUserPool, IUserPoolClient } from 'aws-cdk-lib/aws-cognito';
+import { suppressRules } from '../../../core/checkov.js';
+import { RuntimeConfig } from '../../../core/runtime-config.js';
+import { findWorkspaceRoot } from '../../../core/workspace.js';
 
 export type StoryAgentProps = Omit<
   RuntimeProps,
-  'runtimeName' | 'protocolConfiguration' | 'agentRuntimeArtifact'
->;
+  | 'runtimeName'
+  | 'protocolConfiguration'
+  | 'agentRuntimeArtifact'
+  | 'authorizerConfiguration'
+> & {
+  /**
+   * Identity details for Cognito Authentication
+   */
+  identity: {
+    userPool: IUserPool;
+    userPoolClient: IUserPoolClient;
+  };
+  /**
+   * Removal policy for the session bucket holding the agent's conversation
+   * history. Defaults to retaining it so a stack `destroy` doesn't silently
+   * delete session data — set to `RemovalPolicy.DESTROY` for sandbox/CI teardown.
+   *
+   * @default RemovalPolicy.RETAIN
+   */
+  readonly sessionBucketRemovalPolicy?: RemovalPolicy;
+};
 
-export class StoryAgent extends Construct implements IGrantable {
+export class StoryAgent extends Construct implements IGrantable, IConnectable {
   public readonly dockerImage: AgentRuntimeArtifact;
   public readonly agentCoreRuntime: Runtime;
+  /** Default Gateway target name for this agent. */
+  public readonly agentName = 'story-agent';
+  /** Inbound auth — a fronting Gateway uses this to pick its outbound credential. */
+  public readonly auth = 'cognito';
 
-  constructor(scope: Construct, id: string, props?: StoryAgentProps) {
+  constructor(scope: Construct, id: string, props: StoryAgentProps) {
     super(scope, id);
 
-    this.dockerImage = AgentRuntimeArtifact.fromAsset(
-      path.dirname(url.fileURLToPath(new URL(import.meta.url))),
+    const rc = RuntimeConfig.ensure(this);
+
+    // Resolve the bundle output directory containing the Dockerfile and built artifacts
+    const bundleDir = path.join(
+      findWorkspaceRoot(url.fileURLToPath(new URL(import.meta.url))),
+      'dist/packages/story/docker/story-agent',
+    );
+
+    this.dockerImage = AgentRuntimeArtifact.fromAsset(bundleDir, {
+      platform: Platform.LINUX_ARM64,
+    });
+
+    const {
+      identity,
+      sessionBucketRemovalPolicy = RemovalPolicy.RETAIN,
+      ...restProps
+    } = props ?? {};
+
+    const sessionKey = new Key(this, 'SessionKey', {
+      enableKeyRotation: true,
+    });
+
+    // Allow CloudWatch Logs to use the session key for server access log delivery.
+    const stack = Stack.of(this);
+    sessionKey.addToResourcePolicy(
+      new PolicyStatement({
+        effect: Effect.ALLOW,
+        principals: [
+          new ServicePrincipal(`logs.${stack.region}.amazonaws.com`),
+        ],
+        actions: [
+          'kms:Encrypt',
+          'kms:Decrypt',
+          'kms:ReEncrypt*',
+          'kms:GenerateDataKey*',
+          'kms:DescribeKey',
+        ],
+        resources: ['*'],
+        conditions: {
+          ArnLike: {
+            'kms:EncryptionContext:aws:logs:arn': `arn:aws:logs:${stack.region}:${stack.account}:log-group:*`,
+          },
+        },
+      }),
+    );
+
+    const sessionAccessLogs = new LogGroup(this, 'SessionAccessLogs', {
+      retention: RetentionDays.ONE_YEAR,
+      encryptionKey: sessionKey,
+      removalPolicy: RemovalPolicy.DESTROY,
+    });
+
+    const sessionBucket = new Bucket(this, 'SessionBucket', {
+      enforceSSL: true,
+      removalPolicy: sessionBucketRemovalPolicy,
+      encryption: BucketEncryption.KMS,
+      encryptionKey: sessionKey,
+      blockPublicAccess: BlockPublicAccess.BLOCK_ALL,
+    });
+    suppressRules(
+      sessionBucket,
+      ['CKV_AWS_21'],
+      'Session data does not need versioning enabled',
+    );
+    suppressRules(
+      sessionBucket,
+      ['CKV2_AWS_61'],
+      'Lifecycle configuration not required for session data',
+    );
+    suppressRules(
+      sessionBucket,
+      ['CKV_AWS_144'],
+      'Cross-region replication not required for session data',
+    );
+    suppressRules(
+      sessionBucket,
+      ['CKV2_AWS_62'],
+      'Event notifications not required for session data',
+    );
+    suppressRules(
+      sessionBucket,
+      ['CKV_AWS_18'],
+      'Server access logs are delivered to CloudWatch Logs',
+    );
+
+    const sessionAccessLogsSource: CfnDeliverySource = new CfnDeliverySource(
+      this,
+      'SessionAccessLogsSource',
       {
-        platform: Platform.LINUX_ARM64,
-        extraHash: execSync(
-          `docker inspect dungeon-adventure-story-agent:latest --format '{{.Id}}'`,
-          { encoding: 'utf-8' },
-        ).trim(),
+        name: Lazy.string({
+          produce: () =>
+            Names.uniqueResourceName(sessionAccessLogsSource, {
+              maxLength: 60,
+            }),
+        }),
+        logType: 'S3_SERVER_ACCESS_LOGS',
+        resourceArn: sessionBucket.bucketArn,
       },
     );
+    const sessionBucketPolicy = sessionBucket.policy;
+    if (sessionBucketPolicy) {
+      sessionAccessLogsSource.node.addDependency(sessionBucketPolicy);
+    }
+    const sessionAccessLogsDestination: CfnDeliveryDestination =
+      new CfnDeliveryDestination(this, 'SessionAccessLogsDestination', {
+        name: Lazy.string({
+          produce: () =>
+            Names.uniqueResourceName(sessionAccessLogsDestination, {
+              maxLength: 60,
+            }),
+        }),
+        destinationResourceArn: sessionAccessLogs.logGroupArn,
+      });
+    const sessionAccessLogsDelivery = new CfnDelivery(
+      this,
+      'SessionAccessLogsDelivery',
+      {
+        deliverySourceName: sessionAccessLogsSource.name,
+        deliveryDestinationArn: sessionAccessLogsDestination.attrArn,
+      },
+    );
+    sessionAccessLogsDelivery.addDependency(sessionAccessLogsSource);
 
     this.agentCoreRuntime = new Runtime(this, 'StoryAgent', {
       runtimeName: Lazy.string({
@@ -618,17 +796,88 @@ export class StoryAgent extends Construct implements IGrantable {
       }),
       protocolConfiguration: ProtocolType.HTTP,
       agentRuntimeArtifact: this.dockerImage,
-      ...props,
+      authorizerConfiguration: RuntimeAuthorizerConfiguration.usingCognito(
+        identity.userPool,
+        [identity.userPoolClient],
+      ),
+      // Receive the caller's Authorization header (validated by the authorizer).
+      requestHeaderConfiguration: {
+        allowlistedHeaders: ['Authorization'],
+      },
+      ...restProps,
+      environmentVariables: {
+        RUNTIME_CONFIG_APP_ID: rc.appConfigApplicationId,
+        ...restProps?.environmentVariables,
+      },
+    });
+
+    // Grant access for the agent to invoke bedrock models
+    this.agentCoreRuntime.addToRolePolicy(
+      new PolicyStatement({
+        actions: [
+          'bedrock:InvokeModel',
+          'bedrock:InvokeModelWithResponseStream',
+        ],
+        resources: [
+          'arn:aws:bedrock:*:*:foundation-model/*',
+          'arn:aws:bedrock:*:*:inference-profile/*',
+        ],
+      }),
+    );
+
+    sessionBucket.grantReadWrite(this.agentCoreRuntime);
+
+    rc.grantReadAppConfig(this.agentCoreRuntime);
+
+    rc.set('agentcore', 'agentRuntimes', {
+      ...rc.get('agentcore').agentRuntimes,
+      StoryAgent: {
+        arn: this.agentCoreRuntime.agentRuntimeArn,
+        session: {
+          bucketName: sessionBucket.bucketName,
+        },
+      },
+    });
+
+    rc.set('connection', 'agentRuntimes', {
+      ...rc.get('connection').agentRuntimes,
+      StoryAgent: this.agentCoreRuntime.agentRuntimeArn,
     });
   }
 
+  /**
+   * The principal to grant permissions to.
+   */
   public get grantPrincipal(): IPrincipal {
     return this.agentCoreRuntime.grantPrincipal;
   }
+
+  /**
+   * Network connections for this agent runtime.
+   */
+  public get connections(): Connections {
+    return this.agentCoreRuntime.connections;
+  }
+
+  /**
+   * The HTTPS invocation URL of the runtime.
+   */
+  public get invocationUrl(): string {
+    // The URL must URL-encode the runtime ARN (':' -> '%3A', '/' -> '%2F').
+    // The ARN is a CDK token, so encode at deploy time via Fn.join/Fn.split.
+    const encodedArn = Fn.join(
+      '%2F',
+      Fn.split(
+        '/',
+        Fn.join('%3A', Fn.split(':', this.agentCoreRuntime.agentRuntimeArn)),
+      ),
+    );
+    return `https://bedrock-agentcore.${Stack.of(this).region}.amazonaws.com/runtimes/${encodedArn}/invocations?qualifier=DEFAULT`;
+  }
 }
 ```
 
-This configures a CDK `AgentRuntimeArtifact` which uploads your agent Docker image to ECR, and hosts it using AgentCore Runtime.
+This configures a CDK `AgentRuntimeArtifact` which uploads your agent Docker image to ECR, and hosts it using AgentCore Runtime. Because we chose `--auth=cognito`, the construct requires the user pool/client identity and authorizes AgentCore Runtime invocations through Cognito, forwarding the caller's `Authorization` header. It also provisions the session bucket the Story Agent's `session.py` reads back at runtime — a KMS-encrypted S3 bucket with server access logs delivered to CloudWatch Logs — grants the agent read/write access to it, grants it access to invoke Bedrock models, and registers its ARN and bucket name in `RuntimeConfig` so both the agent (at runtime, via AppConfig) and the Game API (at synth time, via `invocationUrl`) can find it.
 
 You may notice an extra `Dockerfile`, that references the Docker image from the `story` project, allowing us to co-locate the Dockerfile and agent source code.
 

package/docs/get_started/tutorials/dungeon-game/2.mdx

@@ -146,6 +146,14 @@ Delete the `tools` and `resources` directories in `packages/inventory/src/mcp-se
 
 ## Task 3: Update the infrastructure
 
+### Expose the session bucket
+
+The Story Agent's construct already provisions and writes to its own session bucket internally, but doesn't expose it — so nothing outside the agent can be granted access to read from it yet. Since `queryActions` needs to read the conversation history back, for illustration and simplicity we'll expose its internal session bucket as a public property on `common/constructs/src/app/agents/story-agent/story-agent.ts`:
+
+<E2EDiff lang="typescript" before="dungeon-adventure/2/agents/story-agent.ts.old.template" after="dungeon-adventure/2/agents/story-agent.ts.template" />
+
+### Wire up the stack
+
 The `DungeonDb` construct generated by `ts#dynamodb` already provisions our table, so we just need to instantiate it in our stack and grant the Game API and Inventory MCP server the permissions they need. Update `packages/infra/src/stacks/application-stack.ts` as follows:
 
 <E2EDiff lang="typescript" before="dungeon-adventure/1/application-stack.ts.template" after="dungeon-adventure/2/stacks/application-stack.ts.template" />

package/docs/get_started/tutorials/dungeon-game/3.mdx

@@ -3,7 +3,7 @@ title: Implement and configure the Story agent
 description: A walkthrough of how to build an agentic AI-powered dungeon adventure game using the @aws/nx-plugin.
 ---
 
-import { Aside, Code, FileTree, Steps, Tabs, TabItem } from '@astrojs/starlight/components';
+import { Aside, Code, FileTree, Steps } from '@astrojs/starlight/components';
 import { Image } from 'astro:assets';
 import Drawer from '@components/drawer.astro';
 import Link from '@components/link.astro';
@@ -23,25 +23,15 @@ import gameConversationPng from '@assets/game-conversation.png'
 
 ## Task 1: Implement the Story Agent
 
-The Story Agent is a [Strands](https://strandsagents.com/) agent generated with `--protocol=ag-ui` in <Link path="get_started/tutorials/dungeon-game/1">Module 1</Link>, so the UI can stream from it over the [Agent-User Interaction protocol](https://docs.copilotkit.ai/aws-strands/protocol) via CopilotKit. It uses the Inventory MCP Server to manage the player's items, and Strands' built-in [`S3SessionManager`](https://strandsagents.com/latest/documentation/docs/user-guide/concepts/agents/sessions/) to persist conversation history into the sessions bucket we provisioned in Module 2.
+The Story Agent is a [Strands](https://strandsagents.com/) agent generated with `--protocol=ag-ui` in <Link path="get_started/tutorials/dungeon-game/1">Module 1</Link>, so the UI can stream from it over the [Agent-User Interaction protocol](https://docs.copilotkit.ai/aws-strands/protocol) via CopilotKit. It uses the Inventory MCP Server to manage the player's items, and already persists conversation history via the `session.py` the generator provisioned in Module 1 — deployed, that's Strands' built-in [`S3SessionManager`](https://strandsagents.com/latest/documentation/docs/user-guide/concepts/agents/sessions/) writing into the same S3 bucket the Game API's `queryActions` reads from, so the browser can rebuild transcripts on revisit.
 
 ### Agent implementation
 
-Update the following files in `packages/story/dungeon_adventure_story/agent`:
+Update `agent.py` in `packages/story/dungeon_adventure_story/agent`:
 
-<Tabs>
-  <TabItem label="main.py">
-<E2EDiff lang="python" before="dungeon-adventure/3/main.py.old.template" after="dungeon-adventure/3/main.py.template" />
-  </TabItem>
-  <TabItem label="agent.py">
 <E2EDiff lang="python" before="dungeon-adventure/3/agent.py.old.template" after="dungeon-adventure/3/agent.py.template" />
-  </TabItem>
-</Tabs>
 
-The changes are:
-
-- `main.py` adds a `session_manager_provider` that creates an `S3SessionManager` per `thread_id` when deployed, and falls back to an on-disk `FileSessionManager` under `/tmp/strands-sessions` when running under `agent-dev` (`LOCAL_DEV=true`). Deployed, the S3 bucket is the same one the Game API's `queryActions` reads from, so the browser can rebuild transcripts on revisit; locally the agent persists sessions to disk and talks to the local MCP server, without a deployment.
-- `agent.py` drops the sample `subtract` tool and swaps the system prompt for a dungeon-master one that invites the first user message to state the player's name and genre, and uses the Inventory MCP Server's tools.
+`agent.py` drops the sample `subtract` tool and swaps the system prompt for a dungeon-master one that invites the first user message to state the player's name and genre, and uses the Inventory MCP Server's tools. `main.py` needs no changes here — session management was already wired up when the agent was generated.
 
 ## Task 2: Test your Agent locally
 

package/docs/guides/py-agent.mdx

@@ -50,6 +50,7 @@ The generator will add the following files to your existing Python project. The
         - \_\_init\_\_.py Python package initialization
         - init.py FastAPI application setup with CORS and error handling middleware
         - agent.py Main agent definition with sample tools
+        - session.py Resolves the framework-specific session persistence implementation
         - main.py FastAPI entry point for Bedrock AgentCore Runtime
         - Dockerfile Entry point for hosting your agent (excluded when `infra` is set to `None`)
     - pyproject.toml Updated with Strands dependencies
@@ -68,6 +69,7 @@ The entry point exposes your agent over the A2A protocol (Strands uses the [Stra
       - agent/ (or custom name if specified)
         - \_\_init\_\_.py Python package initialization
         - agent.py Main agent definition with sample tools
+        - session.py Resolves the framework-specific session persistence implementation
         - main.py A2A server entry point
         - Dockerfile Entry point for hosting your agent (excluded when `infra` is set to `None`)
     - pyproject.toml Updated with framework and A2A dependencies
@@ -86,6 +88,7 @@ The entry point exposes your agent via the [AG-UI](https://docs.ag-ui.com/) prot
       - agent/ (or custom name if specified)
         - \_\_init\_\_.py Python package initialization
         - agent.py Main agent definition with sample tools
+        - session.py Resolves the framework-specific session persistence implementation
         - main.py AG-UI server entry point
         - Dockerfile Entry point for hosting your agent (excluded when `infra` is set to `None`)
     - pyproject.toml Updated with framework and AG-UI dependencies
@@ -486,6 +489,39 @@ Your agent is automatically configured with observability using the [AWS Distro
 You can find traces in the CloudWatch AWS Console, by selecting "GenAI Observability" in the menu. Note that for traces to be populated you will need to enable [Transaction Search](https://docs.aws.amazon.com/AmazonCloudWatch/latest/monitoring/CloudWatch-Transaction-Search.html).
 
 For more details, refer to the [AgentCore documentation on observability](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/observability-configure.html).
+
+### Session Management
+
+The `session` option maps to a different underlying persistence concept depending on your chosen framework: Strands' [session management](https://strandsagents.com/docs/user-guide/concepts/agents/session-management/) concept for the `strands` framework, or LangGraph's [checkpointer](https://docs.langchain.com/oss/python/langgraph/persistence) concept for the `langchain` framework.
+
+<OptionFilter when={{ framework: 'strands' }} description="Strands session management (session.py)">
+The `session` option controls how your agent persists conversation state (message history, tool state, etc.) across invocations, using the Strands SDK's [SessionManager](https://strandsagents.com/docs/user-guide/concepts/agents/session-management/):
+
+- **`s3`** (default): The CDK/Terraform infrastructure provisions a dedicated S3 bucket for session data, encrypted with a dedicated KMS key and with all public access blocked; server access logs are delivered to a CloudWatch Logs log group via the same key. The agent's IAM role is granted read/write/list/delete access to the bucket and decrypt/generate-data-key access to the key, and the bucket name is registered alongside the agent's ARN in AppConfig runtime configuration.
+- **`in-memory`**: No bucket is provisioned. Conversation state is kept in memory only for the lifetime of the running process and does not survive restarts or scale-in.
+
+This is implemented in the generated `session.py`, which exports a `get_session_manager()` function resolving a `SessionManager` for the current session.
+
+The session ID itself comes from the AgentCore Runtime session (propagated via the `x-amzn-bedrock-agentcore-runtime-session-id` header) and is bound to a [`contextvars.ContextVar`](https://docs.python.org/3/library/contextvars.html)-based context so `get_current_session_id()` can resolve it anywhere in the request — including in any downstream MCP or A2A clients wired up via the <Link path="/guides/connection">`connection` generator</Link>, so the whole call chain shares a consistent session.
+
+:::note[Local Development]
+When running locally (`LOCAL_DEV=true`, set automatically by the `-dev` target), session data is always stored on disk under `tmp/agents/strands/<agent-name>` at the workspace root, regardless of the configured `session` option, for convenience.
+:::
+</OptionFilter>
+
+<OptionFilter when={{ framework: 'langchain' }} description="LangChain session management (session.py)">
+The `session` option controls how your agent's LangGraph checkpointer persists conversation state:
+
+- **`s3`** (default): The deployed agent uses an `S3CheckpointSaver` with the provisioned session bucket, storing checkpoints and pending writes under the `checkpoints/` prefix. This class lives in `s3_checkpoint_saver_langchain.py` in the shared agent-connection project.
+- **`dynamodb-s3`**: The CDK/Terraform infrastructure provisions a DynamoDB table for checkpoints, configured as recommended in the [AWS documentation on using DynamoDB as a checkpoint store for LangGraph agents](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ddb-langgraph-checkpoint.html) (unified `PK`/`SK` schema, `PAY_PER_REQUEST` billing, point-in-time recovery, and a `ttl` attribute), plus an S3 bucket for offloading checkpoints over 350KB. Both are encrypted with a dedicated KMS key; the bucket's server access logs are delivered to a CloudWatch Logs log group via the same key. The agent's IAM role is granted read/write access to the table and bucket, and the table/bucket names are registered alongside the agent's ARN in AppConfig runtime configuration.
+- **`in-memory`**: No table or bucket is provisioned. Conversation state is kept in memory only for the lifetime of the running process and does not survive restarts or scale-in.
+
+This is implemented in the generated `session.py`, which exports a `get_checkpointer()` function called from `agent.py`'s `create_agent(..., checkpointer=get_checkpointer())`.
+
+:::note[Local Development]
+When running locally (`LOCAL_DEV=true`, set automatically by the `-dev` target), `get_checkpointer()` always returns an [`AsyncSqliteSaver`](https://reference.langchain.com/python/langgraph.checkpoint.sqlite/aio/AsyncSqliteSaver) backed by a local SQLite database under `tmp/agents/langchain/<agent-name>` at the workspace root, regardless of the configured `session` option, for convenience.
+:::
+</OptionFilter>
 </OptionFilter>
 
 ## Invoking your Agent

package/docs/guides/ts-agent.mdx

@@ -360,14 +360,6 @@ The `session` option controls how your agent persists conversation state (messag
 
 This is implemented in the generated `session.ts`, which exports a `getSessionManager()` function resolving a `SessionManager` for the current session.
 
-<OptionFilter when={{ protocol: 'ag-ui' }} description="AG-UI session wiring">
-Since AG-UI clones a template agent per conversation thread, `getSessionManager` is wired in as a `sessionManagerProvider` on the `StrandsAgent` config in `index.ts`, so a fresh `SessionManager` is resolved for each thread rather than baked into the shared template agent.
-</OptionFilter>
-
-<OptionFilter when={{ protocol: ['http', 'a2a'] }} description="HTTP/A2A session wiring">
-Since `withSessionId` already caches one `Agent` instance per session, `getSessionManager` is called directly inside `getAgent`'s `new Agent({ sessionManager: await getSessionManager() })` call in `agent.ts`.
-</OptionFilter>
-
 The session ID itself comes from the AgentCore Runtime session (propagated via the `x-amzn-bedrock-agentcore-runtime-session-id` header for A2A/AG-UI, or the WebSocket connection context for HTTP/tRPC) and is bound to an [`AsyncLocalStorage`](https://nodejs.org/api/async_context.html#class-asynclocalstorage)-based context so `getCurrentSessionId()` can resolve it anywhere in the request — including in any downstream MCP or A2A clients wired up via the <Link path="/guides/connection">`connection` generator</Link>, so the whole call chain shares a consistent session.
 
 :::note[Local Development]

package/package.json

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

package/src/py/agent/schema.json

@@ -63,10 +63,10 @@
     },
     "session": {
       "type": "string",
-      "description": "The storage used to persist session for your Agent. Only 'in-memory' is currently supported for Python agents.",
+      "description": "The storage used to persist session for your Agent. LangChain supports 's3' or 'dynamodb-s3'; Strands supports 's3'; 'in-memory' is valid for both.",
       "x-prompt": "How would you like to persist session for your Agent?",
-      "enum": ["in-memory"],
-      "default": "in-memory",
+      "enum": ["s3", "dynamodb-s3", "in-memory"],
+      "default": "s3",
       "x-priority": "important"
     },
     "preferInstallDependencies": {