diff --git a/agent/agent.go b/agent/agent.go index 550392802df..9da4cbea90d 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 153bbaa51ab..578626a208a 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 c032ea1f83a..712a12f308e 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()) +}