Skip to content

Commit dc7f7c3

Browse files
mafredrimafredri
committed
use the correct directory for the sub agent
1 parent 757dc85 commit dc7f7c3

File tree

4 files changed

+101
-29
lines changed

4 files changed

+101
-29
lines changed

‎agent/agent_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2238,6 +2238,10 @@ func TestAgent_DevcontainerAutostart(t *testing.T) {
22382238
require.NoError(t, err, "failed to parse sub-agent ID")
22392239
t.Logf("Connecting to sub-agent: %s (ID: %s)", subAgent.Name, subAgentID)
22402240

2241+
gotDir, err := agentClient.GetSubAgentDirectory(subAgentID)
2242+
require.NoError(t, err, "failed to get sub-agent directory")
2243+
require.Equal(t, "/workspaces/mywork", gotDir, "sub-agent directory should match")
2244+
22412245
subAgentToken, err := uuid.FromBytes(subAgent.GetAuthToken())
22422246
require.NoError(t, err, "failed to parse sub-agent token")
22432247

‎agent/agentcontainers/api.go

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package agentcontainers
22

33
import (
4+
"bytes"
45
"context"
56
"errors"
67
"fmt"
8+
"io"
79
"net/http"
810
"os"
911
"path"
@@ -1075,17 +1077,30 @@ func (api *API) injectSubAgentIntoContainerLocked(ctx context.Context, dc coders
10751077
logger.Warn(ctx, "set CAP_NET_ADMIN on agent binary failed", slog.Error(err))
10761078
}
10771079

1080+
// Detect workspace folder by executing `pwd` in the container.
1081+
// NOTE(mafredri): This is a quick and dirty way to detect the
1082+
// workspace folder inside the container. In the future we will
1083+
// rely more on `devcontainer read-configuration`.
1084+
var pwdBuf bytes.Buffer
1085+
err = api.dccli.Exec(ctx, dc.WorkspaceFolder, dc.ConfigPath, "pwd", []string{},
1086+
WithExecOutput(&pwdBuf, io.Discard),
1087+
WithExecContainerID(container.ID),
1088+
)
1089+
if err != nil {
1090+
return xerrors.Errorf("check workspace folder in container: %w", err)
1091+
}
1092+
directory := strings.TrimSpace(pwdBuf.String())
1093+
if directory == "" {
1094+
logger.Warn(ctx, "detected workspace folder is empty, using default workspace folder",
1095+
slog.F("default_workspace_folder", DevcontainerDefaultContainerWorkspaceFolder))
1096+
directory = DevcontainerDefaultContainerWorkspaceFolder
1097+
}
1098+
10781099
// The preparation of the subagent is done, now we can create the
10791100
// subagent record in the database to receive the auth token.
10801101
createdAgent, err := api.subAgentClient.Create(ctx, SubAgent{
1081-
Name: dc.Name,
1082-
// The default workspaceFolder for devcontainers is /workspaces.
1083-
// However, it can be changed by setting {"workspaceFolder": "/src"}
1084-
// in the devcontainer.json. This information is not encoded into
1085-
// the container labels, so we must rely on the values parsed from
1086-
// the devcontainer.json file on disk.
1087-
// TODO(mafredri): Support custom workspace folders in the future.
1088-
Directory: DevcontainerDefaultContainerWorkspaceFolder,
1102+
Name: dc.Name,
1103+
Directory: directory,
10891104
OperatingSystem: "linux", // Assuming Linux for dev containers.
10901105
Architecture: arch,
10911106
})
@@ -1166,9 +1181,9 @@ func (api *API) runSubAgentInContainer(ctx context.Context, dc codersdk.Workspac
11661181
)
11671182
if err != nil && !errors.Is(err, context.Canceled) {
11681183
logger.Error(ctx, "subagent process failed", slog.Error(err))
1184+
} else {
1185+
logger.Info(ctx, "subagent process finished")
11691186
}
1170-
1171-
logger.Info(ctx, "subagent process finished")
11721187
}
11731188

11741189
func (api *API) Close() error {

‎agent/agentcontainers/api_test.go

Lines changed: 47 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ type fakeDevcontainerCLI struct {
6464
upErr error
6565
upErrC chan error // If set, send to return err, close to return upErr.
6666
execErr error
67-
execErrC chan error // If set, send to return err, close to return execErr.
67+
execErrC chan func(cmd string, args ...string) error // If set, send fn to return err, nil or close to return execErr.
6868
}
6969

7070
func (f *fakeDevcontainerCLI) Up(ctx context.Context, _, _ string, _ ...agentcontainers.DevcontainerCLIUpOptions) (string, error) {
@@ -81,14 +81,14 @@ func (f *fakeDevcontainerCLI) Up(ctx context.Context, _, _ string, _ ...agentcon
8181
return f.upID, f.upErr
8282
}
8383

84-
func (f *fakeDevcontainerCLI) Exec(ctx context.Context, _, _ string, _ string, _ []string, _ ...agentcontainers.DevcontainerCLIExecOptions) error {
84+
func (f *fakeDevcontainerCLI) Exec(ctx context.Context, _, _ string, cmd string, args []string, _ ...agentcontainers.DevcontainerCLIExecOptions) error {
8585
if f.execErrC != nil {
8686
select {
8787
case <-ctx.Done():
8888
return ctx.Err()
89-
case err, ok := <-f.execErrC:
90-
if ok {
91-
return err
89+
case fn, ok := <-f.execErrC:
90+
if ok && fn != nil {
91+
return fn(cmd, args...)
9292
}
9393
}
9494
}
@@ -1239,16 +1239,17 @@ func TestAPI(t *testing.T) {
12391239
}
12401240

12411241
var (
1242-
ctx = testutil.Context(t, testutil.WaitMedium)
1243-
logger = slog.Make()
1244-
mClock = quartz.NewMock(t)
1245-
mCCLI = acmock.NewMockContainerCLI(gomock.NewController(t))
1246-
fakeSAC = &fakeSubAgentClient{
1242+
ctx = testutil.Context(t, testutil.WaitMedium)
1243+
errTestTermination = xerrors.New("test termination")
1244+
logger = slogtest.Make(t, &slogtest.Options{IgnoredErrorIs: []error{errTestTermination}}).Leveled(slog.LevelDebug)
1245+
mClock = quartz.NewMock(t)
1246+
mCCLI = acmock.NewMockContainerCLI(gomock.NewController(t))
1247+
fakeSAC = &fakeSubAgentClient{
12471248
createErrC: make(chan error, 1),
12481249
deleteErrC: make(chan error, 1),
12491250
}
12501251
fakeDCCLI = &fakeDevcontainerCLI{
1251-
execErrC: make(chan error, 1),
1252+
execErrC: make(chan func(cmd string, args ...string) error, 1),
12521253
}
12531254

12541255
testContainer = codersdk.WorkspaceAgentContainer{
@@ -1264,8 +1265,6 @@ func TestAPI(t *testing.T) {
12641265
}
12651266
)
12661267

1267-
defer close(fakeDCCLI.execErrC)
1268-
12691268
coderBin, err := os.Executable()
12701269
require.NoError(t, err)
12711270

@@ -1287,23 +1286,31 @@ func TestAPI(t *testing.T) {
12871286
api := agentcontainers.NewAPI(logger,
12881287
agentcontainers.WithClock(mClock),
12891288
agentcontainers.WithContainerCLI(mCCLI),
1289+
agentcontainers.WithWatcher(watcher.NewNoop()),
12901290
agentcontainers.WithSubAgentClient(fakeSAC),
12911291
agentcontainers.WithSubAgentURL("test-subagent-url"),
12921292
agentcontainers.WithDevcontainerCLI(fakeDCCLI),
12931293
)
12941294
defer api.Close()
12951295

1296-
// Allow initial agent creation to succeed.
1296+
// Close before api.Close() defer to avoid deadlock after test.
1297+
defer close(fakeSAC.createErrC)
1298+
defer close(fakeSAC.deleteErrC)
1299+
defer close(fakeDCCLI.execErrC)
1300+
1301+
// Allow initial agent creation and injection to succeed.
12971302
testutil.RequireSend(ctx, t, fakeSAC.createErrC, nil)
1303+
testutil.RequireSend(ctx, t, fakeDCCLI.execErrC, func(cmd string, args ...string) error {
1304+
assert.Equal(t, "pwd", cmd)
1305+
assert.Empty(t, args)
1306+
return nil
1307+
}) // Exec pwd.
12981308

12991309
// Make sure the ticker function has been registered
13001310
// before advancing the clock.
13011311
tickerTrap.MustWait(ctx).MustRelease(ctx)
13021312
tickerTrap.Close()
13031313

1304-
// Pre-populate for next iteration (also verify previous consumption).
1305-
testutil.RequireSend(ctx, t, fakeSAC.createErrC, nil)
1306-
13071314
// Ensure we only inject the agent once.
13081315
for i := range 3 {
13091316
_, aw := mClock.AdvanceNext()
@@ -1318,6 +1325,8 @@ func TestAPI(t *testing.T) {
13181325
assert.Len(t, fakeSAC.deleted, 0)
13191326
}
13201327

1328+
t.Log("Agent injected successfully, now testing cleanup and reinjection...")
1329+
13211330
// Expect the agent to be reinjected.
13221331
gomock.InOrder(
13231332
mCCLI.EXPECT().DetectArchitecture(gomock.Any(), "test-container-id").Return(runtime.GOARCH, nil),
@@ -1329,17 +1338,36 @@ func TestAPI(t *testing.T) {
13291338
)
13301339

13311340
// Terminate the agent and verify it is deleted.
1332-
testutil.RequireSend(ctx, t, fakeDCCLI.execErrC, nil)
1341+
testutil.RequireSend(ctx, t, fakeDCCLI.execErrC, func(_ string, args ...string) error {
1342+
if len(args) > 0 {
1343+
assert.Equal(t, "agent", args[0])
1344+
} else {
1345+
assert.Fail(t, `want "agent" command argument`)
1346+
}
1347+
return errTestTermination
1348+
})
13331349

13341350
// Allow cleanup to proceed.
1335-
testutil.RequireSend(ctx, t, fakeSAC.deleteErrC, xerrors.New("test termination"))
1351+
testutil.RequireSend(ctx, t, fakeSAC.deleteErrC, nil)
1352+
1353+
t.Log("Waiting for agent recreation...")
1354+
1355+
// Allow agent recreation and reinjection to succeed.
1356+
testutil.RequireSend(ctx, t, fakeSAC.createErrC, nil)
1357+
testutil.RequireSend(ctx, t, fakeDCCLI.execErrC, func(cmd string, args ...string) error {
1358+
assert.Equal(t, "pwd", cmd)
1359+
assert.Empty(t, args)
1360+
return nil
1361+
}) // Exec pwd.
13361362

13371363
// Wait until the agent recreation is started.
13381364
for len(fakeSAC.createErrC) > 0 {
13391365
_, aw := mClock.AdvanceNext()
13401366
aw.MustWait(ctx)
13411367
}
13421368

1369+
t.Log("Agent recreated successfully.")
1370+
13431371
// Verify agent was deleted.
13441372
require.Len(t, fakeSAC.deleted, 1)
13451373
assert.Equal(t, fakeSAC.created[0].ID, fakeSAC.deleted[0])

‎agent/agenttest/client.go

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,10 @@ func (c *Client) GetSubAgents() []*agentproto.SubAgent {
167167
return c.fakeAgentAPI.GetSubAgents()
168168
}
169169

170+
func (c *Client) GetSubAgentDirectory(id uuid.UUID) (string, error) {
171+
return c.fakeAgentAPI.GetSubAgentDirectory(id)
172+
}
173+
170174
type FakeAgentAPI struct {
171175
sync.Mutex
172176
t testing.TB
@@ -182,6 +186,7 @@ type FakeAgentAPI struct {
182186
timings []*agentproto.Timing
183187
connectionReports []*agentproto.ReportConnectionRequest
184188
subAgents map[uuid.UUID]*agentproto.SubAgent
189+
subAgentDirs map[uuid.UUID]string
185190

186191
getAnnouncementBannersFunc func() ([]codersdk.BannerConfig, error)
187192
getResourcesMonitoringConfigurationFunc func() (*agentproto.GetResourcesMonitoringConfigurationResponse, error)
@@ -392,6 +397,10 @@ func (f *FakeAgentAPI) CreateSubAgent(ctx context.Context, req *agentproto.Creat
392397
f.subAgents = make(map[uuid.UUID]*agentproto.SubAgent)
393398
}
394399
f.subAgents[subAgentID] = subAgent
400+
if f.subAgentDirs == nil {
401+
f.subAgentDirs = make(map[uuid.UUID]string)
402+
}
403+
f.subAgentDirs[subAgentID] = req.GetDirectory()
395404

396405
// For a fake implementation, we don't create workspace apps.
397406
// Real implementations would handle req.Apps here.
@@ -452,6 +461,22 @@ func (f *FakeAgentAPI) GetSubAgents() []*agentproto.SubAgent {
452461
return agents
453462
}
454463

464+
func (f *FakeAgentAPI) GetSubAgentDirectory(id uuid.UUID) (string, error) {
465+
f.Lock()
466+
defer f.Unlock()
467+
468+
if f.subAgentDirs == nil {
469+
return "", xerrors.New("no sub-agent directories available")
470+
}
471+
472+
dir, ok := f.subAgentDirs[id]
473+
if !ok {
474+
return "", xerrors.New("sub-agent directory not found")
475+
}
476+
477+
return dir, nil
478+
}
479+
455480
func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest, statsCh chan *agentproto.Stats) *FakeAgentAPI {
456481
return &FakeAgentAPI{
457482
t: t,

0 commit comments

Comments
 (0)