返回 DeepSeek-Reasonix
session_takeover_test.go
根目录 / desktop / session_takeover_test.go
1 package main
2
3 import (
4 "context"
5 "encoding/json"
6 "errors"
7 "net/http"
8 "net/http/httptest"
9 "os"
10 "path/filepath"
11 "strconv"
12 "sync"
13 "sync/atomic"
14 "testing"
15 "time"
16
17 "reasonix/internal/agent"
18 "reasonix/internal/config"
19 "reasonix/internal/control"
20 "reasonix/internal/event"
21 "reasonix/internal/eventwire"
22 "reasonix/internal/remote/bootstrap"
23 )
24
25 func TestTakeoverOwnershipEncodesOpaqueSessionPath(t *testing.T) {
26 want := `C:\Users\测试 User\sessions\a&b%20.jsonl`
27 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
28 if got := r.URL.Query().Get("session"); got != want {
29 t.Errorf("session query = %q, want %q", got, want)
30 }
31 _ = json.NewEncoder(w).Encode(SessionTakeoverView{Holder: "serve"})
32 }))
33 defer srv.Close()
34 if _, err := takeoverOwnership(context.Background(), srv.Client(), srv.URL, want); err != nil {
35 t.Fatal(err)
36 }
37 }
38
39 func TestTakeoverGrantCannotCommitAfterTabRuntimeChanges(t *testing.T) {
40 isolateDesktopUserDirs(t)
41 path := filepath.Join(t.TempDir(), "taken-over.jsonl")
42 source, err := agent.TryAcquireSessionLease(path)
43 if err != nil {
44 t.Fatal(err)
45 }
46 sourceWriter := agent.SessionWriterID()
47 var sourceReleased atomic.Bool
48 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
49 switch r.URL.Path {
50 case "/handoff":
51 if !sourceReleased.CompareAndSwap(false, true) {
52 t.Error("handoff called more than once")
53 w.WriteHeader(http.StatusConflict)
54 return
55 }
56 if err := source.ReleaseForHandoff(sourceWriter, "forward"); err != nil {
57 t.Error(err)
58 w.WriteHeader(http.StatusInternalServerError)
59 return
60 }
61 _ = json.NewEncoder(w).Encode(takeoverGrant{
62 SessionPath: path, MirrorID: "mirror", HandoffID: "forward", ReturnHandoffID: "return",
63 SourceWriterID: sourceWriter, TargetWriterID: sourceWriter,
64 })
65 case "/mirror-end":
66 returned, err := agent.TryAcquireSessionLeaseWithHandoff(path, sourceWriter, "forward")
67 if err != nil {
68 t.Error(err)
69 w.WriteHeader(http.StatusConflict)
70 return
71 }
72 returned.Release()
73 w.WriteHeader(http.StatusNoContent)
74 default:
75 w.WriteHeader(http.StatusNotFound)
76 }
77 }))
78 defer srv.Close()
79
80 app := NewApp()
81 app.ctx = context.Background()
82 tab := app.createTabEntryWithID("global", globalTabWorkspaceRoot(), "", "takeover-tab")
83 tab.SessionPath = path
84 tab.StartupErrLeaseHeld = true
85 tab.StartupErr = (&sessionLeaseBusyError{}).Error()
86 tab.sink = &tabEventSink{tabID: tab.ID, app: app}
87 app.mu.Lock()
88 app.tabs[tab.ID] = tab
89 app.tabOrder = []string{tab.ID}
90 app.activeTabID = tab.ID
91 app.mu.Unlock()
92
93 originalFind := takeoverFindTargetForTest
94 takeoverFindTargetForTest = func(context.Context, *App, string) (takeoverServeRecord, *http.Client, SessionTakeoverView, error) {
95 return takeoverServeRecord{base: srv.URL}, srv.Client(), SessionTakeoverView{Holder: "serve"}, nil
96 }
97 grantSeen := make(chan struct{})
98 runtimeChanged := make(chan struct{})
99 originalHook := takeoverAfterGrantHookForTest
100 takeoverAfterGrantHookForTest = func() {
101 close(grantSeen)
102 <-runtimeChanged
103 }
104 t.Cleanup(func() {
105 takeoverFindTargetForTest = originalFind
106 takeoverAfterGrantHookForTest = originalHook
107 })
108
109 newPath := filepath.Join(t.TempDir(), "replacement.jsonl")
110 replacement := control.New(control.Options{
111 Executor: agent.New(nil, nil, agent.NewSession("system"), agent.Options{}, event.Discard),
112 SessionPath: newPath,
113 Sink: event.Discard,
114 })
115 defer replacement.Close()
116 done := make(chan error, 1)
117 go func() { done <- app.TakeoverSession(tab.ID, "wait") }()
118 <-grantSeen
119 app.runtimeRebuildMu.Lock()
120 app.mu.Lock()
121 tab.SessionPath = newPath
122 tab.Ctrl = replacement
123 tab.StartupErrLeaseHeld = false
124 app.mu.Unlock()
125 app.runtimeRebuildMu.Unlock()
126 close(runtimeChanged)
127 if err := <-done; err == nil {
128 t.Fatal("stale takeover unexpectedly committed")
129 }
130 if tab.Ctrl != replacement || tab.currentSessionPath() != newPath {
131 t.Fatalf("replacement runtime was overwritten: ctrl=%p path=%q", tab.Ctrl, tab.currentSessionPath())
132 }
133 if got := tab.sessionLeaseRuntimeKey(); got != "" {
134 t.Fatalf("stale takeover installed lease %q", got)
135 }
136 if mirror := app.takeoverMirrorForKey(sessionRuntimeKey(path)); mirror != nil {
137 t.Fatal("stale takeover installed a mirror")
138 }
139 }
140
141 func TestDirectAdoptGrantCannotAttachAfterTabRuntimeChanges(t *testing.T) {
142 isolateDesktopUserDirs(t)
143 workspace := t.TempDir()
144 sessionDir := config.ProjectSessionDir(workspace)
145 if err := os.MkdirAll(sessionDir, 0o755); err != nil {
146 t.Fatal(err)
147 }
148 path := filepath.Join(sessionDir, "adopted.jsonl")
149 mirrorEnded := make(chan struct{}, 1)
150 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
151 switch r.URL.Path {
152 case "/auth/token":
153 w.WriteHeader(http.StatusNoContent)
154 case "/ownership":
155 _ = json.NewEncoder(w).Encode(SessionTakeoverView{Holder: "free"})
156 case "/adopt":
157 _ = json.NewEncoder(w).Encode(takeoverGrant{
158 SessionPath: path, MirrorID: "stale-mirror", ReturnHandoffID: "stale-return",
159 SourceWriterID: "serve-writer", TargetWriterID: agent.SessionWriterID(),
160 })
161 case "/mirror-end":
162 mirrorEnded <- struct{}{}
163 w.WriteHeader(http.StatusNoContent)
164 default:
165 w.WriteHeader(http.StatusNotFound)
166 }
167 }))
168 defer srv.Close()
169
170 app := NewApp()
171 app.ctx = context.Background()
172 tab := app.createTabEntryWithID("global", globalTabWorkspaceRoot(), "", "adopt-tab")
173 tab.SessionPath = path
174 oldSink := &tabEventSink{tabID: tab.ID, app: app}
175 tab.sink = oldSink
176 app.mu.Lock()
177 app.tabs[tab.ID] = tab
178 app.tabOrder = []string{tab.ID}
179 app.activeTabID = tab.ID
180 app.newSessionRuntimeLocked(tab, sessionRuntimeKey(path))
181 app.advanceSessionRuntimeEpochLocked(tab)
182 app.mu.Unlock()
183
184 originalDiscover := discoverLocalTakeoverServesForAdopt
185 discoverLocalTakeoverServesForAdopt = func() []takeoverServeRecord {
186 return []takeoverServeRecord{{base: srv.URL, token: "fresh", state: bootstrap.ServeState{Workspace: workspace}}}
187 }
188 grantSeen := make(chan struct{})
189 runtimeChanged := make(chan struct{})
190 originalHook := takeoverAfterAdoptGrantHookForTest
191 takeoverAfterAdoptGrantHookForTest = func() {
192 close(grantSeen)
193 <-runtimeChanged
194 }
195 t.Cleanup(func() {
196 discoverLocalTakeoverServesForAdopt = originalDiscover
197 takeoverAfterAdoptGrantHookForTest = originalHook
198 })
199
200 done := make(chan bool, 1)
201 go func() { done <- app.adoptSessionFromLocalServeOnce(tab.ID, path) }()
202 <-grantSeen
203 newPath := filepath.Join(sessionDir, "replacement.jsonl")
204 newSink := &tabEventSink{tabID: tab.ID, app: app}
205 app.runtimeRebuildMu.Lock()
206 app.mu.Lock()
207 tab.SessionPath = newPath
208 tab.sink = newSink
209 app.advanceSessionRuntimeEpochLocked(tab)
210 app.mu.Unlock()
211 app.runtimeRebuildMu.Unlock()
212 close(runtimeChanged)
213 if !<-done {
214 t.Fatal("stale adoption did not reach a serve")
215 }
216 if mirror := app.takeoverMirrorForKey(sessionRuntimeKey(path)); mirror != nil {
217 t.Fatal("stale adoption installed a mirror")
218 }
219 if oldSink.takeoverMirror.Load() != nil || newSink.takeoverMirror.Load() != nil {
220 t.Fatal("stale adoption attached to an old or replacement event sink")
221 }
222 select {
223 case <-mirrorEnded:
224 case <-time.After(3 * time.Second):
225 t.Fatal("stale mirror generation was not ended")
226 }
227 }
228
229 func TestTakeoverMirrorRetainsFailedReturnUntilReservationSucceeds(t *testing.T) {
230 path := filepath.Join(t.TempDir(), "pending-return.jsonl")
231 lease, err := agent.TryAcquireSessionLease(path)
232 if err != nil {
233 t.Fatal(err)
234 }
235 var mirrorEnded atomic.Bool
236 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
237 if r.URL.Path == "/mirror-end" {
238 mirrorEnded.Store(true)
239 w.WriteHeader(http.StatusNoContent)
240 return
241 }
242 w.WriteHeader(http.StatusNotFound)
243 }))
244 defer srv.Close()
245 m := &takeoverMirror{
246 sessionPath: path,
247 client: srv.Client(),
248 record: takeoverServeRecord{base: srv.URL},
249 grant: takeoverGrant{MirrorID: "mirror", SourceWriterID: "serve", ReturnHandoffID: "return"},
250 wake: make(chan struct{}, 1),
251 }
252 wantErr := errors.New("injected reservation write failure")
253 var calls atomic.Int32
254 m.releaseHandoff = func(candidate *agent.SessionLease, writerID, handoffID string) error {
255 if calls.Add(1) == 1 {
256 return wantErr
257 }
258 return candidate.ReleaseForHandoff(writerID, handoffID)
259 }
260 m.holdPendingReturn(lease)
261 if m.retryPendingReturn(true) {
262 t.Fatal("first retry unexpectedly succeeded")
263 }
264 if mirrorEnded.Load() {
265 t.Fatal("mirror ended before the reverse reservation existed")
266 }
267 if third, err := agent.TryAcquireSessionLease(path); !errors.Is(err, agent.ErrSessionLeaseHeld) {
268 if third != nil {
269 third.Release()
270 }
271 t.Fatalf("third writer acquired pending lease: %v", err)
272 }
273 if !m.retryPendingReturn(true) {
274 t.Fatal("second retry did not publish the reservation")
275 }
276 m.mirrorEnd()
277 if !mirrorEnded.Load() {
278 t.Fatal("mirror-end was not sent after reservation success")
279 }
280 info, err := agent.LoadSessionLeaseInfo(path)
281 if err != nil {
282 t.Fatal(err)
283 }
284 if info.HandoffTo != "serve" || info.HandoffID != "return" {
285 t.Fatalf("return reservation = %+v", info)
286 }
287 }
288
289 func TestTakeoverMirrorRetriesSameDrainedFramesInOrder(t *testing.T) {
290 var mu sync.Mutex
291 var batches [][]eventwire.Event
292 var calls atomic.Int32
293 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
294 var body struct {
295 Frames []eventwire.Event `json:"frames"`
296 }
297 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
298 t.Error(err)
299 w.WriteHeader(http.StatusBadRequest)
300 return
301 }
302 mu.Lock()
303 batches = append(batches, append([]eventwire.Event(nil), body.Frames...))
304 mu.Unlock()
305 if calls.Add(1) == 1 {
306 w.WriteHeader(http.StatusServiceUnavailable)
307 return
308 }
309 _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false})
310 }))
311 defer srv.Close()
312
313 m := &takeoverMirror{
314 sessionPath: "session.jsonl", client: srv.Client(),
315 record: takeoverServeRecord{base: srv.URL}, grant: takeoverGrant{MirrorID: "mirror-1"},
316 wake: make(chan struct{}, 1),
317 }
318 m.forwardEvent(event.Event{Kind: event.Text, Text: "first"})
319 m.forwardEvent(event.Event{Kind: event.Text, Text: "second"})
320 if !m.pushOnce(false) {
321 t.Fatal("forwarder stopped unexpectedly")
322 }
323 if !m.pushOnce(false) {
324 t.Fatal("forwarder stopped unexpectedly")
325 }
326
327 mu.Lock()
328 defer mu.Unlock()
329 if len(batches) != 2 || len(batches[0]) != 2 || len(batches[1]) != 2 {
330 t.Fatalf("batches = %+v, want two complete batches", batches)
331 }
332 for i := range batches[0] {
333 if batches[0][i].Text != batches[1][i].Text {
334 t.Fatalf("retry reordered frame %d: %q != %q", i, batches[0][i].Text, batches[1][i].Text)
335 }
336 }
337 }
338
339 func TestTakeoverMirrorReadoptsAfterServeMoves(t *testing.T) {
340 isolateDesktopUserDirs(t)
341 workspace := t.TempDir()
342 sessionDir := config.ProjectSessionDir(workspace)
343 if err := os.MkdirAll(sessionDir, 0o755); err != nil {
344 t.Fatal(err)
345 }
346 path := filepath.Join(sessionDir, "session.jsonl")
347 deadURL := "http://unreachable.invalid"
348 deadClient := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
349 return nil, errors.New("injected stale serve connection")
350 })}
351 var delivered atomic.Bool
352 newServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
353 switch r.URL.Path {
354 case "/auth/token":
355 w.WriteHeader(http.StatusNoContent)
356 case "/adopt":
357 _ = json.NewEncoder(w).Encode(takeoverGrant{
358 SessionPath: path, MirrorID: "new", ReturnHandoffID: "return", SourceWriterID: "source", TargetWriterID: agent.SessionWriterID(),
359 })
360 case "/external/frames":
361 delivered.Store(true)
362 _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false})
363 default:
364 w.WriteHeader(http.StatusNotFound)
365 }
366 }))
367 defer newServe.Close()
368 originalDiscover := discoverLocalTakeoverServesForMirror
369 discoverLocalTakeoverServesForMirror = func() []takeoverServeRecord {
370 return []takeoverServeRecord{{base: newServe.URL, token: "fresh", state: bootstrap.ServeState{Workspace: workspace}}}
371 }
372 t.Cleanup(func() { discoverLocalTakeoverServesForMirror = originalDiscover })
373 m := &takeoverMirror{
374 sessionPath: path,
375 client: deadClient,
376 record: takeoverServeRecord{base: deadURL},
377 grant: takeoverGrant{MirrorID: "old"},
378 wake: make(chan struct{}, 1),
379 }
380 m.bindingRevision = 1
381 m.forwardEvent(event.Event{Kind: event.Text, Text: "recover"})
382 if !m.pushOnce(false) {
383 t.Fatal("desktop mirror stopped during re-adoption")
384 }
385 if !m.pushOnce(false) {
386 t.Fatal("desktop mirror stopped after re-adoption")
387 }
388 if !delivered.Load() {
389 t.Fatal("desktop mirror did not re-adopt and deliver through the new serve")
390 }
391 _, record, _, grant, revision := m.snapshotBinding()
392 if record.base != newServe.URL || grant.MirrorID != "new" || revision <= 1 {
393 t.Fatalf("re-adopted binding = base %q grant %+v revision %d", record.base, grant, revision)
394 }
395 }
396
397 func TestTakeoverMirrorDemotionReturnsRefreshedGeneration(t *testing.T) {
398 isolateDesktopUserDirs(t)
399 workspace := t.TempDir()
400 sessionDir := config.ProjectSessionDir(workspace)
401 if err := os.MkdirAll(sessionDir, 0o755); err != nil {
402 t.Fatal(err)
403 }
404 path := filepath.Join(sessionDir, "demote.jsonl")
405 oldStarted := make(chan struct{})
406 oldRelease := make(chan struct{})
407 var oldEnds atomic.Int32
408 oldServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
409 switch r.URL.Path {
410 case "/external/frames":
411 close(oldStarted)
412 <-oldRelease
413 w.WriteHeader(http.StatusUnauthorized)
414 case "/mirror-end":
415 oldEnds.Add(1)
416 w.WriteHeader(http.StatusNoContent)
417 default:
418 w.WriteHeader(http.StatusNotFound)
419 }
420 }))
421 defer oldServe.Close()
422 newEnds := make(chan string, 1)
423 newServe := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
424 switch r.URL.Path {
425 case "/auth/token":
426 w.WriteHeader(http.StatusNoContent)
427 case "/adopt":
428 _ = json.NewEncoder(w).Encode(takeoverGrant{
429 SessionPath: path, MirrorID: "mirror-new", ReturnHandoffID: "return-new",
430 SourceWriterID: "serve-new", TargetWriterID: agent.SessionWriterID(),
431 })
432 case "/mirror-end":
433 var body struct {
434 MirrorID string `json:"mirrorId"`
435 }
436 _ = json.NewDecoder(r.Body).Decode(&body)
437 newEnds <- body.MirrorID
438 w.WriteHeader(http.StatusNoContent)
439 default:
440 w.WriteHeader(http.StatusNotFound)
441 }
442 }))
443 defer newServe.Close()
444 originalDiscover := discoverLocalTakeoverServesForMirror
445 discoverLocalTakeoverServesForMirror = func() []takeoverServeRecord {
446 return []takeoverServeRecord{{base: newServe.URL, token: "fresh", state: bootstrap.ServeState{Workspace: workspace}}}
447 }
448 t.Cleanup(func() { discoverLocalTakeoverServesForMirror = originalDiscover })
449
450 lease, err := agent.TryAcquireSessionLease(path)
451 if err != nil {
452 t.Fatal(err)
453 }
454 tab := &WorkspaceTab{ID: "demote-tab"}
455 tab.adoptSessionLease(lease)
456 app := NewApp()
457 m := newTakeoverMirror(app, sessionRuntimeKey(path), tab.ID, path, nil,
458 takeoverServeRecord{base: oldServe.URL}, oldServe.Client(),
459 takeoverGrant{SessionPath: path, MirrorID: "mirror-old", ReturnHandoffID: "return-old", SourceWriterID: "serve-old", TargetWriterID: agent.SessionWriterID()},
460 )
461 m.forwardEvent(event.Event{Kind: event.Text, Text: "generation fence"})
462 pushDone := make(chan bool, 1)
463 go func() { pushDone <- m.pushOnce(false) }()
464 <-oldStarted
465 returnDone := make(chan error, 1)
466 go func() { returnDone <- m.returnLeaseForDemotion(tab) }()
467 close(oldRelease)
468 if !<-pushDone {
469 t.Fatal("forwarder stopped while re-adopting")
470 }
471 if err := <-returnDone; err != nil {
472 t.Fatal(err)
473 }
474 info, err := agent.LoadSessionLeaseInfo(path)
475 if err != nil {
476 t.Fatal(err)
477 }
478 if info == nil || info.HandoffTo != "serve-new" || info.HandoffID != "return-new" {
479 t.Fatalf("reverse reservation = %+v", info)
480 }
481 select {
482 case mirrorID := <-newEnds:
483 if mirrorID != "mirror-new" {
484 t.Fatalf("mirror-end id = %q, want refreshed generation", mirrorID)
485 }
486 case <-time.After(3 * time.Second):
487 t.Fatal("refreshed mirror generation was not ended")
488 }
489 if oldEnds.Load() != 0 || tab.sessionLeaseRuntimeKey() != "" || !m.returned.Load() {
490 t.Fatalf("old ends=%d tab lease=%q returned=%v", oldEnds.Load(), tab.sessionLeaseRuntimeKey(), m.returned.Load())
491 }
492 }
493
494 func TestTakeoverMirrorChunksWithoutDroppingFrames(t *testing.T) {
495 var got []string
496 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
497 var body struct {
498 Frames []eventwire.Event `json:"frames"`
499 }
500 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
501 t.Fatal(err)
502 }
503 if len(body.Frames) > takeoverMirrorMaxQueue {
504 t.Fatalf("batch size = %d", len(body.Frames))
505 }
506 for _, frame := range body.Frames {
507 got = append(got, frame.Text)
508 }
509 _ = json.NewEncoder(w).Encode(map[string]bool{"reclaimRequested": false})
510 }))
511 defer srv.Close()
512 m := &takeoverMirror{
513 sessionPath: "session.jsonl", client: srv.Client(), record: takeoverServeRecord{base: srv.URL},
514 grant: takeoverGrant{MirrorID: "mirror-1"}, wake: make(chan struct{}, 1),
515 }
516 for i := range takeoverMirrorMaxQueue + 37 {
517 m.forwardEvent(event.Event{Kind: event.Text, Text: strconv.Itoa(i)})
518 }
519 if !m.pushOnce(false) {
520 t.Fatal("chunked frame push stopped unexpectedly")
521 }
522 if !m.pushOnce(false) {
523 t.Fatal("chunked frame push stopped unexpectedly")
524 }
525 if len(got) != takeoverMirrorMaxQueue+37 {
526 t.Fatalf("received %d frames", len(got))
527 }
528 for i, text := range got {
529 if text != strconv.Itoa(i) {
530 t.Fatalf("frame %d = %q", i, text)
531 }
532 }
533 }
534
535 func TestRebindWithTakeoverMirrorDoesNotReenterAppLock(t *testing.T) {
536 app, tab, _, _, targetPath, loaded := newAtomicRebindTestApp(t)
537 key := sessionRuntimeKey(targetPath)
538 app.takeoverMu.Lock()
539 if app.takeoverMirrors == nil {
540 app.takeoverMirrors = map[string]*takeoverMirror{}
541 }
542 mirror := &takeoverMirror{app: app, key: key, sessionPath: targetPath}
543 app.takeoverMirrors[key] = mirror
544 app.takeoverMu.Unlock()
545
546 // Rebinding includes controller boot and disk migration, not just locking.
547 // Use the suite timeout for deadlocks; a local timer can fail on slow I/O
548 // and run fixture cleanup while the rebind goroutine still uses its files.
549 if err := app.rebindTabToLoadedSessionPath(tab, targetPath, loaded); err != nil {
550 t.Fatal(err)
551 }
552 if tab.sink.takeoverMirror.Load() != mirror {
553 t.Fatal("replacement sink did not reconnect the takeover mirror")
554 }
555 mirror.mu.Lock()
556 defer mirror.mu.Unlock()
557 if mirror.sink != tab.sink || mirror.tabID != tab.ID {
558 t.Fatal("takeover mirror retained the retired sink or tab binding")
559 }
560 }
561
561 lines GO