pkg/cli/stream.go

@@ -39,6 +39,7 @@ type Stream struct {
 	wrap int
 	col  int // current column position (visible characters)
 	done chan struct{}
+	once sync.Once
 	mu   sync.Mutex
 }
 
@@ -107,12 +108,14 @@ func (s *Stream) WriteFrom(r io.Reader) error {
 
 // Done signals that no more text will arrive.
 func (s *Stream) Done() {
-	s.mu.Lock()
-	if s.col > 0 {
-		fmt.Fprintln(s.out) // ensure trailing newline
-	}
-	s.mu.Unlock()
-	close(s.done)
+	s.once.Do(func() {
+		s.mu.Lock()
+		if s.col > 0 {
+			fmt.Fprintln(s.out) // ensure trailing newline
+		}
+		s.mu.Unlock()
+		close(s.done)
+	})
 }
 
 // Wait blocks until Done is called.

pkg/cli/stream_test.go

@@ -153,6 +153,18 @@ func TestStream_Good(t *testing.T) {
 		assert.Equal(t, "text\n", buf.String()) // no double newline
 	})
 
+	t.Run("Done is idempotent", func(t *testing.T) {
+		var buf bytes.Buffer
+		s := NewStream(WithStreamOutput(&buf))
+
+		s.Write("text")
+		s.Done()
+		s.Done()
+		s.Wait()
+
+		assert.Equal(t, "text\n", buf.String())
+	})
+
 	t.Run("word wrap uses visible width", func(t *testing.T) {
 		var buf bytes.Buffer
 		s := NewStream(WithWordWrap(4), WithStreamOutput(&buf))