oauth.go

  1package dialog
  2
  3import (
  4	"context"
  5	"fmt"
  6	"strings"
  7
  8	"charm.land/bubbles/v2/help"
  9	"charm.land/bubbles/v2/key"
 10	"charm.land/bubbles/v2/spinner"
 11	tea "charm.land/bubbletea/v2"
 12	"charm.land/catwalk/pkg/catwalk"
 13	"charm.land/lipgloss/v2"
 14	"github.com/charmbracelet/crush/internal/config"
 15	"github.com/charmbracelet/crush/internal/oauth"
 16	"github.com/charmbracelet/crush/internal/ui/common"
 17	"github.com/charmbracelet/crush/internal/ui/util"
 18	uv "github.com/charmbracelet/ultraviolet"
 19	"github.com/pkg/browser"
 20)
 21
 22type OAuthProvider interface {
 23	name() string
 24	initiateAuth() tea.Msg
 25	startPolling(deviceCode string, expiresIn int) tea.Cmd
 26	stopPolling() tea.Msg
 27}
 28
 29// OAuthState represents the current state of the device flow.
 30type OAuthState int
 31
 32const (
 33	OAuthStateInitializing OAuthState = iota
 34	OAuthStateDisplay
 35	OAuthStateSuccess
 36	OAuthStateError
 37)
 38
 39// OAuthID is the identifier for the model selection dialog.
 40const OAuthID = "oauth"
 41
 42// OAuth handles the OAuth flow authentication.
 43type OAuth struct {
 44	com          *common.Common
 45	isOnboarding bool
 46
 47	provider      catwalk.Provider
 48	model         config.SelectedModel
 49	modelType     config.SelectedModelType
 50	oAuthProvider OAuthProvider
 51
 52	State OAuthState
 53
 54	spinner spinner.Model
 55	help    help.Model
 56	keyMap  struct {
 57		Copy   key.Binding
 58		Submit key.Binding
 59		Close  key.Binding
 60	}
 61
 62	width           int
 63	deviceCode      string
 64	userCode        string
 65	verificationURL string
 66	expiresIn       int
 67	interval        int
 68	token           *oauth.Token
 69	cancelFunc      context.CancelFunc
 70}
 71
 72var _ Dialog = (*OAuth)(nil)
 73
 74// newOAuth creates a new device flow component.
 75func newOAuth(
 76	com *common.Common,
 77	isOnboarding bool,
 78	provider catwalk.Provider,
 79	model config.SelectedModel,
 80	modelType config.SelectedModelType,
 81	oAuthProvider OAuthProvider,
 82) (*OAuth, tea.Cmd) {
 83	t := com.Styles
 84
 85	m := OAuth{}
 86	m.com = com
 87	m.isOnboarding = isOnboarding
 88	m.provider = provider
 89	m.model = model
 90	m.modelType = modelType
 91	m.oAuthProvider = oAuthProvider
 92	m.width = 60
 93	m.State = OAuthStateInitializing
 94
 95	m.spinner = spinner.New(
 96		spinner.WithSpinner(spinner.Dot),
 97		spinner.WithStyle(t.Base.Foreground(t.GreenLight)),
 98	)
 99
100	m.help = help.New()
101	m.help.Styles = t.DialogHelpStyles()
102
103	m.keyMap.Copy = key.NewBinding(
104		key.WithKeys("c"),
105		key.WithHelp("c", "copy code"),
106	)
107	m.keyMap.Submit = key.NewBinding(
108		key.WithKeys("enter", "ctrl+y"),
109		key.WithHelp("enter", "copy & open"),
110	)
111	m.keyMap.Close = CloseKey
112
113	return &m, tea.Batch(m.spinner.Tick, m.oAuthProvider.initiateAuth)
114}
115
116// ID implements Dialog.
117func (m *OAuth) ID() string {
118	return OAuthID
119}
120
121// HandleMsg handles messages and state transitions.
122func (m *OAuth) HandleMsg(msg tea.Msg) Action {
123	switch msg := msg.(type) {
124	case spinner.TickMsg:
125		switch m.State {
126		case OAuthStateInitializing, OAuthStateDisplay:
127			var cmd tea.Cmd
128			m.spinner, cmd = m.spinner.Update(msg)
129			if cmd != nil {
130				return ActionCmd{cmd}
131			}
132		}
133
134	case tea.KeyPressMsg:
135		switch {
136		case key.Matches(msg, m.keyMap.Copy):
137			cmd := m.copyCode()
138			return ActionCmd{cmd}
139
140		case key.Matches(msg, m.keyMap.Submit):
141			switch m.State {
142			case OAuthStateSuccess:
143				return m.saveKeyAndContinue()
144
145			default:
146				cmd := m.copyCodeAndOpenURL()
147				return ActionCmd{cmd}
148			}
149
150		case key.Matches(msg, m.keyMap.Close):
151			switch m.State {
152			case OAuthStateSuccess:
153				return m.saveKeyAndContinue()
154
155			default:
156				return ActionClose{}
157			}
158		}
159
160	case ActionInitiateOAuth:
161		m.deviceCode = msg.DeviceCode
162		m.userCode = msg.UserCode
163		m.expiresIn = msg.ExpiresIn
164		m.verificationURL = msg.VerificationURL
165		m.interval = msg.Interval
166		m.State = OAuthStateDisplay
167		return ActionCmd{m.oAuthProvider.startPolling(msg.DeviceCode, msg.ExpiresIn)}
168
169	case ActionCompleteOAuth:
170		m.State = OAuthStateSuccess
171		m.token = msg.Token
172		return ActionCmd{m.oAuthProvider.stopPolling}
173
174	case ActionOAuthErrored:
175		m.State = OAuthStateError
176		cmd := tea.Batch(m.oAuthProvider.stopPolling, util.ReportError(msg.Error))
177		return ActionCmd{cmd}
178	}
179	return nil
180}
181
182// View renders the device flow dialog.
183func (m *OAuth) Draw(scr uv.Screen, area uv.Rectangle) *tea.Cursor {
184	var (
185		t           = m.com.Styles
186		dialogStyle = t.Dialog.View.Width(m.width)
187	)
188	if m.isOnboarding {
189		view := m.dialogContent()
190		DrawOnboarding(scr, area, view)
191	} else {
192		view := dialogStyle.Render(m.dialogContent())
193		DrawCenter(scr, area, view)
194	}
195	return nil
196}
197
198func (m *OAuth) dialogContent() string {
199	var (
200		t         = m.com.Styles
201		helpStyle = t.Dialog.HelpView
202	)
203
204	switch m.State {
205	case OAuthStateInitializing:
206		return m.innerDialogContent()
207
208	default:
209		elements := []string{
210			m.headerContent(),
211			m.innerDialogContent(),
212			helpStyle.Render(m.help.View(m)),
213		}
214		return strings.Join(elements, "\n")
215	}
216}
217
218func (m *OAuth) headerContent() string {
219	var (
220		t            = m.com.Styles
221		titleStyle   = t.Dialog.Title
222		textStyle    = t.Dialog.PrimaryText
223		dialogStyle  = t.Dialog.View.Width(m.width)
224		headerOffset = titleStyle.GetHorizontalFrameSize() + dialogStyle.GetHorizontalFrameSize()
225		dialogTitle  = fmt.Sprintf("Authenticate with %s", m.oAuthProvider.name())
226	)
227	if m.isOnboarding {
228		return textStyle.Render(dialogTitle)
229	}
230	return common.DialogTitle(t, titleStyle.Render(dialogTitle), m.width-headerOffset, t.Primary, t.Secondary)
231}
232
233func (m *OAuth) innerDialogContent() string {
234	var (
235		t            = m.com.Styles
236		whiteStyle   = lipgloss.NewStyle().Foreground(t.White)
237		primaryStyle = lipgloss.NewStyle().Foreground(t.Primary)
238		greenStyle   = lipgloss.NewStyle().Foreground(t.GreenLight)
239		linkStyle    = lipgloss.NewStyle().Foreground(t.GreenDark).Underline(true)
240		errorStyle   = lipgloss.NewStyle().Foreground(t.Error)
241		mutedStyle   = lipgloss.NewStyle().Foreground(t.FgMuted)
242	)
243
244	switch m.State {
245	case OAuthStateInitializing:
246		return lipgloss.NewStyle().
247			Margin(1, 1).
248			Width(m.width - 2).
249			Align(lipgloss.Center).
250			Render(
251				greenStyle.Render(m.spinner.View()) +
252					mutedStyle.Render("Initializing..."),
253			)
254
255	case OAuthStateDisplay:
256		instructions := lipgloss.NewStyle().
257			Margin(0, 1).
258			Width(m.width - 2).
259			Render(
260				whiteStyle.Render("Press ") +
261					primaryStyle.Render("enter") +
262					whiteStyle.Render(" to copy the code below and open the browser."),
263			)
264
265		codeBox := lipgloss.NewStyle().
266			Width(m.width-2).
267			Height(7).
268			Align(lipgloss.Center, lipgloss.Center).
269			Background(t.BgBaseLighter).
270			Margin(0, 1).
271			Render(
272				lipgloss.NewStyle().
273					Bold(true).
274					Foreground(t.White).
275					Render(m.userCode),
276			)
277
278		link := linkStyle.Hyperlink(m.verificationURL, "id=oauth-verify").Render(m.verificationURL)
279		url := mutedStyle.
280			Margin(0, 1).
281			Width(m.width - 2).
282			Render("Browser not opening? Refer to\n" + link)
283
284		waiting := lipgloss.NewStyle().
285			Margin(0, 1).
286			Width(m.width - 2).
287			Render(
288				greenStyle.Render(m.spinner.View()) + mutedStyle.Render("Verifying..."),
289			)
290
291		return lipgloss.JoinVertical(
292			lipgloss.Left,
293			"",
294			instructions,
295			"",
296			codeBox,
297			"",
298			url,
299			"",
300			waiting,
301			"",
302		)
303
304	case OAuthStateSuccess:
305		return greenStyle.
306			Margin(1).
307			Width(m.width - 2).
308			Render("Authentication successful!")
309
310	case OAuthStateError:
311		return lipgloss.NewStyle().
312			Margin(1).
313			Width(m.width - 2).
314			Render(errorStyle.Render("Authentication failed."))
315
316	default:
317		return ""
318	}
319}
320
321// FullHelp returns the full help view.
322func (m *OAuth) FullHelp() [][]key.Binding {
323	return [][]key.Binding{m.ShortHelp()}
324}
325
326// ShortHelp returns the full help view.
327func (m *OAuth) ShortHelp() []key.Binding {
328	switch m.State {
329	case OAuthStateError:
330		return []key.Binding{m.keyMap.Close}
331
332	case OAuthStateSuccess:
333		return []key.Binding{
334			key.NewBinding(
335				key.WithKeys("finish", "ctrl+y", "esc"),
336				key.WithHelp("enter", "finish"),
337			),
338		}
339
340	default:
341		return []key.Binding{
342			m.keyMap.Copy,
343			m.keyMap.Submit,
344			m.keyMap.Close,
345		}
346	}
347}
348
349func (d *OAuth) copyCode() tea.Cmd {
350	if d.State != OAuthStateDisplay {
351		return nil
352	}
353	return tea.Sequence(
354		tea.SetClipboard(d.userCode),
355		util.ReportInfo("Code copied to clipboard"),
356	)
357}
358
359func (d *OAuth) copyCodeAndOpenURL() tea.Cmd {
360	if d.State != OAuthStateDisplay {
361		return nil
362	}
363	return tea.Sequence(
364		tea.SetClipboard(d.userCode),
365		func() tea.Msg {
366			if err := browser.OpenURL(d.verificationURL); err != nil {
367				return ActionOAuthErrored{fmt.Errorf("failed to open browser: %w", err)}
368			}
369			return nil
370		},
371		util.ReportInfo("Code copied and URL opened"),
372	)
373}
374
375func (m *OAuth) saveKeyAndContinue() Action {
376	store := m.com.Store()
377
378	err := store.SetProviderAPIKey(config.ScopeGlobal, string(m.provider.ID), m.token)
379	if err != nil {
380		return ActionCmd{util.ReportError(fmt.Errorf("failed to save API key: %w", err))}
381	}
382
383	return ActionSelectModel{
384		Provider:  m.provider,
385		Model:     m.model,
386		ModelType: m.modelType,
387	}
388}