feat: use agent v2 API to fetch manifest

spikecurtis · spikecurtis · commit 71494d4db61f · 2024-01-26T13:24:01.000+04:00
diff --git a/agent/agent.go b/agent/agent.go
@@ -88,7 +88,6 @@ type Options struct {
 }
 
 type Client interface {
-	Manifest(ctx context.Context) (agentsdk.Manifest, error)
 	Listen(ctx context.Context) (drpc.Conn, error)
 	ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
 	PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
@@ -713,12 +712,16 @@ func (a *agent) run(ctx context.Context) error {
 	serviceBanner := agentsdk.ServiceBannerFromProto(sbp)
 	a.serviceBanner.Store(&serviceBanner)
 
-	manifest, err := a.client.Manifest(ctx)
+	mp, err := aAPI.GetManifest(ctx, &proto.GetManifestRequest{})
 	if err != nil {
 		return xerrors.Errorf("fetch metadata: %w", err)
 	}
-	a.logger.Info(ctx, "fetched manifest", slog.F("manifest", manifest))
-
+	a.logger.Info(ctx, "fetched manifest", slog.F("manifest", mp))
+	manifest, err := agentsdk.ManifestFromProto(mp)
+	if err != nil {
+		a.logger.Critical(ctx, "failed to convert manifest", slog.F("manifest", mp), slog.Error(err))
+		return xerrors.Errorf("convert manifest: %w", err)
+	}
 	if manifest.AgentID == uuid.Nil {
 		return xerrors.New("nil agentID returned by manifest")
 	}
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -2039,6 +2039,9 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
 	if metadata.WorkspaceName == "" {
 		metadata.WorkspaceName = "test-workspace"
 	}
+	if metadata.WorkspaceID == uuid.Nil {
+		metadata.WorkspaceID = uuid.New()
+	}
 	coordinator := tailnet.NewCoordinator(logger)
 	t.Cleanup(func() {
 		_ = coordinator.Close()
diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go
@@ -49,7 +49,9 @@ func NewClient(t testing.TB,
 	}
 	err := proto.DRPCRegisterTailnet(mux, drpcService)
 	require.NoError(t, err)
-	fakeAAPI := NewFakeAgentAPI(t, logger)
+	mp, err := agentsdk.ProtoFromManifest(manifest)
+	require.NoError(t, err)
+	fakeAAPI := NewFakeAgentAPI(t, logger, mp)
 	err = agentproto.DRPCRegisterAgent(mux, fakeAAPI)
 	require.NoError(t, err)
 	server := drpcserver.NewWithOptions(mux, drpcserver.Options{
@@ -64,7 +66,6 @@ func NewClient(t testing.TB,
 		t:              t,
 		logger:         logger.Named("client"),
 		agentID:        agentID,
-		manifest:       manifest,
 		statsChan:      statsChan,
 		coordinator:    coordinator,
 		server:         server,
@@ -77,7 +78,6 @@ type Client struct {
 	t                  testing.TB
 	logger             slog.Logger
 	agentID            uuid.UUID
-	manifest           agentsdk.Manifest
 	metadata           map[string]agentsdk.Metadata
 	statsChan          chan *agentsdk.Stats
 	coordinator        tailnet.Coordinator
@@ -98,10 +98,6 @@ func (c *Client) Close() {
 	c.derpMapOnce.Do(func() { close(c.derpMapUpdates) })
 }
 
-func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
-	return c.manifest, nil
-}
-
 func (c *Client) Listen(ctx context.Context) (drpc.Conn, error) {
 	conn, lis := drpcsdk.MemTransportPipe()
 	c.LastWorkspaceAgent = func() {
@@ -252,12 +248,13 @@ type FakeAgentAPI struct {
 	t      testing.TB
 	logger slog.Logger
 
+	manifest *agentproto.Manifest
+
 	getServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
 }
 
-func (*FakeAgentAPI) GetManifest(context.Context, *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
-	// TODO implement me
-	panic("implement me")
+func (f *FakeAgentAPI) GetManifest(context.Context, *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
+	return f.manifest, nil
 }
 
 func (f *FakeAgentAPI) SetServiceBannerFunc(fn func() (codersdk.ServiceBannerConfig, error)) {
@@ -310,9 +307,10 @@ func (*FakeAgentAPI) BatchCreateLogs(context.Context, *agentproto.BatchCreateLog
 	panic("implement me")
 }
 
-func NewFakeAgentAPI(t testing.TB, logger slog.Logger) *FakeAgentAPI {
+func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest) *FakeAgentAPI {
 	return &FakeAgentAPI{
-		t:      t,
-		logger: logger.Named("FakeAgentAPI"),
+		t:        t,
+		logger:   logger.Named("FakeAgentAPI"),
+		manifest: manifest,
 	}
 }
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -135,28 +135,6 @@ type Script struct {
 	Script string `json:"script"`
 }
 
-// Manifest fetches manifest for the currently authenticated workspace agent.
-func (c *Client) Manifest(ctx context.Context) (Manifest, error) {
-	res, err := c.SDK.Request(ctx, http.MethodGet, "/api/v2/workspaceagents/me/manifest", nil)
-	if err != nil {
-		return Manifest{}, err
-	}
-	defer res.Body.Close()
-	if res.StatusCode != http.StatusOK {
-		return Manifest{}, codersdk.ReadBodyAsError(res)
-	}
-	var agentMeta Manifest
-	err = json.NewDecoder(res.Body).Decode(&agentMeta)
-	if err != nil {
-		return Manifest{}, err
-	}
-	err = c.rewriteDerpMap(agentMeta.DERPMap)
-	if err != nil {
-		return Manifest{}, err
-	}
-	return agentMeta, nil
-}
-
 // rewriteDerpMap rewrites the DERP map to use the access URL of the SDK as the
 // "embedded relay" access URL. The passed derp map is modified in place.
 //
