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