package memory import ( "context" "encoding/json" "os" "path/filepath" "testing" "time" "github.com/sipeed/picoclaw/pkg/providers" ) func writeJSONSession( t *testing.T, dir string, filename string, sess jsonSession, ) { t.Helper() data, err := json.MarshalIndent(sess, "", " ") if err != nil { t.Fatalf("marshal session: %v", err) } err = os.WriteFile(filepath.Join(dir, filename), data, 0o644) if err != nil { t.Fatalf("write session file: %v", err) } } func TestMigrateFromJSON_Basic(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() writeJSONSession(t, sessionsDir, "test.json", jsonSession{ Key: "test", Messages: []providers.Message{ {Role: "user", Content: "hello"}, {Role: "assistant", Content: "hi"}, }, Summary: "A greeting.", Created: time.Now(), Updated: time.Now(), }) count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 1 { t.Errorf("expected 1 migrated, got %d", count) } history, err := store.GetHistory(ctx, "test") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 2 { t.Fatalf("expected 2 messages, got %d", len(history)) } if history[0].Content != "hello" || history[1].Content != "hi" { t.Errorf("unexpected messages: %+v", history) } summary, err := store.GetSummary(ctx, "test") if err != nil { t.Fatalf("GetSummary: %v", err) } if summary != "A greeting." { t.Errorf("summary = %q", summary) } } func TestMigrateFromJSON_WithToolCalls(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() writeJSONSession(t, sessionsDir, "tools.json", jsonSession{ Key: "tools", Messages: []providers.Message{ { Role: "assistant", Content: "Searching...", ToolCalls: []providers.ToolCall{ { ID: "call_1", Type: "function", Function: &providers.FunctionCall{ Name: "web_search", Arguments: `{"q":"test"}`, }, }, }, }, { Role: "tool", Content: "result", ToolCallID: "call_1", }, }, Created: time.Now(), Updated: time.Now(), }) count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 1 { t.Errorf("expected 1, got %d", count) } history, err := store.GetHistory(ctx, "tools") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 2 { t.Fatalf("expected 2 messages, got %d", len(history)) } if len(history[0].ToolCalls) != 1 { t.Fatalf("expected 1 tool call, got %d", len(history[0].ToolCalls)) } if history[0].ToolCalls[0].Function.Name != "web_search" { t.Errorf("function = %q", history[0].ToolCalls[0].Function.Name) } if history[1].ToolCallID != "call_1" { t.Errorf("ToolCallID = %q", history[1].ToolCallID) } } func TestMigrateFromJSON_MultipleFiles(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() for i := 0; i < 3; i++ { key := string(rune('a' + i)) writeJSONSession(t, sessionsDir, key+".json", jsonSession{ Key: key, Messages: []providers.Message{{Role: "user", Content: "msg " + key}}, Created: time.Now(), Updated: time.Now(), }) } count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 3 { t.Errorf("expected 3, got %d", count) } for i := 0; i < 3; i++ { key := string(rune('a' + i)) history, histErr := store.GetHistory(ctx, key) if histErr != nil { t.Fatalf("GetHistory(%q): %v", key, histErr) } if len(history) != 1 { t.Errorf("session %q: expected 1 msg, got %d", key, len(history)) } } } func TestMigrateFromJSON_InvalidJSON(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() // One valid, one invalid. writeJSONSession(t, sessionsDir, "good.json", jsonSession{ Key: "good", Messages: []providers.Message{{Role: "user", Content: "ok"}}, Created: time.Now(), Updated: time.Now(), }) err := os.WriteFile( filepath.Join(sessionsDir, "bad.json"), []byte("{invalid json"), 0o644, ) if err != nil { t.Fatalf("write bad file: %v", err) } count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 1 { t.Errorf("expected 1 (bad file skipped), got %d", count) } history, err := store.GetHistory(ctx, "good") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 1 { t.Errorf("expected 1 message, got %d", len(history)) } } func TestMigrateFromJSON_RenamesFiles(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() writeJSONSession(t, sessionsDir, "rename.json", jsonSession{ Key: "rename", Messages: []providers.Message{{Role: "user", Content: "hi"}}, Created: time.Now(), Updated: time.Now(), }) _, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } // Original .json should not exist. _, statErr := os.Stat(filepath.Join(sessionsDir, "rename.json")) if !os.IsNotExist(statErr) { t.Error("rename.json should have been renamed") } // .json.migrated should exist. _, statErr = os.Stat( filepath.Join(sessionsDir, "rename.json.migrated"), ) if statErr != nil { t.Errorf("rename.json.migrated should exist: %v", statErr) } } func TestMigrateFromJSON_Idempotent(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() writeJSONSession(t, sessionsDir, "idem.json", jsonSession{ Key: "idem", Messages: []providers.Message{{Role: "user", Content: "once"}}, Created: time.Now(), Updated: time.Now(), }) count1, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("first migration: %v", err) } if count1 != 1 { t.Errorf("first run: expected 1, got %d", count1) } // Second run should find only .migrated files, skip them. count2, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("second migration: %v", err) } if count2 != 0 { t.Errorf("second run: expected 0, got %d", count2) } history, err := store.GetHistory(ctx, "idem") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 1 { t.Errorf("expected 1 message, got %d", len(history)) } } func TestMigrateFromJSON_ColonInKey(t *testing.T) { sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() // File is named telegram_123 (sanitized), but the key inside is telegram:123. writeJSONSession(t, sessionsDir, "telegram_123.json", jsonSession{ Key: "telegram:123", Messages: []providers.Message{{Role: "user", Content: "from telegram"}}, Created: time.Now(), Updated: time.Now(), }) count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 1 { t.Errorf("expected 1, got %d", count) } // Accessible via the original key "telegram:123". history, err := store.GetHistory(ctx, "telegram:123") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 1 { t.Fatalf("expected 1 message, got %d", len(history)) } if history[0].Content != "from telegram" { t.Errorf("content = %q", history[0].Content) } // In the file-based store, "telegram:123" and "telegram_123" both // sanitize to the same filename, so they share storage. This is // expected — the colon-to-underscore mapping is a one-way function. history2, err := store.GetHistory(ctx, "telegram_123") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history2) != 1 { t.Errorf("expected 1 (same file), got %d", len(history2)) } } func TestMigrateFromJSON_RetryAfterCrash(t *testing.T) { // Simulates a crash during migration: first run writes messages // but doesn't rename the .json file. Second run must replace // (not duplicate) the messages thanks to SetHistory semantics. sessionsDir := t.TempDir() store := newTestStore(t) ctx := context.Background() writeJSONSession(t, sessionsDir, "retry.json", jsonSession{ Key: "retry", Messages: []providers.Message{ {Role: "user", Content: "one"}, {Role: "assistant", Content: "two"}, }, Created: time.Now(), Updated: time.Now(), }) // First migration succeeds — writes messages and renames file. count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("first migration: %v", err) } if count != 1 { t.Fatalf("expected 1, got %d", count) } // Simulate "crash before rename": restore the .json file. src := filepath.Join(sessionsDir, "retry.json.migrated") dst := filepath.Join(sessionsDir, "retry.json") if renameErr := os.Rename(src, dst); renameErr != nil { t.Fatalf("restore .json: %v", renameErr) } // Second migration should re-import without duplicating messages. count, err = MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("second migration: %v", err) } if count != 1 { t.Fatalf("expected 1, got %d", count) } history, err := store.GetHistory(ctx, "retry") if err != nil { t.Fatalf("GetHistory: %v", err) } // Must be exactly 2 messages (not 4 from duplication). if len(history) != 2 { t.Fatalf("expected 2 messages (no duplicates), got %d", len(history)) } if history[0].Content != "one" || history[1].Content != "two" { t.Errorf("unexpected messages: %+v", history) } } func TestMigrateFromJSON_NonexistentDir(t *testing.T) { store := newTestStore(t) ctx := context.Background() count, err := MigrateFromJSON(ctx, "/nonexistent/path", store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 0 { t.Errorf("expected 0, got %d", count) } } func TestMigrateFromJSON_SkipsMetaJSONFiles(t *testing.T) { sessionsDir := t.TempDir() store, err := NewJSONLStore(sessionsDir) if err != nil { t.Fatalf("NewJSONLStore: %v", err) } ctx := context.Background() if addErr := store.AddMessage(ctx, "agent:main:pico:direct:pico:test", "user", "keep me"); addErr != nil { t.Fatalf("AddMessage: %v", addErr) } if summaryErr := store.SetSummary(ctx, "agent:main:pico:direct:pico:test", "keep summary"); summaryErr != nil { t.Fatalf("SetSummary: %v", summaryErr) } metaPath := filepath.Join(sessionsDir, "agent_main_pico_direct_pico_test.meta.json") if _, statErr := os.Stat(metaPath); statErr != nil { t.Fatalf("meta file missing before migration: %v", statErr) } count, err := MigrateFromJSON(ctx, sessionsDir, store) if err != nil { t.Fatalf("MigrateFromJSON: %v", err) } if count != 0 { t.Fatalf("expected 0 migrated, got %d", count) } history, err := store.GetHistory(ctx, "agent:main:pico:direct:pico:test") if err != nil { t.Fatalf("GetHistory: %v", err) } if len(history) != 1 || history[0].Content != "keep me" { t.Fatalf("history = %+v, want preserved single message", history) } summary, err := store.GetSummary(ctx, "agent:main:pico:direct:pico:test") if err != nil { t.Fatalf("GetSummary: %v", err) } if summary != "keep summary" { t.Fatalf("summary = %q, want %q", summary, "keep summary") } if _, statErr := os.Stat(metaPath); statErr != nil { t.Fatalf("meta file should remain in place: %v", statErr) } if _, statErr := os.Stat(metaPath + ".migrated"); !os.IsNotExist(statErr) { t.Fatalf("meta file should not be renamed, stat err = %v", statErr) } }