feat: add support for coder_script (#9584)

* Add basic migrations

* Improve schema

* Refactor agent scripts into it's own package

* Support legacy start and stop script format

* Pipe the scripts!

* Finish the piping

* Fix context usage

* It works!

* Fix sql query

* Fix SQL query

* Rename `LogSourceID` -> `SourceID`

* Fix the FE

* fmt

* Rename migrations

* Fix log tests

* Fix lint err

* Fix gen

* Fix story type

* Rename source to script

* Fix schema jank

* Uncomment test

* Rename proto to TimeoutSeconds

* Fix comments

* Fix comments

* Fix legacy endpoint without specified log_source

* Fix non-blocking by default in agent

* Fix resources tests

* Fix dbfake

* Fix resources

* Fix linting I think

* Add fixtures

* fmt

* Fix startup script behavior

* Fix comments

* Fix context

* Fix cancel

* Fix SQL tests

* Fix e2e tests

* Interrupt on Windows

* Fix agent leaking script process

* Fix migrations

* Fix stories

* Fix duplicate logs appearing

* Gen

* Fix log location

* Fix tests

* Fix tests

* Fix log output

* Show display name in output

* Fix print

* Return timeout on start context

* Gen

* Fix fixture

* Fix the agent status

* Fix startup timeout msg

* Fix command using shared context

* Fix timeout draining

* Change signal type

* Add deterministic colors to startup script logs

---------

Co-authored-by: Muhammad Atif Ali <atif@coder.com>

agent/agent.go

@@ -12,7 +12,6 @@ import (
 	"net/http"
 	"net/netip"
 	"os"
-	"os/exec"
 	"os/user"
 	"path/filepath"
 	"runtime"
@@ -37,6 +36,7 @@ import (
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/agent/agentscripts"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/agent/reconnectingpty"
 	"github.com/coder/coder/v2/buildinfo"
@@ -196,6 +196,7 @@ type agent struct {
 
 	manifest                     atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.
 	reportMetadataInterval       time.Duration
+	scriptRunner                 *agentscripts.Runner
 	serviceBanner                atomic.Pointer[codersdk.ServiceBannerConfig] // serviceBanner is atomic because it is periodically updated.
 	serviceBannerRefreshInterval time.Duration
 	sessionToken                 atomic.Pointer[string]
@@ -238,7 +239,13 @@ func (a *agent) init(ctx context.Context) {
 	sshSrv.Manifest = &a.manifest
 	sshSrv.ServiceBanner = &a.serviceBanner
 	a.sshServer = sshSrv
-
+	a.scriptRunner = agentscripts.New(agentscripts.Options{
+		LogDir:     a.logDir,
+		Logger:     a.logger,
+		SSHServer:  sshSrv,
+		Filesystem: a.filesystem,
+		PatchLogs:  a.client.PatchLogs,
+	})
 	go a.runLoop(ctx)
 }
 
@@ -657,41 +664,29 @@ func (a *agent) run(ctx context.Context) error {
 			}
 		}
 
-		lifecycleState := codersdk.WorkspaceAgentLifecycleReady
-		scriptDone := make(chan error, 1)
+		err = a.scriptRunner.Init(manifest.Scripts)
+		if err != nil {
+			return xerrors.Errorf("init script runner: %w", err)
+		}
 		err = a.trackConnGoroutine(func() {
-			defer close(scriptDone)
-			scriptDone <- a.runStartupScript(ctx, manifest.StartupScript)
+			err := a.scriptRunner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool {
+				return script.RunOnStart
+			})
+			if err != nil {
+				a.logger.Warn(ctx, "startup script failed", slog.Error(err))
+				if errors.Is(err, agentscripts.ErrTimeout) {
+					a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleStartTimeout)
+				} else {
+					a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleStartError)
+				}
+			} else {
+				a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleReady)
+			}
+			a.scriptRunner.StartCron()
 		})
 		if err != nil {
-			return xerrors.Errorf("track startup script: %w", err)
+			return xerrors.Errorf("track conn goroutine: %w", err)
 		}
-		go func() {
-			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 manifest.StartupScriptTimeout > 0 {
-				t := time.NewTimer(manifest.StartupScriptTimeout)
-				defer t.Stop()
-				timeout = t.C
-			}
-
-			var err error
-			select {
-			case err = <-scriptDone:
-			case <-timeout:
-				a.logger.Warn(ctx, "script timed out", slog.F("lifecycle", "startup"), slog.F("timeout", manifest.StartupScriptTimeout))
-				a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleStartTimeout)
-				err = <-scriptDone // The script can still complete after a timeout.
-			}
-			if err != nil {
-				if errors.Is(err, context.Canceled) {
-					return
-				}
-				lifecycleState = codersdk.WorkspaceAgentLifecycleStartError
-			}
-			a.setLifecycle(ctx, lifecycleState)
-		}()
 	}
 
 	// This automatically closes when the context ends!
@@ -1006,93 +1001,6 @@ func (a *agent) runDERPMapSubscriber(ctx context.Context, network *tailnet.Conn)
 	}
 }
 
-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) (err error) {
-	if script == "" {
-		return nil
-	}
-
-	logger := a.logger.With(slog.F("lifecycle", lifecycle))
-
-	logger.Info(ctx, fmt.Sprintf("running %s script", lifecycle), slog.F("script", script))
-	fileWriter, 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 %s script log file: %w", lifecycle, err)
-	}
-	defer func() {
-		err := fileWriter.Close()
-		if err != nil {
-			logger.Warn(ctx, fmt.Sprintf("close %s script log file", lifecycle), slog.Error(err))
-		}
-	}()
-
-	cmdPty, err := a.sshServer.CreateCommand(ctx, script, nil)
-	if err != nil {
-		return xerrors.Errorf("%s script: create command: %w", lifecycle, err)
-	}
-	cmd := cmdPty.AsExec()
-
-	var stdout, stderr io.Writer = fileWriter, fileWriter
-	if lifecycle == "startup" {
-		send, flushAndClose := agentsdk.LogsSender(a.client.PatchLogs, logger)
-		// If ctx is canceled here (or in a writer below), we may be
-		// discarding logs, but that's okay because we're shutting down
-		// anyway. We could consider creating a new context here if we
-		// want better control over flush during shutdown.
-		defer func() {
-			if err := flushAndClose(ctx); err != nil {
-				logger.Warn(ctx, "flush startup logs failed", slog.Error(err))
-			}
-		}()
-
-		infoW := agentsdk.StartupLogsWriter(ctx, send, codersdk.WorkspaceAgentLogSourceStartupScript, codersdk.LogLevelInfo)
-		defer infoW.Close()
-		errW := agentsdk.StartupLogsWriter(ctx, send, codersdk.WorkspaceAgentLogSourceStartupScript, codersdk.LogLevelError)
-		defer errW.Close()
-
-		stdout = io.MultiWriter(fileWriter, infoW)
-		stderr = io.MultiWriter(fileWriter, errW)
-	}
-
-	cmd.Stdout = stdout
-	cmd.Stderr = stderr
-
-	start := time.Now()
-	defer func() {
-		end := time.Now()
-		execTime := end.Sub(start)
-		exitCode := 0
-		if err != nil {
-			exitCode = 255 // Unknown status.
-			var exitError *exec.ExitError
-			if xerrors.As(err, &exitError) {
-				exitCode = exitError.ExitCode()
-			}
-			logger.Warn(ctx, fmt.Sprintf("%s script failed", lifecycle), slog.F("execution_time", execTime), slog.F("exit_code", exitCode), slog.Error(err))
-		} else {
-			logger.Info(ctx, fmt.Sprintf("%s script completed", lifecycle), slog.F("execution_time", execTime), slog.F("exit_code", exitCode))
-		}
-	}()
-
-	err = cmd.Run()
-	if err != nil {
-		// cmd.Run does not return a context canceled error, it returns "signal: killed".
-		if ctx.Err() != nil {
-			return ctx.Err()
-		}
-
-		return xerrors.Errorf("%s script: run: %w", lifecycle, err)
-	}
-	return nil
-}
-
 func (a *agent) handleReconnectingPTY(ctx context.Context, logger slog.Logger, msg codersdk.WorkspaceAgentReconnectingPTYInit, conn net.Conn) (retErr error) {
 	defer conn.Close()
 	a.metrics.connectionsTotal.Add(1)
@@ -1475,39 +1383,23 @@ func (a *agent) Close() error {
 	}
 
 	lifecycleState := codersdk.WorkspaceAgentLifecycleOff
-	if manifest := a.manifest.Load(); manifest != nil && manifest.ShutdownScript != "" {
-		scriptDone := make(chan error, 1)
-		go func() {
-			defer close(scriptDone)
-			scriptDone <- a.runShutdownScript(ctx, manifest.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 manifest.ShutdownScriptTimeout > 0 {
-			t := time.NewTimer(manifest.ShutdownScriptTimeout)
-			defer t.Stop()
-			timeout = t.C
-		}
-
-		var err error
-		select {
-		case err = <-scriptDone:
-		case <-timeout:
-			a.logger.Warn(ctx, "script timed out", slog.F("lifecycle", "shutdown"), slog.F("timeout", manifest.ShutdownScriptTimeout))
-			a.setLifecycle(ctx, codersdk.WorkspaceAgentLifecycleShutdownTimeout)
-			err = <-scriptDone // The script can still complete after a timeout.
-		}
-		if err != nil {
+	err = a.scriptRunner.Execute(ctx, func(script codersdk.WorkspaceAgentScript) bool {
+		return script.RunOnStop
+	})
+	if err != nil {
+		if errors.Is(err, agentscripts.ErrTimeout) {
+			lifecycleState = codersdk.WorkspaceAgentLifecycleShutdownTimeout
+		} else {
 			lifecycleState = codersdk.WorkspaceAgentLifecycleShutdownError
 		}
 	}
-
-	// Set final state and wait for it to be reported because context
-	// cancellation will stop the report loop.
 	a.setLifecycle(ctx, lifecycleState)
 
+	err = a.scriptRunner.Close()
+	if err != nil {
+		a.logger.Error(ctx, "script runner close", slog.Error(err))
+	}
+
 	// 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)

agent/agent_test.go

@@ -50,7 +50,6 @@ import (
 	"github.com/coder/coder/v2/agent/agentproc/agentproctest"
 	"github.com/coder/coder/v2/agent/agentssh"
 	"github.com/coder/coder/v2/agent/agenttest"
-	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/pty"
@@ -1060,84 +1059,6 @@ func TestAgent_SSHConnectionEnvVars(t *testing.T) {
 	}
 }
 
-func TestAgent_StartupScript(t *testing.T) {
-	t.Parallel()
-	output := "something"
-	command := "sh -c 'echo " + output + "'"
-	if runtime.GOOS == "windows" {
-		command = "cmd.exe /c echo " + output
-	}
-	t.Run("Success", func(t *testing.T) {
-		t.Parallel()
-		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-		client := agenttest.NewClient(t,
-			logger,
-			uuid.New(),
-			agentsdk.Manifest{
-				StartupScript: command,
-				DERPMap:       &tailcfg.DERPMap{},
-			},
-			make(chan *agentsdk.Stats),
-			tailnet.NewCoordinator(logger),
-		)
-		closer := agent.New(agent.Options{
-			Client:                 client,
-			Filesystem:             afero.NewMemMapFs(),
-			Logger:                 logger.Named("agent"),
-			ReconnectingPTYTimeout: 0,
-		})
-		t.Cleanup(func() {
-			_ = closer.Close()
-		})
-		assert.Eventually(t, func() bool {
-			got := client.GetLifecycleStates()
-			return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
-		}, testutil.WaitShort, testutil.IntervalMedium)
-
-		require.Len(t, client.GetStartupLogs(), 1)
-		require.Equal(t, output, client.GetStartupLogs()[0].Output)
-	})
-	// This ensures that even when coderd sends back that the startup
-	// script has written too many lines it will still succeed!
-	t.Run("OverflowsAndSkips", func(t *testing.T) {
-		t.Parallel()
-		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-		client := agenttest.NewClient(t,
-			logger,
-			uuid.New(),
-			agentsdk.Manifest{
-				StartupScript: command,
-				DERPMap:       &tailcfg.DERPMap{},
-			},
-			make(chan *agentsdk.Stats, 50),
-			tailnet.NewCoordinator(logger),
-		)
-		client.PatchWorkspaceLogs = func() error {
-			resp := httptest.NewRecorder()
-			httpapi.Write(context.Background(), resp, http.StatusRequestEntityTooLarge, codersdk.Response{
-				Message: "Too many lines!",
-			})
-			res := resp.Result()
-			defer res.Body.Close()
-			return codersdk.ReadBodyAsError(res)
-		}
-		closer := agent.New(agent.Options{
-			Client:                 client,
-			Filesystem:             afero.NewMemMapFs(),
-			Logger:                 logger.Named("agent"),
-			ReconnectingPTYTimeout: 0,
-		})
-		t.Cleanup(func() {
-			_ = closer.Close()
-		})
-		assert.Eventually(t, func() bool {
-			got := client.GetLifecycleStates()
-			return len(got) > 0 && got[len(got)-1] == codersdk.WorkspaceAgentLifecycleReady
-		}, testutil.WaitShort, testutil.IntervalMedium)
-		require.Len(t, client.GetStartupLogs(), 0)
-	})
-}
-
 func TestAgent_Metadata(t *testing.T) {
 	t.Parallel()
 
@@ -1292,8 +1213,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "sleep 3",
-			StartupScriptTimeout: time.Nanosecond,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:     "sleep 3",
+				Timeout:    time.Millisecond,
+				RunOnStart: true,
+			}},
 		}, 0)
 
 		want := []codersdk.WorkspaceAgentLifecycle{
@@ -1314,8 +1238,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "false",
-			StartupScriptTimeout: 30 * time.Second,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:     "false",
+				Timeout:    30 * time.Second,
+				RunOnStart: true,
+			}},
 		}, 0)
 
 		want := []codersdk.WorkspaceAgentLifecycle{
@@ -1336,8 +1263,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "true",
-			StartupScriptTimeout: 30 * time.Second,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:     "true",
+				Timeout:    30 * time.Second,
+				RunOnStart: true,
+			}},
 		}, 0)
 
 		want := []codersdk.WorkspaceAgentLifecycle{
@@ -1358,8 +1288,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
-			ShutdownScript:       "sleep 3",
-			StartupScriptTimeout: 30 * time.Second,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:    "sleep 3",
+				Timeout:   30 * time.Second,
+				RunOnStop: true,
+			}},
 		}, 0)
 
 		assert.Eventually(t, func() bool {
@@ -1396,8 +1329,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
-			ShutdownScript:        "sleep 3",
-			ShutdownScriptTimeout: time.Nanosecond,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:    "sleep 3",
+				Timeout:   time.Millisecond,
+				RunOnStop: true,
+			}},
 		}, 0)
 
 		assert.Eventually(t, func() bool {
@@ -1435,8 +1371,11 @@ func TestAgent_Lifecycle(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, closer := setupAgent(t, agentsdk.Manifest{
-			ShutdownScript:        "false",
-			ShutdownScriptTimeout: 30 * time.Second,
+			Scripts: []codersdk.WorkspaceAgentScript{{
+				Script:    "false",
+				Timeout:   30 * time.Second,
+				RunOnStop: true,
+			}},
 		}, 0)
 
 		assert.Eventually(t, func() bool {
@@ -1480,9 +1419,16 @@ func TestAgent_Lifecycle(t *testing.T) {
 			logger,
 			uuid.New(),
 			agentsdk.Manifest{
-				DERPMap:        derpMap,
-				StartupScript:  "echo 1",
-				ShutdownScript: "echo " + expected,
+				DERPMap: derpMap,
+				Scripts: []codersdk.WorkspaceAgentScript{{
+					LogPath:    "coder-startup-script.log",
+					Script:     "echo 1",
+					RunOnStart: true,
+				}, {
+					LogPath:   "coder-shutdown-script.log",
+					Script:    "echo " + expected,
+					RunOnStop: true,
+				}},
 			},
 			make(chan *agentsdk.Stats, 50),
 			tailnet.NewCoordinator(logger),
@@ -1533,9 +1479,7 @@ func TestAgent_Startup(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "true",
-			StartupScriptTimeout: 30 * time.Second,
-			Directory:            "",
+			Directory: "",
 		}, 0)
 		assert.Eventually(t, func() bool {
 			return client.GetStartup().Version != ""
@@ -1547,9 +1491,7 @@ func TestAgent_Startup(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "true",
-			StartupScriptTimeout: 30 * time.Second,
-			Directory:            "~",
+			Directory: "~",
 		}, 0)
 		assert.Eventually(t, func() bool {
 			return client.GetStartup().Version != ""
@@ -1563,9 +1505,7 @@ func TestAgent_Startup(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "true",
-			StartupScriptTimeout: 30 * time.Second,
-			Directory:            "coder/coder",
+			Directory: "coder/coder",
 		}, 0)
 		assert.Eventually(t, func() bool {
 			return client.GetStartup().Version != ""
@@ -1579,9 +1519,7 @@ func TestAgent_Startup(t *testing.T) {
 		t.Parallel()
 
 		_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
-			StartupScript:        "true",
-			StartupScriptTimeout: 30 * time.Second,
-			Directory:            "$HOME",
+			Directory: "$HOME",
 		}, 0)
 		assert.Eventually(t, func() bool {
 			return client.GetStartup().Version != ""

agent/agentscripts/agentscripts.go

@@ -0,0 +1,269 @@
+package agentscripts
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/robfig/cron/v3"
+	"github.com/spf13/afero"
+	"golang.org/x/sync/errgroup"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/agent/agentssh"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+)
+
+var (
+	// ErrTimeout is returned when a script times out.
+	ErrTimeout = xerrors.New("script timed out")
+
+	parser = cron.NewParser(cron.Second | cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.DowOptional)
+)
+
+// Options are a set of options for the runner.
+type Options struct {
+	LogDir     string
+	Logger     slog.Logger
+	SSHServer  *agentssh.Server
+	Filesystem afero.Fs
+	PatchLogs  func(ctx context.Context, req agentsdk.PatchLogs) error
+}
+
+// New creates a runner for the provided scripts.
+func New(opts Options) *Runner {
+	cronCtx, cronCtxCancel := context.WithCancel(context.Background())
+	return &Runner{
+		Options:       opts,
+		cronCtx:       cronCtx,
+		cronCtxCancel: cronCtxCancel,
+		cron:          cron.New(cron.WithParser(parser)),
+		closed:        make(chan struct{}),
+	}
+}
+
+type Runner struct {
+	Options
+
+	cronCtx       context.Context
+	cronCtxCancel context.CancelFunc
+	cmdCloseWait  sync.WaitGroup
+	closed        chan struct{}
+	closeMutex    sync.Mutex
+	cron          *cron.Cron
+	initialized   atomic.Bool
+	scripts       []codersdk.WorkspaceAgentScript
+}
+
+// Init initializes the runner with the provided scripts.
+// It also schedules any scripts that have a schedule.
+// This function must be called before Execute.
+func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error {
+	if r.initialized.Load() {
+		return xerrors.New("init: already initialized")
+	}
+	r.initialized.Store(true)
+	r.scripts = scripts
+	r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir))
+
+	for _, script := range scripts {
+		if script.Cron == "" {
+			continue
+		}
+		script := script
+		_, err := r.cron.AddFunc(script.Cron, func() {
+			err := r.run(r.cronCtx, script)
+			if err != nil {
+				r.Logger.Warn(context.Background(), "run agent script on schedule", slog.Error(err))
+			}
+		})
+		if err != nil {
+			return xerrors.Errorf("add schedule: %w", err)
+		}
+	}
+	return nil
+}
+
+// StartCron starts the cron scheduler.
+// This is done async to allow for the caller to execute scripts prior.
+func (r *Runner) StartCron() {
+	r.cron.Start()
+}
+
+// Execute runs a set of scripts according to a filter.
+func (r *Runner) Execute(ctx context.Context, filter func(script codersdk.WorkspaceAgentScript) bool) error {
+	if filter == nil {
+		// Execute em' all!
+		filter = func(script codersdk.WorkspaceAgentScript) bool {
+			return true
+		}
+	}
+	var eg errgroup.Group
+	for _, script := range r.scripts {
+		if !filter(script) {
+			continue
+		}
+		script := script
+		eg.Go(func() error {
+			err := r.run(ctx, script)
+			if err != nil {
+				return xerrors.Errorf("run agent script %q: %w", script.LogSourceID, err)
+			}
+			return nil
+		})
+	}
+	return eg.Wait()
+}
+
+// run executes the provided script with the timeout.
+// If the timeout is exceeded, the process is sent an interrupt signal.
+// If the process does not exit after a few seconds, it is forcefully killed.
+// This function immediately returns after a timeout, and does not wait for the process to exit.
+func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript) error {
+	logPath := script.LogPath
+	if logPath == "" {
+		logPath = fmt.Sprintf("coder-script-%s.log", script.LogSourceID)
+	}
+	if !filepath.IsAbs(logPath) {
+		logPath = filepath.Join(r.LogDir, logPath)
+	}
+	logger := r.Logger.With(slog.F("log_path", logPath))
+	logger.Info(ctx, "running agent script", slog.F("script", script.Script))
+
+	fileWriter, err := r.Filesystem.OpenFile(logPath, os.O_CREATE|os.O_RDWR, 0o600)
+	if err != nil {
+		return xerrors.Errorf("open %s script log file: %w", logPath, err)
+	}
+	defer func() {
+		err := fileWriter.Close()
+		if err != nil {
+			logger.Warn(ctx, fmt.Sprintf("close %s script log file", logPath), slog.Error(err))
+		}
+	}()
+
+	var cmd *exec.Cmd
+	cmdCtx := ctx
+	if script.Timeout > 0 {
+		var ctxCancel context.CancelFunc
+		cmdCtx, ctxCancel = context.WithTimeout(ctx, script.Timeout)
+		defer ctxCancel()
+	}
+	cmdPty, err := r.SSHServer.CreateCommand(cmdCtx, script.Script, nil)
+	if err != nil {
+		return xerrors.Errorf("%s script: create command: %w", logPath, err)
+	}
+	cmd = cmdPty.AsExec()
+	cmd.SysProcAttr = cmdSysProcAttr()
+	cmd.WaitDelay = 10 * time.Second
+	cmd.Cancel = cmdCancel(cmd)
+
+	send, flushAndClose := agentsdk.LogsSender(script.LogSourceID, r.PatchLogs, logger)
+	// If ctx is canceled here (or in a writer below), we may be
+	// discarding logs, but that's okay because we're shutting down
+	// anyway. We could consider creating a new context here if we
+	// want better control over flush during shutdown.
+	defer func() {
+		if err := flushAndClose(ctx); err != nil {
+			logger.Warn(ctx, "flush startup logs failed", slog.Error(err))
+		}
+	}()
+
+	infoW := agentsdk.LogsWriter(ctx, send, script.LogSourceID, codersdk.LogLevelInfo)
+	defer infoW.Close()
+	errW := agentsdk.LogsWriter(ctx, send, script.LogSourceID, codersdk.LogLevelError)
+	defer errW.Close()
+	cmd.Stdout = io.MultiWriter(fileWriter, infoW)
+	cmd.Stderr = io.MultiWriter(fileWriter, errW)
+
+	start := time.Now()
+	defer func() {
+		end := time.Now()
+		execTime := end.Sub(start)
+		exitCode := 0
+		if err != nil {
+			exitCode = 255 // Unknown status.
+			var exitError *exec.ExitError
+			if xerrors.As(err, &exitError) {
+				exitCode = exitError.ExitCode()
+			}
+			logger.Warn(ctx, fmt.Sprintf("%s script failed", logPath), slog.F("execution_time", execTime), slog.F("exit_code", exitCode), slog.Error(err))
+		} else {
+			logger.Info(ctx, fmt.Sprintf("%s script completed", logPath), slog.F("execution_time", execTime), slog.F("exit_code", exitCode))
+		}
+	}()
+
+	err = cmd.Start()
+	if err != nil {
+		if errors.Is(err, context.DeadlineExceeded) {
+			return ErrTimeout
+		}
+		return xerrors.Errorf("%s script: start command: %w", logPath, err)
+	}
+
+	cmdDone := make(chan error, 1)
+	err = r.trackCommandGoroutine(func() {
+		cmdDone <- cmd.Wait()
+	})
+	if err != nil {
+		return xerrors.Errorf("%s script: track command goroutine: %w", logPath, err)
+	}
+	select {
+	case <-cmdCtx.Done():
+		// Wait for the command to drain!
+		select {
+		case <-cmdDone:
+		case <-time.After(10 * time.Second):
+		}
+		err = cmdCtx.Err()
+	case err = <-cmdDone:
+	}
+	if errors.Is(err, context.DeadlineExceeded) {
+		err = ErrTimeout
+	}
+	return err
+}
+
+func (r *Runner) Close() error {
+	r.closeMutex.Lock()
+	defer r.closeMutex.Unlock()
+	if r.isClosed() {
+		return nil
+	}
+	close(r.closed)
+	r.cronCtxCancel()
+	r.cron.Stop()
+	r.cmdCloseWait.Wait()
+	return nil
+}
+
+func (r *Runner) trackCommandGoroutine(fn func()) error {
+	r.closeMutex.Lock()
+	defer r.closeMutex.Unlock()
+	if r.isClosed() {
+		return xerrors.New("track command goroutine: closed")
+	}
+	r.cmdCloseWait.Add(1)
+	go func() {
+		defer r.cmdCloseWait.Done()
+		fn()
+	}()
+	return nil
+}
+
+func (r *Runner) isClosed() bool {
+	select {
+	case <-r.closed:
+		return true
+	default:
+		return false
+	}
+}

agent/agentscripts/agentscripts_other.go

@@ -0,0 +1,20 @@
+//go:build !windows
+
+package agentscripts
+
+import (
+	"os/exec"
+	"syscall"
+)
+
+func cmdSysProcAttr() *syscall.SysProcAttr {
+	return &syscall.SysProcAttr{
+		Setsid: true,
+	}
+}
+
+func cmdCancel(cmd *exec.Cmd) func() error {
+	return func() error {
+		return syscall.Kill(-cmd.Process.Pid, syscall.SIGHUP)
+	}
+}

agent/agentscripts/agentscripts_test.go

@@ -0,0 +1,80 @@
+package agentscripts_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+	"go.uber.org/goleak"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/agent/agentscripts"
+	"github.com/coder/coder/v2/agent/agentssh"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func TestExecuteBasic(t *testing.T) {
+	t.Parallel()
+	logs := make(chan agentsdk.PatchLogs, 1)
+	runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error {
+		logs <- req
+		return nil
+	})
+	defer runner.Close()
+	err := runner.Init([]codersdk.WorkspaceAgentScript{{
+		Script: "echo hello",
+	}})
+	require.NoError(t, err)
+	require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
+		return true
+	}))
+	log := <-logs
+	require.Equal(t, "hello", log.Logs[0].Output)
+}
+
+func TestTimeout(t *testing.T) {
+	t.Parallel()
+	runner := setup(t, nil)
+	defer runner.Close()
+	err := runner.Init([]codersdk.WorkspaceAgentScript{{
+		Script:  "sleep infinity",
+		Timeout: time.Millisecond,
+	}})
+	require.NoError(t, err)
+	require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
+}
+
+func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchLogs) error) *agentscripts.Runner {
+	t.Helper()
+	if patchLogs == nil {
+		// noop
+		patchLogs = func(ctx context.Context, req agentsdk.PatchLogs) error {
+			return nil
+		}
+	}
+	fs := afero.NewMemMapFs()
+	logger := slogtest.Make(t, nil)
+	s, err := agentssh.NewServer(context.Background(), logger, prometheus.NewRegistry(), fs, 0, "")
+	require.NoError(t, err)
+	s.AgentToken = func() string { return "" }
+	s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
+	t.Cleanup(func() {
+		_ = s.Close()
+	})
+	return agentscripts.New(agentscripts.Options{
+		LogDir:     t.TempDir(),
+		Logger:     logger,
+		SSHServer:  s,
+		Filesystem: fs,
+		PatchLogs:  patchLogs,
+	})
+}

agent/agentscripts/agentscripts_windows.go

@@ -0,0 +1,17 @@
+package agentscripts
+
+import (
+	"os"
+	"os/exec"
+	"syscall"
+)
+
+func cmdSysProcAttr() *syscall.SysProcAttr {
+	return &syscall.SysProcAttr{}
+}
+
+func cmdCancel(cmd *exec.Cmd) func() error {
+	return func() error {
+		return cmd.Process.Signal(os.Interrupt)
+	}
+}