feat: add checkpoint (#65)

Author: Raezil

README.md

@@ -134,6 +134,9 @@ go run cmd/example/codemode_utcp_workflow/main.go
 # Agent-to-Agent Communication via UTCP
 go run cmd/example/agent_as_tool/main.go
 go run cmd/example/agent_as_utcp_codemode/main.go
+
+# Agent State Persistence (Checkpoint/Restore)
+go run cmd/example/checkpoint/main.go
 ```
 
 #### Example Descriptions
@@ -144,6 +147,7 @@ go run cmd/example/agent_as_utcp_codemode/main.go
 
 - **`cmd/example/agent_as_tool/main.go`**: Demonstrates exposing agents as UTCP tools using `RegisterAsUTCPProvider()`, enabling agent-to-agent communication and hierarchical agent architectures.
 - **`cmd/example/agent_as_utcp_codemode/main.go`**: Shows an agent exposed as a UTCP tool and orchestrated via CodeMode, illustrating natural language to tool call generation.
+- **`cmd/example/checkpoint/main.go`**: Demonstrates how to checkpoint an agent's state to disk and restore it later, preserving conversation history and shared space memberships.
 
 
 ## Project Structure
@@ -340,6 +344,38 @@ fmt.Println(result["response"])
 - **Standardization**: Uses the standard UTCP schema for inputs and outputs.
 - **Zero Overhead**: Uses an in-process transport when running within the same Go application, avoiding network latency.
 
+### Agent State Persistence
+
+Lattice supports **Checkpointing and Restoration**, allowing you to pause agents mid-task, persist their state to disk or a database, and resume them later (even after a crash or restart).
+
+**Key Methods:**
+- `agent.Checkpoint()`: Serializes the agent's state (system prompt, short-term memory, shared space memberships) to a `[]byte`.
+- `agent.Restore(data []byte)`: Rehydrates an agent instance from a checkpoint.
+
+**Example:**
+
+```go
+// 1. Checkpoint the agent
+data, err := agent.Checkpoint()
+if err != nil {
+    log.Fatal(err)
+}
+// Save 'data' to file/DB...
+
+// 2. Restore the agent (later or after crash)
+// Create a fresh agent instance first
+newAgent, err := agent.New(opts) 
+if err != nil {
+    log.Fatal(err)
+}
+
+// Restore state
+if err := newAgent.Restore(data); err != nil {
+    log.Fatal(err)
+}
+// newAgent now has the same memory and context as the original
+```
+
 ## Why Use TOON?
 
 **Token-Oriented Object Notation (TOON)** is integrated into Lattice to dramatically reduce token consumption when passing structured data to and from LLMs. This is especially critical for AI agent workflows where context windows are precious and API costs scale with token usage.

agent.go

@@ -385,6 +385,46 @@ func (a *Agent) Flush(ctx context.Context, sessionID string) error {
 	return a.memory.FlushToLongTerm(ctx, sessionID)
 }
 
+// Checkpoint serializes the agent's current state (system prompt and short-term memory)
+// to a byte slice. This can be saved to disk or a database to pause the agent.
+func (a *Agent) Checkpoint() ([]byte, error) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	state := AgentState{
+		SystemPrompt: a.systemPrompt,
+		ShortTerm:    a.memory.ExportShortTerm(),
+		Timestamp:    time.Now(),
+	}
+
+	if a.Shared != nil {
+		state.JoinedSpaces = a.Shared.ExportJoinedSpaces()
+	}
+
+	return json.Marshal(state)
+}
+
+// Restore rehydrates the agent's state from a checkpoint.
+// It restores the system prompt and short-term memory.
+func (a *Agent) Restore(data []byte) error {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	var state AgentState
+	if err := json.Unmarshal(data, &state); err != nil {
+		return err
+	}
+
+	a.systemPrompt = state.SystemPrompt
+	a.memory.ImportShortTerm(state.ShortTerm)
+
+	if a.Shared != nil && len(state.JoinedSpaces) > 0 {
+		a.Shared.ImportJoinedSpaces(state.JoinedSpaces)
+	}
+
+	return nil
+}
+
 func (a *Agent) executeTool(
 	ctx context.Context,
 	sessionID, toolName string,

agent_checkpoint_test.go

@@ -0,0 +1,159 @@
+package agent
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Protocol-Lattice/go-agent/src/memory"
+)
+
+func TestAgentCheckpointAndRestore(t *testing.T) {
+	ctx := context.Background()
+
+	// 1. Setup initial agent
+	mem := memory.NewSessionMemory(&memory.MemoryBank{}, 10)
+	// Use DummyEmbedder to avoid external calls and ensure speed
+	mem = mem.WithEmbedder(memory.DummyEmbedder{})
+
+	agent, err := New(Options{
+		Model:        &stubModel{response: "ok"},
+		Memory:       mem,
+		SystemPrompt: "Initial Prompt",
+	})
+	if err != nil {
+		t.Fatalf("New returned error: %v", err)
+	}
+
+	// 2. Add some state (memory)
+	sessionID := "test-session"
+	// storeMemory is unexported but accessible in the same package
+	agent.storeMemory(sessionID, "user", "Hello world", nil)
+	agent.storeMemory(sessionID, "assistant", "Hi there", nil)
+
+	// 3. Checkpoint
+	data, err := agent.Checkpoint()
+	if err != nil {
+		t.Fatalf("Checkpoint failed: %v", err)
+	}
+
+	if len(data) == 0 {
+		t.Fatal("Checkpoint returned empty data")
+	}
+
+	// 4. Create new agent (simulate restart)
+	newMem := memory.NewSessionMemory(&memory.MemoryBank{}, 10)
+	newMem = newMem.WithEmbedder(memory.DummyEmbedder{})
+
+	newAgent, err := New(Options{
+		Model:        &stubModel{response: "ok"},
+		Memory:       newMem,
+		SystemPrompt: "Default Prompt", // Different from initial
+	})
+	if err != nil {
+		t.Fatalf("New returned error: %v", err)
+	}
+
+	// 5. Restore
+	if err := newAgent.Restore(data); err != nil {
+		t.Fatalf("Restore failed: %v", err)
+	}
+
+	// 6. Verify
+	if newAgent.systemPrompt != "Initial Prompt" {
+		t.Errorf("System prompt not restored. Got %q, want %q", newAgent.systemPrompt, "Initial Prompt")
+	}
+
+	// Verify memory
+	// We can use RetrieveContext to check if memories are there.
+	records, err := newAgent.memory.RetrieveContext(ctx, sessionID, "", 10)
+	if err != nil {
+		t.Fatalf("RetrieveContext failed: %v", err)
+	}
+
+	if len(records) != 2 {
+		t.Errorf("Expected 2 memory records, got %d", len(records))
+	}
+
+	// Check content
+	foundUser := false
+	foundAssistant := false
+	for _, r := range records {
+		if r.Content == "Hello world" {
+			foundUser = true
+		}
+		if r.Content == "Hi there" {
+			foundAssistant = true
+		}
+	}
+
+	if !foundUser {
+		t.Error("User memory not found")
+	}
+	if !foundAssistant {
+		t.Error("Assistant memory not found")
+	}
+}
+
+func TestAgentCheckpointSharedSpaces(t *testing.T) {
+	// Setup
+	mem := memory.NewSessionMemory(&memory.MemoryBank{}, 10).WithEmbedder(memory.DummyEmbedder{})
+	// Grant permissions in registry
+	mem.Spaces.Grant("team:alpha", "agent-1", memory.SpaceRoleWriter, 0)
+	mem.Spaces.Grant("team:beta", "agent-1", memory.SpaceRoleWriter, 0)
+
+	shared := memory.NewSharedSession(mem, "agent-1", "team:alpha")
+
+	agent, _ := New(Options{
+		Model:  &stubModel{response: "ok"},
+		Memory: mem,
+		Shared: shared,
+	})
+
+	// Join another space
+	if err := agent.Shared.Join("team:beta"); err != nil {
+		t.Fatalf("Join failed: %v", err)
+	}
+
+	// Checkpoint
+	data, err := agent.Checkpoint()
+	if err != nil {
+		t.Fatalf("Checkpoint failed: %v", err)
+	}
+
+	// Restore to new agent
+	newMem := memory.NewSessionMemory(&memory.MemoryBank{}, 10).WithEmbedder(memory.DummyEmbedder{})
+	// Simulate persistent registry: grant permissions again
+	newMem.Spaces.Grant("team:alpha", "agent-1", memory.SpaceRoleWriter, 0)
+	newMem.Spaces.Grant("team:beta", "agent-1", memory.SpaceRoleWriter, 0)
+
+	newShared := memory.NewSharedSession(newMem, "agent-1") // No initial spaces
+	newAgent, _ := New(Options{
+		Model:  &stubModel{response: "ok"},
+		Memory: newMem,
+		Shared: newShared,
+	})
+
+	if err := newAgent.Restore(data); err != nil {
+		t.Fatalf("Restore failed: %v", err)
+	}
+
+	// Verify spaces
+	spaces := newAgent.Shared.Spaces()
+	foundAlpha := false
+	foundBeta := false
+	for _, s := range spaces {
+		if s == "team:alpha" {
+			foundAlpha = true
+		}
+		if s == "team:beta" {
+			foundBeta = true
+		}
+	}
+
+	if !foundAlpha {
+		t.Error("Expected to be joined to team:alpha")
+	}
+	if !foundBeta {
+		t.Error("Expected to be joined to team:beta")
+	}
+}

cmd/example/README.md

@@ -89,6 +89,33 @@ User Input → Manager Agent → agent.researcher (Researcher Agent) → Result
 
 ---
 
+---
+
+### 4. Agent State Persistence (Checkpoint/Restore)
+**File:** `cmd/example/checkpoint/main.go`
+
+Demonstrates how to save an agent's state to disk and restore it later, preserving conversation history and shared space memberships.
+
+**Key Concepts:**
+- `agent.Checkpoint()`: Serializes state to `[]byte`
+- `agent.Restore()`: Rehydrates state from `[]byte`
+- Persisting short-term memory and shared space context
+- Resuming conversations across process restarts
+
+**Run:**
+```bash
+go run cmd/example/checkpoint/main.go
+```
+
+**What it shows:**
+1. Creating an agent and having a conversation
+2. Checkpointing the agent to a JSON file
+3. Creating a fresh agent instance
+4. Restoring the state from the file
+5. Verifying the agent remembers the previous conversation
+
+---
+
 ## CodeMode Pattern Explained
 
 CodeMode is a powerful feature that allows agents to orchestrate tools through generated Go code.