fix(cli): route stream output through injected stdout writer
Some checks are pending
Security Scan / security (push) Waiting to run

Co-Authored-By: Virgil <virgil@lethean.io>
This commit is contained in:
Virgil 2026-04-02 17:13:26 +00:00
parent 050ee5bd8f
commit cdae3a9ac5
3 changed files with 20 additions and 6 deletions

View file

@ -34,7 +34,8 @@ When word-wrap is enabled, the stream tracks the current column position and ins
## Custom Output Writer ## Custom Output Writer
By default, streams write to `os.Stdout`. Redirect to any `io.Writer`: By default, streams write to the CLI stdout writer (`stdoutWriter()`), so tests can
redirect output via `cli.SetStdout` and other callers can provide any `io.Writer`:
```go ```go
var buf strings.Builder var buf strings.Builder
@ -77,7 +78,7 @@ stream.Done()
| Option | Description | | Option | Description |
|--------|-------------| |--------|-------------|
| `WithWordWrap(cols)` | Set the word-wrap column width | | `WithWordWrap(cols)` | Set the word-wrap column width |
| `WithStreamOutput(w)` | Set the output writer (default: `os.Stdout`) | | `WithStreamOutput(w)` | Set the output writer (default: `stdoutWriter()`) |
## Example: LLM Token Streaming ## Example: LLM Token Streaming

View file

@ -3,7 +3,6 @@ package cli
import ( import (
"fmt" "fmt"
"io" "io"
"os"
"strings" "strings"
"sync" "sync"
@ -12,7 +11,8 @@ import (
// StreamOption configures a Stream. // StreamOption configures a Stream.
// //
// stream := cli.NewStream(cli.WithWordWrap(80), cli.WithStreamOutput(os.Stdout)) // stream := cli.NewStream(cli.WithWordWrap(80))
// stream.Wait()
type StreamOption func(*Stream) type StreamOption func(*Stream)
// WithWordWrap sets the word-wrap column width. // WithWordWrap sets the word-wrap column width.
@ -20,7 +20,7 @@ func WithWordWrap(cols int) StreamOption {
return func(s *Stream) { s.wrap = cols } return func(s *Stream) { s.wrap = cols }
} }
// WithStreamOutput sets the output writer (default: os.Stdout). // WithStreamOutput sets the output writer (default: stdoutWriter()).
func WithStreamOutput(w io.Writer) StreamOption { func WithStreamOutput(w io.Writer) StreamOption {
return func(s *Stream) { s.out = w } return func(s *Stream) { s.out = w }
} }
@ -48,7 +48,7 @@ type Stream struct {
// NewStream creates a streaming text renderer. // NewStream creates a streaming text renderer.
func NewStream(opts ...StreamOption) *Stream { func NewStream(opts ...StreamOption) *Stream {
s := &Stream{ s := &Stream{
out: os.Stdout, out: stdoutWriter(),
done: make(chan struct{}), done: make(chan struct{}),
} }
for _, opt := range opts { for _, opt := range opts {

View file

@ -10,6 +10,19 @@ import (
) )
func TestStream_Good(t *testing.T) { func TestStream_Good(t *testing.T) {
t.Run("uses injected stdout by default", func(t *testing.T) {
var buf bytes.Buffer
SetStdout(&buf)
defer SetStdout(nil)
s := NewStream()
s.Write("hello")
s.Done()
s.Wait()
assert.Equal(t, "hello\n", buf.String())
})
t.Run("basic write", func(t *testing.T) { t.Run("basic write", func(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
s := NewStream(WithStreamOutput(&buf)) s := NewStream(WithStreamOutput(&buf))