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>

cli/cliui/agent.go

@@ -80,6 +80,10 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 	if err != nil {
 		return xerrors.Errorf("fetch: %w", err)
 	}
+	logSources := map[uuid.UUID]codersdk.WorkspaceAgentLogSource{}
+	for _, source := range agent.LogSources {
+		logSources[source.ID] = source
+	}
 
 	sw := &stageWriter{w: writer}
 
@@ -123,7 +127,7 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 				return nil
 			}
 
-			stage := "Running workspace agent startup script"
+			stage := "Running workspace agent startup scripts"
 			follow := opts.Wait
 			if !follow {
 				stage += " (non-blocking)"
@@ -173,7 +177,12 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 							return nil
 						}
 						for _, log := range logs {
-							sw.Log(log.CreatedAt, log.Level, log.Output)
+							source, hasSource := logSources[log.SourceID]
+							output := log.Output
+							if hasSource && source.DisplayName != "" {
+								output = source.DisplayName + ": " + output
+							}
+							sw.Log(log.CreatedAt, log.Level, output)
 							lastLog = log
 						}
 					}
@@ -192,16 +201,19 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 			switch agent.LifecycleState {
 			case codersdk.WorkspaceAgentLifecycleReady:
 				sw.Complete(stage, agent.ReadyAt.Sub(*agent.StartedAt))
+			case codersdk.WorkspaceAgentLifecycleStartTimeout:
+				sw.Fail(stage, 0)
+				sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: A startup script timed out and your workspace may be incomplete.")
 			case codersdk.WorkspaceAgentLifecycleStartError:
 				sw.Fail(stage, agent.ReadyAt.Sub(*agent.StartedAt))
 				// Use zero time (omitted) to separate these from the startup logs.
-				sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: The startup script exited with an error and your workspace may be incomplete.")
+				sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: A startup script exited with an error and your workspace may be incomplete.")
 				sw.Log(time.Time{}, codersdk.LogLevelWarn, troubleshootingMessage(agent, "https://coder.com/docs/v2/latest/templates#startup-script-exited-with-an-error"))
 			default:
 				switch {
 				case agent.LifecycleState.Starting():
 					// Use zero time (omitted) to separate these from the startup logs.
-					sw.Log(time.Time{}, codersdk.LogLevelWarn, "Notice: The startup script is still running and your workspace may be incomplete.")
+					sw.Log(time.Time{}, codersdk.LogLevelWarn, "Notice: The startup scripts are still running and your workspace may be incomplete.")
 					sw.Log(time.Time{}, codersdk.LogLevelWarn, troubleshootingMessage(agent, "https://coder.com/docs/v2/latest/templates#your-workspace-may-be-incomplete"))
 					// Note: We don't complete or fail the stage here, it's
 					// intentionally left open to indicate this stage didn't

cli/cliui/agent_test.go

@@ -52,11 +52,36 @@ func TestAgent(t *testing.T) {
 			want: []string{
 				"⧗ Waiting for the workspace agent to connect",
 				"✔ Waiting for the workspace agent to connect",
-				"⧗ Running workspace agent startup script (non-blocking)",
-				"Notice: The startup script is still running and your workspace may be incomplete.",
+				"⧗ Running workspace agent startup scripts (non-blocking)",
+				"Notice: The startup scripts are still running and your workspace may be incomplete.",
 				"For more information and troubleshooting, see",
 			},
 		},
+		{
+			name: "Start timeout",
+			opts: cliui.AgentOptions{
+				FetchInterval: time.Millisecond,
+			},
+			iter: []func(context.Context, *codersdk.WorkspaceAgent, chan []codersdk.WorkspaceAgentLog) error{
+				func(_ context.Context, agent *codersdk.WorkspaceAgent, _ chan []codersdk.WorkspaceAgentLog) error {
+					agent.Status = codersdk.WorkspaceAgentConnecting
+					return nil
+				},
+				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentLog) error {
+					agent.Status = codersdk.WorkspaceAgentConnected
+					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStartTimeout
+					agent.FirstConnectedAt = ptr.Ref(time.Now())
+					return nil
+				},
+			},
+			want: []string{
+				"⧗ Waiting for the workspace agent to connect",
+				"✔ Waiting for the workspace agent to connect",
+				"⧗ Running workspace agent startup scripts (non-blocking)",
+				"✘ Running workspace agent startup scripts (non-blocking)",
+				"Warning: A startup script timed out and your workspace may be incomplete.",
+			},
+		},
 		{
 			name: "Initial connection timeout",
 			opts: cliui.AgentOptions{
@@ -86,8 +111,8 @@ func TestAgent(t *testing.T) {
 				"The workspace agent is having trouble connecting, wait for it to connect or restart your workspace.",
 				"For more information and troubleshooting, see",
 				"✔ Waiting for the workspace agent to connect",
-				"⧗ Running workspace agent startup script (non-blocking)",
-				"✔ Running workspace agent startup script (non-blocking)",
+				"⧗ Running workspace agent startup scripts (non-blocking)",
+				"✔ Running workspace agent startup scripts (non-blocking)",
 			},
 		},
 		{
@@ -120,7 +145,7 @@ func TestAgent(t *testing.T) {
 			},
 		},
 		{
-			name: "Startup script logs",
+			name: "Startup Logs",
 			opts: cliui.AgentOptions{
 				FetchInterval: time.Millisecond,
 				Wait:          true,
@@ -131,10 +156,15 @@ func TestAgent(t *testing.T) {
 					agent.FirstConnectedAt = ptr.Ref(time.Now())
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStarting
 					agent.StartedAt = ptr.Ref(time.Now())
+					agent.LogSources = []codersdk.WorkspaceAgentLogSource{{
+						ID:          uuid.Nil,
+						DisplayName: "testing",
+					}}
 					logs <- []codersdk.WorkspaceAgentLog{
 						{
 							CreatedAt: time.Now(),
 							Output:    "Hello world",
+							SourceID:  uuid.Nil,
 						},
 					}
 					return nil
@@ -152,10 +182,10 @@ func TestAgent(t *testing.T) {
 				},
 			},
 			want: []string{
-				"⧗ Running workspace agent startup script",
-				"Hello world",
+				"⧗ Running workspace agent startup scripts",
+				"testing: Hello world",
 				"Bye now",
-				"✔ Running workspace agent startup script",
+				"✔ Running workspace agent startup scripts",
 			},
 		},
 		{
@@ -181,10 +211,10 @@ func TestAgent(t *testing.T) {
 				},
 			},
 			want: []string{
-				"⧗ Running workspace agent startup script",
+				"⧗ Running workspace agent startup scripts",
 				"Hello world",
-				"✘ Running workspace agent startup script",
-				"Warning: The startup script exited with an error and your workspace may be incomplete.",
+				"✘ Running workspace agent startup scripts",
+				"Warning: A startup script exited with an error and your workspace may be incomplete.",
 				"For more information and troubleshooting, see",
 			},
 		},
@@ -229,9 +259,9 @@ func TestAgent(t *testing.T) {
 				},
 			},
 			want: []string{
-				"⧗ Running workspace agent startup script",
+				"⧗ Running workspace agent startup scripts",
 				"Hello world",
-				"✔ Running workspace agent startup script",
+				"✔ Running workspace agent startup scripts",
 			},
 			wantErr: true,
 		},
@@ -288,11 +318,10 @@ func TestAgent(t *testing.T) {
 
 			var buf bytes.Buffer
 			agent := codersdk.WorkspaceAgent{
-				ID:                    uuid.New(),
-				Status:                codersdk.WorkspaceAgentConnecting,
-				StartupScriptBehavior: codersdk.WorkspaceAgentStartupScriptBehaviorNonBlocking,
-				CreatedAt:             time.Now(),
-				LifecycleState:        codersdk.WorkspaceAgentLifecycleCreated,
+				ID:             uuid.New(),
+				Status:         codersdk.WorkspaceAgentConnecting,
+				CreatedAt:      time.Now(),
+				LifecycleState: codersdk.WorkspaceAgentLifecycleCreated,
 			}
 			logs := make(chan []codersdk.WorkspaceAgentLog, 1)
 
@@ -340,6 +369,9 @@ func TestAgent(t *testing.T) {
 				line := s.Text()
 				t.Log(line)
 				if len(tc.want) == 0 {
+					for i := 0; i < 5; i++ {
+						t.Log(line)
+					}
 					require.Fail(t, "unexpected line", line)
 				}
 				require.Contains(t, line, tc.want[0])

cli/ssh.go

@@ -143,13 +143,11 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			case "no":
 				wait = false
 			case "auto":
-				switch workspaceAgent.StartupScriptBehavior {
-				case codersdk.WorkspaceAgentStartupScriptBehaviorBlocking:
-					wait = true
-				case codersdk.WorkspaceAgentStartupScriptBehaviorNonBlocking:
-					wait = false
-				default:
-					return xerrors.Errorf("unknown startup script behavior %q", workspaceAgent.StartupScriptBehavior)
+				for _, script := range workspaceAgent.Scripts {
+					if script.StartBlocksLogin {
+						wait = true
+						break
+					}
 				}
 			default:
 				return xerrors.Errorf("unknown wait value %q", waitEnum)