11_checkpointing.py

"""11: File checkpointing: let the agent edit a file, then rewind it."""

import asyncio
from pathlib import Path

from claude_agent_sdk import ClaudeAgentOptions, ClaudeSDKClient, ResultMessage, UserMessage
from demo_util import PROJECT

TARGET = Path(PROJECT) / "README.md"


async def main():
    before = TARGET.read_text()
    options = ClaudeAgentOptions(
        cwd=PROJECT,
        enable_file_checkpointing=True,
        permission_mode="acceptEdits",
        extra_args={"replay-user-messages": None},  # needed to receive user-message UUIDs
        max_turns=4,
    )

    checkpoint_id = session_id = None
    async with ClaudeSDKClient(options) as client:
        await client.query("Append a line '## Changelog' followed by '- demo edit' to README.md")
        async for message in client.receive_response():
            if isinstance(message, UserMessage) and message.uuid and not checkpoint_id:
                checkpoint_id = message.uuid          # first user message = restore point
            if isinstance(message, ResultMessage):
                session_id = message.session_id

    after = TARGET.read_text()
    print(f"File changed by agent: {after != before}")
    print("Tail after edit:", repr(after[-60:]))

    # Rewind by resuming the session and calling rewind_files
    async with ClaudeSDKClient(ClaudeAgentOptions(cwd=PROJECT, enable_file_checkpointing=True,
                                                  resume=session_id)) as client:
        await client.query("")
        async for _ in client.receive_response():
            await client.rewind_files(checkpoint_id)
            break

    restored = TARGET.read_text()
    print(f"Rewound to checkpoint {checkpoint_id}")
    print(f"File identical to original after rewind: {restored == before}")


asyncio.run(main())