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