Spinner for "tsh beams add" (#66881)

* "tsh beams add" with spinner

* add recconect message

* add example usage

* minor

* fix re-login

* address comments

* add t.Parallel

lib/utils/spinner/spinner.go

@@ -0,0 +1,83 @@
+/*
+ * Teleport
+ * Copyright (C) 2025  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package spinner
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/charmbracelet/bubbles/spinner"
+	"github.com/charmbracelet/lipgloss"
+)
+
+// Spinner renders an inline animated spinner to a writer.
+type Spinner struct {
+	w     io.Writer
+	model spinner.Spinner
+	style lipgloss.Style
+
+	stop     chan struct{}
+	done     chan struct{}
+	stopOnce sync.Once
+}
+
+// New creates and starts an inline spinner that writes to w.
+// Call Stop to replace the spinner line with a final message.
+func New(w io.Writer, msg string) *Spinner {
+	s := &Spinner{
+		w:    w,
+		stop: make(chan struct{}),
+		done: make(chan struct{}),
+		// ANSI color "6" is cyan. style and model are hard-coded atm but can
+		// potentially become options for spinners.
+		style: lipgloss.NewStyle().Foreground(lipgloss.Color("6")).Bold(true),
+		model: spinner.MiniDot,
+	}
+	go s.run(msg)
+	return s
+}
+
+func (s *Spinner) run(msg string) {
+	defer close(s.done)
+	ticker := time.NewTicker(s.model.FPS)
+	defer ticker.Stop()
+	i := 0
+	for {
+		select {
+		case <-s.stop:
+			fmt.Fprintf(s.w, "\r%s\r", strings.Repeat(" ", len(msg)+2))
+			return
+		case <-ticker.C:
+			frame := s.model.Frames[i%len(s.model.Frames)]
+			fmt.Fprintf(s.w, "\r%s %s", s.style.Render(frame), msg)
+			i++
+		}
+	}
+}
+
+// Stop clears the spinner line. Safe to call multiple times.
+func (s *Spinner) Stop() {
+	s.stopOnce.Do(func() {
+		close(s.stop)
+		<-s.done
+	})
+}

lib/utils/spinner/spinner_test.go

@@ -0,0 +1,51 @@
+/*
+ * Teleport
+ * Copyright (C) 2025  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package spinner
+
+import (
+	"bytes"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSpinner(t *testing.T) {
+	t.Parallel()
+	synctest.Test(t, func(t *testing.T) {
+		var buf bytes.Buffer
+		s := New(&buf, "creating...")
+		t.Cleanup(s.Stop)
+
+		// Sanity check that the bubbletea spinner model defines enough frames
+		// to produce a visible animation. The threshold is arbitrary.
+		require.Greater(t, len(s.model.Frames), 5)
+
+		time.Sleep(s.model.FPS * time.Duration(len(s.model.Frames)*2))
+		s.Stop()
+
+		// Every frame should be preceded by a carriage return after two full cycles.
+		output := buf.String()
+		for _, frame := range s.model.Frames {
+			assert.Contains(t, output, "\r"+s.style.Render(frame)+" creating...")
+		}
+	})
+}

tool/tsh/common/beams_add.go

@@ -20,22 +20,27 @@ package common
 
 import (
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/alecthomas/kingpin/v2"
+	"github.com/charmbracelet/lipgloss"
 	"github.com/gravitational/trace"
 
 	"github.com/gravitational/teleport"
 	beamsv1 "github.com/gravitational/teleport/api/gen/proto/go/teleport/beams/v1"
 	"github.com/gravitational/teleport/lib/client"
 	"github.com/gravitational/teleport/lib/defaults"
+	"github.com/gravitational/teleport/lib/utils"
+	"github.com/gravitational/teleport/lib/utils/spinner"
 	"github.com/gravitational/teleport/tool/common"
 )
 
 type beamsAddCommand struct {
 	*kingpin.CmdClause
-	console bool
-	format  string
+	console             bool
+	format              string
+	isTerminalOverwrite func(io.Writer) bool
 }
 
 func newBeamsAddCommand(parent *kingpin.CmdClause) *beamsAddCommand {
@@ -47,6 +52,7 @@ func newBeamsAddCommand(parent *kingpin.CmdClause) *beamsAddCommand {
 		Short('f').
 		Default(teleport.Text).
 		EnumVar(&cmd.format, defaults.DefaultFormats...)
+	cmd.Alias(beamsAddHelp)
 	return cmd
 }
 
@@ -66,6 +72,12 @@ func (c *beamsAddCommand) run(cf *CLIConf) error {
 		}
 		defer clusterClient.Close()
 
+		// Show spinner after successful connection to avoid spinning during re-login.
+		if c.shouldShowSpinner(cf.Stdout(), c.format) {
+			creatingBeamSpinner := spinner.New(cf.Stdout(), "Creating beam...")
+			defer creatingBeamSpinner.Stop()
+		}
+
 		rootClient, err := clusterClient.ConnectToRootCluster(ctx)
 		if err != nil {
 			return trace.Wrap(err)
@@ -109,9 +121,30 @@ func (c *beamsAddCommand) run(cf *CLIConf) error {
 
 		// Connect to the beam via SSH.
 		if c.console {
-			return trace.Wrap(sshBeam(cf, tc, beam, nil))
+			if err := sshBeam(cf, tc, beam, nil); err != nil {
+				return trace.Wrap(err)
+			}
+			return trace.Wrap(c.printReconnectMessage(cf.Stdout(), beam.GetStatus().GetAlias()))
 		}
 	}
 
 	return nil
 }
+
+func (c *beamsAddCommand) printReconnectMessage(w io.Writer, alias string) error {
+	gray := lipgloss.NewStyle().Foreground(lipgloss.Color("8"))
+	_, err := fmt.Fprintln(w, gray.Render(fmt.Sprintf("\nTo reconnect to this beam, run:\n    tsh beams ssh %s", alias)))
+	return trace.Wrap(err)
+}
+
+func (c *beamsAddCommand) shouldShowSpinner(w io.Writer, format string) bool {
+	switch strings.ToLower(format) {
+	case teleport.JSON, teleport.YAML:
+		return false
+	default:
+		if c.isTerminalOverwrite != nil {
+			return c.isTerminalOverwrite(w)
+		}
+		return utils.IsTerminal(w)
+	}
+}

tool/tsh/common/beams_add_test.go

@@ -0,0 +1,87 @@
+/*
+ * Teleport
+ * Copyright (C) 2026  Gravitational, Inc.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package common
+
+import (
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport"
+)
+
+func Test_beamsAddCommand_shouldShowSpinner(t *testing.T) {
+	t.Parallel()
+	fakeTTY := func(io.Writer) bool { return true }
+	fakeNonTTY := func(io.Writer) bool { return false }
+
+	tests := []struct {
+		name                string
+		format              string
+		isTerminalOverwrite func(io.Writer) bool
+		check               require.BoolAssertionFunc
+	}{
+		{
+			name:                "json format with TTY",
+			format:              teleport.JSON,
+			isTerminalOverwrite: fakeTTY,
+			check:               require.False,
+		},
+		{
+			name:                "yaml format with TTY",
+			format:              teleport.YAML,
+			isTerminalOverwrite: fakeTTY,
+			check:               require.False,
+		},
+		{
+			name:                "text format with TTY",
+			format:              teleport.Text,
+			isTerminalOverwrite: fakeTTY,
+			check:               require.True,
+		},
+		{
+			name:                "text format with non-TTY",
+			format:              teleport.Text,
+			isTerminalOverwrite: fakeNonTTY,
+			check:               require.False,
+		},
+		{
+			name:                "empty format with TTY",
+			format:              "",
+			isTerminalOverwrite: fakeTTY,
+			check:               require.True,
+		},
+		{
+			name:                "empty format with non-TTY",
+			format:              "",
+			isTerminalOverwrite: fakeNonTTY,
+			check:               require.False,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cmd := &beamsAddCommand{
+				isTerminalOverwrite: tt.isTerminalOverwrite,
+			}
+			tt.check(t, cmd.shouldShowSpinner(io.Discard, tt.format))
+		})
+	}
+}

tool/tsh/common/help.go

@@ -126,4 +126,13 @@ Examples:
   Add the database configuration to the specified JSON file
   $ tsh mcp db config --db-user=mydbuser --db-name=mydbname --client-config=my-config.json my-db-resource
 `
+
+	beamsAddHelp = `
+Examples:
+  Create a beam and connect via SSH:
+  $ tsh beams add
+
+  Create a beam without connecting and save its ID:
+  $ BEAM_ID=$(tsh beams add -f json | jq -r '.id')
+  $ tsh beams exec $BEAM_ID -- ls examples`
 )