Move U2M authentication implementation into the CLI

cmd/auth/login.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/browser"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg"
@@ -21,7 +22,6 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/spf13/cobra"
 	"golang.org/x/oauth2"
 )
@@ -47,8 +47,7 @@ const (
 	discoveryFallbackTip    = "\n\nTip: you can specify a workspace directly with: databricks auth login --host <url>"
 	// discoveryHostEnvVar overrides the default https://login.databricks.com
 	// host used by the discovery login flow. Intended for testing and
-	// development against non-production environments. See WithDiscoveryHost
-	// in github.com/databricks/databricks-sdk-go/credentials/u2m.
+	// development against non-production environments.
 	discoveryHostEnvVar = "DATABRICKS_DISCOVERY_HOST"
 )
 
@@ -735,7 +734,8 @@ func discoveryLogin(ctx context.Context, in discoveryLoginInputs) error {
 	// cluster_id, serverless_compute_id) from a prior login to a different host
 	// type must be cleared so they don't leak into the new profile. account_id
 	// and workspace_id are re-added from discovery/introspection results.
-	clearKeys = append(clearKeys,
+	clearKeys = append(
+		clearKeys,
 		"account_id",
 		"workspace_id",
 		databrickscfg.ExperimentalIsUnifiedHostKey,

cmd/auth/login_test.go

@@ -16,11 +16,11 @@ import (
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"

cmd/auth/token.go

@@ -12,6 +12,8 @@ import (
 	"github.com/databricks/cli/cmd/root"
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
+	"github.com/databricks/cli/libs/auth/u2m/cache"
 	"github.com/databricks/cli/libs/browser"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg"
@@ -20,8 +22,6 @@ import (
 	"github.com/databricks/cli/libs/flags"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go/config"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
 	"github.com/spf13/cobra"
 	"golang.org/x/oauth2"
 )

cmd/auth/token_test.go

@@ -11,10 +11,10 @@ import (
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/databrickscfg/profile"
 	"github.com/databricks/cli/libs/env"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
 	"github.com/stretchr/testify/assert"
 	"golang.org/x/oauth2"

libs/auth/arguments.go

@@ -1,8 +1,8 @@
 package auth
 
 import (
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/databricks-sdk-go/config"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 )
 
 // WorkspaceIDNone is a sentinel value persisted to .databrickscfg when the

libs/auth/arguments_test.go

@@ -6,7 +6,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )

libs/auth/credentials.go

@@ -5,11 +5,11 @@ import (
 	"errors"
 
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/config/credentials"
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth"
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 )
 
 // The credentials chain used by the CLI. It is a custom implementation

libs/auth/credentials_test.go

@@ -10,9 +10,9 @@ import (
 	"testing"
 
 	"github.com/databricks/cli/libs/auth/storage"
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/databricks-sdk-go/config"
 	"github.com/databricks/databricks-sdk-go/config/experimental/auth"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"

libs/auth/error.go

@@ -7,9 +7,9 @@ import (
 	"net/http"
 	"strings"
 
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/config"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 )
 
 // Auth type names returned by credential providers.

libs/auth/storage/cache.go

@@ -5,11 +5,11 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/databricks/cli/libs/auth/u2m"
+	"github.com/databricks/cli/libs/auth/u2m/cache"
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/log"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
 )
 
 // storeFactories bundles the constructors ResolveStore depends on. Extracted
@@ -46,9 +46,9 @@ func defaultStoreFactories() storeFactories {
 // fallback does not persist auth_storage = plaintext to [__settings__];
 // pinning happens only on successful login.
 //
-// Every CLI code path that calls u2m.NewPersistentAuth must route the result
-// through u2m.WithTokenCache, otherwise the SDK defaults to the file cache
-// and splits the user's tokens across two backends.
+// Every CLI code path that calls u2m.NewPersistentAuth must supply the cache
+// returned by this package, otherwise U2M uses an in-memory cache and bypasses
+// the user's configured storage backend.
 func ResolveStore(ctx context.Context, override StorageMode) (Store, StorageMode, error) {
 	return resolveStoreForReadWith(ctx, override, defaultStoreFactories())
 }
@@ -75,7 +75,7 @@ func ResolveStoreForLogin(ctx context.Context, override StorageMode) (Store, Sto
 	return resolveStoreForLoginWith(ctx, override, defaultStoreFactories())
 }
 
-// OAuthTokenCache adapts a CLI Store to the SDK's u2m_cache.TokenCache for the
+// OAuthTokenCache adapts a CLI Store to the U2M cache.TokenCache for the
 // U2M PersistentAuth flow, applying the not-found hint so a cache miss carries
 // actionable "run databricks auth login" guidance. Use on read and credential
 // paths. M2M/OIDC callers use the CLI Store directly and must not route through

libs/auth/storage/cache_test.go

@@ -9,9 +9,9 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/databricks/cli/libs/auth/u2m"
 	"github.com/databricks/cli/libs/databrickscfg"
 	"github.com/databricks/cli/libs/env"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"

libs/auth/storage/dual_writing_cache.go

@@ -1,8 +1,8 @@
 package storage
 
 import (
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
-	u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"github.com/databricks/cli/libs/auth/u2m"
+	u2m_cache "github.com/databricks/cli/libs/auth/u2m/cache"
 	"golang.org/x/oauth2"
 )
 
@@ -12,7 +12,7 @@ import (
 // implemented inside PersistentAuth.dualWrite in the SDK, now moved
 // caller-side per the cache-ownership split between SDK and CLI.
 //
-// Mirroring happens inside Store, so every SDK-internal write (Challenge,
+// Mirroring happens inside Store, so every U2M-internal write (Challenge,
 // refresh, discovery) dual-writes without requiring each call site to invoke
 // a helper explicitly.
 type DualWritingTokenCache struct {

libs/auth/storage/dual_writing_cache_test.go

@@ -5,8 +5,8 @@ import (
 	"sync"
 	"testing"
 
-	"github.com/databricks/databricks-sdk-go/credentials/u2m"
-	u2m_cache "github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"github.com/databricks/cli/libs/auth/u2m"
+	u2m_cache "github.com/databricks/cli/libs/auth/u2m/cache"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"

libs/auth/storage/keyring.go

@@ -14,7 +14,7 @@ import (
 
 // keyringServiceName is the service name used for every entry the CLI writes
 // to the OS-native secure store. The account field carries the per-entry
-// cache key the SDK passes through TokenCache.Store / Lookup.
+// cache key the U2M manager passes through TokenCache.Store / Lookup.
 const keyringServiceName = "databricks-cli"
 
 // keyringProbeAccountPrefix is prefixed onto a per-call random suffix to form

libs/auth/storage/not_found_hint.go

@@ -7,8 +7,8 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/databricks/cli/libs/auth/u2m/cache"
 	"github.com/databricks/cli/libs/env"
-	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
 	"golang.org/x/oauth2"
 )
 
@@ -19,7 +19,7 @@ import (
 // their cached credentials are no longer being read.
 //
 // errors.Is(err, cache.ErrNotFound) continues to return true because the
-// wrap uses %w; the SDK's branches on ErrNotFound still fire.
+// wrap uses %w; PersistentAuth's branches on ErrNotFound still fire.
 //
 // Store is delegated unchanged; only Lookup needs the message polish.
 type notFoundHintCache struct {
@@ -45,7 +45,7 @@ func (c *notFoundHintCache) Lookup(key string) (*oauth2.Token, error) {
 
 // notFoundHint replaces cache.ErrNotFound's terse "token not found" string
 // with an actionable message while still satisfying errors.Is(err,
-// cache.ErrNotFound). The SDK's loadToken wraps every cache error with
+// cache.ErrNotFound). PersistentAuth.loadToken wraps every cache error with
 // "cache: %w", and fmt.Errorf("...: %w", ErrNotFound) would tack the
 // original "token not found" onto the end of our hint, producing
 // "cache: <hint>: token not found". A custom type lets us own the

libs/auth/storage/not_found_hint_test.go

@@ -8,7 +8,7 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/databricks/databricks-sdk-go/credentials/u2m/cache"
+	"github.com/databricks/cli/libs/auth/u2m/cache"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/oauth2"

libs/auth/storage/storage.go

@@ -16,8 +16,8 @@ import (
 // ErrNotFound is returned by Store.Lookup when no entry exists for the key, or
 // when a stored entry cannot be decoded by this CLI version (an unknown format
 // is treated as a miss so the caller re-mints rather than failing). It is the
-// CLI-owned counterpart to the SDK's u2m_cache.ErrNotFound; the adapter in
-// ToU2MTokenCache translates between the two.
+// CLI-owned counterpart to the U2M cache.ErrNotFound; the adapter in
+// ToU2MTokenCache translates between the storage and OAuth layers.
 var ErrNotFound = errors.New("token not found")
 
 // Entry is the value held in the CLI token store. It wraps the credential so
@@ -33,7 +33,7 @@ type Entry struct {
 // policy of its own. Implementations are the plaintext file cache and the OS
 // keyring cache. The interface is owned by the CLI rather than the SDK so the
 // entry schema can evolve (metadata, per-entry resilience) without being
-// constrained by the SDK's U2M-internal u2m_cache.TokenCache, which only
+// constrained by the U2M-internal cache.TokenCache, which only
 // carries a bare *oauth2.Token.
 type Store interface {
 	// Put writes e under key, replacing any existing entry.

libs/auth/storage/u2m.go

@@ -0,0 +1,43 @@
+package storage
+
+import (
+	"errors"
+
+	"github.com/databricks/cli/libs/auth/u2m/cache"
+	"golang.org/x/oauth2"
+)
+
+// ToU2MTokenCache adapts a CLI Store to the U2M cache.TokenCache interface.
+// The Store(key, nil) delete convention maps to Store.Delete.
+func ToU2MTokenCache(s Store) cache.TokenCache {
+	return &u2mTokenCache{store: s}
+}
+
+// u2mTokenCache is the ToU2MTokenCache adapter.
+type u2mTokenCache struct {
+	store Store
+}
+
+// Store implements cache.TokenCache. A nil token is the U2M delete
+// signal; everything else is a plain put with no metadata.
+func (tc *u2mTokenCache) Store(key string, t *oauth2.Token) error {
+	if t == nil {
+		return tc.store.Delete(key)
+	}
+	return tc.store.Put(key, Entry{Token: t})
+}
+
+// Lookup implements cache.TokenCache, translating the storage miss sentinel
+// to the U2M package's sentinel.
+func (tc *u2mTokenCache) Lookup(key string) (*oauth2.Token, error) {
+	e, err := tc.store.Lookup(key)
+	if err != nil {
+		if errors.Is(err, ErrNotFound) {
+			return nil, cache.ErrNotFound
+		}
+		return nil, err
+	}
+	return e.Token, nil
+}
+
+var _ cache.TokenCache = (*u2mTokenCache)(nil)