feat(agent): Add shutdown lifecycle states and shutdown_script support (#6139)

* feat(api): Add agent shutdown lifecycle states

* feat(agent): Add shutdown_script support

* feat(agent): Add shutdown_script timeout

* feat(site): Support new agent lifecycle states

---

Co-authored-by: Marcin Tojek <marcin@coder.com>

agent/agent.go

@@ -121,6 +121,7 @@ func New(options Options) io.Closer {
 		logDir:                 options.LogDir,
 		tempDir:                options.TempDir,
 		lifecycleUpdate:        make(chan struct{}, 1),
+		lifecycleReported:      make(chan codersdk.WorkspaceAgentLifecycle, 1),
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
 	}
 	a.init(ctx)
@@ -149,9 +150,10 @@ type agent struct {
 	sessionToken atomic.Pointer[string]
 	sshServer    *ssh.Server
 
-	lifecycleUpdate chan struct{}
-	lifecycleMu     sync.Mutex // Protects following.
-	lifecycleState  codersdk.WorkspaceAgentLifecycle
+	lifecycleUpdate   chan struct{}
+	lifecycleReported chan codersdk.WorkspaceAgentLifecycle
+	lifecycleMu       sync.RWMutex // Protects following.
+	lifecycleState    codersdk.WorkspaceAgentLifecycle
 
 	network       *tailnet.Conn
 	connStatsChan chan *agentsdk.Stats
@@ -207,9 +209,9 @@ func (a *agent) reportLifecycleLoop(ctx context.Context) {
 		}
 
 		for r := retry.New(time.Second, 15*time.Second); r.Wait(ctx); {
-			a.lifecycleMu.Lock()
+			a.lifecycleMu.RLock()
 			state := a.lifecycleState
-			a.lifecycleMu.Unlock()
+			a.lifecycleMu.RUnlock()
 
 			if state == lastReported {
 				break
@@ -222,6 +224,11 @@ func (a *agent) reportLifecycleLoop(ctx context.Context) {
 			})
 			if err == nil {
 				lastReported = state
+				select {
+				case a.lifecycleReported <- state:
+				case <-a.lifecycleReported:
+					a.lifecycleReported <- state
+				}
 				break
 			}
 			if xerrors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded) {
@@ -233,13 +240,20 @@ func (a *agent) reportLifecycleLoop(ctx context.Context) {
 	}
 }
 
+// setLifecycle sets the lifecycle state and notifies the lifecycle loop.
+// The state is only updated if it's a valid state transition.
 func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentLifecycle) {
 	a.lifecycleMu.Lock()
-	defer a.lifecycleMu.Unlock()
-
-	a.logger.Debug(ctx, "set lifecycle state", slog.F("state", state), slog.F("previous", a.lifecycleState))
-
+	lastState := a.lifecycleState
+	if slices.Index(codersdk.WorkspaceAgentLifecycleOrder, lastState) > slices.Index(codersdk.WorkspaceAgentLifecycleOrder, state) {
+		a.logger.Warn(ctx, "attempted to set lifecycle state to a previous state", slog.F("last", lastState), slog.F("state", state))
+		a.lifecycleMu.Unlock()
+		return
+	}
 	a.lifecycleState = state
+	a.logger.Debug(ctx, "set lifecycle state", slog.F("state", state), slog.F("last", lastState))
+	a.lifecycleMu.Unlock()
+
 	select {
 	case a.lifecycleUpdate <- struct{}{}:
 	default:
@@ -299,9 +313,10 @@ func (a *agent) run(ctx context.Context) error {
 			}
 		}
 
+		lifecycleState := codersdk.WorkspaceAgentLifecycleReady
 		scriptDone := make(chan error, 1)
 		scriptStart := time.Now()
-		err := a.trackConnGoroutine(func() {
+		err = a.trackConnGoroutine(func() {
 			defer close(scriptDone)
 			scriptDone <- a.runStartupScript(ctx, metadata.StartupScript)
 		})
@@ -329,16 +344,17 @@ func (a *agent) run(ctx context.Context) error {
 			if errors.Is(err, context.Canceled) {
 				return
 			}
-			execTime := time.Since(scriptStart)
-			lifecycleStatus := codersdk.WorkspaceAgentLifecycleReady
-			if err != nil {
-				a.logger.Warn(ctx, "startup script failed", slog.F("execution_time", execTime), slog.Error(err))
-				lifecycleStatus = codersdk.WorkspaceAgentLifecycleStartError
-			} else {
-				a.logger.Info(ctx, "startup script completed", slog.F("execution_time", execTime))
+			// Only log if there was a startup script.
+			if metadata.StartupScript != "" {
+				execTime := time.Since(scriptStart)
+				if err != nil {
+					a.logger.Warn(ctx, "startup script failed", slog.F("execution_time", execTime), slog.Error(err))
+					lifecycleState = codersdk.WorkspaceAgentLifecycleStartError
+				} else {
+					a.logger.Info(ctx, "startup script completed", slog.F("execution_time", execTime))
+				}
 			}
-
-			a.setLifecycle(ctx, lifecycleStatus)
+			a.setLifecycle(ctx, lifecycleState)
 		}()
 	}
 
@@ -606,14 +622,22 @@ func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error
 }
 
 func (a *agent) runStartupScript(ctx context.Context, script string) error {
+	return a.runScript(ctx, "startup", script)
+}
+
+func (a *agent) runShutdownScript(ctx context.Context, script string) error {
+	return a.runScript(ctx, "shutdown", script)
+}
+
+func (a *agent) runScript(ctx context.Context, lifecycle, script string) error {
 	if script == "" {
 		return nil
 	}
 
-	a.logger.Info(ctx, "running startup script", slog.F("script", script))
-	writer, err := a.filesystem.OpenFile(filepath.Join(a.logDir, "coder-startup-script.log"), os.O_CREATE|os.O_RDWR, 0o600)
+	a.logger.Info(ctx, "running script", slog.F("lifecycle", lifecycle), slog.F("script", script))
+	writer, err := a.filesystem.OpenFile(filepath.Join(a.logDir, fmt.Sprintf("coder-%s-script.log", lifecycle)), os.O_CREATE|os.O_RDWR, 0o600)
 	if err != nil {
-		return xerrors.Errorf("open startup script log file: %w", err)
+		return xerrors.Errorf("open %s script log file: %w", lifecycle, err)
 	}
 	defer func() {
 		_ = writer.Close()
@@ -774,7 +798,7 @@ func (a *agent) createCommand(ctx context.Context, rawCommand string, env []stri
 
 	rawMetadata := a.metadata.Load()
 	if rawMetadata == nil {
-		return nil, xerrors.Errorf("no metadata was provided: %w", err)
+		return nil, xerrors.Errorf("no metadata was provided")
 	}
 	metadata, valid := rawMetadata.(agentsdk.Metadata)
 	if !valid {
@@ -1290,13 +1314,73 @@ func (a *agent) Close() error {
 	if a.isClosed() {
 		return nil
 	}
+
+	ctx := context.Background()
+	a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleShuttingDown)
+
+	lifecycleState := codersdk.WorkspaceAgentLifecycleOff
+	if metadata, ok := a.metadata.Load().(agentsdk.Metadata); ok && metadata.ShutdownScript != "" {
+		scriptDone := make(chan error, 1)
+		scriptStart := time.Now()
+		go func() {
+			defer close(scriptDone)
+			scriptDone <- a.runShutdownScript(ctx, metadata.ShutdownScript)
+		}()
+
+		var timeout <-chan time.Time
+		// If timeout is zero, an older version of the coder
+		// provider was used. Otherwise a timeout is always > 0.
+		if metadata.ShutdownScriptTimeout > 0 {
+			t := time.NewTimer(metadata.ShutdownScriptTimeout)
+			defer t.Stop()
+			timeout = t.C
+		}
+
+		var err error
+		select {
+		case err = <-scriptDone:
+		case <-timeout:
+			a.logger.Warn(ctx, "shutdown script timed out")
+			a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleShutdownTimeout)
+			err = <-scriptDone // The script can still complete after a timeout.
+		}
+		execTime := time.Since(scriptStart)
+		if err != nil {
+			a.logger.Warn(ctx, "shutdown script failed", slog.F("execution_time", execTime), slog.Error(err))
+			lifecycleState = codersdk.WorkspaceAgentLifecycleShutdownError
+		} else {
+			a.logger.Info(ctx, "shutdown script completed", slog.F("execution_time", execTime))
+		}
+	}
+
+	// Set final state and wait for it to be reported because context
+	// cancellation will stop the report loop.
+	a.setLifecycle(ctx, lifecycleState)
+
+	// Wait for the lifecycle to be reported, but don't wait forever so
+	// that we don't break user expectations.
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+lifecycleWaitLoop:
+	for {
+		select {
+		case <-ctx.Done():
+			break lifecycleWaitLoop
+		case s := <-a.lifecycleReported:
+			if s == lifecycleState {
+				break lifecycleWaitLoop
+			}
+		}
+	}
+
 	close(a.closed)
 	a.closeCancel()
+	_ = a.sshServer.Close()
 	if a.network != nil {
 		_ = a.network.Close()
 	}
-	_ = a.sshServer.Close()
 	a.connCloseWait.Wait()
+
 	return nil
 }
 

agent/agent_test.go

@@ -59,7 +59,8 @@ func TestAgent_Stats_SSH(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, _, stats, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	//nolint:dogsled
+	conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
@@ -85,7 +86,8 @@ func TestAgent_Stats_ReconnectingPTY(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, _, stats, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	//nolint:dogsled
+	conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 
 	ptyConn, err := conn.ReconnectingPTY(ctx, uuid.New(), 128, 128, "/bin/bash")
 	require.NoError(t, err)
@@ -114,7 +116,8 @@ func TestAgent_Stats_Magic(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 
-	conn, _, stats, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	//nolint:dogsled
+	conn, _, stats, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -572,7 +575,7 @@ func TestAgent_SFTP(t *testing.T) {
 		home = "/" + strings.ReplaceAll(home, "\\", "/")
 	}
 	//nolint:dogsled
-	conn, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -604,7 +607,7 @@ func TestAgent_SCP(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	defer sshClient.Close()
@@ -709,7 +712,7 @@ func TestAgent_StartupScript(t *testing.T) {
 	}
 	content := "output"
 	//nolint:dogsled
-	_, _, _, fs := setupAgent(t, agentsdk.Metadata{
+	_, _, _, fs, _ := setupAgent(t, agentsdk.Metadata{
 		StartupScript: "echo " + content,
 	}, 0)
 	var gotContent string
@@ -740,10 +743,10 @@ func TestAgent_StartupScript(t *testing.T) {
 func TestAgent_Lifecycle(t *testing.T) {
 	t.Parallel()
 
-	t.Run("Timeout", func(t *testing.T) {
+	t.Run("StartTimeout", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "sleep 5",
 			StartupScriptTimeout: time.Nanosecond,
 		}, 0)
@@ -769,10 +772,10 @@ func TestAgent_Lifecycle(t *testing.T) {
 		}
 	})
 
-	t.Run("Error", func(t *testing.T) {
+	t.Run("StartError", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "false",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -801,7 +804,7 @@ func TestAgent_Lifecycle(t *testing.T) {
 	t.Run("Ready", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 		}, 0)
@@ -826,6 +829,191 @@ func TestAgent_Lifecycle(t *testing.T) {
 			require.Equal(t, want, got)
 		}
 	})
+
+	t.Run("ShuttingDown", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+			ShutdownScript:       "sleep 5",
+			StartupScriptTimeout: 30 * time.Second,
+		}, 0)
+
+		var ready []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			ready = client.getLifecycleStates()
+			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		// Start close asynchronously so that we an inspect the state.
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			err := closer.Close()
+			assert.NoError(t, err)
+		}()
+		t.Cleanup(func() {
+			<-done
+		})
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleShuttingDown,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()[len(ready):]
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		require.Equal(t, want, got)
+	})
+
+	t.Run("ShutdownTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+			ShutdownScript:        "sleep 5",
+			ShutdownScriptTimeout: time.Nanosecond,
+		}, 0)
+
+		var ready []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			ready = client.getLifecycleStates()
+			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		// Start close asynchronously so that we an inspect the state.
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			err := closer.Close()
+			assert.NoError(t, err)
+		}()
+		t.Cleanup(func() {
+			<-done
+		})
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleShuttingDown,
+			codersdk.WorkspaceAgentLifecycleShutdownTimeout,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()[len(ready):]
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		switch len(got) {
+		case 1:
+			// This can happen if lifecycle state updates are
+			// too fast, only the latest one is reported.
+			require.Equal(t, want[1:], got)
+		default:
+			// This is the expected case.
+			require.Equal(t, want, got)
+		}
+	})
+
+	t.Run("ShutdownError", func(t *testing.T) {
+		t.Parallel()
+
+		_, client, _, _, closer := setupAgent(t, agentsdk.Metadata{
+			ShutdownScript:        "false",
+			ShutdownScriptTimeout: 30 * time.Second,
+		}, 0)
+
+		var ready []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			ready = client.getLifecycleStates()
+			return len(ready) > 0 && ready[len(ready)-1] == codersdk.WorkspaceAgentLifecycleReady
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		// Start close asynchronously so that we an inspect the state.
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			err := closer.Close()
+			assert.NoError(t, err)
+		}()
+		t.Cleanup(func() {
+			<-done
+		})
+
+		want := []codersdk.WorkspaceAgentLifecycle{
+			codersdk.WorkspaceAgentLifecycleShuttingDown,
+			codersdk.WorkspaceAgentLifecycleShutdownError,
+		}
+
+		var got []codersdk.WorkspaceAgentLifecycle
+		assert.Eventually(t, func() bool {
+			got = client.getLifecycleStates()[len(ready):]
+			return len(got) > 0 && got[len(got)-1] == want[len(want)-1]
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		switch len(got) {
+		case 1:
+			// This can happen if lifecycle state updates are
+			// too fast, only the latest one is reported.
+			require.Equal(t, want[1:], got)
+		default:
+			// This is the expected case.
+			require.Equal(t, want, got)
+		}
+	})
+
+	t.Run("ShutdownScriptOnce", func(t *testing.T) {
+		t.Parallel()
+
+		expected := "this-is-shutdown"
+		client := &client{
+			t:       t,
+			agentID: uuid.New(),
+			metadata: agentsdk.Metadata{
+				DERPMap:        tailnettest.RunDERPAndSTUN(t),
+				StartupScript:  "echo 1",
+				ShutdownScript: "echo " + expected,
+			},
+			statsChan:   make(chan *agentsdk.Stats),
+			coordinator: tailnet.NewCoordinator(),
+		}
+
+		fs := afero.NewMemMapFs()
+		agent := agent.New(agent.Options{
+			Client:     client,
+			Logger:     slogtest.Make(t, nil).Leveled(slog.LevelInfo),
+			Filesystem: fs,
+		})
+
+		// agent.Close() loads the shutdown script from the agent metadata.
+		// The metadata is populated just before execution of the startup script, so it's mandatory to wait
+		// until the startup starts.
+		require.Eventually(t, func() bool {
+			outputPath := filepath.Join(os.TempDir(), "coder-startup-script.log")
+			content, err := afero.ReadFile(fs, outputPath)
+			if err != nil {
+				t.Logf("read file %q: %s", outputPath, err)
+				return false
+			}
+			return len(content) > 0 // something is in the startup log file
+		}, testutil.WaitShort, testutil.IntervalMedium)
+
+		err := agent.Close()
+		require.NoError(t, err, "agent should be closed successfully")
+
+		outputPath := filepath.Join(os.TempDir(), "coder-shutdown-script.log")
+		logFirstRead, err := afero.ReadFile(fs, outputPath)
+		require.NoError(t, err, "log file should be present")
+		require.Equal(t, expected, string(bytes.TrimSpace(logFirstRead)))
+
+		// Make sure that script can't be executed twice.
+		err = agent.Close()
+		require.NoError(t, err, "don't need to close the agent twice, no effect")
+
+		logSecondRead, err := afero.ReadFile(fs, outputPath)
+		require.NoError(t, err, "log file should be present")
+		require.Equal(t, string(bytes.TrimSpace(logFirstRead)), string(bytes.TrimSpace(logSecondRead)))
+	})
 }
 
 func TestAgent_Startup(t *testing.T) {
@@ -834,7 +1022,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("EmptyDirectory", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "",
@@ -848,7 +1036,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("HomeDirectory", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "~",
@@ -864,7 +1052,7 @@ func TestAgent_Startup(t *testing.T) {
 	t.Run("HomeEnvironmentVariable", func(t *testing.T) {
 		t.Parallel()
 
-		_, client, _, _ := setupAgent(t, agentsdk.Metadata{
+		_, client, _, _, _ := setupAgent(t, agentsdk.Metadata{
 			StartupScript:        "true",
 			StartupScriptTimeout: 30 * time.Second,
 			Directory:            "$HOME",
@@ -891,7 +1079,7 @@ func TestAgent_ReconnectingPTY(t *testing.T) {
 	defer cancel()
 
 	//nolint:dogsled
-	conn, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 	id := uuid.New()
 	netConn, err := conn.ReconnectingPTY(ctx, id, 100, 100, "/bin/bash")
 	require.NoError(t, err)
@@ -993,7 +1181,7 @@ func TestAgent_Dial(t *testing.T) {
 			}()
 
 			//nolint:dogsled
-			conn, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+			conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 			require.True(t, conn.AwaitReachable(context.Background()))
 			conn1, err := conn.DialContext(context.Background(), l.Addr().Network(), l.Addr().String())
 			require.NoError(t, err)
@@ -1015,7 +1203,7 @@ func TestAgent_Speedtest(t *testing.T) {
 	defer cancel()
 	derpMap := tailnettest.RunDERPAndSTUN(t)
 	//nolint:dogsled
-	conn, _, _, _ := setupAgent(t, agentsdk.Metadata{
+	conn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{
 		DERPMap: derpMap,
 	}, 0)
 	defer conn.Close()
@@ -1101,7 +1289,7 @@ func TestAgent_WriteVSCodeConfigs(t *testing.T) {
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
 	//nolint:dogsled
-	agentConn, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
+	agentConn, _, _, _, _ := setupAgent(t, agentsdk.Metadata{}, 0)
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
 	waitGroup := sync.WaitGroup{}
@@ -1148,7 +1336,7 @@ func setupSSHSession(t *testing.T, options agentsdk.Metadata) *ssh.Session {
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 	defer cancel()
 	//nolint:dogsled
-	conn, _, _, _ := setupAgent(t, options, 0)
+	conn, _, _, _, _ := setupAgent(t, options, 0)
 	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	t.Cleanup(func() {
@@ -1173,6 +1361,7 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
 	*client,
 	<-chan *agentsdk.Stats,
 	afero.Fs,
+	io.Closer,
 ) {
 	if metadata.DERPMap == nil {
 		metadata.DERPMap = tailnettest.RunDERPAndSTUN(t)
@@ -1233,7 +1422,7 @@ func setupAgent(t *testing.T, metadata agentsdk.Metadata, ptyTimeout time.Durati
 	if !agentConn.AwaitReachable(ctx) {
 		t.Fatal("agent not reachable")
 	}
-	return agentConn, c, statsCh, fs
+	return agentConn, c, statsCh, fs, closer
 }
 
 var dialTestPayload = []byte("dean-was-here123")