From 021e5d7077a9f12bc9b9e62310a61576f77e23ad Mon Sep 17 00:00:00 2001
From: Aric Camarata <aric.camarata@gmail.com>
Date: Sun, 25 Oct 2026 12:02:23 -0400
Subject: [PATCH] test(ci): make the lost-session and coordinator-restart
 scenarios deterministic

TestForegroundRemoteLostSession: the fixture's drop cancelled the agent
before closing the pipes, so the agent's cancelled result could cross the
wire first and the run (correctly) read it as an operator cancel. Close the
wire first, then stop the agent, as a dropped ssh connection does.

TestPullCoordinatorKillResumesWithoutDuplicates: it asserted the first
token moved to the new session, but after the restart the re-adopted run
offers the lease again with a token of its own and the carrier may switch
to it before or after its reconcile heartbeat. Assert that a token of the
lease is bound to the new session of the same certificate instead.
---
 .../transport/agent/scenarios_test.go         | 26 ++++++++++++++++---
 free/ci/test/sim/foreground_remote_test.go    |  8 ++++--
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/free/ci/internal/transport/agent/scenarios_test.go b/free/ci/internal/transport/agent/scenarios_test.go
index 8a73a35a5..354c06d71 100644
--- a/free/ci/internal/transport/agent/scenarios_test.go
+++ b/free/ci/internal/transport/agent/scenarios_test.go
@@ -177,9 +177,29 @@ func scenarioRestart(t *testing.T) {
 			t.Fatalf("line-%d reached the runs %d times: old %v new %v", i, n, old.lines(), e.sink.lines())
 		}
 	}
-	after, _, _ := e.st.TokenCheck(context.Background(), cl.JTI)
-	if after.Session == before.Session || after.Serial != before.Serial {
-		t.Fatalf("the token must move to the new session of the same certificate: before %+v after %+v", before, after)
+	// Which token carried the lines after the restart is a race of two streams: the re-adopted run offers the lease
+	// again with a token of its own (the carrier then swaps to it), and the agent's reconcile heartbeat may be
+	// answered before or after that offer. In one order the old token moves to the new session and the new one is
+	// never presented; in the other the old one stays on the dead session and the new one is bound fresh. Either
+	// way the token that carried them must belong to the new session of the same certificate.
+	var after store.TokenBinding
+	eventually(t, 10*time.Second, "a token of the lease to be bound to the new session", func() bool {
+		jtis := []string{cl.JTI}
+		if _, held := c.tokenSnapshot(); held[run.l.LeaseID] != "" {
+			if tc, verr := e.tokens.Verify(held[run.l.LeaseID]); verr == nil && tc.JTI != cl.JTI {
+				jtis = append(jtis, tc.JTI)
+			}
+		}
+		for _, jti := range jtis {
+			if b, ok, _ := e.st.TokenCheck(context.Background(), jti); ok && b.Session != before.Session {
+				after = b
+				return true
+			}
+		}
+		return false
+	})
+	if after.Serial != before.Serial {
+		t.Fatalf("the token must stay on the same certificate: before %+v after %+v", before, after)
 	}
 	b := readLog(t, LogFile(e.drv.LogDir, run.l.LeaseID))
 	if strings.Count(b, "\n") != 16 {
diff --git a/free/ci/test/sim/foreground_remote_test.go b/free/ci/test/sim/foreground_remote_test.go
index cb6157344..6438e1db0 100644
--- a/free/ci/test/sim/foreground_remote_test.go
+++ b/free/ci/test/sim/foreground_remote_test.go
@@ -126,10 +126,14 @@ func starter(t *testing.T, nodes []*node) ssh.Starter {
 			_ = errW.Close()
 			done <- e
 		}()
+		// kill is a dropped ssh connection: the wire goes first, so nothing the agent still says reaches the
+		// coordinator, and only then does the agent notice and stop its job. Cancelling the agent first lets its
+		// "cancelled" result cross the pipe before it closes, which the run rightly reads as an operator cancel
+		// and not as a lost session.
 		kill := func() error {
-			cancel()
-			_ = toAgentR.Close()
 			_ = fromAgentW.Close()
+			_ = toAgentR.Close()
+			cancel()
 			return nil
 		}
 		n.mu.Lock()