Merge pull request #2775 from dgageot/fix-tui-control-api-bugs

fix: two TUI control-plane bugs (SSE cancel, IPv6 listen)

Author: dgageot

pkg/runtime/client.go

@@ -476,19 +476,24 @@ func (c *Client) GetAgentToolCount(ctx context.Context, agentFilename, agentName
 }
 
 // StreamSessionEvents streams events for a session as they occur via Server-Sent Events.
+// The returned channel is closed when ctx is cancelled, the stream's max
+// duration is reached, or the server closes the connection.
 func (c *Client) StreamSessionEvents(ctx context.Context, sessionID string) (<-chan Event, error) {
 	endpoint := fmt.Sprintf("/api/sessions/%s/events", sessionID)
 
 	u := *c.baseURL
 	u.Path = path.Join(u.Path, endpoint)
 
-	// Use long timeout for streaming
+	// Bound the maximum lifetime of a single SSE connection. The cancel
+	// must be tied to the goroutine consuming the stream, not to this
+	// function's return: cancelling streamCtx kills the in-flight HTTP
+	// request, which would turn the stream into a one-shot read.
 	timeout := c.timeoutFor("streaming")
 	streamCtx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
 
 	req, err := http.NewRequestWithContext(streamCtx, http.MethodGet, u.String(), http.NoBody)
 	if err != nil {
+		cancel()
 		return nil, fmt.Errorf("creating request: %w", err)
 	}
 
@@ -501,10 +506,12 @@ func (c *Client) StreamSessionEvents(ctx context.Context, sessionID string) (<-c
 
 	resp, err := c.httpClient.Do(req) //nolint:bodyclose // body is closed in the goroutine below
 	if err != nil {
+		cancel()
 		return nil, fmt.Errorf("performing request: %w", err)
 	}
 
 	if resp.StatusCode >= 400 {
+		defer cancel()
 		defer resp.Body.Close()
 		respBody, err := io.ReadAll(resp.Body)
 		if err != nil {
@@ -521,6 +528,7 @@ func (c *Client) StreamSessionEvents(ctx context.Context, sessionID string) (<-c
 	eventChan := make(chan Event, defaultEventChannelCapacity)
 
 	go func() {
+		defer cancel()
 		defer close(eventChan)
 		defer resp.Body.Close()
 

pkg/runtime/client_test.go

@@ -0,0 +1,114 @@
+package runtime
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestClient_StreamSessionEvents_DeliversMultipleEvents verifies that the
+// SSE stream stays open across multiple events instead of being torn down
+// when StreamSessionEvents returns. This is a regression test for a bug
+// where a deferred cancel() on the streaming context killed the in-flight
+// HTTP request as soon as the function returned, turning the stream into
+// a one-shot read.
+func TestClient_StreamSessionEvents_DeliversMultipleEvents(t *testing.T) {
+	t.Parallel()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		flusher, ok := w.(http.Flusher)
+		if !ok {
+			t.Errorf("ResponseWriter must support flushing")
+			return
+		}
+
+		for i := 1; i <= 3; i++ {
+			fmt.Fprintf(w, "data: {\"type\":\"session_title\",\"session_id\":\"s\",\"title\":\"t%d\"}\n\n", i)
+			flusher.Flush()
+			time.Sleep(20 * time.Millisecond)
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	c, err := NewClient(srv.URL)
+	require.NoError(t, err)
+
+	ch, err := c.StreamSessionEvents(t.Context(), "s")
+	require.NoError(t, err)
+
+	var titles []string
+	for ev := range ch {
+		titleEv, ok := ev.(*SessionTitleEvent)
+		if !ok {
+			continue
+		}
+		titles = append(titles, titleEv.Title)
+	}
+
+	assert.Equal(t, []string{"t1", "t2", "t3"}, titles)
+}
+
+// TestClient_StreamSessionEvents_StopsWhenContextCancelled verifies that
+// cancelling the caller's context tears down the stream and closes the
+// returned channel.
+func TestClient_StreamSessionEvents_StopsWhenContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(http.StatusOK)
+		flusher, _ := w.(http.Flusher)
+
+		ticker := time.NewTicker(20 * time.Millisecond)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-r.Context().Done():
+				return
+			case <-ticker.C:
+				fmt.Fprint(w, "data: {\"type\":\"session_title\",\"session_id\":\"s\",\"title\":\"x\"}\n\n")
+				flusher.Flush()
+			}
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	c, err := NewClient(srv.URL)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(t.Context())
+	t.Cleanup(cancel)
+
+	ch, err := c.StreamSessionEvents(ctx, "s")
+	require.NoError(t, err)
+
+	// Drain at least one event to confirm the stream is live.
+	select {
+	case <-ch:
+	case <-time.After(2 * time.Second):
+		t.Fatal("no events received before cancel")
+	}
+
+	cancel()
+
+	// Channel must close in a bounded time after cancel.
+	deadline := time.After(2 * time.Second)
+	for {
+		select {
+		case _, ok := <-ch:
+			if !ok {
+				return
+			}
+		case <-deadline:
+			t.Fatal("channel was not closed after context cancel")
+		}
+	}
+}

pkg/server/listen.go

@@ -62,5 +62,5 @@ func listenUnix(ctx context.Context, path string) (net.Listener, error) {
 
 func listenTCP(ctx context.Context, addr string) (net.Listener, error) {
 	var lc net.ListenConfig
-	return lc.Listen(ctx, "tcp4", addr)
+	return lc.Listen(ctx, "tcp", addr)
 }

pkg/server/listen_test.go

@@ -62,3 +62,42 @@ func TestListen_FD_InvalidDescriptor(t *testing.T) {
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "file descriptor 999999")
 }
+
+// TestListen_TCP_IPv4 verifies that the default TCP listener binds to an
+// IPv4 loopback address. Regression test for the listener being hard-coded
+// to "tcp4" without that being the documented intent.
+func TestListen_TCP_IPv4(t *testing.T) {
+	t.Parallel()
+
+	ln, err := Listen(t.Context(), "127.0.0.1:0")
+	require.NoError(t, err)
+	defer ln.Close()
+
+	tcpAddr, ok := ln.Addr().(*net.TCPAddr)
+	require.True(t, ok)
+	assert.NotNil(t, tcpAddr.IP.To4())
+}
+
+// TestListen_TCP_IPv6 verifies that an IPv6 loopback bind succeeds. The
+// listener used to force "tcp4" which made this fail with
+// "address ::1: non-IPv4 address" on dual-stack hosts.
+func TestListen_TCP_IPv6(t *testing.T) {
+	t.Parallel()
+
+	// Probe whether the host actually has IPv6 before asserting.
+	var probeLC net.ListenConfig
+	probe, probeErr := probeLC.Listen(t.Context(), "tcp6", "[::1]:0")
+	if probeErr != nil {
+		t.Skipf("host does not support IPv6 loopback: %v", probeErr)
+	}
+	_ = probe.Close()
+
+	ln, err := Listen(t.Context(), "[::1]:0")
+	require.NoError(t, err)
+	defer ln.Close()
+
+	tcpAddr, ok := ln.Addr().(*net.TCPAddr)
+	require.True(t, ok)
+	assert.True(t, tcpAddr.IP.IsLoopback())
+	assert.Nil(t, tcpAddr.IP.To4(), "expected an IPv6-only address")
+}