From 2b290eea3f723df1789fef7f7780ac0bd8b4166d Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Tue, 13 Oct 2026 16:01:16 -0700 Subject: [PATCH] fix(BRE2-1216): restore brev connect node status --- pkg/cmd/ls/ls.go | 52 ++++++++++++- pkg/cmd/ls/ls_test.go | 101 +++++++++++++++++++++++++- pkg/store/http_response_error_test.go | 12 +-- 3 files changed, 156 insertions(+), 9 deletions(-) diff --git a/pkg/cmd/ls/ls.go b/pkg/cmd/ls/ls.go index fcf09af0..4b7e9f7e 100644 --- a/pkg/cmd/ls/ls.go +++ b/pkg/cmd/ls/ls.go @@ -10,6 +10,7 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" + "google.golang.org/protobuf/proto" "github.com/brevdev/brev-cli/pkg/analytics" "github.com/brevdev/brev-cli/pkg/auth" @@ -429,7 +430,7 @@ func (ls Ls) RunWorkspaces(cliAuth auth.CLIAuth, org *entity.Organization, showA go func() { defer wg.Done() var err error - nodes, err = ls.listNodes(org) + nodes, err = ls.listNodesWithConnectivity(org) if err != nil { if featureflag.Debug() { _, _ = fmt.Fprintf(os.Stderr, "debug: failed to list Brev Connect machines: %v\n", err) @@ -726,6 +727,44 @@ type NodeInfo struct { Status string `json:"status"` } +const nodeStatusFetchConcurrency = 5 + +// enrichNodesWithConnectivity re-fetches each node via GetNode, which returns +// live connectivity info that ListNodes cannot provide. Fetches run in parallel +// with at most nodeStatusFetchConcurrency in flight. Nodes whose GetNode fails +// keep their list-provided connectivity info. +func (ls Ls) enrichNodesWithConnectivity(ctx context.Context, nodes []*nodev1.ExternalNode) []*nodev1.ExternalNode { + if len(nodes) == 0 { + return nodes + } + client := register.NewNodeServiceClient(ls.lsStore, config.GlobalConfig.GetBrevPublicAPIURL()) + + sem := make(chan struct{}, nodeStatusFetchConcurrency) + var wg sync.WaitGroup + for _, node := range nodes { + wg.Add(1) + go func(node *nodev1.ExternalNode) { + defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() + resp, err := client.GetNode(ctx, connect.NewRequest(&nodev1.GetNodeRequest{ + ExternalNodeId: node.GetExternalNodeId(), + })) + if err != nil { + if featureflag.Debug() { + _, _ = fmt.Fprintf(os.Stderr, "debug: failed to fetch connectivity for node %s: %v\n", node.GetExternalNodeId(), err) + } + return + } + if enriched := resp.Msg.GetExternalNode(); enriched != nil { + proto.Merge(node, enriched) + } + }(node) + } + wg.Wait() + return nodes +} + func (ls Ls) listNodes(org *entity.Organization) ([]*nodev1.ExternalNode, error) { client := register.NewNodeServiceClient(ls.lsStore, config.GlobalConfig.GetBrevPublicAPIURL()) resp, err := client.ListNodes(context.Background(), connect.NewRequest(&nodev1.ListNodesRequest{ @@ -737,9 +776,18 @@ func (ls Ls) listNodes(org *entity.Organization) ([]*nodev1.ExternalNode, error) return resp.Msg.GetItems(), nil } +// TODO should this be optional? +func (ls Ls) listNodesWithConnectivity(org *entity.Organization) ([]*nodev1.ExternalNode, error) { + nodes, err := ls.listNodes(org) + if err != nil { + return nil, err + } + return ls.enrichNodesWithConnectivity(context.Background(), nodes), nil +} + // RunNodes lists Brev Connect machines for the given org. func (ls Ls) RunNodes(org *entity.Organization) error { - nodes, err := ls.listNodes(org) + nodes, err := ls.listNodesWithConnectivity(org) if err != nil { return breverrors.WrapAndTrace(err) } diff --git a/pkg/cmd/ls/ls_test.go b/pkg/cmd/ls/ls_test.go index b2322f58..47d6fd62 100644 --- a/pkg/cmd/ls/ls_test.go +++ b/pkg/cmd/ls/ls_test.go @@ -3,9 +3,11 @@ package ls import ( "context" "encoding/json" + "fmt" "net/http/httptest" "os" "strings" + "sync/atomic" "testing" nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" @@ -712,6 +714,15 @@ func TestOutputWorkspacesJSON_NoNodes(t *testing.T) { type lsFakeNodeService struct { nodev1connect.UnimplementedExternalNodeServiceHandler + // getNodeStatus, when set, is returned as ConnectivityInfo.Status by GetNode. + getNodeStatus nodev1.NetworkMemberStatus + // failNodeID, when set, makes GetNode return an error for that node. + failNodeID string + // getNodeCalls counts GetNode invocations. + getNodeCalls atomic.Int64 + // maxInFlightGetNode tracks peak concurrent GetNode handlers. + maxInFlightGetNode atomic.Int64 + inFlightGetNode atomic.Int64 } func (f *lsFakeNodeService) ListNodes(_ context.Context, req *connect.Request[nodev1.ListNodesRequest]) (*connect.Response[nodev1.ListNodesResponse], error) { @@ -726,6 +737,31 @@ func (f *lsFakeNodeService) ListNodes(_ context.Context, req *connect.Request[no }), nil } +func (f *lsFakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { + if req.Msg.GetExternalNodeId() == f.failNodeID { + return nil, connect.NewError(connect.CodeNotFound, fmt.Errorf("node %s not found", req.Msg.GetExternalNodeId())) + } + f.getNodeCalls.Add(1) + inFlight := f.inFlightGetNode.Add(1) + defer f.inFlightGetNode.Add(-1) + for { + max := f.maxInFlightGetNode.Load() + if inFlight <= max || f.maxInFlightGetNode.CompareAndSwap(max, inFlight) { + break + } + } + return connect.NewResponse(&nodev1.GetNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + Name: "enriched-" + req.Msg.GetExternalNodeId(), + ExternalNodeId: req.Msg.GetExternalNodeId(), + OrganizationId: "org1", + ConnectivityInfo: &nodev1.ConnectivityInfo{ + Status: f.getNodeStatus, + }, + }, + }), nil +} + func withFakeNodeAPI(t *testing.T, fn func()) { t.Helper() _, handler := nodev1connect.NewExternalNodeServiceHandler(&lsFakeNodeService{}) @@ -735,6 +771,67 @@ func withFakeNodeAPI(t *testing.T, fn func()) { fn() } +func withFakeNodeAPIService(t *testing.T, svc *lsFakeNodeService, fn func()) { + t.Helper() + _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) + server := httptest.NewServer(handler) + t.Cleanup(server.Close) + t.Setenv("BREV_PUBLIC_API_URL", server.URL) + fn() +} + +// TestRunNodes_EnrichesConnectivity verifies brev ls nodes shows the live +// GetNode status, not the list-provided "Unknown". +func TestRunNodes_EnrichesConnectivity(t *testing.T) { + term := terminal.New() + + t.Run("GetNode status wins over list status", func(t *testing.T) { + svc := &lsFakeNodeService{getNodeStatus: nodev1.NetworkMemberStatus_NETWORK_MEMBER_STATUS_CONNECTED} + withFakeNodeAPIService(t, svc, func() { + ls := NewLs(newTestStore(), term, false) + if err := ls.RunNodes(&entity.Organization{ID: "org1", Name: "org1"}); err != nil { + t.Fatalf("RunNodes: %v", err) + } + if svc.getNodeCalls.Load() != 2 { + t.Fatalf("expected 2 GetNode calls, got %d", svc.getNodeCalls.Load()) + } + }) + }) + + t.Run("at most 5 GetNode calls in flight", func(t *testing.T) { + svc := &lsFakeNodeService{getNodeStatus: nodev1.NetworkMemberStatus_NETWORK_MEMBER_STATUS_CONNECTED} + withFakeNodeAPIService(t, svc, func() { + ls := NewLs(newTestStore(), term, false) + nodes := make([]*nodev1.ExternalNode, 13) + for i := range nodes { + nodes[i] = &nodev1.ExternalNode{ExternalNodeId: fmt.Sprintf("en-%d", i)} + } + _ = ls.enrichNodesWithConnectivity(context.Background(), nodes) + if svc.getNodeCalls.Load() != 13 { + t.Fatalf("expected 13 GetNode calls, got %d", svc.getNodeCalls.Load()) + } + if max := svc.maxInFlightGetNode.Load(); max > 5 { + t.Fatalf("peak concurrent GetNode calls = %d, want <= 5", max) + } + }) + }) + + t.Run("GetNode failure keeps node with Unknown status", func(t *testing.T) { + svc := &lsFakeNodeService{getNodeStatus: nodev1.NetworkMemberStatus_NETWORK_MEMBER_STATUS_CONNECTED, failNodeID: "en-missing"} + withFakeNodeAPIService(t, svc, func() { + ls := NewLs(newTestStore(), term, false) + nodes := []*nodev1.ExternalNode{{ExternalNodeId: "en-missing"}} + got := ls.enrichNodesWithConnectivity(context.Background(), nodes) + if len(got) != 1 || got[0] != nodes[0] { + t.Fatalf("expected original node retained, got %+v", got) + } + if status := nodeConnectionStatus(got[0]); status != "Unknown" { + t.Fatalf("expected Unknown status on failed enrichment, got %q", status) + } + }) + }) +} + // TestRunLs_ShowAllJSON_WithNodes verifies --all includes nodes from listNodes. func TestRunLs_ShowAllJSON_WithNodes(t *testing.T) { term := terminal.New() @@ -757,8 +854,8 @@ func TestRunLs_ShowAllJSON_WithNodes(t *testing.T) { if err := json.Unmarshal([]byte(out), &res); err != nil { t.Fatalf("parse JSON: %v\n%s", err, out) } - if len(res.Nodes) != 2 || res.Nodes[0].Name != "gpu-node-1" { - t.Fatalf("expected 2 nodes from listNodes, got %+v", res.Nodes) + if len(res.Nodes) != 2 || res.Nodes[0].Name != "enriched-en-123" { + t.Fatalf("expected 2 enriched nodes from listNodes, got %+v", res.Nodes) } }) }) diff --git a/pkg/store/http_response_error_test.go b/pkg/store/http_response_error_test.go index 83e0adce..489a237f 100644 --- a/pkg/store/http_response_error_test.go +++ b/pkg/store/http_response_error_test.go @@ -38,10 +38,11 @@ func TestHTTPResponseError_UnauthorizedWithAPIKeyShowsHintThenOriginalMessage(t msg := err.Error() assert.True(t, strings.HasPrefix(msg, expiredAPIKeyMessage+"\n"), "hint should come first: %q", msg) assert.Contains(t, msg, "brev login --api-key") - // the original error is preserved after the hint + // The base error is preserved after the hint. baseError() renders the + // typed-but-messageless control-plane 401 as the friendly key-expired copy + // (see the UnauthorizedError branch), not the raw "401 Unauthorized" line. assert.True(t, strings.HasSuffix(msg, err.baseError())) - assert.Contains(t, msg, "401 Unauthorized") - assert.Contains(t, msg, "UnauthorizedError") + assert.Contains(t, msg, "your Brev API key is invalid or expired") } func TestHTTPResponseError_UnauthorizedWithoutAPIKeyKeepsOriginalMessage(t *testing.T) { @@ -50,8 +51,9 @@ func TestHTTPResponseError_UnauthorizedWithoutAPIKeyKeepsOriginalMessage(t *test msg := err.Error() assert.NotContains(t, msg, expiredAPIKeyMessage) - assert.Contains(t, msg, "401") - assert.Contains(t, msg, "UnauthorizedError") + // A messageless typed 401 renders the friendly logged-out copy; the raw + // status line only appears for bodies the control-plane mapping can't interpret. + assert.Contains(t, msg, "you're logged out") } func TestHTTPResponseError_UnauthorizedWithNoTokenKeepsOriginalMessage(t *testing.T) {