Agent State Management — infra.bea.ai

An agent that writes a file in step 3 needs that file in step 7. An agent debugging a server needs it to still be running between turns. State management is the infrastructure that preserves agent context across the sequence of LLM calls and tool executions that make up a session.

1. What State Must Be Preserved

An agent session accumulates several types of state:

State type	Examples	Lifetime	Storage size
Conversation history	System prompt, user messages, assistant responses, tool outputs	Entire session	10 KB - 10 MB
Filesystem	Created/modified files, downloaded data, build artifacts	Across tool calls	1 MB - 10 GB
Processes	Running dev servers, database instances, watch scripts	Across tool calls	Memory-resident
Environment	Shell variables, working directory, installed packages	Across tool calls	Negligible
Tool results cache	Previous search results, API responses	Across tool calls	1 KB - 1 MB

Turn 1: User: "Set up a React project"
  → Agent runs: npx create-react-app myapp  (creates 200+ files)
  → STATE: filesystem has /myapp/...

Turn 2: User: "Add a login page"
  → Agent needs /myapp/ to still exist
  → Agent runs: vim src/Login.tsx, npm install axios
  → STATE: filesystem modified, new package installed

Turn 3: User: "Start the dev server and check if it works"
  → Agent runs: npm start &  (background process)
  → Agent runs: curl localhost:3000
  → STATE: process running on port 3000

Turn 4: User: "Fix the CSS warning"
  → Agent needs: filesystem + running server + conversation context

2. Persistence Strategies

Strategy 1: Persistent Sandbox (Keep VM Alive)

The simplest approach — keep the sandbox running between turns:

class PersistentSandbox:
    """Keep sandbox alive across turns. State persists naturally."""
    
    def __init__(self, sandbox_provider, idle_timeout=1800):
        self.sandbox = sandbox_provider.create()
        self.idle_timeout = idle_timeout  # Kill after 30min idle
        self.last_activity = time.time()
    
    async def execute(self, command: str) -> str:
        """Execute in the still-running sandbox."""
        self.last_activity = time.time()
        result = await self.sandbox.run(command)
        return result.stdout + result.stderr
    
    async def is_expired(self) -> bool:
        return time.time() - self.last_activity > self.idle_timeout

Pros: Zero complexity — everything just works (files, processes, environment).
Cons: Expensive — holding a VM/container for 30+ minutes per idle session. At 10K concurrent users with 5-minute active / 25-minute idle, you’re paying for 10K VMs where only ~1.7K are active.

Strategy 2: Filesystem Snapshots (Checkpoint/Restore)

Snapshot the filesystem between turns, restore on next turn:

class SnapshotSandbox:
    """Snapshot filesystem between turns. Processes don't survive."""
    
    def __init__(self, storage_backend, session_id: str):
        self.storage = storage_backend  # S3, local disk, etc.
        self.session_id = session_id
    
    async def start_turn(self):
        """Create sandbox and restore last snapshot."""
        self.sandbox = await create_sandbox()
        
        snapshot = await self.storage.get(f"snapshots/{self.session_id}/latest.tar.gz")
        if snapshot:
            await self.sandbox.upload(snapshot, "/workspace")
            await self.sandbox.run("tar xzf /workspace/latest.tar.gz -C /workspace")
    
    async def end_turn(self):
        """Snapshot filesystem and destroy sandbox."""
        await self.sandbox.run(
            "tar czf /tmp/snapshot.tar.gz -C /workspace ."
        )
        snapshot = await self.sandbox.download("/tmp/snapshot.tar.gz")
        await self.storage.put(f"snapshots/{self.session_id}/latest.tar.gz", snapshot)
        await self.sandbox.destroy()
    
    async def execute(self, command: str) -> str:
        return await self.sandbox.run(command)

Pros: No idle cost — sandbox only exists during active turns.
Cons: Processes die between turns (no running servers). Snapshot/restore adds latency (1-10s depending on size).

Strategy 3: CRIU (Checkpoint/Restore In Userspace)

CRIU can freeze an entire container including running processes and restore it later:

# Checkpoint a running container (including all processes)
criu dump --tree <PID> --images-dir /snapshots/session_123/

# Restore later (processes resume from exact state)
criu restore --images-dir /snapshots/session_123/

Property	Value
Checkpoint time	500ms - 5s (depends on memory size)
Restore time	200ms - 2s
Preserves	Filesystem, processes, network connections, file descriptors
Limitations	Can’t restore across different kernel versions, GPU state not supported

Pros: Full state preservation including running processes.
Cons: Complex, fragile, doesn’t work with GPU workloads.

Comparison

Strategy	Idle cost	Latency per turn	Process persistence	Complexity
Persistent sandbox	High (VM alive)	0ms (already running)	✅ Yes	Low
Filesystem snapshot	None	1-10s (restore)	❌ No	Medium
CRIU	None	0.5-2s (restore)	✅ Yes	High
Hybrid (persist short, snapshot long)	Medium	0ms or 1-10s	Partial	Medium

3. Conversation History Management

The conversation window is the most critical state — it determines what the agent “remembers.”

Window Strategies

class ConversationManager:
    def __init__(self, max_tokens=128_000):
        self.messages = []
        self.max_tokens = max_tokens
    
    def add_message(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})
        self._enforce_window()
    
    def _enforce_window(self):
        """Keep conversation within token budget."""
        while self._total_tokens() > self.max_tokens:
            # Strategy: summarize oldest messages instead of dropping
            oldest = self.messages.pop(1)  # Keep system prompt (index 0)
            
            # Option A: Just drop old messages
            # (simplest, loses information)
            
            # Option B: Summarize dropped messages
            # summary = llm.summarize(oldest["content"])
            # self.messages.insert(1, {"role": "system", "content": f"Summary of earlier: {summary}"})
    
    def _total_tokens(self):
        return sum(token_count(m["content"]) for m in self.messages)

Tool Output Truncation

Tool outputs (especially ls, cat, grep results) can be enormous. Truncate intelligently:

def truncate_tool_output(output: str, max_tokens: int = 4000) -> str:
    """Truncate tool output while preserving useful information."""
    if token_count(output) <= max_tokens:
        return output
    
    lines = output.split("\n")
    
    # Keep first N and last N lines (head + tail)
    head_lines = lines[:50]
    tail_lines = lines[-20:]
    omitted = len(lines) - 70
    
    return "\n".join(head_lines) + f"\n\n... ({omitted} lines omitted) ...\n\n" + "\n".join(tail_lines)

4. Session Resumption

When a user returns after hours or days, restore the session:

class SessionManager:
    def __init__(self, db, storage):
        self.db = db        # Session metadata (Postgres/Redis)
        self.storage = storage  # Filesystem snapshots (S3)
    
    async def save_session(self, session_id: str, state: SessionState):
        """Persist session state for later resumption."""
        await self.db.upsert("sessions", {
            "id": session_id,
            "conversation": json.dumps(state.messages),
            "environment": json.dumps(state.env_vars),
            "working_dir": state.working_dir,
            "packages_installed": state.packages,
            "updated_at": datetime.utcnow(),
        })
        
        if state.filesystem_snapshot:
            await self.storage.upload(
                state.filesystem_snapshot,
                f"sessions/{session_id}/fs.tar.gz",
            )
    
    async def resume_session(self, session_id: str) -> SessionState:
        """Restore a saved session."""
        meta = await self.db.get("sessions", session_id)
        if not meta:
            raise SessionNotFound(session_id)
        
        state = SessionState(
            messages=json.loads(meta["conversation"]),
            env_vars=json.loads(meta["environment"]),
            working_dir=meta["working_dir"],
        )
        
        # Restore filesystem
        snapshot = await self.storage.download(f"sessions/{session_id}/fs.tar.gz")
        if snapshot:
            state.filesystem_snapshot = snapshot
        
        return state
    
    async def cleanup_expired(self, max_age_hours=72):
        """Delete sessions older than max_age."""
        expired = await self.db.query(
            "SELECT id FROM sessions WHERE updated_at < NOW() - INTERVAL '%s hours'",
            max_age_hours,
        )
        for session in expired:
            await self.storage.delete(f"sessions/{session.id}/")
            await self.db.delete("sessions", session.id)

5. State Size Management

Agent sessions can accumulate gigabytes of state. Monitor and limit:

# Per-session state limits
STATE_LIMITS = {
    "conversation_tokens": 128_000,   # Max conversation window
    "filesystem_mb": 2_000,           # 2 GB max workspace
    "snapshot_mb": 500,               # 500 MB max compressed snapshot
    "session_age_hours": 72,          # Auto-expire after 3 days
    "packages_count": 100,            # Max installed packages
}

async def check_state_limits(session: Session):
    """Warn or enforce state limits."""
    fs_size = await session.sandbox.run("du -sm /workspace | cut -f1")
    if int(fs_size) > STATE_LIMITS["filesystem_mb"]:
        logger.warning(f"Session {session.id} workspace exceeds {STATE_LIMITS['filesystem_mb']}MB")
        # Option: block new file creation, or auto-cleanup build artifacts

Self-Check

Categorize the four types of agent state and their persistence requirements
Compare persistent sandbox, filesystem snapshot, and CRIU approaches
Implement conversation window management with token budgets
Design a session resumption system with filesystem restore
Set up state size limits to prevent runaway storage costs

References

E2B Sandbox API: Session Persistence
CRIU: Checkpoint/Restore In Userspace
Modal: Serverless Sandboxes
For sandbox isolation types, see Sandbox Types