Add global org flag

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 <api-key> --org-id <org-id>"
+const (
+	BrevAPIKeyPrefix                              = "bak-"
+	MissingAPIKeyOrgIDMessage                     = "api key auth requires an org id; run brev login --api-key <api-key> --org-id <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)

pkg/cmd/cmd.go

@@ -3,13 +3,15 @@ package cmd
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/brevdev/brev-cli/pkg/analytics"
 	"github.com/brevdev/brev-cli/pkg/auth"
 	"github.com/brevdev/brev-cli/pkg/cmd/agentskill"
 	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"
@@ -71,10 +73,18 @@ import (
 
 var (
 	userFlag      string
+	orgFlag       string
 	printVersion  bool
 	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()
 
@@ -131,6 +141,7 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin
 
 	analytics.SetUserStore(noLoginCmdStore)
 
+	var externalNodeCmdStore *store.AuthHTTPStore
 	cmds := &cobra.Command{
 		SilenceErrors: true,
 		SilenceUsage:  true,
@@ -148,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" {
@@ -176,6 +190,21 @@ 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) != "" {
+				// 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.SetOrganizationOverrideName(orgFlag)
+				}
+			}
 			home, err := fsStore.GetBrevHomePath()
 			if err != nil {
 				fmt.Printf("Warning: %v", err)
@@ -249,7 +278,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()))
 
@@ -264,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
@@ -273,8 +309,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))

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"
 )
@@ -31,6 +32,24 @@ 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)
+	_, 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) {
 	fs := newTestFileStore(t)
 	s := &emailCachingAuthStore{

pkg/cmd/completions/completions.go

@@ -1,8 +1,12 @@
 package completions
 
 import (
+	"fmt"
+	"strings"
+
 	"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"
@@ -20,7 +24,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 +58,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, breverrors.WrapAndTrace(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, breverrors.WrapAndTrace(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) {

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)
+}

pkg/cmd/grantssh/grantssh.go

@@ -55,21 +55,24 @@ func defaultGrantSSHDeps() grantSSHDeps {
 }
 
 func NewCmdGrantSSH(t *terminal.Terminal, store GrantSSHStore) *cobra.Command {
-	var orgFlag string
 	var nodeFlag string
 	var userFlag string
 	var linuxUser string
 	var portIDFlag string
 	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",
 		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)")