Persistent State

import json
import os
from datetime import datetime

class PersistentTaskState:
    """
    State manager that saves to disk after every update.
    Supports resuming interrupted tasks.
    """
    def __init__(self, task_id: str, state_dir: str = "./agent_states"):
        self.task_id = task_id
        self.state_file = os.path.join(state_dir, f"{task_id}.json")
        os.makedirs(state_dir, exist_ok=True)

        if os.path.exists(self.state_file):
            self._state = self._load()
            print(f"Resumed task {task_id} from step {self._state['step']}")
        else:
            self._state = {
                "task_id": task_id,
                "step": 0,
                "status": "running",
                "history": [],
                "data": {},
                "created_at": datetime.now().isoformat()
            }
            self._save()

    def _save(self):
        """Persist state to disk after every change."""
        with open(self.state_file, "w") as f:
            json.dump(self._state, f, indent=2, default=str)

    def _load(self) -> dict:
        with open(self.state_file) as f:
            return json.load(f)

    def record_action(self, action: str, args: dict, observation: dict):
        self._state["history"].append({
            "step": self._state["step"],
            "action": action,
            "args": args,
            "observation": observation,
            "timestamp": datetime.now().isoformat()
        })
        self._state["step"] += 1
        self._save()  # Persist immediately

    def complete(self, result):
        self._state["status"] = "complete"
        self._state["result"] = result
        self._state["completed_at"] = datetime.now().isoformat()
        self._save()

# Usage — will resume if the file already exists
state = PersistentTaskState("TASK-001")
state.record_action("web_search", {"query": "Python"}, {"status": "success"})
print(f"Step: {state._state['step']}, File: {state.state_file}")