tighten Find resolution and SSE error handling

Author: dgageot

cmd/root/run_event_hooks.go

@@ -67,12 +67,46 @@ func withEventHooks(hooks []onEventHook) app.Opt {
 	}
 }
 
+// maxHookOutput caps the diagnostic output we keep for a failed on-event
+// hook. Large enough to be useful, small enough that a chatty or runaway
+// command can't push unbounded data into the agent's heap.
+const maxHookOutput = 4 * 1024
+
 func runEventHook(command string, payload []byte) {
 	shell, argsPrefix := shellpath.DetectShell()
+	// Hooks are detached from the app context on purpose: a hook still
+	// flushing the last event when the user exits the TUI should be allowed
+	// to finish. Each invocation receives a single event on stdin, processes
+	// it, and exits; the spawning goroutine ends with the subprocess.
 	cmd := exec.CommandContext(context.Background(), shell, append(argsPrefix, command)...)
 	cmd.Stdin = bytes.NewReader(payload)
-	out, err := cmd.CombinedOutput()
-	if err != nil {
-		slog.Warn("on-event hook failed", "command", command, "error", err, "output", strings.TrimSpace(string(out)))
+	var out boundedBuffer
+	cmd.Stdout = &out
+	cmd.Stderr = &out
+	if err := cmd.Run(); err != nil {
+		slog.Warn("on-event hook failed", "command", command, "error", err, "output", strings.TrimSpace(out.String()))
 	}
 }
+
+// boundedBuffer captures up to maxHookOutput bytes from a hook subprocess
+// and silently discards the rest. It implements only io.Writer so it can be
+// assigned to exec.Cmd's Stdout/Stderr without forcing exec to buffer the
+// full output internally.
+type boundedBuffer struct {
+	buf bytes.Buffer
+}
+
+func (b *boundedBuffer) Write(p []byte) (int, error) {
+	if remaining := maxHookOutput - b.buf.Len(); remaining > 0 {
+		if len(p) > remaining {
+			b.buf.Write(p[:remaining])
+		} else {
+			b.buf.Write(p)
+		}
+	}
+	return len(p), nil
+}
+
+func (b *boundedBuffer) String() string {
+	return b.buf.String()
+}

cmd/root/run_event_hooks_test.go

@@ -1,6 +1,8 @@
 package root
 
 import (
+	"bytes"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -27,3 +29,28 @@ func TestParseOnEventFlags_BadFormat(t *testing.T) {
 		assert.Error(t, err, "expected error for %q", s)
 	}
 }
+
+func TestBoundedBuffer_CapsAtMaxHookOutput(t *testing.T) {
+	var b boundedBuffer
+
+	n, err := b.Write(bytes.Repeat([]byte("a"), maxHookOutput-3))
+	require.NoError(t, err)
+	assert.Equal(t, maxHookOutput-3, n)
+
+	// A write that straddles the cap is fully accepted from the caller's
+	// perspective (so io.Copy doesn't error) but only the bytes up to the
+	// cap are retained.
+	n, err = b.Write([]byte("bbbbbb"))
+	require.NoError(t, err)
+	assert.Equal(t, 6, n)
+
+	// Subsequent writes past the cap are silently discarded.
+	n, err = b.Write([]byte("ccccc"))
+	require.NoError(t, err)
+	assert.Equal(t, 5, n)
+
+	got := b.String()
+	assert.Len(t, got, maxHookOutput)
+	assert.True(t, strings.HasPrefix(got, strings.Repeat("a", maxHookOutput-3)))
+	assert.True(t, strings.HasSuffix(got, "bbb"))
+}

cmd/root/sse.go

@@ -8,8 +8,14 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 )
 
+// maxErrorBodyBytes caps how much of an error response body we read into
+// memory. SSE error replies should be tiny; this just defends the client
+// against a misbehaving server that streams an unbounded error payload.
+const maxErrorBodyBytes = 4 * 1024
+
 // openEventStream connects to the SSE event stream of a session running on
 // addr and returns the response body. Callers are responsible for closing
 // the body. The body produces standard text/event-stream output with one
@@ -27,16 +33,16 @@ func openEventStream(ctx context.Context, addr, sessionID string) (io.ReadCloser
 		return nil, fmt.Errorf("connecting to %s: %w", url, err)
 	}
 	if resp.StatusCode >= 400 {
-		body, _ := io.ReadAll(resp.Body)
+		body, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBodyBytes))
 		_ = resp.Body.Close()
-		return nil, fmt.Errorf("server returned %s: %s", resp.Status, string(body))
+		return nil, fmt.Errorf("server returned %s: %s", resp.Status, strings.TrimSpace(string(body)))
 	}
 	return resp.Body, nil
 }
 
 // readEventStream reads SSE "data:" lines from r and invokes onEvent with
-// each raw JSON payload. The function returns when ctx is cancelled, the
-// stream ends, or onEvent returns an error.
+// each raw JSON payload. The function returns when ctx is cancelled (with
+// ctx.Err), the stream ends (nil), or onEvent returns an error.
 //
 // Payloads are passed through as json.RawMessage so callers can either
 // forward the bytes verbatim or re-decode them into a typed value without
@@ -46,7 +52,7 @@ func readEventStream(ctx context.Context, r io.Reader, onEvent func(json.RawMess
 	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
 	for scanner.Scan() {
 		if err := ctx.Err(); err != nil {
-			return nil
+			return err
 		}
 		after, ok := bytes.CutPrefix(scanner.Bytes(), []byte("data: "))
 		if !ok {

cmd/root/sse_test.go

@@ -1,6 +1,7 @@
 package root
 
 import (
+	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -67,3 +68,34 @@ func TestOpenEventStream_StreamsSuccess(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, []string{`{"hello":"world"}`}, got)
 }
+
+// TestOpenEventStream_CapsErrorBody guards against a misbehaving server
+// pushing an unbounded error body into client memory.
+func TestOpenEventStream_CapsErrorBody(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = w.Write([]byte(strings.Repeat("x", 10*maxErrorBodyBytes)))
+	}))
+	defer srv.Close()
+
+	_, err := openEventStream(t.Context(), srv.URL, "s1")
+	require.Error(t, err)
+	// The error message embeds at most maxErrorBodyBytes of the body.
+	assert.LessOrEqual(t, len(err.Error()), maxErrorBodyBytes+256)
+}
+
+// TestReadEventStream_ReturnsCtxErr verifies the helper surfaces ctx
+// cancellation so callers can distinguish it from a clean stream end.
+func TestReadEventStream_ReturnsCtxErr(t *testing.T) {
+	ctx, cancel := context.WithCancel(t.Context())
+	body := strings.NewReader("data: {}\n\ndata: {}\n\n")
+
+	calls := 0
+	err := readEventStream(ctx, body, func(json.RawMessage) error {
+		calls++
+		cancel()
+		return nil
+	})
+	require.ErrorIs(t, err, context.Canceled)
+	assert.Equal(t, 1, calls)
+}

cmd/root/watch.go

@@ -1,7 +1,9 @@
 package root
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 
 	"github.com/spf13/cobra"
@@ -56,8 +58,13 @@ func (f *watchFlags) run(cmd *cobra.Command, args []string) (commandErr error) {
 	out.Println("Watching", rec.Addr, "(session", rec.SessionID+")")
 
 	stdout := cmd.OutOrStdout()
-	return readEventStream(ctx, body, func(payload json.RawMessage) error {
+	err = readEventStream(ctx, body, func(payload json.RawMessage) error {
 		_, err := fmt.Fprintln(stdout, string(payload))
 		return err
 	})
+	// Ctrl+C is the normal way to stop watching; don't treat it as failure.
+	if errors.Is(err, context.Canceled) {
+		return nil
+	}
+	return err
 }

pkg/runregistry/registry.go

@@ -145,9 +145,10 @@ var ErrNoRun = errors.New("no live docker-agent run found; start one with: docke
 // An empty target returns the most recently started run. A numeric target is
 // matched by PID; a target starting with "http://" or "https://" is matched
 // against record addresses; anything else is matched as a (possibly partial)
-// session ID. Matching is exact for PID and addr, prefix-or-substring for
-// session ID; ambiguous matches return an error so callers don't act on the
-// wrong session.
+// session ID. PID and address matches are exact. Session-ID matching prefers
+// exact equality and only falls back to substring matching when no record
+// matches exactly; ambiguous substring matches return an error so callers
+// don't act on the wrong session.
 func Find(target string) (Record, error) {
 	target = strings.TrimSpace(target)
 	if target == "" {
@@ -175,7 +176,7 @@ func Find(target string) (Record, error) {
 				return r, nil
 			}
 		}
-		return Record{}, fmt.Errorf("no live run with pid %d", pid)
+		return Record{}, fmt.Errorf("no live run with pid %d: %w", pid, ErrNoRun)
 	}
 
 	if strings.HasPrefix(target, "http://") || strings.HasPrefix(target, "https://") {
@@ -185,18 +186,25 @@ func Find(target string) (Record, error) {
 				return r, nil
 			}
 		}
-		return Record{}, fmt.Errorf("no live run at %s", target)
+		return Record{}, fmt.Errorf("no live run at %s: %w", target, ErrNoRun)
 	}
 
+	// Prefer an exact session-id match: an unambiguous full id must always
+	// resolve, even when other ids contain it as a substring.
+	for _, r := range records {
+		if r.SessionID == target {
+			return r, nil
+		}
+	}
 	var matches []Record
 	for _, r := range records {
-		if r.SessionID == target || strings.HasPrefix(r.SessionID, target) || strings.Contains(r.SessionID, target) {
+		if strings.Contains(r.SessionID, target) {
 			matches = append(matches, r)
 		}
 	}
 	switch len(matches) {
 	case 0:
-		return Record{}, fmt.Errorf("no live run matches %q (pid, http URL, or session id)", target)
+		return Record{}, fmt.Errorf("no live run matches %q (pid, http URL, or session id): %w", target, ErrNoRun)
 	case 1:
 		return matches[0], nil
 	default:

pkg/runregistry/registry_test.go

@@ -135,18 +135,21 @@ func TestFind(t *testing.T) {
 		_, err := Find("999999999")
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "no live run with pid")
+		assert.ErrorIs(t, err, ErrNoRun)
 	})
 
 	t.Run("unknown addr errors", func(t *testing.T) {
 		_, err := Find("http://nope")
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "no live run at")
+		assert.ErrorIs(t, err, ErrNoRun)
 	})
 
 	t.Run("unknown session id errors", func(t *testing.T) {
 		_, err := Find("zzz")
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "no live run matches")
+		assert.ErrorIs(t, err, ErrNoRun)
 	})
 }
 
@@ -162,6 +165,21 @@ func TestFind_AmbiguousSessionID(t *testing.T) {
 	assert.Contains(t, err.Error(), "ambiguous")
 }
 
+// TestFind_ExactMatchBeatsSubstring guards against a regression where an
+// exact session-id match was reported as ambiguous because a longer id
+// contained it as a substring.
+func TestFind_ExactMatchBeatsSubstring(t *testing.T) {
+	withTempDataDir(t)
+
+	pid := os.Getpid()
+	writeRecord(t, "1.json", Record{PID: pid, Addr: "http://a", SessionID: "abc", StartedAt: time.Now()})
+	writeRecord(t, "2.json", Record{PID: pid, Addr: "http://b", SessionID: "abcd", StartedAt: time.Now()})
+
+	rec, err := Find("abc")
+	require.NoError(t, err)
+	assert.Equal(t, "abc", rec.SessionID)
+}
+
 func TestFind_EmptyRegistry(t *testing.T) {
 	withTempDataDir(t)