| 1 | package session |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "testing" |
| 10 | |
| 11 | "reasonix/internal/projectiondb" |
| 12 | "reasonix/internal/provider" |
| 13 | ) |
| 14 | |
| 15 | func appendCompleteMessage(t *testing.T, s *Session, operationID string, message provider.Message) error { |
| 16 | t.Helper() |
| 17 | payload, err := json.Marshal(map[string]any{"message": message}) |
| 18 | if err != nil { |
| 19 | t.Fatal(err) |
| 20 | } |
| 21 | _, err = s.Append(t.Context(), Batch{OperationID: operationID, Events: []Event{{Kind: "message/complete", Payload: payload}}}) |
| 22 | return err |
| 23 | } |
| 24 | |
| 25 | func mustAppendCompleteMessage(t *testing.T, s *Session, operationID string, message provider.Message) { |
| 26 | t.Helper() |
| 27 | if err := appendCompleteMessage(t, s, operationID, message); err != nil { |
| 28 | t.Fatalf("append %s: %v", message.ID, err) |
| 29 | } |
| 30 | if _, err := s.Flush(t.Context()); err != nil { |
| 31 | t.Fatal(err) |
| 32 | } |
| 33 | } |
| 34 | |
| 35 | // A durable write publishes a recovery checkpoint, after which a Service-owned |
| 36 | // runtime holds no message bodies. The writer must still know every live id. |
| 37 | func TestMessageCompleteRefusesIDAlreadyDurable(t *testing.T) { |
| 38 | reopen := map[string]func(t *testing.T, root string, ref SessionRef){ |
| 39 | "same runtime": nil, |
| 40 | "reopened from checkpoint": func(*testing.T, string, SessionRef) {}, |
| 41 | "reopened without checkpoint": removeRecoveryCheckpoints, |
| 42 | } |
| 43 | for name, prepareReopen := range reopen { |
| 44 | t.Run(name, func(t *testing.T) { |
| 45 | root := filepath.Join(t.TempDir(), "sessions") |
| 46 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 47 | if err != nil { |
| 48 | t.Fatal(err) |
| 49 | } |
| 50 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 51 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "dup-complete-write"}) |
| 52 | if err != nil { |
| 53 | t.Fatal(err) |
| 54 | } |
| 55 | mustAppendCompleteMessage(t, runtime.Session(), "user", provider.Message{ID: "claimed", Role: provider.RoleUser, Content: "the user's request"}) |
| 56 | mustAppendCompleteMessage(t, runtime.Session(), "answer", provider.Message{ID: "answer", Role: provider.RoleAssistant, Content: "plan"}) |
| 57 | target := runtime.Session() |
| 58 | if prepareReopen != nil { |
| 59 | ref := runtime.Ref() |
| 60 | if err := service.CloseAll(t.Context()); err != nil { |
| 61 | t.Fatal(err) |
| 62 | } |
| 63 | prepareReopen(t, root, ref) |
| 64 | service, err = NewService("local", NewFilesystemPersistence(root)) |
| 65 | if err != nil { |
| 66 | t.Fatal(err) |
| 67 | } |
| 68 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 69 | opened, err := service.Open(t.Context(), ref) |
| 70 | if err != nil { |
| 71 | t.Fatal(err) |
| 72 | } |
| 73 | t.Cleanup(func() { _ = opened.Release(context.Background()) }) |
| 74 | target = opened.Runtime().Session() |
| 75 | } |
| 76 | err = appendCompleteMessage(t, target, "reused", provider.Message{ID: "claimed", Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: "plan approved"}) |
| 77 | if !errors.Is(err, ErrDuplicateMessageID) { |
| 78 | t.Fatalf("append of an id already durable: %v, want ErrDuplicateMessageID", err) |
| 79 | } |
| 80 | if err := appendCompleteMessage(t, target, "fresh", provider.Message{ID: "fresh", Role: provider.RoleUser, Content: "next"}); err != nil { |
| 81 | t.Fatalf("a refused duplicate must not poison later writes: %v", err) |
| 82 | } |
| 83 | }) |
| 84 | } |
| 85 | } |
| 86 | |
| 87 | func removeRecoveryCheckpoints(t *testing.T, root string, ref SessionRef) { |
| 88 | t.Helper() |
| 89 | path := filepath.Join(recoveryCacheDir(filepath.Join(root, ref.SessionID)), recoveryDBName) |
| 90 | if _, err := os.Stat(path); err != nil { |
| 91 | t.Fatalf("recovery checkpoint store: %v", err) |
| 92 | } |
| 93 | if err := os.Remove(path); err != nil { |
| 94 | t.Fatal(err) |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | // commitDuplicateComplete writes the logged shape: a second message/complete |
| 99 | // for an id the log already holds. Forgetting the id reproduces a writer that |
| 100 | // only saw the ids accepted since its last checkpoint. |
| 101 | func commitDuplicateComplete(t *testing.T, s *Session, message provider.Message) { |
| 102 | t.Helper() |
| 103 | s.mu.Lock() |
| 104 | delete(s.messageIDs, message.ID) |
| 105 | s.mu.Unlock() |
| 106 | mustAppendCompleteMessage(t, s, "duplicate-"+message.ID, message) |
| 107 | } |
| 108 | |
| 109 | func TestDuplicateCompleteOnDiskKeepsFirstOccurrence(t *testing.T) { |
| 110 | root := filepath.Join(t.TempDir(), "sessions") |
| 111 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 112 | if err != nil { |
| 113 | t.Fatal(err) |
| 114 | } |
| 115 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 116 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "dup-complete-read"}) |
| 117 | if err != nil { |
| 118 | t.Fatal(err) |
| 119 | } |
| 120 | s := runtime.Session() |
| 121 | mustAppendCompleteMessage(t, s, "user", provider.Message{ID: "claimed", Role: provider.RoleUser, Content: "the user's request"}) |
| 122 | mustAppendCompleteMessage(t, s, "answer", provider.Message{ID: "answer", Role: provider.RoleAssistant, Content: "plan"}) |
| 123 | compacted, err := json.Marshal(map[string]any{"messages": []provider.Message{{ID: "answer", Role: provider.RoleAssistant, Content: "plan"}}, "reason": "compaction"}) |
| 124 | if err != nil { |
| 125 | t.Fatal(err) |
| 126 | } |
| 127 | if _, err := s.Append(t.Context(), Batch{OperationID: "compact", Events: []Event{{Kind: "model/context-replace", Payload: compacted}}}); err != nil { |
| 128 | t.Fatal(err) |
| 129 | } |
| 130 | commitDuplicateComplete(t, s, provider.Message{ID: "claimed", Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: "plan approved"}) |
| 131 | mustAppendCompleteMessage(t, s, "tail", provider.Message{ID: "tail", Role: provider.RoleAssistant, Content: "done"}) |
| 132 | ref := runtime.Ref() |
| 133 | want := []string{"claimed", "answer", "tail"} |
| 134 | firstKept := func(where string, messages []provider.Message) { |
| 135 | t.Helper() |
| 136 | if got := messageIDs(messages); !idsEqual(got, want) { |
| 137 | t.Fatalf("%s ids = %v, want %v", where, got, want) |
| 138 | } |
| 139 | if messages[0].Content != "the user's request" { |
| 140 | t.Fatalf("%s kept %q, want the first occurrence", where, messages[0].Content) |
| 141 | } |
| 142 | } |
| 143 | if err := service.CloseAll(t.Context()); err != nil { |
| 144 | t.Fatal(err) |
| 145 | } |
| 146 | for _, removeCaches := range []bool{false, true} { |
| 147 | if removeCaches { |
| 148 | for _, cache := range []string{".query-cache", ".recovery-cache"} { |
| 149 | if err := os.RemoveAll(filepath.Join(root, cache)); err != nil { |
| 150 | t.Fatal(err) |
| 151 | } |
| 152 | } |
| 153 | } |
| 154 | reopened, err := NewService("local", NewFilesystemPersistence(root)) |
| 155 | if err != nil { |
| 156 | t.Fatal(err) |
| 157 | } |
| 158 | query := reopened.Query() |
| 159 | snapshot, err := query.Snapshot(t.Context(), ref) |
| 160 | if err != nil { |
| 161 | t.Fatalf("snapshot (caches removed=%v): %v", removeCaches, err) |
| 162 | } |
| 163 | firstKept("snapshot", snapshot.Projection.Messages) |
| 164 | info, err := query.RefreshMetadata(t.Context(), ref) |
| 165 | if err != nil { |
| 166 | t.Fatalf("catalog metadata (caches removed=%v): %v", removeCaches, err) |
| 167 | } |
| 168 | if info.Preview != "the user's request" { |
| 169 | t.Fatalf("catalog preview (caches removed=%v) = %q, want the first occurrence", removeCaches, info.Preview) |
| 170 | } |
| 171 | if _, _, err := query.prepareHistoryIndex(t.Context(), ref); err != nil { |
| 172 | t.Fatalf("prepare history index (caches removed=%v): %v", removeCaches, err) |
| 173 | } |
| 174 | page, err := query.ReadHistoryWindow(t.Context(), ref, HistoryWindowRequest{Anchor: "newest"}) |
| 175 | if err != nil || page.Status != "ready" { |
| 176 | t.Fatalf("read history (caches removed=%v): %+v, %v", removeCaches, page, err) |
| 177 | } |
| 178 | if got := windowIDs(t, page); !idsEqual(got, want) { |
| 179 | t.Fatalf("history ids (caches removed=%v) = %v, want %v", removeCaches, got, want) |
| 180 | } |
| 181 | if search := searchHistoryReady(t, query, ref, "approved", "", 10); len(search.Hits) != 0 { |
| 182 | t.Fatalf("search found the discarded occurrence (caches removed=%v): %+v", removeCaches, search) |
| 183 | } |
| 184 | opened, err := reopened.Open(t.Context(), ref) |
| 185 | if err != nil { |
| 186 | t.Fatalf("reopen (caches removed=%v): %v", removeCaches, err) |
| 187 | } |
| 188 | firstKept("reopened", opened.Runtime().Session().Snapshot().Projection.Messages) |
| 189 | // Without a checkpoint the open replays the whole log, as the first open |
| 190 | // after an upgrade does. |
| 191 | if removeCaches { |
| 192 | if got := opened.Runtime().Session().CatalogMetadata().Preview; got != "the user's request" { |
| 193 | t.Fatalf("replayed preview = %q, want the first occurrence", got) |
| 194 | } |
| 195 | if got := messageIDs(opened.Runtime().Session().DeriveMessages()); !idsEqual(got, []string{"answer", "tail"}) { |
| 196 | t.Fatalf("replayed model workset = %v, want the compacted context without the repeat", got) |
| 197 | } |
| 198 | } |
| 199 | if err := opened.Release(t.Context()); err != nil { |
| 200 | t.Fatal(err) |
| 201 | } |
| 202 | if err := reopened.CloseAll(t.Context()); err != nil { |
| 203 | t.Fatal(err) |
| 204 | } |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | // A checkpoint older than the repeat replays it in the tail. An older build |
| 209 | // demoting the current checkpoint and appending the repeat reaches this. |
| 210 | func TestDuplicateCompleteInCheckpointTailKeepsFirstOccurrence(t *testing.T) { |
| 211 | root := filepath.Join(t.TempDir(), "sessions") |
| 212 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 213 | if err != nil { |
| 214 | t.Fatal(err) |
| 215 | } |
| 216 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 217 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "dup-complete-tail"}) |
| 218 | if err != nil { |
| 219 | t.Fatal(err) |
| 220 | } |
| 221 | ref := runtime.Ref() |
| 222 | mustAppendCompleteMessage(t, runtime.Session(), "user", provider.Message{ID: "claimed", Role: provider.RoleUser, Content: "the user's request"}) |
| 223 | mustAppendCompleteMessage(t, runtime.Session(), "answer", provider.Message{ID: "answer", Role: provider.RoleAssistant, Content: "plan"}) |
| 224 | if err := service.CloseAll(t.Context()); err != nil { |
| 225 | t.Fatal(err) |
| 226 | } |
| 227 | checkpoints := filepath.Join(recoveryCacheDir(filepath.Join(root, ref.SessionID)), recoveryDBName) |
| 228 | stale, err := os.ReadFile(checkpoints) |
| 229 | if err != nil { |
| 230 | t.Fatal(err) |
| 231 | } |
| 232 | |
| 233 | service, err = NewService("local", NewFilesystemPersistence(root)) |
| 234 | if err != nil { |
| 235 | t.Fatal(err) |
| 236 | } |
| 237 | opened, err := service.Open(t.Context(), ref) |
| 238 | if err != nil { |
| 239 | t.Fatal(err) |
| 240 | } |
| 241 | commitDuplicateComplete(t, opened.Runtime().Session(), provider.Message{ID: "claimed", Role: provider.RoleUser, Origin: provider.MessageOriginHost, Content: "plan approved"}) |
| 242 | mustAppendCompleteMessage(t, opened.Runtime().Session(), "tail", provider.Message{ID: "tail", Role: provider.RoleAssistant, Content: "done"}) |
| 243 | if err := opened.Release(t.Context()); err != nil { |
| 244 | t.Fatal(err) |
| 245 | } |
| 246 | if err := service.CloseAll(t.Context()); err != nil { |
| 247 | t.Fatal(err) |
| 248 | } |
| 249 | if err := os.WriteFile(checkpoints, stale, 0o600); err != nil { |
| 250 | t.Fatal(err) |
| 251 | } |
| 252 | |
| 253 | service, err = NewService("local", NewFilesystemPersistence(root)) |
| 254 | if err != nil { |
| 255 | t.Fatal(err) |
| 256 | } |
| 257 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 258 | reopened, err := service.Open(t.Context(), ref) |
| 259 | if err != nil { |
| 260 | t.Fatal(err) |
| 261 | } |
| 262 | t.Cleanup(func() { _ = reopened.Release(context.Background()) }) |
| 263 | s := reopened.Runtime().Session() |
| 264 | model := s.DeriveMessages() |
| 265 | if got := messageIDs(model); !idsEqual(got, []string{"claimed", "answer", "tail"}) || model[0].Content != "the user's request" { |
| 266 | t.Fatalf("model workset after the checkpoint tail = %v (%q), want the first occurrence only", got, model[0].Content) |
| 267 | } |
| 268 | if got := s.CatalogMetadata().Preview; got != "the user's request" { |
| 269 | t.Fatalf("preview after the checkpoint tail = %q, want the first occurrence", got) |
| 270 | } |
| 271 | } |
| 272 | |
| 273 | // The repeat of an earlier turn's answer must not become a later turn's final |
| 274 | // reply in the history index. |
| 275 | func TestDuplicateCompleteDoesNotBecomeALaterTurnsFinal(t *testing.T) { |
| 276 | root := filepath.Join(t.TempDir(), "sessions") |
| 277 | service, err := NewService("local", NewFilesystemPersistence(root)) |
| 278 | if err != nil { |
| 279 | t.Fatal(err) |
| 280 | } |
| 281 | t.Cleanup(func() { _ = service.CloseAll(context.Background()) }) |
| 282 | runtime, err := service.Create(t.Context(), CreateOptions{SessionID: "dup-complete-turn"}) |
| 283 | if err != nil { |
| 284 | t.Fatal(err) |
| 285 | } |
| 286 | s := runtime.Session() |
| 287 | appendTurn := func(turnID string, messages ...provider.Message) { |
| 288 | t.Helper() |
| 289 | if _, err := s.Append(t.Context(), Batch{OperationID: turnID + "-start", TurnID: turnID, Events: []Event{{Kind: "turn/start"}}}); err != nil { |
| 290 | t.Fatal(err) |
| 291 | } |
| 292 | for _, message := range messages { |
| 293 | if message.ID == "answer" && turnID == "second" { |
| 294 | s.mu.Lock() |
| 295 | delete(s.messageIDs, message.ID) |
| 296 | s.mu.Unlock() |
| 297 | } |
| 298 | payload, err := json.Marshal(map[string]any{"message": message}) |
| 299 | if err != nil { |
| 300 | t.Fatal(err) |
| 301 | } |
| 302 | if _, err := s.Append(t.Context(), Batch{OperationID: turnID + "-" + message.ID, TurnID: turnID, Events: []Event{{Kind: "message/complete", Payload: payload}}}); err != nil { |
| 303 | t.Fatal(err) |
| 304 | } |
| 305 | } |
| 306 | if _, err := s.Append(t.Context(), Batch{OperationID: turnID + "-end", TurnID: turnID, Events: []Event{{Kind: "turn/end", Payload: []byte(`{"status":"completed"}`)}}}); err != nil { |
| 307 | t.Fatal(err) |
| 308 | } |
| 309 | } |
| 310 | appendTurn("first", provider.Message{ID: "question", Role: provider.RoleUser, Content: "q"}, provider.Message{ID: "answer", Role: provider.RoleAssistant, Content: "a", WorkDurationMs: 1000}) |
| 311 | appendTurn("second", provider.Message{ID: "follow-up", Role: provider.RoleUser, Content: "again"}, provider.Message{ID: "answer", Role: provider.RoleAssistant, Content: "repeat", WorkDurationMs: 999999}) |
| 312 | if _, err := s.Flush(t.Context()); err != nil { |
| 313 | t.Fatal(err) |
| 314 | } |
| 315 | _, path, err := service.Query().prepareHistoryIndex(t.Context(), runtime.Ref()) |
| 316 | if err != nil { |
| 317 | t.Fatal(err) |
| 318 | } |
| 319 | handle, err := projectiondb.Open(t.Context(), projectiondb.OpenOptions{Path: path, Migrations: historyMigrations, RequireDisk: true, MaxOpenConns: 1}) |
| 320 | if err != nil { |
| 321 | t.Fatal(err) |
| 322 | } |
| 323 | defer handle.DB.Close() |
| 324 | var final string |
| 325 | if err := handle.DB.QueryRowContext(t.Context(), `SELECT COALESCE(final_message_id,'') FROM turn_summaries WHERE turn_id=?`, "second").Scan(&final); err != nil { |
| 326 | t.Fatal(err) |
| 327 | } |
| 328 | if final == "answer" { |
| 329 | t.Fatalf("the repeated answer became the second turn's final reply") |
| 330 | } |
| 331 | } |
| 332 |