Allow the fake data to be streamed with a small delay

Signed-off-by: David Gageot <david.gageot@docker.com>

Author: dgageot

cmd/root/api.go

@@ -98,7 +98,7 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 	}
 
 	// Start fake proxy if --fake is specified
-	cleanup, err := setupFakeProxy(f.fakeResponses, &f.runConfig)
+	cleanup, err := setupFakeProxy(f.fakeResponses, 0, &f.runConfig)
 	if err != nil {
 		return err
 	}

cmd/root/record.go

@@ -11,16 +11,25 @@ import (
 )
 
 // setupFakeProxy starts a fake proxy if fakeResponses is non-empty.
+// streamDelayMs controls simulated streaming: 0 = disabled, >0 = delay in milliseconds between chunks.
 // It returns a cleanup function that must be called when done (typically via defer).
-func setupFakeProxy(fakeResponses string, runConfig *config.RuntimeConfig) (cleanup func() error, err error) {
+func setupFakeProxy(fakeResponses string, streamDelayMs int, runConfig *config.RuntimeConfig) (cleanup func() error, err error) {
 	if fakeResponses == "" {
 		return func() error { return nil }, nil
 	}
 
 	// Normalize path by stripping .yaml suffix (go-vcr adds it automatically)
 	fakeResponses = strings.TrimSuffix(fakeResponses, ".yaml")
 
-	proxyURL, cleanupFn, err := fake.StartProxy(fakeResponses)
+	var opts []fake.ProxyOption
+	if streamDelayMs > 0 {
+		opts = append(opts,
+			fake.WithSimulateStream(true),
+			fake.WithStreamChunkDelay(time.Duration(streamDelayMs)*time.Millisecond),
+		)
+	}
+
+	proxyURL, cleanupFn, err := fake.StartProxy(fakeResponses, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to start fake proxy: %w", err)
 	}

cmd/root/run.go

@@ -8,7 +8,6 @@ import (
 	"os"
 	"path/filepath"
 	"runtime/pprof"
-	"strings"
 
 	"github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
@@ -37,6 +36,7 @@ type runExecFlags struct {
 	sessionID         string
 	recordPath        string
 	fakeResponses     string
+	fakeStreamDelay   int
 	exitAfterResponse bool
 	cpuProfile        string
 	forceTUI          bool
@@ -86,6 +86,8 @@ func addRunOrExecFlags(cmd *cobra.Command, flags *runExecFlags) {
 	cmd.PersistentFlags().StringVarP(&flags.sessionDB, "session-db", "s", filepath.Join(paths.GetHomeDir(), ".cagent", "session.db"), "Path to the session database")
 	cmd.PersistentFlags().StringVar(&flags.sessionID, "session", "", "Continue from a previous session by ID")
 	cmd.PersistentFlags().StringVar(&flags.fakeResponses, "fake", "", "Replay AI responses from cassette file (for testing)")
+	cmd.PersistentFlags().IntVar(&flags.fakeStreamDelay, "fake-stream", 0, "Simulate streaming with delay in ms between chunks (default 15ms if no value given)")
+	cmd.Flag("fake-stream").NoOptDefVal = "15" // --fake-stream without value uses 15ms
 	cmd.PersistentFlags().StringVar(&flags.recordPath, "record", "", "Record AI API interactions to cassette file (auto-generates filename if empty)")
 	cmd.PersistentFlags().Lookup("record").NoOptDefVal = "true"
 	cmd.PersistentFlags().BoolVar(&flags.exitAfterResponse, "exit-after-response", false, "Exit TUI after first assistant response completes")
@@ -142,7 +144,7 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
 	}
 
 	// Start fake proxy if --fake is specified
-	fakeCleanup, err := setupFakeProxy(f.fakeResponses, &f.runConfig)
+	fakeCleanup, err := setupFakeProxy(f.fakeResponses, f.fakeStreamDelay, &f.runConfig)
 	if err != nil {
 		return err
 	}

e2e/proxy_test.go

@@ -29,6 +29,7 @@ func startRecordingAIProxy(t *testing.T) (*httptest.Server, *config.RuntimeConfi
 		recorder.ModeRecordOnce,
 		matcher,
 		fake.APIKeyHeaderUpdater,
+		nil,
 	)
 	require.NoError(t, err)
 

pkg/fake/proxy.go

@@ -3,6 +3,7 @@
 package fake
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"fmt"
@@ -21,10 +22,42 @@ import (
 	"gopkg.in/dnaeon/go-vcr.v4/pkg/recorder"
 )
 
+// ProxyOptions configures the fake proxy behavior.
+type ProxyOptions struct {
+	// SimulateStream adds delays between SSE chunks to simulate real streaming.
+	SimulateStream bool
+	// StreamChunkDelay is the delay between SSE chunks when SimulateStream is true.
+	// Defaults to 15ms if not set.
+	StreamChunkDelay time.Duration
+}
+
+// ProxyOption is a function that configures ProxyOptions.
+type ProxyOption func(*ProxyOptions)
+
+// WithSimulateStream enables simulated streaming with delays between chunks.
+func WithSimulateStream(enabled bool) ProxyOption {
+	return func(o *ProxyOptions) {
+		o.SimulateStream = enabled
+	}
+}
+
+// WithStreamChunkDelay sets the delay between SSE chunks.
+func WithStreamChunkDelay(d time.Duration) ProxyOption {
+	return func(o *ProxyOptions) {
+		o.StreamChunkDelay = d
+	}
+}
+
 // StartProxy starts an internal HTTP proxy that replays cassette responses.
 // It returns the proxy URL and a cleanup function that should be called when done.
-func StartProxy(cassettePath string) (string, func() error, error) {
-	return StartProxyWithOptions(cassettePath, recorder.ModeReplayOnly, nil, nil)
+func StartProxy(cassettePath string, opts ...ProxyOption) (string, func() error, error) {
+	options := &ProxyOptions{
+		StreamChunkDelay: 15 * time.Millisecond,
+	}
+	for _, opt := range opts {
+		opt(options)
+	}
+	return StartProxyWithOptions(cassettePath, recorder.ModeReplayOnly, nil, nil, options)
 }
 
 // StartRecordingProxy starts a proxy that records AI API interactions to a cassette file.
@@ -51,7 +84,7 @@ func StartStreamingRecordingProxy(
 	e := echo.New()
 	e.HideBanner = true
 	e.HidePort = true
-	e.Any("/*", Handle(streamRec, headerUpdater))
+	e.Any("/*", Handle(streamRec, headerUpdater, nil))
 
 	httpServer := httptest.NewServer(e)
 
@@ -102,17 +135,23 @@ func APIKeyHeaderUpdater(host string, req *http.Request) {
 // - mode: recorder mode (ModeReplayOnly, ModeRecordOnce, etc.)
 // - matcher: custom matcher function (nil uses DefaultMatcher)
 // - headerUpdater: optional function to update request headers (for recording with real API keys)
+// - options: proxy options for stream simulation, etc.
 func StartProxyWithOptions(
 	cassettePath string,
 	mode recorder.Mode,
 	matcher recorder.MatcherFunc,
 	headerUpdater func(host string, req *http.Request),
+	options *ProxyOptions,
 ) (string, func() error, error) {
 	hasMatcher := matcher != nil
 	if !hasMatcher {
 		matcher = DefaultMatcher(nil)
 	}
 
+	if options == nil {
+		options = &ProxyOptions{}
+	}
+
 	transport, err := recorder.New(cassettePath,
 		recorder.WithMode(mode),
 		recorder.WithMatcher(matcher),
@@ -127,7 +166,7 @@ func StartProxyWithOptions(
 	e := echo.New()
 	e.HideBanner = true
 	e.HidePort = true
-	e.Any("/*", Handle(transport, headerUpdater))
+	e.Any("/*", Handle(transport, headerUpdater, options))
 
 	httpServer := httptest.NewServer(e)
 
@@ -227,7 +266,11 @@ func TargetURLForHost(host string) func(req *http.Request) string {
 
 // Handle creates an echo handler that proxies requests through the VCR transport.
 // The headerUpdater is called with the host and request to update headers (e.g., for adding API keys).
-func Handle(transport http.RoundTripper, headerUpdater func(host string, req *http.Request)) echo.HandlerFunc {
+// The options parameter controls streaming simulation behavior.
+func Handle(transport http.RoundTripper, headerUpdater func(host string, req *http.Request), options *ProxyOptions) echo.HandlerFunc {
+	if options == nil {
+		options = &ProxyOptions{}
+	}
 	return func(c echo.Context) error {
 		ctx := c.Request().Context()
 
@@ -269,6 +312,9 @@ func Handle(transport http.RoundTripper, headerUpdater func(host string, req *ht
 		c.Response().WriteHeader(resp.StatusCode)
 
 		if IsStreamResponse(resp) {
+			if options.SimulateStream {
+				return SimulatedStreamCopy(c, resp, options.StreamChunkDelay)
+			}
 			return StreamCopy(c, resp)
 		}
 
@@ -277,6 +323,44 @@ func Handle(transport http.RoundTripper, headerUpdater func(host string, req *ht
 	}
 }
 
+// SimulatedStreamCopy copies a streaming SSE response to the client with artificial delays
+// between events to simulate real-time streaming behavior.
+func SimulatedStreamCopy(c echo.Context, resp *http.Response, chunkDelay time.Duration) error {
+	ctx := c.Request().Context()
+	writer := c.Response().Writer
+
+	scanner := bufio.NewScanner(resp.Body)
+	// SSE events can be large, increase buffer size
+	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
+
+	for scanner.Scan() {
+		select {
+		case <-ctx.Done():
+			slog.WarnContext(ctx, "client disconnected, stop streaming")
+			return nil
+		default:
+		}
+
+		line := scanner.Text()
+		// Write the line with newline
+		if _, err := writer.Write([]byte(line + "\n")); err != nil {
+			return err
+		}
+		c.Response().Flush()
+
+		// Add delay after data lines (SSE events start with "data:")
+		if strings.HasPrefix(line, "data:") {
+			select {
+			case <-ctx.Done():
+				return nil
+			case <-time.After(chunkDelay):
+			}
+		}
+	}
+
+	return scanner.Err()
+}
+
 // streamReadResult holds the result of a streaming read operation.
 type streamReadResult struct {
 	n   int64
@@ -330,6 +414,8 @@ func StreamCopy(c echo.Context, resp *http.Response) error {
 }
 
 // IsStreamResponse checks if the response should be streamed.
+// It checks Content-Type headers first, then falls back to peeking at the body
+// for SSE format (useful when headers are stripped in recorded cassettes).
 func IsStreamResponse(resp *http.Response) bool {
 	ct := strings.ToLower(resp.Header.Get("Content-Type"))
 	if strings.Contains(ct, "text/event-stream") {
@@ -341,7 +427,46 @@ func IsStreamResponse(resp *http.Response) bool {
 		return true
 	}
 
-	return strings.Contains(ct, "application/octet-stream") ||
+	if strings.Contains(ct, "application/octet-stream") ||
 		strings.Contains(ct, "application/x-ndjson") ||
-		strings.Contains(ct, "application/stream+json")
+		strings.Contains(ct, "application/stream+json") {
+		return true
+	}
+
+	// If no streaming headers detected, peek at the body to check for SSE format.
+	// This handles cassettes where headers were stripped during recording.
+	if resp.Body != nil {
+		// Read enough to detect SSE prefixes ("data:" or "event:")
+		peek := make([]byte, 6)
+		n, err := resp.Body.Read(peek)
+		if err == nil || n > 0 {
+			// Reconstruct the body with the peeked bytes prepended
+			resp.Body = &peekReader{peeked: peek[:n], rest: resp.Body}
+			// Check for SSE format markers
+			if bytes.HasPrefix(peek[:n], []byte("data:")) || bytes.HasPrefix(peek[:n], []byte("event:")) {
+				return true
+			}
+		}
+	}
+
+	return false
+}
+
+// peekReader wraps a reader with already-peeked bytes.
+type peekReader struct {
+	peeked []byte
+	rest   io.ReadCloser
+}
+
+func (p *peekReader) Read(b []byte) (int, error) {
+	if len(p.peeked) > 0 {
+		n := copy(b, p.peeked)
+		p.peeked = p.peeked[n:]
+		return n, nil
+	}
+	return p.rest.Read(b)
+}
+
+func (p *peekReader) Close() error {
+	return p.rest.Close()
 }