ccb1a64 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

refactor: make oauth dialog generic and move provider logic to interface

ccb1a64Andrey Nering9 months ago

internal/ui/dialog/oauth.go+19−76

compare with main

unifiedsplit

@@ -4,7 +4,6 @@ import (

4 "context"

5 "fmt"

6 "strings"

7− "time"

9 "charm.land/bubbles/v2/help"

8 "charm.land/bubbles/v2/help"

10 "charm.land/bubbles/v2/key"

9 "charm.land/bubbles/v2/key"

@@ -14,13 +13,19 @@ import (

14 "github.com/charmbracelet/catwalk/pkg/catwalk"

13 "github.com/charmbracelet/catwalk/pkg/catwalk"

15 "github.com/charmbracelet/crush/internal/config"

14 "github.com/charmbracelet/crush/internal/config"

16 "github.com/charmbracelet/crush/internal/oauth"

15 "github.com/charmbracelet/crush/internal/oauth"

17− "github.com/charmbracelet/crush/internal/oauth/hyper"

18 "github.com/charmbracelet/crush/internal/ui/common"

16 "github.com/charmbracelet/crush/internal/ui/common"

19 "github.com/charmbracelet/crush/internal/uiutil"

17 "github.com/charmbracelet/crush/internal/uiutil"

20 uv "github.com/charmbracelet/ultraviolet"

18 uv "github.com/charmbracelet/ultraviolet"

21 "github.com/pkg/browser"

19 "github.com/pkg/browser"

22)

20)

22+type OAuthProvider interface {

23+ name() string

24+ initiateAuth() tea.Msg

25+ startPolling(deviceCode string, expiresIn int) tea.Cmd

26+ stopPolling() tea.Msg

27+}

28+

24// OAuthState represents the current state of the device flow.

29// OAuthState represents the current state of the device flow.

25type OAuthState int

30type OAuthState int

@@ -38,9 +43,10 @@ const OAuthID = "oauth"

38type OAuth struct {

43type OAuth struct {

39 com *common.Common

44 com *common.Common

41− provider catwalk.Provider

46+ provider catwalk.Provider

42− model config.SelectedModel

47+ model config.SelectedModel

43− modelType config.SelectedModelType

48+ modelType config.SelectedModelType

49+ oAuthProvider OAuthProvider

45 State OAuthState

51 State OAuthState

@@ -63,8 +69,8 @@ type OAuth struct {

64var _ Dialog = (*OAuth)(nil)

70var _ Dialog = (*OAuth)(nil)

66−// NewOAuth creates a new device flow component.

72+// newOAuth creates a new device flow component.

67−func NewOAuth(com *common.Common, provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) (*OAuth, error) {

73+func newOAuth(com *common.Common, provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType, oAuthProvider OAuthProvider) (*OAuth, error) {

68 t := com.Styles

74 t := com.Styles

70 m := OAuth{}

76 m := OAuth{}

@@ -72,6 +78,7 @@ func NewOAuth(com *common.Common, provider catwalk.Provider, model config.Select

72 m.provider = provider

78 m.provider = provider

73 m.model = model

79 m.model = model

74 m.modelType = modelType

80 m.modelType = modelType

81+ m.oAuthProvider = oAuthProvider

75 m.width = 60

82 m.width = 60

76 m.State = OAuthStateInitializing

83 m.State = OAuthStateInitializing

@@ -103,7 +110,7 @@ func (m *OAuth) ID() string {

103

110

104// Init implements Dialog.

111// Init implements Dialog.

105func (m *OAuth) Init() tea.Cmd {

112func (m *OAuth) Init() tea.Cmd {

106− return tea.Batch(m.spinner.Tick, m.initiateDeviceAuth)

113+ return tea.Batch(m.spinner.Tick, m.oAuthProvider.initiateAuth)

107}

114}

108

115

109// HandleMsg handles messages and state transitions.

116// HandleMsg handles messages and state transitions.

@@ -151,16 +158,16 @@ func (m *OAuth) HandleMsg(msg tea.Msg) Action {

151 m.expiresIn = msg.ExpiresIn

158 m.expiresIn = msg.ExpiresIn

152 m.verificationURL = msg.VerificationURL

159 m.verificationURL = msg.VerificationURL

153 m.State = OAuthStateDisplay

160 m.State = OAuthStateDisplay

154− return ActionCmd{m.startPolling(msg.DeviceCode)}

161+ return ActionCmd{m.oAuthProvider.startPolling(msg.DeviceCode, msg.ExpiresIn)}

155

162

156 case ActionCompleteOAuth:

163 case ActionCompleteOAuth:

157 m.State = OAuthStateSuccess

164 m.State = OAuthStateSuccess

158 m.token = msg.Token

165 m.token = msg.Token

159− return ActionCmd{m.stopPolling}

166+ return ActionCmd{m.oAuthProvider.stopPolling}

160

167

161 case ActionOAuthErrored:

168 case ActionOAuthErrored:

162 m.State = OAuthStateError

169 m.State = OAuthStateError

163− cmd := tea.Batch(m.stopPolling, uiutil.ReportError(msg.Error))

170+ cmd := tea.Batch(m.oAuthProvider.stopPolling, uiutil.ReportError(msg.Error))

164 return ActionCmd{cmd}

171 return ActionCmd{cmd}

165 }

172 }

166 return nil

173 return nil

@@ -204,7 +211,7 @@ func (m *OAuth) headerContent() string {

204 dialogStyle = t.Dialog.View.Width(m.width)

211 dialogStyle = t.Dialog.View.Width(m.width)

205 headerOffset = titleStyle.GetHorizontalFrameSize() + dialogStyle.GetHorizontalFrameSize()

212 headerOffset = titleStyle.GetHorizontalFrameSize() + dialogStyle.GetHorizontalFrameSize()

206 )

213 )

207− return common.DialogTitle(t, titleStyle.Render("Authenticate with Hyper"), m.width-headerOffset)

214+ return common.DialogTitle(t, titleStyle.Render("Authenticate with "+m.oAuthProvider.name()), m.width-headerOffset)

208}

215}

209

216

210func (m *OAuth) innerDialogContent() string {

217func (m *OAuth) innerDialogContent() string {

@@ -357,67 +364,3 @@ func (m *OAuth) saveKeyAndContinue() Action {

357 ModelType: m.modelType,

364 ModelType: m.modelType,

358 }

365 }

359}

366}

360−

361−func (m *OAuth) initiateDeviceAuth() tea.Msg {

362− minimumWait := 750 * time.Millisecond

363− startTime := time.Now()

364−

365− ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)

366− defer cancel()

367−

368− authResp, err := hyper.InitiateDeviceAuth(ctx)

369−

370− ellapsed := time.Since(startTime)

371− if ellapsed < minimumWait {

372− time.Sleep(minimumWait - ellapsed)

373− }

374−

375− if err != nil {

376− return ActionOAuthErrored{fmt.Errorf("failed to initiate device auth: %w", err)}

377− }

378−

379− return ActionInitiateOAuth{

380− DeviceCode: authResp.DeviceCode,

381− UserCode: authResp.UserCode,

382− ExpiresIn: authResp.ExpiresIn,

383− VerificationURL: authResp.VerificationURL,

384− }

385−}

386−

387−// startPolling starts polling for the device token.

388−func (m *OAuth) startPolling(deviceCode string) tea.Cmd {

389− return func() tea.Msg {

390− ctx, cancel := context.WithCancel(context.Background())

391− m.cancelFunc = cancel

392−

393− refreshToken, err := hyper.PollForToken(ctx, deviceCode, m.expiresIn)

394− if err != nil {

395− if ctx.Err() != nil {

396− return nil

397− }

398− return ActionOAuthErrored{err}

399− }

400−

401− token, err := hyper.ExchangeToken(ctx, refreshToken)

402− if err != nil {

403− return ActionOAuthErrored{fmt.Errorf("token exchange failed: %w", err)}

404− }

405−

406− introspect, err := hyper.IntrospectToken(ctx, token.AccessToken)

407− if err != nil {

408− return ActionOAuthErrored{fmt.Errorf("token introspection failed: %w", err)}

409− }

410− if !introspect.Active {

411− return ActionOAuthErrored{fmt.Errorf("access token is not active")}

412− }

413−

414− return ActionCompleteOAuth{token}

415− }

416−}

417−

418−func (m *OAuth) stopPolling() tea.Msg {

419− if m.cancelFunc != nil {

420− m.cancelFunc()

421− }

422− return nil

423−}

internal/ui/dialog/oauth_hyper.gonew file+90

@@ -0,0 +1,90 @@

1+package dialog

3+import (

4+ "context"

5+ "fmt"

6+ "time"

8+ tea "charm.land/bubbletea/v2"

9+ "github.com/charmbracelet/catwalk/pkg/catwalk"

10+ "github.com/charmbracelet/crush/internal/config"

11+ "github.com/charmbracelet/crush/internal/oauth/hyper"

12+ "github.com/charmbracelet/crush/internal/ui/common"

13+)

14+

15+func NewOAuthHyper(com *common.Common, provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) (*OAuth, error) {

16+ return newOAuth(com, provider, model, modelType, &OAuthHyper{})

17+}

18+

19+type OAuthHyper struct {

20+ cancelFunc func()

21+}

22+

23+var _ OAuthProvider = (*OAuthHyper)(nil)

24+

25+func (m *OAuthHyper) name() string {

26+ return "Hyper"

27+}

28+

29+func (m *OAuthHyper) initiateAuth() tea.Msg {

30+ minimumWait := 750 * time.Millisecond

31+ startTime := time.Now()

32+

33+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)

34+ defer cancel()

35+

36+ authResp, err := hyper.InitiateDeviceAuth(ctx)

37+

38+ ellapsed := time.Since(startTime)

39+ if ellapsed < minimumWait {

40+ time.Sleep(minimumWait - ellapsed)

41+ }

42+

43+ if err != nil {

44+ return ActionOAuthErrored{fmt.Errorf("failed to initiate device auth: %w", err)}

45+ }

46+

47+ return ActionInitiateOAuth{

48+ DeviceCode: authResp.DeviceCode,

49+ UserCode: authResp.UserCode,

50+ ExpiresIn: authResp.ExpiresIn,

51+ VerificationURL: authResp.VerificationURL,

52+ }

53+}

54+

55+func (m *OAuthHyper) startPolling(deviceCode string, expiresIn int) tea.Cmd {

56+ return func() tea.Msg {

57+ ctx, cancel := context.WithCancel(context.Background())

58+ m.cancelFunc = cancel

59+

60+ refreshToken, err := hyper.PollForToken(ctx, deviceCode, expiresIn)

61+ if err != nil {

62+ if ctx.Err() != nil {

63+ return nil

64+ }

65+ return ActionOAuthErrored{err}

66+ }

67+

68+ token, err := hyper.ExchangeToken(ctx, refreshToken)

69+ if err != nil {

70+ return ActionOAuthErrored{fmt.Errorf("token exchange failed: %w", err)}

71+ }

72+

73+ introspect, err := hyper.IntrospectToken(ctx, token.AccessToken)

74+ if err != nil {

75+ return ActionOAuthErrored{fmt.Errorf("token introspection failed: %w", err)}

76+ }

77+ if !introspect.Active {

78+ return ActionOAuthErrored{fmt.Errorf("access token is not active")}

79+ }

80+

81+ return ActionCompleteOAuth{token}

82+ }

83+}

84+

85+func (m *OAuthHyper) stopPolling() tea.Msg {

86+ if m.cancelFunc != nil {

87+ m.cancelFunc()

88+ }

89+ return nil

90+}

internal/ui/model/ui.go+8−4

@@ -1020,9 +1020,9 @@ func substituteArgs(content string, args map[string]string) string {

1020func (m *UI) openAuthenticationDialog(provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) tea.Cmd {

1021 switch provider.ID {

1022 case "hyper":

1023− return m.openOAuthDialog(provider, model, modelType)

1023+ return m.openOAuthHyperDialog(provider, model, modelType)

1024 case catwalk.InferenceProviderCopilot:

1025− return m.openOAuthDialog(provider, model, modelType)

1025+ return m.openOAuthCopilotDialog(provider, model, modelType)

1026 default:

1027 return m.openAPIKeyInputDialog(provider, model, modelType)

1028 }

@@ -1042,13 +1042,13 @@ func (m *UI) openAPIKeyInputDialog(provider catwalk.Provider, model config.Selec

1042 return nil

1043}

1044

1045−func (m *UI) openOAuthDialog(provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) tea.Cmd {

1045+func (m *UI) openOAuthHyperDialog(provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) tea.Cmd {

1046 if m.dialog.ContainsDialog(dialog.OAuthID) {

1047 m.dialog.BringToFront(dialog.OAuthID)

1048 return nil

1049 }

1050

1051− oAuthDialog, err := dialog.NewOAuth(m.com, provider, model, modelType)

1051+ oAuthDialog, err := dialog.NewOAuthHyper(m.com, provider, model, modelType)

1052 if err != nil {

1053 return uiutil.ReportError(err)

1054 }

@@ -1057,6 +1057,10 @@ func (m *UI) openOAuthDialog(provider catwalk.Provider, model config.SelectedMod

1057 return oAuthDialog.Init()

1058}

1059

1060+func (m *UI) openOAuthCopilotDialog(provider catwalk.Provider, model config.SelectedModel, modelType config.SelectedModelType) tea.Cmd {

1061+ panic("TODO")

1062+}

1063+

1060func (m *UI) handleKeyPressMsg(msg tea.KeyPressMsg) tea.Cmd {

1064func (m *UI) handleKeyPressMsg(msg tea.KeyPressMsg) tea.Cmd {

1061 var cmds []tea.Cmd

1065 var cmds []tea.Cmd

1062

1066