Skip to content

Commit 2ec5a58

Browse files
mafredrimafredri
committed
fix(agent/agentcontainers): reduce need to recreate sub agents
Updates #18332
1 parent c033618 commit 2ec5a58

File tree

4 files changed

+145
-82
lines changed

4 files changed

+145
-82
lines changed

‎agent/agent_test.go‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2080,6 +2080,10 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
20802080
subAgentConnected := make(chan subAgentRequestPayload, 1)
20812081
subAgentReady := make(chan struct{}, 1)
20822082
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
2083+
if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/api/v2/workspaceagents/me/") {
2084+
return
2085+
}
2086+
20832087
t.Logf("Sub-agent request received: %s %s", r.Method, r.URL.Path)
20842088

20852089
if r.Method != http.MethodPost {
@@ -2226,11 +2230,22 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
22262230
// Ensure the container update routine runs.
22272231
tickerFuncTrap.MustWait(ctx).MustRelease(ctx)
22282232
tickerFuncTrap.Close()
2229-
_, next := mClock.AdvanceNext()
2230-
next.MustWait(ctx)
22312233

2232-
// Verify that a subagent was created.
2233-
subAgents := agentClient.GetSubAgents()
2234+
// Since the agent does RefreshContainers, and the ticker function
2235+
// is set to skip instead of queue, we must advance the clock
2236+
// multiple times to ensure that the sub-agent is created.
2237+
var subAgents []*proto.SubAgent
2238+
for {
2239+
_, next := mClock.AdvanceNext()
2240+
next.MustWait(ctx)
2241+
2242+
// Verify that a subagent was created.
2243+
subAgents = agentClient.GetSubAgents()
2244+
if len(subAgents) > 0 {
2245+
t.Logf("Found sub-agents: %d", len(subAgents))
2246+
break
2247+
}
2248+
}
22342249
require.Len(t, subAgents, 1, "expected one sub agent")
22352250

22362251
subAgent := subAgents[0]

‎agent/agentcontainers/api.go‎

Lines changed: 56 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -977,7 +977,7 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
977977
)
978978

979979
// Check if subagent already exists for this devcontainer.
980-
recreateSubAgent := false
980+
maybeRecreateSubAgent := false
981981
proc, injected := api.injectedSubAgentProcs[dc.WorkspaceFolder]
982982
if injected {
983983
if proc.containerID == container.ID && proc.ctx.Err() == nil {
@@ -992,12 +992,15 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
992992
logger.Debug(ctx, "container ID changed, injecting subagent into new container",
993993
slog.F("old_container_id", proc.containerID),
994994
)
995-
recreateSubAgent = true
995+
maybeRecreateSubAgent = true
996996
}
997997

998998
// Container ID changed or the subagent process is not running,
999999
// stop the existing subagent context to replace it.
10001000
proc.stop()
1001+
} else {
1002+
// Set SubAgent defaults.
1003+
proc.agent.OperatingSystem = "linux" // Assuming Linux for devcontainers.
10011004
}
10021005

10031006
// Prepare the subAgentProcess to be used when running the subagent.
@@ -1090,36 +1093,29 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
10901093
// logger.Warn(ctx, "set CAP_NET_ADMIN on agent binary failed", slog.Error(err))
10911094
// }
10921095

1093-
// Detect workspace folder by executing `pwd` in the container.
1094-
// NOTE(mafredri): This is a quick and dirty way to detect the
1095-
// workspace folder inside the container. In the future we will
1096-
// rely more on `devcontainer read-configuration`.
1097-
var pwdBuf bytes.Buffer
1098-
err = api.dccli.Exec(ctx, dc.WorkspaceFolder, dc.ConfigPath, "pwd", []string{},
1099-
WithExecOutput(&pwdBuf, io.Discard),
1100-
WithExecContainerID(container.ID),
1101-
)
1102-
if err != nil {
1103-
return xerrors.Errorf("check workspace folder in container: %w", err)
1104-
}
1105-
directory := strings.TrimSpace(pwdBuf.String())
1106-
if directory == "" {
1107-
logger.Warn(ctx, "detected workspace folder is empty, using default workspace folder",
1108-
slog.F("default_workspace_folder", DevcontainerDefaultContainerWorkspaceFolder),
1096+
subAgentConfig := proc.agent.CloneConfig(dc)
1097+
if proc.agent.ID == uuid.Nil || maybeRecreateSubAgent {
1098+
// Detect workspace folder by executing `pwd` in the container.
1099+
// NOTE(mafredri): This is a quick and dirty way to detect the
1100+
// workspace folder inside the container. In the future we will
1101+
// rely more on `devcontainer read-configuration`.
1102+
var pwdBuf bytes.Buffer
1103+
err = api.dccli.Exec(ctx, dc.WorkspaceFolder, dc.ConfigPath, "pwd", []string{},
1104+
WithExecOutput(&pwdBuf, io.Discard),
1105+
WithExecContainerID(container.ID),
11091106
)
1110-
directory = DevcontainerDefaultContainerWorkspaceFolder
1111-
}
1112-
1113-
if proc.agent.ID != uuid.Nil && recreateSubAgent {
1114-
logger.Debug(ctx, "deleting existing subagent for recreation", slog.F("agent_id", proc.agent.ID))
1115-
client := *api.subAgentClient.Load()
1116-
err = client.Delete(ctx, proc.agent.ID)
11171107
if err != nil {
1118-
return xerrors.Errorf("delete existing subagent failed: %w", err)
1108+
return xerrors.Errorf("check workspace folder in container: %w", err)
11191109
}
1120-
proc.agent = SubAgent{}
1121-
}
1122-
if proc.agent.ID == uuid.Nil {
1110+
directory := strings.TrimSpace(pwdBuf.String())
1111+
if directory == "" {
1112+
logger.Warn(ctx, "detected workspace folder is empty, using default workspace folder",
1113+
slog.F("default_workspace_folder", DevcontainerDefaultContainerWorkspaceFolder),
1114+
)
1115+
directory = DevcontainerDefaultContainerWorkspaceFolder
1116+
}
1117+
subAgentConfig.Directory = directory
1118+
11231119
displayAppsMap := map[codersdk.DisplayApp]bool{
11241120
// NOTE(DanielleMaywood):
11251121
// We use the same defaults here as set in terraform-provider-coder.
@@ -1138,6 +1134,13 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11381134

11391135
for _, customization := range coderCustomization {
11401136
for app, enabled := range customization.DisplayApps {
1137+
if _, ok := displayAppsMap[app]; !ok {
1138+
logger.Warn(ctx, "unknown display app in devcontainer customization, ignoring",
1139+
slog.F("app", app),
1140+
slog.F("enabled", enabled),
1141+
)
1142+
continue
1143+
}
11411144
displayAppsMap[app] = enabled
11421145
}
11431146
}
@@ -1149,26 +1152,41 @@ func (api *API) maybeInjectSubAgentIntoContainerLocked(ctx context.Context, dc c
11491152
displayApps = append(displayApps, app)
11501153
}
11511154
}
1155+
slices.Sort(displayApps)
11521156

1157+
subAgentConfig.DisplayApps = displayApps
1158+
}
1159+
1160+
deleteSubAgent := maybeRecreateSubAgent && !proc.agent.EqualConfig(subAgentConfig)
1161+
if deleteSubAgent {
1162+
logger.Debug(ctx, "deleting existing subagent for recreation", slog.F("agent_id", proc.agent.ID))
1163+
client := *api.subAgentClient.Load()
1164+
err = client.Delete(ctx, proc.agent.ID)
1165+
if err != nil {
1166+
return xerrors.Errorf("delete existing subagent failed: %w", err)
1167+
}
1168+
proc.agent = SubAgent{} // Clear agent to signal that we need to create a new one.
1169+
}
1170+
1171+
if proc.agent.ID == uuid.Nil {
11531172
logger.Debug(ctx, "creating new subagent",
1154-
slog.F("directory", directory),
1155-
slog.F("display_apps", displayApps),
1173+
slog.F("directory", subAgentConfig.Directory),
1174+
slog.F("display_apps", subAgentConfig.DisplayApps),
11561175
)
11571176

11581177
// Create new subagent record in the database to receive the auth token.
11591178
client := *api.subAgentClient.Load()
1160-
proc.agent, err = client.Create(ctx, SubAgent{
1161-
Name: dc.Name,
1162-
Directory: directory,
1163-
OperatingSystem: "linux", // Assuming Linux for devcontainers.
1164-
Architecture: arch,
1165-
DisplayApps: displayApps,
1166-
})
1179+
newSubAgent, err := client.Create(ctx, subAgentConfig)
11671180
if err != nil {
11681181
return xerrors.Errorf("create subagent failed: %w", err)
11691182
}
1183+
proc.agent = newSubAgent
11701184

11711185
logger.Info(ctx, "created new subagent", slog.F("agent_id", proc.agent.ID))
1186+
} else {
1187+
logger.Debug(ctx, "subagent already exists, skipping recreation",
1188+
slog.F("agent_id", proc.agent.ID),
1189+
)
11721190
}
11731191

11741192
api.mu.Lock() // Re-lock to update the agent.

‎agent/agentcontainers/api_test.go‎

Lines changed: 49 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@ func (w *fakeWatcher) sendEventWaitNextCalled(ctx context.Context, event fsnotif
212212

213213
// fakeSubAgentClient implements SubAgentClient for testing purposes.
214214
type fakeSubAgentClient struct {
215+
logger slog.Logger
215216
agents map[uuid.UUID]agentcontainers.SubAgent
216217

217218
listErrC chan error // If set, send to return error, close to return nil.
@@ -240,6 +241,7 @@ func (m *fakeSubAgentClient) List(ctx context.Context) ([]agentcontainers.SubAge
240241
}
241242

242243
func (m *fakeSubAgentClient) Create(ctx context.Context, agent agentcontainers.SubAgent) (agentcontainers.SubAgent, error) {
244+
m.logger.Debug(ctx, "creating sub agent", slog.F("agent", agent))
243245
if m.createErrC != nil {
244246
select {
245247
case <-ctx.Done():
@@ -261,6 +263,7 @@ func (m *fakeSubAgentClient) Create(ctx context.Context, agent agentcontainers.S
261263
}
262264

263265
func (m *fakeSubAgentClient) Delete(ctx context.Context, id uuid.UUID) error {
266+
m.logger.Debug(ctx, "deleting sub agent", slog.F("id", id.String()))
264267
if m.deleteErrC != nil {
265268
select {
266269
case <-ctx.Done():
@@ -1245,6 +1248,7 @@ func TestAPI(t *testing.T) {
12451248
mClock = quartz.NewMock(t)
12461249
mCCLI = acmock.NewMockContainerCLI(gomock.NewController(t))
12471250
fakeSAC = &fakeSubAgentClient{
1251+
logger: logger.Named("fakeSubAgentClient"),
12481252
createErrC: make(chan error, 1),
12491253
deleteErrC: make(chan error, 1),
12501254
}
@@ -1270,7 +1274,7 @@ func TestAPI(t *testing.T) {
12701274

12711275
mCCLI.EXPECT().List(gomock.Any()).Return(codersdk.WorkspaceAgentListContainersResponse{
12721276
Containers: []codersdk.WorkspaceAgentContainer{testContainer},
1273-
}, nil).Times(1 + 3) // 1 initial call + 3 updates.
1277+
}, nil).Times(3) // 1 initial call + 2 updates.
12741278
gomock.InOrder(
12751279
mCCLI.EXPECT().DetectArchitecture(gomock.Any(), "test-container-id").Return(runtime.GOARCH, nil),
12761280
mCCLI.EXPECT().ExecAs(gomock.Any(), "test-container-id", "root", "mkdir", "-p", "/.coder-agent").Return(nil, nil),
@@ -1315,19 +1319,20 @@ func TestAPI(t *testing.T) {
13151319
tickerTrap.MustWait(ctx).MustRelease(ctx)
13161320
tickerTrap.Close()
13171321

1318-
// Ensure we only inject the agent once.
1319-
for i := range 3 {
1320-
_, aw := mClock.AdvanceNext()
1321-
aw.MustWait(ctx)
1322+
// Refresh twice to ensure idempotency of agent creation.
1323+
err = api.RefreshContainers(ctx)
1324+
require.NoError(t, err, "refresh containers should not fail")
1325+
t.Logf("Agents created: %d, deleted: %d", len(fakeSAC.created), len(fakeSAC.deleted))
13221326

1323-
t.Logf("Iteration %d: agents created: %d", i+1, len(fakeSAC.created))
1327+
err = api.RefreshContainers(ctx)
1328+
require.NoError(t, err, "refresh containers should not fail")
1329+
t.Logf("Agents created: %d, deleted: %d", len(fakeSAC.created), len(fakeSAC.deleted))
13241330

1325-
// Verify agent was created.
1326-
require.Len(t, fakeSAC.created, 1)
1327-
assert.Equal(t, "test-container", fakeSAC.created[0].Name)
1328-
assert.Equal(t, "/workspaces", fakeSAC.created[0].Directory)
1329-
assert.Len(t, fakeSAC.deleted, 0)
1330-
}
1331+
// Verify agent was created.
1332+
require.Len(t, fakeSAC.created, 1)
1333+
assert.Equal(t, "test-container", fakeSAC.created[0].Name)
1334+
assert.Equal(t, "/workspaces", fakeSAC.created[0].Directory)
1335+
assert.Len(t, fakeSAC.deleted, 0)
13311336

13321337
t.Log("Agent injected successfully, now testing reinjection into the same container...")
13331338

@@ -1349,32 +1354,23 @@ func TestAPI(t *testing.T) {
13491354
// Expect the agent to be reinjected.
13501355
mCCLI.EXPECT().List(gomock.Any()).Return(codersdk.WorkspaceAgentListContainersResponse{
13511356
Containers: []codersdk.WorkspaceAgentContainer{testContainer},
1352-
}, nil).Times(3) // 3 updates.
1357+
}, nil).Times(1) // 1 update.
13531358
gomock.InOrder(
13541359
mCCLI.EXPECT().DetectArchitecture(gomock.Any(), "test-container-id").Return(runtime.GOARCH, nil),
13551360
mCCLI.EXPECT().ExecAs(gomock.Any(), "test-container-id", "root", "mkdir", "-p", "/.coder-agent").Return(nil, nil),
13561361
mCCLI.EXPECT().Copy(gomock.Any(), "test-container-id", coderBin, "/.coder-agent/coder").Return(nil),
13571362
mCCLI.EXPECT().ExecAs(gomock.Any(), "test-container-id", "root", "chmod", "0755", "/.coder-agent", "/.coder-agent/coder").Return(nil, nil),
13581363
)
13591364

1360-
// Allow agent reinjection to succeed.
1361-
testutil.RequireSend(ctx, t, fakeDCCLI.execErrC, func(cmd string, args ...string) error {
1362-
assert.Equal(t, "pwd", cmd)
1363-
assert.Empty(t, args)
1364-
return nil
1365-
}) // Exec pwd.
1366-
1367-
// Ensure we only inject the agent once.
1368-
for i := range 3 {
1369-
_, aw := mClock.AdvanceNext()
1370-
aw.MustWait(ctx)
1365+
// Agent reinjection will succeed and we will not re-create the
1366+
// agent, nor re-probe pwd.
1367+
err = api.RefreshContainers(ctx)
1368+
require.NoError(t, err, "refresh containers should not fail")
1369+
t.Logf("Agents created: %d, deleted: %d", len(fakeSAC.created), len(fakeSAC.deleted))
13711370

1372-
t.Logf("Iteration %d: agents created: %d", i+1, len(fakeSAC.created))
1373-
1374-
// Verify that the agent was reused.
1375-
require.Len(t, fakeSAC.created, 1)
1376-
assert.Len(t, fakeSAC.deleted, 0)
1377-
}
1371+
// Verify that the agent was reused.
1372+
require.Len(t, fakeSAC.created, 1)
1373+
assert.Len(t, fakeSAC.deleted, 0)
13781374

13791375
t.Log("Agent reinjected successfully, now testing agent deletion and recreation...")
13801376

@@ -1383,7 +1379,7 @@ func TestAPI(t *testing.T) {
13831379
// Expect the agent to be injected.
13841380
mCCLI.EXPECT().List(gomock.Any()).Return(codersdk.WorkspaceAgentListContainersResponse{
13851381
Containers: []codersdk.WorkspaceAgentContainer{testContainer},
1386-
}, nil).Times(3) // 3 updates.
1382+
}, nil).Times(1) // 1 update.
13871383
gomock.InOrder(
13881384
mCCLI.EXPECT().DetectArchitecture(gomock.Any(), "new-test-container-id").Return(runtime.GOARCH, nil),
13891385
mCCLI.EXPECT().ExecAs(gomock.Any(), "new-test-container-id", "root", "mkdir", "-p", "/.coder-agent").Return(nil, nil),
@@ -1404,7 +1400,20 @@ func TestAPI(t *testing.T) {
14041400
})
14051401
<-terminated
14061402

1407-
// Simulate the agent deletion.
1403+
fakeDCCLI.readConfig.MergedConfiguration.Customizations.Coder = []agentcontainers.CoderCustomization{
1404+
{
1405+
DisplayApps: map[codersdk.DisplayApp]bool{
1406+
codersdk.DisplayAppSSH: true,
1407+
codersdk.DisplayAppWebTerminal: true,
1408+
codersdk.DisplayAppVSCodeDesktop: true,
1409+
codersdk.DisplayAppVSCodeInsiders: true,
1410+
codersdk.DisplayAppPortForward: true,
1411+
},
1412+
},
1413+
}
1414+
1415+
// Simulate the agent deletion (this happens because the
1416+
// devcontainer configuration changed).
14081417
testutil.RequireSend(ctx, t, fakeSAC.deleteErrC, nil)
14091418
// Expect the agent to be recreated.
14101419
testutil.RequireSend(ctx, t, fakeSAC.createErrC, nil)
@@ -1414,13 +1423,9 @@ func TestAPI(t *testing.T) {
14141423
return nil
14151424
}) // Exec pwd.
14161425

1417-
// Advance the clock to run updaterLoop.
1418-
for i := range 3 {
1419-
_, aw := mClock.AdvanceNext()
1420-
aw.MustWait(ctx)
1421-
1422-
t.Logf("Iteration %d: agents created: %d, deleted: %d", i+1, len(fakeSAC.created), len(fakeSAC.deleted))
1423-
}
1426+
err = api.RefreshContainers(ctx)
1427+
require.NoError(t, err, "refresh containers should not fail")
1428+
t.Logf("Agents created: %d, deleted: %d", len(fakeSAC.created), len(fakeSAC.deleted))
14241429

14251430
// Verify the agent was deleted and recreated.
14261431
require.Len(t, fakeSAC.deleted, 1, "there should be one deleted agent after recreation")
@@ -1453,6 +1458,7 @@ func TestAPI(t *testing.T) {
14531458
mClock = quartz.NewMock(t)
14541459
mCCLI = acmock.NewMockContainerCLI(gomock.NewController(t))
14551460
fakeSAC = &fakeSubAgentClient{
1461+
logger: logger.Named("fakeSubAgentClient"),
14561462
agents: map[uuid.UUID]agentcontainers.SubAgent{
14571463
existingAgentID: existingAgent,
14581464
},
@@ -1577,7 +1583,10 @@ func TestAPI(t *testing.T) {
15771583
logger = testutil.Logger(t)
15781584
mClock = quartz.NewMock(t)
15791585
mCCLI = acmock.NewMockContainerCLI(gomock.NewController(t))
1580-
fSAC = &fakeSubAgentClient{createErrC: make(chan error, 1)}
1586+
fSAC = &fakeSubAgentClient{
1587+
logger: logger.Named("fakeSubAgentClient"),
1588+
createErrC: make(chan error, 1),
1589+
}
15811590
fDCCLI = &fakeDevcontainerCLI{
15821591
readConfig: agentcontainers.DevcontainerConfig{
15831592
MergedConfiguration: agentcontainers.DevcontainerConfiguration{

0 commit comments

Comments
 (0)