| 1 | package cli |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "errors" |
| 7 | "io" |
| 8 | "net/http" |
| 9 | "net/http/httptest" |
| 10 | "path/filepath" |
| 11 | "strconv" |
| 12 | "strings" |
| 13 | "sync" |
| 14 | "sync/atomic" |
| 15 | "testing" |
| 16 | "time" |
| 17 | |
| 18 | "reasonix/internal/agent" |
| 19 | "reasonix/internal/control" |
| 20 | "reasonix/internal/event" |
| 21 | "reasonix/internal/eventwire" |
| 22 | "reasonix/internal/jobs" |
| 23 | "reasonix/internal/provider" |
| 24 | ) |
| 25 | |
| 26 | type takeoverRecordSink struct { |
| 27 | mu sync.Mutex |
| 28 | events []event.Event |
| 29 | } |
| 30 | |
| 31 | func (s *takeoverRecordSink) Emit(e event.Event) { |
| 32 | s.mu.Lock() |
| 33 | s.events = append(s.events, e) |
| 34 | s.mu.Unlock() |
| 35 | } |
| 36 | |
| 37 | func TestCLITakeoverManagerRetriesSameFrameBatchInOrder(t *testing.T) { |
| 38 | var mu sync.Mutex |
| 39 | var requests [][]eventwire.Event |
| 40 | var calls atomic.Int32 |
| 41 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 42 | var body struct { |
| 43 | Frames []eventwire.Event `json:"frames"` |
| 44 | } |
| 45 | if err := json.NewDecoder(r.Body).Decode(&body); err != nil { |
| 46 | t.Error(err) |
| 47 | w.WriteHeader(http.StatusBadRequest) |
| 48 | return |
| 49 | } |
| 50 | mu.Lock() |
| 51 | requests = append(requests, append([]eventwire.Event(nil), body.Frames...)) |
| 52 | mu.Unlock() |
| 53 | if calls.Add(1) == 1 { |
| 54 | w.WriteHeader(http.StatusServiceUnavailable) |
| 55 | return |
| 56 | } |
| 57 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 58 | })) |
| 59 | defer srv.Close() |
| 60 | |
| 61 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 62 | m.binding = &cliTakeoverBinding{ |
| 63 | path: "session.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 64 | grant: cliTakeoverGrant{MirrorID: "mirror-1"}, |
| 65 | } |
| 66 | m.Emit(event.Event{Kind: event.Text, Text: "first"}) |
| 67 | m.Emit(event.Event{Kind: event.Text, Text: "second"}) |
| 68 | if !m.push(false) { |
| 69 | t.Fatal("frame push stopped unexpectedly") |
| 70 | } |
| 71 | if !m.push(false) { |
| 72 | t.Fatal("frame push stopped unexpectedly") |
| 73 | } |
| 74 | |
| 75 | mu.Lock() |
| 76 | defer mu.Unlock() |
| 77 | if len(requests) != 2 || len(requests[0]) != 2 || len(requests[1]) != 2 { |
| 78 | t.Fatalf("request batches = %+v, want two complete two-frame batches", requests) |
| 79 | } |
| 80 | for i := range requests[0] { |
| 81 | if requests[0][i].Text != requests[1][i].Text { |
| 82 | t.Fatalf("retry reordered frame %d: %q != %q", i, requests[0][i].Text, requests[1][i].Text) |
| 83 | } |
| 84 | } |
| 85 | } |
| 86 | |
| 87 | func TestCLITakeoverManagerChunksWithoutDroppingFrames(t *testing.T) { |
| 88 | var got []string |
| 89 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 90 | var body struct { |
| 91 | Frames []eventwire.Event `json:"frames"` |
| 92 | } |
| 93 | if err := json.NewDecoder(r.Body).Decode(&body); err != nil { |
| 94 | t.Fatal(err) |
| 95 | } |
| 96 | if len(body.Frames) > cliTakeoverMaxFrames { |
| 97 | t.Fatalf("batch size = %d", len(body.Frames)) |
| 98 | } |
| 99 | for _, frame := range body.Frames { |
| 100 | got = append(got, frame.Text) |
| 101 | } |
| 102 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 103 | })) |
| 104 | defer srv.Close() |
| 105 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 106 | m.binding = &cliTakeoverBinding{ |
| 107 | path: "session.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 108 | grant: cliTakeoverGrant{MirrorID: "mirror-1"}, |
| 109 | } |
| 110 | for i := range cliTakeoverMaxFrames + 37 { |
| 111 | m.Emit(event.Event{Kind: event.Text, Text: strconv.Itoa(i)}) |
| 112 | } |
| 113 | if !m.push(false) { |
| 114 | t.Fatal("chunked frame push stopped unexpectedly") |
| 115 | } |
| 116 | if !m.push(false) { |
| 117 | t.Fatal("chunked frame push stopped unexpectedly") |
| 118 | } |
| 119 | if len(got) != cliTakeoverMaxFrames+37 { |
| 120 | t.Fatalf("received %d frames", len(got)) |
| 121 | } |
| 122 | for i, text := range got { |
| 123 | if text != strconv.Itoa(i) { |
| 124 | t.Fatalf("frame %d = %q", i, text) |
| 125 | } |
| 126 | } |
| 127 | } |
| 128 | |
| 129 | func TestCLITakeoverManagerHeartbeatSendsEmptyFrameBatch(t *testing.T) { |
| 130 | got := make(chan int, 1) |
| 131 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 132 | var body struct { |
| 133 | Frames []eventwire.Event `json:"frames"` |
| 134 | } |
| 135 | if err := json.NewDecoder(r.Body).Decode(&body); err != nil { |
| 136 | t.Error(err) |
| 137 | } |
| 138 | got <- len(body.Frames) |
| 139 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 140 | })) |
| 141 | defer srv.Close() |
| 142 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 143 | m.binding = &cliTakeoverBinding{ |
| 144 | path: "session.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 145 | grant: cliTakeoverGrant{MirrorID: "mirror-1"}, |
| 146 | } |
| 147 | if !m.push(true) { |
| 148 | t.Fatal("heartbeat stopped the manager") |
| 149 | } |
| 150 | if frames := <-got; frames != 0 { |
| 151 | t.Fatalf("heartbeat frames = %d, want 0", frames) |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | func TestCLITakeoverManagerReclaimReturnsLeaseAndSignalsExit(t *testing.T) { |
| 156 | dir := t.TempDir() |
| 157 | path := filepath.Join(dir, "active.jsonl") |
| 158 | session := agent.NewSession("system") |
| 159 | session.Add(provider.Message{Role: provider.RoleUser, Content: "hello"}) |
| 160 | if err := session.Save(path); err != nil { |
| 161 | t.Fatal(err) |
| 162 | } |
| 163 | loaded, err := agent.LoadSession(path) |
| 164 | if err != nil { |
| 165 | t.Fatal(err) |
| 166 | } |
| 167 | exec := agent.New(nil, nil, loaded, agent.Options{}, &takeoverRecordSink{}) |
| 168 | ctrl := newOwnedTestController(t, control.Options{Executor: exec, SessionDir: dir, SessionPath: path}) |
| 169 | defer ctrl.Close() |
| 170 | leases := control.NewSessionLeaseKeeper() |
| 171 | defer leases.Release() |
| 172 | if err := leases.Rebind(path); err != nil { |
| 173 | t.Fatal(err) |
| 174 | } |
| 175 | if err := leases.BindControllerAuthority(ctrl); err != nil { |
| 176 | t.Fatal(err) |
| 177 | } |
| 178 | |
| 179 | var mirrorEnded atomic.Bool |
| 180 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 181 | switch r.URL.Path { |
| 182 | case "/external/frames": |
| 183 | _ = json.NewEncoder(w).Encode(map[string]any{"reclaimRequested": true, "reclaimMode": "wait"}) |
| 184 | case "/mirror-end": |
| 185 | mirrorEnded.Store(true) |
| 186 | w.WriteHeader(http.StatusNoContent) |
| 187 | default: |
| 188 | w.WriteHeader(http.StatusNotFound) |
| 189 | } |
| 190 | })) |
| 191 | defer srv.Close() |
| 192 | |
| 193 | exited := make(chan struct{}, 1) |
| 194 | m := newCLITakeoverManager(&takeoverRecordSink{}, leases) |
| 195 | m.AttachController(ctrl) |
| 196 | m.SetYieldCallback(func() { exited <- struct{}{} }) |
| 197 | m.Activate(&cliTakeoverBinding{ |
| 198 | path: path, record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 199 | grant: cliTakeoverGrant{ |
| 200 | MirrorID: "mirror-1", SourceWriterID: "serve-writer", ReturnHandoffID: "return-1", |
| 201 | }, |
| 202 | }) |
| 203 | m.Emit(event.Event{Kind: event.Text, Text: "answer"}) |
| 204 | select { |
| 205 | case <-exited: |
| 206 | case <-time.After(5 * time.Second): |
| 207 | t.Fatal("reclaim did not return the lease") |
| 208 | } |
| 209 | if !m.Returned() || !mirrorEnded.Load() { |
| 210 | t.Fatalf("returned=%v mirrorEnded=%v", m.Returned(), mirrorEnded.Load()) |
| 211 | } |
| 212 | info, err := agent.LoadSessionLeaseInfo(path) |
| 213 | if err != nil { |
| 214 | t.Fatal(err) |
| 215 | } |
| 216 | if info == nil || info.HandoffTo != "serve-writer" || info.HandoffID != "return-1" { |
| 217 | t.Fatalf("reverse reservation = %+v", info) |
| 218 | } |
| 219 | if err := m.Close(); err != nil { |
| 220 | t.Fatal(err) |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | func TestCLITakeoverManagerReclaimWaitsForBackgroundJobs(t *testing.T) { |
| 225 | dir := t.TempDir() |
| 226 | path := filepath.Join(dir, "background.jsonl") |
| 227 | session := agent.NewSession("system") |
| 228 | if err := session.Save(path); err != nil { |
| 229 | t.Fatal(err) |
| 230 | } |
| 231 | manager := jobs.NewManager(event.Discard) |
| 232 | ctrl := newOwnedTestController(t, control.Options{ |
| 233 | Executor: agent.New(nil, nil, session, agent.Options{}, event.Discard), |
| 234 | Jobs: manager, |
| 235 | SessionDir: dir, |
| 236 | SessionPath: path, |
| 237 | }) |
| 238 | defer ctrl.Close() |
| 239 | jobStarted := make(chan struct{}) |
| 240 | releaseJob := make(chan struct{}) |
| 241 | manager.StartForSession(agent.BranchID(path), "task", "background", func(ctx context.Context, _ io.Writer) (string, error) { |
| 242 | close(jobStarted) |
| 243 | select { |
| 244 | case <-releaseJob: |
| 245 | return "done", nil |
| 246 | case <-ctx.Done(): |
| 247 | return "", ctx.Err() |
| 248 | } |
| 249 | }) |
| 250 | <-jobStarted |
| 251 | |
| 252 | leases := control.NewSessionLeaseKeeper() |
| 253 | defer leases.Release() |
| 254 | if err := leases.Rebind(path); err != nil { |
| 255 | t.Fatal(err) |
| 256 | } |
| 257 | if err := leases.BindControllerAuthority(ctrl); err != nil { |
| 258 | t.Fatal(err) |
| 259 | } |
| 260 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 261 | switch r.URL.Path { |
| 262 | case "/mirror-end": |
| 263 | w.WriteHeader(http.StatusNoContent) |
| 264 | default: |
| 265 | w.WriteHeader(http.StatusNotFound) |
| 266 | } |
| 267 | })) |
| 268 | defer srv.Close() |
| 269 | |
| 270 | exited := make(chan struct{}, 1) |
| 271 | m := newCLITakeoverManager(event.Discard, leases) |
| 272 | m.AttachController(ctrl) |
| 273 | m.SetYieldCallback(func() { exited <- struct{}{} }) |
| 274 | binding := &cliTakeoverBinding{ |
| 275 | path: path, record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 276 | grant: cliTakeoverGrant{MirrorID: "mirror", SourceWriterID: "serve-writer", ReturnHandoffID: "return"}, |
| 277 | } |
| 278 | m.mu.Lock() |
| 279 | m.binding = binding |
| 280 | m.revision = 1 |
| 281 | m.mu.Unlock() |
| 282 | m.requestYieldFor(binding, 1, false) |
| 283 | select { |
| 284 | case <-exited: |
| 285 | t.Fatal("reclaim returned the lease while a background job was active") |
| 286 | case <-time.After(150 * time.Millisecond): |
| 287 | } |
| 288 | contender, err := agent.TryAcquireSessionLease(path) |
| 289 | if contender != nil { |
| 290 | contender.Release() |
| 291 | } |
| 292 | if !errors.Is(err, agent.ErrSessionLeaseHeld) { |
| 293 | t.Fatalf("background reclaim lease error = %v, want held", err) |
| 294 | } |
| 295 | |
| 296 | close(releaseJob) |
| 297 | select { |
| 298 | case <-exited: |
| 299 | case <-time.After(5 * time.Second): |
| 300 | t.Fatal("reclaim did not finish after the background job completed") |
| 301 | } |
| 302 | if err := m.Close(); err != nil { |
| 303 | t.Fatal(err) |
| 304 | } |
| 305 | } |
| 306 | |
| 307 | func TestCLITakeoverManagerTransitionsBetweenMirrorsAtomically(t *testing.T) { |
| 308 | dir := t.TempDir() |
| 309 | current := filepath.Join(dir, "current.jsonl") |
| 310 | nextPath := filepath.Join(dir, "next.jsonl") |
| 311 | leases := control.NewSessionLeaseKeeper() |
| 312 | defer leases.Release() |
| 313 | if err := leases.Rebind(current); err != nil { |
| 314 | t.Fatal(err) |
| 315 | } |
| 316 | nextSource, err := agent.TryAcquireSessionLease(nextPath) |
| 317 | if err != nil { |
| 318 | t.Fatal(err) |
| 319 | } |
| 320 | if err := nextSource.ReleaseForHandoff(agent.SessionWriterID(), "forward-next"); err != nil { |
| 321 | t.Fatal(err) |
| 322 | } |
| 323 | |
| 324 | var oldEnded atomic.Bool |
| 325 | oldServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 326 | if r.URL.Path == "/mirror-end" { |
| 327 | oldEnded.Store(true) |
| 328 | w.WriteHeader(http.StatusNoContent) |
| 329 | return |
| 330 | } |
| 331 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 332 | })) |
| 333 | defer oldServe.Close() |
| 334 | m := newCLITakeoverManager(&takeoverRecordSink{}, leases) |
| 335 | m.binding = &cliTakeoverBinding{ |
| 336 | path: current, record: cliServeRecord{base: oldServe.URL}, client: oldServe.Client(), |
| 337 | grant: cliTakeoverGrant{MirrorID: "old-mirror", SourceWriterID: "old-serve", ReturnHandoffID: "return-old"}, |
| 338 | } |
| 339 | next := &cliTakeoverBinding{ |
| 340 | path: nextPath, |
| 341 | grant: cliTakeoverGrant{MirrorID: "next-mirror", SourceWriterID: agent.SessionWriterID(), HandoffID: "forward-next"}, |
| 342 | } |
| 343 | next.previous, err = leases.RebindDetachingWithHandoff(nextPath, agent.SessionWriterID(), "forward-next") |
| 344 | if err != nil { |
| 345 | t.Fatal(err) |
| 346 | } |
| 347 | next.priorMirror = m.binding |
| 348 | if err := next.commitPrevious(m); err != nil { |
| 349 | t.Fatalf("commitPrevious: %v", err) |
| 350 | } |
| 351 | if got := leases.HeldPath(); got != agent.CanonicalSessionPath(nextPath) { |
| 352 | t.Fatalf("held path = %q, want next", got) |
| 353 | } |
| 354 | if !oldEnded.Load() { |
| 355 | t.Fatal("old mirror was not ended after the atomic keeper transition") |
| 356 | } |
| 357 | info, err := agent.LoadSessionLeaseInfo(current) |
| 358 | if err != nil { |
| 359 | t.Fatal(err) |
| 360 | } |
| 361 | if info == nil || info.HandoffTo != "old-serve" || info.HandoffID != "return-old" { |
| 362 | t.Fatalf("old reverse reservation = %+v", info) |
| 363 | } |
| 364 | } |
| 365 | |
| 366 | func TestCLIFailedTakeoverRestoresSourceAndReturnsTarget(t *testing.T) { |
| 367 | dir := t.TempDir() |
| 368 | current := filepath.Join(dir, "current.jsonl") |
| 369 | target := filepath.Join(dir, "target.jsonl") |
| 370 | leases := control.NewSessionLeaseKeeper() |
| 371 | defer leases.Release() |
| 372 | if err := leases.Rebind(current); err != nil { |
| 373 | t.Fatal(err) |
| 374 | } |
| 375 | targetSource, err := agent.TryAcquireSessionLease(target) |
| 376 | if err != nil { |
| 377 | t.Fatal(err) |
| 378 | } |
| 379 | if err := targetSource.ReleaseForHandoff(agent.SessionWriterID(), "forward-target"); err != nil { |
| 380 | t.Fatal(err) |
| 381 | } |
| 382 | previous, err := leases.RebindDetachingWithHandoff(target, agent.SessionWriterID(), "forward-target") |
| 383 | if err != nil { |
| 384 | t.Fatal(err) |
| 385 | } |
| 386 | var mirrorEnded atomic.Bool |
| 387 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 388 | if r.URL.Path == "/mirror-end" { |
| 389 | mirrorEnded.Store(true) |
| 390 | w.WriteHeader(http.StatusNoContent) |
| 391 | return |
| 392 | } |
| 393 | w.WriteHeader(http.StatusNotFound) |
| 394 | })) |
| 395 | defer srv.Close() |
| 396 | binding := &cliTakeoverBinding{ |
| 397 | path: target, previous: previous, record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 398 | grant: cliTakeoverGrant{ |
| 399 | MirrorID: "target-mirror", SourceWriterID: "target-serve", ReturnHandoffID: "return-target", |
| 400 | }, |
| 401 | } |
| 402 | manager := newCLITakeoverManager(&takeoverRecordSink{}, leases) |
| 403 | if err := cliReturnFailedTakeover(binding, leases, manager); err != nil { |
| 404 | t.Fatal(err) |
| 405 | } |
| 406 | if got := leases.HeldPath(); got != agent.CanonicalSessionPath(current) { |
| 407 | t.Fatalf("held path = %q, want restored source", got) |
| 408 | } |
| 409 | info, err := agent.LoadSessionLeaseInfo(target) |
| 410 | if err != nil { |
| 411 | t.Fatal(err) |
| 412 | } |
| 413 | if info == nil || info.HandoffTo != "target-serve" || info.HandoffID != "return-target" { |
| 414 | t.Fatalf("target reverse reservation = %+v", info) |
| 415 | } |
| 416 | if !mirrorEnded.Load() { |
| 417 | t.Fatal("failed target mirror was not ended") |
| 418 | } |
| 419 | } |
| 420 | |
| 421 | func TestCLITakeoverManagerSerializesRequestsAndDiscardsStaleReclaim(t *testing.T) { |
| 422 | firstEntered := make(chan struct{}) |
| 423 | releaseFirst := make(chan struct{}) |
| 424 | var inFlight atomic.Int32 |
| 425 | var maxInFlight atomic.Int32 |
| 426 | var mu sync.Mutex |
| 427 | var batches [][]eventwire.Event |
| 428 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 429 | current := inFlight.Add(1) |
| 430 | defer inFlight.Add(-1) |
| 431 | for { |
| 432 | observed := maxInFlight.Load() |
| 433 | if current <= observed || maxInFlight.CompareAndSwap(observed, current) { |
| 434 | break |
| 435 | } |
| 436 | } |
| 437 | var body struct { |
| 438 | Frames []eventwire.Event `json:"frames"` |
| 439 | } |
| 440 | if err := json.NewDecoder(r.Body).Decode(&body); err != nil { |
| 441 | t.Error(err) |
| 442 | return |
| 443 | } |
| 444 | mu.Lock() |
| 445 | batches = append(batches, append([]eventwire.Event(nil), body.Frames...)) |
| 446 | call := len(batches) |
| 447 | mu.Unlock() |
| 448 | if call == 1 { |
| 449 | close(firstEntered) |
| 450 | <-releaseFirst |
| 451 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": true}) |
| 452 | return |
| 453 | } |
| 454 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 455 | })) |
| 456 | defer srv.Close() |
| 457 | |
| 458 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 459 | firstBinding := &cliTakeoverBinding{ |
| 460 | path: "a.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 461 | grant: cliTakeoverGrant{MirrorID: "mirror-a"}, |
| 462 | } |
| 463 | m.binding = firstBinding |
| 464 | m.revision = 1 |
| 465 | m.Emit(event.Event{Kind: event.Text, Text: "A"}) |
| 466 | firstDone := make(chan struct{}) |
| 467 | go func() { |
| 468 | m.push(false) |
| 469 | close(firstDone) |
| 470 | }() |
| 471 | <-firstEntered |
| 472 | |
| 473 | // Model a completed generation switch while the old HTTP response is in |
| 474 | // flight. The response must be ignored before it can request a reclaim. |
| 475 | m.mu.Lock() |
| 476 | m.binding = &cliTakeoverBinding{ |
| 477 | path: "b.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 478 | grant: cliTakeoverGrant{MirrorID: "mirror-b"}, |
| 479 | } |
| 480 | m.revision++ |
| 481 | m.mu.Unlock() |
| 482 | m.Emit(event.Event{Kind: event.Text, Text: "B"}) |
| 483 | secondDone := make(chan struct{}) |
| 484 | go func() { |
| 485 | m.push(false) |
| 486 | close(secondDone) |
| 487 | }() |
| 488 | close(releaseFirst) |
| 489 | <-firstDone |
| 490 | <-secondDone |
| 491 | |
| 492 | if got := maxInFlight.Load(); got != 1 { |
| 493 | t.Fatalf("max requests in flight = %d, want 1", got) |
| 494 | } |
| 495 | mu.Lock() |
| 496 | defer mu.Unlock() |
| 497 | if len(batches) != 2 || len(batches[0]) != 1 || batches[0][0].Text != "A" || len(batches[1]) != 1 || batches[1][0].Text != "B" { |
| 498 | t.Fatalf("batches = %+v, want A then B", batches) |
| 499 | } |
| 500 | if m.Reclaiming() { |
| 501 | t.Fatal("stale reclaim response affected the new binding") |
| 502 | } |
| 503 | } |
| 504 | |
| 505 | func TestCLITakeoverManagerReadoptsOnAuthFailureAndServerMove(t *testing.T) { |
| 506 | var delivered atomic.Int32 |
| 507 | newServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 508 | switch r.URL.Path { |
| 509 | case "/auth/token": |
| 510 | w.WriteHeader(http.StatusNoContent) |
| 511 | case "/adopt": |
| 512 | _ = json.NewEncoder(w).Encode(cliTakeoverGrant{ |
| 513 | SessionPath: "session.jsonl", MirrorID: "mirror-new", ReturnHandoffID: "return-new", |
| 514 | SourceWriterID: "serve-new", TargetWriterID: agent.SessionWriterID(), |
| 515 | }) |
| 516 | case "/external/frames": |
| 517 | delivered.Add(1) |
| 518 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 519 | default: |
| 520 | w.WriteHeader(http.StatusNotFound) |
| 521 | } |
| 522 | })) |
| 523 | defer newServe.Close() |
| 524 | oldServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { |
| 525 | w.WriteHeader(http.StatusUnauthorized) |
| 526 | })) |
| 527 | defer oldServe.Close() |
| 528 | originalDiscover := discoverCLIServesForTakeover |
| 529 | discoverCLIServesForTakeover = func() []cliServeRecord { |
| 530 | return []cliServeRecord{{base: newServe.URL, token: "fresh"}} |
| 531 | } |
| 532 | t.Cleanup(func() { discoverCLIServesForTakeover = originalDiscover }) |
| 533 | |
| 534 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 535 | prior := &cliTakeoverBinding{path: "prior.jsonl"} |
| 536 | m.binding = &cliTakeoverBinding{ |
| 537 | path: "session.jsonl", record: cliServeRecord{base: oldServe.URL}, client: oldServe.Client(), |
| 538 | grant: cliTakeoverGrant{MirrorID: "mirror-old"}, canonical: true, priorMirror: prior, |
| 539 | } |
| 540 | m.revision = 1 |
| 541 | m.Emit(event.Event{Kind: event.Text, Text: "recover me"}) |
| 542 | if !m.push(false) { |
| 543 | t.Fatal("manager stopped during re-adopt") |
| 544 | } |
| 545 | if !m.push(false) { |
| 546 | t.Fatal("manager stopped during re-adopt") |
| 547 | } |
| 548 | if delivered.Load() != 1 { |
| 549 | t.Fatalf("new serve received %d frame batches, want 1", delivered.Load()) |
| 550 | } |
| 551 | binding, _, _, revision := m.snapshot() |
| 552 | if binding == nil || binding.grant.MirrorID != "mirror-new" || !binding.canonical || binding.priorMirror != prior || revision <= 1 { |
| 553 | t.Fatalf("binding after re-adopt = %+v revision=%d", binding, revision) |
| 554 | } |
| 555 | } |
| 556 | |
| 557 | func TestCLITakeoverManagerReadoptsAfterConnectionRefused(t *testing.T) { |
| 558 | deadServe := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) |
| 559 | deadURL := deadServe.URL |
| 560 | deadClient := deadServe.Client() |
| 561 | deadServe.Close() |
| 562 | var delivered atomic.Bool |
| 563 | newServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 564 | switch r.URL.Path { |
| 565 | case "/auth/token": |
| 566 | w.WriteHeader(http.StatusNoContent) |
| 567 | case "/adopt": |
| 568 | _ = json.NewEncoder(w).Encode(cliTakeoverGrant{SessionPath: "session.jsonl", MirrorID: "new", ReturnHandoffID: "return", SourceWriterID: "source", TargetWriterID: agent.SessionWriterID()}) |
| 569 | case "/external/frames": |
| 570 | delivered.Store(true) |
| 571 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 572 | default: |
| 573 | w.WriteHeader(http.StatusNotFound) |
| 574 | } |
| 575 | })) |
| 576 | defer newServe.Close() |
| 577 | originalDiscover := discoverCLIServesForTakeover |
| 578 | discoverCLIServesForTakeover = func() []cliServeRecord { |
| 579 | return []cliServeRecord{{base: newServe.URL, token: "fresh"}} |
| 580 | } |
| 581 | t.Cleanup(func() { discoverCLIServesForTakeover = originalDiscover }) |
| 582 | m := newCLITakeoverManager(nil, nil) |
| 583 | m.binding = &cliTakeoverBinding{path: "session.jsonl", record: cliServeRecord{base: deadURL}, client: deadClient, grant: cliTakeoverGrant{MirrorID: "old"}} |
| 584 | m.revision = 1 |
| 585 | m.Emit(event.Event{Kind: event.Text, Text: "recover"}) |
| 586 | if !m.push(false) { |
| 587 | t.Fatal("connection refusal stopped the manager during rediscovery") |
| 588 | } |
| 589 | if !m.push(false) { |
| 590 | t.Fatal("connection refusal stopped the manager after rediscovery") |
| 591 | } |
| 592 | if !delivered.Load() { |
| 593 | t.Fatal("connection refusal did not rediscover and deliver through the new serve") |
| 594 | } |
| 595 | } |
| 596 | |
| 597 | func TestCLITakeoverManagerKeepsActualRequestBelowEightMiB(t *testing.T) { |
| 598 | var maxBody atomic.Int64 |
| 599 | var delivered atomic.Int32 |
| 600 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 601 | body, err := io.ReadAll(r.Body) |
| 602 | if err != nil { |
| 603 | t.Error(err) |
| 604 | return |
| 605 | } |
| 606 | for { |
| 607 | observed := maxBody.Load() |
| 608 | if int64(len(body)) <= observed || maxBody.CompareAndSwap(observed, int64(len(body))) { |
| 609 | break |
| 610 | } |
| 611 | } |
| 612 | if len(body) > eventwire.MirrorBatchMaxBytes { |
| 613 | t.Errorf("request body = %d bytes", len(body)) |
| 614 | } |
| 615 | var payload struct { |
| 616 | Frames []eventwire.Event `json:"frames"` |
| 617 | } |
| 618 | if err := json.Unmarshal(body, &payload); err != nil { |
| 619 | t.Error(err) |
| 620 | return |
| 621 | } |
| 622 | delivered.Add(int32(len(payload.Frames))) |
| 623 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 624 | })) |
| 625 | defer srv.Close() |
| 626 | m := newCLITakeoverManager(nil, nil) |
| 627 | m.binding = &cliTakeoverBinding{path: "session.jsonl", record: cliServeRecord{base: srv.URL}, client: srv.Client(), grant: cliTakeoverGrant{MirrorID: "mirror"}} |
| 628 | m.revision = 1 |
| 629 | chunk := strings.Repeat("x", 1<<20) |
| 630 | for range 10 { |
| 631 | m.Emit(event.Event{Kind: event.Text, Text: chunk}) |
| 632 | } |
| 633 | for delivered.Load() < 10 { |
| 634 | if !m.push(false) { |
| 635 | t.Fatal("manager stopped while chunking large request") |
| 636 | } |
| 637 | } |
| 638 | if maxBody.Load() > eventwire.MirrorBatchMaxBytes { |
| 639 | t.Fatalf("largest body = %d", maxBody.Load()) |
| 640 | } |
| 641 | } |
| 642 | |
| 643 | func TestCLITakeoverManagerRediscoverAfterThreeServerErrors(t *testing.T) { |
| 644 | oldServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { |
| 645 | w.WriteHeader(http.StatusServiceUnavailable) |
| 646 | })) |
| 647 | defer oldServe.Close() |
| 648 | newServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 649 | switch r.URL.Path { |
| 650 | case "/auth/token": |
| 651 | w.WriteHeader(http.StatusNoContent) |
| 652 | case "/adopt": |
| 653 | _ = json.NewEncoder(w).Encode(cliTakeoverGrant{SessionPath: "session.jsonl", MirrorID: "new", ReturnHandoffID: "return", SourceWriterID: "source", TargetWriterID: agent.SessionWriterID()}) |
| 654 | default: |
| 655 | _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false}) |
| 656 | } |
| 657 | })) |
| 658 | defer newServe.Close() |
| 659 | var discoveries atomic.Int32 |
| 660 | originalDiscover := discoverCLIServesForTakeover |
| 661 | discoverCLIServesForTakeover = func() []cliServeRecord { |
| 662 | discoveries.Add(1) |
| 663 | return []cliServeRecord{{base: newServe.URL, token: "fresh"}} |
| 664 | } |
| 665 | t.Cleanup(func() { discoverCLIServesForTakeover = originalDiscover }) |
| 666 | |
| 667 | m := newCLITakeoverManager(&takeoverRecordSink{}, nil) |
| 668 | m.binding = &cliTakeoverBinding{path: "session.jsonl", record: cliServeRecord{base: oldServe.URL}, client: oldServe.Client(), grant: cliTakeoverGrant{MirrorID: "old"}} |
| 669 | m.revision = 1 |
| 670 | m.Emit(event.Event{Kind: event.Text, Text: "retry"}) |
| 671 | for range cliTakeoverRediscoverFailures { |
| 672 | if !m.push(false) { |
| 673 | t.Fatal("manager stopped before rediscovery") |
| 674 | } |
| 675 | } |
| 676 | if discoveries.Load() != 1 { |
| 677 | t.Fatalf("discoveries = %d, want 1 after threshold", discoveries.Load()) |
| 678 | } |
| 679 | } |
| 680 | |
| 681 | func TestCLITakeoverManagerRetainsPendingReturnUntilReservationSucceeds(t *testing.T) { |
| 682 | dir := t.TempDir() |
| 683 | source := filepath.Join(dir, "source.jsonl") |
| 684 | target := filepath.Join(dir, "target.jsonl") |
| 685 | leases := control.NewSessionLeaseKeeper() |
| 686 | defer leases.Release() |
| 687 | if err := leases.Rebind(source); err != nil { |
| 688 | t.Fatal(err) |
| 689 | } |
| 690 | targetKeeper := control.NewSessionLeaseKeeper() |
| 691 | if err := targetKeeper.Rebind(target); err != nil { |
| 692 | t.Fatal(err) |
| 693 | } |
| 694 | var mirrorEnded atomic.Bool |
| 695 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 696 | if r.URL.Path == "/mirror-end" { |
| 697 | mirrorEnded.Store(true) |
| 698 | w.WriteHeader(http.StatusNoContent) |
| 699 | return |
| 700 | } |
| 701 | w.WriteHeader(http.StatusNotFound) |
| 702 | })) |
| 703 | defer srv.Close() |
| 704 | binding := &cliTakeoverBinding{ |
| 705 | path: target, record: cliServeRecord{base: srv.URL}, client: srv.Client(), |
| 706 | grant: cliTakeoverGrant{MirrorID: "mirror", SourceWriterID: "serve", ReturnHandoffID: "return"}, |
| 707 | } |
| 708 | m := newCLITakeoverManager(&takeoverRecordSink{}, leases) |
| 709 | wantErr := errors.New("injected reservation failure") |
| 710 | var calls atomic.Int32 |
| 711 | m.retirePending = func(keeper *control.SessionLeaseKeeper, writerID, handoffID string) error { |
| 712 | if calls.Add(1) == 1 { |
| 713 | return wantErr |
| 714 | } |
| 715 | return keeper.RetireDetachedForHandoff(writerID, handoffID) |
| 716 | } |
| 717 | m.pending = []*cliPendingReturn{{keeper: targetKeeper, binding: binding, nextTry: time.Now()}} |
| 718 | m.retryPendingReturns(true) |
| 719 | if mirrorEnded.Load() { |
| 720 | t.Fatal("mirror ended before the reservation succeeded") |
| 721 | } |
| 722 | if got := leases.HeldPath(); got != agent.CanonicalSessionPath(source) { |
| 723 | t.Fatalf("source keeper moved to %q", got) |
| 724 | } |
| 725 | if third, err := agent.TryAcquireSessionLease(target); !errors.Is(err, agent.ErrSessionLeaseHeld) { |
| 726 | if third != nil { |
| 727 | third.Release() |
| 728 | } |
| 729 | t.Fatalf("third writer acquired pending target: %v", err) |
| 730 | } |
| 731 | m.retryPendingReturns(true) |
| 732 | if !mirrorEnded.Load() { |
| 733 | t.Fatal("mirror did not end after reservation retry succeeded") |
| 734 | } |
| 735 | info, err := agent.LoadSessionLeaseInfo(target) |
| 736 | if err != nil { |
| 737 | t.Fatal(err) |
| 738 | } |
| 739 | if info.HandoffTo != "serve" || info.HandoffID != "return" { |
| 740 | t.Fatalf("target reservation = %+v", info) |
| 741 | } |
| 742 | } |
| 743 | |
| 744 | // TestCLISessionTakeoverCandidateRequiresAResidentServe keeps the /takeover |
| 745 | // offer honest: the command takes a session from a resident serve, so a lease |
| 746 | // held by another CLI on a machine running no serve must not be advertised as |
| 747 | // takeable. The holder's PID still need not match a discovered serve, because |
| 748 | // state-file PIDs drift across serve restarts. |
| 749 | func TestCLISessionTakeoverCandidateRequiresAResidentServe(t *testing.T) { |
| 750 | held := &agent.SessionLeaseError{Path: "/tmp/held.jsonl", Info: &agent.SessionLeaseInfo{PID: 4242, WriterID: "other-cli"}} |
| 751 | previous := discoverCLIServesForTakeover |
| 752 | t.Cleanup(func() { discoverCLIServesForTakeover = previous }) |
| 753 | |
| 754 | discoverCLIServesForTakeover = func() []cliServeRecord { return nil } |
| 755 | if cliSessionTakeoverCandidate(held) { |
| 756 | t.Fatal("/takeover was offered for a holder with no resident serve to take it from") |
| 757 | } |
| 758 | |
| 759 | discoverCLIServesForTakeover = func() []cliServeRecord { |
| 760 | return []cliServeRecord{{pid: 99, base: "http://127.0.0.1:1", token: "tok"}} |
| 761 | } |
| 762 | if !cliSessionTakeoverCandidate(held) { |
| 763 | t.Fatal("/takeover was withheld although a resident serve could hand the session over") |
| 764 | } |
| 765 | if cliSessionTakeoverCandidate(errors.New("unrelated failure")) { |
| 766 | t.Fatal("/takeover was offered for an error carrying no lease holder") |
| 767 | } |
| 768 | if cliSessionTakeoverCandidate(&agent.SessionLeaseError{Path: "/tmp/held.jsonl"}) { |
| 769 | t.Fatal("/takeover was offered for a lease error with no holder info") |
| 770 | } |
| 771 | } |
| 772 |