coder · sreya · May 18, 2023 · May 17, 2023 · May 17, 2023 · May 17, 2023
diff --git a/agent/agent.go b/agent/agent.go
@@ -62,6 +62,7 @@ type Options struct {
 	IgnorePorts            map[int]string
 	SSHMaxTimeout          time.Duration
 	TailnetListenPort      uint16
+	Subsystem              codersdk.AgentSubsystem
 }
 
 type Client interface {
@@ -119,6 +120,7 @@ func New(options Options) Agent {
 		ignorePorts:            options.IgnorePorts,
 		connStatsChan:          make(chan *agentsdk.Stats, 1),
 		sshMaxTimeout:          options.SSHMaxTimeout,
+		subsystem:              options.Subsystem,
 	}
 	a.init(ctx)
 	return a
@@ -136,6 +138,7 @@ type agent struct {
 	// listing all listening ports. This is helpful to hide ports that
 	// are used by the agent, that the user does not care about.
 	ignorePorts map[int]string
+	subsystem   codersdk.AgentSubsystem
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
@@ -488,6 +491,7 @@ func (a *agent) run(ctx context.Context) error {
 	err = a.client.PostStartup(ctx, agentsdk.PostStartupRequest{
 		Version:           buildinfo.Version(),
 		ExpandedDirectory: manifest.Directory,
+		Subsystem:         a.subsystem,
 	})
 	if err != nil {
 		return xerrors.Errorf("update workspace agent version: %w", err)
@@ -1455,3 +1459,8 @@ func expandDirectory(dir string) (string, error) {
 	}
 	return dir, nil
 }
+
+// EnvAgentSubsystem is the environment variable used to denote the
+// specialized environment in which the agent is running
+// (e.g. envbox, envbuilder).
+const EnvAgentSubsystem = "CODER_AGENT_SUBSYSTEM"
diff --git a/cli/agent.go b/cli/agent.go
@@ -26,6 +26,7 @@ import (
 	"github.com/coder/coder/agent/reaper"
 	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/cli/clibase"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/codersdk/agentsdk"
 )
 
@@ -197,6 +198,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				return xerrors.Errorf("add executable to $PATH: %w", err)
 			}
 
+			subsystem := inv.Environ.Get(agent.EnvAgentSubsystem)
 			agnt := agent.New(agent.Options{
 				Client:            client,
 				Logger:            logger,
@@ -218,6 +220,7 @@ func (r *RootCmd) workspaceAgent() *clibase.Cmd {
 				},
 				IgnorePorts:   ignorePorts,
 				SSHMaxTimeout: sshMaxTimeout,
+				Subsystem:     codersdk.AgentSubsystem(subsystem),
 			})
 
 			debugSrvClose := ServeHandler(ctx, logger, agnt.HTTPDebug(), debugAddress, "debug")

diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -12,8 +12,10 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/clitest"
 	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionersdk/proto"
 	"github.com/coder/coder/pty/ptytest"
@@ -235,4 +237,43 @@ func TestWorkspaceAgent(t *testing.T) {
 		_, err = uuid.Parse(strings.TrimSpace(string(token)))
 		require.NoError(t, err)
 	})
+
+	t.Run("PostStartup", func(t *testing.T) {
+		t.Parallel()
+
+		authToken := uuid.NewString()
+		client := coderdtest.New(t, &coderdtest.Options{
+			IncludeProvisionerDaemon: true,
+		})
+		user := coderdtest.CreateFirstUser(t, client)
+		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:          echo.ParseComplete,
+			ProvisionApply: echo.ProvisionApplyWithAgent(authToken),
+		})
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+
+		logDir := t.TempDir()
+		inv, _ := clitest.New(t,
+			"agent",
+			"--auth", "token",
+			"--agent-token", authToken,
+			"--agent-url", client.URL.String(),
+			"--log-dir", logDir,
+		)
+		// Set the subsystem for the agent.
+		inv.Environ.Set(agent.EnvAgentSubsystem, string(codersdk.AgentSubsystemEnvbox))
+
+		pty := ptytest.New(t).Attach(inv)
+
+		clitest.Start(t, inv)
+		pty.ExpectMatch("starting agent")
+
+		resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+		require.Len(t, resources, 1)
+		require.Len(t, resources[0].Agents, 1)
+		require.Equal(t, codersdk.AgentSubsystemEnvbox, resources[0].Agents[0].Subsystem)
+	})
 }
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/coderd/database/dbauthz/querier_test.go b/coderd/database/dbauthz/querier_test.go
@@ -1029,7 +1029,8 @@ func (s *MethodTestSuite) TestWorkspace() {
 		res := dbgen.WorkspaceResource(s.T(), db, database.WorkspaceResource{JobID: build.JobID})
 		agt := dbgen.WorkspaceAgent(s.T(), db, database.WorkspaceAgent{ResourceID: res.ID})
 		check.Args(database.UpdateWorkspaceAgentStartupByIDParams{
-			ID: agt.ID,
+			ID:        agt.ID,
+			Subsystem: database.WorkspaceAgentSubsystemNone,
 		}).Asserts(ws, rbac.ActionUpdate).Returns()
 	}))
 	s.Run("GetWorkspaceAgentStartupLogsAfter", s.Subtest(func(db database.Store, check *expects) {

diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
@@ -3702,6 +3702,7 @@ func (q *fakeQuerier) UpdateWorkspaceAgentStartupByID(_ context.Context, arg dat
 
 		agent.Version = arg.Version
 		agent.ExpandedDirectory = arg.ExpandedDirectory
+		agent.Subsystem = arg.Subsystem
 		q.workspaceAgents[index] = agent
 		return nil
 	}

diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql
diff --git a/coderd/database/migrations/000123_workspace_agent_subsystem.down.sql b/coderd/database/migrations/000123_workspace_agent_subsystem.down.sql
@@ -0,0 +1,4 @@
+BEGIN;
+ALTER TABLE workspace_agents DROP COLUMN subsystem;
+DROP TYPE workspace_agent_subsystem;
+COMMIT;
diff --git a/coderd/database/migrations/000123_workspace_agent_subsystem.up.sql b/coderd/database/migrations/000123_workspace_agent_subsystem.up.sql
@@ -0,0 +1,4 @@
+BEGIN;
+CREATE TYPE workspace_agent_subsystem AS ENUM ('envbuilder', 'envbox', 'none');
+ALTER TABLE workspace_agents ADD COLUMN subsystem workspace_agent_subsystem NOT NULL default 'none';
+COMMIT;
diff --git a/coderd/database/models.go b/coderd/database/models.go