Skip to content

Commit 1b9e884

Browse files
mafredrimafredri
committed
fix(agent/agentcontainers): update sub agent client on reconnect
Fixes coder/internal#697
1 parent 97474bb commit 1b9e884

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()
@@ -623,9 +632,9 @@ func safeFriendlyName(name string) string {
623632
return name
624633
}
625634

626-
// refreshContainers triggers an immediate update of the container list
635+
// RefreshContainers triggers an immediate update of the container list
627636
// and waits for it to complete.
628-
func (api *API) refreshContainers(ctx context.Context) (err error) {
637+
func (api *API) RefreshContainers(ctx context.Context) (err error) {
629638
defer func() {
630639
if err != nil {
631640
err = xerrors.Errorf("refresh containers failed: %w", err)
@@ -860,7 +869,7 @@ func (api *API) recreateDevcontainer(dc codersdk.WorkspaceAgentDevcontainer, con
860869

861870
// Ensure an immediate refresh to accurately reflect the
862871
// devcontainer state after recreation.
863-
if err := api.refreshContainers(ctx); err != nil {
872+
if err := api.RefreshContainers(ctx); err != nil {
864873
logger.Error(ctx, "failed to trigger immediate refresh after devcontainer recreation", slog.Error(err))
865874
}
866875
}
@@ -904,7 +913,8 @@ func (api *API) markDevcontainerDirty(configPath string, modifiedAt time.Time) {
904913
// slate. This method has an internal timeout to prevent blocking
905914
// indefinitely if something goes wrong with the subagent deletion.
906915
func (api *API) cleanupSubAgents(ctx context.Context) error {
907-
agents, err := api.subAgentClient.List(ctx)
916+
client := *api.subAgentClient.Load()
917+
agents, err := client.List(ctx)
908918
if err != nil {
909919
return xerrors.Errorf("list agents: %w", err)
910920
}
@@ -927,7 +937,8 @@ func (api *API) cleanupSubAgents(ctx context.Context) error {
927937
if injected[agent.ID] {
928938
continue
929939
}
930-
err := api.subAgentClient.Delete(ctx, agent.ID)
940+
client := *api.subAgentClient.Load()
941+
err := client.Delete(ctx, agent.ID)
931942
if err != nil {
932943
api.logger.Error(ctx, "failed to delete agent",
933944
slog.Error(err),
@@ -1101,7 +1112,8 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11011112

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

11461158
// Create new subagent record in the database to receive the auth token.
1147-
proc.agent, err = api.subAgentClient.Create(ctx, SubAgent{
1159+
client := *api.subAgentClient.Load()
1160+
proc.agent, err = client.Create(ctx, SubAgent{
11481161
Name: dc.Name,
11491162
Directory: directory,
11501163
OperatingSystem: "linux", // Assuming Linux for devcontainers.
@@ -1163,7 +1176,8 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11631176
if api.closed {
11641177
deleteCtx, deleteCancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
11651178
defer deleteCancel()
1166-
err := api.subAgentClient.Delete(deleteCtx, proc.agent.ID)
1179+
client := *api.subAgentClient.Load()
1180+
err := client.Delete(deleteCtx, proc.agent.ID)
11671181
if err != nil {
11681182
return xerrors.Errorf("delete existing subagent failed after API closed: %w", err)
11691183
}
@@ -1249,8 +1263,9 @@ func (api *API) Close() error {
12491263
// Note: We can't use api.ctx here because it's canceled.
12501264
deleteCtx, deleteCancel := context.WithTimeout(context.Background(), defaultOperationTimeout)
12511265
defer deleteCancel()
1266+
client := *api.subAgentClient.Load()
12521267
for _, id := range subAgentIDs {
1253-
err := api.subAgentClient.Delete(deleteCtx, id)
1268+
err := client.Delete(deleteCtx, id)
12541269
if err != nil {
12551270
api.logger.Error(api.ctx, "delete subagent record during shutdown failed",
12561271
slog.Error(err),

0 commit comments

Comments
 (0)