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 index a7ce115ea..58d09960e 100644 --- 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 @@ -21,9 +21,9 @@ extracted knowledge intact. │ │ dual-write at write time │ dual-write│ ="SKIP") history only, │ │ ▼ │ │ NO re-extraction │ │ LTM extraction │ │ │ -│ │ │ │ BatchCreateMemoryRecords │ -│ ▼ │ │ (requestIdentifier = │ -│ Kinesis record stream ─────────┼── LTM ─────┼─▶ source memoryRecordId) ◀── │ +│ │ │ │ Batch Create/Update/Delete │ +│ ▼ │ │ MemoryRecords, matched by │ +│ Kinesis record stream ─────────┼── LTM ─────┼─▶ source record ID ◀── │ │ (MEMORY_RECORDS/FULL_CONTENT) │ streaming │ consumer Lambda │ └────────────────────────────────┘ └──────────────────────────────┘ ``` @@ -36,9 +36,10 @@ Two AgentCore building blocks make this work: 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). + every record change to a Kinesis Data Stream, which the consumer applies to the + target: creates, updates, and deletes. Creates send a `clientToken` derived from + the stream event, so re-delivery can't duplicate a record, and a DynamoDB record + map tracks which target record each source record became. ## Repository layout @@ -46,6 +47,7 @@ Two AgentCore building blocks make this work: . ├── agentcore_replication/ # reusable Python package │ ├── stream_consumer.py # LTM: consume Kinesis record stream -> target +│ ├── record_map.py # LTM: source -> target record ID map (DynamoDB) │ └── dual_writer.py # STM: dual-write CreateEvent (SKIP on target) ├── lambda/ │ └── stream_handler.py # Kinesis-triggered LTM consumer (production) @@ -56,6 +58,7 @@ Two AgentCore building blocks make this work: │ ├── 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 +├── tests/ # unit tests: python -m pytest tests └── requirements.txt ``` @@ -136,12 +139,12 @@ 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. +Lambda → the target (create, update, or delete). > **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). +> with `ListMemoryRecords` → `BatchCreateMemoryRecords`, and write each copy to the +> record map so that later updates and deletes find it. ## Failover @@ -149,8 +152,6 @@ This is active-passive. On primary-region failure, point your agents at the repl 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`. @@ -164,9 +165,9 @@ missing. ## 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. +* **Record map coverage**: updates and deletes rely on the record map. A record + that reached the target without a map entry (for example, an unrecorded backfill) + is re-created on its next update, and its delete is skipped. * **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). 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 index 267305e2e..7712b3bb3 100644 --- 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 @@ -5,9 +5,10 @@ * **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. + A consumer (:mod:`stream_consumer`, run as a Lambda or locally) applies each + create, update, and delete to the target. A :mod:`record_map` tracks which + target record each source record became, and creates send a ``clientToken`` + derived from the stream event, so replays don't duplicate records. * **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 @@ -20,6 +21,8 @@ See ``README.md`` for deployment instructions. """ +from .dual_writer import DualRegionEventWriter +from .record_map import DynamoDBRecordMap, InMemoryRecordMap, MappedRecord, RecordMap from .stream_consumer import ( StreamStats, make_target_client, @@ -27,15 +30,18 @@ replicate_stream_event, stream_delivery_resources, ) -from .dual_writer import DualRegionEventWriter __all__ = [ + # STM via dual-write CreateEvent (extractionMode="SKIP" on target) + "DualRegionEventWriter", # LTM via record streaming + "DynamoDBRecordMap", + "InMemoryRecordMap", + "MappedRecord", + "RecordMap", "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 index 1d7823e9f..fcfdef09e 100644 --- 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 @@ -24,11 +24,14 @@ 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. + +To retry a turn, pass the same ``client_token`` and ``event_timestamp`` again (for +example, values derived from your own turn ID). If you let ``record_turn`` +generate them, a retry after a timeout writes the turn a second time. """ import logging import uuid -from typing import Optional import boto3 from botocore.exceptions import ClientError @@ -57,7 +60,7 @@ def __init__( target_memory_id: str, source_region: str, target_region: str, - session: Optional[boto3.Session] = None, + session: boto3.Session | None = None, ): self.source_memory_id = source_memory_id self.target_memory_id = target_memory_id @@ -71,8 +74,8 @@ def record_turn( session_id: str, role: str, text: str, - event_timestamp: Optional[float] = None, - client_token: Optional[str] = None, + event_timestamp: float | None = None, + client_token: str | None = None, ) -> dict: """Record one conversation turn in BOTH regions. @@ -135,14 +138,14 @@ def _create_event_idempotent( 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, - ) + kwargs = { + "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: diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/record_map.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/record_map.py new file mode 100644 index 000000000..7f993d000 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/agentcore_replication/record_map.py @@ -0,0 +1,73 @@ +"""Source-to-target record ID map for LTM replication. + +The target memory assigns its own ``memoryRecordId`` to every record that the +consumer creates, and a ``MemoryRecordDeleted`` stream event carries only the +source IDs. To apply updates and deletes, the consumer has to remember which +target record each source record became, and the time of the last change that +it applied (so a redelivered or out-of-order event is skipped, not re-applied). + +Two implementations share one small interface: + +* :class:`DynamoDBRecordMap` — the Lambda path, backed by the table in + ``infra/streaming-stack.yaml``. +* :class:`InMemoryRecordMap` — the local demo and the unit tests. +""" + +from dataclasses import dataclass +from decimal import Decimal +from typing import Protocol + + +@dataclass(frozen=True) +class MappedRecord: + """What the consumer knows about one replicated source record.""" + + target_record_id: str | None # None for a tombstone of a never-replicated record + event_time: float # epoch seconds of the last change applied + deleted: bool = False + + +class RecordMap(Protocol): + def get(self, source_record_id: str) -> MappedRecord | None: ... + + def put(self, source_record_id: str, record: MappedRecord) -> None: ... + + +class InMemoryRecordMap: + """Process-local map: fine for a demo loop, not for Lambda (each container has its own).""" + + def __init__(self): + self._records = {} + + def get(self, source_record_id: str) -> MappedRecord | None: + return self._records.get(source_record_id) + + def put(self, source_record_id: str, record: MappedRecord) -> None: + self._records[source_record_id] = record + + +class DynamoDBRecordMap: + """Durable map in a DynamoDB table keyed by ``sourceRecordId`` (a boto3 ``Table`` resource).""" + + def __init__(self, table): + self._table = table + + def get(self, source_record_id: str) -> MappedRecord | None: + item = self._table.get_item(Key={"sourceRecordId": source_record_id}, ConsistentRead=True).get("Item") + if not item: + return None + return MappedRecord( + target_record_id=item.get("targetRecordId"), + event_time=float(item["eventTime"]), + deleted=bool(item.get("deleted", False)), + ) + + def put(self, source_record_id: str, record: MappedRecord) -> None: + item = { + "sourceRecordId": source_record_id, + "eventTime": Decimal(str(record.event_time)), + "deleted": record.deleted, + } + if record.target_record_id: + item["targetRecordId"] = record.target_record_id + self._table.put_item(Item=item) 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 index 8e4a7bfc1..fb3ea3461 100644 --- 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 @@ -3,12 +3,17 @@ 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. +those stream events and applies each change to the target region: creates with +``BatchCreateMemoryRecords``, updates with ``BatchUpdateMemoryRecords``, and +deletes with ``BatchDeleteMemoryRecords``. + +Replays are safe for two reasons. Each create sends a ``clientToken`` derived +from the stream event, so a redelivered event can't create a second record +(``BatchCreateMemoryRecords`` deduplicates on ``clientToken``; ``requestIdentifier`` +is for tracking only). And a :class:`~agentcore_replication.record_map.RecordMap` +remembers which target record each source record became, plus the time of the +last change applied, so updates and deletes reach the right record and stale or +redelivered events are skipped. The same core function powers two consumers: @@ -35,23 +40,28 @@ """ import base64 +import hashlib import json import logging -import uuid +import re from dataclasses import dataclass, field from datetime import datetime, timezone from botocore.exceptions import ClientError +from .record_map import MappedRecord, RecordMap + logger = logging.getLogger(__name__) -# Stream event types we replicate vs. ignore. -REPLICABLE_EVENTS = {"MemoryRecordCreated", "MemoryRecordUpdated"} -SKIP_EVENTS = {"StreamingEnabled", "MemoryRecordDeleted"} +# Stream event types we replicate. Anything else (StreamingEnabled, or a type +# added later) is acknowledged and skipped. +REPLICABLE_EVENTS = {"MemoryRecordCreated", "MemoryRecordUpdated", "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"} +# ``failedRecords[].errorCode`` is an integer. Throttling (429) and server-side +# errors (5xx) are worth retrying (let the ESM redeliver / a local loop re-poll); +# anything else is terminal for that record. +NOT_FOUND = 404 +THROTTLED = 429 def stream_delivery_resources(stream_arn: str) -> dict: @@ -98,20 +108,44 @@ def _to_epoch(event_time) -> float: if isinstance(event_time, (int, float)): return float(event_time) if isinstance(event_time, str) and event_time: + # The stream sends nanoseconds; before Python 3.11, fromisoformat accepts at most 6 fractional digits. + iso = re.sub(r"(\.\d{6})\d+", r"\1", event_time.replace("Z", "+00:00")) try: - dt = datetime.fromisoformat(event_time.replace("Z", "+00:00")) + dt = datetime.fromisoformat(iso) return dt.timestamp() except ValueError: pass return datetime.now(timezone.utc).timestamp() +def _client_token(stream_event: dict) -> str: + """Derive the create ``clientToken`` from the event, so a redelivered event sends the same token. + + The consumer sends one record per call, so the token depends only on this event. + """ + key = "|".join(str(stream_event.get(k, "")) for k in ("memoryId", "memoryRecordId", "eventType", "eventTime")) + return hashlib.sha256(key.encode("utf-8")).hexdigest() + + +def _raise_for_failed(resp: dict, operation: str, record_id: str, ignore_not_found: bool = False) -> None: + """Raise on ``failedRecords``: ``ClientError`` if retryable (so the batch is retried), else ``RuntimeError``.""" + for failed in resp.get("failedRecords", []): + code = failed.get("errorCode") or 0 + if ignore_not_found and code == NOT_FOUND: + continue + msg = f"record {record_id}: {code} {failed.get('errorMessage')}" + if code == THROTTLED or code >= 500: + raise ClientError({"Error": {"Code": str(code), "Message": msg}}, operation) + raise RuntimeError(msg) + + def replicate_stream_event( stream_event: dict, target_client, target_memory_id: str, + record_map: RecordMap, ) -> str: - """Replicate a single decoded ``memoryStreamEvent`` to the target region. + """Apply 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. @@ -119,13 +153,34 @@ def replicate_stream_event( 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. + if event_type not in REPLICABLE_EVENTS or not record_id: + # StreamingEnabled is a control event; unknown types are ignored forward-compatibly. logger.info("skip stream event type=%s id=%s", event_type, record_id) return "skipped" + event_time = _to_epoch(stream_event.get("eventTime")) + mapped = record_map.get(record_id) + if mapped and (event_time <= mapped.event_time or mapped.deleted): + # Redelivered or out of order: this change, or a newer one, was already applied. + logger.info("skip stale %s for %s", event_type, record_id) + return "skipped" + + if event_type == "MemoryRecordDeleted": + # Delete events carry only IDs, so the map supplies the target record. + outcome = "skipped" + if mapped and mapped.target_record_id: + resp = target_client.batch_delete_memory_records( + memoryId=target_memory_id, + records=[{"memoryRecordId": mapped.target_record_id}], + ) + _raise_for_failed(resp, "BatchDeleteMemoryRecords", record_id, ignore_not_found=True) + outcome = "replicated" + # Keep a tombstone so a late or redelivered create for this record is skipped. + target_record_id = mapped.target_record_id if mapped else None + record_map.put(record_id, MappedRecord(target_record_id, event_time, deleted=True)) + logger.info("deleted record %s from %s", record_id, target_memory_id) + return outcome + text = stream_event.get("memoryRecordText") if not text: logger.warning("skip record %s: no memoryRecordText", record_id) @@ -134,44 +189,49 @@ def replicate_stream_event( # 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, - } + # target and the Batch APIs would reject it. AgentCore associates the + # replicated record by its namespaces instead. + namespaces = stream_event.get("namespaces") or [] + if mapped: + # Already replicated: apply the newer content to the same target record. + resp = target_client.batch_update_memory_records( + memoryId=target_memory_id, + records=[ + { + "memoryRecordId": mapped.target_record_id, + "timestamp": event_time, + "content": {"text": text}, + "namespaces": namespaces, + } + ], + ) + _raise_for_failed(resp, "BatchUpdateMemoryRecords", record_id) + record_map.put(record_id, MappedRecord(mapped.target_record_id, event_time)) + logger.info("updated record %s in %s", record_id, target_memory_id) + return "replicated" + + # Not replicated yet (a create, or an update whose create predates streaming). resp = target_client.batch_create_memory_records( memoryId=target_memory_id, - records=[record], - clientToken=uuid.uuid4().hex, + records=[ + { + "requestIdentifier": record_id, + "content": {"text": text}, + "namespaces": namespaces, + "timestamp": event_time, + } + ], + clientToken=_client_token(stream_event), ) - - 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) - + _raise_for_failed(resp, "BatchCreateMemoryRecords", record_id) + created = resp.get("successfulRecords", []) + if not created: + raise RuntimeError(f"record {record_id}: BatchCreateMemoryRecords returned no record ID") + record_map.put(record_id, MappedRecord(created[0]["memoryRecordId"], event_time)) logger.info("replicated record %s -> %s", record_id, target_memory_id) return "replicated" @@ -180,6 +240,7 @@ def process_kinesis_records( kinesis_records, target_client, target_memory_id: str, + record_map: RecordMap, stats: "StreamStats | None" = None, ) -> StreamStats: """Decode and replicate a batch of raw Kinesis records. @@ -205,7 +266,7 @@ def process_kinesis_records( continue try: - outcome = replicate_stream_event(stream_event, target_client, target_memory_id) + outcome = replicate_stream_event(stream_event, target_client, target_memory_id, record_map) if outcome == "replicated": stats.replicated += 1 else: 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 index 6c5c8856d..82ad9f928 100644 --- 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 @@ -3,8 +3,9 @@ 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. + applies each create, update, and delete to the TARGET region's memory. A + DynamoDB table maps each source record ID to the record ID that the target + assigned, so updates and deletes reach the right record. STM is replicated separately by the application at write time (dual-write CreateEvent with extractionMode=SKIP on the target) — see @@ -76,7 +77,26 @@ Resources: MessageRetentionPeriod: 1209600 # 14 days VisibilityTimeout: 300 - # 4) Consumer Lambda execution role: read the stream, write the target memory. + # 4) Map of source record ID -> target record ID and last applied event time, + # so the consumer can apply updates and deletes and skip redelivered events. + RecordMapTable: + Type: AWS::DynamoDB::Table + Properties: + TableName: !Sub "agentcore-ltm-record-map-${TargetRegion}" + BillingMode: PAY_PER_REQUEST + AttributeDefinitions: + - AttributeName: sourceRecordId + AttributeType: S + KeySchema: + - AttributeName: sourceRecordId + KeyType: HASH + PointInTimeRecoverySpecification: + PointInTimeRecoveryEnabled: true + SSESpecification: + SSEEnabled: true + + # 5) Consumer Lambda execution role: read the stream, write the target memory + # and the record map. ConsumerRole: Type: AWS::IAM::Role Properties: @@ -109,10 +129,19 @@ Resources: Resource: !GetAtt ReplicationDLQ.Arn - Sid: TargetMemoryWrite Effect: Allow - Action: bedrock-agentcore:BatchCreateMemoryRecords + Action: + - bedrock-agentcore:BatchCreateMemoryRecords + - bedrock-agentcore:BatchUpdateMemoryRecords + - bedrock-agentcore:BatchDeleteMemoryRecords Resource: - !Sub "arn:aws:bedrock-agentcore:${TargetRegion}:${AWS::AccountId}:memory/${TargetMemoryId}" - !Sub "arn:aws:bedrock-agentcore:${TargetRegion}:${AWS::AccountId}:memory/${TargetMemoryId}/*" + - Sid: RecordMapReadWrite + Effect: Allow + Action: + - dynamodb:GetItem + - dynamodb:PutItem + Resource: !GetAtt RecordMapTable.Arn ConsumerFunction: Type: AWS::Lambda::Function @@ -130,8 +159,9 @@ Resources: Variables: TARGET_MEMORY_ID: !Ref TargetMemoryId TARGET_REGION: !Ref TargetRegion + RECORD_MAP_TABLE: !Ref RecordMapTable - # 5) Wire the stream to the consumer. + # 6) Wire the stream to the consumer. StreamMapping: Type: AWS::Lambda::EventSourceMapping Properties: @@ -191,3 +221,6 @@ Outputs: Value: !Ref ConsumerFunction DlqUrl: Value: !Ref ReplicationDLQ + RecordMapTableName: + Description: Source-to-target record ID map used by the consumer + Value: !Ref RecordMapTable 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 index a10442c66..67f4c4dc3 100644 --- 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 @@ -3,8 +3,9 @@ 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). +create, update, and delete is applied to the TARGET region. A DynamoDB record map +tracks which target record each source record became, so updates and deletes +reach the right record and redelivered events are skipped. STM is replicated separately by the application at write time (dual-write ``CreateEvent`` with ``extractionMode="SKIP"`` on the target) — see @@ -14,6 +15,7 @@ ---------------------- TARGET_MEMORY_ID : target (replica) memory resource ID TARGET_REGION : target AWS region +RECORD_MAP_TABLE : DynamoDB table for the source-to-target record ID map A retryable failure raises, so the ESM retries the batch (configure BisectBatchOnFunctionError + an SQS DLQ on the mapping for poison records). @@ -24,7 +26,7 @@ import os import boto3 - +from agentcore_replication.record_map import DynamoDBRecordMap from agentcore_replication.stream_consumer import ( StreamStats, make_target_client, @@ -36,9 +38,11 @@ TARGET_MEMORY_ID = os.environ["TARGET_MEMORY_ID"] TARGET_REGION = os.environ["TARGET_REGION"] +RECORD_MAP_TABLE = os.environ["RECORD_MAP_TABLE"] -# Build the target client once per container (cold start) and reuse it. +# Build the clients once per container (cold start) and reuse them. _target_client = make_target_client(boto3.Session(), TARGET_REGION) +_record_map = DynamoDBRecordMap(boto3.resource("dynamodb").Table(RECORD_MAP_TABLE)) def lambda_handler(event, context): @@ -47,6 +51,7 @@ def lambda_handler(event, context): event.get("Records", []), target_client=_target_client, target_memory_id=TARGET_MEMORY_ID, + record_map=_record_map, stats=stats, ) result = stats.as_dict() 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 index ac31913c8..8d633b947 100755 --- 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 @@ -109,6 +109,17 @@ aws bedrock-agentcore-control update-memory --region "$SOURCE_REGION" \ --memory-id "$SOURCE_MEMORY_ID" \ --memory-execution-role-arn "$EXEC_ROLE" >/dev/null +# UpdateMemory leaves the memory UPDATING for a minute or two, and the next +# update is rejected until it's ACTIVE again. +echo "==> Waiting for $SOURCE_MEMORY_ID to become ACTIVE" +while true; do + STATUS="$(aws bedrock-agentcore-control get-memory --region "$SOURCE_REGION" \ + --memory-id "$SOURCE_MEMORY_ID" --query memory.status --output text)" + [[ "$STATUS" == ACTIVE ]] && break + [[ "$STATUS" == FAILED ]] && { echo "Memory $SOURCE_MEMORY_ID is FAILED" >&2; exit 1; } + sleep 10 +done + # 6. Enable record streaming on the source memory -> Kinesis. echo "==> Enabling record streaming on $SOURCE_MEMORY_ID" python "$ROOT/scripts/enable_streaming.py" \ 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 index 34ce605e6..707615de6 100644 --- 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 @@ -43,6 +43,7 @@ sys.path.insert(0, ".") from agentcore_replication import ( DualRegionEventWriter, + InMemoryRecordMap, StreamStats, make_target_client, process_kinesis_records, @@ -96,12 +97,12 @@ def wait_active(ctl, mem_id, label, timeout=600): 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, - ) + kwargs = { + "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"] @@ -191,18 +192,19 @@ def consume_stream(kinesis, stream_name, target_client, target_memory_id, expect ] stats = StreamStats() + record_map = InMemoryRecordMap() # the Lambda path uses DynamoDBRecordMap instead 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) + process_kinesis_records(records, target_client, target_memory_id, record_map, 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: + if stats.replicated >= expected > 0: break time.sleep(5) return stats @@ -229,7 +231,8 @@ def count_events(data, mem_id): def list_records(data, mem_id): recs, token = [], None while True: - kw = {"memoryId": mem_id, "namespace": "/", "maxResults": 100} + # namespacePath="/" lists every namespace; namespace="/" returns no records. + kw = {"memoryId": mem_id, "namespacePath": "/", "maxResults": 100} if token: kw["nextToken"] = token resp = data.list_memory_records(**kw) @@ -369,8 +372,8 @@ def main(): 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.)" + " (Replays are safe: creates send a clientToken derived from the " + "stream event, and the record map skips events already applied.)" ) else: ok = False diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/conftest.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/conftest.py new file mode 100644 index 000000000..ef49a2880 --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/conftest.py @@ -0,0 +1,5 @@ +import sys +from pathlib import Path + +# Make the sample's agentcore_replication package importable when pytest runs from the sample directory. +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) diff --git a/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/test_stream_consumer.py b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/test_stream_consumer.py new file mode 100644 index 000000000..55a90ae3f --- /dev/null +++ b/01-features/04-manage-context-of-your-agent/memory/06-production-patterns/00-multi-region-replication/tests/test_stream_consumer.py @@ -0,0 +1,177 @@ +"""Unit tests for the LTM stream consumer. + +Run from the sample directory: ``python -m pytest tests``. No AWS access is needed; +a fake target client stands in for ``bedrock-agentcore``. +""" + +import base64 +import json + +import pytest +from agentcore_replication.record_map import DynamoDBRecordMap, InMemoryRecordMap, MappedRecord +from agentcore_replication.stream_consumer import _to_epoch, process_kinesis_records, replicate_stream_event +from botocore.exceptions import ClientError + +TARGET = "mem-target" + + +class FakeTargetClient: + """Records calls and mimics AgentCore: a repeated clientToken returns the first response.""" + + def __init__(self, failed_records=None): + self.calls = [] + self.records = {} + self._responses_by_token = {} + self._failed_records = failed_records or [] + + def batch_create_memory_records(self, memoryId, records, clientToken): + self.calls.append(("create", clientToken, records)) + if self._failed_records: + return {"successfulRecords": [], "failedRecords": self._failed_records} + if clientToken not in self._responses_by_token: + record_id = f"tgt-{len(self.records) + 1}" + self.records[record_id] = records[0]["content"]["text"] + self._responses_by_token[clientToken] = { + "successfulRecords": [ + {"memoryRecordId": record_id, "requestIdentifier": records[0]["requestIdentifier"]} + ], + "failedRecords": [], + } + return self._responses_by_token[clientToken] + + def batch_update_memory_records(self, memoryId, records): + self.calls.append(("update", None, records)) + self.records[records[0]["memoryRecordId"]] = records[0]["content"]["text"] + return {"successfulRecords": [{"memoryRecordId": records[0]["memoryRecordId"]}], "failedRecords": []} + + def batch_delete_memory_records(self, memoryId, records): + self.calls.append(("delete", None, records)) + record_id = records[0]["memoryRecordId"] + if self.records.pop(record_id, None) is None: + return { + "successfulRecords": [], + "failedRecords": [{"memoryRecordId": record_id, "status": "FAILED", "errorCode": 404}], + } + return {"successfulRecords": [{"memoryRecordId": record_id}], "failedRecords": []} + + +def event(event_type, record_id="src-1", text="likes tea", event_time="2026-10-09T12:00:00Z"): + return { + "eventType": event_type, + "memoryId": "mem-source", + "memoryRecordId": record_id, + "memoryRecordText": text, + "namespaces": ["/facts/user-1"], + "eventTime": event_time, + } + + +def replicate(stream_event, client, record_map): + return replicate_stream_event(stream_event, client, TARGET, record_map) + + +def test_a_redelivered_create_reuses_the_same_client_token(): + first, second = FakeTargetClient(), FakeTargetClient() + replicate(event("MemoryRecordCreated"), first, InMemoryRecordMap()) + replicate(event("MemoryRecordCreated"), second, InMemoryRecordMap()) + assert first.calls[0][1] == second.calls[0][1] + + +def test_a_redelivered_create_is_skipped_once_the_record_is_mapped(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + assert replicate(event("MemoryRecordCreated"), client, record_map) == "replicated" + assert replicate(event("MemoryRecordCreated"), client, record_map) == "skipped" + assert len(client.records) == 1 + + +def test_an_update_is_applied_to_the_mapped_standby_record(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + replicate(event("MemoryRecordCreated"), client, record_map) + replicate(event("MemoryRecordUpdated", text="likes coffee", event_time="2026-10-09T12:05:00Z"), client, record_map) + assert client.calls[-1][0] == "update" + assert client.records == {"tgt-1": "likes coffee"} + + +def test_an_update_that_is_not_newer_than_the_last_applied_change_is_skipped(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + replicate(event("MemoryRecordCreated", event_time="2026-10-09T12:05:00Z"), client, record_map) + outcome = replicate(event("MemoryRecordUpdated", text="stale"), client, record_map) + assert outcome == "skipped" + assert client.records == {"tgt-1": "likes tea"} + + +def test_an_update_for_an_unmapped_record_creates_it(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + replicate(event("MemoryRecordUpdated", text="likes coffee"), client, record_map) + assert client.calls[0][0] == "create" + assert record_map.get("src-1").target_record_id == "tgt-1" + + +def test_a_delete_removes_the_mapped_standby_record(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + replicate(event("MemoryRecordCreated"), client, record_map) + replicate(event("MemoryRecordDeleted", event_time="2026-10-09T12:10:00Z"), client, record_map) + assert client.records == {} + assert record_map.get("src-1").deleted + + +def test_a_create_that_arrives_after_its_delete_is_skipped(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + replicate(event("MemoryRecordDeleted", event_time="2026-10-09T12:10:00Z"), client, record_map) + assert replicate(event("MemoryRecordCreated"), client, record_map) == "skipped" + assert client.records == {} + + +def test_a_delete_of_a_record_already_gone_from_the_standby_counts_as_done(): + client, record_map = FakeTargetClient(), InMemoryRecordMap() + record_map.put("src-1", MappedRecord(target_record_id="tgt-missing", event_time=0.0)) + assert replicate(event("MemoryRecordDeleted"), client, record_map) == "replicated" + assert record_map.get("src-1").deleted + + +def test_a_stream_event_time_with_nanoseconds_is_parsed_rather_than_replaced_by_now(): + # The stream sends 9 fractional digits; datetime.fromisoformat rejects them before Python 3.11. + assert _to_epoch("2026-10-09T15:55:36.241618788Z") == pytest.approx(1791561336.241618, abs=1e-3) + + +def test_streaming_enabled_and_unknown_events_are_skipped_without_failing(): + client = FakeTargetClient() + raw = [ + {"kinesis": {"data": base64.b64encode(json.dumps({"memoryStreamEvent": e}).encode()).decode()}} + for e in ({"eventType": "StreamingEnabled"}, {"eventType": "SomethingNew"}) + ] + stats = process_kinesis_records(raw, client, TARGET, InMemoryRecordMap()) + assert (stats.skipped, stats.failed, client.calls) == (2, 0, []) + + +@pytest.mark.parametrize("code", [429, 500, 503]) +def test_a_throttled_or_server_side_failed_record_raises_so_the_batch_is_retried(code): + client = FakeTargetClient(failed_records=[{"status": "FAILED", "errorCode": code, "errorMessage": "try again"}]) + with pytest.raises(ClientError): + replicate(event("MemoryRecordCreated"), client, InMemoryRecordMap()) + + +def test_a_client_side_failed_record_is_terminal_and_leaves_the_record_unmapped(): + client, record_map = FakeTargetClient(failed_records=[{"status": "FAILED", "errorCode": 400}]), InMemoryRecordMap() + with pytest.raises(RuntimeError): + replicate(event("MemoryRecordCreated"), client, record_map) + assert record_map.get("src-1") is None + + +class FakeTable: + def __init__(self): + self.items = {} + + def get_item(self, Key, ConsistentRead): + item = self.items.get(Key["sourceRecordId"]) + return {"Item": item} if item else {} + + def put_item(self, Item): + self.items[Item["sourceRecordId"]] = Item + + +def test_the_dynamodb_record_map_round_trips_a_tombstone(): + record_map = DynamoDBRecordMap(FakeTable()) + record_map.put("src-1", MappedRecord(target_record_id=None, event_time=1.5, deleted=True)) + assert record_map.get("src-1") == MappedRecord(target_record_id=None, event_time=1.5, deleted=True) + assert record_map.get("src-2") is None diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 8956040a9..9b6b3cf7b 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -140,3 +140,4 @@ - Mark Roy (markproy) - Bent Krause (krausexb) - Georgios Tsoukas (gtsoukasaws) +- Gilbert Lepadatu (gilbertlep)