From 063e236cdfd441e78919512190dc1f5d8ec6ae1a Mon Sep 17 00:00:00 2001 From: Pratik Patel Date: Tue, 13 Oct 2026 14:15:21 -0700 Subject: [PATCH] fix status and use devplane as source of truth --- Makefile | 2 - go.mod | 4 +- go.sum | 8 +- pkg/cmd/deregister/deregister_test.go | 3 +- pkg/cmd/enablessh/enablessh_test.go | 3 +- pkg/cmd/gpucreate/gpucreate.go | 118 +++-- pkg/cmd/gpucreate/gpucreate_test.go | 716 +++++++++++++++----------- pkg/cmd/grantssh/grantssh_test.go | 3 +- pkg/cmd/launch/launch.go | 5 +- pkg/cmd/launch/launch_test.go | 11 + pkg/cmd/launch/secrets.go | 8 +- pkg/cmd/ls/ls.go | 89 +++- pkg/cmd/ls/ls_test.go | 61 ++- pkg/cmd/mintcert/mintcert.go | 7 +- pkg/cmd/ports/close.go | 6 +- pkg/cmd/ports/open.go | 14 +- pkg/cmd/ports/ports.go | 8 +- pkg/cmd/ports/ports_test.go | 6 + pkg/cmd/ports/update.go | 18 +- pkg/cmd/refresh/refresh.go | 5 +- pkg/cmd/refresh/refresh_test.go | 6 + pkg/cmd/refresh/sshaccess.go | 4 +- pkg/cmd/register/providers.go | 3 +- pkg/cmd/register/register.go | 73 +++ pkg/cmd/register/register_test.go | 301 ++++++++++- pkg/cmd/register/rpcclient.go | 145 ------ pkg/cmd/register/rpcclient_test.go | 387 -------------- pkg/cmd/revokessh/revokessh_test.go | 3 +- pkg/cmd/util/externalnode.go | 14 +- pkg/cmd/util/externalnode_test.go | 6 + pkg/config/config.go | 10 +- pkg/environment/status.go | 76 +++ pkg/environment/status_test.go | 92 ++++ pkg/store/devplane_client.go | 78 +++ pkg/store/devplane_client_test.go | 64 +++ pkg/store/environment.go | 54 ++ pkg/store/http.go | 10 +- pkg/store/http_response_error_test.go | 12 +- pkg/store/instancetypes.go | 12 +- pkg/store/managedsecret.go | 7 +- pkg/store/organization.go | 28 +- 41 files changed, 1471 insertions(+), 1009 deletions(-) delete mode 100644 pkg/cmd/register/rpcclient.go delete mode 100644 pkg/cmd/register/rpcclient_test.go create mode 100644 pkg/environment/status.go create mode 100644 pkg/environment/status_test.go create mode 100644 pkg/store/devplane_client.go create mode 100644 pkg/store/devplane_client_test.go create mode 100644 pkg/store/environment.go diff --git a/Makefile b/Makefile index 8cbf9ee81..ba454f037 100644 --- a/Makefile +++ b/Makefile @@ -42,11 +42,9 @@ ifdef env ifeq ($(env),local-stack) @echo 'export BREV_API_URL="http://localhost:8080"' >> brev @echo 'export BREV_PUBLIC_API_URL="http://localhost:8081"' >> brev - @echo 'export BREV_GRPC_URL="localhost:8081"' >> brev else @echo 'export BREV_API_URL="https://bd.$(env).brev.nvidia.com"' >> brev @echo 'export BREV_PUBLIC_API_URL="https://api.$(env).brev.nvidia.com"' >> brev - @echo 'export BREV_GRPC_URL="api.$(env).brev.nvidia.com:443"' >> brev endif @echo 'exec "$$(cd "$$(dirname "$$0")" && pwd)/brev-local" "$$@"' >> brev @chmod +x brev diff --git a/go.mod b/go.mod index 4d0829606..e826383b0 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,8 @@ module github.com/brevdev/brev-cli go 1.25.0 require ( - buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20260915224432-6a5d0c858507.1 - buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260915224432-6a5d0c858507.2 + buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20261002013339-6f4f902cdf2d.1 + buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20261002013339-6f4f902cdf2d.2 connectrpc.com/connect v1.21.0 github.com/NVIDIA/go-nvml v0.13.0-1 github.com/alessio/shellescape v1.4.1 diff --git a/go.sum b/go.sum index 294c9e6c4..f73d8a8d7 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,7 @@ -buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20260915224432-6a5d0c858507.1 h1:CFXTDUzEoMYQzJXAG+hnMO5Cs6h3zVq734k/VnyWuCM= -buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20260915224432-6a5d0c858507.1/go.mod h1:dS34xJi+QXqWDdKwUHTBfzkcbQFNK3u4UH3F3L6izx0= -buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260915224432-6a5d0c858507.2 h1:CjGsO8kcUR0YcKzjKy/L5m2jIhnNW6VW6ArYliA2p6o= -buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20260915224432-6a5d0c858507.2/go.mod h1:SBVLUYGc/4vgMlhOJL5SkwIrSh/FL9EEimlpI8zQSXo= +buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20261002013339-6f4f902cdf2d.1 h1:6YxQ+ovFXTH0qhavCXrRAsa/+CgyJviIBBPvuar5fAU= +buf.build/gen/go/brevdev/devplane/connectrpc/go v1.21.0-20261002013339-6f4f902cdf2d.1/go.mod h1:0EoTlhUve4UaoJRDO8DXLA79x0HYGLnFchd8ITsIn/4= +buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20261002013339-6f4f902cdf2d.2 h1:FoyCt/oxRt8ZFE4JmtAclKTlGOT7zFbyyhY/xeEmyjQ= +buf.build/gen/go/brevdev/devplane/protocolbuffers/go v1.36.12-20261002013339-6f4f902cdf2d.2/go.mod h1:SBVLUYGc/4vgMlhOJL5SkwIrSh/FL9EEimlpI8zQSXo= buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.12-20220906235457-8b4922735da5.2 h1:W5SVs96P8ZgLElYzL93rJu+S+XmkdGQSaVWzPvR5lvg= buf.build/gen/go/brevdev/protoc-gen-gotag/protocolbuffers/go v1.36.12-20220906235457-8b4922735da5.2/go.mod h1:SacJAYqnICCQAsBA46cSA/hxhqhxYkiYzseucf6/fhQ= cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= diff --git a/pkg/cmd/deregister/deregister_test.go b/pkg/cmd/deregister/deregister_test.go index eb02932c0..0d0342e94 100644 --- a/pkg/cmd/deregister/deregister_test.go +++ b/pkg/cmd/deregister/deregister_test.go @@ -19,6 +19,7 @@ import ( "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" "github.com/brevdev/brev-cli/pkg/sshcert" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -141,7 +142,7 @@ type mockNodeClientFactory struct { } func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient { - return register.NewNodeServiceClient(provider, m.serverURL) + return store.NewDevPlaneClient(provider, m.serverURL).ExternalNodes } type mockSSHKeyRemover struct { diff --git a/pkg/cmd/enablessh/enablessh_test.go b/pkg/cmd/enablessh/enablessh_test.go index c747a46d5..a4b6f7dac 100644 --- a/pkg/cmd/enablessh/enablessh_test.go +++ b/pkg/cmd/enablessh/enablessh_test.go @@ -17,6 +17,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -222,7 +223,7 @@ func Test_RemoveAuthorizedKey_DoesNotRemoveOtherBrevKeys(t *testing.T) { type mockNodeClientFactory struct{ serverURL string } func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient { - return register.NewNodeServiceClient(provider, m.serverURL) + return store.NewDevPlaneClient(provider, m.serverURL).ExternalNodes } type mockEnableSSHStore struct { diff --git a/pkg/cmd/gpucreate/gpucreate.go b/pkg/cmd/gpucreate/gpucreate.go index cb6b8f123..ea235d62b 100644 --- a/pkg/cmd/gpucreate/gpucreate.go +++ b/pkg/cmd/gpucreate/gpucreate.go @@ -2,7 +2,9 @@ package gpucreate import ( + "context" "encoding/json" + "errors" "fmt" "io" "math/rand/v2" @@ -20,11 +22,14 @@ import ( "time" "unicode" + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/util" "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/environment" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/featureflag" "github.com/brevdev/brev-cli/pkg/names" @@ -109,7 +114,7 @@ type GPUCreateStore interface { GetActiveOrganizationOrDefault() (*entity.Organization, error) GetAuthTokens() (*entity.AuthTokens, error) GetCurrentUser() (*entity.User, error) - GetWorkspace(workspaceID string) (*entity.Workspace, error) + GetEnvironment(ctx context.Context, environmentID string) (*devplanev1.Environment, error) CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error) DeleteWorkspace(workspaceID string) (*entity.Workspace, error) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) @@ -266,7 +271,7 @@ func NewCmdGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore) *cobra return runDryRun(t, gpuCreateStore, opts.InstanceTypes, &filters) } - return RunGPUCreate(t, gpuCreateStore, opts) + return RunGPUCreate(cmd.Context(), t, gpuCreateStore, opts) }, } @@ -1130,25 +1135,53 @@ func (c *createContext) cleanupExtraInstances(workspaces []*entity.Workspace) [] return workspaces[:c.opts.Count] } -// waitForInstances waits for all instances to be ready -func (c *createContext) waitForInstances(workspaces []*entity.Workspace) { +// errReadyTimeout marks a pollUntilReady timeout. It is non-fatal: waiting is a +// convenience and the instance may still finish provisioning after the window. +var errReadyTimeout = errors.New("timeout waiting for instance to be ready") + +type workspaceFailedError struct{ message string } + +func (e *workspaceFailedError) Error() string { return e.message } + +// waitForInstances blocks until each workspace is ready. It returns an error +// when any instance reaches a terminal failure so the command exits non-zero; +// timeouts and transient polling errors are reported but remain non-fatal. +func (c *createContext) waitForInstances(parentCtx context.Context, workspaces []*entity.Workspace) error { if c.opts.Detached { - return + return nil } c.logf("\nWaiting for instance(s) to be ready...\n") c.logf("You can safely ctrl+c to exit\n") + var failed []string for _, ws := range workspaces { - err := c.pollUntilReady(ws.ID) - if err != nil { - if strings.Contains(err.Error(), "timeout waiting") { - c.logf(" %s: Timeout waiting for ready state\n", ws.Name) - } else { - c.logf(" %s: %s\n", ws.Name, c.colorize(err.Error(), c.t.Red)) - } - } + err := c.pollUntilReady(parentCtx, ws.ID) + var failedErr *workspaceFailedError + switch { + case err == nil: + case errors.Is(err, errReadyTimeout): + c.logf(" %s: Timeout waiting for ready state\n", ws.Name) + case parentCtx.Err() != nil: + // The parent context expired or was canceled while waiting. + // Propagate instead of logging, or a dead parent deadline would + // still produce a success summary. + return err + case errors.As(err, &failedErr): + c.logf(" %s: %s\n", ws.Name, c.colorize(err.Error(), c.t.Red)) + failed = append(failed, err.Error()) + default: + c.logf(" %s: %s\n", ws.Name, c.colorize(err.Error(), c.t.Red)) + } + } + + if len(failed) > 0 { + return breverrors.NewValidationError(fmt.Sprintf( + "%d/%d instance(s) failed to become ready: %s", + len(failed), len(workspaces), strings.Join(failed, "; "), + )) } + return nil } // printSummary outputs the final creation summary @@ -1172,8 +1205,9 @@ func (c *createContext) printSummary(workspaces []*entity.Workspace) { } } -// RunGPUCreate executes the GPU create with retry logic -func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions) error { +// RunGPUCreate executes the GPU create with retry logic. parentCtx carries the +// caller's deadline into the readiness polls. +func RunGPUCreate(parentCtx context.Context, t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUCreateOptions) error { ctx, err := newCreateContext(t, gpuCreateStore, opts) if err != nil { return err @@ -1230,7 +1264,9 @@ func RunGPUCreate(t *terminal.Terminal, gpuCreateStore GPUCreateStore, opts GPUC } successfulWorkspaces = ctx.cleanupExtraInstances(successfulWorkspaces) - ctx.waitForInstances(successfulWorkspaces) + if err := ctx.waitForInstances(parentCtx, successfulWorkspaces); err != nil { + return err + } ctx.printSummary(successfulWorkspaces) return nil @@ -1609,31 +1645,55 @@ func resolveWorkspaceUserOptions(options *store.CreateWorkspacesOptions, user *e } // pollUntilReady waits for a workspace to reach the running state -func (c *createContext) pollUntilReady(wsID string) error { +func (c *createContext) pollUntilReady(parentCtx context.Context, wsID string) error { deadline := time.Now().Add(c.opts.Timeout) - for time.Now().Before(deadline) { - ws, err := c.store.GetWorkspace(wsID) + pollCtx, cancel := context.WithDeadline(parentCtx, deadline) + defer cancel() + + for { + // An expired parent deadline (or cancellation) must propagate as an + // error, not be classified as this instance's nonfatal readiness + // timeout — the caller's context ran out, not the per-instance wait. + if err := parentCtx.Err(); err != nil { + return breverrors.WrapAndTrace(err) + } + if err := pollCtx.Err(); err != nil { + if errors.Is(err, context.DeadlineExceeded) { + return errReadyTimeout + } + return breverrors.WrapAndTrace(err) + } + + env, err := c.store.GetEnvironment(pollCtx, wsID) if err != nil { return breverrors.WrapAndTrace(err) } - if ws.Status == entity.Running { - c.logf(" %s: %s\n", ws.Name, c.colorize("Ready", c.t.Green)) + name := wsID + + switch { + case env.GetStatus() == devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING: + c.logf(" %s: %s\n", name, c.colorize("Ready", c.t.Green)) return nil + case environment.IsFailed(env): + return &workspaceFailedError{message: instanceFailureMessage(name, environment.FailureMessage(env))} } - if ws.Status == entity.Failure { - if ws.StatusMessage != "" { - return breverrors.NewValidationError(fmt.Sprintf("instance %s failed: %s", ws.Name, ws.StatusMessage)) - } - return breverrors.NewValidationError(fmt.Sprintf("instance %s failed", ws.Name)) + // Wait for the next poll interval; cancellation and deadline are + // handled by the pollCtx check at the top of the loop. + select { + case <-time.After(5 * time.Second): + case <-pollCtx.Done(): } - - time.Sleep(5 * time.Second) } +} - return breverrors.NewValidationError("timeout waiting for instance to be ready") +func instanceFailureMessage(name, statusMessage string) string { + if statusMessage != "" { + return fmt.Sprintf("instance %s failed: %s", name, statusMessage) + } + return fmt.Sprintf("instance %s failed", name) } // displayConnectBreadCrumb shows connection instructions diff --git a/pkg/cmd/gpucreate/gpucreate_test.go b/pkg/cmd/gpucreate/gpucreate_test.go index e259dc06e..8f8d0f1ba 100644 --- a/pkg/cmd/gpucreate/gpucreate_test.go +++ b/pkg/cmd/gpucreate/gpucreate_test.go @@ -1,6 +1,7 @@ package gpucreate import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -8,6 +9,8 @@ import ( "testing" "time" + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" @@ -17,14 +20,30 @@ import ( "github.com/stretchr/testify/require" ) +// backgroundContext and canceledContext are context builders for table-driven +// tests that vary the parent context. +func backgroundContext() (context.Context, context.CancelFunc) { + return context.WithCancel(context.Background()) +} + +func canceledContext() (context.Context, context.CancelFunc) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + return ctx, cancel +} + +// expiredDeadlineContext returns a parent whose deadline has already passed. +func expiredDeadlineContext() (context.Context, context.CancelFunc) { + return context.WithDeadline(context.Background(), time.Now().Add(-time.Minute)) +} + // MockGPUCreateStore is a mock implementation of GPUCreateStore for testing type MockGPUCreateStore struct { User *entity.User Org *entity.Organization - Workspaces map[string]*entity.Workspace - CreateError error CreateErrorTypes map[string]error // Errors for specific instance types - DeleteError error + EnvironmentStatuses map[string]devplanev1.EnvironmentStatus + EnvironmentStatusMessages map[string]string CreatedOptions []*store.CreateWorkspacesOptions CreatedWorkspaces []*entity.Workspace DeletedWorkspaceIDs []string @@ -41,8 +60,8 @@ func NewMockGPUCreateStore() *MockGPUCreateStore { ID: "org-123", Name: "test-org", }, - Workspaces: make(map[string]*entity.Workspace), CreateErrorTypes: make(map[string]error), + EnvironmentStatuses: make(map[string]devplanev1.EnvironmentStatus), CreatedWorkspaces: []*entity.Workspace{}, DeletedWorkspaceIDs: []string{}, } @@ -60,14 +79,17 @@ func (m *MockGPUCreateStore) GetActiveOrganizationOrDefault() (*entity.Organizat return m.Org, nil } -func (m *MockGPUCreateStore) GetWorkspace(workspaceID string) (*entity.Workspace, error) { - if ws, ok := m.Workspaces[workspaceID]; ok { - return ws, nil +func (m *MockGPUCreateStore) GetEnvironment(_ context.Context, environmentID string) (*devplanev1.Environment, error) { + env := &devplanev1.Environment{EnvironmentId: environmentID} + if status, ok := m.EnvironmentStatuses[environmentID]; ok { + env.Status = status } - return &entity.Workspace{ - ID: workspaceID, - Status: entity.Running, - }, nil + if msg, ok := m.EnvironmentStatusMessages[environmentID]; ok { + env.Instance = &devplanev1.Instance{ + ProvisionStatus: &devplanev1.CreateAttemptStatus{Message: msg}, + } + } + return env, nil } func (m *MockGPUCreateStore) CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error) { @@ -76,31 +98,20 @@ func (m *MockGPUCreateStore) CreateWorkspace(organizationID string, options *sto return nil, err } - if m.CreateError != nil { - return nil, m.CreateError - } - ws := &entity.Workspace{ ID: "ws-" + options.Name, Name: options.Name, InstanceType: options.InstanceType, Status: entity.Running, } - m.Workspaces[ws.ID] = ws m.CreatedOptions = append(m.CreatedOptions, options) m.CreatedWorkspaces = append(m.CreatedWorkspaces, ws) return ws, nil } func (m *MockGPUCreateStore) DeleteWorkspace(workspaceID string) (*entity.Workspace, error) { - if m.DeleteError != nil { - return nil, m.DeleteError - } - m.DeletedWorkspaceIDs = append(m.DeletedWorkspaceIDs, workspaceID) - ws := m.Workspaces[workspaceID] - delete(m.Workspaces, workspaceID) - return ws, nil + return nil, nil } func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID string) ([]entity.Workspace, error) { @@ -108,7 +119,12 @@ func (m *MockGPUCreateStore) GetWorkspaceByNameOrID(orgID string, nameOrID strin } func (m *MockGPUCreateStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { - return nil, nil + return &gpusearch.AllInstanceTypesResponse{ + AllInstanceTypes: []gpusearch.InstanceType{ + {Type: "test.fail.quota", CloudCredID: "testkube-brev-test"}, + {Type: "g5.xlarge", CloudCredID: "cc-shadeform"}, + }, + }, nil } func (m *MockGPUCreateStore) GetLaunchable(launchableID string) (*store.LaunchableResponse, error) { @@ -793,129 +809,58 @@ func TestParseInstanceTypesStdin(t *testing.T) { assert.Error(t, err) } -func TestGPUCreateOptions(t *testing.T) { - opts := GPUCreateOptions{ - Name: "my-instance", - InstanceTypes: []InstanceSpec{ - {Type: "g5.xlarge", DiskGB: 500}, - {Type: "g5.2xlarge"}, +func TestResolveWorkspaceUserOptions(t *testing.T) { + cases := []struct { + name string + globalUserType string + existing *store.CreateWorkspacesOptions + wantTemplate string + wantClass string + }{ + { + name: "standard user gets user template", + globalUserType: "Standard", + wantTemplate: store.UserWorkspaceTemplateID, + wantClass: store.UserWorkspaceClassID, + }, + { + name: "admin user gets dev template", + globalUserType: "Admin", + wantTemplate: store.DevWorkspaceTemplateID, + wantClass: store.DevWorkspaceClassID, + }, + { + name: "existing values are preserved", + globalUserType: "Standard", + existing: &store.CreateWorkspacesOptions{ + WorkspaceTemplateID: "custom-template", + WorkspaceClassID: "custom-class", + }, + wantTemplate: "custom-template", + wantClass: "custom-class", }, - Count: 2, - Parallel: 3, - Detached: true, - } - - assert.Equal(t, "my-instance", opts.Name) - assert.Len(t, opts.InstanceTypes, 2) - assert.Equal(t, "g5.xlarge", opts.InstanceTypes[0].Type) - assert.Equal(t, 500.0, opts.InstanceTypes[0].DiskGB) - assert.Equal(t, "g5.2xlarge", opts.InstanceTypes[1].Type) - assert.Equal(t, 0.0, opts.InstanceTypes[1].DiskGB) - assert.Equal(t, 2, opts.Count) - assert.Equal(t, 3, opts.Parallel) - assert.True(t, opts.Detached) -} - -func TestResolveWorkspaceUserOptionsStandard(t *testing.T) { - user := &entity.User{ - ID: "user-123", - GlobalUserType: "Standard", - } - - options := &store.CreateWorkspacesOptions{} - result := resolveWorkspaceUserOptions(options, user) - - assert.Equal(t, store.UserWorkspaceTemplateID, result.WorkspaceTemplateID) - assert.Equal(t, store.UserWorkspaceClassID, result.WorkspaceClassID) -} - -func TestResolveWorkspaceUserOptionsAdmin(t *testing.T) { - user := &entity.User{ - ID: "user-123", - GlobalUserType: "Admin", } - options := &store.CreateWorkspacesOptions{} - result := resolveWorkspaceUserOptions(options, user) - - assert.Equal(t, store.DevWorkspaceTemplateID, result.WorkspaceTemplateID) - assert.Equal(t, store.DevWorkspaceClassID, result.WorkspaceClassID) -} + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + options := tc.existing + if options == nil { + options = &store.CreateWorkspacesOptions{} + } -func TestResolveWorkspaceUserOptionsPreserveExisting(t *testing.T) { - user := &entity.User{ - ID: "user-123", - GlobalUserType: "Standard", - } + result := resolveWorkspaceUserOptions(options, &entity.User{ + ID: "user-123", + GlobalUserType: entity.GlobalUserType(tc.globalUserType), + }) - options := &store.CreateWorkspacesOptions{ - WorkspaceTemplateID: "custom-template", - WorkspaceClassID: "custom-class", + assert.Equal(t, tc.wantTemplate, result.WorkspaceTemplateID) + assert.Equal(t, tc.wantClass, result.WorkspaceClassID) + }) } - result := resolveWorkspaceUserOptions(options, user) - - // Should preserve existing values - assert.Equal(t, "custom-template", result.WorkspaceTemplateID) - assert.Equal(t, "custom-class", result.WorkspaceClassID) -} - -func TestMockGPUCreateStoreBasics(t *testing.T) { - mock := NewMockGPUCreateStore() - - user, err := mock.GetCurrentUser() - assert.NoError(t, err) - assert.Equal(t, "user-123", user.ID) - - org, err := mock.GetActiveOrganizationOrDefault() - assert.NoError(t, err) - assert.Equal(t, "org-123", org.ID) -} - -func TestMockGPUCreateStoreCreateWorkspace(t *testing.T) { - mock := NewMockGPUCreateStore() - - options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") - options.WithInstanceType("g5.xlarge") - - ws, err := mock.CreateWorkspace("org-123", options) - assert.NoError(t, err) - assert.Equal(t, "test-instance", ws.Name) - assert.Equal(t, "g5.xlarge", ws.InstanceType) - assert.Len(t, mock.CreatedWorkspaces, 1) -} - -func TestMockGPUCreateStoreDeleteWorkspace(t *testing.T) { - mock := NewMockGPUCreateStore() - - // First create a workspace - options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") - ws, _ := mock.CreateWorkspace("org-123", options) - - // Then delete it - _, err := mock.DeleteWorkspace(ws.ID) - assert.NoError(t, err) - assert.Contains(t, mock.DeletedWorkspaceIDs, ws.ID) -} - -func TestMockGPUCreateStoreTypeSpecificError(t *testing.T) { - mock := NewMockGPUCreateStore() - mock.CreateErrorTypes["g5.xlarge"] = assert.AnError - - options := store.NewCreateWorkspacesOptions("cluster-1", "test-instance") - options.WithInstanceType("g5.xlarge") - - _, err := mock.CreateWorkspace("org-123", options) - assert.Error(t, err) - - // Different type should work - options2 := store.NewCreateWorkspacesOptions("cluster-1", "test-instance-2") - options2.WithInstanceType("g5.2xlarge") - - ws, err := mock.CreateWorkspace("org-123", options2) - assert.NoError(t, err) - assert.NotNil(t, ws) } +// TestCreateDryRunWithExplicitTypesDoesNotProvision covers --dry-run end to +// end: flags parse, no CreateWorkspace call is made. func TestCreateDryRunWithExplicitTypesDoesNotProvision(t *testing.T) { mock := NewMockGPUCreateStore() term := terminal.New() @@ -962,47 +907,52 @@ func TestGetFilteredInstanceTypesNoMatch(t *testing.T) { assert.Len(t, specs, 0) } +// TestParseTableInput covers both `brev search` table shapes (GPU and CPU). func TestParseTableInput(t *testing.T) { - tableInput := strings.Join([]string{ - "TYPE TARGET_DISK GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR", - "g5.xlarge 500 A10G 1 24 GB 24 GB 8.6 4 $1.01", - "g5.2xlarge 500 A10G 1 24 GB 24 GB 8.6 8 $1.21", - "p4d.24xlarge 1000 A100 8 40 GB 320 GB 8.0 96 $32.77", - "", - "Found 3 GPU instance types", - }, "\n") - - specs := parseTableInput(tableInput) - - assert.Len(t, specs, 3) - assert.Equal(t, "g5.xlarge", specs[0].Type) - assert.Equal(t, 500.0, specs[0].DiskGB) - assert.Equal(t, "g5.2xlarge", specs[1].Type) - assert.Equal(t, 500.0, specs[1].DiskGB) - assert.Equal(t, "p4d.24xlarge", specs[2].Type) - assert.Equal(t, 1000.0, specs[2].DiskGB) -} - -func TestParseTableInputCPU(t *testing.T) { - // Simulated plain table output from `brev search cpu` - tableInput := strings.Join([]string{ - " TYPE TARGET_DISK PROVIDER VCPUS RAM ARCH DISK $/GB/MO BOOT FEATURES $/HR", - " n2d-highcpu-2 10 gcp 2 2 x86_64 10GB-16TB $0.13 7m SP $0.05", - " n1-standard-1 10 gcp 1 4 x86_64 10GB-16TB $0.14 7m SP $0.06", - " m8i-flex.8xlarge 500 aws 32 128 x86_64 10GB-16TB $0.10 7m SRP $1.93", - "", - "Found 3 CPU instance types", - }, "\n") - - specs := parseTableInput(tableInput) + cases := []struct { + name string + input string + want []InstanceSpec + }{ + { + name: "gpu table", + input: strings.Join([]string{ + "TYPE TARGET_DISK GPU COUNT VRAM/GPU TOTAL VRAM CAPABILITY VCPUs $/HR", + "g5.xlarge 500 A10G 1 24 GB 24 GB 8.6 4 $1.01", + "g5.2xlarge 500 A10G 1 24 GB 24 GB 8.6 8 $1.21", + "p4d.24xlarge 1000 A100 8 40 GB 320 GB 8.0 96 $32.77", + "", + "Found 3 GPU instance types", + }, "\n"), + want: []InstanceSpec{ + {Type: "g5.xlarge", DiskGB: 500}, + {Type: "g5.2xlarge", DiskGB: 500}, + {Type: "p4d.24xlarge", DiskGB: 1000}, + }, + }, + { + name: "cpu table", + input: strings.Join([]string{ + " TYPE TARGET_DISK PROVIDER VCPUS RAM ARCH DISK $/GB/MO BOOT FEATURES $/HR", + " n2d-highcpu-2 10 gcp 2 2 x86_64 10GB-16TB $0.13 7m SP $0.05", + " n1-standard-1 10 gcp 1 4 x86_64 10GB-16TB $0.14 7m SP $0.06", + " m8i-flex.8xlarge 500 aws 32 128 x86_64 10GB-16TB $0.10 7m SRP $1.93", + "", + "Found 3 CPU instance types", + }, "\n"), + want: []InstanceSpec{ + {Type: "n2d-highcpu-2", DiskGB: 10}, + {Type: "n1-standard-1", DiskGB: 10}, + {Type: "m8i-flex.8xlarge", DiskGB: 500}, + }, + }, + } - assert.Len(t, specs, 3) - assert.Equal(t, "n2d-highcpu-2", specs[0].Type) - assert.Equal(t, 10.0, specs[0].DiskGB) - assert.Equal(t, "n1-standard-1", specs[1].Type) - assert.Equal(t, 10.0, specs[1].DiskGB) - assert.Equal(t, "m8i-flex.8xlarge", specs[2].Type) - assert.Equal(t, 500.0, specs[2].DiskGB) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + assert.Equal(t, tc.want, parseTableInput(tc.input)) + }) + } } func TestParseJSONInput(t *testing.T) { @@ -1055,23 +1005,230 @@ func TestFormatInstanceSpecs(t *testing.T) { assert.Equal(t, "g5.xlarge (1000GB disk), p4d.24xlarge, g6.xlarge (500GB disk)", result) } -func TestPollUntilReadyReportsWorkspaceFailureMessage(t *testing.T) { - store := NewMockGPUCreateStore() - store.Workspaces["ws-failed"] = &entity.Workspace{ - ID: "ws-failed", - Name: "test", - Status: entity.Failure, - StatusMessage: "unexpected end of JSON input", +// TestPollUntilReady covers the readiness-wait behavior across ctx states and +// environment status outcomes. +func TestPollUntilReady(t *testing.T) { + cases := []struct { + name string + status devplanev1.EnvironmentStatus + instanceMsg string + parentCtx func() (context.Context, context.CancelFunc) + wantErrIs error + wantErrHas string + }{ + { + name: "cancellation propagates immediately", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING, + parentCtx: canceledContext, + wantErrIs: context.Canceled, + }, + { + name: "expired parent deadline is not a readiness timeout", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING, + parentCtx: expiredDeadlineContext, + wantErrIs: context.DeadlineExceeded, + }, + { + name: "terminal failure returns message", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + instanceMsg: "unexpected end of JSON input", + parentCtx: backgroundContext, + wantErrHas: "instance ws-1 failed: unexpected end of JSON input", + }, + { + name: "failure without provider message falls back to label", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + parentCtx: backgroundContext, + wantErrHas: "instance ws-1 failed", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + store := NewMockGPUCreateStore() + store.EnvironmentStatuses["ws-1"] = tc.status + if tc.instanceMsg != "" { + store.EnvironmentStatusMessages = map[string]string{"ws-1": tc.instanceMsg} + } + + ctx := &createContext{ + t: terminal.New(), + store: store, + opts: GPUCreateOptions{Timeout: time.Minute}, + logf: func(string, ...interface{}) {}, + } + + parentCtx, cancel := tc.parentCtx() + defer cancel() + + err := ctx.pollUntilReady(parentCtx, "ws-1") + + require.Error(t, err) + if tc.wantErrIs != nil { + assert.ErrorIs(t, err, tc.wantErrIs) + } else { + assert.ErrorContains(t, err, tc.wantErrHas) + } + }) + } +} + +// TestWaitForInstances covers how the wait aggregates per-instance outcomes. +func TestWaitForInstances(t *testing.T) { + quietLogf := func(string, ...interface{}) {} + + cases := []struct { + name string + workspaces []*entity.Workspace + statuses map[string]devplanev1.EnvironmentStatus + instanceMsgs map[string]string + timeout time.Duration + parentCtx func() (context.Context, context.CancelFunc) + wantErr error + wantErrHas []string + wantNoError bool + }{ + { + name: "timeout on every instance is a warning, not an error", + workspaces: []*entity.Workspace{{ID: "ws-1", Name: "one"}}, + timeout: 0, // deadline already elapsed: immediate timeout + parentCtx: backgroundContext, + wantNoError: true, + }, + { + name: "terminal failure is reported with instance and message", + workspaces: []*entity.Workspace{{ID: "ws-failed", Name: "failed-instance"}}, + statuses: map[string]devplanev1.EnvironmentStatus{ + "ws-failed": devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + }, + instanceMsgs: map[string]string{ + "ws-failed": "out of quota in the region fulfill the request, InsufficientQuota", + }, + timeout: time.Second, + parentCtx: backgroundContext, + wantErrHas: []string{"1/1 instance(s) failed to become ready", "ws-failed", "InsufficientQuota"}, + }, + { + name: "cancellation aborts remaining instances", + workspaces: []*entity.Workspace{{ID: "ws-1", Name: "first"}, {ID: "ws-2", Name: "second"}}, + timeout: time.Minute, + parentCtx: canceledContext, + wantErr: context.Canceled, + }, + { + name: "expired parent deadline is not swallowed as warning", + workspaces: []*entity.Workspace{{ID: "ws-1", Name: "one"}}, + timeout: time.Minute, + parentCtx: expiredDeadlineContext, + wantErr: context.DeadlineExceeded, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + store := NewMockGPUCreateStore() + store.EnvironmentStatuses = tc.statuses + store.EnvironmentStatusMessages = tc.instanceMsgs + + ctx := &createContext{ + t: terminal.New(), + store: store, + opts: GPUCreateOptions{Timeout: tc.timeout}, + logf: quietLogf, + } + + parentCtx, cancel := tc.parentCtx() + defer cancel() + + err := ctx.waitForInstances(parentCtx, tc.workspaces) + + if tc.wantNoError { + assert.NoError(t, err) + return + } + require.Error(t, err) + if tc.wantErr != nil { + assert.ErrorIs(t, err, tc.wantErr) + } + for _, want := range tc.wantErrHas { + assert.ErrorContains(t, err, want) + } + }) } +} - ctx := &createContext{ - store: store, - opts: GPUCreateOptions{Timeout: time.Second}, +// TestRunGPUCreateLifecycle covers the end-to-end create flow per environment +// status outcome, including how ctx cancellation surfaces. +func TestRunGPUCreateLifecycle(t *testing.T) { + cases := []struct { + name string + status devplanev1.EnvironmentStatus + instanceMsg string + parentCtx func() (context.Context, context.CancelFunc) + wantErrIs error + wantErrHas []string + wantNoError bool + }{ + { + name: "ready environment succeeds", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING, + parentCtx: backgroundContext, + wantNoError: true, + }, + { + name: "failed environment reports failure", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + instanceMsg: "out of quota in the region fulfill the request, InsufficientQuota", + parentCtx: backgroundContext, + wantErrHas: []string{"failed to become ready", "InsufficientQuota"}, + }, + { + name: "canceled context aborts before the success summary", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING, + parentCtx: canceledContext, + wantErrIs: context.Canceled, + }, + { + name: "expired parent deadline aborts before the success summary", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING, + parentCtx: expiredDeadlineContext, + wantErrIs: context.DeadlineExceeded, + }, } - err := ctx.pollUntilReady("ws-failed") + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + mock := NewMockGPUCreateStore() + mock.EnvironmentStatuses["ws-repro"] = tc.status + if tc.instanceMsg != "" { + mock.EnvironmentStatusMessages = map[string]string{"ws-repro": tc.instanceMsg} + } - assert.ErrorContains(t, err, "instance test failed: unexpected end of JSON input") + parentCtx, cancel := tc.parentCtx() + defer cancel() + + err := RunGPUCreate(parentCtx, terminal.New(), mock, GPUCreateOptions{ + Name: "repro", + Count: 1, + Parallel: 1, + InstanceTypes: []InstanceSpec{{Type: "g5.xlarge"}}, + Timeout: time.Minute, + }) + + if tc.wantNoError { + require.NoError(t, err) + require.Len(t, mock.CreatedWorkspaces, 1) + return + } + require.Error(t, err) + if tc.wantErrIs != nil { + assert.ErrorIs(t, err, tc.wantErrIs) + } + for _, want := range tc.wantErrHas { + assert.ErrorContains(t, err, want) + } + }) + } } func TestInlineLaunchableLifeCycleScript(t *testing.T) { @@ -1212,116 +1369,93 @@ func TestValidateInstanceTypeAvailability(t *testing.T) { }) } -func TestCreateInstancesWithTypeSkipsInvalidType(t *testing.T) { - mock := NewMockGPUCreateStore() - ctx := &createContext{ - t: terminal.New(), - store: mock, - opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4"}, - org: mock.Org, - user: mock.User, - piped: true, - allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ - AllInstanceTypes: []gpusearch.InstanceType{ - { - Type: "hyperstack_H100_sxm5x8", - CloudCredID: "cc-shadeform", - }, - }, - }, +// TestCreateInstancesWithType covers per-type pre-flight validation outcomes +// and how the cloud credential is chosen. +func TestCreateInstancesWithType(t *testing.T) { + catalog := func(types ...gpusearch.InstanceType) *gpusearch.AllInstanceTypesResponse { + return &gpusearch.AllInstanceTypesResponse{AllInstanceTypes: types} } - ctx.logf = func(_ string, _ ...interface{}) {} - - result := ctx.createInstancesWithType(InstanceSpec{Type: "hyperstack_H100x8_one"}, 0, 1) - - assert.True(t, result.hadFailure, "expected hadFailure for an invalid instance type") - assert.Empty(t, result.successes, "expected no successes for invalid type") - assert.NoError(t, result.fatalError, "invalid type should not be fatal — caller may try the next type") - assert.Empty(t, mock.CreatedWorkspaces, "CreateWorkspace must not be called when the type is unrecognized") -} -func TestCreateInstancesWithTypeSkipsUnavailableType(t *testing.T) { - mock := NewMockGPUCreateStore() - ctx := &createContext{ - t: terminal.New(), - store: mock, - opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4"}, - org: mock.Org, - user: mock.User, - piped: true, - allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ - AllInstanceTypes: []gpusearch.InstanceType{ - {Type: "hyperstack_H100x8_NVLINK"}, - }, + cases := []struct { + name string + spec InstanceSpec + opts GPUCreateOptions + catalogTypes *gpusearch.AllInstanceTypesResponse + wantHadFail bool + wantSuccesses int + wantCreated int + wantCredID string + }{ + { + name: "unrecognized type fails without creating", + spec: InstanceSpec{Type: "hyperstack_H100x8_one"}, + catalogTypes: catalog(gpusearch.InstanceType{Type: "hyperstack_H100_sxm5x8", CloudCredID: "cc-shadeform"}), + wantHadFail: true, }, - } - ctx.logf = func(_ string, _ ...interface{}) {} - - result := ctx.createInstancesWithType(InstanceSpec{Type: "hyperstack_H100x8_NVLINK"}, 0, 1) - - assert.True(t, result.hadFailure, "expected hadFailure for an unavailable instance type") - assert.Empty(t, result.successes) - assert.Empty(t, mock.CreatedWorkspaces, "CreateWorkspace must not be called when no cloud credential is available") -} - -func TestCreateInstancesWithTypeSetsCloudCredIDFromCatalog(t *testing.T) { - mock := NewMockGPUCreateStore() - ctx := &createContext{ - t: terminal.New(), - store: mock, - opts: GPUCreateOptions{Count: 1, Parallel: 1, Name: "jt-4"}, - org: mock.Org, - user: mock.User, - piped: true, - allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ - AllInstanceTypes: []gpusearch.InstanceType{ - { - Type: "hyperstack_H100_sxm5x8", - CloudCredID: "cc-shadeform", + { + name: "type without cloud credential fails without creating", + spec: InstanceSpec{Type: "hyperstack_H100x8_NVLINK"}, + catalogTypes: catalog(gpusearch.InstanceType{Type: "hyperstack_H100x8_NVLINK"}), + wantHadFail: true, + }, + { + name: "catalog credential is applied to the create request", + spec: InstanceSpec{Type: "hyperstack_H100_sxm5x8"}, + catalogTypes: catalog(gpusearch.InstanceType{Type: "hyperstack_H100_sxm5x8", CloudCredID: "cc-shadeform"}), + wantSuccesses: 1, + wantCreated: 1, + wantCredID: "cc-shadeform", + }, + { + name: "launchable bypasses validation and uses launchable credential", + spec: InstanceSpec{Type: "n2-standard-4"}, + opts: GPUCreateOptions{ + LaunchableID: "env-abc", + LaunchableInfo: &store.LaunchableResponse{ + ID: "env-abc", + Name: "test-launchable", + CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{ + CloudCredID: "cc-from-launchable", + InstanceType: "n2-standard-4", + }, }, }, + catalogTypes: catalog(), // launchable's type is not in the org listing + wantSuccesses: 1, + wantCreated: 1, + wantCredID: "cc-from-launchable", }, } - ctx.logf = func(_ string, _ ...interface{}) {} - result := ctx.createInstancesWithType(InstanceSpec{Type: "hyperstack_H100_sxm5x8"}, 0, 1) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + mock := NewMockGPUCreateStore() + opts := tc.opts + opts.Count, opts.Parallel, opts.Name = 1, 1, "jt-4" + + ctx := &createContext{ + t: terminal.New(), + store: mock, + opts: opts, + org: mock.Org, + user: mock.User, + piped: true, + allInstanceTypes: tc.catalogTypes, + logf: func(string, ...interface{}) {}, + } - assert.False(t, result.hadFailure) - require.Len(t, mock.CreatedOptions, 1) - assert.Equal(t, "cc-shadeform", mock.CreatedOptions[0].CloudCredID) -} + result := ctx.createInstancesWithType(tc.spec, 0, 1) -func TestCreateInstancesWithTypeBypassesValidationForLaunchable(t *testing.T) { - mock := NewMockGPUCreateStore() - ctx := &createContext{ - t: terminal.New(), - store: mock, - opts: GPUCreateOptions{ - Count: 1, - Parallel: 1, - Name: "jt-4", - LaunchableID: "env-abc", - LaunchableInfo: &store.LaunchableResponse{ - ID: "env-abc", - Name: "test-launchable", - CreateWorkspaceRequest: store.LaunchableWorkspaceRequest{ - CloudCredID: "cc-from-launchable", - InstanceType: "n2-standard-4", - }, - }, - }, - org: mock.Org, - user: mock.User, - piped: true, - allInstanceTypes: &gpusearch.AllInstanceTypesResponse{ - AllInstanceTypes: []gpusearch.InstanceType{}, // launchable's type is not in the org listing - }, + assert.Equal(t, tc.wantHadFail, result.hadFailure) + assert.Len(t, result.successes, tc.wantSuccesses) + assert.Len(t, mock.CreatedWorkspaces, tc.wantCreated) + if tc.wantCredID != "" { + require.Len(t, mock.CreatedOptions, 1) + assert.Equal(t, tc.wantCredID, mock.CreatedOptions[0].CloudCredID) + } + if tc.wantHadFail { + assert.NoError(t, result.fatalError, "per-type failure should not be fatal — caller may try the next type") + } + }) } - ctx.logf = func(_ string, _ ...interface{}) {} - - result := ctx.createInstancesWithType(InstanceSpec{Type: "n2-standard-4"}, 0, 1) - - assert.False(t, result.hadFailure, "launchable should not be blocked by pre-flight validation") - assert.Len(t, result.successes, 1, "expected the launchable instance to be created") - assert.Len(t, mock.CreatedWorkspaces, 1) } diff --git a/pkg/cmd/grantssh/grantssh_test.go b/pkg/cmd/grantssh/grantssh_test.go index 1b2f6aa2e..c3b76e77b 100644 --- a/pkg/cmd/grantssh/grantssh_test.go +++ b/pkg/cmd/grantssh/grantssh_test.go @@ -15,6 +15,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -47,7 +48,7 @@ type mockNodeClientFactory struct { } func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient { - return register.NewNodeServiceClient(provider, m.serverURL) + return store.NewDevPlaneClient(provider, m.serverURL).ExternalNodes } // mockRegistrationStore satisfies register.RegistrationStore. diff --git a/pkg/cmd/launch/launch.go b/pkg/cmd/launch/launch.go index 40e5bf12c..275be4427 100644 --- a/pkg/cmd/launch/launch.go +++ b/pkg/cmd/launch/launch.go @@ -27,6 +27,7 @@ const defaultLaunchTimeoutSeconds = 300 type Store interface { gpucreate.GPUCreateStore GetAccessToken() (string, error) + DevPlane() *store.DevPlaneClient } type commandOptions struct { @@ -51,6 +52,7 @@ type launchCommandArgs struct { } type remoteLaunchArgs struct { + ctx context.Context terminal *terminal.Terminal store Store launchableID string @@ -178,6 +180,7 @@ func runLaunchCommand(ctx context.Context, args launchCommandArgs) error { }) } return runRemoteLaunch(remoteLaunchArgs{ + ctx: ctx, terminal: args.terminal, store: args.store, launchableID: launchableID, @@ -277,7 +280,7 @@ func runRemoteLaunch(args remoteLaunchArgs) error { if err != nil { return err } - err = gpucreate.RunGPUCreate(args.terminal, args.store, gpucreate.GPUCreateOptions{ + err = gpucreate.RunGPUCreate(args.ctx, args.terminal, args.store, gpucreate.GPUCreateOptions{ Name: args.name, InstanceTypes: instanceTypes, Count: 1, diff --git a/pkg/cmd/launch/launch_test.go b/pkg/cmd/launch/launch_test.go index ac60c398c..71541bbf9 100644 --- a/pkg/cmd/launch/launch_test.go +++ b/pkg/cmd/launch/launch_test.go @@ -13,7 +13,10 @@ import ( "strings" "testing" + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" @@ -119,6 +122,14 @@ type fakeLaunchStore struct { func (f *fakeLaunchStore) GetAccessToken() (string, error) { return "token", nil } +func (f *fakeLaunchStore) DevPlane() *store.DevPlaneClient { + return store.NewDevPlaneClient(f, config.GlobalConfig.GetBrevPublicAPIURL()) +} + +func (f *fakeLaunchStore) GetEnvironment(_ context.Context, environmentID string) (*devplanev1.Environment, error) { + return &devplanev1.Environment{EnvironmentId: environmentID}, nil +} + func (f *fakeLaunchStore) GetCurrentUser() (*entity.User, error) { return &entity.User{ID: "user-1", GlobalUserType: "Standard"}, nil } diff --git a/pkg/cmd/launch/secrets.go b/pkg/cmd/launch/secrets.go index 2e23afdcf..26cfe62a1 100644 --- a/pkg/cmd/launch/secrets.go +++ b/pkg/cmd/launch/secrets.go @@ -3,8 +3,6 @@ package launch import ( "context" - "github.com/brevdev/brev-cli/pkg/cmd/register" - "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/managedsecret" "github.com/brevdev/brev-cli/pkg/store" @@ -20,11 +18,9 @@ type devplaneManagedSecretResolver struct { client managedsecret.Client } -func newDevplaneManagedSecretResolver(provider Store) managedSecretResolver { +func newDevplaneManagedSecretResolver(launchStore Store) managedSecretResolver { return devplaneManagedSecretResolver{ - client: managedsecret.NewClient( - register.NewManagedSecretServiceClient(provider, config.GlobalConfig.GetBrevPublicAPIURL()), - ), + client: managedsecret.NewClient(launchStore.DevPlane().ManagedSecrets), } } diff --git a/pkg/cmd/ls/ls.go b/pkg/cmd/ls/ls.go index fcf09af0d..28bc863cb 100644 --- a/pkg/cmd/ls/ls.go +++ b/pkg/cmd/ls/ls.go @@ -8,7 +8,7 @@ import ( "os" "sync" - nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" "github.com/brevdev/brev-cli/pkg/analytics" @@ -18,11 +18,11 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/hello" - "github.com/brevdev/brev-cli/pkg/cmd/register" cmdutil "github.com/brevdev/brev-cli/pkg/cmd/util" "github.com/brevdev/brev-cli/pkg/cmdcontext" "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/environment" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/featureflag" "github.com/brevdev/brev-cli/pkg/store" @@ -41,6 +41,8 @@ type LsStore interface { GetWorkspace(workspaceID string) (*entity.Workspace, error) GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error) externalnode.TokenProvider + DevPlane() *store.DevPlaneClient + GetEnvironmentsByOrg(ctx context.Context, orgID string) ([]*devplanev1.Environment, error) GetAuthTokens() (*entity.AuthTokens, error) GetInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, error) hello.HelloStore @@ -89,7 +91,7 @@ with other commands like stop, start, or delete.`, if err != nil { return breverrors.WrapAndTrace(err) } - err = RunLs(t, cliAuth, loginLsStore, args, showAll, jsonOutput) + err = RunLs(cmd.Context(), t, cliAuth, loginLsStore, args, showAll, jsonOutput) if err != nil { return breverrors.WrapAndTrace(err) } @@ -130,7 +132,39 @@ func getOrgForRunLs(lsStore LsStore) (*entity.Organization, error) { return org, nil } -func RunLs(t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []string, showAll bool, jsonOutput bool) error { +// overlayStatuses applies the environment lifecycle status to the workspace +// rows, returning any status-fetch error rather than reporting a stale or +// missing status. +func (ls Ls) overlayStatuses(workspaces []entity.Workspace, environments []*devplanev1.Environment, fetchErr error) error { + if fetchErr != nil { + return breverrors.WrapAndTrace(fetchErr) + } + applyStatuses(workspaces, environments) + return nil +} + +// applyStatuses sets each workspace's status and message from the matching +// environment, keyed by workspace ID. +func applyStatuses(workspaces []entity.Workspace, environments []*devplanev1.Environment) { + byID := make(map[string]*devplanev1.Environment, len(environments)) + for _, env := range environments { + byID[env.GetEnvironmentId()] = env + } + for i := range workspaces { + env, ok := byID[workspaces[i].ID] + if !ok { + continue + } + if display := environment.Display(env); display != "" { + workspaces[i].Status = display + } + if message := environment.FailureMessage(env); message != "" { + workspaces[i].StatusMessage = message + } + } +} + +func RunLs(ctx context.Context, t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []string, showAll bool, jsonOutput bool) error { ls := NewLs(lsStore, t, jsonOutput) org, err := getOrgForRunLs(lsStore) @@ -142,12 +176,12 @@ func RunLs(t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []s } if len(args) == 1 { //nolint:gocritic // don't want to switch - err = handleLsArg(ls, cliAuth, args[0], org, showAll) + err = handleLsArg(ctx, ls, cliAuth, args[0], org, showAll) if err != nil { return breverrors.WrapAndTrace(err) } } else if len(args) == 0 { - err = ls.RunWorkspaces(cliAuth, org, showAll) + err = ls.RunWorkspaces(ctx, cliAuth, org, showAll) if err != nil { return breverrors.WrapAndTrace(err) } @@ -158,7 +192,7 @@ func RunLs(t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []s return nil } -func handleLsArg(ls *Ls, cliAuth auth.CLIAuth, arg string, org *entity.Organization, showAll bool) error { +func handleLsArg(ctx context.Context, ls *Ls, cliAuth auth.CLIAuth, arg string, org *entity.Organization, showAll bool) error { switch classifyLsArg(arg) { case lsArgOrgs: if cliAuth.IsAPIKey() { @@ -166,7 +200,7 @@ func handleLsArg(ls *Ls, cliAuth auth.CLIAuth, arg string, org *entity.Organizat } return wrapLsRun(ls.RunOrgs()) case lsArgWorkspaces: - return wrapLsRun(ls.RunWorkspaces(cliAuth, org, showAll)) + return wrapLsRun(ls.RunWorkspaces(ctx, cliAuth, org, showAll)) case lsArgUsers: return runAdminLsArg(cliAuth, "users", func() error { return ls.RunUser(showAll) @@ -178,7 +212,7 @@ func handleLsArg(ls *Ls, cliAuth auth.CLIAuth, arg string, org *entity.Organizat case lsArgNodes: return wrapLsRun(ls.RunNodes(org)) case lsArgInstances: - return wrapLsRun(ls.RunInstances(cliAuth, org, showAll)) + return wrapLsRun(ls.RunInstances(ctx, cliAuth, org, showAll)) default: return nil } @@ -405,15 +439,16 @@ func buildGPULookup(s LsStore) map[string]string { return lookup } -func (ls Ls) RunWorkspaces(cliAuth auth.CLIAuth, org *entity.Organization, showAll bool) error { +func (ls Ls) RunWorkspaces(ctx context.Context, cliAuth auth.CLIAuth, org *entity.Organization, showAll bool) error { // Fetch workspaces and instance types concurrently var allWorkspaces []entity.Workspace var wsErr error var gpuLookup map[string]string - var nodes []*nodev1.ExternalNode - + var nodes []*devplanev1.ExternalNode + var environments []*devplanev1.Environment + var envErr error var wg sync.WaitGroup - wg.Add(2) + wg.Add(3) if showAll { wg.Add(1) } @@ -425,6 +460,10 @@ func (ls Ls) RunWorkspaces(cliAuth auth.CLIAuth, org *entity.Organization, showA defer wg.Done() gpuLookup = buildGPULookup(ls.lsStore) }() + go func() { + defer wg.Done() + environments, envErr = ls.lsStore.GetEnvironmentsByOrg(ctx, org.ID) + }() if showAll { go func() { defer wg.Done() @@ -443,6 +482,10 @@ func (ls Ls) RunWorkspaces(cliAuth auth.CLIAuth, org *entity.Organization, showA return breverrors.WrapAndTrace(wsErr) } + if err := ls.overlayStatuses(allWorkspaces, environments, envErr); err != nil { + return err + } + // Determine which workspaces to show var workspacesToShow []entity.Workspace switch { @@ -490,8 +533,8 @@ func (ls Ls) RunWorkspaces(cliAuth auth.CLIAuth, org *entity.Organization, showA return nil } -func (ls Ls) RunInstances(cliAuth auth.CLIAuth, org *entity.Organization, showAll bool) error { - if err := ls.RunWorkspaces(cliAuth, org, showAll); err != nil { +func (ls Ls) RunInstances(ctx context.Context, cliAuth auth.CLIAuth, org *entity.Organization, showAll bool) error { + if err := ls.RunWorkspaces(ctx, cliAuth, org, showAll); err != nil { return err } return nil @@ -539,7 +582,7 @@ func getInstanceTypeAndKind(w entity.Workspace, gpuLookup map[string]string) (st return "", "" } -func toNodeInfos(nodes []*nodev1.ExternalNode) []NodeInfo { +func toNodeInfos(nodes []*devplanev1.ExternalNode) []NodeInfo { var infos []NodeInfo for _, n := range nodes { infos = append(infos, NodeInfo{ @@ -551,7 +594,7 @@ func toNodeInfos(nodes []*nodev1.ExternalNode) []NodeInfo { return infos } -func (ls Ls) outputWorkspacesJSON(workspaces []entity.Workspace, gpuLookup map[string]string, nodes []*nodev1.ExternalNode) error { +func (ls Ls) outputWorkspacesJSON(workspaces []entity.Workspace, gpuLookup map[string]string, nodes []*devplanev1.ExternalNode) error { var wsInfos []WorkspaceInfo for _, w := range workspaces { instanceType, instanceKind := getInstanceTypeAndKind(w, gpuLookup) @@ -726,9 +769,9 @@ type NodeInfo struct { Status string `json:"status"` } -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{ +func (ls Ls) listNodes(org *entity.Organization) ([]*devplanev1.ExternalNode, error) { + client := ls.lsStore.DevPlane().ExternalNodes + resp, err := client.ListNodes(context.Background(), connect.NewRequest(&devplanev1.ListNodesRequest{ OrganizationId: org.ID, })) if err != nil { @@ -766,7 +809,7 @@ func (ls Ls) RunNodes(org *entity.Organization) error { return nil } -func (ls Ls) outputNodesJSON(nodes []*nodev1.ExternalNode) error { +func (ls Ls) outputNodesJSON(nodes []*devplanev1.ExternalNode) error { output, err := json.MarshalIndent(toNodeInfos(nodes), "", " ") if err != nil { return breverrors.WrapAndTrace(err) @@ -775,7 +818,7 @@ func (ls Ls) outputNodesJSON(nodes []*nodev1.ExternalNode) error { return nil } -func displayNodesTable(t *terminal.Terminal, nodes []*nodev1.ExternalNode, isPiped bool) { +func displayNodesTable(t *terminal.Terminal, nodes []*devplanev1.ExternalNode, isPiped bool) { ta := table.NewWriter() ta.SetOutputMirror(os.Stdout) ta.Style().Options = getBrevTableOptions() @@ -790,7 +833,7 @@ func displayNodesTable(t *terminal.Terminal, nodes []*nodev1.ExternalNode, isPip ta.Render() } -func nodeConnectionStatus(n *nodev1.ExternalNode) string { +func nodeConnectionStatus(n *devplanev1.ExternalNode) string { ci := n.GetConnectivityInfo() if ci == nil { return "Unknown" diff --git a/pkg/cmd/ls/ls_test.go b/pkg/cmd/ls/ls_test.go index b2322f58e..8c6b5c240 100644 --- a/pkg/cmd/ls/ls_test.go +++ b/pkg/cmd/ls/ls_test.go @@ -8,20 +8,47 @@ import ( "strings" "testing" - nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" - nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + devplanev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" authpkg "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" + "github.com/stretchr/testify/assert" ) const testAPIKey = authpkg.BrevAPIKeyPrefix + "test-key" +func TestApplyStatusesFeedsDisplayFromEnvironment(t *testing.T) { + workspaces := []entity.Workspace{ + {ID: "env-failed", Name: "failed", Status: entity.Running}, + {ID: "env-running", Name: "running", Status: entity.Deploying}, + {ID: "env-unknown", Name: "unknown", Status: entity.Deploying}, + } + failedEnv := &devplanev1.Environment{EnvironmentId: "env-failed", Status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED} + failedEnv.Instance = &devplanev1.Instance{ + ProvisionStatus: &devplanev1.CreateAttemptStatus{Message: "out of quota"}, + } + environments := []*devplanev1.Environment{ + failedEnv, + {EnvironmentId: "env-running", Status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING}, + } + + applyStatuses(workspaces, environments) + + // The environment status wins even when the workspace row said otherwise. + assert.Equal(t, entity.Failure, workspaces[0].Status) + assert.Equal(t, "out of quota", workspaces[0].StatusMessage) + assert.Equal(t, entity.Running, workspaces[1].Status) + // Workspaces without a matching environment keep their (stale) row status. + assert.Equal(t, entity.Deploying, workspaces[2].Status) +} + // mockLsStore implements LsStore (including the embedded hello.HelloStore) for // testing the ls command routing without real API calls. type mockLsStore struct { @@ -48,6 +75,14 @@ func (m *mockLsStore) GetAccessToken() (string, error) { return "tok", nil } +func (m *mockLsStore) DevPlane() *store.DevPlaneClient { + return store.NewDevPlaneClient(m, config.GlobalConfig.GetBrevPublicAPIURL()) +} + +func (m *mockLsStore) GetEnvironmentsByOrg(_ context.Context, _ string) ([]*devplanev1.Environment, error) { + return nil, nil +} + func (m *mockLsStore) GetWorkspace(_ string) (*entity.Workspace, error) { return nil, nil } @@ -112,7 +147,7 @@ func resolveTestCLIAuth(t *testing.T, s *mockLsStore) authpkg.CLIAuth { func runLs(t *testing.T, term *terminal.Terminal, s *mockLsStore, args []string, showAll bool) error { t.Helper() - return RunLs(term, resolveTestCLIAuth(t, s), s, args, showAll, true) + return RunLs(context.Background(), term, resolveTestCLIAuth(t, s), s, args, showAll, true) } func TestRunLs_APIKeyJSONSkipsUserAndOrgList(t *testing.T) { @@ -424,7 +459,7 @@ func TestRunLs_ShowAllTable(t *testing.T) { term := terminal.New() out := captureStdout(t, func() { - err := RunLs(term, resolveTestCLIAuth(t, s), s, nil, true, false) + err := RunLs(context.Background(), term, resolveTestCLIAuth(t, s), s, nil, true, false) if err != nil { t.Fatalf("RunLs --all returned error: %v", err) } @@ -598,7 +633,7 @@ func TestHandleLsArg_Routing(t *testing.T) { "workspace", "workspaces", } for _, arg := range successArgs { - if err := handleLsArg(ls, resolveTestCLIAuth(t, s), arg, s.org, false); err != nil { + if err := handleLsArg(context.Background(), ls, resolveTestCLIAuth(t, s), arg, s.org, false); err != nil { t.Errorf("handleLsArg(%q) returned unexpected error: %v", arg, err) } } @@ -606,7 +641,7 @@ func TestHandleLsArg_Routing(t *testing.T) { // "node"/"nodes" route to RunNodes which calls the gRPC client — verify // it attempts the path (error expected due to no real client). for _, arg := range []string{"node", "nodes"} { - _ = handleLsArg(ls, resolveTestCLIAuth(t, s), arg, s.org, false) + _ = handleLsArg(context.Background(), ls, resolveTestCLIAuth(t, s), arg, s.org, false) } } @@ -626,7 +661,7 @@ func TestOutputWorkspacesJSON_WithNodes(t *testing.T) { CreatedByUserID: "u1", }, } - nodes := []*nodev1.ExternalNode{ + nodes := []*devplanev1.ExternalNode{ { Name: "gpu-node-1", ExternalNodeId: "en-123", @@ -711,15 +746,15 @@ func TestOutputWorkspacesJSON_NoNodes(t *testing.T) { } type lsFakeNodeService struct { - nodev1connect.UnimplementedExternalNodeServiceHandler + devplanev1connect.UnimplementedExternalNodeServiceHandler } -func (f *lsFakeNodeService) ListNodes(_ context.Context, req *connect.Request[nodev1.ListNodesRequest]) (*connect.Response[nodev1.ListNodesResponse], error) { +func (f *lsFakeNodeService) ListNodes(_ context.Context, req *connect.Request[devplanev1.ListNodesRequest]) (*connect.Response[devplanev1.ListNodesResponse], error) { if req.Msg.GetOrganizationId() != "org1" { - return connect.NewResponse(&nodev1.ListNodesResponse{}), nil + return connect.NewResponse(&devplanev1.ListNodesResponse{}), nil } - return connect.NewResponse(&nodev1.ListNodesResponse{ - Items: []*nodev1.ExternalNode{ + return connect.NewResponse(&devplanev1.ListNodesResponse{ + Items: []*devplanev1.ExternalNode{ {Name: "gpu-node-1", ExternalNodeId: "en-123", OrganizationId: "org1"}, {Name: "gpu-node-2", ExternalNodeId: "en-456", OrganizationId: "org1"}, }, @@ -728,7 +763,7 @@ func (f *lsFakeNodeService) ListNodes(_ context.Context, req *connect.Request[no func withFakeNodeAPI(t *testing.T, fn func()) { t.Helper() - _, handler := nodev1connect.NewExternalNodeServiceHandler(&lsFakeNodeService{}) + _, handler := devplanev1connect.NewExternalNodeServiceHandler(&lsFakeNodeService{}) server := httptest.NewServer(handler) t.Cleanup(server.Close) t.Setenv("BREV_PUBLIC_API_URL", server.URL) diff --git a/pkg/cmd/mintcert/mintcert.go b/pkg/cmd/mintcert/mintcert.go index f2b7d4a6f..a0714aef4 100644 --- a/pkg/cmd/mintcert/mintcert.go +++ b/pkg/cmd/mintcert/mintcert.go @@ -15,11 +15,11 @@ import ( "github.com/spf13/afero" "github.com/spf13/cobra" - "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/externalnode" "github.com/brevdev/brev-cli/pkg/sshcert" + "github.com/brevdev/brev-cli/pkg/store" ) const timeout = 15 * time.Second @@ -191,8 +191,9 @@ func runMintCertWith(store Store, fs afero.Fs, issuer CertIssuer, req mintCertRe } func newCertIssuer(provider externalnode.TokenProvider, baseURL string) CertIssuer { + devPlane := store.NewDevPlaneClient(provider, baseURL) return rpcCertIssuer{ - client: register.NewEnvironmentServiceClient(provider, baseURL), - nodeClient: register.NewNodeServiceClient(provider, baseURL), + client: devPlane.Environments, + nodeClient: devPlane.ExternalNodes, } } diff --git a/pkg/cmd/ports/close.go b/pkg/cmd/ports/close.go index b35b989f4..07fbf0e44 100644 --- a/pkg/cmd/ports/close.go +++ b/pkg/cmd/ports/close.go @@ -11,9 +11,7 @@ import ( "github.com/spf13/cobra" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/register" cmdutil "github.com/brevdev/brev-cli/pkg/cmd/util" - "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" ) @@ -146,14 +144,14 @@ func closePort( portID string, ) error { if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments _, err := client.ClosePort(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceClosePortRequest{ PortId: portID, })) return breverrors.WrapAndTrace(err) } if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes _, err := client.ClosePort(ctx, connect.NewRequest(&devplanev1.ClosePortRequest{ PortId: portID, })) diff --git a/pkg/cmd/ports/open.go b/pkg/cmd/ports/open.go index 02edd0886..cea2dc85f 100644 --- a/pkg/cmd/ports/open.go +++ b/pkg/cmd/ports/open.go @@ -15,9 +15,7 @@ import ( "github.com/spf13/cobra" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/register" cmdutil "github.com/brevdev/brev-cli/pkg/cmd/util" - "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" ) @@ -190,7 +188,7 @@ func OpenHTTP( if err != nil { return err } - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.OpenHTTPPort(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceOpenHTTPPortRequest{ EnvironmentId: target.Workspace.ID, PortNumber: portNumber, @@ -210,7 +208,7 @@ func OpenHTTP( if err != nil { return err } - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.OpenHTTPPort(ctx, connect.NewRequest(&devplanev1.OpenHTTPPortRequest{ ExternalNodeId: target.Node.GetExternalNodeId(), PortNumber: portNumber, @@ -251,7 +249,7 @@ func Open( var openedPort *devplanev1.Port if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.OpenPort(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceOpenPortRequest{ EnvironmentId: target.Workspace.ID, Protocol: protocol, @@ -265,7 +263,7 @@ func Open( openedPort = resp.Msg.GetPort() } } else if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.OpenPort(ctx, connect.NewRequest(&devplanev1.OpenPortRequest{ ExternalNodeId: target.Node.GetExternalNodeId(), Protocol: protocol, @@ -306,7 +304,7 @@ func OpenSequential( var openedPorts []*devplanev1.Port if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.OpenSequentialPorts(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceOpenSequentialPortsRequest{ EnvironmentId: target.Workspace.ID, Protocol: protocol, @@ -321,7 +319,7 @@ func OpenSequential( openedPorts = resp.Msg.GetPorts() } } else if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.OpenSequentialPorts(ctx, connect.NewRequest(&devplanev1.OpenSequentialPortsRequest{ ExternalNodeId: target.Node.GetExternalNodeId(), Protocol: protocol, diff --git a/pkg/cmd/ports/ports.go b/pkg/cmd/ports/ports.go index 7339e3dc5..277d8a7cb 100644 --- a/pkg/cmd/ports/ports.go +++ b/pkg/cmd/ports/ports.go @@ -17,9 +17,7 @@ import ( "github.com/spf13/cobra" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/register" cmdutil "github.com/brevdev/brev-cli/pkg/cmd/util" - "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" ) @@ -121,7 +119,7 @@ func resolveTargetPorts( var apiPorts []*devplanev1.Port if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.GetNetworkInfo(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceGetNetworkInfoRequest{ EnvironmentId: target.Workspace.ID, })) @@ -167,7 +165,7 @@ func resolvePortID( //nolint:gocyclo // A port may belong to either supported ow if err != nil { return nil, nil, breverrors.WrapAndTrace(err) } - environmentClient := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + environmentClient := portStore.DevPlane().Environments var lookupErr error for i := range workspaces { workspace := &workspaces[i] @@ -192,7 +190,7 @@ func resolvePortID( //nolint:gocyclo // A port may belong to either supported ow if err != nil { return nil, nil, breverrors.WrapAndTrace(err) } - nodeClient := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + nodeClient := portStore.DevPlane().ExternalNodes nodesResp, err := nodeClient.ListNodes(ctx, connect.NewRequest(&devplanev1.ListNodesRequest{ OrganizationId: org.ID, Options: &devplanev1.ListNodesOptions{ diff --git a/pkg/cmd/ports/ports_test.go b/pkg/cmd/ports/ports_test.go index e47f83cb8..fe3f9626f 100644 --- a/pkg/cmd/ports/ports_test.go +++ b/pkg/cmd/ports/ports_test.go @@ -12,7 +12,9 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/store" ) type fakeStore struct { @@ -45,6 +47,10 @@ func (s *fakeStore) GetAccessToken() (string, error) { return "test-token", nil } +func (s *fakeStore) DevPlane() *store.DevPlaneClient { + return store.NewDevPlaneClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) +} + type fakeEnvironmentService struct { devplanev1connect.UnimplementedEnvironmentServiceHandler t *testing.T diff --git a/pkg/cmd/ports/update.go b/pkg/cmd/ports/update.go index cc4a2183e..45074e69d 100644 --- a/pkg/cmd/ports/update.go +++ b/pkg/cmd/ports/update.go @@ -10,9 +10,7 @@ import ( "github.com/spf13/cobra" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/register" cmdutil "github.com/brevdev/brev-cli/pkg/cmd/util" - "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" ) @@ -266,7 +264,7 @@ func setPortTarget( destinationPort int32, ) (*devplanev1.Port, error) { if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.SetPortTarget(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceSetPortTargetRequest{ PortId: portID, PortNumber: destinationPort, })) @@ -279,7 +277,7 @@ func setPortTarget( return resp.Msg.GetPort(), nil } if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.SetPortTarget(ctx, connect.NewRequest(&devplanev1.SetPortTargetRequest{ PortId: portID, PortNumber: destinationPort, })) @@ -302,7 +300,7 @@ func setPortAllowedSources( allowedSources []string, ) (*devplanev1.Port, error) { if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.SetPortAllowedSources(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceSetPortAllowedSourcesRequest{ PortId: portID, AllowedSources: &devplanev1.EnvironmentServiceSetPortAllowedSourcesRequestAllowedSources{ @@ -318,7 +316,7 @@ func setPortAllowedSources( return resp.Msg.GetPort(), nil } if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.SetPortAllowedSources(ctx, connect.NewRequest(&devplanev1.SetPortAllowedSourcesRequest{ PortId: portID, AllowedSources: allowedSources, })) @@ -342,7 +340,7 @@ func setHTTPPortProtocol( protocol devplanev1.HttpPortProtocol, ) (*devplanev1.Port, error) { if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.SetHTTPPortProtocol(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceSetHTTPPortProtocolRequest{ PortId: portID, HttpProtocol: protocol, })) @@ -355,7 +353,7 @@ func setHTTPPortProtocol( return resp.Msg.GetPort(), nil } if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.SetHTTPPortProtocol(ctx, connect.NewRequest(&devplanev1.SetHTTPPortProtocolRequest{ PortId: portID, HttpProtocol: protocol, })) @@ -379,7 +377,7 @@ func setHTTPPortAccess( public bool, ) (*devplanev1.Port, error) { if target.Workspace != nil { - client := register.NewEnvironmentServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().Environments resp, err := client.SetHTTPPortAccess(ctx, connect.NewRequest(&devplanev1.EnvironmentServiceSetHTTPPortAccessRequest{ PortId: portID, AuthorizedEmails: &devplanev1.EnvironmentServiceSetHTTPPortAccessRequestAuthorizedEmails{ @@ -396,7 +394,7 @@ func setHTTPPortAccess( return resp.Msg.GetPort(), nil } if target.Node != nil { - client := register.NewNodeServiceClient(portStore, config.GlobalConfig.GetBrevPublicAPIURL()) + client := portStore.DevPlane().ExternalNodes resp, err := client.SetHTTPPortAccess(ctx, connect.NewRequest(&devplanev1.SetHTTPPortAccessRequest{ PortId: portID, AuthorizedEmails: authorizedEmails, diff --git a/pkg/cmd/refresh/refresh.go b/pkg/cmd/refresh/refresh.go index 07f0d496e..073730eec 100644 --- a/pkg/cmd/refresh/refresh.go +++ b/pkg/cmd/refresh/refresh.go @@ -12,10 +12,8 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" - "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/cmd/util" "github.com/brevdev/brev-cli/pkg/cmdcontext" - "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/ssh" @@ -37,6 +35,7 @@ type RefreshStore interface { MkdirAll(string, fs.FileMode) error GetBrevCloudflaredBinaryPath() (string, error) Create(string) (io.WriteCloser, error) + DevPlane() *store.DevPlaneClient } func NewCmdRefresh(t *terminal.Terminal, store RefreshStore) *cobra.Command { @@ -176,7 +175,7 @@ func getExternalNodeSSHEntries(store RefreshStore, identity refreshIdentity) []s } org, user := identity.org, identity.user - client := register.NewNodeServiceClient(store, config.GlobalConfig.GetBrevPublicAPIURL()) + client := store.DevPlane().ExternalNodes resp, err := client.ListNodes(context.Background(), connect.NewRequest(&nodev1.ListNodesRequest{ OrganizationId: org.ID, Options: &nodev1.ListNodesOptions{ diff --git a/pkg/cmd/refresh/refresh_test.go b/pkg/cmd/refresh/refresh_test.go index a3c8bce06..9ccdffbeb 100644 --- a/pkg/cmd/refresh/refresh_test.go +++ b/pkg/cmd/refresh/refresh_test.go @@ -8,7 +8,9 @@ import ( "github.com/stretchr/testify/require" "github.com/brevdev/brev-cli/pkg/cmd/util" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/store" ) func strPtr(s string) *string { return &s } @@ -40,6 +42,10 @@ func (s *identityCountingStore) GetContextWorkspaces() ([]entity.Workspace, erro return s.workspaces, nil } +func (s *identityCountingStore) DevPlane() *store.DevPlaneClient { + return store.NewDevPlaneClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) +} + func (s *identityCountingStore) GetContextWorkspacesFor(orgID, userID string) ([]entity.Workspace, error) { s.forIdentity = append(s.forIdentity, orgID+"/"+userID) return s.workspaces, nil diff --git a/pkg/cmd/refresh/sshaccess.go b/pkg/cmd/refresh/sshaccess.go index 8ddb4e4c7..caee70e32 100644 --- a/pkg/cmd/refresh/sshaccess.go +++ b/pkg/cmd/refresh/sshaccess.go @@ -9,8 +9,6 @@ import ( devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" - "github.com/brevdev/brev-cli/pkg/cmd/register" - "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/sshcert" @@ -48,7 +46,7 @@ func (s workspaceSSHStore) GetContextWorkspaces() ([]entity.Workspace, error) { } } - client := register.NewEnvironmentServiceClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) + client := s.DevPlane().Environments ctx, cancel := context.WithTimeout(context.Background(), sshAccessRefreshBudget) defer cancel() diff --git a/pkg/cmd/register/providers.go b/pkg/cmd/register/providers.go index 6970b80e1..e2bdbc17f 100644 --- a/pkg/cmd/register/providers.go +++ b/pkg/cmd/register/providers.go @@ -11,6 +11,7 @@ import ( nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -93,5 +94,5 @@ func (ShellSetupRunner) RunSetup(script string) error { return runSetupCommand(s type DefaultNodeClientFactory struct{} func (DefaultNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, baseURL string) nodev1connect.ExternalNodeServiceClient { - return NewNodeServiceClient(provider, baseURL) + return store.NewDevPlaneClient(provider, baseURL).ExternalNodes } diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 9da4b170f..98c6e57ca 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -516,3 +516,76 @@ func resumeRegistration(ctx context.Context, t *terminal.Terminal, s RegisterSto org := &entity.Organization{ID: pending.OrgID, Name: pending.OrgName} return runRegisterSteps(ctx, t, s, pending.DisplayName, org, deps, pending.DeviceID, registrationToken) } + +// toProtoNodeSpec converts the local HardwareProfile (used for collection, display, +// persistence) to the generated proto NodeSpec for RPC calls. +func toProtoNodeSpec(hw *HardwareProfile) *nodev1.NodeSpec { + if hw == nil { + return nil + } + + proto := &nodev1.NodeSpec{ + RamBytes: hw.RAMBytes, + CpuCount: hw.CPUCount, + } + + for _, st := range hw.Storage { + storageSpec := &nodev1.StorageSpec{ + StorageBytes: st.StorageBytes, + StorageType: st.StorageType, + } + if st.Name != "" { + storageSpec.Device = &st.Name + } + proto.Storage = append(proto.Storage, storageSpec) + } + + if hw.Architecture != "" { + proto.Architecture = &hw.Architecture + } + if hw.OS != "" { + proto.Os = &hw.OS + } + if hw.OSVersion != "" { + proto.OsVersion = &hw.OSVersion + } + + if hw.ProductName != "" { + proto.ProductName = &hw.ProductName + } + + for _, g := range hw.GPUs { + pg := &nodev1.GPUSpec{ + Model: g.Model, + Count: g.Count, + MemoryBytes: g.MemoryBytes, + } + if g.Architecture != "" { + pg.GpuArchitecture = &g.Architecture + } + proto.Gpus = append(proto.Gpus, pg) + } + + for _, ic := range hw.Interconnects { + spec := &nodev1.InterconnectSpec{Device: ic.Device} + switch ic.Type { + case "NVLink": + spec.Details = &nodev1.InterconnectSpec_Nvlink{ + Nvlink: &nodev1.NVLinkDetails{ + ActiveLinks: int32(ic.ActiveLinks), + Version: ic.Version, + }, + } + case "PCIe": + spec.Details = &nodev1.InterconnectSpec_Pcie{ + Pcie: &nodev1.PCIeDetails{ + Generation: int32(ic.Generation), + Width: int32(ic.Width), + }, + } + } + proto.Interconnects = append(proto.Interconnects, spec) + } + + return proto +} diff --git a/pkg/cmd/register/register_test.go b/pkg/cmd/register/register_test.go index 68a7a95b8..df32729c6 100644 --- a/pkg/cmd/register/register_test.go +++ b/pkg/cmd/register/register_test.go @@ -15,6 +15,7 @@ import ( "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/sudo" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -154,7 +155,7 @@ type mockNodeClientFactory struct { } func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient { - return NewNodeServiceClient(provider, m.serverURL) + return store.NewDevPlaneClient(provider, m.serverURL).ExternalNodes } // testHardwareProfile returns a realistic HardwareProfile for use in tests. @@ -984,3 +985,301 @@ func Test_isAuthenticated(t *testing.T) { } }) } + +type staticTokenProvider struct{ token string } + +func (p staticTokenProvider) GetAccessToken() (string, error) { return p.token, nil } + +func Test_toProtoNodeSpec(t *testing.T) { + cpuCount := int32(12) + ramBytes := int64(137438953472) + memBytes := int64(137438953472) + + local := &HardwareProfile{ + GPUs: []GPU{ + {Model: "NVIDIA GB10", Architecture: "Blackwell", Count: 2, MemoryBytes: &memBytes}, + }, + RAMBytes: &ramBytes, + CPUCount: &cpuCount, + Architecture: "arm64", + Storage: []StorageDevice{ + {Name: "nvme0n1", StorageBytes: 500107862016, StorageType: "SSD"}, + }, + OS: "Ubuntu", + OSVersion: "24.04", + ProductName: "DGX Spark", + Interconnects: []Interconnect{ + {Type: "NVLink", Device: "GPU 0", ActiveLinks: 4, Version: 4}, + }, + } + + proto := toProtoNodeSpec(local) + + if proto.GetCpuCount() != 12 { + t.Errorf("expected CpuCount 12, got %d", proto.GetCpuCount()) + } + if proto.GetRamBytes() != 137438953472 { + t.Errorf("expected RamBytes, got %d", proto.GetRamBytes()) + } + if proto.GetArchitecture() != "arm64" { + t.Errorf("expected arm64, got %s", proto.GetArchitecture()) + } + if proto.GetOs() != "Ubuntu" { + t.Errorf("expected Ubuntu, got %s", proto.GetOs()) + } + if proto.GetOsVersion() != "24.04" { + t.Errorf("expected 24.04, got %s", proto.GetOsVersion()) + } + if len(proto.GetStorage()) != 1 { + t.Fatalf("expected 1 storage entry, got %d", len(proto.GetStorage())) + } + if proto.GetStorage()[0].GetStorageBytes() != 500107862016 { + t.Errorf("expected StorageBytes 500107862016, got %d", proto.GetStorage()[0].GetStorageBytes()) + } + if proto.GetStorage()[0].GetStorageType() != "SSD" { + t.Errorf("expected SSD, got %s", proto.GetStorage()[0].GetStorageType()) + } + if len(proto.GetGpus()) != 1 { + t.Fatalf("expected 1 GPU, got %d", len(proto.GetGpus())) + } + gpu := proto.GetGpus()[0] + if gpu.GetModel() != "NVIDIA GB10" { + t.Errorf("expected NVIDIA GB10, got %s", gpu.GetModel()) + } + if gpu.GetCount() != 2 { + t.Errorf("expected count 2, got %d", gpu.GetCount()) + } + if gpu.GetMemoryBytes() != 137438953472 { + t.Errorf("expected memory bytes, got %d", gpu.GetMemoryBytes()) + } +} + +func Test_toProtoNodeSpec_PCIeInterconnect(t *testing.T) { + local := &HardwareProfile{ + Architecture: "amd64", + Interconnects: []Interconnect{ + {Type: "PCIe", Device: "GPU 0", Generation: 4, Width: 16}, + }, + } + + proto := toProtoNodeSpec(local) + + if len(proto.GetInterconnects()) != 1 { + t.Fatalf("expected 1 interconnect, got %d", len(proto.GetInterconnects())) + } + ic := proto.GetInterconnects()[0] + if ic.GetDevice() != "GPU 0" { + t.Errorf("expected device 'GPU 0', got %s", ic.GetDevice()) + } + pcie := ic.GetPcie() + if pcie == nil { + t.Fatal("expected PCIe details, got nil") + } + if pcie.GetGeneration() != 4 { + t.Errorf("expected PCIe Gen 4, got %d", pcie.GetGeneration()) + } + if pcie.GetWidth() != 16 { + t.Errorf("expected PCIe Width 16, got %d", pcie.GetWidth()) + } +} + +func Test_toProtoNodeSpec_Nil(t *testing.T) { + if toProtoNodeSpec(nil) != nil { + t.Error("expected nil for nil input") + } +} + +func Test_toProtoNodeSpec_MinimalFields(t *testing.T) { + local := &HardwareProfile{ + Architecture: "amd64", + } + proto := toProtoNodeSpec(local) + if proto.GetArchitecture() != "amd64" { + t.Errorf("expected amd64, got %s", proto.GetArchitecture()) + } + if proto.RamBytes != nil { + t.Error("expected nil RamBytes") + } + if proto.CpuCount != nil { + t.Error("expected nil CpuCount") + } + if len(proto.GetGpus()) != 0 { + t.Error("expected no GPUs") + } +} + +// fakeNodeService implements the server side of ExternalNodeService for testing. +type fakeNodeService struct { + nodev1connect.UnimplementedExternalNodeServiceHandler + addNodeFn func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) + removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) + getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) + openPortFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) + grantNodeSSHAccessFn func(*nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) +} + +func (f *fakeNodeService) AddNode(_ context.Context, req *connect.Request[nodev1.AddNodeRequest]) (*connect.Response[nodev1.AddNodeResponse], error) { + resp, err := f.addNodeFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (f *fakeNodeService) RemoveNode(_ context.Context, req *connect.Request[nodev1.RemoveNodeRequest]) (*connect.Response[nodev1.RemoveNodeResponse], error) { + resp, err := f.removeNodeFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { + if f.getNodeFn == nil { + return nil, connect.NewError(connect.CodeUnimplemented, nil) + } + resp, err := f.getNodeFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) { + if f.openPortFn == nil { + return connect.NewResponse(&nodev1.OpenPortResponse{ + Port: &nodev1.Port{ + PortId: "port_ssh", + Protocol: req.Msg.GetProtocol(), + PortNumber: req.Msg.GetPortNumber(), + ServerPort: req.Msg.GetPortNumber(), + }, + }), nil + } + resp, err := f.openPortFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, req *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) { + if f.grantNodeSSHAccessFn == nil { + return nil, connect.NewError(connect.CodeUnimplemented, nil) + } + resp, err := f.grantNodeSSHAccessFn(req.Msg) + if err != nil { + return nil, err + } + return connect.NewResponse(resp), nil +} + +func Test_ExternalNodeServiceClient_AddNode(t *testing.T) { + svc := &fakeNodeService{ + addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + if req.GetOrganizationId() != "org_123" { + t.Errorf("unexpected org ID: %s", req.GetOrganizationId()) + } + if req.GetName() != "My Spark" { + t.Errorf("unexpected name: %s", req.GetName()) + } + return &nodev1.AddNodeResponse{ + ExternalNode: &nodev1.ExternalNode{ + ExternalNodeId: "unode_abc", + OrganizationId: "org_123", + Name: req.GetName(), + DeviceId: req.GetDeviceId(), + }, + }, nil + }, + } + + _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) + server := httptest.NewServer(handler) + defer server.Close() + + client := store.NewDevPlaneClient(&staticTokenProvider{token: "tok"}, server.URL).ExternalNodes + + resp, err := client.AddNode(context.Background(), connect.NewRequest(&nodev1.AddNodeRequest{ + OrganizationId: "org_123", + Name: "My Spark", + DeviceId: "dev-uuid", + NodeSpec: &nodev1.NodeSpec{Architecture: strPtr("arm64")}, + })) + if err != nil { + t.Fatalf("AddNode failed: %v", err) + } + if resp.Msg.GetExternalNode().GetExternalNodeId() != "unode_abc" { + t.Errorf("unexpected node ID: %s", resp.Msg.GetExternalNode().GetExternalNodeId()) + } +} + +func Test_ExternalNodeServiceClient_AddNode_ServerError(t *testing.T) { + svc := &fakeNodeService{ + addNodeFn: func(_ *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { + return nil, connect.NewError(connect.CodeInternal, nil) + }, + } + + _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) + server := httptest.NewServer(handler) + defer server.Close() + + client := store.NewDevPlaneClient(&staticTokenProvider{token: "tok"}, server.URL).ExternalNodes + + _, err := client.AddNode(context.Background(), connect.NewRequest(&nodev1.AddNodeRequest{ + OrganizationId: "org_123", + Name: "Test", + DeviceId: "dev", + })) + if err == nil { + t.Fatal("expected error for server error response") + } +} + +func Test_ExternalNodeServiceClient_RemoveNode(t *testing.T) { + svc := &fakeNodeService{ + removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + if req.GetExternalNodeId() != "unode_abc" { + t.Errorf("unexpected node ID: %s", req.GetExternalNodeId()) + } + return &nodev1.RemoveNodeResponse{}, nil + }, + } + + _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) + server := httptest.NewServer(handler) + defer server.Close() + + client := store.NewDevPlaneClient(&staticTokenProvider{token: "tok"}, server.URL).ExternalNodes + + _, err := client.RemoveNode(context.Background(), connect.NewRequest(&nodev1.RemoveNodeRequest{ + ExternalNodeId: "unode_abc", + })) + if err != nil { + t.Fatalf("RemoveNode failed: %v", err) + } +} + +func Test_ExternalNodeServiceClient_RemoveNode_ServerError(t *testing.T) { + svc := &fakeNodeService{ + removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { + return nil, connect.NewError(connect.CodeNotFound, nil) + }, + } + + _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) + server := httptest.NewServer(handler) + defer server.Close() + + client := store.NewDevPlaneClient(&staticTokenProvider{token: "tok"}, server.URL).ExternalNodes + + _, err := client.RemoveNode(context.Background(), connect.NewRequest(&nodev1.RemoveNodeRequest{ + ExternalNodeId: "unode_missing", + })) + if err == nil { + t.Fatal("expected error for not found response") + } +} + +func strPtr(s string) *string { return &s } diff --git a/pkg/cmd/register/rpcclient.go b/pkg/cmd/register/rpcclient.go deleted file mode 100644 index 3689a199e..000000000 --- a/pkg/cmd/register/rpcclient.go +++ /dev/null @@ -1,145 +0,0 @@ -package register - -import ( - "net/http" - - nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" - nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" - - breverrors "github.com/brevdev/brev-cli/pkg/errors" - "github.com/brevdev/brev-cli/pkg/externalnode" - "github.com/brevdev/brev-cli/pkg/store" -) - -// bearerTokenTransport injects a Bearer token into every request. -// We use a custom RoundTripper instead of setting headers on individual -// client.Do() calls because ConnectRPC owns the HTTP requests internally — -// this is the only hook we have to add auth headers. -type bearerTokenTransport struct { - provider externalnode.TokenProvider - base http.RoundTripper -} - -func (t *bearerTokenTransport) RoundTrip(req *http.Request) (*http.Response, error) { - token, err := t.provider.GetAccessToken() - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - req = req.Clone(req.Context()) - req.Header.Set("Authorization", "Bearer "+token) - store.AddCLIAttributionParams(req) - resp, err := t.base.RoundTrip(req) - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - return resp, nil -} - -// newAuthenticatedHTTPClient creates an http.Client that injects the bearer token -// from the given provider on every request. -func newAuthenticatedHTTPClient(provider externalnode.TokenProvider) *http.Client { - return &http.Client{ - Transport: &bearerTokenTransport{ - provider: provider, - base: http.DefaultTransport, - }, - } -} - -// NewNodeServiceClient creates a ConnectRPC ExternalNodeServiceClient using the -// given token provider for authentication. -func NewNodeServiceClient(provider externalnode.TokenProvider, baseURL string) nodev1connect.ExternalNodeServiceClient { - return nodev1connect.NewExternalNodeServiceClient( - newAuthenticatedHTTPClient(provider), - baseURL, - ) -} - -// NewEnvironmentServiceClient creates an authenticated ConnectRPC EnvironmentServiceClient. -func NewEnvironmentServiceClient(provider externalnode.TokenProvider, baseURL string) nodev1connect.EnvironmentServiceClient { - return nodev1connect.NewEnvironmentServiceClient( - newAuthenticatedHTTPClient(provider), - baseURL, - ) -} - -// NewManagedSecretServiceClient creates an authenticated ConnectRPC managed-secret client. -func NewManagedSecretServiceClient(provider externalnode.TokenProvider, baseURL string) nodev1connect.ManagedSecretServiceClient { - return nodev1connect.NewManagedSecretServiceClient( - newAuthenticatedHTTPClient(provider), - baseURL, - ) -} - -// toProtoNodeSpec converts the local HardwareProfile (used for collection, display, -// persistence) to the generated proto NodeSpec for RPC calls. -func toProtoNodeSpec(hw *HardwareProfile) *nodev1.NodeSpec { - if hw == nil { - return nil - } - - proto := &nodev1.NodeSpec{ - RamBytes: hw.RAMBytes, - CpuCount: hw.CPUCount, - } - - for _, st := range hw.Storage { - storageSpec := &nodev1.StorageSpec{ - StorageBytes: st.StorageBytes, - StorageType: st.StorageType, - } - if st.Name != "" { - storageSpec.Device = &st.Name - } - proto.Storage = append(proto.Storage, storageSpec) - } - - if hw.Architecture != "" { - proto.Architecture = &hw.Architecture - } - if hw.OS != "" { - proto.Os = &hw.OS - } - if hw.OSVersion != "" { - proto.OsVersion = &hw.OSVersion - } - - if hw.ProductName != "" { - proto.ProductName = &hw.ProductName - } - - for _, g := range hw.GPUs { - pg := &nodev1.GPUSpec{ - Model: g.Model, - Count: g.Count, - MemoryBytes: g.MemoryBytes, - } - if g.Architecture != "" { - pg.GpuArchitecture = &g.Architecture - } - proto.Gpus = append(proto.Gpus, pg) - } - - for _, ic := range hw.Interconnects { - spec := &nodev1.InterconnectSpec{Device: ic.Device} - switch ic.Type { - case "NVLink": - spec.Details = &nodev1.InterconnectSpec_Nvlink{ - Nvlink: &nodev1.NVLinkDetails{ - ActiveLinks: int32(ic.ActiveLinks), - Version: ic.Version, - }, - } - case "PCIe": - spec.Details = &nodev1.InterconnectSpec_Pcie{ - Pcie: &nodev1.PCIeDetails{ - Generation: int32(ic.Generation), - Width: int32(ic.Width), - }, - } - } - proto.Interconnects = append(proto.Interconnects, spec) - } - - return proto -} diff --git a/pkg/cmd/register/rpcclient_test.go b/pkg/cmd/register/rpcclient_test.go deleted file mode 100644 index 78ffdee42..000000000 --- a/pkg/cmd/register/rpcclient_test.go +++ /dev/null @@ -1,387 +0,0 @@ -package register - -import ( - "context" - "net/http" - "net/http/httptest" - "net/url" - "runtime" - "testing" - - nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" - nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" - "connectrpc.com/connect" - - "github.com/brevdev/brev-cli/pkg/cmd/version" -) - -type mockTokenProvider struct { - token string - err error -} - -func (m *mockTokenProvider) GetAccessToken() (string, error) { - return m.token, m.err -} - -func Test_bearerTokenTransport_InjectsHeader(t *testing.T) { - var gotAuth string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotAuth = r.Header.Get("Authorization") - w.WriteHeader(http.StatusOK) - })) - defer server.Close() - - provider := &mockTokenProvider{token: "test-token-123"} - client := newAuthenticatedHTTPClient(provider) - - req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL, nil) - resp, err := client.Do(req) - if err != nil { - t.Fatalf("request failed: %v", err) - } - defer resp.Body.Close() //nolint:errcheck // test - - if gotAuth != "Bearer test-token-123" { - t.Errorf("expected 'Bearer test-token-123', got %q", gotAuth) - } -} - -func Test_bearerTokenTransport_AddsCLIAttributionParams(t *testing.T) { - var gotQuery url.Values - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - gotQuery = r.URL.Query() - w.WriteHeader(http.StatusOK) - })) - defer server.Close() - - client := newAuthenticatedHTTPClient(&mockTokenProvider{token: "t"}) - - req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, server.URL+"/devplaneapi.v1.ExternalNodeService/ListNodes?connect=v1", nil) - resp, err := client.Do(req) - if err != nil { - t.Fatalf("request failed: %v", err) - } - defer resp.Body.Close() //nolint:errcheck // test - - // Same attribution the REST client sends, so dev-plane sees one CLI identity. - if got := gotQuery.Get("utm_source"); got != "cli" { - t.Errorf("expected utm_source=cli, got %q", got) - } - if got := gotQuery.Get("cli_version"); got != version.Version { - t.Errorf("expected cli_version=%q, got %q", version.Version, got) - } - if got := gotQuery.Get("os"); got != runtime.GOOS { - t.Errorf("expected os=%q, got %q", runtime.GOOS, got) - } - // ConnectRPC's own protocol params must survive. - if got := gotQuery.Get("connect"); got != "v1" { - t.Errorf("expected existing query param to be preserved, got %q", got) - } -} - -func Test_bearerTokenTransport_PropagatesTokenError(t *testing.T) { - provider := &mockTokenProvider{err: http.ErrAbortHandler} - client := newAuthenticatedHTTPClient(provider) - - req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://localhost", nil) - resp, err := client.Do(req) - if err == nil { - resp.Body.Close() //nolint:errcheck // test - t.Fatal("expected error from token provider") - } -} - -func Test_toProtoNodeSpec(t *testing.T) { - cpuCount := int32(12) - ramBytes := int64(137438953472) - memBytes := int64(137438953472) - - local := &HardwareProfile{ - GPUs: []GPU{ - {Model: "NVIDIA GB10", Architecture: "Blackwell", Count: 2, MemoryBytes: &memBytes}, - }, - RAMBytes: &ramBytes, - CPUCount: &cpuCount, - Architecture: "arm64", - Storage: []StorageDevice{ - {Name: "nvme0n1", StorageBytes: 500107862016, StorageType: "SSD"}, - }, - OS: "Ubuntu", - OSVersion: "24.04", - ProductName: "DGX Spark", - Interconnects: []Interconnect{ - {Type: "NVLink", Device: "GPU 0", ActiveLinks: 4, Version: 4}, - }, - } - - proto := toProtoNodeSpec(local) - - if proto.GetCpuCount() != 12 { - t.Errorf("expected CpuCount 12, got %d", proto.GetCpuCount()) - } - if proto.GetRamBytes() != 137438953472 { - t.Errorf("expected RamBytes, got %d", proto.GetRamBytes()) - } - if proto.GetArchitecture() != "arm64" { - t.Errorf("expected arm64, got %s", proto.GetArchitecture()) - } - if proto.GetOs() != "Ubuntu" { - t.Errorf("expected Ubuntu, got %s", proto.GetOs()) - } - if proto.GetOsVersion() != "24.04" { - t.Errorf("expected 24.04, got %s", proto.GetOsVersion()) - } - if len(proto.GetStorage()) != 1 { - t.Fatalf("expected 1 storage entry, got %d", len(proto.GetStorage())) - } - if proto.GetStorage()[0].GetStorageBytes() != 500107862016 { - t.Errorf("expected StorageBytes 500107862016, got %d", proto.GetStorage()[0].GetStorageBytes()) - } - if proto.GetStorage()[0].GetStorageType() != "SSD" { - t.Errorf("expected SSD, got %s", proto.GetStorage()[0].GetStorageType()) - } - if len(proto.GetGpus()) != 1 { - t.Fatalf("expected 1 GPU, got %d", len(proto.GetGpus())) - } - gpu := proto.GetGpus()[0] - if gpu.GetModel() != "NVIDIA GB10" { - t.Errorf("expected NVIDIA GB10, got %s", gpu.GetModel()) - } - if gpu.GetCount() != 2 { - t.Errorf("expected count 2, got %d", gpu.GetCount()) - } - if gpu.GetMemoryBytes() != 137438953472 { - t.Errorf("expected memory bytes, got %d", gpu.GetMemoryBytes()) - } -} - -func Test_toProtoNodeSpec_PCIeInterconnect(t *testing.T) { - local := &HardwareProfile{ - Architecture: "amd64", - Interconnects: []Interconnect{ - {Type: "PCIe", Device: "GPU 0", Generation: 4, Width: 16}, - }, - } - - proto := toProtoNodeSpec(local) - - if len(proto.GetInterconnects()) != 1 { - t.Fatalf("expected 1 interconnect, got %d", len(proto.GetInterconnects())) - } - ic := proto.GetInterconnects()[0] - if ic.GetDevice() != "GPU 0" { - t.Errorf("expected device 'GPU 0', got %s", ic.GetDevice()) - } - pcie := ic.GetPcie() - if pcie == nil { - t.Fatal("expected PCIe details, got nil") - } - if pcie.GetGeneration() != 4 { - t.Errorf("expected PCIe Gen 4, got %d", pcie.GetGeneration()) - } - if pcie.GetWidth() != 16 { - t.Errorf("expected PCIe Width 16, got %d", pcie.GetWidth()) - } -} - -func Test_toProtoNodeSpec_Nil(t *testing.T) { - if toProtoNodeSpec(nil) != nil { - t.Error("expected nil for nil input") - } -} - -func Test_toProtoNodeSpec_MinimalFields(t *testing.T) { - local := &HardwareProfile{ - Architecture: "amd64", - } - proto := toProtoNodeSpec(local) - if proto.GetArchitecture() != "amd64" { - t.Errorf("expected amd64, got %s", proto.GetArchitecture()) - } - if proto.RamBytes != nil { - t.Error("expected nil RamBytes") - } - if proto.CpuCount != nil { - t.Error("expected nil CpuCount") - } - if len(proto.GetGpus()) != 0 { - t.Error("expected no GPUs") - } -} - -// fakeNodeService implements the server side of ExternalNodeService for testing. -type fakeNodeService struct { - nodev1connect.UnimplementedExternalNodeServiceHandler - addNodeFn func(*nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) - removeNodeFn func(*nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) - getNodeFn func(*nodev1.GetNodeRequest) (*nodev1.GetNodeResponse, error) - openPortFn func(*nodev1.OpenPortRequest) (*nodev1.OpenPortResponse, error) - grantNodeSSHAccessFn func(*nodev1.GrantNodeSSHAccessRequest) (*nodev1.GrantNodeSSHAccessResponse, error) -} - -func (f *fakeNodeService) AddNode(_ context.Context, req *connect.Request[nodev1.AddNodeRequest]) (*connect.Response[nodev1.AddNodeResponse], error) { - resp, err := f.addNodeFn(req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - -func (f *fakeNodeService) RemoveNode(_ context.Context, req *connect.Request[nodev1.RemoveNodeRequest]) (*connect.Response[nodev1.RemoveNodeResponse], error) { - resp, err := f.removeNodeFn(req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - -func (f *fakeNodeService) GetNode(_ context.Context, req *connect.Request[nodev1.GetNodeRequest]) (*connect.Response[nodev1.GetNodeResponse], error) { - if f.getNodeFn == nil { - return nil, connect.NewError(connect.CodeUnimplemented, nil) - } - resp, err := f.getNodeFn(req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - -func (f *fakeNodeService) OpenPort(_ context.Context, req *connect.Request[nodev1.OpenPortRequest]) (*connect.Response[nodev1.OpenPortResponse], error) { - if f.openPortFn == nil { - return connect.NewResponse(&nodev1.OpenPortResponse{ - Port: &nodev1.Port{ - PortId: "port_ssh", - Protocol: req.Msg.GetProtocol(), - PortNumber: req.Msg.GetPortNumber(), - ServerPort: req.Msg.GetPortNumber(), - }, - }), nil - } - resp, err := f.openPortFn(req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - -func (f *fakeNodeService) GrantNodeSSHAccess(_ context.Context, req *connect.Request[nodev1.GrantNodeSSHAccessRequest]) (*connect.Response[nodev1.GrantNodeSSHAccessResponse], error) { - if f.grantNodeSSHAccessFn == nil { - return nil, connect.NewError(connect.CodeUnimplemented, nil) - } - resp, err := f.grantNodeSSHAccessFn(req.Msg) - if err != nil { - return nil, err - } - return connect.NewResponse(resp), nil -} - -func Test_NewNodeServiceClient_AddNode(t *testing.T) { - svc := &fakeNodeService{ - addNodeFn: func(req *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { - if req.GetOrganizationId() != "org_123" { - t.Errorf("unexpected org ID: %s", req.GetOrganizationId()) - } - if req.GetName() != "My Spark" { - t.Errorf("unexpected name: %s", req.GetName()) - } - return &nodev1.AddNodeResponse{ - ExternalNode: &nodev1.ExternalNode{ - ExternalNodeId: "unode_abc", - OrganizationId: "org_123", - Name: req.GetName(), - DeviceId: req.GetDeviceId(), - }, - }, nil - }, - } - - _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) - server := httptest.NewServer(handler) - defer server.Close() - - client := NewNodeServiceClient(&mockTokenProvider{token: "tok"}, server.URL) - - resp, err := client.AddNode(context.Background(), connect.NewRequest(&nodev1.AddNodeRequest{ - OrganizationId: "org_123", - Name: "My Spark", - DeviceId: "dev-uuid", - NodeSpec: &nodev1.NodeSpec{Architecture: strPtr("arm64")}, - })) - if err != nil { - t.Fatalf("AddNode failed: %v", err) - } - if resp.Msg.GetExternalNode().GetExternalNodeId() != "unode_abc" { - t.Errorf("unexpected node ID: %s", resp.Msg.GetExternalNode().GetExternalNodeId()) - } -} - -func Test_NewNodeServiceClient_AddNode_ServerError(t *testing.T) { - svc := &fakeNodeService{ - addNodeFn: func(_ *nodev1.AddNodeRequest) (*nodev1.AddNodeResponse, error) { - return nil, connect.NewError(connect.CodeInternal, nil) - }, - } - - _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) - server := httptest.NewServer(handler) - defer server.Close() - - client := NewNodeServiceClient(&mockTokenProvider{token: "tok"}, server.URL) - - _, err := client.AddNode(context.Background(), connect.NewRequest(&nodev1.AddNodeRequest{ - OrganizationId: "org_123", - Name: "Test", - DeviceId: "dev", - })) - if err == nil { - t.Fatal("expected error for server error response") - } -} - -func Test_NewNodeServiceClient_RemoveNode(t *testing.T) { - svc := &fakeNodeService{ - removeNodeFn: func(req *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { - if req.GetExternalNodeId() != "unode_abc" { - t.Errorf("unexpected node ID: %s", req.GetExternalNodeId()) - } - return &nodev1.RemoveNodeResponse{}, nil - }, - } - - _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) - server := httptest.NewServer(handler) - defer server.Close() - - client := NewNodeServiceClient(&mockTokenProvider{token: "tok"}, server.URL) - - _, err := client.RemoveNode(context.Background(), connect.NewRequest(&nodev1.RemoveNodeRequest{ - ExternalNodeId: "unode_abc", - })) - if err != nil { - t.Fatalf("RemoveNode failed: %v", err) - } -} - -func Test_NewNodeServiceClient_RemoveNode_ServerError(t *testing.T) { - svc := &fakeNodeService{ - removeNodeFn: func(_ *nodev1.RemoveNodeRequest) (*nodev1.RemoveNodeResponse, error) { - return nil, connect.NewError(connect.CodeNotFound, nil) - }, - } - - _, handler := nodev1connect.NewExternalNodeServiceHandler(svc) - server := httptest.NewServer(handler) - defer server.Close() - - client := NewNodeServiceClient(&mockTokenProvider{token: "tok"}, server.URL) - - _, err := client.RemoveNode(context.Background(), connect.NewRequest(&nodev1.RemoveNodeRequest{ - ExternalNodeId: "unode_missing", - })) - if err == nil { - t.Fatal("expected error for not found response") - } -} - -func strPtr(s string) *string { return &s } diff --git a/pkg/cmd/revokessh/revokessh_test.go b/pkg/cmd/revokessh/revokessh_test.go index b0f971aa6..d0945f572 100644 --- a/pkg/cmd/revokessh/revokessh_test.go +++ b/pkg/cmd/revokessh/revokessh_test.go @@ -13,6 +13,7 @@ import ( "github.com/brevdev/brev-cli/pkg/cmd/register" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/externalnode" + "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" ) @@ -31,7 +32,7 @@ type mockNodeClientFactory struct { } func (m mockNodeClientFactory) NewNodeClient(provider externalnode.TokenProvider, _ string) nodev1connect.ExternalNodeServiceClient { - return register.NewNodeServiceClient(provider, m.serverURL) + return store.NewDevPlaneClient(provider, m.serverURL).ExternalNodes } type mockRegistrationStore struct { diff --git a/pkg/cmd/util/externalnode.go b/pkg/cmd/util/externalnode.go index 7d37e43fc..f00ba3a6e 100644 --- a/pkg/cmd/util/externalnode.go +++ b/pkg/cmd/util/externalnode.go @@ -8,11 +8,10 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" - "github.com/brevdev/brev-cli/pkg/cmd/register" - "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/ssh" + "github.com/brevdev/brev-cli/pkg/store" ) // ExternalNodeStore is the minimal interface needed for Brev Connect machine lookup and SSH resolution. @@ -20,6 +19,7 @@ type ExternalNodeStore interface { GetActiveOrganizationOrDefault() (*entity.Organization, error) GetAccessToken() (string, error) GetCurrentUser() (*entity.User, error) + DevPlane() *store.DevPlaneClient } type WorkspaceOrNodeResolver interface { @@ -129,8 +129,8 @@ func resolvePortForSSHAccess(node *nodev1.ExternalNode, access *nodev1.SSHAccess // OpenPort calls the OpenPort RPC to open a port on a Brev Connect machine via netbird. // This must be called before attempting to connect to a non-SSH port on a node. -func OpenPort(store ExternalNodeStore, nodeID string, portNumber int32, protocol nodev1.PortProtocol) (*nodev1.Port, error) { - client := register.NewNodeServiceClient(store, config.GlobalConfig.GetBrevPublicAPIURL()) +func OpenPort(nodeStore ExternalNodeStore, nodeID string, portNumber int32, protocol nodev1.PortProtocol) (*nodev1.Port, error) { + client := nodeStore.DevPlane().ExternalNodes resp, err := client.OpenPort(context.Background(), connect.NewRequest(&nodev1.OpenPortRequest{ ExternalNodeId: nodeID, Protocol: protocol, @@ -151,12 +151,12 @@ func FindExternalNode(store ExternalNodeStore, nameOrID string) (*nodev1.Externa // FindExternalNodeWithContext searches for a Brev Connect machine by name or ID in the user's active // organization. Exact IDs take precedence over case-insensitive names. // Returns (nil, nil) if no matching node is found. -func FindExternalNodeWithContext(ctx context.Context, store ExternalNodeStore, nameOrID string) (*nodev1.ExternalNode, error) { - org, err := store.GetActiveOrganizationOrDefault() +func FindExternalNodeWithContext(ctx context.Context, nodeStore ExternalNodeStore, nameOrID string) (*nodev1.ExternalNode, error) { + org, err := nodeStore.GetActiveOrganizationOrDefault() if err != nil { return nil, breverrors.WrapAndTrace(err) } - client := register.NewNodeServiceClient(store, config.GlobalConfig.GetBrevPublicAPIURL()) + client := nodeStore.DevPlane().ExternalNodes resp, err := client.ListNodes(ctx, connect.NewRequest(&nodev1.ListNodesRequest{ OrganizationId: org.ID, Options: &nodev1.ListNodesOptions{ diff --git a/pkg/cmd/util/externalnode_test.go b/pkg/cmd/util/externalnode_test.go index 2674b9ebd..0e3fea8b7 100644 --- a/pkg/cmd/util/externalnode_test.go +++ b/pkg/cmd/util/externalnode_test.go @@ -12,8 +12,10 @@ import ( nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" + "github.com/brevdev/brev-cli/pkg/store" ) // mockExternalNodeStore satisfies the shared node and workspace lookup interfaces. @@ -31,6 +33,10 @@ func (m *mockExternalNodeStore) GetActiveOrganizationOrDefault() (*entity.Organi func (m *mockExternalNodeStore) GetAccessToken() (string, error) { return "tok", nil } +func (m *mockExternalNodeStore) DevPlane() *store.DevPlaneClient { + return store.NewDevPlaneClient(m, config.GlobalConfig.GetBrevPublicAPIURL()) +} + func (m *mockExternalNodeStore) GetAuthTokens() (*entity.AuthTokens, error) { return nil, nil } func (m *mockExternalNodeStore) GetWorkspaceByNameOrID(_, _ string) ([]entity.Workspace, error) { diff --git a/pkg/config/config.go b/pkg/config/config.go index 7d492335d..74be86c0f 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -8,7 +8,6 @@ type EnvVarName string // should be caps with underscore const ( brevAPIURL EnvVarName = "BREV_API_URL" - brevGRPCURL EnvVarName = "BREV_GRPC_URL" brevAuthURL EnvVarName = "BREV_AUTH_URL" brevAuthIssuerURL EnvVarName = "BREV_AUTH_ISSUER_URL" brevConsoleURL EnvVarName = "BREV_CONSOLE_URL" @@ -33,16 +32,13 @@ func (c ConstantsConfig) GetBrevAPIURl() string { return getEnvOrDefault(brevAPIURL, "https://brevapi.us-west-2-prod.control-plane.brev.dev") } -// GetBrevPublicAPIURL returns the public REST API URL (no auth required) +// GetBrevPublicAPIURL returns the base URL of dev-plane's public API server +// (BREV_PUBLIC_API_URL). That server serves both the public HTTP routes (/v1/...) +// and the ConnectRPC services the CLI calls (pkg/cmd/register, pkg/environment). func (c ConstantsConfig) GetBrevPublicAPIURL() string { return getEnvOrDefault(brevPublicAPIURL, "https://api.brev.dev") } -func (c ConstantsConfig) GetBrevGRPCURL() string { - // GRPC does not use https:// prefix - return getEnvOrDefault(brevGRPCURL, "api.brev.dev:443") -} - func (c ConstantsConfig) GetBrevAuthURL() string { return getEnvOrDefault(brevAuthURL, "https://api.ngc.nvidia.com") } diff --git a/pkg/environment/status.go b/pkg/environment/status.go new file mode 100644 index 000000000..efffcdcf5 --- /dev/null +++ b/pkg/environment/status.go @@ -0,0 +1,76 @@ +// Package environment provides shared operations for DevPlane environments. +package environment + +import ( + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + + "github.com/brevdev/brev-cli/pkg/entity" +) + +// Display maps environment status onto workspace status +func Display(env *devplanev1.Environment) string { + switch env.GetStatus() { + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING: + return entity.Running + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPED: + return entity.Stopped + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING: + return entity.Starting + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING: + return entity.Stopping + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATED: + return entity.Deleting + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_INSTANCE_LOST: + return entity.Failure + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PENDING, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING: + return entity.Deploying + } + return "" +} + +func FailureMessage(env *devplanev1.Environment) string { + if message := env.GetInstance().GetProvisionStatus().GetMessage(); message != "" { + return message + } + switch env.GetStatus() { + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING_FAILED: + return "preparation failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED: + return "provisioning failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING_FAILED: + return "build failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING_FAILED: + return "start failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING_FAILED: + return "stop failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING_FAILED: + return "termination failed" + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_INSTANCE_LOST: + return "instance lost" + } + return "" +} + +func IsFailed(env *devplanev1.Environment) bool { + switch env.GetStatus() { + case devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING_FAILED, + devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_INSTANCE_LOST: + return true + } + return false +} diff --git a/pkg/environment/status_test.go b/pkg/environment/status_test.go new file mode 100644 index 000000000..031a83c53 --- /dev/null +++ b/pkg/environment/status_test.go @@ -0,0 +1,92 @@ +package environment + +import ( + "testing" + + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "github.com/stretchr/testify/assert" + + "github.com/brevdev/brev-cli/pkg/entity" +) + +func TestStatusDisplayMapsEnvironmentStatuses(t *testing.T) { + cases := []struct { + name string + status devplanev1.EnvironmentStatus + display string + }{ + {"pending", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PENDING, entity.Deploying}, + {"preparing", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING, entity.Deploying}, + {"provisioning", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING, entity.Deploying}, + {"building", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING, entity.Deploying}, + {"running", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING, entity.Running}, + {"stopping", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING, entity.Stopping}, + {"stopped", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPED, entity.Stopped}, + {"starting", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING, entity.Starting}, + {"terminating", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING, entity.Deleting}, + {"terminated", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATED, entity.Deleting}, + {"preparing_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PREPARING_FAILED, entity.Failure}, + {"provisioning_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, entity.Failure}, + {"building_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_BUILDING_FAILED, entity.Failure}, + {"starting_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STARTING_FAILED, entity.Failure}, + {"stopping_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_STOPPING_FAILED, entity.Failure}, + {"terminating_failed", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_TERMINATING_FAILED, entity.Failure}, + {"instance_lost", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_INSTANCE_LOST, entity.Failure}, + {"unspecified", devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_UNSPECIFIED, ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + display := Display(&devplanev1.Environment{ + EnvironmentId: "env-1", + Name: "box", + Status: tc.status, + }) + assert.Equal(t, tc.display, display) + }) + } +} + +// TestFailureMessage covers the provider-message preference and fallbacks. +func TestFailureMessage(t *testing.T) { + cases := []struct { + name string + status devplanev1.EnvironmentStatus + instanceMsg string + want string + }{ + { + name: "provider message wins", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_PROVISIONING_FAILED, + instanceMsg: "out of quota", + want: "out of quota", + }, + { + name: "instance lost falls back to label", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_INSTANCE_LOST, + want: "instance lost", + }, + { + name: "non-failure status has no message", + status: devplanev1.EnvironmentStatus_ENVIRONMENT_STATUS_RUNNING, + want: "", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + env := &devplanev1.Environment{ + EnvironmentId: "env-1", + Name: "box", + Status: tc.status, + } + if tc.instanceMsg != "" { + env.Instance = &devplanev1.Instance{ + ProvisionStatus: &devplanev1.CreateAttemptStatus{Message: tc.instanceMsg}, + } + } + + assert.Equal(t, tc.want, FailureMessage(env)) + }) + } +} diff --git a/pkg/store/devplane_client.go b/pkg/store/devplane_client.go new file mode 100644 index 000000000..50cadda56 --- /dev/null +++ b/pkg/store/devplane_client.go @@ -0,0 +1,78 @@ +package store + +import ( + "net/http" + + devplanev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" + + breverrors "github.com/brevdev/brev-cli/pkg/errors" +) + +// DevPlaneClient bundles dev-plane's ConnectRPC service clients over a single HTTP client +type DevPlaneClient struct { + Instances devplanev1connect.InstanceServiceClient + Environments devplanev1connect.EnvironmentServiceClient + Organizations devplanev1connect.OrganizationServiceClient + ExternalNodes devplanev1connect.ExternalNodeServiceClient + ManagedSecrets devplanev1connect.ManagedSecretServiceClient +} + +// NewAuthenticatedHTTPClient builds the HTTP client shared by every authenticated +// dev-plane ConnectRPC service. It injects the bearer token from auth and the CLI +// attribution params on each request. +func NewAuthenticatedHTTPClient(auth Auth) *http.Client { + return &http.Client{Transport: &authenticatedTransport{provider: auth, base: http.DefaultTransport}} +} + +// NewPublicHTTPClient builds the HTTP client used by unauthenticated dev-plane +// public routes; it only adds the CLI attribution params. +func NewPublicHTTPClient() *http.Client { + return &http.Client{Transport: attributionTransport{base: http.DefaultTransport}} +} + +// NewDevPlaneClient builds an authenticated bundle over one shared HTTP client. +func NewDevPlaneClient(auth Auth, baseURL string) *DevPlaneClient { + return newDevPlaneClient(NewAuthenticatedHTTPClient(auth), baseURL) +} + +// NewPublicDevPlaneClient builds an unauthenticated bundle for dev-plane's public routes. +func NewPublicDevPlaneClient(baseURL string) *DevPlaneClient { + return newDevPlaneClient(NewPublicHTTPClient(), baseURL) +} + +func newDevPlaneClient(httpClient *http.Client, baseURL string) *DevPlaneClient { + return &DevPlaneClient{ + Instances: devplanev1connect.NewInstanceServiceClient(httpClient, baseURL), + Environments: devplanev1connect.NewEnvironmentServiceClient(httpClient, baseURL), + Organizations: devplanev1connect.NewOrganizationServiceClient(httpClient, baseURL), + ExternalNodes: devplanev1connect.NewExternalNodeServiceClient(httpClient, baseURL), + ManagedSecrets: devplanev1connect.NewManagedSecretServiceClient(httpClient, baseURL), + } +} + +// DevPlane returns the store's authenticated dev-plane client bundle. It is +// built in WithAuthHTTPClient, so commands and store methods share one HTTP +// client and one base URL. Expected to have been created via WithAuthHTTPClient +func (s *AuthHTTPStore) DevPlane() *DevPlaneClient { + return s.devPlaneClient +} + +type authenticatedTransport struct { + provider Auth + base http.RoundTripper +} + +func (t *authenticatedTransport) RoundTrip(req *http.Request) (*http.Response, error) { + token, err := t.provider.GetAccessToken() + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + req = req.Clone(req.Context()) + req.Header.Set("Authorization", "Bearer "+token) + AddCLIAttributionParams(req) + resp, err := t.base.RoundTrip(req) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + return resp, nil +} diff --git a/pkg/store/devplane_client_test.go b/pkg/store/devplane_client_test.go new file mode 100644 index 000000000..8196d8092 --- /dev/null +++ b/pkg/store/devplane_client_test.go @@ -0,0 +1,64 @@ +package store + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "runtime" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/brevdev/brev-cli/pkg/cmd/version" +) + +type failingAuth struct{} + +func (failingAuth) GetAccessToken() (string, error) { + return "", errors.New("token unavailable") +} + +// TestAuthenticatedHTTPClientRequestAnnotation covers what the client adds to +// every outgoing request. +func TestAuthenticatedHTTPClientRequestAnnotation(t *testing.T) { + var gotAuth string + var gotQuery url.Values + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotQuery = r.URL.Query() + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + client := NewAuthenticatedHTTPClient(MockAuth{}) + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, + server.URL+"/devplaneapi.v1.ExternalNodeService/ListNodes?connect=v1", nil) + require.NoError(t, err) + resp, err := client.Do(req) + require.NoError(t, err) + defer resp.Body.Close() //nolint:errcheck // test + + assert.Equal(t, "Bearer mock-token", gotAuth) + // Same attribution the REST client sends, so dev-plane sees one CLI identity. + assert.Equal(t, "cli", gotQuery.Get("utm_source")) + assert.Equal(t, version.Version, gotQuery.Get("cli_version")) + assert.Equal(t, runtime.GOOS, gotQuery.Get("os")) + // ConnectRPC's own protocol params must survive. + assert.Equal(t, "v1", gotQuery.Get("connect")) +} + +func TestAuthenticatedHTTPClientPropagatesTokenError(t *testing.T) { + client := NewAuthenticatedHTTPClient(failingAuth{}) + + req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "http://localhost", nil) + require.NoError(t, err) + resp, err := client.Do(req) + if err == nil { + resp.Body.Close() //nolint:errcheck // test + t.Fatal("expected error from token provider") + } +} diff --git a/pkg/store/environment.go b/pkg/store/environment.go new file mode 100644 index 000000000..92be278b6 --- /dev/null +++ b/pkg/store/environment.go @@ -0,0 +1,54 @@ +package store + +import ( + "context" + + devplanev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" + "connectrpc.com/connect" + + breverrors "github.com/brevdev/brev-cli/pkg/errors" +) + +const maxEnvironmentPageSize = 100 + +// GetEnvironment reads a single environment with its instance attached. +// Callers derive display state with environment.Display/FailureMessage. +func (s *AuthHTTPStore) GetEnvironment(ctx context.Context, environmentID string) (*devplanev1.Environment, error) { + env, err := s.DevPlane().Environments.GetEnvironment(ctx, connect.NewRequest(&devplanev1.GetEnvironmentRequest{ + EnvironmentId: environmentID, + AttachedDataOptions: &devplanev1.GetEnvironmentAttachedDataOptions{Instance: true}, + })) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + return env.Msg.GetEnvironment(), nil +} + +// GetEnvironmentsByOrg reads every environment in the organization. +// Callers derive display state with environment.Display/FailureMessage. +func (s *AuthHTTPStore) GetEnvironmentsByOrg(ctx context.Context, orgID string) ([]*devplanev1.Environment, error) { + var environments []*devplanev1.Environment + pageToken := "" + for { + response, err := s.DevPlane().Organizations.ListOrganizationComputeResources(ctx, connect.NewRequest(&devplanev1.ListOrganizationComputeResourcesRequest{ + OrganizationId: orgID, + PageParams: &devplanev1.PageParams{PageSize: maxEnvironmentPageSize, PageToken: pageToken}, + Options: &devplanev1.ListOrganizationComputeResourcesOptions{ + Type: devplanev1.ComputeResourceType_COMPUTE_RESOURCE_TYPE_ENVIRONMENT, + EnvironmentOptions: &devplanev1.ListEnvironmentAttachedDataOptions{Instance: true}, + }, + })) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + for _, item := range response.Msg.GetItems() { + if env := item.GetEnvironment(); env != nil { + environments = append(environments, env) + } + } + pageToken = response.Msg.GetNextPageToken() + if pageToken == "" { + return environments, nil + } + } +} diff --git a/pkg/store/http.go b/pkg/store/http.go index 3d27fe7fc..5c452865e 100644 --- a/pkg/store/http.go +++ b/pkg/store/http.go @@ -13,6 +13,7 @@ import ( "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/cmd/version" + "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/featureflag" @@ -103,6 +104,7 @@ type AuthHTTPStore struct { isRefreshTokenHandlerSet bool organizationOverride *entity.Organization organizationOverrideName string + devPlaneClient *DevPlaneClient BasicStore } @@ -145,7 +147,9 @@ func (f *FileStore) WithAuthHTTPClient(c *AuthHTTPClient) *AuthHTTPStore { c.restyClient.SetQueryParam("local", "true") } na := f.WithNoAuthHTTPClient(NewNoAuthHTTPClient(c.restyClient.BaseURL)) - return &AuthHTTPStore{NoAuthHTTPStore: *na, authHTTPClient: c} + s := &AuthHTTPStore{NoAuthHTTPStore: *na, authHTTPClient: c} + s.devPlaneClient = NewDevPlaneClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) + return s } func (n *NoAuthHTTPStore) WithAuthHTTPClient(c *AuthHTTPClient) *AuthHTTPStore { @@ -154,7 +158,9 @@ func (n *NoAuthHTTPStore) WithAuthHTTPClient(c *AuthHTTPClient) *AuthHTTPStore { if id == "" { c.restyClient.SetQueryParam("local", "true") } - return &AuthHTTPStore{NoAuthHTTPStore: *n, authHTTPClient: c} + s := &AuthHTTPStore{NoAuthHTTPStore: *n, authHTTPClient: c} + s.devPlaneClient = NewDevPlaneClient(s, config.GlobalConfig.GetBrevPublicAPIURL()) + return s } func (n *NoAuthHTTPStore) WithAuth(auth Auth, options ...Option) *AuthHTTPStore { diff --git a/pkg/store/http_response_error_test.go b/pkg/store/http_response_error_test.go index 83e0adce9..489a237f1 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) { diff --git a/pkg/store/instancetypes.go b/pkg/store/instancetypes.go index 81acefd60..6fccdf94f 100644 --- a/pkg/store/instancetypes.go +++ b/pkg/store/instancetypes.go @@ -2,9 +2,7 @@ package store import ( "context" - "net/http" - "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" devplaneapiv1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" @@ -29,10 +27,7 @@ func (s AuthHTTPStore) GetInstanceTypes(includeCPU bool) (*gpusearch.InstanceTyp // fetchInstanceTypes fetches instance types from dev-plane's public Connect API. func fetchInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, error) { - client := devplaneapiv1connect.NewInstanceServiceClient( - &http.Client{Transport: attributionTransport{base: http.DefaultTransport}}, - config.NewConstants().GetBrevPublicAPIURL(), - ) + client := NewPublicDevPlaneClient(config.GlobalConfig.GetBrevPublicAPIURL()).Instances skipAccessFilter := false res, err := client.ListPublicInstanceType(context.Background(), connect.NewRequest(&devplaneapiv1.ListPublicInstanceTypeRequest{ Options: &devplaneapiv1.ListInstanceTypeOptions{ @@ -50,10 +45,7 @@ func fetchInstanceTypes(includeCPU bool) (*gpusearch.InstanceTypesResponse, erro // GetAllInstanceTypesWithCloudCreds fetches org-scoped instance types from dev-plane's public Connect API. func (s AuthHTTPStore) GetAllInstanceTypesWithCloudCreds(orgID string) (*gpusearch.AllInstanceTypesResponse, error) { - client := devplaneapiv1connect.NewInstanceServiceClient( - &http.Client{Transport: &authHTTPStoreTransport{store: &s, base: http.DefaultTransport}}, - config.NewConstants().GetBrevPublicAPIURL(), - ) + client := s.DevPlane().Instances includeUnavailable := false includePreemptible := false includeCPU := true diff --git a/pkg/store/managedsecret.go b/pkg/store/managedsecret.go index 78422c7c6..48fc2edf3 100644 --- a/pkg/store/managedsecret.go +++ b/pkg/store/managedsecret.go @@ -4,21 +4,16 @@ import ( "context" "errors" "fmt" - "net/http" nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" - "github.com/brevdev/brev-cli/pkg/config" breverrors "github.com/brevdev/brev-cli/pkg/errors" ) func (s *AuthHTTPStore) newManagedSecretServiceClient() nodev1connect.ManagedSecretServiceClient { - return nodev1connect.NewManagedSecretServiceClient( - &http.Client{Transport: &authHTTPStoreTransport{store: s, base: http.DefaultTransport}}, - config.GlobalConfig.GetBrevPublicAPIURL(), - ) + return s.DevPlane().ManagedSecrets } func (s *AuthHTTPStore) CreateManagedSecret(organizationID, name, value string) (*nodev1.ManagedSecret, error) { diff --git a/pkg/store/organization.go b/pkg/store/organization.go index 14430383a..21b02eccf 100644 --- a/pkg/store/organization.go +++ b/pkg/store/organization.go @@ -3,15 +3,12 @@ package store import ( "context" "fmt" - "net/http" "strings" - nodev1connect "buf.build/gen/go/brevdev/devplane/connectrpc/go/devplaneapi/v1/devplaneapiv1connect" nodev1 "buf.build/gen/go/brevdev/devplane/protocolbuffers/go/devplaneapi/v1" "connectrpc.com/connect" "github.com/brevdev/brev-cli/pkg/auth" - "github.com/brevdev/brev-cli/pkg/config" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/files" @@ -311,31 +308,8 @@ func (s AuthHTTPStore) CreateInviteLink(organizationID string) (string, error) { return result, nil } -type authHTTPStoreTransport struct { - store *AuthHTTPStore - base http.RoundTripper -} - -func (t *authHTTPStoreTransport) RoundTrip(req *http.Request) (*http.Response, error) { - token, err := t.store.GetAccessToken() - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - req = req.Clone(req.Context()) - req.Header.Set("Authorization", "Bearer "+token) - AddCLIAttributionParams(req) - resp, err := t.base.RoundTrip(req) - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - return resp, nil -} - func (s *AuthHTTPStore) ListOrganizationMembers(ctx context.Context, orgID string) ([]*nodev1.OrganizationMember, error) { - client := nodev1connect.NewOrganizationServiceClient( - &http.Client{Transport: &authHTTPStoreTransport{store: s, base: http.DefaultTransport}}, - config.GlobalConfig.GetBrevPublicAPIURL(), - ) + client := s.DevPlane().Organizations var members []*nodev1.OrganizationMember var pageToken string