1  package session
  2  
  3  import (
  4  	"encoding/json"
  5  	"os"
  6  	"path/filepath"
  7  	"testing"
  8  	"time"
  9  
 10  	"github.com/Kocoro-lab/ShanClaw/internal/client"
 11  )
 12  
 13  func TestStore_SaveLoad(t *testing.T) {
 14  	dir := t.TempDir()
 15  	store := NewStore(dir)
 16  	defer store.Close()
 17  
 18  	sess := &Session{
 19  		ID:    "test-123",
 20  		Title: "Test session",
 21  		CWD:   "/tmp/test",
 22  		Messages: []client.Message{
 23  			{Role: "user", Content: client.NewTextContent("hello")},
 24  			{Role: "assistant", Content: client.NewTextContent("hi there")},
 25  		},
 26  	}
 27  
 28  	if err := store.Save(sess); err != nil {
 29  		t.Fatalf("save failed: %v", err)
 30  	}
 31  
 32  	loaded, err := store.Load("test-123")
 33  	if err != nil {
 34  		t.Fatalf("load failed: %v", err)
 35  	}
 36  	if loaded.Title != "Test session" {
 37  		t.Errorf("expected 'Test session', got %q", loaded.Title)
 38  	}
 39  	if len(loaded.Messages) != 2 {
 40  		t.Errorf("expected 2 messages, got %d", len(loaded.Messages))
 41  	}
 42  }
 43  
 44  func TestStore_List(t *testing.T) {
 45  	dir := t.TempDir()
 46  	store := NewStore(dir)
 47  	defer store.Close()
 48  
 49  	store.Save(&Session{ID: "aaa", Title: "First"})
 50  	store.Save(&Session{ID: "bbb", Title: "Second"})
 51  
 52  	sessions, err := store.List()
 53  	if err != nil {
 54  		t.Fatalf("list failed: %v", err)
 55  	}
 56  	if len(sessions) != 2 {
 57  		t.Errorf("expected 2 sessions, got %d", len(sessions))
 58  	}
 59  }
 60  
 61  func TestStore_Delete(t *testing.T) {
 62  	dir := t.TempDir()
 63  	store := NewStore(dir)
 64  	defer store.Close()
 65  
 66  	store.Save(&Session{ID: "del-me", Title: "Delete me"})
 67  
 68  	if err := store.Delete("del-me"); err != nil {
 69  		t.Fatalf("delete failed: %v", err)
 70  	}
 71  
 72  	if _, err := store.Load("del-me"); err == nil {
 73  		t.Error("expected error loading deleted session")
 74  	}
 75  
 76  	// Verify file is gone
 77  	path := filepath.Join(dir, "del-me.json")
 78  	if fileExists(path) {
 79  		t.Error("session file should be deleted")
 80  	}
 81  }
 82  
 83  func TestStore_SaveLoadWithImageContent(t *testing.T) {
 84  	dir := t.TempDir()
 85  	store := NewStore(dir)
 86  	defer store.Close()
 87  
 88  	sess := &Session{
 89  		ID:    "vision-test",
 90  		Title: "Vision test",
 91  		CWD:   "/tmp",
 92  		Messages: []client.Message{
 93  			{Role: "user", Content: client.NewTextContent("take a screenshot")},
 94  			{Role: "user", Content: client.NewBlockContent([]client.ContentBlock{
 95  				{Type: "text", Text: "Screenshot captured"},
 96  				{Type: "image", Source: &client.ImageSource{
 97  					Type:      "base64",
 98  					MediaType: "image/png",
 99  					Data:      "iVBORfake",
100  				}},
101  			})},
102  			{Role: "assistant", Content: client.NewTextContent("I see a desktop")},
103  		},
104  	}
105  
106  	if err := store.Save(sess); err != nil {
107  		t.Fatalf("save failed: %v", err)
108  	}
109  
110  	loaded, err := store.Load("vision-test")
111  	if err != nil {
112  		t.Fatalf("load failed: %v", err)
113  	}
114  
115  	if len(loaded.Messages) != 3 {
116  		t.Fatalf("expected 3 messages, got %d", len(loaded.Messages))
117  	}
118  
119  	// First message: plain string
120  	if loaded.Messages[0].Content.Text() != "take a screenshot" {
121  		t.Errorf("msg[0] text mismatch: %q", loaded.Messages[0].Content.Text())
122  	}
123  
124  	// Second message: content blocks with image
125  	if !loaded.Messages[1].Content.HasBlocks() {
126  		t.Fatal("msg[1] should have blocks")
127  	}
128  	blocks := loaded.Messages[1].Content.Blocks()
129  	if len(blocks) != 2 {
130  		t.Fatalf("expected 2 blocks, got %d", len(blocks))
131  	}
132  	if blocks[1].Source == nil || blocks[1].Source.Data != "iVBORfake" {
133  		t.Error("image block data not preserved")
134  	}
135  
136  	// Third message: plain string
137  	if loaded.Messages[2].Content.Text() != "I see a desktop" {
138  		t.Errorf("msg[2] text mismatch: %q", loaded.Messages[2].Content.Text())
139  	}
140  }
141  
142  func TestStore_SaveLoadWithUsageSummary(t *testing.T) {
143  	dir := t.TempDir()
144  	store := NewStore(dir)
145  	defer store.Close()
146  
147  	sess := &Session{
148  		ID:    "usage-test",
149  		Title: "Usage test",
150  		CWD:   "/tmp",
151  		Usage: &UsageSummary{
152  			LLMCalls:              3,
153  			InputTokens:           150,
154  			OutputTokens:          45,
155  			TotalTokens:           195,
156  			CostUSD:               0.67,
157  			CacheReadTokens:       80,
158  			CacheCreationTokens:   300,
159  			CacheCreation5mTokens: 100,
160  			CacheCreation1hTokens: 200,
161  			Model:                 "claude-test",
162  		},
163  	}
164  
165  	if err := store.Save(sess); err != nil {
166  		t.Fatalf("save failed: %v", err)
167  	}
168  
169  	loaded, err := store.Load("usage-test")
170  	if err != nil {
171  		t.Fatalf("load failed: %v", err)
172  	}
173  	if loaded.Usage == nil {
174  		t.Fatal("expected usage summary to be loaded")
175  	}
176  	if loaded.Usage.CacheCreationTokens != 300 {
177  		t.Fatalf("expected legacy cache creation total 300, got %d", loaded.Usage.CacheCreationTokens)
178  	}
179  	if loaded.Usage.CacheCreation5mTokens != 100 || loaded.Usage.CacheCreation1hTokens != 200 {
180  		t.Fatalf("expected split cache creation 100/200, got %d/%d", loaded.Usage.CacheCreation5mTokens, loaded.Usage.CacheCreation1hTokens)
181  	}
182  }
183  
184  func TestStore_SearchIntegration(t *testing.T) {
185  	dir := t.TempDir()
186  	store := NewStore(dir) // auto-creates index + rebuilds (nothing to rebuild)
187  	defer store.Close()
188  
189  	sess := &Session{
190  		ID:    "int-test",
191  		Title: "Integration",
192  		CWD:   "/tmp",
193  		Messages: []client.Message{
194  			{Role: "user", Content: client.NewTextContent("deploy the kubernetes cluster")},
195  			{Role: "assistant", Content: client.NewTextContent("I'll help you deploy k8s")},
196  		},
197  	}
198  	store.Save(sess)
199  
200  	// Search should find it
201  	results, err := store.Search("kubernetes", 10)
202  	if err != nil {
203  		t.Fatalf("Search: %v", err)
204  	}
205  	if len(results) != 1 {
206  		t.Fatalf("expected 1 result, got %d", len(results))
207  	}
208  
209  	// List should use index (fast path)
210  	summaries, err := store.List()
211  	if err != nil {
212  		t.Fatalf("List: %v", err)
213  	}
214  	if len(summaries) != 1 {
215  		t.Fatalf("expected 1 session, got %d", len(summaries))
216  	}
217  
218  	// Delete should clean up index
219  	store.Delete("int-test")
220  	results, _ = store.Search("kubernetes", 10)
221  	if len(results) != 0 {
222  		t.Errorf("expected 0 results after delete, got %d", len(results))
223  	}
224  }
225  
226  func TestStore_GracefulDegradation(t *testing.T) {
227  	dir := t.TempDir()
228  	store := &Store{dir: dir, index: nil} // simulate index failure
229  
230  	sess := &Session{ID: "no-idx", Title: "No index", CWD: "/tmp"}
231  	if err := store.Save(sess); err != nil {
232  		t.Fatalf("Save should work without index: %v", err)
233  	}
234  
235  	summaries, err := store.List()
236  	if err != nil {
237  		t.Fatalf("List should fall back to JSON scan: %v", err)
238  	}
239  	if len(summaries) != 1 {
240  		t.Errorf("expected 1 session from JSON fallback, got %d", len(summaries))
241  	}
242  
243  	_, err = store.Search("anything", 10)
244  	if err == nil {
245  		t.Error("Search should return error when index is nil")
246  	}
247  }
248  
249  func TestStore_ListEmptyDir(t *testing.T) {
250  	dir := t.TempDir()
251  	store := NewStore(dir)
252  	defer store.Close()
253  
254  	summaries, err := store.List()
255  	if err != nil {
256  		t.Fatalf("List on empty dir: %v", err)
257  	}
258  	if len(summaries) != 0 {
259  		t.Errorf("expected 0 sessions, got %d", len(summaries))
260  	}
261  }
262  
263  func TestStore_FirstLaunchMigration(t *testing.T) {
264  	dir := t.TempDir()
265  
266  	// Write JSON files WITHOUT an index (simulate pre-SQLite sessions)
267  	rawStore := &Store{dir: dir, index: nil}
268  	rawStore.Save(&Session{
269  		ID:    "legacy-1",
270  		Title: "Legacy session",
271  		CWD:   "/tmp",
272  		Messages: []client.Message{
273  			{Role: "user", Content: client.NewTextContent("legacy migration test")},
274  		},
275  	})
276  
277  	// Now create store normally — should detect empty index and rebuild
278  	store := NewStore(dir)
279  	defer store.Close()
280  
281  	// Should be searchable after migration
282  	results, err := store.Search("legacy", 10)
283  	if err != nil {
284  		t.Fatalf("Search after migration: %v", err)
285  	}
286  	if len(results) != 1 {
287  		t.Fatalf("expected 1 result after migration, got %d", len(results))
288  	}
289  }
290  
291  func TestSessionMessageMetaSerialization(t *testing.T) {
292  	sess := &Session{
293  		ID:        "test-meta",
294  		CreatedAt: time.Now(),
295  		UpdatedAt: time.Now(),
296  		Title:     "Test",
297  		Messages: []client.Message{
298  			{Role: "user", Content: client.NewTextContent("hello")},
299  			{Role: "assistant", Content: client.NewTextContent("hi")},
300  		},
301  		MessageMeta: []MessageMeta{
302  			{Source: "slack"},
303  			{Source: "slack"},
304  		},
305  	}
306  
307  	data, err := json.Marshal(sess)
308  	if err != nil {
309  		t.Fatal(err)
310  	}
311  
312  	var loaded Session
313  	if err := json.Unmarshal(data, &loaded); err != nil {
314  		t.Fatal(err)
315  	}
316  
317  	if len(loaded.MessageMeta) != 2 {
318  		t.Fatalf("expected 2 meta entries, got %d", len(loaded.MessageMeta))
319  	}
320  	if loaded.MessageMeta[0].Source != "slack" {
321  		t.Fatalf("expected source 'slack', got %q", loaded.MessageMeta[0].Source)
322  	}
323  }
324  
325  func TestSessionMessageMetaBackwardCompat(t *testing.T) {
326  	// Old session JSON without message_meta should deserialize cleanly
327  	oldJSON := `{"id":"old","created_at":"2026-01-01T00:00:00Z","updated_at":"2026-01-01T00:00:00Z","title":"Old","messages":[]}`
328  
329  	var sess Session
330  	if err := json.Unmarshal([]byte(oldJSON), &sess); err != nil {
331  		t.Fatal(err)
332  	}
333  	if sess.MessageMeta != nil {
334  		t.Fatal("expected nil MessageMeta for old session")
335  	}
336  }
337  
338  func TestSessionSourceAtSafety(t *testing.T) {
339  	sess := &Session{
340  		Messages: []client.Message{
341  			{Role: "user", Content: client.NewTextContent("hello")},
342  		},
343  		// No MessageMeta — simulates legacy session
344  	}
345  
346  	// Negative index should not panic
347  	if sess.SourceAt(-1) != "unknown" {
348  		t.Fatal("expected 'unknown' for negative index")
349  	}
350  
351  	// Out of bounds should not panic
352  	if sess.SourceAt(999) != "unknown" {
353  		t.Fatal("expected 'unknown' for out-of-bounds index")
354  	}
355  
356  	// Valid index but no meta
357  	if sess.SourceAt(0) != "unknown" {
358  		t.Fatal("expected 'unknown' for missing meta")
359  	}
360  
361  	// With meta
362  	sess.MessageMeta = []MessageMeta{{Source: "slack"}}
363  	if sess.SourceAt(0) != "slack" {
364  		t.Fatalf("expected 'slack', got %q", sess.SourceAt(0))
365  	}
366  }
367  
368  func TestStore_LoadLegacyStringContent(t *testing.T) {
369  	dir := t.TempDir()
370  	legacyJSON := `{
371  		"id": "legacy-test",
372  		"title": "Legacy",
373  		"cwd": "/tmp",
374  		"messages": [
375  			{"role": "user", "content": "hello"},
376  			{"role": "assistant", "content": "hi there"}
377  		]
378  	}`
379  	os.WriteFile(filepath.Join(dir, "legacy-test.json"), []byte(legacyJSON), 0600)
380  
381  	store := NewStore(dir)
382  	defer store.Close()
383  	loaded, err := store.Load("legacy-test")
384  	if err != nil {
385  		t.Fatalf("load legacy failed: %v", err)
386  	}
387  	if loaded.Messages[0].Content.Text() != "hello" {
388  		t.Errorf("expected 'hello', got %q", loaded.Messages[0].Content.Text())
389  	}
390  	if loaded.Messages[1].Content.Text() != "hi there" {
391  		t.Errorf("expected 'hi there', got %q", loaded.Messages[1].Content.Text())
392  	}
393  }
394  
395  func TestHistoryForLoop_FiltersInjected(t *testing.T) {
396  	sess := &Session{
397  		Messages: []client.Message{
398  			{Role: "user", Content: client.NewTextContent("real user 1")},
399  			{Role: "assistant", Content: client.NewTextContent("real assistant 1")},
400  			{Role: "user", Content: client.NewTextContent("STOP. You wrote out tool calls as text…")},
401  			{Role: "assistant", Content: client.NewTextContent("sorry, trying again")},
402  			{Role: "user", Content: client.NewTextContent("real user 2")},
403  		},
404  		MessageMeta: []MessageMeta{
405  			{Source: "local"},
406  			{Source: "local"},
407  			{Source: "local", SystemInjected: true},
408  			{Source: "local"},
409  			{Source: "local"},
410  		},
411  	}
412  	got := sess.HistoryForLoop()
413  	if len(got) != 4 {
414  		t.Fatalf("got %d messages, want 4 (injected nudge must be filtered)", len(got))
415  	}
416  	for _, m := range got {
417  		if m.Content.Text() == "STOP. You wrote out tool calls as text…" {
418  			t.Errorf("injected nudge leaked into HistoryForLoop output: %q", m.Content.Text())
419  		}
420  	}
421  	// Real user/assistant messages must survive intact in order.
422  	want := []string{"real user 1", "real assistant 1", "sorry, trying again", "real user 2"}
423  	for i, m := range got {
424  		if m.Content.Text() != want[i] {
425  			t.Errorf("msg[%d] = %q, want %q", i, m.Content.Text(), want[i])
426  		}
427  	}
428  }
429  
430  func TestHistoryForLoop_NoMeta(t *testing.T) {
431  	// Legacy sessions with no meta must pass through unchanged.
432  	sess := &Session{
433  		Messages: []client.Message{
434  			{Role: "user", Content: client.NewTextContent("hi")},
435  			{Role: "assistant", Content: client.NewTextContent("hello")},
436  		},
437  	}
438  	got := sess.HistoryForLoop()
439  	if len(got) != 2 {
440  		t.Errorf("got %d messages, want 2", len(got))
441  	}
442  }
443  
444  func TestHistoryForLoop_ShortMeta(t *testing.T) {
445  	// Meta shorter than Messages (partial legacy migration): keep unannotated tail.
446  	sess := &Session{
447  		Messages: []client.Message{
448  			{Role: "user", Content: client.NewTextContent("legacy1")},
449  			{Role: "user", Content: client.NewTextContent("legacy2")},
450  			{Role: "user", Content: client.NewTextContent("new nudge")},
451  		},
452  		MessageMeta: []MessageMeta{
453  			{}, // legacy1 — no flag
454  			// legacy2 and new nudge have no meta entries at all
455  		},
456  	}
457  	got := sess.HistoryForLoop()
458  	if len(got) != 3 {
459  		t.Errorf("got %d messages, want 3 (unannotated positions must survive)", len(got))
460  	}
461  }
462  
463  func TestFilterInjected_NoFlagsFastPath(t *testing.T) {
464  	// When nothing is flagged, FilterInjected takes the fast path: it aliases
465  	// the input slice (no allocation) but caps the capacity so a later append
466  	// on the result cannot silently extend into the caller's backing array
467  	// past the visible length.
468  	backing := make([]client.Message, 2, 10) // extra capacity on purpose
469  	backing[0] = client.Message{Role: "user", Content: client.NewTextContent("a")}
470  	backing[1] = client.Message{Role: "assistant", Content: client.NewTextContent("b")}
471  	meta := []MessageMeta{{Source: "x"}, {Source: "x"}}
472  
473  	got := FilterInjected(backing, meta)
474  	if len(got) != 2 {
475  		t.Fatalf("got len %d, want 2", len(got))
476  	}
477  	if &got[0] != &backing[0] {
478  		t.Error("expected fast path to alias the original slice (no allocation)")
479  	}
480  	if cap(got) != len(got) {
481  		t.Errorf("expected capped capacity %d, got %d — append on result could corrupt caller's backing array", len(got), cap(got))
482  	}
483  }
484  
485  func TestFilterInjected_NoMetaFastPath(t *testing.T) {
486  	// The len(meta) == 0 branch must also cap capacity, same reasoning.
487  	backing := make([]client.Message, 1, 5)
488  	backing[0] = client.Message{Role: "user", Content: client.NewTextContent("a")}
489  	got := FilterInjected(backing, nil)
490  	if cap(got) != len(got) {
491  		t.Errorf("expected capped capacity %d, got %d", len(got), cap(got))
492  	}
493  }