Skip to content

Commit 6ae3e46

Browse files
spikecurtisspikecurtis
committed
feat: use agent v2 API to fetch manifest
1 parent 0eff646 commit 6ae3e46

8 files changed

Lines changed: 91 additions & 83 deletions

File tree

‎agent/agent.go‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,14 +88,14 @@ type Options struct {
8888
}
8989

9090
type Client interface {
91-
Manifest(ctx context.Context) (agentsdk.Manifest, error)
9291
Listen(ctx context.Context) (drpc.Conn, error)
9392
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
9493
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
9594
PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
9695
PostStartup(ctx context.Context, req agentsdk.PostStartupRequest) error
9796
PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequest) error
9897
PatchLogs(ctx context.Context, req agentsdk.PatchLogs) error
98+
RewriteDERPMap(derpMap *tailcfg.DERPMap)
9999
}
100100

101101
type Agent interface {
@@ -713,15 +713,20 @@ func (a *agent) run(ctx context.Context) error {
713713
serviceBanner := agentsdk.ServiceBannerFromProto(sbp)
714714
a.serviceBanner.Store(&serviceBanner)
715715

716-
manifest, err := a.client.Manifest(ctx)
716+
mp, err := aAPI.GetManifest(ctx, &proto.GetManifestRequest{})
717717
if err != nil {
718718
return xerrors.Errorf("fetch metadata: %w", err)
719719
}
720-
a.logger.Info(ctx, "fetched manifest", slog.F("manifest", manifest))
721-
720+
a.logger.Info(ctx, "fetched manifest", slog.F("manifest", mp))
721+
manifest, err := agentsdk.ManifestFromProto(mp)
722+
if err != nil {
723+
a.logger.Critical(ctx, "failed to convert manifest", slog.F("manifest", mp), slog.Error(err))
724+
return xerrors.Errorf("convert manifest: %w", err)
725+
}
722726
if manifest.AgentID == uuid.Nil {
723727
return xerrors.New("nil agentID returned by manifest")
724728
}
729+
a.client.RewriteDERPMap(manifest.DERPMap)
725730

726731
// Expand the directory and send it back to coderd so external
727732
// applications that rely on the directory can use it.
@@ -1124,6 +1129,7 @@ func (a *agent) runDERPMapSubscriber(ctx context.Context, conn drpc.Conn, networ
11241129
return xerrors.Errorf("recv DERPMap error: %w", err)
11251130
}
11261131
dm := tailnet.DERPMapFromProto(dmp)
1132+
a.client.RewriteDERPMap(dm)
11271133
network.SetDERPMap(dm)
11281134
}
11291135
}

‎agent/agent_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2043,6 +2043,9 @@ func setupAgent(t *testing.T, metadata agentsdk.Manifest, ptyTimeout time.Durati
20432043
if metadata.WorkspaceName == "" {
20442044
metadata.WorkspaceName = "test-workspace"
20452045
}
2046+
if metadata.WorkspaceID == uuid.Nil {
2047+
metadata.WorkspaceID = uuid.New()
2048+
}
20462049
coordinator := tailnet.NewCoordinator(logger)
20472050
t.Cleanup(func() {
20482051
_ = coordinator.Close()

‎agent/agenttest/client.go‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,9 @@ func NewClient(t testing.TB,
4949
}
5050
err := proto.DRPCRegisterTailnet(mux, drpcService)
5151
require.NoError(t, err)
52-
fakeAAPI := NewFakeAgentAPI(t, logger)
52+
mp, err := agentsdk.ProtoFromManifest(manifest)
53+
require.NoError(t, err)
54+
fakeAAPI := NewFakeAgentAPI(t, logger, mp)
5355
err = agentproto.DRPCRegisterAgent(mux, fakeAAPI)
5456
require.NoError(t, err)
5557
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
@@ -64,7 +66,6 @@ func NewClient(t testing.TB,
6466
t: t,
6567
logger: logger.Named("client"),
6668
agentID: agentID,
67-
manifest: manifest,
6869
statsChan: statsChan,
6970
coordinator: coordinator,
7071
server: server,
@@ -77,7 +78,6 @@ type Client struct {
7778
t testing.TB
7879
logger slog.Logger
7980
agentID uuid.UUID
80-
manifest agentsdk.Manifest
8181
metadata map[string]agentsdk.Metadata
8282
statsChan chan *agentsdk.Stats
8383
coordinator tailnet.Coordinator
@@ -94,14 +94,12 @@ type Client struct {
9494
derpMapOnce sync.Once
9595
}
9696

97+
func (*Client) RewriteDERPMap(*tailcfg.DERPMap) {}
98+
9799
func (c *Client) Close() {
98100
c.derpMapOnce.Do(func() { close(c.derpMapUpdates) })
99101
}
100102

101-
func (c *Client) Manifest(_ context.Context) (agentsdk.Manifest, error) {
102-
return c.manifest, nil
103-
}
104-
105103
func (c *Client) Listen(ctx context.Context) (drpc.Conn, error) {
106104
conn, lis := drpcsdk.MemTransportPipe()
107105
c.LastWorkspaceAgent = func() {
@@ -252,12 +250,13 @@ type FakeAgentAPI struct {
252250
t testing.TB
253251
logger slog.Logger
254252

253+
manifest *agentproto.Manifest
254+
255255
getServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
256256
}
257257

258-
func (*FakeAgentAPI) GetManifest(context.Context, *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
259-
// TODO implement me
260-
panic("implement me")
258+
func (f *FakeAgentAPI) GetManifest(context.Context, *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
259+
return f.manifest, nil
261260
}
262261

263262
func (f *FakeAgentAPI) SetServiceBannerFunc(fn func() (codersdk.ServiceBannerConfig, error)) {
@@ -310,9 +309,10 @@ func (*FakeAgentAPI) BatchCreateLogs(context.Context, *agentproto.BatchCreateLog
310309
panic("implement me")
311310
}
312311

313-
func NewFakeAgentAPI(t testing.TB, logger slog.Logger) *FakeAgentAPI {
312+
func NewFakeAgentAPI(t testing.TB, logger slog.Logger, manifest *agentproto.Manifest) *FakeAgentAPI {
314313
return &FakeAgentAPI{
315-
t: t,
316-
logger: logger.Named("FakeAgentAPI"),
314+
t: t,
315+
logger: logger.Named("FakeAgentAPI"),
316+
manifest: manifest,
317317
}
318318
}

‎coderd/agentapi/api.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,16 @@ func New(opts Options) *API {
9797
AgentFn: api.agent,
9898
Database: opts.Database,
9999
DerpMapFn: opts.DerpMapFn,
100+
WorkspaceIDFn: func(ctx context.Context, wa *database.WorkspaceAgent) (uuid.UUID, error) {
101+
if opts.WorkspaceID != uuid.Nil {
102+
return opts.WorkspaceID, nil
103+
}
104+
ws, err := opts.Database.GetWorkspaceByAgentID(ctx, wa.ID)
105+
if err != nil {
106+
return uuid.Nil, err
107+
}
108+
return ws.Workspace.ID, nil
109+
},
100110
}
101111

102112
api.ServiceBannerAPI = &ServiceBannerAPI{

‎coderd/workspaceagents_test.go‎

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
"cdr.dev/slog/sloggers/slogtest"
2424
"github.com/coder/coder/v2/agent"
2525
"github.com/coder/coder/v2/agent/agenttest"
26+
agentproto "github.com/coder/coder/v2/agent/proto"
2627
"github.com/coder/coder/v2/coderd"
2728
"github.com/coder/coder/v2/coderd/coderdtest"
2829
"github.com/coder/coder/v2/coderd/coderdtest/oidctest"
@@ -500,8 +501,7 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
500501
// Verify that the manifest has DisableDirectConnections set to true.
501502
agentClient := agentsdk.New(client.URL)
502503
agentClient.SetSessionToken(r.AgentToken)
503-
manifest, err := agentClient.Manifest(ctx)
504-
require.NoError(t, err)
504+
manifest := requireGetManifest(ctx, t, agentClient)
505505
require.True(t, manifest.DisableDirectConnections)
506506

507507
_ = agenttest.New(t, client.URL, r.AgentToken)
@@ -825,11 +825,10 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
825825
agentClient := agentsdk.New(client.URL)
826826
agentClient.SetSessionToken(r.AgentToken)
827827

828-
manifest, err := agentClient.Manifest(ctx)
829-
require.NoError(t, err)
828+
manifest := requireGetManifest(ctx, t, agentClient)
830829
require.EqualValues(t, codersdk.WorkspaceAppHealthDisabled, manifest.Apps[0].Health)
831830
require.EqualValues(t, codersdk.WorkspaceAppHealthInitializing, manifest.Apps[1].Health)
832-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{})
831+
err := agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{})
833832
require.Error(t, err)
834833
// empty
835834
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{})
@@ -855,8 +854,7 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
855854
},
856855
})
857856
require.NoError(t, err)
858-
manifest, err = agentClient.Manifest(ctx)
859-
require.NoError(t, err)
857+
manifest = requireGetManifest(ctx, t, agentClient)
860858
require.EqualValues(t, codersdk.WorkspaceAppHealthHealthy, manifest.Apps[1].Health)
861859
// update to unhealthy
862860
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{
@@ -865,8 +863,7 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
865863
},
866864
})
867865
require.NoError(t, err)
868-
manifest, err = agentClient.Manifest(ctx)
869-
require.NoError(t, err)
866+
manifest = requireGetManifest(ctx, t, agentClient)
870867
require.EqualValues(t, codersdk.WorkspaceAppHealthUnhealthy, manifest.Apps[1].Health)
871868
}
872869

@@ -1110,8 +1107,7 @@ func TestWorkspaceAgent_Metadata(t *testing.T) {
11101107

11111108
ctx := testutil.Context(t, testutil.WaitMedium)
11121109

1113-
manifest, err := agentClient.Manifest(ctx)
1114-
require.NoError(t, err)
1110+
manifest := requireGetManifest(ctx, t, agentClient)
11151111

11161112
// Verify manifest API response.
11171113
require.Equal(t, workspace.ID, manifest.WorkspaceID)
@@ -1282,8 +1278,7 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
12821278

12831279
ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitSuperLong))
12841280

1285-
manifest, err := agentClient.Manifest(ctx)
1286-
require.NoError(t, err)
1281+
manifest := requireGetManifest(ctx, t, agentClient)
12871282

12881283
post := func(ctx context.Context, key, value string) error {
12891284
return agentClient.PostMetadata(ctx, agentsdk.PostMetadataRequest{
@@ -1630,3 +1625,18 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
16301625
require.Equal(t, 1, validateCalls, "validate calls duplicated on same token")
16311626
})
16321627
}
1628+
1629+
func requireGetManifest(ctx context.Context, t testing.TB, client agent.Client) agentsdk.Manifest {
1630+
conn, err := client.Listen(ctx)
1631+
require.NoError(t, err)
1632+
defer func() {
1633+
cErr := conn.Close()
1634+
require.NoError(t, cErr)
1635+
}()
1636+
aAPI := agentproto.NewDRPCAgentClient(conn)
1637+
mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
1638+
require.NoError(t, err)
1639+
manifest, err := agentsdk.ManifestFromProto(mp)
1640+
require.NoError(t, err)
1641+
return manifest
1642+
}

‎coderd/workspaceapps/apptest/setup.go‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"cdr.dev/slog"
2121
"cdr.dev/slog/sloggers/slogtest"
2222
"github.com/coder/coder/v2/agent"
23+
agentproto "github.com/coder/coder/v2/agent/proto"
2324
"github.com/coder/coder/v2/coderd/coderdtest"
2425
"github.com/coder/coder/v2/coderd/workspaceapps"
2526
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
@@ -397,7 +398,10 @@ func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.U
397398
primaryAppHost, err := client.AppHost(appHostCtx)
398399
require.NoError(t, err)
399400
if primaryAppHost.Host != "" {
400-
manifest, err := agentClient.Manifest(appHostCtx)
401+
rpcConn, err := agentClient.Listen(appHostCtx)
402+
require.NoError(t, err)
403+
aAPI := agentproto.NewDRPCAgentClient(rpcConn)
404+
manifest, err := aAPI.GetManifest(appHostCtx, &agentproto.GetManifestRequest{})
401405
require.NoError(t, err)
402406

403407
appHost := appurl.ApplicationURL{
@@ -408,7 +412,9 @@ func createWorkspaceWithApps(t *testing.T, client *codersdk.Client, orgID uuid.U
408412
Username: me.Username,
409413
}
410414
proxyURL := "http://" + appHost.String() + strings.ReplaceAll(primaryAppHost.Host, "*", "")
411-
require.Equal(t, manifest.VSCodePortProxyURI, proxyURL)
415+
require.Equal(t, manifest.VsCodePortProxyUri, proxyURL)
416+
err = rpcConn.Close()
417+
require.NoError(t, err)
412418
}
413419
agentCloser := agent.New(agent.Options{
414420
Client: agentClient,

‎codersdk/agentsdk/agentsdk.go‎

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -135,35 +135,13 @@ type Script struct {
135135
Script string `json:"script"`
136136
}
137137

138-
// Manifest fetches manifest for the currently authenticated workspace agent.
139-
func (c *Client) Manifest(ctx context.Context) (Manifest, error) {
140-
res, err := c.SDK.Request(ctx, http.MethodGet, "/api/v2/workspaceagents/me/manifest", nil)
141-
if err != nil {
142-
return Manifest{}, err
143-
}
144-
defer res.Body.Close()
145-
if res.StatusCode != http.StatusOK {
146-
return Manifest{}, codersdk.ReadBodyAsError(res)
147-
}
148-
var agentMeta Manifest
149-
err = json.NewDecoder(res.Body).Decode(&agentMeta)
150-
if err != nil {
151-
return Manifest{}, err
152-
}
153-
err = c.rewriteDerpMap(agentMeta.DERPMap)
154-
if err != nil {
155-
return Manifest{}, err
156-
}
157-
return agentMeta, nil
158-
}
159-
160-
// rewriteDerpMap rewrites the DERP map to use the access URL of the SDK as the
138+
// RewriteDERPMap rewrites the DERP map to use the access URL of the SDK as the
161139
// "embedded relay" access URL. The passed derp map is modified in place.
162140
//
163141
// Agents can provide an arbitrary access URL that may be different that the
164142
// globally configured one. This breaks the built-in DERP, which would continue
165143
// to reference the global access URL.
166-
func (c *Client) rewriteDerpMap(derpMap *tailcfg.DERPMap) error {
144+
func (c *Client) RewriteDERPMap(derpMap *tailcfg.DERPMap) {
167145
accessingPort := c.SDK.URL.Port()
168146
if accessingPort == "" {
169147
accessingPort = "80"
@@ -173,7 +151,9 @@ func (c *Client) rewriteDerpMap(derpMap *tailcfg.DERPMap) error {
173151
}
174152
accessPort, err := strconv.Atoi(accessingPort)
175153
if err != nil {
176-
return xerrors.Errorf("convert accessing port %q: %w", accessingPort, err)
154+
// this should never happen because URL.Port() returns the empty string if the port is not
155+
// valid.
156+
c.SDK.Logger().Critical(context.Background(), "failed to parse URL port", slog.F("port", accessingPort))
177157
}
178158
for _, region := range derpMap.Regions {
179159
if !region.EmbeddedRelay {
@@ -189,7 +169,6 @@ func (c *Client) rewriteDerpMap(derpMap *tailcfg.DERPMap) error {
189169
node.ForceHTTP = c.SDK.URL.Scheme == "http"
190170
}
191171
}
192-
return nil
193172
}
194173

195174
// Listen connects to the workspace agent API WebSocket

‎codersdk/workspaceagents_test.go‎

Lines changed: 19 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ import (
55
"net/http"
66
"net/http/httptest"
77
"net/url"
8-
"strconv"
98
"sync/atomic"
109
"testing"
1110
"time"
@@ -20,37 +19,32 @@ import (
2019
"github.com/coder/coder/v2/testutil"
2120
)
2221

23-
func TestWorkspaceAgentMetadata(t *testing.T) {
22+
func TestWorkspaceRewriteDERPMap(t *testing.T) {
2423
t.Parallel()
25-
// This test ensures that the DERP map returned properly
26-
// mutates built-in DERPs with the client access URL.
27-
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
28-
httpapi.Write(context.Background(), w, http.StatusOK, agentsdk.Manifest{
29-
DERPMap: &tailcfg.DERPMap{
30-
Regions: map[int]*tailcfg.DERPRegion{
31-
1: {
32-
EmbeddedRelay: true,
33-
RegionID: 1,
34-
Nodes: []*tailcfg.DERPNode{{
35-
HostName: "bananas.org",
36-
DERPPort: 1,
37-
}},
38-
},
39-
},
24+
// This test ensures that RewriteDERPMap mutates built-in DERPs with the
25+
// client access URL.
26+
dm := &tailcfg.DERPMap{
27+
Regions: map[int]*tailcfg.DERPRegion{
28+
1: {
29+
EmbeddedRelay: true,
30+
RegionID: 1,
31+
Nodes: []*tailcfg.DERPNode{{
32+
HostName: "bananas.org",
33+
DERPPort: 1,
34+
}},
4035
},
41-
})
42-
}))
43-
parsed, err := url.Parse(srv.URL)
36+
},
37+
}
38+
parsed, err := url.Parse("https://coconuts.org:44558")
4439
require.NoError(t, err)
4540
client := agentsdk.New(parsed)
46-
manifest, err := client.Manifest(context.Background())
47-
require.NoError(t, err)
48-
region := manifest.DERPMap.Regions[1]
41+
client.RewriteDERPMap(dm)
42+
region := dm.Regions[1]
4943
require.True(t, region.EmbeddedRelay)
5044
require.Len(t, region.Nodes, 1)
5145
node := region.Nodes[0]
52-
require.Equal(t, parsed.Hostname(), node.HostName)
53-
require.Equal(t, parsed.Port(), strconv.Itoa(node.DERPPort))
46+
require.Equal(t, "coconuts.org", node.HostName)
47+
require.Equal(t, 44558, node.DERPPort)
5448
}
5549

5650
func TestAgentReportStats(t *testing.T) {

0 commit comments

Comments
 (0)