chore: use tea.WithContext(inv.Context()) for tea program

DanielleMaywood · DanielleMaywood · commit 907824252b10 · 2024-09-04T10:18:49.000Z
diff --git a/cli/cliui/select.go b/cli/cliui/select.go
@@ -92,18 +92,25 @@ func Select(inv *serpent.Invocation, opts SelectOptions) (string, error) {
 
 	m, err := tea.NewProgram(
 		initialModel,
+		tea.WithContext(inv.Context()),
 		tea.WithInput(inv.Stdin),
 		tea.WithOutput(inv.Stdout),
 	).Run()
 
-	var value string
-	if m, ok := m.(selectModel); ok {
-		if m.canceled {
-			return value, Canceled
-		}
-		value = m.selected
+	if err != nil {
+		return "", err
+	}
+
+	model, ok := m.(selectModel)
+	if !ok {
+		return "", xerrors.New(fmt.Sprintf("unknown model found %T (%+v)", m, m))
 	}
-	return value, err
+
+	if model.canceled {
+		return "", Canceled
+	}
+
+	return model.selected, err
 }
 
 type selectModel struct {
@@ -286,19 +293,25 @@ func MultiSelect(inv *serpent.Invocation, opts MultiSelectOptions) ([]string, er
 
 	m, err := tea.NewProgram(
 		initialModel,
+		tea.WithContext(inv.Context()),
 		tea.WithInput(inv.Stdin),
 		tea.WithOutput(inv.Stdout),
 	).Run()
 
-	values := []string{}
-	if m, ok := m.(multiSelectModel); ok {
-		if m.canceled {
-			return values, Canceled
-		}
+	if err != nil {
+		return nil, err
+	}
 
-		values = m.selectedOptions()
+	model, ok := m.(multiSelectModel)
+	if !ok {
+		return nil, xerrors.New(fmt.Sprintf("unknown model found %T (%+v)", m, m))
 	}
-	return values, err
+
+	if model.canceled {
+		return nil, Canceled
+	}
+
+	return model.selectedOptions(), err
 }
 
 type multiSelectOption struct {
