feat(BRE2-915): api key as auth method

Author: patelspratik

pkg/auth/auth.go

@@ -101,6 +101,24 @@ type Auth struct {
 	shouldLogin          func() (bool, error)
 }
 
+const BrevAPIKeyPrefix = "brev_api_"
+
+type TokenProvider interface {
+	GetAccessToken() (string, error)
+}
+
+func IsBrevAPIKey(token string) bool {
+	return strings.HasPrefix(strings.TrimSpace(token), BrevAPIKeyPrefix)
+}
+
+func IsAPIKeyAuthStore(tokenProvider TokenProvider) bool {
+	token, err := tokenProvider.GetAccessToken()
+	if err != nil {
+		return false
+	}
+	return IsBrevAPIKey(token)
+}
+
 func NewAuth(authStore AuthStore, oauth OAuth) *Auth {
 	return &Auth{
 		authStore:            authStore,
@@ -146,6 +164,14 @@ func (t Auth) GetFreshAccessTokenOrNil() (string, error) {
 		return "", nil
 	}
 
+	if tokens.APIKey != "" {
+		apiKey := strings.TrimSpace(tokens.APIKey)
+		if apiKey == "" {
+			return "", breverrors.NewValidationError("api key is empty")
+		}
+		return apiKey, nil
+	}
+
 	// should always at least have access token?
 	if tokens.AccessToken == "" {
 		breverrors.GetDefaultErrorReporter().ReportMessage("access token is an empty string but shouldn't be")
@@ -222,6 +248,31 @@ func (t Auth) LoginWithToken(token string) error {
 	return nil
 }
 
+func (t Auth) LoginWithAPIKey(apiKey string) error {
+	apiKey = strings.TrimSpace(apiKey)
+	if apiKey == "" {
+		return breverrors.NewValidationError("api key is empty")
+	}
+	if !strings.HasPrefix(apiKey, BrevAPIKeyPrefix) {
+		return breverrors.NewValidationError(fmt.Sprintf("api key must start with %s", BrevAPIKeyPrefix))
+	}
+
+	tokens, err := t.getSavedTokensOrNil()
+	if err != nil {
+		return breverrors.WrapAndTrace(err)
+	}
+	if tokens == nil {
+		tokens = &entity.AuthTokens{}
+	}
+	tokens.APIKey = apiKey
+
+	err = t.authStore.SaveAuthTokens(*tokens)
+	if err != nil {
+		return breverrors.WrapAndTrace(err)
+	}
+	return nil
+}
+
 // showLoginURL displays the login link and CLI alternative for manual navigation.
 func showLoginURL(url string) {
 	urlType := color.New(color.FgCyan, color.Bold).SprintFunc()
@@ -313,7 +364,7 @@ func (t Auth) getSavedTokensOrNil() (*entity.AuthTokens, error) {
 		}
 		return nil, breverrors.WrapAndTrace(err)
 	}
-	if tokens != nil && tokens.AccessToken == "" && tokens.RefreshToken == "" {
+	if tokens != nil && tokens.AccessToken == "" && tokens.RefreshToken == "" && tokens.APIKey == "" {
 		return nil, nil
 	}
 	return tokens, nil
@@ -415,7 +466,7 @@ func AuthProviderFlagToCredentialProvider(authProviderFlag string) entity.Creden
 func StandardLogin(authProvider string, email string, tokens *entity.AuthTokens) OAuth {
 	// Set KAS as the default authenticator
 	shouldPromptEmail := false
-	if email == "" && tokens != nil && tokens.AccessToken != "" {
+	if email == "" && tokens != nil && tokens.AccessToken != "" && tokens.APIKey == "" {
 		email = GetEmailFromToken(tokens.AccessToken)
 		shouldPromptEmail = true
 	}
@@ -445,7 +496,7 @@ func StandardLogin(authProvider string, email string, tokens *entity.AuthTokens)
 		kasAuthenticator,
 	})
 
-	if tokens != nil && tokens.AccessToken != "" {
+	if tokens != nil && tokens.AccessToken != "" && tokens.APIKey == "" {
 		authenticatorFromToken, errr := authRetriever.GetByToken(tokens.AccessToken)
 		if errr != nil {
 			fmt.Printf("%v\n", errr)

pkg/auth/auth_test.go

@@ -1,6 +1,8 @@
 package auth
 
 import (
+	"io"
+	"os"
 	"testing"
 
 	"github.com/brevdev/brev-cli/pkg/entity"
@@ -38,10 +40,12 @@ func TestIsAccessTokenValid(t *testing.T) {
 
 type MockAuthStore struct {
 	authTokens *entity.AuthTokens
+	saved      entity.AuthTokens
 	didSave    bool
 }
 
-func (m *MockAuthStore) SaveAuthTokens(_ entity.AuthTokens) error {
+func (m *MockAuthStore) SaveAuthTokens(tokens entity.AuthTokens) error {
+	m.saved = tokens
 	m.didSave = true
 	return nil
 }
@@ -79,6 +83,116 @@ func (m MockOauth) GetNewAuthTokensWithRefresh(_ string) (*entity.AuthTokens, er
 
 const validToken = "abc"
 
+func TestGetFreshAccessTokenOrNil_APIKeySkipsJWTValidationAndRefresh(t *testing.T) {
+	s := MockAuthStore{authTokens: &entity.AuthTokens{
+		AccessToken:  "expired-jwt",
+		APIKey:       "brev_api_test-key",
+		RefreshToken: "should-not-refresh",
+	}}
+	a := Auth{
+		&s,
+		&MockOauth{}, func(_ string) (bool, error) {
+			t.Fatal("api keys must not be parsed as JWTs")
+			return false, nil
+		},
+		func() (bool, error) {
+			t.Fatal("api keys must not trigger login")
+			return false, nil
+		},
+	}
+
+	res, err := a.GetFreshAccessTokenOrNil()
+	assert.NoError(t, err)
+	assert.Equal(t, "brev_api_test-key", res)
+	assert.False(t, s.didSave)
+}
+
+func TestGetFreshAccessTokenOrNil_APIKeyOnlyCredentialReturnsAPIKey(t *testing.T) {
+	s := MockAuthStore{authTokens: &entity.AuthTokens{
+		APIKey: "brev_api_test-key",
+	}}
+	a := Auth{
+		&s,
+		&MockOauth{}, func(_ string) (bool, error) {
+			t.Fatal("api keys must not be parsed as JWTs")
+			return false, nil
+		},
+		func() (bool, error) {
+			t.Fatal("api keys must not trigger login")
+			return false, nil
+		},
+	}
+
+	res, err := a.GetFreshAccessTokenOrNil()
+	assert.NoError(t, err)
+	assert.Equal(t, "brev_api_test-key", res)
+	assert.False(t, s.didSave)
+}
+
+func TestLoginWithAPIKey_SavesTypedCredential(t *testing.T) {
+	s := MockAuthStore{}
+	a := Auth{
+		authStore: &s,
+		oauth:     &MockOauth{},
+	}
+
+	err := a.LoginWithAPIKey("brev_api_test-key")
+	assert.NoError(t, err)
+	assert.True(t, s.didSave)
+	assert.Equal(t, entity.AuthTokens{
+		APIKey: "brev_api_test-key",
+	}, s.saved)
+}
+
+func TestLoginWithAPIKey_PreservesExistingJWT(t *testing.T) {
+	s := MockAuthStore{authTokens: &entity.AuthTokens{
+		AccessToken:  "existing-jwt",
+		RefreshToken: "existing-refresh",
+	}}
+	a := Auth{
+		authStore: &s,
+		oauth:     &MockOauth{},
+	}
+
+	err := a.LoginWithAPIKey("brev_api_test-key")
+	assert.NoError(t, err)
+	assert.Equal(t, entity.AuthTokens{
+		AccessToken:  "existing-jwt",
+		RefreshToken: "existing-refresh",
+		APIKey:       "brev_api_test-key",
+	}, s.saved)
+}
+
+func TestLoginWithAPIKey_EmptyKeyReturnsError(t *testing.T) {
+	s := MockAuthStore{}
+	a := Auth{
+		authStore: &s,
+		oauth:     &MockOauth{},
+	}
+
+	err := a.LoginWithAPIKey("")
+	assert.Error(t, err)
+	assert.False(t, s.didSave)
+}
+
+func TestStandardLogin_APIKeyCredentialDoesNotProbeOAuthProviders(t *testing.T) {
+	oldStdout := os.Stdout
+	readPipe, writePipe, err := os.Pipe()
+	assert.NoError(t, err)
+	os.Stdout = writePipe
+
+	_ = StandardLogin("", "", &entity.AuthTokens{
+		AccessToken: "existing-jwt",
+		APIKey:      "brev_api_test-key",
+	})
+
+	assert.NoError(t, writePipe.Close())
+	os.Stdout = oldStdout
+	out, err := io.ReadAll(readPipe)
+	assert.NoError(t, err)
+	assert.Empty(t, string(out))
+}
+
 func TestSuccessNoRefreshGetFreshAccessTokenOrLogin(t *testing.T) {
 	s := MockAuthStore{authTokens: &entity.AuthTokens{
 		AccessToken:  validToken,

pkg/cmd/completions/completions.go

@@ -1,6 +1,7 @@
 package completions
 
 import (
+	"github.com/brevdev/brev-cli/pkg/auth"
 	"github.com/brevdev/brev-cli/pkg/entity"
 	"github.com/brevdev/brev-cli/pkg/store"
 	"github.com/brevdev/brev-cli/pkg/terminal"
@@ -18,12 +19,6 @@ 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) {
-		user, err := completionStore.GetCurrentUser()
-		if err != nil {
-			t.Errprint(err, "")
-			return nil, cobra.ShellCompDirectiveError
-		}
-
 		org, err := completionStore.GetActiveOrganizationOrDefault()
 		if err != nil {
 			t.Errprint(err, "")
@@ -33,7 +28,17 @@ func GetAllWorkspaceNameCompletionHandler(completionStore CompletionStore, t *te
 			return []string{}, cobra.ShellCompDirectiveDefault
 		}
 
-		workspaces, err := completionStore.GetWorkspaces(org.ID, &store.GetWorkspacesOptions{UserID: user.ID})
+		var options *store.GetWorkspacesOptions
+		if tokenProvider, ok := completionStore.(auth.TokenProvider); !ok || !auth.IsAPIKeyAuthStore(tokenProvider) {
+			user, err := completionStore.GetCurrentUser()
+			if err != nil {
+				t.Errprint(err, "")
+				return nil, cobra.ShellCompDirectiveError
+			}
+			options = &store.GetWorkspacesOptions{UserID: user.ID}
+		}
+
+		workspaces, err := completionStore.GetWorkspaces(org.ID, options)
 		if err != nil {
 			t.Errprint(err, "")
 			return nil, cobra.ShellCompDirectiveError

pkg/cmd/delete/delete.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/brevdev/brev-cli/pkg/auth"
 	"github.com/brevdev/brev-cli/pkg/cmd/completions"
 	"github.com/brevdev/brev-cli/pkg/cmd/util"
 	"github.com/brevdev/brev-cli/pkg/entity"
@@ -99,6 +100,9 @@ func deleteWorkspace(workspaceName string, t *terminal.Terminal, deleteStore Del
 
 func handleAdminUser(err error, deleteStore DeleteStore, piped bool) error {
 	if strings.Contains(err.Error(), "not found") {
+		if tokenProvider, ok := deleteStore.(auth.TokenProvider); ok && auth.IsAPIKeyAuthStore(tokenProvider) {
+			return breverrors.WrapAndTrace(err)
+		}
 		user, err1 := deleteStore.GetCurrentUser()
 		if err1 != nil {
 			return breverrors.WrapAndTrace(err1)

pkg/cmd/gpucreate/gpucreate.go

@@ -14,6 +14,7 @@ import (
 	"time"
 	"unicode"
 
+	"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"
@@ -97,6 +98,7 @@ type GPUCreateStore interface {
 	util.GetWorkspaceByNameOrIDErrStore
 	gpusearch.GPUSearchStore
 	GetActiveOrganizationOrDefault() (*entity.Organization, error)
+	GetAccessToken() (string, error)
 	GetCurrentUser() (*entity.User, error)
 	GetWorkspace(workspaceID string) (*entity.Workspace, error)
 	CreateWorkspace(organizationID string, options *store.CreateWorkspacesOptions) (*entity.Workspace, error)
@@ -720,10 +722,13 @@ func newCreateContext(t *terminal.Terminal, store GPUCreateStore, opts GPUCreate
 		}
 	}
 
-	// Get user
-	user, err := store.GetCurrentUser()
-	if err != nil {
-		return nil, breverrors.WrapAndTrace(err)
+	user := &entity.User{}
+	if !auth.IsAPIKeyAuthStore(store) {
+		var err error
+		user, err = store.GetCurrentUser()
+		if err != nil {
+			return nil, breverrors.WrapAndTrace(err)
+		}
 	}
 	ctx.user = user
 
@@ -877,7 +882,11 @@ func (c *createContext) waitForInstances(workspaces []*entity.Workspace) {
 	for _, ws := range workspaces {
 		err := c.pollUntilReady(ws.ID)
 		if err != nil {
-			c.logf("  %s: Timeout waiting for ready state\n", ws.Name)
+			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))
+			}
 		}
 	}
 }
@@ -1220,6 +1229,9 @@ func (c *createContext) pollUntilReady(wsID string) error {
 		}
 
 		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))
 		}
 

pkg/cmd/gpucreate/gpucreate_test.go

@@ -3,6 +3,7 @@ package gpucreate
 import (
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/brevdev/brev-cli/pkg/cmd/gpusearch"
 	"github.com/brevdev/brev-cli/pkg/entity"
@@ -44,6 +45,10 @@ func (m *MockGPUCreateStore) GetCurrentUser() (*entity.User, error) {
 	return m.User, nil
 }
 
+func (m *MockGPUCreateStore) GetAccessToken() (string, error) {
+	return "", nil
+}
+
 func (m *MockGPUCreateStore) GetActiveOrganizationOrDefault() (*entity.Organization, error) {
 	return m.Org, nil
 }
@@ -641,3 +646,22 @@ func TestFormatInstanceSpecs(t *testing.T) {
 	result := formatInstanceSpecs(specs)
 	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",
+	}
+
+	ctx := &createContext{
+		store: store,
+		opts:  GPUCreateOptions{Timeout: time.Second},
+	}
+
+	err := ctx.pollUntilReady("ws-failed")
+
+	assert.ErrorContains(t, err, "instance test failed: unexpected end of JSON input")
+}