diff --git a/cmd/cmd.go b/cmd/cmd.go index ce7187aee..91b278db9 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -8,6 +8,7 @@ import ( "crypto/ed25519" "crypto/rand" "crypto/sha256" + "encoding/base64" "encoding/pem" "errors" "fmt" @@ -16,6 +17,7 @@ import ( "math" "net" "net/http" + "net/url" "os" "os/signal" "path/filepath" @@ -30,6 +32,7 @@ import ( "github.com/containerd/console" "github.com/mattn/go-runewidth" "github.com/olekukonko/tablewriter" + "github.com/pkg/browser" "github.com/spf13/cobra" "golang.org/x/crypto/ssh" "golang.org/x/term" @@ -516,7 +519,26 @@ func RunHandler(cmd *cobra.Command, args []string) error { return generate(cmd, opts) } -func errFromUnknownKey(unknownKeyErr error) error { +func generateFingerprint(key string) string { + hash := sha256.Sum256([]byte(key)) + fingerprint := base64.RawURLEncoding.EncodeToString(hash[:6]) + + var formatted strings.Builder + for i, char := range fingerprint { + if i > 0 && i%2 == 0 { + formatted.WriteRune('-') + } + formatted.WriteRune(char) + } + + return formatted.String() +} + +// tryConnect handles key validation when a connection fails due to an unknown key. +// It attempts to open the browser for interactive sessions to let users connect their key, +// falling back to command-line instructions for non-interactive sessions. +// Returns nil if browser flow succeeds, or an error with connection instructions otherwise. +func tryConnect(unknownKeyErr error) error { // find SSH public key in the error message // TODO (brucemacd): the API should return structured errors so that this message parsing isn't needed sshKeyPattern := `ssh-\w+ [^\s"]+` @@ -545,14 +567,23 @@ func errFromUnknownKey(unknownKeyErr error) error { return unknownKeyErr } - var msg strings.Builder - msg.WriteString(unknownKeyErr.Error()) - msg.WriteString("\n\nYour ollama key is:\n") - msg.WriteString(localPubKey) - msg.WriteString("\nAdd your key at:\n") - msg.WriteString("https://ollama.com/settings/keys") + if term.IsTerminal(int(os.Stdout.Fd())) { + // URL encode the key and device name for the browser URL + encodedKey := base64.RawURLEncoding.EncodeToString([]byte(localPubKey)) + d, _ := os.Hostname() + encodedDevice := url.QueryEscape(d) + browserURL := fmt.Sprintf("https://ollama.com/connect?host=%s&key=%s", encodedDevice, encodedKey) - return errors.New(msg.String()) + if err := browser.OpenURL(browserURL); err == nil { + fmt.Printf("\nOpening browser to add your key...\n") + fmt.Printf("\nCheck that this code matches what is shown in your browser:\n") + fmt.Printf("\n %s\n", generateFingerprint(localPubKey)) + return nil + } + } + + // only return error for non-interactive terminals or if browser opening failed + return fmt.Errorf("%s\nAdd your key at:\nhttps://ollama.com/settings/keys", unknownKeyErr.Error()) } return unknownKeyErr @@ -611,13 +642,16 @@ func PushHandler(cmd *cobra.Command, args []string) error { if spinner != nil { spinner.Stop() } + if p != nil { + p.Stop() + } if strings.Contains(err.Error(), "access denied") { return errors.New("you are not authorized to push to this namespace, create the model under a namespace you own") } if strings.Contains(err.Error(), errtypes.UnknownOllamaKeyErrMsg) && isOllamaHost { // the user has not added their ollama key to ollama.com // return an error with a more user-friendly message - return errFromUnknownKey(err) + return tryConnect(err) } return err } diff --git a/go.mod b/go.mod index 496d8d3ab..f56a221bb 100644 --- a/go.mod +++ b/go.mod @@ -22,6 +22,7 @@ require ( github.com/mattn/go-runewidth v0.0.14 github.com/nlpodyssey/gopickle v0.3.0 github.com/pdevine/tensor v0.0.0-20240510204454-f88f4562727c + github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c golang.org/x/image v0.22.0 ) diff --git a/go.sum b/go.sum index b3093ceb9..04332cf15 100644 --- a/go.sum +++ b/go.sum @@ -159,6 +159,8 @@ github.com/phpdave11/gofpdf v1.4.2/go.mod h1:zpO6xFn9yxo3YLyMvW8HcKWVdbNqgIfOOp2 github.com/phpdave11/gofpdi v1.0.12/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI= github.com/pierrec/lz4/v4 v4.1.8 h1:ieHkV+i2BRzngO4Wd/3HGowuZStgq6QkPsD1eolNAO4= github.com/pierrec/lz4/v4 v4.1.8/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=