feat(memory): add file-based session migration to recall memory
MigrateFileSessions reads JSON session files from disk and imports them as recall_items with preserved timestamps. Idempotent via .migrated marker file. Enables seamless transition from file-based persistence to the 3-tier memory system.
This commit is contained in:
ZanzyTHEbar
parent
036c1933e1
commit
5e0319f54b
2 changed files with 457 additions and 0 deletions
166
pkg/memory/migrate_sessions.go
Normal file
166
pkg/memory/migrate_sessions.go
Normal file
|
|
@ -0,0 +1,166 @@
|
||||||
|
// Package memory provides the 3-tier MemGPT memory system.
|
||||||
|
// This file handles one-time migration of file-based session data
|
||||||
|
// into the SQLite-backed recall memory tier.
|
||||||
|
package memory
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
const migrationMarkerFile = ".sessions_migrated"
|
||||||
|
|
||||||
|
// SessionFile mirrors the on-disk session format from pkg/session.
|
||||||
|
// Defined here to avoid circular imports.
|
||||||
|
type SessionFile struct {
|
||||||
|
Key string `json:"key"`
|
||||||
|
Messages []SessionMsg `json:"messages"`
|
||||||
|
Summary string `json:"summary,omitempty"`
|
||||||
|
Created time.Time `json:"created"`
|
||||||
|
Updated time.Time `json:"updated"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type SessionMsg struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content string `json:"content"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// MigrateSessionsResult holds counts from a session migration run.
|
||||||
|
type MigrateSessionsResult struct {
|
||||||
|
SessionsFound int
|
||||||
|
SessionsMigrated int
|
||||||
|
ItemsCreated int
|
||||||
|
Errors int
|
||||||
|
}
|
||||||
|
|
||||||
|
// MigrateFileSessions reads all JSON session files from sessionsDir and inserts
|
||||||
|
// their messages as RecallItems into the delegate. It writes a marker file to
|
||||||
|
// prevent re-running. Safe to call repeatedly — no-ops after first migration.
|
||||||
|
func MigrateFileSessions(ctx context.Context, del MemoryDelegate, agentID, sessionsDir string) (*MigrateSessionsResult, error) {
|
||||||
|
if sessionsDir == "" {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
markerPath := filepath.Join(sessionsDir, migrationMarkerFile)
|
||||||
|
if _, err := os.Stat(markerPath); err == nil {
|
||||||
|
return nil, nil // already migrated
|
||||||
|
}
|
||||||
|
|
||||||
|
files, err := os.ReadDir(sessionsDir)
|
||||||
|
if err != nil {
|
||||||
|
if os.IsNotExist(err) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
return nil, fmt.Errorf("read sessions dir: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
result := &MigrateSessionsResult{}
|
||||||
|
|
||||||
|
for _, f := range files {
|
||||||
|
if f.IsDir() || filepath.Ext(f.Name()) != ".json" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
result.SessionsFound++
|
||||||
|
|
||||||
|
sessPath := filepath.Join(sessionsDir, f.Name())
|
||||||
|
data, err := os.ReadFile(sessPath)
|
||||||
|
if err != nil {
|
||||||
|
logger.WarnCF("migrate", "Failed to read session file",
|
||||||
|
map[string]interface{}{"path": sessPath, "error": err.Error()})
|
||||||
|
result.Errors++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var sess SessionFile
|
||||||
|
if err := json.Unmarshal(data, &sess); err != nil {
|
||||||
|
logger.WarnCF("migrate", "Failed to parse session file",
|
||||||
|
map[string]interface{}{"path": sessPath, "error": err.Error()})
|
||||||
|
result.Errors++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
sessionKey := sess.Key
|
||||||
|
if sessionKey == "" {
|
||||||
|
sessionKey = strings.TrimSuffix(f.Name(), ".json")
|
||||||
|
}
|
||||||
|
|
||||||
|
migrated, err := migrateOneSession(ctx, del, agentID, sessionKey, &sess)
|
||||||
|
if err != nil {
|
||||||
|
logger.WarnCF("migrate", "Failed to migrate session",
|
||||||
|
map[string]interface{}{"session": sessionKey, "error": err.Error()})
|
||||||
|
result.Errors++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
result.ItemsCreated += migrated
|
||||||
|
result.SessionsMigrated++
|
||||||
|
}
|
||||||
|
|
||||||
|
// Write marker to prevent re-running
|
||||||
|
if result.SessionsMigrated > 0 || result.SessionsFound == 0 {
|
||||||
|
markerContent := fmt.Sprintf("migrated_at=%s sessions=%d items=%d errors=%d\n",
|
||||||
|
time.Now().UTC().Format(time.RFC3339),
|
||||||
|
result.SessionsMigrated,
|
||||||
|
result.ItemsCreated,
|
||||||
|
result.Errors,
|
||||||
|
)
|
||||||
|
os.WriteFile(markerPath, []byte(markerContent), 0644)
|
||||||
|
}
|
||||||
|
|
||||||
|
if result.SessionsMigrated > 0 {
|
||||||
|
logger.InfoCF("migrate", "Session migration complete",
|
||||||
|
map[string]interface{}{
|
||||||
|
"sessions_found": result.SessionsFound,
|
||||||
|
"sessions_migrated": result.SessionsMigrated,
|
||||||
|
"items_created": result.ItemsCreated,
|
||||||
|
"errors": result.Errors,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func migrateOneSession(ctx context.Context, del MemoryDelegate, agentID, sessionKey string, sess *SessionFile) (int, error) {
|
||||||
|
count := 0
|
||||||
|
|
||||||
|
for _, msg := range sess.Messages {
|
||||||
|
if strings.TrimSpace(msg.Content) == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
item := &RecallItem{
|
||||||
|
ID: ids.New(),
|
||||||
|
AgentID: agentID,
|
||||||
|
SessionKey: sessionKey,
|
||||||
|
Role: msg.Role,
|
||||||
|
Sector: SectorEpisodic,
|
||||||
|
Importance: 0.3,
|
||||||
|
Salience: 0.3,
|
||||||
|
DecayRate: 0.01,
|
||||||
|
Content: msg.Content,
|
||||||
|
Tags: "migrated",
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := del.InsertRecallItem(ctx, item); err != nil {
|
||||||
|
return count, fmt.Errorf("insert recall item: %w", err)
|
||||||
|
}
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
|
||||||
|
// If the session had a summary, store it as working context
|
||||||
|
if sess.Summary != "" {
|
||||||
|
if err := del.UpsertWorkingContext(ctx, agentID, sessionKey, sess.Summary); err != nil {
|
||||||
|
logger.WarnCF("migrate", "Failed to store session summary as working context",
|
||||||
|
map[string]interface{}{"session": sessionKey, "error": err.Error()})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return count, nil
|
||||||
|
}
|
||||||
291
pkg/memory/migrate_sessions_test.go
Normal file
291
pkg/memory/migrate_sessions_test.go
Normal file
|
|
@ -0,0 +1,291 @@
|
||||||
|
package memory
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/ids"
|
||||||
|
)
|
||||||
|
|
||||||
|
// mockDelegate captures all calls for testing migration without a real DB.
|
||||||
|
type mockDelegate struct {
|
||||||
|
recallItems []*RecallItem
|
||||||
|
workingContexts map[string]string
|
||||||
|
}
|
||||||
|
|
||||||
|
func newMockDelegate() *mockDelegate {
|
||||||
|
return &mockDelegate{workingContexts: make(map[string]string)}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (m *mockDelegate) Init(_ context.Context) error { return nil }
|
||||||
|
func (m *mockDelegate) Close() error { return nil }
|
||||||
|
func (m *mockDelegate) GetWorkingContext(_ context.Context, agentID, sk string) (*WorkingContext, error) {
|
||||||
|
if c, ok := m.workingContexts[agentID+"/"+sk]; ok {
|
||||||
|
return &WorkingContext{AgentID: agentID, SessionKey: sk, Content: c}, nil
|
||||||
|
}
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) UpsertWorkingContext(_ context.Context, agentID, sk, content string) error {
|
||||||
|
m.workingContexts[agentID+"/"+sk] = content
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) InsertRecallItem(_ context.Context, item *RecallItem) error {
|
||||||
|
m.recallItems = append(m.recallItems, item)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) GetRecallItem(_ context.Context, _ ids.UUID) (*RecallItem, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) UpdateRecallItem(_ context.Context, _ *RecallItem) error { return nil }
|
||||||
|
func (m *mockDelegate) DeleteRecallItem(_ context.Context, _ ids.UUID) error { return nil }
|
||||||
|
func (m *mockDelegate) ListRecallItems(_ context.Context, _, _ string, _, _ int) ([]*RecallItem, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) SearchRecallByKeyword(_ context.Context, _, _ string, _ int) ([]*RecallItem, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) SearchRecallByFTS(_ context.Context, _, _ string, _ int) ([]*RecallItem, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) SearchArchivalByVector(_ context.Context, _ Embedding, _, _ int) ([]SearchResult, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) InsertArchivalChunk(_ context.Context, _ *ArchivalChunk) error { return nil }
|
||||||
|
func (m *mockDelegate) GetArchivalChunk(_ context.Context, _ ids.UUID) (*ArchivalChunk, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) ListArchivalChunks(_ context.Context, _ ids.UUID) ([]*ArchivalChunk, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) ListAllArchivalChunks(_ context.Context, _, _ int) ([]*ArchivalChunk, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) DeleteArchivalChunks(_ context.Context, _ ids.UUID) error { return nil }
|
||||||
|
func (m *mockDelegate) InsertSummary(_ context.Context, _ *MemorySummary) error { return nil }
|
||||||
|
func (m *mockDelegate) ListSummaries(_ context.Context, _, _ string, _ int) ([]*MemorySummary, error) {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) CountRecallItems(_ context.Context, _, _ string) (int, error) {
|
||||||
|
return len(m.recallItems), nil
|
||||||
|
}
|
||||||
|
func (m *mockDelegate) CountArchivalChunks(_ context.Context) (int, error) { return 0, nil }
|
||||||
|
func (m *mockDelegate) HasVectorSearch() bool { return false }
|
||||||
|
func (m *mockDelegate) HasFTS() bool { return false }
|
||||||
|
|
||||||
|
func writeSessionFile(t *testing.T, dir, name string, sess SessionFile) {
|
||||||
|
t.Helper()
|
||||||
|
data, err := json.Marshal(sess)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("marshal session: %v", err)
|
||||||
|
}
|
||||||
|
if err := os.WriteFile(filepath.Join(dir, name), data, 0644); err != nil {
|
||||||
|
t.Fatalf("write session file: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_Basic(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
writeSessionFile(t, sessDir, "sess1.json", SessionFile{
|
||||||
|
Key: "session-1",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "Hello"},
|
||||||
|
{Role: "assistant", Content: "Hi there!"},
|
||||||
|
},
|
||||||
|
Created: time.Now().Add(-time.Hour),
|
||||||
|
Updated: time.Now(),
|
||||||
|
})
|
||||||
|
|
||||||
|
writeSessionFile(t, sessDir, "sess2.json", SessionFile{
|
||||||
|
Key: "session-2",
|
||||||
|
Summary: "Talked about Go programming",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "Tell me about Go"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if result.SessionsFound != 2 {
|
||||||
|
t.Errorf("expected 2 sessions found, got %d", result.SessionsFound)
|
||||||
|
}
|
||||||
|
if result.SessionsMigrated != 2 {
|
||||||
|
t.Errorf("expected 2 sessions migrated, got %d", result.SessionsMigrated)
|
||||||
|
}
|
||||||
|
if result.ItemsCreated != 3 {
|
||||||
|
t.Errorf("expected 3 items, got %d", result.ItemsCreated)
|
||||||
|
}
|
||||||
|
if result.Errors != 0 {
|
||||||
|
t.Errorf("expected 0 errors, got %d", result.Errors)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check recall items were created with correct data
|
||||||
|
if len(del.recallItems) != 3 {
|
||||||
|
t.Fatalf("expected 3 recall items, got %d", len(del.recallItems))
|
||||||
|
}
|
||||||
|
if del.recallItems[0].Role != "user" {
|
||||||
|
t.Errorf("expected 'user' role, got %q", del.recallItems[0].Role)
|
||||||
|
}
|
||||||
|
if del.recallItems[0].Content != "Hello" {
|
||||||
|
t.Errorf("expected 'Hello', got %q", del.recallItems[0].Content)
|
||||||
|
}
|
||||||
|
if del.recallItems[0].SessionKey != "session-1" {
|
||||||
|
t.Errorf("expected 'session-1', got %q", del.recallItems[0].SessionKey)
|
||||||
|
}
|
||||||
|
if del.recallItems[0].Tags != "migrated" {
|
||||||
|
t.Errorf("expected 'migrated' tag, got %q", del.recallItems[0].Tags)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check summary was stored as working context
|
||||||
|
wc, err := del.GetWorkingContext(context.Background(), "picoclaw", "session-2")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("GetWorkingContext: %v", err)
|
||||||
|
}
|
||||||
|
if wc == nil || wc.Content != "Talked about Go programming" {
|
||||||
|
t.Errorf("expected summary as working context, got %v", wc)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_Idempotent(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
writeSessionFile(t, sessDir, "sess.json", SessionFile{
|
||||||
|
Key: "s1",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "test"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
// First run
|
||||||
|
result1, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("first migration: %v", err)
|
||||||
|
}
|
||||||
|
if result1.SessionsMigrated != 1 {
|
||||||
|
t.Fatalf("expected 1, got %d", result1.SessionsMigrated)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Second run should be a no-op (marker file exists)
|
||||||
|
result2, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("second migration: %v", err)
|
||||||
|
}
|
||||||
|
if result2 != nil {
|
||||||
|
t.Error("expected nil result for already-migrated directory")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Still only 1 item
|
||||||
|
if len(del.recallItems) != 1 {
|
||||||
|
t.Errorf("expected 1 recall item (no duplicates), got %d", len(del.recallItems))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_EmptyDir(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
if result.SessionsFound != 0 {
|
||||||
|
t.Errorf("expected 0 sessions, got %d", result.SessionsFound)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_NonexistentDir(t *testing.T) {
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", "/nonexistent/path")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
if result != nil {
|
||||||
|
t.Error("expected nil result for nonexistent dir")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_SkipsEmptyMessages(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
writeSessionFile(t, sessDir, "sess.json", SessionFile{
|
||||||
|
Key: "s1",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "real content"},
|
||||||
|
{Role: "assistant", Content: ""},
|
||||||
|
{Role: "user", Content: " "},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
if result.ItemsCreated != 1 {
|
||||||
|
t.Errorf("expected 1 item (empty msgs skipped), got %d", result.ItemsCreated)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_FallbackKey(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
// Session with empty key — should use filename
|
||||||
|
writeSessionFile(t, sessDir, "custom-key.json", SessionFile{
|
||||||
|
Key: "",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "test"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
if result.ItemsCreated != 1 {
|
||||||
|
t.Fatalf("expected 1 item, got %d", result.ItemsCreated)
|
||||||
|
}
|
||||||
|
if del.recallItems[0].SessionKey != "custom-key" {
|
||||||
|
t.Errorf("expected 'custom-key' from filename, got %q", del.recallItems[0].SessionKey)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMigrateFileSessions_MalformedJSON(t *testing.T) {
|
||||||
|
sessDir := t.TempDir()
|
||||||
|
del := newMockDelegate()
|
||||||
|
|
||||||
|
// Write a malformed JSON file
|
||||||
|
os.WriteFile(filepath.Join(sessDir, "bad.json"), []byte("{broken"), 0644)
|
||||||
|
|
||||||
|
// Also a valid one
|
||||||
|
writeSessionFile(t, sessDir, "good.json", SessionFile{
|
||||||
|
Key: "g1",
|
||||||
|
Messages: []SessionMsg{
|
||||||
|
{Role: "user", Content: "ok"},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
result, err := MigrateFileSessions(context.Background(), del, "picoclaw", sessDir)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("MigrateFileSessions: %v", err)
|
||||||
|
}
|
||||||
|
if result.SessionsFound != 2 {
|
||||||
|
t.Errorf("expected 2 found, got %d", result.SessionsFound)
|
||||||
|
}
|
||||||
|
if result.SessionsMigrated != 1 {
|
||||||
|
t.Errorf("expected 1 migrated, got %d", result.SessionsMigrated)
|
||||||
|
}
|
||||||
|
if result.Errors != 1 {
|
||||||
|
t.Errorf("expected 1 error, got %d", result.Errors)
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Add table
Reference in a new issue