Skip to content

Commit dc5c132

Browse files
spikecurtisspikecurtis
committed
feat: use agent v2 API to post startup
1 parent ee908ab commit dc5c132

File tree

11 files changed

+102
-71
lines changed

11 files changed

+102
-71
lines changed

‎agent/agent.go

Lines changed: 9 additions & 5 deletions
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

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1395,56 +1395,52 @@ func TestAgent_Startup(t *testing.T) {
13951395

13961396
t.Run("EmptyDirectory", func(t *testing.T) {
13971397
t.Parallel()
1398+
ctx := testutil.Context(t, testutil.WaitShort)
13981399

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

14081407
t.Run("HomeDirectory", func(t *testing.T) {
14091408
t.Parallel()
1409+
ctx := testutil.Context(t, testutil.WaitShort)
14101410

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

14221420
t.Run("NotAbsoluteDirectory", func(t *testing.T) {
14231421
t.Parallel()
1422+
ctx := testutil.Context(t, testutil.WaitShort)
14241423

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

14361433
t.Run("HomeEnvironmentVariable", func(t *testing.T) {
14371434
t.Parallel()
1435+
ctx := testutil.Context(t, testutil.WaitShort)
14381436

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

‎agent/agenttest/client.go

Lines changed: 11 additions & 20 deletions
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
}

‎coderd/agentapi/api.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,6 @@ import (
2929
tailnetproto "github.com/coder/coder/v2/tailnet/proto"
3030
)
3131

32-
const AgentAPIVersionDRPC = "2.0"
33-
3432
// API implements the DRPC agent API interface from agent/proto. This struct is
3533
// instantiated once per agent connection and kept alive for the duration of the
3634
// session.

‎coderd/agentapi/lifecycle.go

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"time"
77

88
"github.com/google/uuid"
9+
"golang.org/x/exp/slices"
910
"golang.org/x/mod/semver"
1011
"golang.org/x/xerrors"
1112
"google.golang.org/protobuf/types/known/timestamppb"
@@ -16,6 +17,12 @@ import (
1617
"github.com/coder/coder/v2/coderd/database/dbtime"
1718
)
1819

20+
type contextKeyAPIVersion struct{}
21+
22+
func WithAPIVersion(ctx context.Context, version string) context.Context {
23+
return context.WithValue(ctx, contextKeyAPIVersion{}, version)
24+
}
25+
1926
type LifecycleAPI struct {
2027
AgentFn func(context.Context) (database.WorkspaceAgent, error)
2128
WorkspaceIDFn func(context.Context, *database.WorkspaceAgent) (uuid.UUID, error)
@@ -123,6 +130,10 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
123130
}
124131

125132
func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {
133+
apiVersion, ok := ctx.Value(contextKeyAPIVersion{}).(string)
134+
if !ok {
135+
return nil, xerrors.Errorf("internal error; api version unspecified")
136+
}
126137
workspaceAgent, err := a.AgentFn(ctx)
127138
if err != nil {
128139
return nil, err
@@ -164,13 +175,14 @@ func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.Update
164175
dbSubsystems = append(dbSubsystems, dbSubsystem)
165176
}
166177
}
178+
slices.Sort(dbSubsystems)
167179

168180
err = a.Database.UpdateWorkspaceAgentStartupByID(ctx, database.UpdateWorkspaceAgentStartupByIDParams{
169181
ID: workspaceAgent.ID,
170182
Version: req.Startup.Version,
171183
ExpandedDirectory: req.Startup.ExpandedDirectory,
172184
Subsystems: dbSubsystems,
173-
APIVersion: AgentAPIVersionDRPC,
185+
APIVersion: apiVersion,
174186
})
175187
if err != nil {
176188
return nil, xerrors.Errorf("update workspace agent startup in database: %w", err)

‎coderd/agentapi/lifecycle_test.go

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -382,10 +382,11 @@ func TestUpdateStartup(t *testing.T) {
382382
database.WorkspaceAgentSubsystemEnvbuilder,
383383
database.WorkspaceAgentSubsystemExectrace,
384384
},
385-
APIVersion: agentapi.AgentAPIVersionDRPC,
385+
APIVersion: "2.0",
386386
}).Return(nil)
387387

388-
resp, err := api.UpdateStartup(context.Background(), &agentproto.UpdateStartupRequest{
388+
ctx := agentapi.WithAPIVersion(context.Background(), "2.0")
389+
resp, err := api.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{
389390
Startup: startup,
390391
})
391392
require.NoError(t, err)
@@ -416,7 +417,8 @@ func TestUpdateStartup(t *testing.T) {
416417
Subsystems: []agentproto.Startup_Subsystem{},
417418
}
418419

419-
resp, err := api.UpdateStartup(context.Background(), &agentproto.UpdateStartupRequest{
420+
ctx := agentapi.WithAPIVersion(context.Background(), "2.0")
421+
resp, err := api.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{
420422
Startup: startup,
421423
})
422424
require.Error(t, err)
@@ -451,7 +453,8 @@ func TestUpdateStartup(t *testing.T) {
451453
},
452454
}
453455

454-
resp, err := api.UpdateStartup(context.Background(), &agentproto.UpdateStartupRequest{
456+
ctx := agentapi.WithAPIVersion(context.Background(), "2.0")
457+
resp, err := api.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{
455458
Startup: startup,
456459
})
457460
require.Error(t, err)

‎coderd/workspaceagents_test.go

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import (
2424
"github.com/coder/coder/v2/agent"
2525
"github.com/coder/coder/v2/agent/agenttest"
2626
agentproto "github.com/coder/coder/v2/agent/proto"
27-
"github.com/coder/coder/v2/coderd"
2827
"github.com/coder/coder/v2/coderd/coderdtest"
2928
"github.com/coder/coder/v2/coderd/coderdtest/oidctest"
3029
"github.com/coder/coder/v2/coderd/database"
@@ -1389,13 +1388,13 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
13891388
}
13901389
)
13911390

1392-
err := agentClient.PostStartup(ctx, agentsdk.PostStartupRequest{
1391+
err := postStartup(ctx, t, agentClient, &agentproto.Startup{
13931392
Version: expectedVersion,
13941393
ExpandedDirectory: expectedDir,
1395-
Subsystems: []codersdk.AgentSubsystem{
1394+
Subsystems: []agentproto.Startup_Subsystem{
13961395
// Not sorted.
1397-
expectedSubsystems[1],
1398-
expectedSubsystems[0],
1396+
agentproto.Startup_EXECTRACE,
1397+
agentproto.Startup_ENVBOX,
13991398
},
14001399
})
14011400
require.NoError(t, err)
@@ -1409,7 +1408,7 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
14091408
require.Equal(t, expectedDir, wsagent.ExpandedDirectory)
14101409
// Sorted
14111410
require.Equal(t, expectedSubsystems, wsagent.Subsystems)
1412-
require.Equal(t, coderd.AgentAPIVersionREST, wsagent.APIVersion)
1411+
require.Equal(t, agentproto.CurrentVersion.String(), wsagent.APIVersion)
14131412
})
14141413

14151414
t.Run("InvalidSemver", func(t *testing.T) {
@@ -1427,13 +1426,10 @@ func TestWorkspaceAgent_Startup(t *testing.T) {
14271426

14281427
ctx := testutil.Context(t, testutil.WaitMedium)
14291428

1430-
err := agentClient.PostStartup(ctx, agentsdk.PostStartupRequest{
1429+
err := postStartup(ctx, t, agentClient, &agentproto.Startup{
14311430
Version: "1.2.3",
14321431
})
1433-
require.Error(t, err)
1434-
cerr, ok := codersdk.AsError(err)
1435-
require.True(t, ok)
1436-
require.Equal(t, http.StatusBadRequest, cerr.StatusCode())
1432+
require.ErrorContains(t, err, "invalid agent semver version")
14371433
})
14381434
}
14391435

@@ -1640,3 +1636,15 @@ func requireGetManifest(ctx context.Context, t testing.TB, client agent.Client)
16401636
require.NoError(t, err)
16411637
return manifest
16421638
}
1639+
1640+
func postStartup(ctx context.Context, t testing.TB, client agent.Client, startup *agentproto.Startup) error {
1641+
conn, err := client.Listen(ctx)
1642+
require.NoError(t, err)
1643+
defer func() {
1644+
cErr := conn.Close()
1645+
require.NoError(t, cErr)
1646+
}()
1647+
aAPI := agentproto.NewDRPCAgentClient(conn)
1648+
_, err = aAPI.UpdateStartup(ctx, &agentproto.UpdateStartupRequest{Startup: startup})
1649+
return err
1650+
}

‎coderd/workspaceagentsrpc.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
154154
Auth: tailnet.AgentTunnelAuth{},
155155
}
156156
ctx = tailnet.WithStreamID(ctx, streamID)
157+
ctx = agentapi.WithAPIVersion(ctx, version)
157158
err = agentAPI.Serve(ctx, mux)
158159
if err != nil {
159160
api.Logger.Warn(ctx, "workspace agent RPC listen error", slog.Error(err))

‎codersdk/agentsdk/agentsdk.go

Lines changed: 0 additions & 12 deletions
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

Lines changed: 12 additions & 0 deletions
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+
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy