From dea7328db0fa861ca243945ea0525aff84c9cd39 Mon Sep 17 00:00:00 2001 From: Akarsha Sehwag Date: Fri, 26 Jun 2026 10:16:44 -0400 Subject: [PATCH] feat(memory): add multi-region-replication example (#1747) * feat(memory): add multi-region-replication example * feat(memory): update func * chore: ruff formatting --- .../00-multi-region-replication/README.md | 179 +++++++ .../agentcore_replication/__init__.py | 41 ++ .../agentcore_replication/dual_writer.py | 155 ++++++ .../agentcore_replication/stream_consumer.py | 235 +++++++++ .../infra/streaming-stack.yaml | 193 ++++++++ .../lambda/stream_handler.py | 58 +++ .../requirements.txt | 1 + .../scripts/create_memories.py | 96 ++++ .../scripts/deploy_streaming.sh | 121 +++++ .../scripts/enable_streaming.py | 75 +++ .../scripts/full_demo.py | 447 ++++++++++++++++++ 11 files changed, 1601 insertions(+) create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/README.md create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/__init__.py create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/dual_writer.py create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/stream_consumer.py create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/infra/streaming-stack.yaml create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/lambda/stream_handler.py create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/requirements.txt create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/create_memories.py create mode 100755 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/deploy_streaming.sh create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/enable_streaming.py create mode 100644 01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/full_demo.py diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/README.md b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/README.md new file mode 100644 index 00000000..eb6ebd85 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/README.md @@ -0,0 +1,179 @@ +# Cross-Region Replication for Amazon Bedrock AgentCore Memory + +A deployable sample for **customer-driven, cross-region replication** of AgentCore +Memory. It replicates **both** memory layers from a primary region to a replica +region, so you can fail over to the replica with conversation history *and* +extracted knowledge intact. + +| Layer | What it holds | How it replicates | +| --- | --- | --- | +| **LTM** (long-term records) | extracted facts / knowledge | **record streaming** — source memory → Kinesis Data Stream → consumer Lambda → `BatchCreateMemoryRecords` in the target | +| **STM** (short-term events) | raw conversation turns | **dual-write at write time** — `CreateEvent` to the source (normal) and to the target with `extractionMode="SKIP"` | + +## How it works + +``` + Source Region (us-east-1) Target Region (us-west-2) +┌────────────────────────────────┐ ┌──────────────────────────────┐ +│ Memory (primary) │ │ Memory (replica) │ +│ │ │ │ +│ CreateEvent (normal) ──────────┼── STM ─────┼─▶ CreateEvent(extractionMode │ +│ │ dual-write at write time │ dual-write│ ="SKIP") history only, │ +│ ▼ │ │ NO re-extraction │ +│ LTM extraction │ │ │ +│ │ │ │ BatchCreateMemoryRecords │ +│ ▼ │ │ (requestIdentifier = │ +│ Kinesis record stream ─────────┼── LTM ─────┼─▶ source memoryRecordId) ◀── │ +│ (MEMORY_RECORDS/FULL_CONTENT) │ streaming │ consumer Lambda │ +└────────────────────────────────┘ └──────────────────────────────┘ +``` + +Two AgentCore building blocks make this work: + +* **`extractionMode="SKIP"` on `CreateEvent`** — copy events into the target for + history / failover replay **without** re-triggering LTM extraction there. This + is the key to keeping the two paths from colliding: the source's extracted + records already arrive over the stream, so the target must *not* re-extract the + replicated events into duplicate records. +* **Record streaming (`streamDeliveryResources`)** — the source memory publishes + every extracted record to a Kinesis Data Stream, which the consumer replays into + the target. Using the source `memoryRecordId` as the target `requestIdentifier` + makes replays idempotent (re-delivery is a conditional no-op, not a duplicate). + +## Repository layout + +``` +. +├── agentcore_replication/ # reusable Python package +│ ├── stream_consumer.py # LTM: consume Kinesis record stream -> target +│ └── dual_writer.py # STM: dual-write CreateEvent (SKIP on target) +├── lambda/ +│ └── stream_handler.py # Kinesis-triggered LTM consumer (production) +├── infra/ +│ └── streaming-stack.yaml # Kinesis + consumer Lambda + ESM + IAM + DLQ +├── scripts/ +│ ├── create_memories.py # create matching source + target memories +│ ├── enable_streaming.py # turn record streaming on/off (UpdateMemory) +│ ├── full_demo.py # end-to-end STM+LTM demo (the headline sample) +│ └── deploy_streaming.sh # deploy the LTM streaming path +└── requirements.txt +``` + +## Prerequisites + +* AWS CLI v2, configured for an account with AgentCore access in **both** regions. +* Python 3.10+ and `pip install -r requirements.txt` (`boto3 >= 1.43.36`). +* IAM permissions to create CloudFormation stacks, Lambda functions, IAM roles, + Kinesis streams, and S3 buckets in the source region. + +## Run the full demo first + +The fastest way to *see* both layers replicate is the self-contained demo. With +**real AWS resources** it creates two throwaway memories, enables record +streaming on the source, dual-writes a conversation, consumes the Kinesis stream +locally (the same code the Lambda runs), verifies both layers landed in the +target, and proves `extractionMode="SKIP"` prevented re-extraction. It tears +everything down afterward (`--keep` leaves it in place). + +```bash +pip install -r requirements.txt + +python scripts/full_demo.py \ + --source-region us-east-1 \ + --target-region us-west-2 + +# leave the memories/stream/role in place to inspect: +python scripts/full_demo.py --keep + +# extraction + streaming are async; allow more time on a cold account: +python scripts/full_demo.py --extraction-timeout 600 --stream-poll-timeout 600 +``` + +## Build it into your own pipeline + +### 1. Create matching memories in both regions + +```bash +python scripts/create_memories.py \ + --name my-agent-memory \ + --source-region us-east-1 \ + --target-region us-west-2 +# writes memories.json with both IDs +``` + +This writes the source and target memory IDs to `memories.json`. Use those two +IDs in place of the `mem-...` placeholders in steps 2 and 3. + +### 2. STM — dual-write events from your agent + +Replace your single `create_event` call with `DualRegionEventWriter`, which writes +to both regions (target with `extractionMode="SKIP"`): + +```python +from agentcore_replication import DualRegionEventWriter + +writer = DualRegionEventWriter( + source_memory_id="mem-aaaaaaaaaa", + target_memory_id="mem-bbbbbbbbbb", + source_region="us-east-1", + target_region="us-west-2", +) +writer.record_turn(actor_id="user-1", session_id="sess-1", + role="USER", text="I'm vegetarian") +``` + +### 3. LTM — deploy the streaming consumer + +```bash +scripts/deploy_streaming.sh \ + --source-memory-id mem-aaaaaaaaaa \ + --target-memory-id mem-bbbbbbbbbb \ + --source-region us-east-1 \ + --target-region us-west-2 +``` + +This creates the Kinesis stream + consumer Lambda (`infra/streaming-stack.yaml`) +in the **source** region, wires the Event Source Mapping with a DLQ and CloudWatch +alarms, attaches a streaming execution role to the source memory, and enables +record streaming. From then on, extracted LTM records flow source → Kinesis → +Lambda → `BatchCreateMemoryRecords` in the target. + +> **Note:** record streaming only carries records created *after* you enable it. +> Enable streaming before your agents start writing, or backfill pre-existing data +> with `ListMemoryRecords` → `BatchCreateMemoryRecords` (the same idempotent +> `requestIdentifier` path the consumer uses). + +## Failover + +This is active-passive. On primary-region failure, point your agents at the replica +memory in the target region. To replicate the other direction afterward, swap +source/target: enable streaming on the new primary (`enable_streaming.py`), point +your dual-writer the other way, and deploy a consumer in the new source region. +Because LTM writes are idempotent, the first reverse pass safely lands only what's +missing. + +> **Pause replication** during failover with +> `scripts/enable_streaming.py --memory-id --region --disable`. + +## Cost + +* **Kinesis** — one shard in the source region (~$11/mo) plus PUT payload units. +* **Lambda** — invoked per stream batch; pennies at typical memory write rates. +* **AgentCore Memory** storage in the second region. +* **STM dual-write** adds one extra `CreateEvent` per turn — no standing infra. + +## Limitations + +* **Deletes** are not replicated (`MemoryRecordDeleted` events are skipped). + AgentCore consolidation handles stale records in the target; call + `DeleteMemoryRecord` explicitly if you need exact parity. +* **STM latency**: STM is replicated synchronously at write time, so a target + outage surfaces on the write path — wrap `record_turn` to tolerate target + failures (the source write is what your agent depends on). +* **LTM latency**: RPO ≈ extraction time + Kinesis/Lambda lag (seconds). +* **Single AWS account** in this sample. Cross-account requires resource policies / + assumed roles on the target memory and stream. +* **Strategy IDs**: both memories are created with identical strategy config so + namespaces align. The source `memoryStrategyId` is not forwarded (it is + per-memory); AgentCore associates replicated records by namespace. If your + regions differ, map source strategy IDs to target IDs in the consumer. diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/__init__.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/__init__.py new file mode 100644 index 00000000..267305e2 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/__init__.py @@ -0,0 +1,41 @@ +"""Customer-driven cross-region replication for Amazon Bedrock AgentCore Memory. + +Replicates **both** memory layers from a source region to a target region: + +* **LTM (long-term records) via record streaming** — the source memory is + configured with ``streamDeliveryResources`` (``MEMORY_RECORDS`` / + ``FULL_CONTENT``), so extracted records are published to a Kinesis Data Stream. + A consumer (:mod:`stream_consumer`, run as a Lambda or locally) re-creates each + record in the target via ``BatchCreateMemoryRecords``, using the source + ``memoryRecordId`` as the ``requestIdentifier`` so replays are idempotent. +* **STM (short-term events) via dual-write ``CreateEvent``** — + :class:`DualRegionEventWriter` writes every conversation turn to both regions: + the source normally (triggering extraction, which feeds the stream above) and + the target with ``extractionMode="SKIP"`` (history only, no re-extraction). + +``extractionMode="SKIP"`` on the target STM write is what keeps the two paths +from colliding: LTM arrives via the stream, so the target must NOT re-extract the +replicated events into duplicate records. + +See ``README.md`` for deployment instructions. +""" + +from .stream_consumer import ( + StreamStats, + make_target_client, + process_kinesis_records, + replicate_stream_event, + stream_delivery_resources, +) +from .dual_writer import DualRegionEventWriter + +__all__ = [ + # LTM via record streaming + "StreamStats", + "make_target_client", + "process_kinesis_records", + "replicate_stream_event", + "stream_delivery_resources", + # STM via dual-write CreateEvent (extractionMode="SKIP" on target) + "DualRegionEventWriter", +] diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/dual_writer.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/dual_writer.py new file mode 100644 index 00000000..1d7823e9 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/dual_writer.py @@ -0,0 +1,155 @@ +"""STM replication via dual-write ``CreateEvent`` (the ``extractionMode`` story). + +Short-term memory (the raw conversation events) is replicated at write time, not +after the fact: when the agent records a turn, it writes the **same event to both +regions** — + +* **source region** — a normal ``CreateEvent`` (no ``extractionMode``). The event + lands in STM *and* feeds the long-term extraction pipeline, which distills + durable facts into LTM. Those LTM records are what the Kinesis record stream + replicates to the target (see :mod:`stream_consumer`). +* **target region** — ``CreateEvent(extractionMode="SKIP")``. The event lands in + the target's STM for conversation history / failover replay, but is **excluded + from long-term extraction** — because the corresponding LTM records are already + arriving over the stream. Without ``SKIP`` the target would re-extract the same + facts, producing duplicate LTM records and double extraction cost. + +This split is the whole point of the sample: **LTM replicates via streaming, STM +replicates via dual-write CreateEvent, and ``extractionMode="SKIP"`` on the +target is what keeps the two paths from colliding.** + +Cross-region event identity +---------------------------- +AgentCore derives an event's sort key from ``(eventTimestamp, clientToken)``. To +get a stable, idempotent event identity in both regions, pass an explicit +``clientToken`` and the same ``eventTimestamp`` to both writes — don't rely on +the SDK auto-generating either. :class:`DualRegionEventWriter` does this for you. +""" + +import logging +import uuid +from typing import Optional + +import boto3 +from botocore.exceptions import ClientError + +logger = logging.getLogger(__name__) + +_EXTRACTION_MODE_SKIP = "SKIP" + + +class DualRegionEventWriter: + """Write conversation events to a source memory and a target replica. + + Parameters + ---------- + source_memory_id, target_memory_id: + Memory resource IDs in each region. + source_region, target_region: + AWS regions, e.g. ``us-east-1`` / ``us-west-2``. + session: + Optional pre-configured ``boto3.Session``. + """ + + def __init__( + self, + source_memory_id: str, + target_memory_id: str, + source_region: str, + target_region: str, + session: Optional[boto3.Session] = None, + ): + self.source_memory_id = source_memory_id + self.target_memory_id = target_memory_id + session = session or boto3.Session() + self.source = session.client("bedrock-agentcore", region_name=source_region) + self.target = session.client("bedrock-agentcore", region_name=target_region) + + def record_turn( + self, + actor_id: str, + session_id: str, + role: str, + text: str, + event_timestamp: Optional[float] = None, + client_token: Optional[str] = None, + ) -> dict: + """Record one conversation turn in BOTH regions. + + The source write triggers normal LTM extraction; the target write uses + ``extractionMode="SKIP"`` so it stores history only. Returns a dict with + both event responses and the shared ``clientToken`` / ``eventTimestamp`` + used (so callers can correlate the two regions). + """ + if event_timestamp is None: + # A single timestamp shared by both writes keeps the event identity + # aligned across regions. (Date.now-style call is fine here — this is + # the live app path, not a replayable workflow.) + import time + + event_timestamp = time.time() + # Shared, explicit client token => deterministic, idempotent identity in + # both regions (AgentCore keys events on (eventTimestamp, clientToken)). + client_token = client_token or f"turn-{uuid.uuid4().hex}" + + payload = [{"conversational": {"content": {"text": text}, "role": role}}] + + # 1) Source: normal write -> STM + triggers extraction (-> stream -> LTM). + source_resp = self.source.create_event( + memoryId=self.source_memory_id, + actorId=actor_id, + sessionId=session_id, + eventTimestamp=event_timestamp, + payload=payload, + clientToken=client_token, + ) + + # 2) Target: SKIP extraction -> STM history only (LTM comes via stream). + target_resp = self._create_event_idempotent( + self.target, + self.target_memory_id, + actor_id, + session_id, + event_timestamp, + payload, + client_token, + extraction_mode=_EXTRACTION_MODE_SKIP, + ) + + return { + "clientToken": client_token, + "eventTimestamp": event_timestamp, + "source_event": source_resp.get("event", {}), + "target_event": (target_resp or {}).get("event", {}), + } + + @staticmethod + def _create_event_idempotent( + client, + memory_id, + actor_id, + session_id, + event_timestamp, + payload, + client_token, + extraction_mode=None, + ): + """CreateEvent that treats an idempotent "already exists" collision as OK.""" + kwargs = dict( + memoryId=memory_id, + actorId=actor_id, + sessionId=session_id, + eventTimestamp=event_timestamp, + payload=payload, + clientToken=client_token, + ) + if extraction_mode: + kwargs["extractionMode"] = extraction_mode + try: + return client.create_event(**kwargs) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code", "") + if code in ("ConflictException", "ValidationException") and ("exist" in str(exc).lower()): + logger.info("event already exists (idempotent), token=%s", client_token) + return None + raise diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/stream_consumer.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/stream_consumer.py new file mode 100644 index 00000000..8e4a7bfc --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/stream_consumer.py @@ -0,0 +1,235 @@ +"""LTM replication via **record streaming**. + +This is the real-time long-term-memory path. The source memory is configured +with ``streamDeliveryResources`` (``MEMORY_RECORDS`` / ``FULL_CONTENT``) so every +extracted/updated record is published to a Kinesis Data Stream. A consumer reads +those stream events and re-creates each record in the target region with +``BatchCreateMemoryRecords``. + +Idempotency comes from ``requestIdentifier``: we set it to the source record's +``memoryRecordId``, so replaying the same stream record is a conditional no-op in +the target rather than a duplicate write. + +The same core function powers two consumers: + +* ``lambda/stream_handler.py`` — production: a Kinesis Event Source Mapping + invokes the Lambda with a batch of records. +* ``scripts/full_demo.py`` — demo: a local ``GetRecords`` loop feeds the exact + same logic, so the demo exercises the real stream end-to-end without deploying + the Lambda. + +Stream event shape (decoded Kinesis ``data``):: + + { + "memoryStreamEvent": { + "eventType": "MemoryRecordCreated" | "MemoryRecordUpdated" + | "MemoryRecordDeleted" | "StreamingEnabled", + "memoryId": "mem-...", + "memoryRecordId": "mem-rec-...", + "memoryRecordText": "the extracted fact text", + "namespaces": ["/facts/demo-user", ...], + "eventTime": "2026-06-25T12:34:56.789Z", + "memoryStrategyId": "strat-..." # when present + } + } +""" + +import base64 +import json +import logging +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone + +from botocore.exceptions import ClientError + +logger = logging.getLogger(__name__) + +# Stream event types we replicate vs. ignore. +REPLICABLE_EVENTS = {"MemoryRecordCreated", "MemoryRecordUpdated"} +SKIP_EVENTS = {"StreamingEnabled", "MemoryRecordDeleted"} + +# Errors worth retrying (let the ESM redeliver / a local loop re-poll); anything +# else is terminal for that record. +RETRYABLE_ERRORS = {"ThrottledException", "ServiceException", "RetryableConflictException"} + + +def stream_delivery_resources(stream_arn: str) -> dict: + """Build the ``streamDeliveryResources`` payload for full-content LTM records. + + This is the source-side contract for the record stream this module consumes: + pass it to ``UpdateMemory``/``CreateMemory`` to publish every extracted record + (``MEMORY_RECORDS`` / ``FULL_CONTENT``) to the given Kinesis Data Stream. + """ + return { + "resources": [ + { + "kinesis": { + "dataStreamArn": stream_arn, + "contentConfigurations": [{"type": "MEMORY_RECORDS", "level": "FULL_CONTENT"}], + } + } + ] + } + + +@dataclass +class StreamStats: + """Counters returned by a stream-consumption pass.""" + + received: int = 0 + replicated: int = 0 + skipped: int = 0 + failed: int = 0 + errors: list = field(default_factory=list) + + def as_dict(self) -> dict: + return { + "received": self.received, + "replicated": self.replicated, + "skipped": self.skipped, + "failed": self.failed, + "errors": self.errors, + } + + +def _to_epoch(event_time) -> float: + """Normalize a stream ``eventTime`` (ISO-8601 string) to epoch seconds.""" + if isinstance(event_time, (int, float)): + return float(event_time) + if isinstance(event_time, str) and event_time: + try: + dt = datetime.fromisoformat(event_time.replace("Z", "+00:00")) + return dt.timestamp() + except ValueError: + pass + return datetime.now(timezone.utc).timestamp() + + +def replicate_stream_event( + stream_event: dict, + target_client, + target_memory_id: str, +) -> str: + """Replicate a single decoded ``memoryStreamEvent`` to the target region. + + Returns ``"replicated"`` or ``"skipped"``, or raises ``ClientError`` on a + retryable failure so the caller (ESM or local loop) can retry the batch. + """ + event_type = stream_event.get("eventType", "Unknown") + record_id = stream_event.get("memoryRecordId", "") + + if event_type in SKIP_EVENTS or event_type not in REPLICABLE_EVENTS: + # StreamingEnabled is a control event; MemoryRecordDeleted is not + # replicated (consolidation cleans up the target). Unknown types are + # ignored forward-compatibly. + logger.info("skip stream event type=%s id=%s", event_type, record_id) + return "skipped" + + text = stream_event.get("memoryRecordText") + if not text: + logger.warning("skip record %s: no memoryRecordText", record_id) + return "skipped" + + # Preserve namespaces verbatim so vector search behaves identically in both + # regions. This is active-passive one-way replication, so no loop-prevention + # prefix is needed (the target memory does not stream back). + namespaces = stream_event.get("namespaces") or [] + timestamp = _to_epoch(stream_event.get("eventTime")) + + # Use the source memoryRecordId as the requestIdentifier so replays of the + # same stream record are idempotent (the target de-dups on requestIdentifier + # rather than creating a second record). Fall back to a fresh id if the + # stream event somehow lacks one. + # + # NOTE: the source's memoryStrategyId is intentionally NOT forwarded. Strategy + # IDs are generated per-memory, so the source's ID does not exist in the + # target and BatchCreateMemoryRecords would reject it. AgentCore associates + # the replicated record by its namespaces instead. + record = { + "requestIdentifier": record_id or uuid.uuid4().hex, + "content": {"text": text}, + "namespaces": namespaces, + "timestamp": timestamp, + } + + resp = target_client.batch_create_memory_records( + memoryId=target_memory_id, + records=[record], + clientToken=uuid.uuid4().hex, + ) + + failed = resp.get("failedRecords", []) + if failed: + f = failed[0] + code = f.get("errorCode", "") + msg = f"record {record_id}: {code} {f.get('errorMessage')}" + if code in RETRYABLE_ERRORS: + # Raise so the ESM/loop retries the whole batch. + raise ClientError( + {"Error": {"Code": code, "Message": msg}}, + "BatchCreateMemoryRecords", + ) + raise RuntimeError(msg) + + logger.info("replicated record %s -> %s", record_id, target_memory_id) + return "replicated" + + +def process_kinesis_records( + kinesis_records, + target_client, + target_memory_id: str, + stats: "StreamStats | None" = None, +) -> StreamStats: + """Decode and replicate a batch of raw Kinesis records. + + ``kinesis_records`` is a list of dicts each shaped like a Lambda Kinesis + record (``{"kinesis": {"data": ""}}``) or a raw ``GetRecords`` + record (``{"Data": b"..."}``). Both shapes are handled. + + A retryable ``ClientError`` is re-raised (so the ESM retries the batch); + terminal errors are recorded in ``stats`` and skipped. + """ + stats = stats or StreamStats() + for rec in kinesis_records: + stats.received += 1 + try: + raw = _extract_data(rec) + payload = json.loads(raw) + stream_event = payload.get("memoryStreamEvent", payload) + except Exception as exc: # noqa: BLE001 - malformed record, never crash + stats.failed += 1 + stats.errors.append(f"decode: {exc}") + logger.error("malformed stream record: %s", exc) + continue + + try: + outcome = replicate_stream_event(stream_event, target_client, target_memory_id) + if outcome == "replicated": + stats.replicated += 1 + else: + stats.skipped += 1 + except ClientError: + # Retryable: bubble up so the batch is redelivered. + raise + except Exception as exc: # noqa: BLE001 - terminal for this record + stats.failed += 1 + stats.errors.append(str(exc)) + logger.error("replicate failed: %s", exc) + return stats + + +def _extract_data(rec) -> str: + """Return the decoded UTF-8 JSON string from a Kinesis record (either shape).""" + if "kinesis" in rec: # Lambda ESM event shape + return base64.b64decode(rec["kinesis"]["data"]).decode("utf-8") + data = rec.get("Data") # raw GetRecords shape (boto3 returns bytes) + if isinstance(data, (bytes, bytearray)): + return data.decode("utf-8") + return data + + +def make_target_client(session, region_name: str): + """Build a bedrock-agentcore client for the target region.""" + return session.client("bedrock-agentcore", region_name=region_name) diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/infra/streaming-stack.yaml b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/infra/streaming-stack.yaml new file mode 100644 index 00000000..6c5c8856 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/infra/streaming-stack.yaml @@ -0,0 +1,193 @@ +AWSTemplateFormatVersion: "2010-09-09" +Description: > + Cross-region replication for Amazon Bedrock AgentCore Memory — LTM via record + streaming. Deployed in the SOURCE region: a Kinesis Data Stream receives the + source memory's record-stream events; a consumer Lambda (Event Source Mapping) + replicates each record into the TARGET region's memory via + BatchCreateMemoryRecords, preserving the original memoryRecordId. + + STM is replicated separately by the application at write time (dual-write + CreateEvent with extractionMode=SKIP on the target) — see + agentcore_replication.dual_writer. This stack handles only the LTM path. + + After deploying, enable streaming on the source memory with: + scripts/enable_streaming.py --memory-id --region \ + --stream-arn + +Parameters: + TargetMemoryId: + Type: String + Description: Memory resource ID in the TARGET region (replica). + TargetRegion: + Type: String + Description: Target AWS region (where the replica memory lives). + CodeS3Bucket: + Type: String + Description: S3 bucket (in this/source region) holding the packaged Lambda zip. + CodeS3Key: + Type: String + Default: agentcore-stream-replicator.zip + Description: S3 key of the packaged Lambda zip. + KinesisStreamName: + Type: String + Default: agentcore-ltm-stream + Description: Name of the Kinesis Data Stream the source memory streams into. + +Resources: + # 1) Kinesis stream the SOURCE memory delivers MEMORY_RECORDS/FULL_CONTENT into. + MemoryStream: + Type: AWS::Kinesis::Stream + Properties: + Name: !Ref KinesisStreamName + ShardCount: 1 + RetentionPeriodHours: 24 + + # 2) Role AgentCore assumes to PUT records into the stream. Attach its ARN to + # the source memory's memoryExecutionRoleArn (CreateMemory/UpdateMemory). + MemoryStreamingRole: + Type: AWS::IAM::Role + Properties: + RoleName: !Sub "agentcore-memory-streaming-${AWS::Region}" + AssumeRolePolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Principal: + Service: bedrock-agentcore.amazonaws.com + Action: sts:AssumeRole + Policies: + - PolicyName: KinesisPut + PolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Action: + - kinesis:PutRecord + - kinesis:PutRecords + - kinesis:DescribeStream + - kinesis:DescribeStreamSummary + Resource: !GetAtt MemoryStream.Arn + + # 3) DLQ for poison records the consumer can't replicate. + ReplicationDLQ: + Type: AWS::SQS::Queue + Properties: + QueueName: !Sub "agentcore-stream-replication-dlq-${AWS::Region}" + MessageRetentionPeriod: 1209600 # 14 days + VisibilityTimeout: 300 + + # 4) Consumer Lambda execution role: read the stream, write the target memory. + ConsumerRole: + Type: AWS::IAM::Role + Properties: + AssumeRolePolicyDocument: + Version: "2012-10-17" + Statement: + - Effect: Allow + Principal: + Service: lambda.amazonaws.com + Action: sts:AssumeRole + ManagedPolicyArns: + - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole + Policies: + - PolicyName: StreamConsumerPolicy + PolicyDocument: + Version: "2012-10-17" + Statement: + - Sid: KinesisRead + Effect: Allow + Action: + - kinesis:GetRecords + - kinesis:GetShardIterator + - kinesis:DescribeStream + - kinesis:DescribeStreamSummary + - kinesis:ListShards + Resource: !GetAtt MemoryStream.Arn + - Sid: DlqWrite + Effect: Allow + Action: sqs:SendMessage + Resource: !GetAtt ReplicationDLQ.Arn + - Sid: TargetMemoryWrite + Effect: Allow + Action: bedrock-agentcore:BatchCreateMemoryRecords + Resource: + - !Sub "arn:aws:bedrock-agentcore:${TargetRegion}:${AWS::AccountId}:memory/${TargetMemoryId}" + - !Sub "arn:aws:bedrock-agentcore:${TargetRegion}:${AWS::AccountId}:memory/${TargetMemoryId}/*" + + ConsumerFunction: + Type: AWS::Lambda::Function + Properties: + FunctionName: !Sub "agentcore-ltm-stream-replicator-${TargetRegion}" + Runtime: python3.12 + Handler: stream_handler.lambda_handler + Timeout: 120 + MemorySize: 256 + Role: !GetAtt ConsumerRole.Arn + Code: + S3Bucket: !Ref CodeS3Bucket + S3Key: !Ref CodeS3Key + Environment: + Variables: + TARGET_MEMORY_ID: !Ref TargetMemoryId + TARGET_REGION: !Ref TargetRegion + + # 5) Wire the stream to the consumer. + StreamMapping: + Type: AWS::Lambda::EventSourceMapping + Properties: + EventSourceArn: !GetAtt MemoryStream.Arn + FunctionName: !Ref ConsumerFunction + StartingPosition: TRIM_HORIZON + BatchSize: 10 + BisectBatchOnFunctionError: true + MaximumRetryAttempts: 3 + MaximumRecordAgeInSeconds: 3600 + DestinationConfig: + OnFailure: + Destination: !GetAtt ReplicationDLQ.Arn + + ConsumerErrorAlarm: + Type: AWS::CloudWatch::Alarm + Properties: + AlarmName: !Sub "agentcore-ltm-stream-replicator-errors" + AlarmDescription: LTM stream consumer reported errors. + Namespace: AWS/Lambda + MetricName: Errors + Dimensions: + - Name: FunctionName + Value: !Ref ConsumerFunction + Statistic: Sum + Period: 300 + EvaluationPeriods: 1 + Threshold: 1 + ComparisonOperator: GreaterThanOrEqualToThreshold + TreatMissingData: notBreaching + + ReplicationLagAlarm: + Type: AWS::CloudWatch::Alarm + Properties: + AlarmName: !Sub "agentcore-ltm-stream-replication-lag" + AlarmDescription: Kinesis iterator age (replication lag) is high. + Namespace: AWS/Lambda + MetricName: IteratorAge + Dimensions: + - Name: FunctionName + Value: !Ref ConsumerFunction + Statistic: Maximum + Period: 300 + EvaluationPeriods: 1 + Threshold: 900000 # 15 min in ms + ComparisonOperator: GreaterThanThreshold + TreatMissingData: notBreaching + +Outputs: + MemoryStreamArn: + Description: Pass this to scripts/enable_streaming.py --stream-arn + Value: !GetAtt MemoryStream.Arn + MemoryStreamingRoleArn: + Description: Set as the SOURCE memory's memoryExecutionRoleArn + Value: !GetAtt MemoryStreamingRole.Arn + ConsumerFunctionName: + Value: !Ref ConsumerFunction + DlqUrl: + Value: !Ref ReplicationDLQ diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/lambda/stream_handler.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/lambda/stream_handler.py new file mode 100644 index 00000000..a10442c6 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/lambda/stream_handler.py @@ -0,0 +1,58 @@ +"""Kinesis-triggered Lambda: replicate LTM records to the target region. + +This is the production LTM path. An Event Source Mapping invokes this handler +with a batch of Kinesis records published by the SOURCE memory's record stream +(``streamDeliveryResources`` with ``MEMORY_RECORDS`` / ``FULL_CONTENT``). Each +record is re-created in the TARGET region via ``BatchCreateMemoryRecords``, using +the source ``memoryRecordId`` as the ``requestIdentifier`` (idempotent). + +STM is replicated separately by the application at write time (dual-write +``CreateEvent`` with ``extractionMode="SKIP"`` on the target) — see +``agentcore_replication.dual_writer``. This handler only handles LTM. + +Environment variables +---------------------- +TARGET_MEMORY_ID : target (replica) memory resource ID +TARGET_REGION : target AWS region + +A retryable failure raises, so the ESM retries the batch (configure +BisectBatchOnFunctionError + an SQS DLQ on the mapping for poison records). +""" + +import json +import logging +import os + +import boto3 + +from agentcore_replication.stream_consumer import ( + StreamStats, + make_target_client, + process_kinesis_records, +) + +logging.getLogger().setLevel(logging.INFO) +logger = logging.getLogger(__name__) + +TARGET_MEMORY_ID = os.environ["TARGET_MEMORY_ID"] +TARGET_REGION = os.environ["TARGET_REGION"] + +# Build the target client once per container (cold start) and reuse it. +_target_client = make_target_client(boto3.Session(), TARGET_REGION) + + +def lambda_handler(event, context): + stats = StreamStats() + process_kinesis_records( + event.get("Records", []), + target_client=_target_client, + target_memory_id=TARGET_MEMORY_ID, + stats=stats, + ) + result = stats.as_dict() + logger.info("LTM stream replication: %s", json.dumps(result)) + # Retryable errors already raised inside process_kinesis_records (so the ESM + # retries). Terminal per-record failures are reported but don't fail the + # batch — they'd otherwise block the shard forever. Route them to a DLQ via + # the ESM's DestinationConfig.OnFailure if you need to capture them. + return result diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/requirements.txt b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/requirements.txt new file mode 100644 index 00000000..c26e4b57 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/requirements.txt @@ -0,0 +1 @@ +boto3>=1.43.36 diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/create_memories.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/create_memories.py new file mode 100644 index 00000000..fcc1513d --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/create_memories.py @@ -0,0 +1,96 @@ +"""Create matching source + target memories for the replication sample. + +Creates an AgentCore memory in the source region and an identically-configured +memory in the target region, using the bedrock-agentcore-control API. Prints the +two memory IDs, which feed into ``deploy_streaming.sh`` and the dual-writer. + +Both memories use the same strategy config so that namespaces line up across +regions. (Strategy IDs are per-memory and are not forwarded; replicated records +are associated by namespace.) + +Usage:: + + python scripts/create_memories.py \ + --name my-agent-memory \ + --source-region us-east-1 \ + --target-region us-west-2 +""" + +import argparse +import json +import time + +import boto3 +from botocore.exceptions import ClientError + + +def _wait_active(client, memory_id, label, timeout=600): + deadline = time.time() + timeout + while time.time() < deadline: + status = client.get_memory(memoryId=memory_id)["memory"]["status"] + print(f" {label}: {status}") + if status == "ACTIVE": + return + if status == "FAILED": + raise SystemExit(f"{label} memory creation FAILED") + time.sleep(15) + raise SystemExit(f"Timed out waiting for {label} to become ACTIVE") + + +def create_memory(region, name, strategies): + client = boto3.client("bedrock-agentcore-control", region_name=region) + try: + resp = client.create_memory( + name=name, + description="Cross-region replication sample", + memoryStrategies=strategies, + eventExpiryDuration=90, + ) + except ClientError as exc: + raise SystemExit(f"create_memory failed in {region}: {exc}") + memory_id = resp["memory"]["id"] + print(f"Created memory {memory_id} in {region}") + _wait_active(client, memory_id, region) + return memory_id + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--name", required=True) + p.add_argument("--source-region", required=True) + p.add_argument("--target-region", required=True) + p.add_argument( + "--out", + default="memories.json", + help="where to write the resulting memory IDs", + ) + args = p.parse_args() + + # A simple semantic strategy; mirror your real config here. Identical config + # in both regions keeps namespaces / strategy IDs aligned for replication. + strategies = [ + { + "semanticMemoryStrategy": { + "name": "semantic", + "namespaces": ["/facts/{actorId}"], + } + } + ] + + source_id = create_memory(args.source_region, f"{args.name}-source", strategies) + target_id = create_memory(args.target_region, f"{args.name}-target", strategies) + + out = { + "source_memory_id": source_id, + "target_memory_id": target_id, + "source_region": args.source_region, + "target_region": args.target_region, + } + with open(args.out, "w") as f: + json.dump(out, f, indent=2) + print("\nMemory IDs written to", args.out) + print(json.dumps(out, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/deploy_streaming.sh b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/deploy_streaming.sh new file mode 100755 index 00000000..ac31913c --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/deploy_streaming.sh @@ -0,0 +1,121 @@ +#!/usr/bin/env bash +# +# Deploy the LTM record-streaming replication path (Kinesis -> consumer Lambda). +# +# The stack is deployed in the SOURCE region (where the Kinesis stream lives and +# the source memory streams its records). The consumer Lambda writes cross-region +# into the TARGET memory. After deploying, this script also enables streaming on +# the source memory and points it at the new stream. +# +# STM replication is handled by the application at write time (dual-write +# CreateEvent with extractionMode=SKIP) — see agentcore_replication.dual_writer. +# +# Usage: +# scripts/deploy_streaming.sh \ +# --source-memory-id mem-aaaaaaaaaa \ +# --target-memory-id mem-bbbbbbbbbb \ +# --source-region us-east-1 \ +# --target-region us-west-2 \ +# [--memory-execution-role arn:aws:iam::...:role/...] \ +# [--bucket my-deploy-bucket] \ +# [--stream-name agentcore-ltm-stream] +# +set -euo pipefail + +SOURCE_MEMORY_ID="" TARGET_MEMORY_ID="" +SOURCE_REGION="" TARGET_REGION="" +BUCKET="" STREAM_NAME="agentcore-ltm-stream" +MEMORY_EXEC_ROLE="" +STACK_NAME="agentcore-ltm-stream-replicator" + +while [[ $# -gt 0 ]]; do + case "$1" in + --source-memory-id) SOURCE_MEMORY_ID="$2"; shift 2 ;; + --target-memory-id) TARGET_MEMORY_ID="$2"; shift 2 ;; + --source-region) SOURCE_REGION="$2"; shift 2 ;; + --target-region) TARGET_REGION="$2"; shift 2 ;; + --bucket) BUCKET="$2"; shift 2 ;; + --stream-name) STREAM_NAME="$2"; shift 2 ;; + --memory-execution-role) MEMORY_EXEC_ROLE="$2"; shift 2 ;; + --stack-name) STACK_NAME="$2"; shift 2 ;; + *) echo "Unknown arg: $1" >&2; exit 1 ;; + esac +done + +for v in SOURCE_MEMORY_ID TARGET_MEMORY_ID SOURCE_REGION TARGET_REGION; do + if [[ -z "${!v}" ]]; then echo "Missing --${v,,} (replace _ with -)" >&2; exit 1; fi +done + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +ACCOUNT_ID="$(aws sts get-caller-identity --query Account --output text)" +[[ -z "$BUCKET" ]] && BUCKET="agentcore-replicator-${ACCOUNT_ID}-${SOURCE_REGION}" + +echo "==> Account: $ACCOUNT_ID" +echo "==> Source: $SOURCE_MEMORY_ID @ $SOURCE_REGION (stream lives here)" +echo "==> Target: $TARGET_MEMORY_ID @ $TARGET_REGION (replica)" +echo "==> Kinesis stream: $STREAM_NAME" +echo "==> Deploy bucket: $BUCKET" + +# 1. Deploy bucket in the SOURCE region. +if ! aws s3api head-bucket --bucket "$BUCKET" --region "$SOURCE_REGION" 2>/dev/null; then + echo "==> Creating deploy bucket $BUCKET" + if [[ "$SOURCE_REGION" == "us-east-1" ]]; then + aws s3api create-bucket --bucket "$BUCKET" --region "$SOURCE_REGION" + else + aws s3api create-bucket --bucket "$BUCKET" --region "$SOURCE_REGION" \ + --create-bucket-configuration "LocationConstraint=$SOURCE_REGION" + fi +fi + +# 2. Package the consumer Lambda (stream_handler.py + the package). +BUILD="$(mktemp -d)"; trap 'rm -rf "$BUILD"' EXIT +cp "$ROOT/lambda/stream_handler.py" "$BUILD/" +cp -r "$ROOT/agentcore_replication" "$BUILD/agentcore_replication" +ZIP="$BUILD/agentcore-stream-replicator.zip" +( cd "$BUILD" && zip -qr "$ZIP" stream_handler.py agentcore_replication ) +echo "==> Built $(du -h "$ZIP" | cut -f1) package" +CODE_KEY="agentcore-stream-replicator-$(date +%s).zip" +aws s3 cp "$ZIP" "s3://$BUCKET/$CODE_KEY" --region "$SOURCE_REGION" + +# 3. Deploy the stack in the SOURCE region. +echo "==> Deploying stack $STACK_NAME in $SOURCE_REGION" +aws cloudformation deploy \ + --region "$SOURCE_REGION" \ + --stack-name "$STACK_NAME" \ + --template-file "$ROOT/infra/streaming-stack.yaml" \ + --capabilities CAPABILITY_NAMED_IAM \ + --parameter-overrides \ + TargetMemoryId="$TARGET_MEMORY_ID" \ + TargetRegion="$TARGET_REGION" \ + CodeS3Bucket="$BUCKET" \ + CodeS3Key="$CODE_KEY" \ + KinesisStreamName="$STREAM_NAME" + +# 4. Read stack outputs. +get_out() { + aws cloudformation describe-stacks --region "$SOURCE_REGION" \ + --stack-name "$STACK_NAME" \ + --query "Stacks[0].Outputs[?OutputKey=='$1'].OutputValue" --output text +} +STREAM_ARN="$(get_out MemoryStreamArn)" +STREAM_ROLE_ARN="$(get_out MemoryStreamingRoleArn)" +echo "==> Stream ARN: $STREAM_ARN" +echo "==> Streaming role: $STREAM_ROLE_ARN" + +# 5. Ensure the source memory has an execution role that can write to Kinesis. +EXEC_ROLE="${MEMORY_EXEC_ROLE:-$STREAM_ROLE_ARN}" +echo "==> Setting source memory execution role -> $EXEC_ROLE" +aws bedrock-agentcore-control update-memory --region "$SOURCE_REGION" \ + --memory-id "$SOURCE_MEMORY_ID" \ + --memory-execution-role-arn "$EXEC_ROLE" >/dev/null + +# 6. Enable record streaming on the source memory -> Kinesis. +echo "==> Enabling record streaming on $SOURCE_MEMORY_ID" +python "$ROOT/scripts/enable_streaming.py" \ + --memory-id "$SOURCE_MEMORY_ID" \ + --region "$SOURCE_REGION" \ + --stream-arn "$STREAM_ARN" + +echo "==> Done." +echo " LTM: source memory -> Kinesis -> consumer Lambda -> target memory." +echo " STM: have your app dual-write via agentcore_replication.DualRegionEventWriter." diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/enable_streaming.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/enable_streaming.py new file mode 100644 index 00000000..0ef13fd8 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/enable_streaming.py @@ -0,0 +1,75 @@ +"""Enable / disable AgentCore Memory record streaming to a Kinesis stream. + +Record streaming is configured via ``streamDeliveryResources`` on the memory +(set at ``CreateMemory`` or via ``UpdateMemory``). This script flips it on or off +for an existing memory, pointing it at a Kinesis Data Stream ARN with a +``MEMORY_RECORDS`` / ``FULL_CONTENT`` content configuration — exactly what the +record-streaming LTM replication path consumes. + +The memory's execution role (``memoryExecutionRoleArn``) must allow +``kinesis:PutRecord*`` / ``kinesis:DescribeStream`` on the target stream. + +Usage:: + + # turn streaming ON, delivering source LTM records to the Kinesis stream + python scripts/enable_streaming.py \ + --memory-id mem-aaaaaaaaaa \ + --region us-east-1 \ + --stream-arn arn:aws:kinesis:us-east-1:123456789012:stream/agentcore-ltm + + # turn streaming OFF (e.g. when failing over / pausing replication) + python scripts/enable_streaming.py \ + --memory-id mem-aaaaaaaaaa --region us-east-1 --disable +""" + +import argparse +import json +import sys + +import boto3 +from botocore.exceptions import ClientError + +sys.path.insert(0, ".") +from agentcore_replication import stream_delivery_resources + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--memory-id", required=True) + p.add_argument("--region", required=True) + p.add_argument("--stream-arn", help="Kinesis Data Stream ARN (required unless --disable)") + p.add_argument("--disable", action="store_true", help="turn streaming OFF") + args = p.parse_args() + + if not args.disable and not args.stream_arn: + raise SystemExit("--stream-arn is required unless --disable is set") + + ctl = boto3.client("bedrock-agentcore-control", region_name=args.region) + + if args.disable: + # An empty resources list detaches the stream. + resources = {"resources": []} + print(f"Disabling record streaming on {args.memory_id} ({args.region})") + else: + resources = stream_delivery_resources(args.stream_arn) + print(f"Enabling record streaming on {args.memory_id} ({args.region})") + print(f" -> {args.stream_arn} (MEMORY_RECORDS / FULL_CONTENT)") + + try: + resp = ctl.update_memory( + memoryId=args.memory_id, + streamDeliveryResources=resources, + ) + except ClientError as exc: + raise SystemExit(f"update_memory failed: {exc}") + + print( + json.dumps( + {"streamDeliveryResources": resp["memory"].get("streamDeliveryResources", {})}, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/full_demo.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/full_demo.py new file mode 100644 index 00000000..34ce605e --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/scripts/full_demo.py @@ -0,0 +1,447 @@ +"""Full two-region demo: LTM via record streaming + STM via dual-write CreateEvent. + +This is the headline sample. It replicates **both** memory layers across two +regions using the production architecture (not polling): + +* **STM (short-term events)** — replicated at *write time* by dual-writing each + conversation turn with :class:`DualRegionEventWriter`: + - SOURCE region: normal ``CreateEvent`` -> STM + triggers LTM extraction. + - TARGET region: ``CreateEvent(extractionMode="SKIP")`` -> STM history only, + NO re-extraction (the LTM records arrive via the stream instead). +* **LTM (long-term records)** — replicated via **record streaming**: the source + memory streams extracted records to a Kinesis Data Stream; this script consumes + that stream locally (the same code the Lambda runs) and re-creates each record + in the target with its original ``memoryRecordId``. + +End to end: + 1. Create source + target memories (identical semantic-strategy config). + 2. Create a Kinesis stream in the source region and enable record streaming on + the source memory (``streamDeliveryResources`` = MEMORY_RECORDS/FULL_CONTENT). + 3. Drive a conversation through ``DualRegionEventWriter`` (STM replicates now; + source extraction kicks off and will publish LTM records to the stream). + 4. Consume the Kinesis stream and replicate LTM records to the target region. + 5. Verify the target holds the same STM events AND LTM records (same IDs), and + prove ``extractionMode="SKIP"`` kept the target from re-extracting (its LTM + count matches what the stream delivered — it didn't double). + 6. Tear everything down (``--keep`` leaves it in place). + +Usage: + python scripts/full_demo.py --source-region us-east-1 --target-region us-west-2 + python scripts/full_demo.py --keep + python scripts/full_demo.py --extraction-timeout 600 --stream-poll-timeout 600 +""" + +import argparse +import json +import sys +import time +import uuid + +import boto3 +from botocore.exceptions import ClientError + +sys.path.insert(0, ".") +from agentcore_replication import ( + DualRegionEventWriter, + StreamStats, + make_target_client, + process_kinesis_records, + stream_delivery_resources, +) + +ACTOR_ID = "demo-user" +SESSION_ID = "demo-session-1" +NAMESPACE_TEMPLATE = "/facts/{actorId}" + +CONVERSATION = [ + ("USER", "Hi! I'm planning a trip and wanted to set up my travel profile."), + ("ASSISTANT", "Great — tell me your preferences and I'll remember them."), + ("USER", "I'm vegetarian, and I always want a window seat on flights."), + ("ASSISTANT", "Noted: vegetarian meals and window seats."), + ("USER", "I live in Seattle and I prefer morning departures before 10am."), + ("ASSISTANT", "Got it — Seattle home base, morning flights before 10am."), + ("USER", "Also I'm allergic to peanuts, please flag that for meals."), + ("ASSISTANT", "I'll flag the peanut allergy on every booking."), +] + + +def log(msg=""): + print(msg, flush=True) + + +def banner(title): + log() + log("=" * 72) + log(f" {title}") + log("=" * 72) + + +# --------------------------------------------------------------------------- # +# Control-plane helpers +# --------------------------------------------------------------------------- # + + +def wait_active(ctl, mem_id, label, timeout=600): + deadline = time.time() + timeout + while time.time() < deadline: + st = ctl.get_memory(memoryId=mem_id)["memory"]["status"] + log(f" [{label}] {mem_id} status={st}") + if st == "ACTIVE": + return + if st == "FAILED": + raise SystemExit(f"{label} memory FAILED") + time.sleep(15) + raise SystemExit(f"{label} memory not ACTIVE in {timeout}s") + + +def create_memory(ctl, region, name, execution_role_arn=None): + strategies = [{"semanticMemoryStrategy": {"name": "semantic", "namespaces": [NAMESPACE_TEMPLATE]}}] + kwargs = dict( + name=name, + description="STM+LTM cross-region replication demo", + eventExpiryDuration=90, + memoryStrategies=strategies, + ) + if execution_role_arn: + kwargs["memoryExecutionRoleArn"] = execution_role_arn + mem_id = ctl.create_memory(**kwargs)["memory"]["id"] + log(f"Created {name} -> {mem_id} ({region})") + wait_active(ctl, mem_id, region) + return mem_id + + +# --------------------------------------------------------------------------- # +# Kinesis + streaming setup (source region) +# --------------------------------------------------------------------------- # + + +def create_stream(kinesis, name): + try: + kinesis.create_stream(StreamName=name, ShardCount=1) + except kinesis.exceptions.ResourceInUseException: + pass + waiter = kinesis.get_waiter("stream_exists") + waiter.wait(StreamName=name) + arn = kinesis.describe_stream_summary(StreamName=name)["StreamDescriptionSummary"]["StreamARN"] + log(f"Kinesis stream ready: {arn}") + return arn + + +def ensure_memory_stream_role(iam, stream_arn, role_name): + """Create (or reuse) an execution role AgentCore uses to write to Kinesis.""" + assume = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": {"Service": "bedrock-agentcore.amazonaws.com"}, + "Action": "sts:AssumeRole", + } + ], + } + try: + arn = iam.create_role( + RoleName=role_name, + AssumeRolePolicyDocument=json.dumps(assume), + Description="AgentCore memory -> Kinesis streaming (demo)", + )["Role"]["Arn"] + except iam.exceptions.EntityAlreadyExistsException: + arn = iam.get_role(RoleName=role_name)["Role"]["Arn"] + iam.put_role_policy( + RoleName=role_name, + PolicyName="kinesis-put", + PolicyDocument=json.dumps( + { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "kinesis:PutRecord", + "kinesis:PutRecords", + "kinesis:DescribeStream", + "kinesis:DescribeStreamSummary", + ], + "Resource": stream_arn, + } + ], + } + ), + ) + log(f"Memory streaming role ready: {arn}") + # IAM propagation is eventually consistent; give it a moment. + time.sleep(10) + return arn + + +def enable_streaming(ctl, memory_id, stream_arn): + ctl.update_memory( + memoryId=memory_id, + streamDeliveryResources=stream_delivery_resources(stream_arn), + ) + log(f"Record streaming enabled on {memory_id} -> {stream_arn}") + + +def consume_stream(kinesis, stream_name, target_client, target_memory_id, expected, timeout): + """Poll the Kinesis stream and replicate LTM records until `expected` land.""" + log(f"\nConsuming Kinesis stream (up to {timeout}s, expecting ~{expected} records)...") + shard_id = kinesis.describe_stream(StreamName=stream_name)["StreamDescription"]["Shards"][0]["ShardId"] + shard_iter = kinesis.get_shard_iterator(StreamName=stream_name, ShardId=shard_id, ShardIteratorType="TRIM_HORIZON")[ + "ShardIterator" + ] + + stats = StreamStats() + deadline = time.time() + timeout + while time.time() < deadline: + resp = kinesis.get_records(ShardIterator=shard_iter, Limit=100) + shard_iter = resp["NextShardIterator"] + records = resp.get("Records", []) + if records: + process_kinesis_records(records, target_client, target_memory_id, stats=stats) + log( + f" stream: received={stats.received} replicated={stats.replicated} " + f"skipped={stats.skipped} failed={stats.failed}" + ) + if stats.replicated >= expected and expected > 0: + break + time.sleep(5) + return stats + + +# --------------------------------------------------------------------------- # +# Data-plane verification helpers +# --------------------------------------------------------------------------- # + + +def count_events(data, mem_id): + total, token = 0, None + while True: + kw = {"memoryId": mem_id, "actorId": ACTOR_ID, "sessionId": SESSION_ID, "maxResults": 100} + if token: + kw["nextToken"] = token + resp = data.list_events(**kw) + total += len(resp.get("events", [])) + token = resp.get("nextToken") + if not token: + return total + + +def list_records(data, mem_id): + recs, token = [], None + while True: + kw = {"memoryId": mem_id, "namespace": "/", "maxResults": 100} + if token: + kw["nextToken"] = token + resp = data.list_memory_records(**kw) + recs.extend(resp.get("memoryRecordSummaries", [])) + token = resp.get("nextToken") + if not token: + return recs + + +def wait_for_extraction(data, mem_id, label, timeout): + log(f"\nWaiting up to {timeout}s for SOURCE LTM extraction (async)...") + deadline = time.time() + timeout + recs = [] + while time.time() < deadline: + recs = list_records(data, mem_id) + if recs: + time.sleep(20) # let multi-fact extraction finish + return list_records(data, mem_id) + log(f" [{label}] 0 LTM records yet; waiting...") + time.sleep(15) + return recs + + +def show_records(recs, label): + log(f"\n{label}: {len(recs)} LTM record(s)") + for r in recs: + text = (r.get("content") or {}).get("text", "") + log(f" - {r['memoryRecordId']}: {text}") + + +# --------------------------------------------------------------------------- # +# Main +# --------------------------------------------------------------------------- # + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--source-region", default="us-east-1") + p.add_argument("--target-region", default="us-west-2") + p.add_argument( + "--extraction-timeout", type=int, default=420, help="seconds to wait for async source LTM extraction" + ) + p.add_argument("--stream-poll-timeout", type=int, default=420, help="seconds to consume the Kinesis stream") + p.add_argument("--keep", action="store_true", help="don't delete resources") + args = p.parse_args() + + suffix = uuid.uuid4().hex[:8] + src_ctl = boto3.client("bedrock-agentcore-control", region_name=args.source_region) + tgt_ctl = boto3.client("bedrock-agentcore-control", region_name=args.target_region) + kinesis = boto3.client("kinesis", region_name=args.source_region) + iam = boto3.client("iam") + + src_id = tgt_id = None + stream_name = f"agentcore-ltm-demo-{suffix}" + role_name = f"agentcore-mem-stream-demo-{suffix}" + created_stream = created_role = False + + try: + banner("STEP 1 — Create matching memories in both regions") + tgt_id = create_memory(tgt_ctl, args.target_region, f"replDemoTgt{suffix}") + src_id = create_memory(src_ctl, args.source_region, f"replDemoSrc{suffix}") + + banner("STEP 2 — Create Kinesis stream + enable record streaming on SOURCE") + stream_arn = create_stream(kinesis, stream_name) + created_stream = True + role_arn = ensure_memory_stream_role(iam, stream_arn, role_name) + created_role = True + # Re-create source memory WITH the execution role so it can write to + # Kinesis (memoryExecutionRoleArn can't be added by this demo's + # update_memory call surface uniformly, so set it up front via update). + src_ctl.update_memory(memoryId=src_id, memoryExecutionRoleArn=role_arn) + wait_active(src_ctl, src_id, f"{args.source_region}/role") + enable_streaming(src_ctl, src_id, stream_arn) + + banner("STEP 3 — Dual-write the conversation (STM replicates NOW)") + log( + "SOURCE: CreateEvent (normal) -> STM + triggers LTM extraction\n" + 'TARGET: CreateEvent(extractionMode="SKIP") -> STM history only\n' + ) + writer = DualRegionEventWriter(src_id, tgt_id, args.source_region, args.target_region) + for i, (role, text) in enumerate(CONVERSATION): + writer.record_turn(ACTOR_ID, SESSION_ID, role, text, event_timestamp=time.time() + i * 0.001) + log(f" + [{role}] {text}") + + src_data = boto3.client("bedrock-agentcore", region_name=args.source_region) + tgt_data = boto3.client("bedrock-agentcore", region_name=args.target_region) + src_events = count_events(src_data, src_id) + log( + f"\nSource STM: {src_events} events. " + f"Target STM (via SKIP dual-write): {count_events(tgt_data, tgt_id)} events." + ) + + banner("STEP 4 — Wait for SOURCE extraction, then replicate LTM via STREAM") + src_recs = wait_for_extraction(src_data, src_id, args.source_region, args.extraction_timeout) + show_records(src_recs, "SOURCE") + + target_client = make_target_client(boto3.Session(), args.target_region) + stream_stats = consume_stream( + kinesis, stream_name, target_client, tgt_id, expected=len(src_recs), timeout=args.stream_poll_timeout + ) + log(f"\nStream replication stats: {stream_stats.as_dict()}") + + banner("STEP 5 — Verify TARGET (both layers) + prove extractionMode=SKIP") + time.sleep(8) + tgt_events = count_events(tgt_data, tgt_id) + log(f"TARGET STM: {tgt_events} events (source had {src_events}).") + + # Freshly-written LTM records are not queryable instantly (the same async + # indexing lag we wait out on the source). Poll until the records the + # stream replicated become visible, so we measure a real baseline rather + # than a transient 0. + expected_ltm = stream_stats.replicated + log(f"\nWaiting for the {expected_ltm} stream-replicated record(s) to be queryable in the target...") + tgt_recs = [] + deadline = time.time() + 120 + while time.time() < deadline: + tgt_recs = list_records(tgt_data, tgt_id) + if len(tgt_recs) >= expected_ltm: + break + time.sleep(5) + show_records(tgt_recs, "TARGET") + + banner("RESULT") + ok = True + + if src_events and tgt_events >= src_events: + log(f'✅ STM: {tgt_events} events dual-written to target with extractionMode="SKIP".') + else: + ok = False + log(f"❌ STM: expected >= {src_events}, found {tgt_events}.") + + if src_recs: + src_texts = {(r.get("content") or {}).get("text") for r in src_recs} + tgt_texts = [(r.get("content") or {}).get("text") for r in tgt_recs] + tgt_text_set = set(tgt_texts) + + if src_texts.issubset(tgt_text_set): + log(f"✅ LTM: all {len(src_texts)} source records replicated via the Kinesis record stream.") + log( + " (Replays are idempotent: the source memoryRecordId is used " + "as the target requestIdentifier, so re-delivery is a no-op.)" + ) + else: + ok = False + log(f"❌ LTM: missing in target: {src_texts - tgt_text_set}") + + # SKIP proof: the SKIP'd STM events were dual-written back in Step 3, + # minutes ago — long enough that, if SKIP were ignored, the target + # would already have re-extracted them. Re-extraction would surface as + # records BEYOND the ones the stream delivered, i.e. duplicate facts + # (the same text appearing twice). So: no duplicate texts AND the + # target count does not exceed what the stream replicated. + duplicates = len(tgt_texts) != len(tgt_text_set) + if not duplicates and len(tgt_recs) <= expected_ltm: + log( + f"✅ SKIP: target holds exactly the {len(tgt_recs)} " + "stream-replicated record(s), no duplicates — the dual-written " + "events were NOT re-extracted. LTM came only from the stream, " + 'exactly as extractionMode="SKIP" guarantees.' + ) + else: + ok = False + log( + f"❌ SKIP: target has {len(tgt_recs)} records " + f"(expected {expected_ltm}), duplicates={duplicates} — the " + "SKIP'd events appear to have been re-extracted." + ) + else: + log( + "ℹ️ LTM: source produced no records in time; STM path still " + "verified. Re-run with a larger --extraction-timeout." + ) + + log() + log( + "🎉 DEMO PASSED: STM (dual-write) + LTM (record streaming) replicated." + if ok + else "DEMO FAILED — see ❌ lines above." + ) + if not ok: + sys.exit(1) + + finally: + if not args.keep: + log("\nCleaning up...") + for ctl, mid, rg in [(src_ctl, src_id, args.source_region), (tgt_ctl, tgt_id, args.target_region)]: + if mid: + try: + ctl.delete_memory(memoryId=mid) + log(f" deleted memory {mid} ({rg})") + except ClientError as e: + log(f" warn {mid}: {e}") + if created_stream: + try: + kinesis.delete_stream(StreamName=stream_name, EnforceConsumerDeletion=True) + log(f" deleted stream {stream_name}") + except ClientError as e: + log(f" warn stream: {e}") + if created_role: + try: + iam.delete_role_policy(RoleName=role_name, PolicyName="kinesis-put") + iam.delete_role(RoleName=role_name) + log(f" deleted role {role_name}") + except ClientError as e: + log(f" warn role: {e}") + else: + log("\n--keep set; resources left in place:") + log(f" source memory: {src_id} ({args.source_region})") + log(f" target memory: {tgt_id} ({args.target_region})") + log(f" kinesis stream: {stream_name} ({args.source_region})") + log(f" streaming role: {role_name}") + + +if __name__ == "__main__": + main()