From 5981e9978e83524e37756f03da6d6498dc30812e Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Fri, 10 Jul 2026 10:19:25 +0000 Subject: [PATCH] feat(agent): register scripts as sync units for lifecycle tracking The script runner now registers each script as a sync unit during Init() so all scripts are immediately visible via 'coder exp sync list' and appear in 'coder exp sync timeline'. As each script executes, its unit status transitions from pending to started, then to completed. Changes: - agent/agentscripts/agentscripts.go: Options gains a UnitManager field. Init() registers scripts using DisplayName as the unit ID. run() sets unit status to started before execution and completed after. - agent/agent.go: Creates a shared unit.Manager and passes it to both the script runner and the socket server (via WithUnitManager). - Tests: added TestScriptUnitsRegistered and TestScriptUnitsLifecycle. Rebased onto the sync-unit-timeline branch: the socket server already exposes a WithUnitManager option there, so this revives PR #26459 without re-introducing its NewServer signature change. --- agent/agent.go | 5 ++ agent/agentscripts/agentscripts.go | 50 +++++++++++ agent/agentscripts/agentscripts_test.go | 113 ++++++++++++++++++++++++ 3 files changed, 168 insertions(+) diff --git a/agent/agent.go b/agent/agent.go index 550392802df5a..9da4cbea90daa 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -52,6 +52,7 @@ import ( "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/agent/proto/resourcesmonitor" "github.com/coder/coder/v2/agent/reconnectingpty" + "github.com/coder/coder/v2/agent/unit" "github.com/coder/coder/v2/agent/usershell" "github.com/coder/coder/v2/agent/x/agentdesktop" "github.com/coder/coder/v2/agent/x/agentmcp" @@ -367,6 +368,7 @@ type agent struct { socketServerEnabled bool socketPath string socketServer *agentsocket.Server + unitManager *unit.Manager derpTLSConfig *tls.Config } @@ -449,6 +451,7 @@ func (a *agent) init() { panic(err) } a.sshServer = sshSrv + a.unitManager = unit.NewManager() a.scriptRunner = agentscripts.New(agentscripts.Options{ LogDir: a.logDir, DataDirBase: a.scriptDataDir, @@ -458,6 +461,7 @@ func (a *agent) init() { GetScriptLogger: func(logSourceID uuid.UUID) agentscripts.ScriptLogger { return a.logSender.GetScriptLogger(logSourceID) }, + UnitManager: a.unitManager, }) // Register runner metrics. If the prom registry is nil, the metrics // will not report anywhere. @@ -566,6 +570,7 @@ func (a *agent) initSocketServer() { a.logger.Named("socket"), agentsocket.WithPath(a.socketPath), agentsocket.WithContextManager(a.contextManager), + agentsocket.WithUnitManager(a.unitManager), ) if err != nil { a.logger.Error(a.hardCtx, "failed to create socket server", slog.Error(err), slog.F("path", a.socketPath)) diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go index 153bbaa51abaa..578626a208a6a 100644 --- a/agent/agentscripts/agentscripts.go +++ b/agent/agentscripts/agentscripts.go @@ -23,6 +23,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/agent/agentssh" "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/agent/unit" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" @@ -56,6 +57,9 @@ type Options struct { SSHServer *agentssh.Server Filesystem afero.Fs GetScriptLogger func(logSourceID uuid.UUID) ScriptLogger + // UnitManager, when set, causes the runner to register each script as a + // sync unit during Init and to update unit status as scripts execute. + UnitManager *unit.Manager } // New creates a runner for the provided scripts. @@ -139,6 +143,28 @@ func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript, scriptCompleted S } r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir)) + // Register all scripts as sync units so they are immediately visible + // via "coder exp sync list". Registration happens before any script + // executes, giving observers a complete picture of pending work. + if r.UnitManager != nil { + for _, script := range scripts { + if script.DisplayName == "" { + r.Logger.Warn(r.cronCtx, "skipping unit registration for script with empty display name", + slog.F("script_id", script.ID)) + continue + } + if err := r.UnitManager.Register(unit.ID(script.DisplayName)); err != nil { + if errors.Is(err, unit.ErrUnitAlreadyRegistered) { + r.Logger.Warn(r.cronCtx, "duplicate unit name during script registration", + slog.F("display_name", script.DisplayName), + slog.F("script_id", script.ID)) + continue + } + return xerrors.Errorf("register script unit %q: %w", script.DisplayName, err) + } + } + } + err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700) if err != nil { return xerrors.Errorf("create script bin dir: %w", err) @@ -335,6 +361,9 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript, cmd.Stdout = io.MultiWriter(fileWriter, infoW) cmd.Stderr = io.MultiWriter(fileWriter, errW) + // Mark the script's sync unit as started before execution begins. + r.setUnitStatus(ctx, script, unit.StatusStarted) + start := dbtime.Now() defer func() { end := dbtime.Now() @@ -404,6 +433,10 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript, if err != nil { logger.Warn(ctx, "reporting script completed: track command goroutine", slog.Error(err)) } + + // Mark the script's sync unit as completed after execution finishes, + // regardless of whether the script succeeded or failed. + r.setUnitStatus(ctx, script, unit.StatusComplete) }() err = cmd.Start() @@ -488,3 +521,20 @@ func (r *Runner) isClosed() bool { return false } } + +// setUnitStatus updates the sync unit status for a script. +// It is a no-op when no UnitManager is configured or when the script +// has no DisplayName. +func (r *Runner) setUnitStatus(ctx context.Context, script codersdk.WorkspaceAgentScript, status unit.Status) { + if r.UnitManager == nil || script.DisplayName == "" { + return + } + if err := r.UnitManager.UpdateStatus(unit.ID(script.DisplayName), status); err != nil { + // Log but do not fail; unit tracking is best-effort relative + // to script execution. + r.Logger.Warn(ctx, "failed to update unit status", + slog.F("display_name", script.DisplayName), + slog.F("status", string(status)), + slog.Error(err)) + } +} diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go index c032ea1f83a1a..712a12f308eba 100644 --- a/agent/agentscripts/agentscripts_test.go +++ b/agent/agentscripts/agentscripts_test.go @@ -19,6 +19,7 @@ import ( "github.com/coder/coder/v2/agent/agentscripts" "github.com/coder/coder/v2/agent/agentssh" "github.com/coder/coder/v2/agent/agenttest" + "github.com/coder/coder/v2/agent/unit" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/testutil" @@ -355,3 +356,115 @@ func (*fakeScriptLogger) Flush(context.Context) error { func newFakeScriptLogger() *fakeScriptLogger { return &fakeScriptLogger{make(chan agentsdk.Log, 100)} } + +func setupWithUnitManager(t *testing.T, getScriptLogger func(uuid.UUID) agentscripts.ScriptLogger) (*agentscripts.Runner, *unit.Manager) { + t.Helper() + logger := testutil.Logger(t) + fs := afero.NewOsFs() + s, err := agentssh.NewServer(context.Background(), logger, prometheus.NewRegistry(), fs, agentexec.DefaultExecer, nil) + require.NoError(t, err) + t.Cleanup(func() { + _ = s.Close() + }) + m := unit.NewManager() + r := agentscripts.New(agentscripts.Options{ + LogDir: t.TempDir(), + DataDirBase: t.TempDir(), + Logger: logger, + SSHServer: s, + Filesystem: fs, + GetScriptLogger: getScriptLogger, + UnitManager: m, + }) + return r, m +} + +func TestScriptUnitsRegistered(t *testing.T) { + t.Parallel() + + runner, mgr := setupWithUnitManager(t, func(_ uuid.UUID) agentscripts.ScriptLogger { + return noopScriptLogger{} + }) + defer runner.Close() + + scripts := []codersdk.WorkspaceAgentScript{ + { + ID: uuid.New(), + LogSourceID: uuid.New(), + DisplayName: "Install Tools", + Script: "echo install", + RunOnStart: true, + }, + { + ID: uuid.New(), + LogSourceID: uuid.New(), + DisplayName: "Configure Environment", + Script: "echo configure", + RunOnStart: true, + }, + { + ID: uuid.New(), + LogSourceID: uuid.New(), + DisplayName: "", // No display name; should be skipped. + Script: "echo nameless", + RunOnStart: true, + }, + } + + aAPI := agenttest.NewFakeAgentAPI(t, testutil.Logger(t), nil, nil) + err := runner.Init(scripts, aAPI.ScriptCompleted) + require.NoError(t, err) + + // All scripts with a display name should be registered as pending. + units := mgr.ListUnits() + require.Len(t, units, 2) + + unitsByName := make(map[unit.ID]unit.Unit) + for _, u := range units { + unitsByName[u.ID()] = u + } + + for _, name := range []string{"Install Tools", "Configure Environment"} { + u, ok := unitsByName[unit.ID(name)] + require.True(t, ok, "expected unit %q to be registered", name) + require.Equal(t, unit.StatusPending, u.Status(), "unit %q should be pending", name) + } +} + +func TestScriptUnitsLifecycle(t *testing.T) { + t.Parallel() + + runner, mgr := setupWithUnitManager(t, func(_ uuid.UUID) agentscripts.ScriptLogger { + return noopScriptLogger{} + }) + defer runner.Close() + + scriptName := "My Script" + scripts := []codersdk.WorkspaceAgentScript{ + { + ID: uuid.New(), + LogSourceID: uuid.New(), + DisplayName: scriptName, + Script: "echo lifecycle", + RunOnStart: true, + }, + } + + aAPI := agenttest.NewFakeAgentAPI(t, testutil.Logger(t), nil, nil) + err := runner.Init(scripts, aAPI.ScriptCompleted) + require.NoError(t, err) + + // Before execution: pending. + u, err := mgr.Unit(unit.ID(scriptName)) + require.NoError(t, err) + require.Equal(t, unit.StatusPending, u.Status()) + + // Execute the script. + err = runner.Execute(context.Background(), agentscripts.ExecuteStartScripts) + require.NoError(t, err) + + // After execution: completed. + u, err = mgr.Unit(unit.ID(scriptName)) + require.NoError(t, err) + require.Equal(t, unit.StatusComplete, u.Status()) +}