diff --git a/src/dynavec/__init__.py b/src/dynavec/__init__.py index daed567..b384daf 100644 --- a/src/dynavec/__init__.py +++ b/src/dynavec/__init__.py @@ -23,6 +23,12 @@ from .async_client import AsyncDynavec from .bm25 import BM25Index from .cache import BaseCache, DynamoDBCache, RedisCache, SemanticCache, warm_cache +from .checkpoint import ( + BaseCheckpointer, + Checkpoint, + DynamoDBCheckpointer, + MemoryCheckpointer, +) from .client import Dynavec from .config import DynavecConfig from .credentials import AWSCredentials @@ -178,6 +184,11 @@ "get_run_logger", "set_default_run_logger", "log_run_event", + # Orchestration & Checkpointing (dynaflow) + "Checkpoint", + "BaseCheckpointer", + "MemoryCheckpointer", + "DynamoDBCheckpointer", # Ingestion "Record", "ingest", diff --git a/src/dynavec/agents/__init__.py b/src/dynavec/agents/__init__.py index d9fa20a..262cfa9 100644 --- a/src/dynavec/agents/__init__.py +++ b/src/dynavec/agents/__init__.py @@ -2,6 +2,12 @@ from __future__ import annotations +from ..checkpoint import ( + BaseCheckpointer, + Checkpoint, + DynamoDBCheckpointer, + MemoryCheckpointer, +) from .base import ( AgentResult, AgentStep, @@ -35,4 +41,8 @@ "mcp_tools_from_session", "mcp_tools_from_session_async", "tool", + "Checkpoint", + "BaseCheckpointer", + "MemoryCheckpointer", + "DynamoDBCheckpointer", ] diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py index c6313c0..3bd9546 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -3,18 +3,69 @@ from __future__ import annotations from collections.abc import Callable, Sequence -from typing import Any +from typing import TYPE_CHECKING, Any from ..chat.base import ChatModel, Message, Tool, ToolCall from .base import AgentResult, AgentStep, AgentTool from .registry import ToolRegistry +if TYPE_CHECKING: + from ..checkpoint import BaseCheckpointer + DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. You solve tasks step-by-step using a ReAct (Reason + Act) approach. When you need to look up information or perform calculations, invoke the appropriate tools. When you have the final answer to the user's goal, provide the clear and complete answer directly without calling any tools.""" +def _message_to_dict(msg: Message) -> dict[str, Any]: + return { + "role": msg.role, + "content": msg.content, + "tool_call_id": msg.tool_call_id, + "tool_calls": [ + {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in msg.tool_calls + ], + } + + +def _dict_to_message(d: dict[str, Any]) -> Message: + tcs = [ + ToolCall(id=t["id"], name=t["name"], arguments=t["arguments"]) + for t in d.get("tool_calls", []) + ] + return Message( + role=d["role"], + content=d.get("content"), + tool_calls=tcs, + tool_call_id=d.get("tool_call_id"), + ) + + +def _step_to_dict(step: AgentStep) -> dict[str, Any]: + return { + "step_number": step.step_number, + "thought": step.thought, + "tool_calls": [ + {"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in step.tool_calls + ], + "observations": list(step.observations), + } + + +def _dict_to_step(d: dict[str, Any]) -> AgentStep: + tcs = [ + ToolCall(id=t["id"], name=t["name"], arguments=t["arguments"]) + for t in d.get("tool_calls", []) + ] + return AgentStep( + step_number=int(d["step_number"]), + thought=d.get("thought"), + tool_calls=tcs, + observations=list(d.get("observations") or []), + ) + + class ReActAgent: """A ReAct-style agent that alternates between thinking, calling tools, and observing results.""" @@ -24,12 +75,14 @@ def __init__( tools: Sequence[AgentTool | Callable[..., Any]] | ToolRegistry | None = None, system_prompt: str | None = None, max_steps: int = 10, + checkpointer: BaseCheckpointer | None = None, ) -> None: self.model = model self.system_prompt = ( system_prompt if system_prompt is not None else DEFAULT_REACT_SYSTEM_PROMPT ) self.max_steps = max(1, max_steps) + self.checkpointer = checkpointer self._tools_map: dict[str, AgentTool] = {} self._chat_tools: list[Tool] = [] @@ -115,13 +168,47 @@ def _finalize_result( tool_calls_count=total_tool_calls, ) - def run(self, goal: str, **kwargs: Any) -> AgentResult: + def run( + self, + goal: str, + *, + thread_id: str | None = None, + resume_from: str | None = None, + **kwargs: Any, + ) -> AgentResult: """Execute the ReAct loop synchronously until goal completion or max_steps.""" - messages = self._init_messages(goal) + messages: list[Message] = [] steps: list[AgentStep] = [] total_tool_calls = 0 + start_step = 1 + + if self.checkpointer is not None and thread_id is not None: + existing_cp = self.checkpointer.get(thread_id, resume_from) + if existing_cp is not None: + state = existing_cp.state + if "messages" in state: + messages = [_dict_to_message(m) for m in state["messages"]] + if "steps" in state: + steps = [_dict_to_step(s) for s in state["steps"]] + total_tool_calls = int(state.get("total_tool_calls", 0)) + start_step = existing_cp.step + 1 - for step_idx in range(1, self.max_steps + 1): + if not messages: + messages = self._init_messages(goal) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [], + "total_tool_calls": 0, + }, + node_id="init", + step=0, + ) + + for step_idx in range(start_step, self.max_steps + 1): chat_res = self.model.invoke( messages, tools=self._chat_tools if self._chat_tools else None, @@ -138,6 +225,20 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: observations=[], ) steps.append(step) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "output": msg.content or "", + "finished": True, + }, + node_id=f"step_{step_idx}", + step=step_idx, + ) return AgentResult( output=msg.content or "", steps=steps, @@ -158,15 +259,78 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: ) steps.append(step) - return self._finalize_result(steps, total_tool_calls) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "finished": False, + }, + node_id=f"step_{step_idx}", + step=step_idx, + ) + + res = self._finalize_result(steps, total_tool_calls) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "output": res.output, + "finished": False, + }, + node_id="max_steps_reached", + step=self.max_steps, + ) + return res - async def arun(self, goal: str, **kwargs: Any) -> AgentResult: + async def arun( + self, + goal: str, + *, + thread_id: str | None = None, + resume_from: str | None = None, + **kwargs: Any, + ) -> AgentResult: """Execute the ReAct loop asynchronously.""" - messages = self._init_messages(goal) + messages: list[Message] = [] steps: list[AgentStep] = [] total_tool_calls = 0 + start_step = 1 - for step_idx in range(1, self.max_steps + 1): + if self.checkpointer is not None and thread_id is not None: + existing_cp = self.checkpointer.get(thread_id, resume_from) + if existing_cp is not None: + state = existing_cp.state + if "messages" in state: + messages = [_dict_to_message(m) for m in state["messages"]] + if "steps" in state: + steps = [_dict_to_step(s) for s in state["steps"]] + total_tool_calls = int(state.get("total_tool_calls", 0)) + start_step = existing_cp.step + 1 + + if not messages: + messages = self._init_messages(goal) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [], + "total_tool_calls": 0, + }, + node_id="init", + step=0, + ) + + for step_idx in range(start_step, self.max_steps + 1): chat_res = await self.model.ainvoke( messages, tools=self._chat_tools if self._chat_tools else None, @@ -183,6 +347,20 @@ async def arun(self, goal: str, **kwargs: Any) -> AgentResult: observations=[], ) steps.append(step) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "output": msg.content or "", + "finished": True, + }, + node_id=f"step_{step_idx}", + step=step_idx, + ) return AgentResult( output=msg.content or "", steps=steps, @@ -205,4 +383,33 @@ async def arun(self, goal: str, **kwargs: Any) -> AgentResult: ) steps.append(step) - return self._finalize_result(steps, total_tool_calls) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "finished": False, + }, + node_id=f"step_{step_idx}", + step=step_idx, + ) + + res = self._finalize_result(steps, total_tool_calls) + if self.checkpointer is not None and thread_id is not None: + self.checkpointer.put( + thread_id, + state={ + "goal": goal, + "messages": [_message_to_dict(m) for m in messages], + "steps": [_step_to_dict(s) for s in steps], + "total_tool_calls": total_tool_calls, + "output": res.output, + "finished": False, + }, + node_id="max_steps_reached", + step=self.max_steps, + ) + return res diff --git a/src/dynavec/checkpoint.py b/src/dynavec/checkpoint.py new file mode 100644 index 0000000..f5c4afa --- /dev/null +++ b/src/dynavec/checkpoint.py @@ -0,0 +1,282 @@ +"""Run and thread state checkpointer for agent orchestration (dynaflow). + +Provides persistent snapshots and time-travel resumability for agent runs +keyed by thread_id and checkpoint_id. Supports both an in-memory checkpointer +for local execution/testing and a serverless DynamoDB-backed checkpointer +for durable, in-account agent execution. +""" + +from __future__ import annotations + +import threading +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any + +from .stores.dynamodb import _from_dynamo, _to_dynamo + + +@dataclass +class Checkpoint: + """A point-in-time snapshot of an agent or workflow run state.""" + + thread_id: str + checkpoint_id: str = field( + default_factory=lambda: ( + f"cp_{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}" + ) + ) + parent_checkpoint_id: str | None = None + node_id: str | None = None + step: int = 0 + state: dict[str, Any] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + + def to_dict(self) -> dict[str, Any]: + """Convert checkpoint to a dictionary.""" + return { + "thread_id": self.thread_id, + "checkpoint_id": self.checkpoint_id, + "parent_checkpoint_id": self.parent_checkpoint_id, + "node_id": self.node_id, + "step": self.step, + "state": self.state, + "metadata": self.metadata, + "timestamp": self.timestamp, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> Checkpoint: + """Construct a Checkpoint from a dictionary.""" + return cls( + thread_id=data["thread_id"], + checkpoint_id=data["checkpoint_id"], + parent_checkpoint_id=data.get("parent_checkpoint_id"), + node_id=data.get("node_id"), + step=int(data.get("step", 0)), + state=dict(data.get("state") or {}), + metadata=dict(data.get("metadata") or {}), + timestamp=data.get("timestamp") or datetime.now(timezone.utc).isoformat(), + ) + + +class BaseCheckpointer(ABC): + """Abstract interface for storing and retrieving thread checkpoints.""" + + @abstractmethod + def save(self, checkpoint: Checkpoint) -> Checkpoint: + """Persist a checkpoint for a thread.""" + + @abstractmethod + def get(self, thread_id: str, checkpoint_id: str | None = None) -> Checkpoint | None: + """Load the latest checkpoint for a thread, or a specific checkpoint by ID.""" + + @abstractmethod + def list(self, thread_id: str, limit: int | None = None) -> list[Checkpoint]: + """List checkpoints for a thread in chronological order.""" + + def put( + self, + thread_id: str, + state: dict[str, Any], + *, + node_id: str | None = None, + step: int | None = None, + parent_checkpoint_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> Checkpoint: + """Convenience method to construct and save a new checkpoint.""" + if parent_checkpoint_id is None or step is None: + latest = self.get(thread_id) + if latest is not None: + if parent_checkpoint_id is None: + parent_checkpoint_id = latest.checkpoint_id + if step is None: + step = latest.step + 1 + else: + if step is None: + step = 0 + + cp = Checkpoint( + thread_id=thread_id, + parent_checkpoint_id=parent_checkpoint_id, + node_id=node_id, + step=step, + state=state, + metadata=metadata or {}, + ) + return self.save(cp) + + +class MemoryCheckpointer(BaseCheckpointer): + """Thread-safe, in-memory checkpointer for local runs and unit testing.""" + + def __init__(self) -> None: + self._threads: dict[str, list[Checkpoint]] = {} + self._lock = threading.Lock() + + def save(self, checkpoint: Checkpoint) -> Checkpoint: + with self._lock: + history = self._threads.setdefault(checkpoint.thread_id, []) + history.append(checkpoint) + return checkpoint + + def get(self, thread_id: str, checkpoint_id: str | None = None) -> Checkpoint | None: + with self._lock: + history = self._threads.get(thread_id, []) + if not history: + return None + if checkpoint_id is None: + return history[-1] + for cp in reversed(history): + if cp.checkpoint_id == checkpoint_id: + return cp + return None + + def list(self, thread_id: str, limit: int | None = None) -> list[Checkpoint]: + with self._lock: + history = list(self._threads.get(thread_id, [])) + if limit is not None: + history = history[-limit:] + return history + + def clear(self, thread_id: str | None = None) -> None: + """Clear all checkpoints, or checkpoints for a specific thread.""" + with self._lock: + if thread_id is not None: + self._threads.pop(thread_id, None) + else: + self._threads.clear() + + +class DynamoDBCheckpointer(BaseCheckpointer): + """DynamoDB-backed persistent checkpointer. + + Compatible with dynavec's single-table schema: + - Individual checkpoint: pk = "{namespace}#checkpoint#{thread_id}#{checkpoint_id}" + - Thread manifest: pk = "{namespace}#thread#{thread_id}" + """ + + def __init__( + self, + table: Any, + *, + namespace: str = "default", + ) -> None: + self.table = table + self.namespace = namespace + + def _cp_pk(self, thread_id: str, checkpoint_id: str) -> str: + return f"{self.namespace}#checkpoint#{thread_id}#{checkpoint_id}" + + def _thread_pk(self, thread_id: str) -> str: + return f"{self.namespace}#thread#{thread_id}" + + def save(self, checkpoint: Checkpoint) -> Checkpoint: + """Save a checkpoint item and update the thread manifest in DynamoDB.""" + item: dict[str, Any] = { + "pk": self._cp_pk(checkpoint.thread_id, checkpoint.checkpoint_id), + "ns": self.namespace, + "type": "checkpoint", + "thread_id": checkpoint.thread_id, + "checkpoint_id": checkpoint.checkpoint_id, + "parent_checkpoint_id": checkpoint.parent_checkpoint_id, + "node_id": checkpoint.node_id, + "step": checkpoint.step, + "state": checkpoint.state, + "metadata": checkpoint.metadata, + "timestamp": checkpoint.timestamp, + } + self.table.put_item(Item=_to_dynamo(item)) + + # Update thread manifest + manifest_pk = self._thread_pk(checkpoint.thread_id) + summary = { + "checkpoint_id": checkpoint.checkpoint_id, + "parent_checkpoint_id": checkpoint.parent_checkpoint_id, + "node_id": checkpoint.node_id, + "step": checkpoint.step, + "timestamp": checkpoint.timestamp, + } + + try: + self.table.update_item( + Key={"pk": manifest_pk}, + UpdateExpression=( + "SET latest_checkpoint_id = :cid, " + "updated_at = :ts, " + "#tp = :type, " + "history = list_append(if_not_exists(history, :empty_list), :new_entry)" + ), + ExpressionAttributeNames={"#tp": "type"}, + ExpressionAttributeValues={ + ":cid": checkpoint.checkpoint_id, + ":ts": checkpoint.timestamp, + ":type": "thread_manifest", + ":empty_list": [], + ":new_entry": [_to_dynamo(summary)], + }, + ) + except Exception: + # Fallback for mock environments without full update_item list_append support + resp = self.table.get_item(Key={"pk": manifest_pk}) + manifest = resp.get("Item") or { + "pk": manifest_pk, + "ns": self.namespace, + "type": "thread_manifest", + "thread_id": checkpoint.thread_id, + "history": [], + } + manifest["latest_checkpoint_id"] = checkpoint.checkpoint_id + manifest["updated_at"] = checkpoint.timestamp + history = list(manifest.get("history") or []) + history.append(_to_dynamo(summary)) + manifest["history"] = history + self.table.put_item(Item=_to_dynamo(manifest)) + + return checkpoint + + def get(self, thread_id: str, checkpoint_id: str | None = None) -> Checkpoint | None: + """Fetch the latest checkpoint for a thread or a specific checkpoint by ID.""" + target_cid = checkpoint_id + if target_cid is None: + manifest_resp = self.table.get_item(Key={"pk": self._thread_pk(thread_id)}) + manifest_item = manifest_resp.get("Item") + if not manifest_item or "latest_checkpoint_id" not in manifest_item: + return None + target_cid = str(manifest_item["latest_checkpoint_id"]) + + cp_resp = self.table.get_item(Key={"pk": self._cp_pk(thread_id, target_cid)}) + cp_item = cp_resp.get("Item") + if not cp_item: + return None + + clean = _from_dynamo(cp_item) + return Checkpoint.from_dict(clean) + + def list(self, thread_id: str, limit: int | None = None) -> list[Checkpoint]: + """List checkpoints for a thread in chronological order.""" + manifest_resp = self.table.get_item(Key={"pk": self._thread_pk(thread_id)}) + manifest_item = manifest_resp.get("Item") + if not manifest_item: + return [] + + history = _from_dynamo(manifest_item.get("history", [])) + if not history: + return [] + + if limit is not None: + history = history[-limit:] + + checkpoints: list[Checkpoint] = [] + for entry in history: + cid = entry.get("checkpoint_id") + if cid: + cp = self.get(thread_id, cid) + if cp: + checkpoints.append(cp) + + return checkpoints diff --git a/tests/test_checkpointer.py b/tests/test_checkpointer.py new file mode 100644 index 0000000..5705663 --- /dev/null +++ b/tests/test_checkpointer.py @@ -0,0 +1,372 @@ +"""Unit tests for agent & workflow state checkpointers (Issue #280).""" + +from __future__ import annotations + +import boto3 +from moto import mock_aws + +from dynavec.agents import ReActAgent +from dynavec.chat.base import ChatModel, ChatResult, Message, Tool, ToolCall +from dynavec.checkpoint import ( + BaseCheckpointer, + Checkpoint, + DynamoDBCheckpointer, + MemoryCheckpointer, +) + + +class ScriptedChatModel(ChatModel): + """A deterministic test fake for ChatModel that replays configured responses.""" + + def __init__(self, responses: list[ChatResult | Message | str]) -> None: + self.responses = list(responses) + self.call_count = 0 + self.received_messages: list[list[Message]] = [] + + def invoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ) -> ChatResult: + self.received_messages.append(list(messages)) + if self.call_count >= len(self.responses): + return ChatResult( + message=Message(role="assistant", content="Responses exhausted."), + finish_reason="stop", + ) + resp = self.responses[self.call_count] + self.call_count += 1 + if isinstance(resp, ChatResult): + return resp + if isinstance(resp, Message): + return ChatResult(message=resp, finish_reason="stop") + return ChatResult( + message=Message(role="assistant", content=resp), + finish_reason="stop", + ) + + async def ainvoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ) -> ChatResult: + return self.invoke(messages, tools, **kwargs) + + def stream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ): + raise NotImplementedError + + async def astream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: object, + ): + raise NotImplementedError + + +def test_checkpoint_dataclass_round_trip(): + """Verify Checkpoint serialization and deserialization round-trip.""" + cp = Checkpoint( + thread_id="th-001", + checkpoint_id="cp-001", + parent_checkpoint_id=None, + node_id="start_node", + step=1, + state={"query": "test", "count": 42, "scores": [0.9, 0.8]}, + metadata={"user": "alice"}, + ) + + d = cp.to_dict() + assert d["thread_id"] == "th-001" + assert d["checkpoint_id"] == "cp-001" + assert d["state"]["count"] == 42 + + restored = Checkpoint.from_dict(d) + assert restored.thread_id == cp.thread_id + assert restored.checkpoint_id == cp.checkpoint_id + assert restored.state == cp.state + assert restored.metadata == cp.metadata + + +def test_memory_checkpointer_crud(): + """Verify MemoryCheckpointer save, get, list, put, and clear operations.""" + chk = MemoryCheckpointer() + assert isinstance(chk, BaseCheckpointer) + + # Empty thread returns None + assert chk.get("thread-1") is None + assert chk.list("thread-1") == [] + + # Put initial checkpoint + cp1 = chk.put( + "thread-1", + state={"val": 10}, + node_id="node_a", + step=0, + ) + assert cp1.step == 0 + assert cp1.state["val"] == 10 + + # Put second checkpoint (auto-chains parent and step) + cp2 = chk.put( + "thread-1", + state={"val": 20}, + node_id="node_b", + ) + assert cp2.step == 1 + assert cp2.parent_checkpoint_id == cp1.checkpoint_id + + # Get latest + latest = chk.get("thread-1") + assert latest is not None + assert latest.checkpoint_id == cp2.checkpoint_id + + # Get specific historical checkpoint (time travel) + first = chk.get("thread-1", cp1.checkpoint_id) + assert first is not None + assert first.checkpoint_id == cp1.checkpoint_id + assert first.state["val"] == 10 + + # List + all_cps = chk.list("thread-1") + assert len(all_cps) == 2 + assert [c.step for c in all_cps] == [0, 1] + + # List with limit + limited = chk.list("thread-1", limit=1) + assert len(limited) == 1 + assert limited[0].checkpoint_id == cp2.checkpoint_id + + # Clear + chk.clear("thread-1") + assert chk.get("thread-1") is None + + +def test_dynamodb_checkpointer_moto(): + """Verify DynamoDBCheckpointer save, get, list with mocked AWS DynamoDB table.""" + with mock_aws(): + ddb = boto3.resource("dynamodb", region_name="us-east-1") + table = ddb.create_table( + TableName="test_checkpoints", + KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + + chk = DynamoDBCheckpointer(table=table, namespace="tenant_1") + + # Initial lookup is empty + assert chk.get("th-ddb-1") is None + assert chk.list("th-ddb-1") == [] + + # Save checkpoint 1 + cp1 = Checkpoint( + thread_id="th-ddb-1", + checkpoint_id="cp-1", + node_id="node_1", + step=1, + state={"status": "in_progress", "ratio": 0.75, "tags": ["a", "b"]}, + metadata={"source": "api"}, + ) + chk.save(cp1) + + # Get latest + got_latest = chk.get("th-ddb-1") + assert got_latest is not None + assert got_latest.checkpoint_id == "cp-1" + assert got_latest.step == 1 + assert got_latest.state["status"] == "in_progress" + assert got_latest.state["ratio"] == 0.75 + assert got_latest.state["tags"] == ["a", "b"] + + # Save checkpoint 2 + cp2 = Checkpoint( + thread_id="th-ddb-1", + checkpoint_id="cp-2", + parent_checkpoint_id="cp-1", + node_id="node_2", + step=2, + state={"status": "completed", "ratio": 1.0}, + ) + chk.save(cp2) + + # Get latest now points to cp-2 + latest2 = chk.get("th-ddb-1") + assert latest2 is not None + assert latest2.checkpoint_id == "cp-2" + assert latest2.step == 2 + assert latest2.state["status"] == "completed" + + # Time-travel get returns cp-1 + hist1 = chk.get("th-ddb-1", "cp-1") + assert hist1 is not None + assert hist1.checkpoint_id == "cp-1" + assert hist1.step == 1 + + # List checkpoints + all_cps = chk.list("th-ddb-1") + assert len(all_cps) == 2 + assert [c.checkpoint_id for c in all_cps] == ["cp-1", "cp-2"] + + # List with limit + lim = chk.list("th-ddb-1", limit=1) + assert len(lim) == 1 + assert lim[0].checkpoint_id == "cp-2" + + +def test_dynamodb_checkpointer_thread_isolation(): + """Verify separate threads do not interfere with each other's checkpoints.""" + with mock_aws(): + ddb = boto3.resource("dynamodb", region_name="us-east-1") + table = ddb.create_table( + TableName="test_checkpoints_iso", + KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + + chk = DynamoDBCheckpointer(table=table, namespace="tenant_iso") + + chk.put("thread_A", state={"agent": "A1"}, node_id="init", step=0) + chk.put("thread_B", state={"agent": "B1"}, node_id="init", step=0) + chk.put("thread_A", state={"agent": "A2"}, node_id="step_1", step=1) + + list_a = chk.list("thread_A") + list_b = chk.list("thread_B") + + assert len(list_a) == 2 + assert len(list_b) == 1 + assert list_a[-1].state["agent"] == "A2" + assert list_b[0].state["agent"] == "B1" + + +def test_react_agent_checkpointing(): + """Verify ReActAgent records checkpoints after each step when checkpointer is configured.""" + model = ScriptedChatModel( + [ + # Step 1: Model requests tool call + ChatResult( + message=Message( + role="assistant", + content="Let me check the weather.", + tool_calls=[ + ToolCall( + id="tc-1", + name="get_weather", + arguments='{"location": "Seattle"}', + ) + ], + ) + ), + # Step 2: Model gives final answer + ChatResult( + message=Message( + role="assistant", + content="The weather in Seattle is 65F and sunny.", + ) + ), + ] + ) + + def get_weather(location: str) -> str: + return f"{location}: 65F and sunny" + + checkpointer = MemoryCheckpointer() + agent = ReActAgent( + model=model, + tools=[get_weather], + checkpointer=checkpointer, + ) + + result = agent.run("What's the weather in Seattle?", thread_id="thread-weather") + + assert result.finished is True + assert "65F and sunny" in result.output + + # Check that checkpoints were written: step 0 (init), step 1 (tool), step 2 (completion) + cps = checkpointer.list("thread-weather") + assert len(cps) == 3 + assert [c.step for c in cps] == [0, 1, 2] + assert cps[0].node_id == "init" + assert cps[1].node_id == "step_1" + assert cps[2].node_id == "step_2" + assert cps[2].state["finished"] is True + + +def test_react_agent_resumption(): + """Verify resuming an agent from a previous thread checkpoint reproduces state and finishes.""" + # First agent execution stops after step 1 + model_part1 = ScriptedChatModel( + [ + ChatResult( + message=Message( + role="assistant", + content="Looking up data.", + tool_calls=[ + ToolCall( + id="tc-1", + name="fetch_data", + arguments='{"item": "widget"}', + ) + ], + ) + ), + ] + ) + + def fetch_data(item: str) -> str: + return f"Item {item} inventory: 42 units" + + checkpointer = MemoryCheckpointer() + agent1 = ReActAgent( + model=model_part1, + tools=[fetch_data], + max_steps=1, # Intentional stop after step 1 + checkpointer=checkpointer, + ) + + res1 = agent1.run("Check inventory for widget", thread_id="thread-resume") + assert res1.finished is False + assert res1.termination_reason == "max_steps_reached" + + # Verify checkpointer recorded step 1 + cp_step1 = checkpointer.get("thread-resume") + assert cp_step1 is not None + assert cp_step1.step == 1 + + # Second agent picks up from checkpoint on the same thread + model_part2 = ScriptedChatModel( + [ + ChatResult( + message=Message( + role="assistant", + content="We have 42 widgets available in stock.", + ) + ), + ] + ) + + agent2 = ReActAgent( + model=model_part2, + tools=[fetch_data], + max_steps=5, + checkpointer=checkpointer, + ) + + res2 = agent2.run("Check inventory for widget", thread_id="thread-resume") + + assert res2.finished is True + assert res2.output == "We have 42 widgets available in stock." + + # Final checkpointer state + cps = checkpointer.list("thread-resume") + assert len(cps) >= 3 + assert cps[-1].step == 2 + assert cps[-1].state["finished"] is True