From a0f9774df1c199e3300b568935a3cb7a909b153e Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Wed, 12 Aug 2026 16:40:29 -0700 Subject: [PATCH 1/6] Add org flag to more commands --- pkg/cmd/cmd.go | 32 +++++++++++++-- pkg/cmd/cmd_test.go | 7 ++++ pkg/cmd/grantssh/grantssh.go | 6 ++- pkg/cmd/invite/invite.go | 49 +++++------------------ pkg/cmd/ls/ls.go | 59 +++++---------------------- pkg/cmd/ls/ls_test.go | 6 +-- pkg/cmd/org/org.go | 4 +- pkg/cmd/org/set.go | 11 +---- pkg/cmd/redeem/redeem.go | 46 +++++---------------- pkg/cmd/register/register.go | 6 ++- pkg/cmd/revokessh/revokessh.go | 6 ++- pkg/cmd/start/start.go | 73 +++++++++------------------------- pkg/store/http.go | 8 ++++ pkg/store/organization.go | 3 ++ pkg/store/organization_test.go | 10 +++++ 15 files changed, 121 insertions(+), 205 deletions(-) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 49b5254d7..b4914257a 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -3,6 +3,7 @@ package cmd import ( "fmt" + "strings" "github.com/brevdev/brev-cli/pkg/analytics" "github.com/brevdev/brev-cli/pkg/auth" @@ -71,6 +72,7 @@ import ( var ( userFlag string + orgFlag string printVersion bool noCheckLatest bool ) @@ -82,6 +84,7 @@ func NewDefaultBrevCommand() *cobra.Command { cmd.PersistentFlags().BoolP("help", "h", false, "Help for Brev") cmd.PersistentFlags().StringVar(&userFlag, "user", "", "Non root user to use for per user configuration of commands run as root") + cmd.PersistentFlags().StringVarP(&orgFlag, "org", "o", "", "Organization to use for this command (does not change the active organization)") cmd.PersistentFlags().BoolVar(&printVersion, "version", false, "Print version output") cmd.PersistentFlags().BoolVar(&noCheckLatest, "no-check-latest", false, "Do not check for the latest version when printing version") @@ -131,6 +134,7 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin analytics.SetUserStore(noLoginCmdStore) + var externalNodeCmdStore *store.AuthHTTPStore cmds := &cobra.Command{ SilenceErrors: true, SilenceUsage: true, @@ -176,6 +180,28 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin } } + loginCmdStore.SetOrganizationOverride(nil) + noLoginCmdStore.SetOrganizationOverride(nil) + if externalNodeCmdStore != nil { + externalNodeCmdStore.SetOrganizationOverride(nil) + } + if strings.TrimSpace(orgFlag) != "" { + orgs, orgErr := loginCmdStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgFlag}) + if orgErr != nil { + return breverrors.WrapAndTrace(orgErr) + } + if len(orgs) == 0 { + return breverrors.NewValidationError(fmt.Sprintf("no org found with name %s", orgFlag)) + } + if len(orgs) > 1 { + return breverrors.NewValidationError(fmt.Sprintf("more than one org found with name %s", orgFlag)) + } + loginCmdStore.SetOrganizationOverride(&orgs[0]) + noLoginCmdStore.SetOrganizationOverride(&orgs[0]) + if externalNodeCmdStore != nil { + externalNodeCmdStore.SetOrganizationOverride(&orgs[0]) + } + } home, err := fsStore.GetBrevHomePath() if err != nil { fmt.Printf("Warning: %v", err) @@ -249,7 +275,7 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin memLoginAuth := auth.NewLoginAuth(memAuthStore, memAuthenticator) memLoginAuth.WithShouldLogin(func() (bool, error) { return true, nil }) - externalNodeCmdStore := fsStore.WithNoAuthHTTPClient( + externalNodeCmdStore = fsStore.WithNoAuthHTTPClient( store.NewNoAuthHTTPClient(conf.GetBrevAPIURl()), ).WithAuth(memLoginAuth, store.WithDebug(conf.GetDebugHTTP())) @@ -273,8 +299,8 @@ func createCmdTree(cmd *cobra.Command, t *terminal.Terminal, loginCmdStore *stor cmd.AddCommand(set.NewCmdSet(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(ls.NewCmdLs(t, loginCmdStore, noLoginCmdStore)) cmd.AddCommand(org.NewCmdOrg(t, loginCmdStore, noLoginCmdStore)) - cmd.AddCommand(invite.NewCmdInvite(t, loginCmdStore, noLoginCmdStore)) - cmd.AddCommand(redeem.NewCmdRedeem(t, loginCmdStore, noLoginCmdStore)) + cmd.AddCommand(invite.NewCmdInvite(t, loginCmdStore)) + cmd.AddCommand(redeem.NewCmdRedeem(t, loginCmdStore)) cmd.AddCommand(portforward.NewCmdPortForwardSSH(loginCmdStore, t)) cmd.AddCommand(login.NewCmdLogin(t, noLoginCmdStore, loginAuth)) cmd.AddCommand(logout.NewCmdLogout(loginAuth, noLoginCmdStore)) diff --git a/pkg/cmd/cmd_test.go b/pkg/cmd/cmd_test.go index c9289a4c3..eb0dd57d0 100644 --- a/pkg/cmd/cmd_test.go +++ b/pkg/cmd/cmd_test.go @@ -31,6 +31,13 @@ func newTestFileStore(t *testing.T) *store.FileStore { ) } +func TestNewDefaultBrevCommandRegistersGlobalOrgFlag(t *testing.T) { + cmd := NewDefaultBrevCommand() + flag := cmd.PersistentFlags().Lookup("org") + require.NotNil(t, flag) + assert.Equal(t, "o", flag.Shorthand) +} + func TestEmailCachingAuthStore_SaveCachesEmail(t *testing.T) { fs := newTestFileStore(t) s := &emailCachingAuthStore{ diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 35985834e..47719eaef 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -55,7 +55,6 @@ func defaultGrantSSHDeps() grantSSHDeps { } func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { - var orgFlag string var nodeFlag string var userFlag string var linuxUser string @@ -70,6 +69,10 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { Long: "Grant SSH access to a node for another member of your organization. Interactive: no flags, prompts for org, node, port, and user. Non-interactive: --org, --node, --user, --linux-user, and --port-id required.", Example: " brev grant-ssh\n brev grant-ssh --org my-org --node my-node --user user@example.com --linux-user ubuntu --port-id port_abc --approve", RunE: func(cmd *cobra.Command, args []string) error { + orgFlag, err := cmd.Flags().GetString("org") + if err != nil { + return breverrors.WrapAndTrace(err) + } interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" opts := grantSSHOpts{ interactive: interactive, @@ -84,7 +87,6 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode)") cmd.Flags().StringVarP(&nodeFlag, "node", "n", "", "node name (required in non-interactive mode)") cmd.Flags().StringVarP(&userFlag, "user", "u", "", "Brev user ID or email to grant (required in non-interactive mode)") cmd.Flags().StringVar(&linuxUser, "linux-user", "", "Linux username on the target node (required in non-interactive mode)") diff --git a/pkg/cmd/invite/invite.go b/pkg/cmd/invite/invite.go index 901319031..7003a6e8b 100644 --- a/pkg/cmd/invite/invite.go +++ b/pkg/cmd/invite/invite.go @@ -2,30 +2,24 @@ package invite import ( - "fmt" - "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/completions" "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/store" "github.com/brevdev/brev-cli/pkg/terminal" "github.com/spf13/cobra" ) type InviteStore interface { - completions.CompletionStore GetUsers(queryParams map[string]string) ([]entity.User, error) GetWorkspace(workspaceID string) (*entity.Workspace, error) CreateInviteLink(organizationID string) (string, error) + GetActiveOrganizationOrDefault() (*entity.Organization, error) } -func NewCmdInvite(t *terminal.Terminal, loginInviteStore InviteStore, noLoginInviteStore InviteStore) *cobra.Command { - var org string - +func NewCmdInvite(t *terminal.Terminal, loginInviteStore InviteStore) *cobra.Command { cmd := &cobra.Command{ Annotations: map[string]string{"organization": ""}, Use: "invite", @@ -45,7 +39,7 @@ func NewCmdInvite(t *terminal.Terminal, loginInviteStore InviteStore, noLoginInv }, Args: cmderrors.TransformToValidationError(cobra.NoArgs), RunE: func(cmd *cobra.Command, args []string) error { - err := RunInvite(t, loginInviteStore, org) + err := RunInvite(t, loginInviteStore) if err != nil { return breverrors.WrapAndTrace(err) } @@ -53,39 +47,16 @@ func NewCmdInvite(t *terminal.Terminal, loginInviteStore InviteStore, noLoginInv }, } - cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)") - err := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginInviteStore, t)) - if err != nil { - breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) - fmt.Print(breverrors.WrapAndTrace(err)) - } - return cmd } -func RunInvite(t *terminal.Terminal, inviteStore InviteStore, orgflag string) error { - var org *entity.Organization - if orgflag != "" { - orgs, err := inviteStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgflag}) - if err != nil { - return breverrors.WrapAndTrace(err) - } - if len(orgs) == 0 { - return fmt.Errorf("no org found with name %s", orgflag) - } else if len(orgs) > 1 { - return fmt.Errorf("more than one org found with name %s", orgflag) - } - - org = &orgs[0] - } else { - currOrg, err := inviteStore.GetActiveOrganizationOrDefault() - if err != nil { - return breverrors.WrapAndTrace(err) - } - if currOrg == nil { - return fmt.Errorf("no orgs exist") - } - org = currOrg +func RunInvite(t *terminal.Terminal, inviteStore InviteStore) error { + org, err := inviteStore.GetActiveOrganizationOrDefault() + if err != nil { + return breverrors.WrapAndTrace(err) + } + if org == nil { + return breverrors.NewValidationError("no orgs exist") } token, err := inviteStore.CreateInviteLink(org.ID) diff --git a/pkg/cmd/ls/ls.go b/pkg/cmd/ls/ls.go index 996897cff..f2d8667c5 100644 --- a/pkg/cmd/ls/ls.go +++ b/pkg/cmd/ls/ls.go @@ -16,7 +16,6 @@ import ( "github.com/brevdev/brev-cli/pkg/externalnode" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/completions" "github.com/brevdev/brev-cli/pkg/cmd/gpusearch" "github.com/brevdev/brev-cli/pkg/cmd/hello" "github.com/brevdev/brev-cli/pkg/cmd/register" @@ -49,7 +48,6 @@ type LsStore interface { func NewCmdLs(t *terminal.Terminal, loginLsStore LsStore, noLoginLsStore LsStore) *cobra.Command { var showAll bool - var org string var jsonOutput bool cmd := &cobra.Command{ @@ -91,7 +89,7 @@ with other commands like stop, start, or delete.`, if err != nil { return breverrors.WrapAndTrace(err) } - err = RunLs(t, cliAuth, loginLsStore, args, org, showAll, jsonOutput) + err = RunLs(t, cliAuth, loginLsStore, args, showAll, jsonOutput) if err != nil { return breverrors.WrapAndTrace(err) } @@ -102,13 +100,6 @@ with other commands like stop, start, or delete.`, }, } - cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)") - err := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginLsStore, t)) - if err != nil { - breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) - fmt.Print(breverrors.WrapAndTrace(err)) - } - cmd.Flags().BoolVar(&showAll, "all", false, "show all instances and external nodes in org") cmd.Flags().BoolVar(&jsonOutput, "json", false, "output as JSON") @@ -128,53 +119,21 @@ func trackLsAnalytics(cliAuth auth.CLIAuth) { _ = analytics.TrackEvent(data) } -func getOrgForRunLs(cliAuth auth.CLIAuth, lsStore LsStore, orgflag string) (*entity.Organization, error) { - var org *entity.Organization - if cliAuth.IsAPIKey() { - if orgflag != "" { - return nil, breverrors.NewValidationError("api key auth is scoped to the org saved during login; --org is not supported") - } - org, err := lsStore.GetActiveOrganizationOrDefault() - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - if org == nil { - return nil, breverrors.NewValidationError("no orgs exist") - } - return org, nil +func getOrgForRunLs(lsStore LsStore) (*entity.Organization, error) { + org, err := lsStore.GetActiveOrganizationOrDefault() + if err != nil { + return nil, breverrors.WrapAndTrace(err) } - - if orgflag != "" { - var orgs []entity.Organization - orgs, err := lsStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgflag}) - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - if len(orgs) == 0 { - return nil, breverrors.NewValidationError(fmt.Sprintf("no org found with name %s", orgflag)) - } else if len(orgs) > 1 { - return nil, breverrors.NewValidationError(fmt.Sprintf("more than one org found with name %s", orgflag)) - } - - org = &orgs[0] - } else { - var currOrg *entity.Organization - currOrg, err := lsStore.GetActiveOrganizationOrDefault() - if err != nil { - return nil, breverrors.WrapAndTrace(err) - } - if currOrg == nil { - return nil, breverrors.NewValidationError("no orgs exist") - } - org = currOrg + if org == nil { + return nil, breverrors.NewValidationError("no orgs exist") } return org, nil } -func RunLs(t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []string, orgflag string, showAll bool, jsonOutput bool) error { +func RunLs(t *terminal.Terminal, cliAuth auth.CLIAuth, lsStore LsStore, args []string, showAll bool, jsonOutput bool) error { ls := NewLs(lsStore, t, jsonOutput) - org, err := getOrgForRunLs(cliAuth, lsStore, orgflag) + org, err := getOrgForRunLs(lsStore) if err != nil { return breverrors.WrapAndTrace(err) } diff --git a/pkg/cmd/ls/ls_test.go b/pkg/cmd/ls/ls_test.go index 0574e4e06..cedca34ed 100644 --- a/pkg/cmd/ls/ls_test.go +++ b/pkg/cmd/ls/ls_test.go @@ -112,7 +112,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(term, resolveTestCLIAuth(t, s), s, args, showAll, true) } func TestRunLs_APIKeyJSONSkipsUserAndOrgList(t *testing.T) { @@ -192,7 +192,7 @@ func TestGetOrgForRunLs_APIKeyUsesActiveOrgDisplayName(t *testing.T) { s.authTokens = &entity.AuthTokens{APIKey: testAPIKey, APIKeyOrgID: "org-login"} s.org = &entity.Organization{ID: "org-login", Name: "friendly-org"} - org, err := getOrgForRunLs(resolveTestCLIAuth(t, s), s, "") + org, err := getOrgForRunLs(s) if err != nil { t.Fatalf("getOrgForRunLs returned error: %v", err) } @@ -427,7 +427,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(term, resolveTestCLIAuth(t, s), s, nil, true, false) if err != nil { t.Fatalf("RunLs --all returned error: %v", err) } diff --git a/pkg/cmd/org/org.go b/pkg/cmd/org/org.go index 7e0a8dfbc..6d7677f0f 100644 --- a/pkg/cmd/org/org.go +++ b/pkg/cmd/org/org.go @@ -63,10 +63,10 @@ func NewCmdOrg(t *terminal.Terminal, orgcmdStore OrgCmdStore, noorgcmdStore OrgC }, } - cmd.AddCommand(NewCmdOrgSet(t, orgcmdStore, noorgcmdStore)) + cmd.AddCommand(NewCmdOrgSet(t, orgcmdStore)) cmd.AddCommand(NewCmdOrgLs(t, orgcmdStore)) cmd.AddCommand(NewCmdOrgCreate(t, orgcmdStore)) - cmd.AddCommand(invite.NewCmdInvite(t, orgcmdStore, noorgcmdStore)) + cmd.AddCommand(invite.NewCmdInvite(t, orgcmdStore)) // Use default cobra template to show all subcommands cmd.SetUsageTemplate(`Usage:{{if .Runnable}} diff --git a/pkg/cmd/org/set.go b/pkg/cmd/org/set.go index 55dda0992..d9cb5b3f8 100644 --- a/pkg/cmd/org/set.go +++ b/pkg/cmd/org/set.go @@ -4,7 +4,6 @@ import ( "fmt" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/completions" "github.com/brevdev/brev-cli/pkg/cmdcontext" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/store" @@ -13,9 +12,8 @@ import ( "github.com/spf13/cobra" ) -func NewCmdOrgSet(t *terminal.Terminal, orgcmdStore OrgCmdStore, noorgcmdStore OrgCmdStore) *cobra.Command { +func NewCmdOrgSet(t *terminal.Terminal, orgcmdStore OrgCmdStore) *cobra.Command { var showAll bool - var org string cmd := &cobra.Command{ Annotations: map[string]string{"orgsubcommand": ""}, @@ -44,13 +42,6 @@ func NewCmdOrgSet(t *terminal.Terminal, orgcmdStore OrgCmdStore, noorgcmdStore O }, } - cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org)") - err := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noorgcmdStore, t)) - if err != nil { - breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) - fmt.Print(breverrors.WrapAndTrace(err)) - } - cmd.Flags().BoolVar(&showAll, "all", false, "show all workspaces in org") return cmd diff --git a/pkg/cmd/redeem/redeem.go b/pkg/cmd/redeem/redeem.go index c3acac377..9af1770d2 100644 --- a/pkg/cmd/redeem/redeem.go +++ b/pkg/cmd/redeem/redeem.go @@ -5,7 +5,6 @@ import ( "time" "github.com/brevdev/brev-cli/pkg/cmd/cmderrors" - "github.com/brevdev/brev-cli/pkg/cmd/completions" "github.com/brevdev/brev-cli/pkg/cmdcontext" "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" @@ -17,14 +16,10 @@ import ( type RedeemStore interface { GetActiveOrganizationOrDefault() (*entity.Organization, error) - GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error) RedeemCouponCode(organizationID string, code string) (*store.RedeemCouponCodeResponse, error) - completions.CompletionStore } -func NewCmdRedeem(t *terminal.Terminal, redeemStore RedeemStore, noRedeemStore RedeemStore) *cobra.Command { - var orgFlag string - +func NewCmdRedeem(t *terminal.Terminal, redeemStore RedeemStore) *cobra.Command { cmd := &cobra.Command{ Annotations: map[string]string{"organization": ""}, Use: "redeem ", @@ -45,7 +40,7 @@ func NewCmdRedeem(t *terminal.Terminal, redeemStore RedeemStore, noRedeemStore R }, Args: cmderrors.TransformToValidationError(cobra.ExactArgs(1)), RunE: func(cmd *cobra.Command, args []string) error { - err := RunRedeem(t, redeemStore, args[0], orgFlag) + err := RunRedeem(t, redeemStore, args[0]) if err != nil { return breverrors.WrapAndTrace(err) } @@ -53,41 +48,18 @@ func NewCmdRedeem(t *terminal.Terminal, redeemStore RedeemStore, noRedeemStore R }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization (will override active org)") - err := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noRedeemStore, t)) - if err != nil { - breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) - fmt.Print(breverrors.WrapAndTrace(err)) - } - return cmd } -func RunRedeem(t *terminal.Terminal, redeemStore RedeemStore, code string, orgFlag string) error { +func RunRedeem(t *terminal.Terminal, redeemStore RedeemStore, code string) error { startTime := time.Now() - var org *entity.Organization - if orgFlag != "" { - orgs, err := redeemStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgFlag}) - if err != nil { - return breverrors.WrapAndTrace(err) - } - if len(orgs) == 0 { - return fmt.Errorf("no org found with name %s", orgFlag) - } else if len(orgs) > 1 { - return fmt.Errorf("more than one org found with name %s", orgFlag) - } - - org = &orgs[0] - } else { - currOrg, err := redeemStore.GetActiveOrganizationOrDefault() - if err != nil { - return breverrors.WrapAndTrace(err) - } - if currOrg == nil { - return fmt.Errorf("no orgs exist") - } - org = currOrg + org, err := redeemStore.GetActiveOrganizationOrDefault() + if err != nil { + return breverrors.WrapAndTrace(err) + } + if org == nil { + return breverrors.NewValidationError("no orgs exist") } result, err := redeemStore.RedeemCouponCode(org.ID, code) diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 2ad88b435..017cf262e 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -96,7 +96,6 @@ Two modes are supported: ) func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { - var orgFlag string var nameFlag string var sshPort int var approveFlag bool @@ -110,6 +109,10 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { Example: registerExample, Args: cobra.NoArgs, RunE: func(cmd *cobra.Command, args []string) error { + orgFlag, err := cmd.Flags().GetString("org") + if err != nil { + return breverrors.WrapAndTrace(err) + } interactive := nameFlag == "" && orgFlag == "" && sshPort == 0 opts := registerOpts{ interactive: interactive, @@ -122,7 +125,6 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required when using non-interactive mode)") cmd.Flags().StringVarP(&nameFlag, "name", "n", "", "device name (required when using non-interactive mode)") cmd.Flags().IntVarP(&sshPort, "ssh-port", "p", 0, "SSH port (if ssh access is desired)") cmd.Flags().BoolVar(&approveFlag, "approve", false, "skip all confirmation prompts (assume yes)") diff --git a/pkg/cmd/revokessh/revokessh.go b/pkg/cmd/revokessh/revokessh.go index f40b37824..6ff09ffa5 100644 --- a/pkg/cmd/revokessh/revokessh.go +++ b/pkg/cmd/revokessh/revokessh.go @@ -48,7 +48,6 @@ func defaultRevokeSSHDeps() revokeSSHDeps { } func NewCmdRevokeSSH(t *terminal.Terminal, store RevokeSSHStore) *cobra.Command { - var orgFlag string var nodeFlag string var userFlag string var linuxUserFlag string @@ -63,6 +62,10 @@ func NewCmdRevokeSSH(t *terminal.Terminal, store RevokeSSHStore) *cobra.Command Long: "Revoke SSH access to a node for a member of your organization. Interactive: no flags, prompts for org, node, and which access entry to revoke. Non-interactive: --org, --node, --user, --linux-user, and --port-id required.", Example: " brev revoke-ssh\n brev revoke-ssh --org my-org --node my-node --user user@example.com --linux-user ubuntu --port-id port_abc --approve", RunE: func(cmd *cobra.Command, args []string) error { + orgFlag, err := cmd.Flags().GetString("org") + if err != nil { + return breverrors.WrapAndTrace(err) + } interactive := orgFlag == "" && nodeFlag == "" && userFlag == "" && linuxUserFlag == "" && portIDFlag == "" opts := revokeSSHOpts{ interactive: interactive, @@ -77,7 +80,6 @@ func NewCmdRevokeSSH(t *terminal.Terminal, store RevokeSSHStore) *cobra.Command }, } - cmd.Flags().StringVarP(&orgFlag, "org", "o", "", "organization name (required in non-interactive mode)") cmd.Flags().StringVarP(&nodeFlag, "node", "n", "", "node name (required in non-interactive mode)") cmd.Flags().StringVarP(&userFlag, "user", "u", "", "Brev user ID or email to revoke (required in non-interactive mode)") cmd.Flags().StringVar(&linuxUserFlag, "linux-user", "", "Linux username on the target node (required in non-interactive mode)") diff --git a/pkg/cmd/start/start.go b/pkg/cmd/start/start.go index dd98d2b6c..086df062a 100644 --- a/pkg/cmd/start/start.go +++ b/pkg/cmd/start/start.go @@ -49,7 +49,6 @@ type StartStore interface { } func NewCmdStart(t *terminal.Terminal, startStore StartStore, noLoginStartStore StartStore) *cobra.Command { - var org string var name string var detached bool var empty bool @@ -81,11 +80,11 @@ func NewCmdStart(t *terminal.Terminal, startStore StartStore, noLoginStartStore // If stdin is piped, handle multiple instances (only start existing stopped instances) if stdinPiped && len(names) > 0 { - return runBatchStart(t, names, org, setupScript, setupRepo, setupPath, cpu, gpu, piped, startStore) + return runBatchStart(t, names, setupScript, setupRepo, setupPath, cpu, gpu, piped, startStore) } // Single instance mode (original behavior) - return runSingleStart(t, names, name, org, setupScript, setupRepo, setupPath, cpu, gpu, detached, piped, startStore) + return runSingleStart(t, names, name, setupScript, setupRepo, setupPath, cpu, gpu, detached, piped, startStore) }, } cmd.Flags().BoolVarP(&detached, "detached", "d", false, "run the command in the background instead of blocking the shell") @@ -96,21 +95,14 @@ func NewCmdStart(t *terminal.Terminal, startStore StartStore, noLoginStartStore cmd.Flags().StringVarP(&setupScript, "setup-script", "s", "", "takes a raw gist url to an env setup script") cmd.Flags().StringVarP(&setupRepo, "setup-repo", "r", "", "repo that holds env setup script. you must pass in --setup-path if you use this argument") cmd.Flags().StringVarP(&setupPath, "setup-path", "p", "", "path to env setup script. If you include --setup-repo we will apply this argument to that repo") - cmd.Flags().StringVarP(&org, "org", "o", "", "organization (will override active org if creating a workspace)") // GPU options cmd.Flags().StringVarP(&gpu, "gpu", "g", "n1-highmem-4:nvidia-tesla-t4:1", "GPU instance type. Refer to https://docs.nvidia.com/brev/latest/quick-start.html#select-your-compute for more information") - err := cmd.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(noLoginStartStore, t)) - if err != nil { - breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) - fmt.Print(breverrors.WrapAndTrace(err)) - } return cmd } type StartOptions struct { RepoOrPathOrNameOrID string // todo make invidual options Name string - OrgName string SetupScript string SetupRepo string SetupPath string @@ -291,30 +283,16 @@ func createEmptyWorkspace(user *entity.User, apiKeyAuth bool, t *terminal.Termin // ensure org var orgID string - if options.OrgName == "" { - activeorg, err := startStore.GetActiveOrganizationOrDefault() - if err != nil { - return breverrors.WrapAndTrace(err) - } - if activeorg == nil { - return breverrors.NewValidationError("no org exist") - } - orgID = activeorg.ID - } else { - orgs, err := startStore.GetOrganizations(&store.GetOrganizationsOptions{Name: options.OrgName}) - if err != nil { - return breverrors.WrapAndTrace(err) - } - if len(orgs) == 0 { - return breverrors.NewValidationError(fmt.Sprintf("no org with name %s", options.OrgName)) - } else if len(orgs) > 1 { - return breverrors.NewValidationError(fmt.Sprintf("more than one org with name %s", options.OrgName)) - } - orgID = orgs[0].ID + activeorg, err := startStore.GetActiveOrganizationOrDefault() + if err != nil { + return breverrors.WrapAndTrace(err) } + if activeorg == nil { + return breverrors.NewValidationError("no org exist") + } + orgID = activeorg.ID var setupScriptContents string - var err error if len(options.SetupScript) > 0 { contents, err1 := startStore.GetSetupScriptContentsByURL(options.SetupScript) setupScriptContents += "\n" + contents @@ -490,27 +468,14 @@ func createNewWorkspaceFromGit(user *entity.User, apiKeyAuth bool, t *terminal.T } var orgID string - if startOptions.OrgName == "" { - activeorg, err2 := startStore.GetActiveOrganizationOrDefault() - if err2 != nil { - return breverrors.WrapAndTrace(err2) - } - if activeorg == nil { - return breverrors.NewValidationError("no org exist") - } - orgID = activeorg.ID - } else { - orgs, err2 := startStore.GetOrganizations(&store.GetOrganizationsOptions{Name: startOptions.OrgName}) - if err2 != nil { - return breverrors.WrapAndTrace(err2) - } - if len(orgs) == 0 { - return breverrors.NewValidationError(fmt.Sprintf("no org with name %s", startOptions.OrgName)) - } else if len(orgs) > 1 { - return breverrors.NewValidationError(fmt.Sprintf("more than one org with name %s", startOptions.OrgName)) - } - orgID = orgs[0].ID + activeorg, err2 := startStore.GetActiveOrganizationOrDefault() + if err2 != nil { + return breverrors.WrapAndTrace(err2) + } + if activeorg == nil { + return breverrors.NewValidationError("no org exist") } + orgID = activeorg.ID err := createWorkspace(user, apiKeyAuth, t, newWorkspace, orgID, startStore, startOptions) if err != nil { @@ -650,14 +615,13 @@ func pollUntil(t *terminal.Terminal, wsid string, state string, startStore Start // runBatchStart handles starting multiple instances when stdin is piped. // Exit codes: 0 = all succeeded, 1 = all failed, 2 = partial failure. -func runBatchStart(t *terminal.Terminal, names []string, org, setupScript, setupRepo, setupPath, cpu, gpu string, piped bool, startStore StartStore) error { +func runBatchStart(t *terminal.Terminal, names []string, setupScript, setupRepo, setupPath, cpu, gpu string, piped bool, startStore StartStore) error { var startedNames []string var errs error for _, instanceName := range names { err := runStartWorkspace(t, StartOptions{ RepoOrPathOrNameOrID: instanceName, Name: "", - OrgName: org, SetupScript: setupScript, SetupRepo: setupRepo, SetupPath: setupPath, @@ -685,7 +649,7 @@ func runBatchStart(t *terminal.Terminal, names []string, org, setupScript, setup } // runSingleStart handles starting a single instance (original behavior) -func runSingleStart(t *terminal.Terminal, names []string, name, org, setupScript, setupRepo, setupPath, cpu, gpu string, detached, piped bool, startStore StartStore) error { +func runSingleStart(t *terminal.Terminal, names []string, name, setupScript, setupRepo, setupPath, cpu, gpu string, detached, piped bool, startStore StartStore) error { repoOrPathOrNameOrID := "" if len(names) > 0 { repoOrPathOrNameOrID = names[0] @@ -694,7 +658,6 @@ func runSingleStart(t *terminal.Terminal, names []string, name, org, setupScript err := runStartWorkspace(t, StartOptions{ RepoOrPathOrNameOrID: repoOrPathOrNameOrID, Name: name, - OrgName: org, SetupScript: setupScript, SetupRepo: setupRepo, SetupPath: setupPath, diff --git a/pkg/store/http.go b/pkg/store/http.go index 60884f810..5c206492e 100644 --- a/pkg/store/http.go +++ b/pkg/store/http.go @@ -8,6 +8,7 @@ import ( "strings" "github.com/brevdev/brev-cli/pkg/cmd/version" + "github.com/brevdev/brev-cli/pkg/entity" breverrors "github.com/brevdev/brev-cli/pkg/errors" "github.com/brevdev/brev-cli/pkg/featureflag" resty "github.com/go-resty/resty/v2" @@ -50,9 +51,16 @@ type AuthHTTPStore struct { NoAuthHTTPStore authHTTPClient *AuthHTTPClient isRefreshTokenHandlerSet bool + organizationOverride *entity.Organization BasicStore } +// SetOrganizationOverride selects an organization for the lifetime of this +// store. It deliberately does not update the user's persisted active org. +func (s *AuthHTTPStore) SetOrganizationOverride(org *entity.Organization) { + s.organizationOverride = org +} + func (n *NoAuthHTTPStore) GetWindowsDir() (string, error) { return n.GetWSLHostHomeDir() } diff --git a/pkg/store/organization.go b/pkg/store/organization.go index bc02aedb4..630d0e15d 100644 --- a/pkg/store/organization.go +++ b/pkg/store/organization.go @@ -131,6 +131,9 @@ func (s AuthHTTPStore) GetActiveOrganizationOrNil() (*entity.Organization, error // returns the 'set'/active organization or the default one or nil if no orgs exist func (s AuthHTTPStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) { + if s.organizationOverride != nil { + return s.organizationOverride, nil + } org, err := s.GetActiveOrganizationOrNil() if err != nil { return nil, breverrors.WrapAndTrace(err) diff --git a/pkg/store/organization_test.go b/pkg/store/organization_test.go index ec7f0cebb..6e2704205 100644 --- a/pkg/store/organization_test.go +++ b/pkg/store/organization_test.go @@ -44,6 +44,16 @@ func TestGetActiveOrganization(t *testing.T) { } } +func TestGetActiveOrganizationOrDefault_UsesInvocationOverride(t *testing.T) { + fs := MakeMockAuthHTTPStore() + override := &entity.Organization{ID: "org-other", Name: "other"} + fs.SetOrganizationOverride(override) + + org, err := fs.GetActiveOrganizationOrDefault() + require.NoError(t, err) + assert.Same(t, override, org) +} + func TestGetOrganizations(t *testing.T) { fs := MakeMockAuthHTTPStore() httpmock.ActivateNonDefault(fs.authHTTPClient.restyClient.GetClient()) From 30d042d73f99ab9390e85635e7dc9d7394e33e92 Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Wed, 12 Aug 2026 16:57:05 -0700 Subject: [PATCH 2/6] resolve global org override with command-specific auth --- pkg/cmd/cmd.go | 19 ++++++------------- pkg/store/http.go | 10 ++++++++++ pkg/store/organization.go | 13 +++++++++++++ pkg/store/organization_test.go | 15 +++++++++++++++ 4 files changed, 44 insertions(+), 13 deletions(-) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index b4914257a..f81740988 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -186,20 +186,13 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin externalNodeCmdStore.SetOrganizationOverride(nil) } if strings.TrimSpace(orgFlag) != "" { - orgs, orgErr := loginCmdStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgFlag}) - if orgErr != nil { - return breverrors.WrapAndTrace(orgErr) - } - if len(orgs) == 0 { - return breverrors.NewValidationError(fmt.Sprintf("no org found with name %s", orgFlag)) - } - if len(orgs) > 1 { - return breverrors.NewValidationError(fmt.Sprintf("more than one org found with name %s", orgFlag)) - } - loginCmdStore.SetOrganizationOverride(&orgs[0]) - noLoginCmdStore.SetOrganizationOverride(&orgs[0]) + // Resolve the override lazily so commands using an in-memory + // authenticator (such as register) do not authenticate through + // loginCmdStore before their command handler runs. + loginCmdStore.SetOrganizationOverrideName(orgFlag) + noLoginCmdStore.SetOrganizationOverrideName(orgFlag) if externalNodeCmdStore != nil { - externalNodeCmdStore.SetOrganizationOverride(&orgs[0]) + externalNodeCmdStore.SetOrganizationOverrideName(orgFlag) } } home, err := fsStore.GetBrevHomePath() diff --git a/pkg/store/http.go b/pkg/store/http.go index 5c206492e..0df0daec5 100644 --- a/pkg/store/http.go +++ b/pkg/store/http.go @@ -52,6 +52,7 @@ type AuthHTTPStore struct { authHTTPClient *AuthHTTPClient isRefreshTokenHandlerSet bool organizationOverride *entity.Organization + organizationOverrideName string BasicStore } @@ -59,6 +60,15 @@ type AuthHTTPStore struct { // store. It deliberately does not update the user's persisted active org. func (s *AuthHTTPStore) SetOrganizationOverride(org *entity.Organization) { s.organizationOverride = org + s.organizationOverrideName = "" +} + +// SetOrganizationOverrideName selects an organization by name for the lifetime +// of this store. Resolution is deferred until the organization is needed so +// the request uses this store's authentication flow. +func (s *AuthHTTPStore) SetOrganizationOverrideName(name string) { + s.organizationOverride = nil + s.organizationOverrideName = strings.TrimSpace(name) } func (n *NoAuthHTTPStore) GetWindowsDir() (string, error) { diff --git a/pkg/store/organization.go b/pkg/store/organization.go index 630d0e15d..d098b4b8a 100644 --- a/pkg/store/organization.go +++ b/pkg/store/organization.go @@ -134,6 +134,19 @@ func (s AuthHTTPStore) GetActiveOrganizationOrDefault() (*entity.Organization, e if s.organizationOverride != nil { return s.organizationOverride, nil } + if s.organizationOverrideName != "" { + orgs, err := s.GetOrganizations(&GetOrganizationsOptions{Name: s.organizationOverrideName}) + if err != nil { + return nil, breverrors.WrapAndTrace(err) + } + if len(orgs) == 0 { + return nil, breverrors.NewValidationError(fmt.Sprintf("no org found with name %s", s.organizationOverrideName)) + } + if len(orgs) > 1 { + return nil, breverrors.NewValidationError(fmt.Sprintf("more than one org found with name %s", s.organizationOverrideName)) + } + return &orgs[0], nil + } org, err := s.GetActiveOrganizationOrNil() if err != nil { return nil, breverrors.WrapAndTrace(err) diff --git a/pkg/store/organization_test.go b/pkg/store/organization_test.go index 6e2704205..3f0c397b6 100644 --- a/pkg/store/organization_test.go +++ b/pkg/store/organization_test.go @@ -54,6 +54,21 @@ func TestGetActiveOrganizationOrDefault_UsesInvocationOverride(t *testing.T) { assert.Same(t, override, org) } +func TestGetActiveOrganizationOrDefault_ResolvesNamedInvocationOverride(t *testing.T) { + fs := MakeMockAuthHTTPStore() + fs.SetOrganizationOverrideName("other") + httpmock.ActivateNonDefault(fs.authHTTPClient.restyClient.GetClient()) + defer httpmock.DeactivateAndReset() + + expected := entity.Organization{ID: "org-other", Name: "other"} + url := fmt.Sprintf("%s/%s", fs.authHTTPClient.restyClient.BaseURL, orgPath) + httpmock.RegisterResponder("GET", url, httpmock.NewJsonResponderOrPanic(http.StatusOK, []entity.Organization{expected})) + + org, err := fs.GetActiveOrganizationOrDefault() + require.NoError(t, err) + assert.Equal(t, &expected, org) +} + func TestGetOrganizations(t *testing.T) { fs := MakeMockAuthHTTPStore() httpmock.ActivateNonDefault(fs.authHTTPClient.restyClient.GetClient()) From 930ba06242e2aa95f68c83fb8b9a1d4361a93a92 Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Wed, 12 Aug 2026 17:06:54 -0700 Subject: [PATCH 3/6] Do not support --org flag when using api key auth --- pkg/store/organization.go | 3 +++ pkg/store/organization_test.go | 17 +++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/pkg/store/organization.go b/pkg/store/organization.go index d098b4b8a..458e28f3e 100644 --- a/pkg/store/organization.go +++ b/pkg/store/organization.go @@ -131,6 +131,9 @@ func (s AuthHTTPStore) GetActiveOrganizationOrNil() (*entity.Organization, error // returns the 'set'/active organization or the default one or nil if no orgs exist func (s AuthHTTPStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) { + if (s.organizationOverride != nil || s.organizationOverrideName != "") && auth.IsAPIKeyAuthStore(&s) { + return nil, breverrors.NewValidationError("api key auth is scoped to the org saved during login; --org is not supported") + } if s.organizationOverride != nil { return s.organizationOverride, nil } diff --git a/pkg/store/organization_test.go b/pkg/store/organization_test.go index 3f0c397b6..274c0a009 100644 --- a/pkg/store/organization_test.go +++ b/pkg/store/organization_test.go @@ -159,6 +159,23 @@ func TestGetActiveOrganization_APIKeyUsesCredentialOrg(t *testing.T) { assert.Equal(t, "org-api-key", org.Name) } +func TestGetActiveOrganization_APIKeyRejectsInvocationOverride(t *testing.T) { + apiKey := authpkg.BrevAPIKeyPrefix + "test-key" + fileStore, _, _ := newAuthTokenTestStore(t) + s := fileStore.WithAuthHTTPClient(NewAuthHTTPClient(MockAuth{token: &apiKey}, "https://api.test")) + require.NoError(t, s.SaveAuthTokens(entity.AuthTokens{ + APIKey: apiKey, + APIKeyOrgID: "org-api-key", + })) + s.SetOrganizationOverrideName("other") + + org, err := s.GetActiveOrganizationOrDefault() + + assert.Nil(t, org) + require.Error(t, err) + assert.Contains(t, err.Error(), "api key auth is scoped") +} + func TestGetActiveOrganization_APIKeyUsesCredentialOrgNameWhenAvailable(t *testing.T) { apiKey := authpkg.BrevAPIKeyPrefix + "test-key" fileStore, _, _ := newAuthTokenTestStore(t) From e86a0841f9a63aeb664fa70361a3eeb356d38af0 Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Thu, 13 Aug 2026 14:06:01 -0700 Subject: [PATCH 4/6] Refine global org override auth handling --- pkg/auth/auth.go | 8 +++++--- pkg/cmd/cmd.go | 19 ++++++++++++++++++- pkg/cmd/cmd_test.go | 12 ++++++++++++ pkg/cmd/grantssh/grantssh.go | 2 +- pkg/cmd/login/login.go | 9 +++++++-- pkg/cmd/login/login_test.go | 19 +++++++++++++++++++ pkg/cmd/register/register.go | 2 +- pkg/cmd/revokessh/revokessh.go | 2 +- pkg/store/organization.go | 2 +- 9 files changed, 65 insertions(+), 10 deletions(-) diff --git a/pkg/auth/auth.go b/pkg/auth/auth.go index 21de0d855..f1acd7c4c 100644 --- a/pkg/auth/auth.go +++ b/pkg/auth/auth.go @@ -100,9 +100,11 @@ type Auth struct { shouldLogin func() (bool, error) } -const BrevAPIKeyPrefix = "bak-" - -const MissingAPIKeyOrgIDMessage = "api key auth requires an org id; run brev login --api-key --org-id " +const ( + BrevAPIKeyPrefix = "bak-" + MissingAPIKeyOrgIDMessage = "api key auth requires an org id; run brev login --api-key --org-id " + APIKeyOrganizationOverrideNotSupportedMessage = "api key auth is scoped to the org saved during login; --org is not supported" +) type APIKeyAuthStore interface { GetAuthTokens() (*entity.AuthTokens, error) diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index f81740988..63ee6ad20 100644 --- a/pkg/cmd/cmd.go +++ b/pkg/cmd/cmd.go @@ -11,6 +11,7 @@ import ( analyticscmd "github.com/brevdev/brev-cli/pkg/cmd/analytics" "github.com/brevdev/brev-cli/pkg/cmd/background" "github.com/brevdev/brev-cli/pkg/cmd/clipboard" + "github.com/brevdev/brev-cli/pkg/cmd/completions" "github.com/brevdev/brev-cli/pkg/cmd/configureenvvars" "github.com/brevdev/brev-cli/pkg/cmd/connect" "github.com/brevdev/brev-cli/pkg/cmd/copy" @@ -77,6 +78,13 @@ var ( noCheckLatest bool ) +const externalNodeAuthAnnotation = "external-node-auth" + +func allowsOrgOverrideWithExternalAuth(cmd *cobra.Command) bool { + _, ok := cmd.Annotations[externalNodeAuthAnnotation] + return ok +} + func NewDefaultBrevCommand() *cobra.Command { cmd := NewBrevCommand() @@ -84,7 +92,6 @@ func NewDefaultBrevCommand() *cobra.Command { cmd.PersistentFlags().BoolP("help", "h", false, "Help for Brev") cmd.PersistentFlags().StringVar(&userFlag, "user", "", "Non root user to use for per user configuration of commands run as root") - cmd.PersistentFlags().StringVarP(&orgFlag, "org", "o", "", "Organization to use for this command (does not change the active organization)") cmd.PersistentFlags().BoolVar(&printVersion, "version", false, "Print version output") cmd.PersistentFlags().BoolVar(&noCheckLatest, "no-check-latest", false, "Do not check for the latest version when printing version") @@ -152,6 +159,9 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin PersistentPreRunE: func(cmd *cobra.Command, args []string) error { analytics.RecordCommandStart(cmd, args) breverrors.GetDefaultErrorReporter().AddTag("command", cmd.Name()) + if strings.TrimSpace(orgFlag) != "" && !allowsOrgOverrideWithExternalAuth(cmd) && auth.IsAPIKeyAuthStore(loginCmdStore) { + return breverrors.NewValidationError(auth.APIKeyOrganizationOverrideNotSupportedMessage) + } // version info gets in the way of the output for // configure-env-vars, since shells are going to eval it if featureflag.ShowVersionOnRun() && !printVersion && cmd.Name() != "configure-env-vars" { @@ -283,6 +293,13 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin fmt.Printf("%v\n", err) } + cmds.PersistentFlags().StringVarP(&orgFlag, "org", "o", "", "Organization to use for this command (does not change the active organization)") + err = cmds.RegisterFlagCompletionFunc("org", completions.GetOrgsNameCompletionHandler(loginCmdStore, t)) + if err != nil { + breverrors.GetDefaultErrorReporter().ReportError(breverrors.WrapAndTrace(err)) + fmt.Print(breverrors.WrapAndTrace(err)) + } + createCmdTree(cmds, t, loginCmdStore, noLoginCmdStore, loginAuth, externalNodeCmdStore) return cmds diff --git a/pkg/cmd/cmd_test.go b/pkg/cmd/cmd_test.go index eb0dd57d0..a66b7f44e 100644 --- a/pkg/cmd/cmd_test.go +++ b/pkg/cmd/cmd_test.go @@ -8,6 +8,7 @@ import ( "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/store" "github.com/spf13/afero" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -36,6 +37,17 @@ func TestNewDefaultBrevCommandRegistersGlobalOrgFlag(t *testing.T) { flag := cmd.PersistentFlags().Lookup("org") require.NotNil(t, flag) assert.Equal(t, "o", flag.Shorthand) + _, registered := cmd.GetFlagCompletionFunc("org") + assert.True(t, registered) +} + +func TestAllowsOrgOverrideWithExternalAuth(t *testing.T) { + for _, commandName := range []string{"register", "grant-ssh", "revoke-ssh"} { + assert.True(t, allowsOrgOverrideWithExternalAuth(&cobra.Command{Use: commandName, Annotations: map[string]string{externalNodeAuthAnnotation: ""}})) + } + for _, commandName := range []string{"deregister", "enable-ssh", "start"} { + assert.False(t, allowsOrgOverrideWithExternalAuth(&cobra.Command{Use: commandName})) + } } func TestEmailCachingAuthStore_SaveCachesEmail(t *testing.T) { diff --git a/pkg/cmd/grantssh/grantssh.go b/pkg/cmd/grantssh/grantssh.go index 47719eaef..535fa5b49 100644 --- a/pkg/cmd/grantssh/grantssh.go +++ b/pkg/cmd/grantssh/grantssh.go @@ -62,7 +62,7 @@ func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command { var approveFlag bool cmd := &cobra.Command{ - Annotations: map[string]string{"configuration": ""}, + Annotations: map[string]string{"configuration": "", "external-node-auth": ""}, Use: "grant-ssh", DisableFlagsInUseLine: true, Short: "Grant SSH access to a node for another org member", diff --git a/pkg/cmd/login/login.go b/pkg/cmd/login/login.go index 51d044c42..0c18ded8c 100644 --- a/pkg/cmd/login/login.go +++ b/pkg/cmd/login/login.go @@ -52,9 +52,9 @@ type Auth interface { } // loginStore must be a no prompt store -func NewCmdLogin(t *terminal.Terminal, loginStore LoginStore, auth Auth) *cobra.Command { +func NewCmdLogin(t *terminal.Terminal, loginStore LoginStore, loginAuth Auth) *cobra.Command { opts := LoginOptions{ - Auth: auth, + Auth: loginAuth, LoginStore: loginStore, } @@ -75,6 +75,11 @@ func NewCmdLogin(t *terminal.Terminal, loginStore LoginStore, auth Auth) *cobra. Args: cmderrors.TransformToValidationError(cobra.NoArgs), RunE: func(cmd *cobra.Command, args []string) error { apiKeyLogin := strings.TrimSpace(apiKey) != "" + if apiKeyLogin { + if orgFlag := cmd.Flags().Lookup("org"); orgFlag != nil && strings.TrimSpace(orgFlag.Value.String()) != "" { + return breverrors.NewValidationError(auth.APIKeyOrganizationOverrideNotSupportedMessage) + } + } err := opts.RunLogin(t, loginToken, apiKey, apiKeyOrgID, skipBrowser, emailFlag, authProviderFlag) if err != nil { // if err is ImportIDEConfigError, log err with sentry but continue diff --git a/pkg/cmd/login/login_test.go b/pkg/cmd/login/login_test.go index 46fcec712..b0453d906 100644 --- a/pkg/cmd/login/login_test.go +++ b/pkg/cmd/login/login_test.go @@ -8,6 +8,7 @@ import ( "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/store" "github.com/brevdev/brev-cli/pkg/terminal" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -181,6 +182,24 @@ func TestNewCmdLoginWithAPIKey_SkipsPostLoginHooks(t *testing.T) { assert.Equal(t, 0, loginStore.userHomeDirCalls) } +func TestNewCmdLoginWithAPIKeyRejectsGlobalOrgOverride(t *testing.T) { + auth := &mockLoginAuth{} + loginStore := &mockLoginStore{} + loginCmd := NewCmdLogin(terminal.New(), loginStore, auth) + rootCmd := &cobra.Command{Use: "brev"} + rootCmd.PersistentFlags().String("org", "", "organization") + rootCmd.AddCommand(loginCmd) + rootCmd.SetOut(&bytes.Buffer{}) + rootCmd.SetErr(&bytes.Buffer{}) + rootCmd.SetArgs([]string{"login", "--api-key", testAPIKey, "--org-id", "org-test", "--org", "other"}) + + err := rootCmd.Execute() + + require.Error(t, err) + assert.Contains(t, err.Error(), "--org is not supported") + assert.Equal(t, 0, auth.apiKeyCalls) +} + func TestNewCmdLogin_HidesAPIKeyFlagsFromHelp(t *testing.T) { cmd := NewCmdLogin(terminal.New(), &mockLoginStore{}, &mockLoginAuth{}) var out bytes.Buffer diff --git a/pkg/cmd/register/register.go b/pkg/cmd/register/register.go index 017cf262e..babbcee34 100644 --- a/pkg/cmd/register/register.go +++ b/pkg/cmd/register/register.go @@ -101,7 +101,7 @@ func NewCmdRegister(t *terminal.Terminal, store RegisterStore) *cobra.Command { var approveFlag bool cmd := &cobra.Command{ - Annotations: map[string]string{"configuration": ""}, + Annotations: map[string]string{"configuration": "", "external-node-auth": ""}, Use: "register", DisableFlagsInUseLine: true, Short: "Register this device with Brev", diff --git a/pkg/cmd/revokessh/revokessh.go b/pkg/cmd/revokessh/revokessh.go index 6ff09ffa5..7fc66d8d2 100644 --- a/pkg/cmd/revokessh/revokessh.go +++ b/pkg/cmd/revokessh/revokessh.go @@ -55,7 +55,7 @@ func NewCmdRevokeSSH(t *terminal.Terminal, store RevokeSSHStore) *cobra.Command var approveFlag bool cmd := &cobra.Command{ - Annotations: map[string]string{"configuration": ""}, + Annotations: map[string]string{"configuration": "", "external-node-auth": ""}, Use: "revoke-ssh", DisableFlagsInUseLine: true, Short: "Revoke SSH access to a node for an org member", diff --git a/pkg/store/organization.go b/pkg/store/organization.go index 458e28f3e..fc532d8c1 100644 --- a/pkg/store/organization.go +++ b/pkg/store/organization.go @@ -132,7 +132,7 @@ func (s AuthHTTPStore) GetActiveOrganizationOrNil() (*entity.Organization, error // returns the 'set'/active organization or the default one or nil if no orgs exist func (s AuthHTTPStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) { if (s.organizationOverride != nil || s.organizationOverrideName != "") && auth.IsAPIKeyAuthStore(&s) { - return nil, breverrors.NewValidationError("api key auth is scoped to the org saved during login; --org is not supported") + return nil, breverrors.NewValidationError(auth.APIKeyOrganizationOverrideNotSupportedMessage) } if s.organizationOverride != nil { return s.organizationOverride, nil From d949f98910c4c90d32e0af7a28cb964e6931d3e1 Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Thu, 13 Aug 2026 14:27:24 -0700 Subject: [PATCH 5/6] If the --org flag is set, use that org for completions --- pkg/cmd/completions/completions.go | 32 +++++++++++- pkg/cmd/completions/completions_test.go | 69 +++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 pkg/cmd/completions/completions_test.go diff --git a/pkg/cmd/completions/completions.go b/pkg/cmd/completions/completions.go index 1dcf989bc..40d7ed7d7 100644 --- a/pkg/cmd/completions/completions.go +++ b/pkg/cmd/completions/completions.go @@ -1,6 +1,9 @@ package completions import ( + "fmt" + "strings" + "github.com/brevdev/brev-cli/pkg/auth" "github.com/brevdev/brev-cli/pkg/entity" "github.com/brevdev/brev-cli/pkg/store" @@ -20,7 +23,7 @@ type CompletionHandler func(cmd *cobra.Command, args []string, toComplete string func GetAllWorkspaceNameCompletionHandler(completionStore CompletionStore, t *terminal.Terminal) CompletionHandler { return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { - org, err := completionStore.GetActiveOrganizationOrDefault() + org, err := getOrganizationForCompletion(cmd, completionStore) if err != nil { t.Errprint(err, "") return nil, cobra.ShellCompDirectiveError @@ -54,6 +57,33 @@ func GetAllWorkspaceNameCompletionHandler(completionStore CompletionStore, t *te } } +func getOrganizationForCompletion(cmd *cobra.Command, completionStore CompletionStore) (*entity.Organization, error) { + orgFlag := cmd.Flag("org") + if orgFlag == nil || strings.TrimSpace(orgFlag.Value.String()) == "" { + org, err := completionStore.GetActiveOrganizationOrDefault() + if err != nil { + return nil, err + } + return org, nil + } + if auth.IsAPIKeyAuthStore(completionStore) { + return nil, fmt.Errorf("%s", auth.APIKeyOrganizationOverrideNotSupportedMessage) + } + + orgName := strings.TrimSpace(orgFlag.Value.String()) + orgs, err := completionStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgName}) + if err != nil { + return nil, err + } + if len(orgs) == 0 { + return nil, fmt.Errorf("no org found with name %s", orgName) + } + if len(orgs) > 1 { + return nil, fmt.Errorf("more than one org found with name %s", orgName) + } + return &orgs[0], nil +} + func GetOrgsNameCompletionHandler(completionStore CompletionStore, t *terminal.Terminal) CompletionHandler { return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { if auth.IsAPIKeyAuthStore(completionStore) { diff --git a/pkg/cmd/completions/completions_test.go b/pkg/cmd/completions/completions_test.go new file mode 100644 index 000000000..92ec353da --- /dev/null +++ b/pkg/cmd/completions/completions_test.go @@ -0,0 +1,69 @@ +package completions + +import ( + "testing" + + "github.com/brevdev/brev-cli/pkg/entity" + "github.com/brevdev/brev-cli/pkg/store" + "github.com/brevdev/brev-cli/pkg/terminal" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type mockCompletionStore struct { + activeOrg *entity.Organization + organizations []entity.Organization + workspaces map[string][]entity.Workspace + requestedOrg string +} + +func (m *mockCompletionStore) GetAuthTokens() (*entity.AuthTokens, error) { return nil, nil } + +func (m *mockCompletionStore) GetWorkspaces(orgID string, _ *store.GetWorkspacesOptions) ([]entity.Workspace, error) { + m.requestedOrg = orgID + return m.workspaces[orgID], nil +} + +func (m *mockCompletionStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) { + return m.activeOrg, nil +} + +func (m *mockCompletionStore) GetCurrentUser() (*entity.User, error) { + return &entity.User{ID: "user-1"}, nil +} + +func (m *mockCompletionStore) GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error) { + if options == nil || options.Name == "" { + return m.organizations, nil + } + var matches []entity.Organization + for _, org := range m.organizations { + if org.Name == options.Name { + matches = append(matches, org) + } + } + return matches, nil +} + +func TestWorkspaceCompletionUsesOrgOverride(t *testing.T) { + completionStore := &mockCompletionStore{ + activeOrg: &entity.Organization{ID: "org-a", Name: "orgA"}, + organizations: []entity.Organization{{ID: "org-a", Name: "orgA"}, {ID: "org-b", Name: "orgB"}}, + workspaces: map[string][]entity.Workspace{ + "org-a": {{Name: "instance-a"}}, + "org-b": {{Name: "instance-b"}}, + }, + } + rootCmd := &cobra.Command{Use: "brev"} + rootCmd.PersistentFlags().String("org", "", "organization") + shellCmd := &cobra.Command{Use: "shell"} + rootCmd.AddCommand(shellCmd) + require.NoError(t, rootCmd.PersistentFlags().Set("org", "orgB")) + + names, directive := GetAllWorkspaceNameCompletionHandler(completionStore, terminal.New())(shellCmd, nil, "") + + assert.Equal(t, []string{"instance-b"}, names) + assert.Equal(t, cobra.ShellCompDirectiveDefault, directive) + assert.Equal(t, "org-b", completionStore.requestedOrg) +} From e0b934bde8446ec9408e2d362e28afb8304f6496 Mon Sep 17 00:00:00 2001 From: Stephanie Ahart Date: Thu, 13 Aug 2026 14:32:22 -0700 Subject: [PATCH 6/6] Fix lint errors --- pkg/cmd/completions/completions.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pkg/cmd/completions/completions.go b/pkg/cmd/completions/completions.go index 40d7ed7d7..12d687507 100644 --- a/pkg/cmd/completions/completions.go +++ b/pkg/cmd/completions/completions.go @@ -6,6 +6,7 @@ import ( "github.com/brevdev/brev-cli/pkg/auth" "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/spf13/cobra" @@ -62,7 +63,7 @@ func getOrganizationForCompletion(cmd *cobra.Command, completionStore Completion if orgFlag == nil || strings.TrimSpace(orgFlag.Value.String()) == "" { org, err := completionStore.GetActiveOrganizationOrDefault() if err != nil { - return nil, err + return nil, breverrors.WrapAndTrace(err) } return org, nil } @@ -73,7 +74,7 @@ func getOrganizationForCompletion(cmd *cobra.Command, completionStore Completion orgName := strings.TrimSpace(orgFlag.Value.String()) orgs, err := completionStore.GetOrganizations(&store.GetOrganizationsOptions{Name: orgName}) if err != nil { - return nil, err + return nil, breverrors.WrapAndTrace(err) } if len(orgs) == 0 { return nil, fmt.Errorf("no org found with name %s", orgName)