Skip to content

Commit 54c9521

Browse files
mafredrimafredri
committed
fix(agent/agentcontainers): update sub agent client on reconnect
Fixes coder/internal#697
1 parent 2b22ee2 commit 54c9521

File tree

2 files changed

+41
-12
lines changed

2 files changed

+41
-12
lines changed

‎agent/agent.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1188,6 +1188,14 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context,
11881188
}
11891189
a.metrics.startupScriptSeconds.WithLabelValues(label).Set(dur)
11901190
a.scriptRunner.StartCron()
1191+
1192+
// If the container API is enabled, trigger an immediate refresh
1193+
// for quick sub agent injection.
1194+
if cAPI := a.containerAPI.Load(); cAPI != nil {
1195+
if err := cAPI.RefreshContainers(ctx); err != nil {
1196+
a.logger.Error(ctx, "failed to refresh containers", slog.Error(err))
1197+
}
1198+
}
11911199
})
11921200
if err != nil {
11931201
return xerrors.Errorf("track conn goroutine: %w", err)
@@ -1253,6 +1261,12 @@ func (a *agent) createOrUpdateNetwork(manifestOK, networkOK *checkpoint) func(co
12531261
network.SetDERPMap(manifest.DERPMap)
12541262
network.SetDERPForceWebSockets(manifest.DERPForceWebSockets)
12551263
network.SetBlockEndpoints(manifest.DisableDirectConnections)
1264+
1265+
// Update the subagent client if the container API is available.
1266+
if cAPI := a.containerAPI.Load(); cAPI != nil {
1267+
client := agentcontainers.NewSubAgentClientFromAPI(a.logger, aAPI)
1268+
cAPI.UpdateSubAgentClient(client)
1269+
}
12561270
}
12571271
return nil
12581272
}

‎agent/agentcontainers/api.go

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"slices"
1515
"strings"
1616
"sync"
17+
"sync/atomic"
1718
"time"
1819

1920
"github.com/fsnotify/fsnotify"
@@ -59,7 +60,7 @@ type API struct {
5960
dccli DevcontainerCLI
6061
clock quartz.Clock
6162
scriptLogger func(logSourceID uuid.UUID) ScriptLogger
62-
subAgentClient SubAgentClient
63+
subAgentClient atomic.Pointer[SubAgentClient]
6364
subAgentURL string
6465
subAgentEnv []string
6566

@@ -133,7 +134,7 @@ func WithDevcontainerCLI(dccli DevcontainerCLI) Option {
133134
// This is used to list, create and delete devcontainer agents.
134135
func WithSubAgentClient(client SubAgentClient) Option {
135136
return func(api *API) {
136-
api.subAgentClient = client
137+
api.subAgentClient.Store(&client)
137138
}
138139
}
139140

@@ -230,7 +231,6 @@ func NewAPI(logger slog.Logger, options ...Option) *API {
230231
logger: logger,
231232
clock: quartz.NewReal(),
232233
execer: agentexec.DefaultExecer,
233-
subAgentClient: noopSubAgentClient{},
234234
containerLabelIncludeFilter: make(map[string]string),
235235
devcontainerNames: make(map[string]bool),
236236
knownDevcontainers: make(map[string]codersdk.WorkspaceAgentDevcontainer),
@@ -259,6 +259,10 @@ func NewAPI(logger slog.Logger, options ...Option) *API {
259259
api.watcher = watcher.NewNoop()
260260
}
261261
}
262+
if api.subAgentClient.Load() == nil {
263+
var c SubAgentClient = noopSubAgentClient{}
264+
api.subAgentClient.Store(&c)
265+
}
262266

263267
go api.watcherLoop()
264268
go api.updaterLoop()
@@ -375,6 +379,11 @@ func (api *API) updaterLoop() {
375379
}
376380
}
377381

382+
// UpdateSubAgentClient updates the `SubAgentClient` for the API.
383+
func (api *API) UpdateSubAgentClient(client SubAgentClient) {
384+
api.subAgentClient.Store(&client)
385+
}
386+
378387
// Routes returns the HTTP handler for container-related routes.
379388
func (api *API) Routes() http.Handler {
380389
r := chi.NewRouter()
@@ -622,9 +631,9 @@ func safeFriendlyName(name string) string {
622631
return name
623632
}
624633

625-
// refreshContainers triggers an immediate update of the container list
634+
// RefreshContainers triggers an immediate update of the container list
626635
// and waits for it to complete.
627-
func (api *API) refreshContainers(ctx context.Context) (err error) {
636+
func (api *API) RefreshContainers(ctx context.Context) (err error) {
628637
defer func() {
629638
if err != nil {
630639
err = xerrors.Errorf("refresh containers failed: %w", err)
@@ -859,7 +868,7 @@ func (api *API) recreateDevcontainer(dc codersdk.WorkspaceAgentDevcontainer, con
859868

860869
// Ensure an immediate refresh to accurately reflect the
861870
// devcontainer state after recreation.
862-
if err := api.refreshContainers(ctx); err != nil {
871+
if err := api.RefreshContainers(ctx); err != nil {
863872
logger.Error(ctx, "failed to trigger immediate refresh after devcontainer recreation", slog.Error(err))
864873
}
865874
}
@@ -903,7 +912,8 @@ func (api *API) markDevcontainerDirty(configPath string, modifiedAt time.Time) {
903912
// slate. This method has an internal timeout to prevent blocking
904913
// indefinitely if something goes wrong with the subagent deletion.
905914
func (api *API) cleanupSubAgents(ctx context.Context) error {
906-
agents, err := api.subAgentClient.List(ctx)
915+
client := *api.subAgentClient.Load()
916+
agents, err := client.List(ctx)
907917
if err != nil {
908918
return xerrors.Errorf("list agents: %w", err)
909919
}
@@ -926,7 +936,8 @@ func (api *API) cleanupSubAgents(ctx context.Context) error {
926936
if injected[agent.ID] {
927937
continue
928938
}
929-
err := api.subAgentClient.Delete(ctx, agent.ID)
939+
client := *api.subAgentClient.Load()
940+
err := client.Delete(ctx, agent.ID)
930941
if err != nil {
931942
api.logger.Error(ctx, "failed to delete agent",
932943
slog.Error(err),
@@ -1100,7 +1111,8 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11001111

11011112
if proc.agent.ID != uuid.Nil && recreateSubAgent {
11021113
logger.Debug(ctx, "deleting existing subagent for recreation", slog.F("agent_id", proc.agent.ID))
1103-
err = api.subAgentClient.Delete(ctx, proc.agent.ID)
1114+
client := *api.subAgentClient.Load()
1115+
err = client.Delete(ctx, proc.agent.ID)
11041116
if err != nil {
11051117
return xerrors.Errorf("delete existing subagent failed: %w", err)
11061118
}
@@ -1143,7 +1155,8 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11431155
)
11441156

11451157
// Create new subagent record in the database to receive the auth token.
1146-
proc.agent, err = api.subAgentClient.Create(ctx, SubAgent{
1158+
client := *api.subAgentClient.Load()
1159+
proc.agent, err = client.Create(ctx, SubAgent{
11471160
Name: dc.Name,
11481161
Directory: directory,
11491162
OperatingSystem: "linux", // Assuming Linux for devcontainers.
@@ -1162,7 +1175,8 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11621175
if api.closed {
11631176
deleteCtx, deleteCancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
11641177
defer deleteCancel()
1165-
err := api.subAgentClient.Delete(deleteCtx, proc.agent.ID)
1178+
client := *api.subAgentClient.Load()
1179+
err := client.Delete(deleteCtx, proc.agent.ID)
11661180
if err != nil {
11671181
return xerrors.Errorf("delete existing subagent failed after API closed: %w", err)
11681182
}
@@ -1248,8 +1262,9 @@ func (api *API) Close() error {
12481262
// Note: We can't use api.ctx here because it's canceled.
12491263
deleteCtx, deleteCancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
12501264
defer deleteCancel()
1265+
client := *api.subAgentClient.Load()
12511266
for _, id := range subAgentIDs {
1252-
err := api.subAgentClient.Delete(deleteCtx, id)
1267+
err := client.Delete(deleteCtx, id)
12531268
if err != nil {
12541269
api.logger.Error(api.ctx, "delete subagent record during shutdown failed",
12551270
slog.Error(err),

0 commit comments

Comments
 (0)