| 1 | package session |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "strings" |
| 8 | "testing" |
| 9 | |
| 10 | "reasonix/internal/provider" |
| 11 | ) |
| 12 | |
| 13 | func TestTitleMessagesReadAuthoredHistoryAcrossRestart(t *testing.T) { |
| 14 | service, err := NewService("local", NewFilesystemPersistence(t.TempDir())) |
| 15 | if err != nil { |
| 16 | t.Fatal(err) |
| 17 | } |
| 18 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 19 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "title-history"}) |
| 20 | if err != nil { |
| 21 | t.Fatal(err) |
| 22 | } |
| 23 | if messages, err := service.Query().TitleMessages(t.Context(), runtime.Ref(), 3); err != nil || len(messages) != 0 { |
| 24 | t.Fatalf("empty history = %v, %v", messages, err) |
| 25 | } |
| 26 | want := strings.Repeat("large user text ", 80000) |
| 27 | for _, message := range []provider.Message{ |
| 28 | {ID: "host", Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: "hidden host instructions"}, |
| 29 | {ID: "one", Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: want}, |
| 30 | {ID: "assistant", Role: provider.RoleAssistant, Content: "reply"}, |
| 31 | {ID: "two", Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "second request"}, |
| 32 | {ID: "three", Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "third request"}, |
| 33 | {ID: "four", Role: provider.RoleUser, Origin: provider.MessageOriginUser, Content: "fourth request"}, |
| 34 | } { |
| 35 | payload, _ := json.Marshal(map[string]any{"message": message}) |
| 36 | if _, err := runtime.Session().AppendBatch(t.Context(), message.ID, []Event{{Kind: "message/complete", Payload: payload}}); err != nil { |
| 37 | t.Fatal(err) |
| 38 | } |
| 39 | } |
| 40 | if err := service.Close(t.Context(), runtime.Ref()); err != nil { |
| 41 | t.Fatal(err) |
| 42 | } |
| 43 | messages, err := service.Query().TitleMessages(t.Context(), runtime.Ref(), 3) |
| 44 | if err != nil || len(messages) != 3 { |
| 45 | t.Fatalf("title messages = %d, %v", len(messages), err) |
| 46 | } |
| 47 | if messages[0].Content != want || messages[1].Content != "second request" || messages[2].Content != "third request" { |
| 48 | t.Fatalf("title message identities = %q, %q, %q", messages[0].ID, messages[1].ID, messages[2].ID) |
| 49 | } |
| 50 | cancelled, cancel := context.WithCancel(t.Context()) |
| 51 | cancel() |
| 52 | if _, err := service.Query().TitleMessages(cancelled, runtime.Ref(), 3); !errors.Is(err, context.Canceled) { |
| 53 | t.Fatalf("cancelled title read = %v", err) |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | func TestSetTitleIfSequenceIsAtomicWithManualTitleCommit(t *testing.T) { |
| 58 | service, err := NewService("local", NewFilesystemPersistence(t.TempDir())) |
| 59 | if err != nil { |
| 60 | t.Fatal(err) |
| 61 | } |
| 62 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 63 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "title-race"}) |
| 64 | if err != nil { |
| 65 | t.Fatal(err) |
| 66 | } |
| 67 | expected := uint64(0) |
| 68 | prepared, err := runtime.Session().PrepareBatchContext(t.Context(), "delayed-title", Batch{Events: []Event{{Kind: "session/title", Payload: []byte(`{"title":"stale AI title"}`)}}}) |
| 69 | if err != nil { |
| 70 | t.Fatal(err) |
| 71 | } |
| 72 | defer prepared.Release() |
| 73 | if err := service.SetTitle(t.Context(), runtime.Ref(), "manual title"); err != nil { |
| 74 | t.Fatal(err) |
| 75 | } |
| 76 | before := runtime.Session().EventSequence() |
| 77 | cancelled, cancel := context.WithCancel(t.Context()) |
| 78 | cancel() |
| 79 | if err := service.SetTitle(cancelled, runtime.Ref(), "cancelled title"); !errors.Is(err, context.Canceled) { |
| 80 | t.Fatalf("cancelled title write = %v", err) |
| 81 | } |
| 82 | if _, err := runtime.Session().commitPrepared(prepared, &expected); !errors.Is(err, ErrSessionTitleChanged) { |
| 83 | t.Fatalf("interleaved title commit = %v", err) |
| 84 | } |
| 85 | if runtime.Session().EventSequence() != before { |
| 86 | t.Fatal("rejected title consumed an event sequence") |
| 87 | } |
| 88 | if err := service.Close(t.Context(), runtime.Ref()); err != nil { |
| 89 | t.Fatal(err) |
| 90 | } |
| 91 | if err := service.SetTitleIfSequence(t.Context(), runtime.Ref(), 0, "stale AI title"); !errors.Is(err, ErrSessionTitleChanged) { |
| 92 | t.Fatalf("cold stale title = %v", err) |
| 93 | } |
| 94 | info, err := service.Query().Stat(t.Context(), runtime.Ref()) |
| 95 | if err != nil { |
| 96 | t.Fatal(err) |
| 97 | } |
| 98 | if err := service.SetTitleIfSequence(t.Context(), runtime.Ref(), info.TitleSequence, "new AI title"); err != nil { |
| 99 | t.Fatal(err) |
| 100 | } |
| 101 | if info, err := service.Query().Stat(t.Context(), runtime.Ref()); err != nil || info.Title != "new AI title" { |
| 102 | t.Fatalf("persisted title = %+v, %v", info, err) |
| 103 | } |
| 104 | } |
| 105 | |
| 106 | func TestSetTitleIfSequenceRejectsABAAndSameValueManualSave(t *testing.T) { |
| 107 | service, err := NewService("local", NewFilesystemPersistence(t.TempDir())) |
| 108 | if err != nil { |
| 109 | t.Fatal(err) |
| 110 | } |
| 111 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 112 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "title-aba"}) |
| 113 | if err != nil { |
| 114 | t.Fatal(err) |
| 115 | } |
| 116 | if err := service.SetTitle(t.Context(), runtime.Ref(), "A"); err != nil { |
| 117 | t.Fatal(err) |
| 118 | } |
| 119 | original := runtime.Session().Snapshot().Projection.TitleSequence |
| 120 | if err := service.SetTitle(t.Context(), runtime.Ref(), "B"); err != nil { |
| 121 | t.Fatal(err) |
| 122 | } |
| 123 | if err := service.SetTitle(t.Context(), runtime.Ref(), "A"); err != nil { |
| 124 | t.Fatal(err) |
| 125 | } |
| 126 | if err := service.SetTitleIfSequence(t.Context(), runtime.Ref(), original, "stale AI"); !errors.Is(err, ErrSessionTitleChanged) { |
| 127 | t.Fatalf("ABA title commit = %v", err) |
| 128 | } |
| 129 | current := runtime.Session().Snapshot().Projection.TitleSequence |
| 130 | if err := service.SetTitle(t.Context(), runtime.Ref(), "A"); err != nil { |
| 131 | t.Fatal(err) |
| 132 | } |
| 133 | if err := service.SetTitleIfSequence(t.Context(), runtime.Ref(), current, "stale same-value AI"); !errors.Is(err, ErrSessionTitleChanged) { |
| 134 | t.Fatalf("same-value title commit = %v", err) |
| 135 | } |
| 136 | } |
| 137 |