fix(agent/agentcontainers): reduce need to recreate sub agents

mafredri · mafredri · commit d4967717f5ab · 2025-06-17T10:56:19.000Z
Updates #18332
diff --git a/agent/agentcontainers/api.go b/agent/agentcontainers/api.go
@@ -977,7 +977,7 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
 	)
 
 	// Check if subagent already exists for this devcontainer.
-	recreateSubAgent := false
+	maybeRecreateSubAgent := false
 	proc, injected := api.injectedSubAgentProcs[dc.WorkspaceFolder]
 	if injected {
 		if proc.containerID == container.ID && proc.ctx.Err() == nil {
@@ -992,12 +992,15 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
 			logger.Debug(ctx, "container ID changed, injecting subagent into new container",
 				slog.F("old_container_id", proc.containerID),
 			)
-			recreateSubAgent = true
+			maybeRecreateSubAgent = true
 		}
 
 		// Container ID changed or the subagent process is not running,
 		// stop the existing subagent context to replace it.
 		proc.stop()
+	} else {
+		// Set SubAgent defaults.
+		proc.agent.OperatingSystem = "linux" // Assuming Linux for devcontainers.
 	}
 
 	// Prepare the subAgentProcess to be used when running the subagent.
@@ -1090,36 +1093,33 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
 	// 	logger.Warn(ctx, "set CAP_NET_ADMIN on agent binary failed", slog.Error(err))
 	// }
 
+	updatedSubAgent := proc.agent
+
 	// Detect workspace folder by executing `pwd` in the container.
 	// NOTE(mafredri): This is a quick and dirty way to detect the
 	// workspace folder inside the container. In the future we will
 	// rely more on `devcontainer read-configuration`.
-	var pwdBuf bytes.Buffer
-	err = api.dccli.Exec(ctx, dc.WorkspaceFolder, dc.ConfigPath, "pwd", []string{},
-		WithExecOutput(&pwdBuf, io.Discard),
-		WithExecContainerID(container.ID),
-	)
-	if err != nil {
-		return xerrors.Errorf("check workspace folder in container: %w", err)
-	}
-	directory := strings.TrimSpace(pwdBuf.String())
-	if directory == "" {
-		logger.Warn(ctx, "detected workspace folder is empty, using default workspace folder",
-			slog.F("default_workspace_folder", DevcontainerDefaultContainerWorkspaceFolder),
+	if proc.agent.Directory == "" || maybeRecreateSubAgent {
+		var pwdBuf bytes.Buffer
+		err = api.dccli.Exec(ctx, dc.WorkspaceFolder, dc.ConfigPath, "pwd", []string{},
+			WithExecOutput(&pwdBuf, io.Discard),
+			WithExecContainerID(container.ID),
 		)
-		directory = DevcontainerDefaultContainerWorkspaceFolder
-	}
-
-	if proc.agent.ID != uuid.Nil && recreateSubAgent {
-		logger.Debug(ctx, "deleting existing subagent for recreation", slog.F("agent_id", proc.agent.ID))
-		client := *api.subAgentClient.Load()
-		err = client.Delete(ctx, proc.agent.ID)
 		if err != nil {
-			return xerrors.Errorf("delete existing subagent failed: %w", err)
+			return xerrors.Errorf("check workspace folder in container: %w", err)
+		}
+		directory := strings.TrimSpace(pwdBuf.String())
+		if directory == "" {
+			logger.Warn(ctx, "detected workspace folder is empty, using default workspace folder",
+				slog.F("default_workspace_folder", DevcontainerDefaultContainerWorkspaceFolder),
+			)
+			directory = DevcontainerDefaultContainerWorkspaceFolder
 		}
-		proc.agent = SubAgent{}
+
+		updatedSubAgent.Directory = directory
 	}
-	if proc.agent.ID == uuid.Nil {
+
+	if maybeRecreateSubAgent {
 		displayAppsMap := map[codersdk.DisplayApp]bool{
 			// NOTE(DanielleMaywood):
 			// We use the same defaults here as set in terraform-provider-coder.
@@ -1138,6 +1138,9 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
 
 			for _, customization := range coderCustomization {
 				for app, enabled := range customization.DisplayApps {
+					if _, ok := displayAppsMap[app]; !ok {
+						continue // Ignore unknown display apps.
+					}
 					displayAppsMap[app] = enabled
 				}
 			}
@@ -1149,24 +1152,37 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
 				displayApps = append(displayApps, app)
 			}
 		}
+		slices.Sort(displayApps)
+
+		updatedSubAgent.Name = dc.Name
+		updatedSubAgent.DisplayApps = displayApps
+	}
+
+	recreateSubAgent := maybeRecreateSubAgent && !proc.agent.Equal(updatedSubAgent)
+	if recreateSubAgent {
+		logger.Debug(ctx, "deleting existing subagent for recreation", slog.F("agent_id", proc.agent.ID))
+		client := *api.subAgentClient.Load()
+		err = client.Delete(ctx, proc.agent.ID)
+		if err != nil {
+			return xerrors.Errorf("delete existing subagent failed: %w", err)
+		}
+		proc.agent = SubAgent{} // In case of error, make sure this is cleared.
+
+		// Reset fields that are not part of the request.
+		updatedSubAgent.ID = uuid.Nil
+		updatedSubAgent.AuthToken = uuid.Nil
 
 		logger.Debug(ctx, "creating new subagent",
-			slog.F("directory", directory),
-			slog.F("display_apps", displayApps),
+			slog.F("directory", updatedSubAgent.Directory),
+			slog.F("display_apps", updatedSubAgent.DisplayApps),
 		)
 
 		// Create new subagent record in the database to receive the auth token.
-		client := *api.subAgentClient.Load()
-		proc.agent, err = client.Create(ctx, SubAgent{
-			Name:            dc.Name,
-			Directory:       directory,
-			OperatingSystem: "linux", // Assuming Linux for devcontainers.
-			Architecture:    arch,
-			DisplayApps:     displayApps,
-		})
+		newSubAgent, err := client.Create(ctx, updatedSubAgent)
 		if err != nil {
 			return xerrors.Errorf("create subagent failed: %w", err)
 		}
+		proc.agent = newSubAgent
 
 		logger.Info(ctx, "created new subagent", slog.F("agent_id", proc.agent.ID))
 	}
diff --git a/agent/agentcontainers/subagent.go b/agent/agentcontainers/subagent.go
@@ -2,6 +2,7 @@ package agentcontainers
 
 import (
 	"context"
+	"slices"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -23,6 +24,14 @@ type SubAgent struct {
 	DisplayApps     []codersdk.DisplayApp
 }
 
+func (s SubAgent) Equal(other SubAgent) bool {
+	return s.Name == other.Name &&
+		s.Directory == other.Directory &&
+		s.Architecture == other.Architecture &&
+		s.OperatingSystem == other.OperatingSystem &&
+		slices.Equal(s.DisplayApps, other.DisplayApps)
+}
+
 // SubAgentClient is an interface for managing sub agents and allows
 // changing the implementation without having to deal with the
 // agentproto package directly.
