refactor: replace custom DynamoDB checkpointer with langgraph-checkpoint-aws

docs/aws-store.md

@@ -19,16 +19,16 @@ The coding agent supports an optional AWS-backed state store using **DynamoDB**
 └──────────┘       └───────────┘       └────────────┘
 ```
 
-- **DynamoDB** stores checkpoint metadata (thread ID, checkpoint ID, parent, serialization type, S3 key pointer) and intermediate writes (small channel values)
-- **S3** stores the serialized checkpoint blob (channel_values, channel_versions, etc.)
-- Single-table DynamoDB design with composite sort keys (`cp#...` for checkpoints, `wr#...` for writes)
+- **DynamoDB** stores checkpoint metadata, intermediate writes, and small checkpoint payloads.
+- **S3** stores larger checkpoint payloads (typically >350KB) that exceed DynamoDB's efficient item size limits.
+- Unified storage design managed by `langgraph-checkpoint-aws.DynamoDBSaver`.
 
 ## Prerequisites
 
 - AWS account with permissions to create DynamoDB tables, S3 buckets, and IAM policies
 - [Terraform](https://www.terraform.io/) >= 1.5 (for resource provisioning)
 - AWS credentials configured (env vars, IAM role, or AWS profile)
-- Python `aioboto3` package: `pip install coding-agent[aws]`
+- Python `langgraph-checkpoint-aws` package: `pip install coding-agent[aws]`
 
 ## Setup
 

pyproject.toml

@@ -23,14 +23,15 @@ dependencies = [
 
 [project.optional-dependencies]
 aws = [
-    "aioboto3>=13.0",
+    "langgraph-checkpoint-aws>=1.0",
 ]
 dev = [
     "pytest>=8.0",
     "pytest-asyncio>=0.24",
     "ruff>=0.9",
     "moto[dynamodb,s3]>=5.0",
     "interrogate>=1.7.0",
+    "aioboto3>=13.0",
 ]
 
 [project.scripts]

scripts/clear_aws_store.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env uv run python3
+"""Script to empty DynamoDB table and S3 bucket using aioboto3."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import sys
+from pathlib import Path
+
+try:
+    import aioboto3
+except ImportError:
+    print("Error: aioboto3 not found. Please install it with 'pip install aioboto3'.")
+    sys.exit(1)
+
+def load_env():
+    """Load .env file from the project root."""
+    env_path = Path(".env")
+    if env_path.exists():
+        print(f"Loading environment from {env_path}")
+        with open(env_path) as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith("#") and "=" in line:
+                    key, value = line.split("=", 1)
+                    key = key.strip()
+                    value = value.strip()
+                    if len(value) >= 2 and (
+                        (value.startswith('"') and value.endswith('"')) or
+                        (value.startswith("'") and value.endswith("'"))
+                    ):
+                        value = value[1:-1]
+                    os.environ.setdefault(key, value)
+    else:
+        print("No .env file found in current directory.")
+
+async def empty_dynamodb(table_name, region, endpoint_url):
+    """Scan and delete all items from a DynamoDB table."""
+    session = aioboto3.Session()
+    async with session.resource("dynamodb", region_name=region, endpoint_url=endpoint_url) as dynamodb:
+        table = await dynamodb.Table(table_name)
+
+        try:
+            # Getting key schema directly from describe_table to avoid property coroutine issues
+            desc = await table.meta.client.describe_table(TableName=table_name)
+            key_names = [k['AttributeName'] for k in desc['Table']['KeySchema']]
+        except Exception as e:
+            print(f"Error accessing table {table_name}: {e}")
+            return
+
+        print(f"Scanning table {table_name}...")
+
+        count = 0
+        scan_params = {}
+
+        while True:
+            response = await table.scan(**scan_params)
+            items = response.get('Items', [])
+
+            if items:
+                async with table.batch_writer() as batch:
+                    for item in items:
+                        await batch.delete_item(Key={k: item[k] for k in key_names})
+                        count += 1
+
+            if 'LastEvaluatedKey' not in response:
+                break
+            scan_params['ExclusiveStartKey'] = response['LastEvaluatedKey']
+
+        print(f"Deleted {count} items from DynamoDB table {table_name}.")
+
+async def empty_s3(bucket_name, region, endpoint_url):
+    """Delete all objects and versions from an S3 bucket."""
+    session = aioboto3.Session()
+    async with session.resource("s3", region_name=region, endpoint_url=endpoint_url) as s3:
+        try:
+            bucket = await s3.Bucket(bucket_name)
+            print(f"Deleting all objects and versions from bucket {bucket_name}...")
+
+            deleted_count = 0
+            delete_batch = []
+
+            # Delete all versions (including current ones and delete markers)
+            async for version in bucket.object_versions.all():
+                # In aioboto3 resource collections, meta.data is usually pre-populated.
+                # Accessing raw attributes from meta.data avoids coroutine property issues.
+                key = version.meta.data.get('Key')
+                version_id = version.meta.data.get('VersionId')
+                delete_batch.append({'Key': key, 'VersionId': version_id})
+
+                if len(delete_batch) >= 1000:
+                    await bucket.delete_objects(Delete={'Objects': delete_batch})
+                    deleted_count += len(delete_batch)
+                    delete_batch = []
+
+            # Final batch
+            if delete_batch:
+                await bucket.delete_objects(Delete={'Objects': delete_batch})
+                deleted_count += len(delete_batch)
+
+            print(f"Deleted {deleted_count} object(s)/version(s) from S3 bucket {bucket_name}.")
+        except Exception as e:
+            print(f"Error emptying bucket {bucket_name}: {e}")
+            import traceback
+            traceback.print_exc()
+
+async def main():
+    load_env()
+
+    table_name = os.environ.get("DYNAMODB_TABLE_NAME")
+    bucket_name = os.environ.get("S3_CHECKPOINT_BUCKET")
+    region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") or "us-east-1"
+    endpoint_url = os.environ.get("AWS_ENDPOINT_URL")
+
+    if not table_name and not bucket_name:
+        print("Error: DYNAMODB_TABLE_NAME or S3_CHECKPOINT_BUCKET not found in .env or environment.")
+        sys.exit(1)
+
+    tasks = []
+    if table_name:
+        print(f"--- DynamoDB ({region}) ---")
+        tasks.append(empty_dynamodb(table_name, region, endpoint_url))
+
+    if bucket_name:
+        print(f"\n--- S3 ({region}) ---")
+        tasks.append(empty_s3(bucket_name, region, endpoint_url))
+
+    if tasks:
+        await asyncio.gather(*tasks)
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/main.py

@@ -26,6 +26,8 @@ def main() -> None:
         level=logging.DEBUG if args.verbose else logging.INFO,
         format="%(asctime)s %(levelname)-8s %(name)s: %(message)s",
     )
+    # Reduce noise from langgraph_checkpoint_aws
+    logging.getLogger("langgraph_checkpoint_aws").setLevel(logging.WARNING)
 
     settings = load_settings(args.config)