Skip to content

Commit 36d34da

Browse files
spikecurtisspikecurtis
committed
feat: use agent v2 API to post startup
1 parent 8504a42 commit 36d34da

File tree

6 files changed

+61
-53
lines changed

6 files changed

+61
-53
lines changed

‎agent/agent.go

+9-5
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,6 @@ type Client interface {
9292
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
9393
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
9494
PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
95-
PostStartup(ctx context.Context, req agentsdk.PostStartupRequest) error
9695
PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequest) error
9796
PatchLogs(ctx context.Context, req agentsdk.PatchLogs) error
9897
RewriteDERPMap(derpMap *tailcfg.DERPMap)
@@ -737,13 +736,18 @@ func (a *agent) run(ctx context.Context) error {
737736
if err != nil {
738737
return xerrors.Errorf("expand directory: %w", err)
739738
}
740-
err = a.client.PostStartup(ctx, agentsdk.PostStartupRequest{
739+
subsys, err := agentsdk.ProtoFromSubsystems(a.subsystems)
740+
if err != nil {
741+
a.logger.Critical(ctx, "failed to convert subsystems", slog.Error(err))
742+
return xerrors.Errorf("failed to convert subsystems: %w", err)
743+
}
744+
_, err = aAPI.UpdateStartup(ctx, &proto.UpdateStartupRequest{Startup: &proto.Startup{
741745
Version: buildinfo.Version(),
742746
ExpandedDirectory: manifest.Directory,
743-
Subsystems: a.subsystems,
744-
})
747+
Subsystems: subsys,
748+
}})
745749
if err != nil {
746-
return xerrors.Errorf("update workspace agent version: %w", err)
750+
return xerrors.Errorf("update workspace agent startup: %w", err)
747751
}
748752

749753
oldManifest := a.manifest.Swap(&manifest)

‎agent/agent_test.go

+12-16
Original file line numberDiff line numberDiff line change
@@ -1394,56 +1394,52 @@ func TestAgent_Startup(t *testing.T) {
13941394

13951395
t.Run("EmptyDirectory", func(t *testing.T) {
13961396
t.Parallel()
1397+
ctx := testutil.Context(t, testutil.WaitShort)
13971398

13981399
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
13991400
Directory: "",
14001401
}, 0)
1401-
assert.Eventually(t, func() bool {
1402-
return client.GetStartup().Version != ""
1403-
}, testutil.WaitShort, testutil.IntervalFast)
1404-
require.Equal(t, "", client.GetStartup().ExpandedDirectory)
1402+
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
1403+
require.Equal(t, "", startup.GetExpandedDirectory())
14051404
})
14061405

14071406
t.Run("HomeDirectory", func(t *testing.T) {
14081407
t.Parallel()
1408+
ctx := testutil.Context(t, testutil.WaitShort)
14091409

14101410
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
14111411
Directory: "~",
14121412
}, 0)
1413-
assert.Eventually(t, func() bool {
1414-
return client.GetStartup().Version != ""
1415-
}, testutil.WaitShort, testutil.IntervalFast)
1413+
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
14161414
homeDir, err := os.UserHomeDir()
14171415
require.NoError(t, err)
1418-
require.Equal(t, homeDir, client.GetStartup().ExpandedDirectory)
1416+
require.Equal(t, homeDir, startup.GetExpandedDirectory())
14191417
})
14201418

14211419
t.Run("NotAbsoluteDirectory", func(t *testing.T) {
14221420
t.Parallel()
1421+
ctx := testutil.Context(t, testutil.WaitShort)
14231422

14241423
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
14251424
Directory: "coder/coder",
14261425
}, 0)
1427-
assert.Eventually(t, func() bool {
1428-
return client.GetStartup().Version != ""
1429-
}, testutil.WaitShort, testutil.IntervalFast)
1426+
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
14301427
homeDir, err := os.UserHomeDir()
14311428
require.NoError(t, err)
1432-
require.Equal(t, filepath.Join(homeDir, "coder/coder"), client.GetStartup().ExpandedDirectory)
1429+
require.Equal(t, filepath.Join(homeDir, "coder/coder"), startup.GetExpandedDirectory())
14331430
})
14341431

14351432
t.Run("HomeEnvironmentVariable", func(t *testing.T) {
14361433
t.Parallel()
1434+
ctx := testutil.Context(t, testutil.WaitShort)
14371435

14381436
_, client, _, _, _ := setupAgent(t, agentsdk.Manifest{
14391437
Directory: "$HOME",
14401438
}, 0)
1441-
assert.Eventually(t, func() bool {
1442-
return client.GetStartup().Version != ""
1443-
}, testutil.WaitShort, testutil.IntervalFast)
1439+
startup := testutil.RequireRecvCtx(ctx, t, client.GetStartup())
14441440
homeDir, err := os.UserHomeDir()
14451441
require.NoError(t, err)
1446-
require.Equal(t, homeDir, client.GetStartup().ExpandedDirectory)
1442+
require.Equal(t, homeDir, startup.GetExpandedDirectory())
14471443
})
14481444
}
14491445

‎agent/agenttest/client.go

+11-20
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,6 @@ type Client struct {
8888

8989
mu sync.Mutex // Protects following.
9090
lifecycleStates []codersdk.WorkspaceAgentLifecycle
91-
startup agentsdk.PostStartupRequest
9291
logs []agentsdk.Log
9392
derpMapUpdates chan *tailcfg.DERPMap
9493
derpMapOnce sync.Once
@@ -173,10 +172,8 @@ func (c *Client) PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsR
173172
return nil
174173
}
175174

176-
func (c *Client) GetStartup() agentsdk.PostStartupRequest {
177-
c.mu.Lock()
178-
defer c.mu.Unlock()
179-
return c.startup
175+
func (c *Client) GetStartup() <-chan *agentproto.Startup {
176+
return c.fakeAgentAPI.startupCh
180177
}
181178

182179
func (c *Client) GetMetadata() map[string]agentsdk.Metadata {
@@ -198,14 +195,6 @@ func (c *Client) PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequ
198195
return nil
199196
}
200197

201-
func (c *Client) PostStartup(ctx context.Context, startup agentsdk.PostStartupRequest) error {
202-
c.mu.Lock()
203-
defer c.mu.Unlock()
204-
c.startup = startup
205-
c.logger.Debug(ctx, "post startup", slog.F("req", startup))
206-
return nil
207-
}
208-
209198
func (c *Client) GetStartupLogs() []agentsdk.Log {
210199
c.mu.Lock()
211200
defer c.mu.Unlock()
@@ -250,7 +239,8 @@ type FakeAgentAPI struct {
250239
t testing.TB
251240
logger slog.Logger
252241

253-
manifest *agentproto.Manifest
242+
manifest *agentproto.Manifest
243+
startupCh chan *agentproto.Startup
254244

255245
getServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
256246
}
@@ -294,9 +284,9 @@ func (*FakeAgentAPI) BatchUpdateAppHealths(context.Context, *agentproto.BatchUpd
294284
panic("implement me")
295285
}
296286

297-
func (*FakeAgentAPI) UpdateStartup(context.Context, *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {
298-
// TODO implement me
299-
panic("implement me")
287+
func (f *FakeAgentAPI) UpdateStartup(_ context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {
288+
f.startupCh <- req.GetStartup()
289+
return req.GetStartup(), nil
300290
}
301291

302292
func (*FakeAgentAPI) BatchUpdateMetadata(context.Context, *agentproto.BatchUpdateMetadataRequest) (*agentproto.BatchUpdateMetadataResponse, error) {
@@ -311,8 +301,9 @@ func (*FakeAgentAPI) BatchCreateLogs(context.Context, *agentproto.BatchCreateLog
311301

312302
func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest) *FakeAgentAPI {
313303
return &FakeAgentAPI{
314-
t: t,
315-
logger: logger.Named("FakeAgentAPI"),
316-
manifest: manifest,
304+
t: t,
305+
logger: logger.Named("FakeAgentAPI"),
306+
manifest: manifest,
307+
startupCh: make(chan *agentproto.Startup, 100),
317308
}
318309
}

‎codersdk/agentsdk/agentsdk.go

-12
Original file line numberDiff line numberDiff line change
@@ -556,18 +556,6 @@ type PostStartupRequest struct {
556556
Subsystems []codersdk.AgentSubsystem `json:"subsystems"`
557557
}
558558

559-
func (c *Client) PostStartup(ctx context.Context, req PostStartupRequest) error {
560-
res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/startup", req)
561-
if err != nil {
562-
return err
563-
}
564-
defer res.Body.Close()
565-
if res.StatusCode != http.StatusOK {
566-
return codersdk.ReadBodyAsError(res)
567-
}
568-
return nil
569-
}
570-
571559
type Log struct {
572560
CreatedAt time.Time `json:"created_at"`
573561
Output string `json:"output"`

‎codersdk/agentsdk/convert.go

+12
Original file line numberDiff line numberDiff line change
@@ -266,3 +266,15 @@ func ProtoFromServiceBanner(sb codersdk.ServiceBannerConfig) *proto.ServiceBanne
266266
BackgroundColor: sb.BackgroundColor,
267267
}
268268
}
269+
270+
func ProtoFromSubsystems(ss []codersdk.AgentSubsystem) ([]proto.Startup_Subsystem, error) {
271+
ret := make([]proto.Startup_Subsystem, len(ss))
272+
for i, s := range ss {
273+
pi, ok := proto.Startup_Subsystem_value[strings.ToUpper(string(s))]
274+
if !ok {
275+
return nil, xerrors.Errorf("unknown subsystem: %s", s)
276+
}
277+
ret[i] = proto.Startup_Subsystem(pi)
278+
}
279+
return ret, nil
280+
}

‎codersdk/agentsdk/convert_test.go

+17
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import (
44
"testing"
55
"time"
66

7+
"github.com/coder/coder/v2/agent/proto"
8+
79
"github.com/google/uuid"
810
"github.com/stretchr/testify/require"
911
"tailscale.com/tailcfg"
@@ -144,3 +146,18 @@ func TestManifest(t *testing.T) {
144146
require.Equal(t, manifest.Metadata, back.Metadata)
145147
require.Equal(t, manifest.Scripts, back.Scripts)
146148
}
149+
150+
func TestSubsystems(t *testing.T) {
151+
ss := []codersdk.AgentSubsystem{
152+
codersdk.AgentSubsystemEnvbox,
153+
codersdk.AgentSubsystemEnvbuilder,
154+
codersdk.AgentSubsystemExectrace,
155+
}
156+
ps, err := agentsdk.ProtoFromSubsystems(ss)
157+
require.NoError(t, err)
158+
require.Equal(t, ps, []proto.Startup_Subsystem{
159+
proto.Startup_ENVBOX,
160+
proto.Startup_ENVBUILDER,
161+
proto.Startup_EXECTRACE,
162+
})
163+
}

0 commit comments

Comments
 (0)