fix(cli/exp_agents_diff): degrade oversized watcher payloads gracefully

ibetitsmike · ibetitsmike · commit fcd8b7b85828 · 2026-04-21T04:43:19.000Z
Addresses Codex P2 on head 0d3f417: the local chat git watcher set a 4 MiB websocket read limit but agentgit caps each repository's UnifiedDiff at 3 MiB (maxTotalDiffSize), so a Changes message aggregating two or more maxed-out repos legitimately exceeds 4 MiB and triggers a coder/websocket close with StatusMessageTooBig. Because wsjson.Decoder swallowed the read error (log at Debug, close the channel), fetchLocalChatDiffContents surfaced a generic "git watch connection closed" error and the diff drawer crashed with a hard error in exactly the large-multi-repo case the local fallback was meant to improve: a regression from the pre-fallback behavior. Fix it with a narrowly scoped, two-layered degrade path: 1. Bump the client-side read limit from 4 MiB to 32 MiB (localChatDiffReadLimit), which covers ~10 maxed-out repos and makes realistic multi-repo worktrees just work without any fallback. 2. For pathological payloads beyond 32 MiB, treat the specific StatusMessageTooBig close as ignorable. fetchLocalChatDiffContents now reads from the websocket directly (bypassing wsjson.NewStream) so it can inspect websocket.CloseStatus(err), returning the narrow errLocalDiffMessageTooLarge sentinel on read-limit violations. shouldIgnoreLocalDiffFallbackError maps that sentinel onto the same remote-empty-diff fallback we already use for 404/403/400. Generic websocket close / decode / protocol errors continue to surface as hard errors so real regressions do not silently disappear behind the fallback. Tests: - IgnoresWatcherMessageTooBigCloses accepts the websocket, reads the refresh, then closes with StatusMessageTooBig. The client must degrade to the empty remote diff without error. - SurfacesUnexpectedWatcherCloseErrors closes with StatusInternalError instead and pins that the fallback is scoped to MessageTooBig only, so a future attempt to blanket-ignore every close reason immediately breaks the test. Pre-commit (gen/fmt/lint/build) passed. Full cli/ tests pass in 60.9s. > Mux is Mike's coding agent. This commit was made on Mike's behalf.
diff --git a/cli/exp_agents_diff.go b/cli/exp_agents_diff.go
@@ -2,6 +2,7 @@ package cli
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -13,12 +14,30 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/codersdk"
-	"github.com/coder/coder/v2/codersdk/wsjson"
 	"github.com/coder/websocket"
 )
 
 const localChatDiffWatchTimeout = 5 * time.Second
 
+// localChatDiffReadLimit bounds the size of the Changes message the
+// client is willing to receive from the chat git watcher. agentgit
+// caps each repository's UnifiedDiff at ~3 MiB (maxTotalDiffSize),
+// and a Changes payload can aggregate many repos plus metadata, so
+// 4 MiB is too tight for realistic multi-repo worktrees. 32 MiB
+// covers ~10 maxed-out repos; pathological payloads beyond that still
+// fall back to the remote empty diff via
+// errLocalDiffMessageTooLarge / shouldIgnoreLocalDiffFallbackError.
+const localChatDiffReadLimit = 32 << 20 // 32 MiB
+
+// errLocalDiffMessageTooLarge is returned when the chat git watcher
+// produces a Changes payload that exceeds localChatDiffReadLimit.
+// shouldIgnoreLocalDiffFallbackError treats it as ignorable so the
+// TUI degrades to the remote empty diff rather than surfacing a hard
+// error. Generic websocket close / decode errors are intentionally
+// still surfaced so real protocol regressions do not silently
+// disappear behind the fallback.
+var errLocalDiffMessageTooLarge = xerrors.New("chat git watcher payload exceeded client read limit")
+
 func fetchChatDiffContents(
 	ctx context.Context,
 	client *codersdk.ExperimentalClient,
@@ -74,6 +93,17 @@ func fetchChatDiffContents(
 // exactly one contributing repository. The caller uses singleRepo to
 // decide whether it is safe to backfill remote-only metadata onto the
 // local diff. All error paths return singleRepo=false.
+//
+// This intentionally bypasses wsjson.NewStream and reads the websocket
+// directly so we can inspect the close status: an oversized Changes
+// payload must degrade to the remote empty diff via
+// errLocalDiffMessageTooLarge + shouldIgnoreLocalDiffFallbackError,
+// but wsjson.Decoder swallows the read error (logs at debug) and
+// closes the channel, which would collapse that specific case into
+// the same generic "connection closed" bucket as server crashes or
+// decode failures. Reading directly lets us narrowly fall back only
+// for read-limit violations while still surfacing real protocol
+// regressions.
 func fetchLocalChatDiffContents(
 	parentCtx context.Context,
 	client *codersdk.ExperimentalClient,
@@ -89,41 +119,57 @@ func fetchLocalChatDiffContents(
 	defer func() {
 		_ = conn.Close(websocket.StatusNormalClosure, "")
 	}()
-	conn.SetReadLimit(1 << 22) // 4MiB
+	conn.SetReadLimit(localChatDiffReadLimit)
 
-	stream := wsjson.NewStream[
-		codersdk.WorkspaceAgentGitServerMessage,
-		codersdk.WorkspaceAgentGitClientMessage,
-	](conn, websocket.MessageText, websocket.MessageText, client.Logger())
-	if err := stream.Send(codersdk.WorkspaceAgentGitClientMessage{
+	refreshPayload, err := json.Marshal(codersdk.WorkspaceAgentGitClientMessage{
 		Type: codersdk.WorkspaceAgentGitClientMessageTypeRefresh,
-	}); err != nil {
+	})
+	if err != nil {
+		return codersdk.ChatDiffContents{}, false, xerrors.Errorf("marshal git refresh: %w", err)
+	}
+	if err := conn.Write(ctx, websocket.MessageText, refreshPayload); err != nil {
 		return codersdk.ChatDiffContents{}, false, xerrors.Errorf("request git refresh: %w", err)
 	}
 
-	messages := stream.Chan()
 	for {
-		select {
-		case <-ctx.Done():
-			return codersdk.ChatDiffContents{}, false, xerrors.Errorf("watch chat git: %w", ctx.Err())
-		case msg, ok := <-messages:
-			if !ok {
-				if ctx.Err() != nil {
-					return codersdk.ChatDiffContents{}, false, xerrors.Errorf("watch chat git: %w", ctx.Err())
-				}
-				return codersdk.ChatDiffContents{}, false, xerrors.New("git watch connection closed")
+		msgType, payload, err := conn.Read(ctx)
+		if err != nil {
+			// Context expiration gets its own wrapping so it threads
+			// cleanly through shouldIgnoreLocalDiffFallbackError's
+			// context.DeadlineExceeded case.
+			if ctxErr := ctx.Err(); ctxErr != nil {
+				return codersdk.ChatDiffContents{}, false, xerrors.Errorf("watch chat git: %w", ctxErr)
+			}
+			// A Changes payload that exceeds localChatDiffReadLimit
+			// causes coder/websocket to close the connection with
+			// StatusMessageTooBig. Surface that as the narrow
+			// sentinel so the caller can fall back to the remote
+			// empty diff instead of surfacing a hard error.
+			if websocket.CloseStatus(err) == websocket.StatusMessageTooBig {
+				return codersdk.ChatDiffContents{}, false, errLocalDiffMessageTooLarge
 			}
-			switch msg.Type {
-			case codersdk.WorkspaceAgentGitServerMessageTypeError:
-				message := strings.TrimSpace(msg.Message)
-				if message == "" {
-					message = "git watch returned an unknown error"
-				}
-				return codersdk.ChatDiffContents{}, false, xerrors.New(message)
-			case codersdk.WorkspaceAgentGitServerMessageTypeChanges:
-				diff, singleRepo := buildLocalChatDiffContents(chatID, msg.Repositories)
-				return diff, singleRepo, nil
+			return codersdk.ChatDiffContents{}, false, xerrors.Errorf("read git watch: %w", err)
+		}
+		// Ignore unexpected frame types instead of erroring; the
+		// watcher only emits text frames today and a future binary
+		// heartbeat should not break the overlay.
+		if msgType != websocket.MessageText {
+			continue
+		}
+		var msg codersdk.WorkspaceAgentGitServerMessage
+		if err := json.Unmarshal(payload, &msg); err != nil {
+			return codersdk.ChatDiffContents{}, false, xerrors.Errorf("decode git watch message: %w", err)
+		}
+		switch msg.Type {
+		case codersdk.WorkspaceAgentGitServerMessageTypeError:
+			message := strings.TrimSpace(msg.Message)
+			if message == "" {
+				message = "git watch returned an unknown error"
 			}
+			return codersdk.ChatDiffContents{}, false, xerrors.New(message)
+		case codersdk.WorkspaceAgentGitServerMessageTypeChanges:
+			diff, singleRepo := buildLocalChatDiffContents(chatID, msg.Repositories)
+			return diff, singleRepo, nil
 		}
 	}
 }
@@ -220,6 +266,15 @@ func shouldIgnoreLocalDiffFallbackError(err error) bool {
 	if errors.Is(err, context.DeadlineExceeded) {
 		return true
 	}
+	// An oversized Changes payload is a best-effort degradation
+	// point: the remote /diff endpoint already returns the empty
+	// placeholder in this case, so fall back to it instead of
+	// surfacing a hard error. Scoped narrowly to the explicit
+	// StatusMessageTooBig sentinel so generic websocket close /
+	// decode errors still reach the user.
+	if errors.Is(err, errLocalDiffMessageTooLarge) {
+		return true
+	}
 
 	sdkErr, ok := codersdk.AsError(err)
 	if !ok {
diff --git a/cli/exp_agents_diff_test.go b/cli/exp_agents_diff_test.go
@@ -406,6 +406,85 @@ func TestFetchChatDiffContents(t *testing.T) {
 		require.Equal(t, remotePR, *diff.PullRequestURL)
 	})
 
+	t.Run("IgnoresWatcherMessageTooBigCloses", func(t *testing.T) {
+		t.Parallel()
+
+		// agentgit caps each repository's UnifiedDiff at ~3 MiB and a
+		// Changes payload aggregates every repo plus metadata, so a
+		// realistic multi-repo workspace can legitimately produce a
+		// payload that exceeds the client's websocket read limit.
+		// When that happens coder/websocket closes the connection
+		// with StatusMessageTooBig. fetchChatDiffContents must map
+		// that specific close status onto errLocalDiffMessageTooLarge
+		// and fall back to the remote empty diff rather than
+		// surfacing a hard error to the TUI. Without this subtest,
+		// removing the StatusMessageTooBig branch in
+		// fetchLocalChatDiffContents or the errLocalDiffMessageTooLarge
+		// branch in shouldIgnoreLocalDiffFallbackError would
+		// silently regress the large-multi-repo case this feature is
+		// meant to improve.
+		ctx := t.Context()
+		chatID := uuid.New()
+		path := fmt.Sprintf("/api/experimental/chats/%s", chatID)
+		client := newTestExperimentalClient(t, http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			switch r.URL.Path {
+			case path + "/diff":
+				rw.Header().Set("Content-Type", "application/json")
+				require.NoError(t, json.NewEncoder(rw).Encode(codersdk.ChatDiffContents{ChatID: chatID}))
+			case path + "/stream/git":
+				conn, err := websocket.Accept(rw, r, nil)
+				require.NoError(t, err)
+				// Drain the refresh before closing so the client
+				// surfaces the close status from its next Read, not
+				// an unrelated write error.
+				_, _, err = conn.Read(ctx)
+				require.NoError(t, err)
+				require.NoError(t, conn.Close(websocket.StatusMessageTooBig, "too big"))
+			default:
+				http.NotFound(rw, r)
+			}
+		}))
+
+		diff, err := fetchChatDiffContents(ctx, client, chatID)
+		require.NoError(t, err)
+		require.Equal(t, chatID, diff.ChatID)
+		require.Empty(t, diff.Diff)
+	})
+
+	t.Run("SurfacesUnexpectedWatcherCloseErrors", func(t *testing.T) {
+		t.Parallel()
+
+		// The StatusMessageTooBig fallback is intentionally narrow:
+		// a generic websocket close (for example the server
+		// crashing and closing with StatusInternalError) should
+		// surface as an error rather than silently degrading,
+		// because that would hide real protocol regressions behind
+		// the best-effort fallback. This subtest pins that
+		// distinction so a future attempt to blanket-ignore every
+		// close reason immediately breaks the test.
+		ctx := t.Context()
+		chatID := uuid.New()
+		path := fmt.Sprintf("/api/experimental/chats/%s", chatID)
+		client := newTestExperimentalClient(t, http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+			switch r.URL.Path {
+			case path + "/diff":
+				rw.Header().Set("Content-Type", "application/json")
+				require.NoError(t, json.NewEncoder(rw).Encode(codersdk.ChatDiffContents{ChatID: chatID}))
+			case path + "/stream/git":
+				conn, err := websocket.Accept(rw, r, nil)
+				require.NoError(t, err)
+				_, _, err = conn.Read(ctx)
+				require.NoError(t, err)
+				require.NoError(t, conn.Close(websocket.StatusInternalError, "boom"))
+			default:
+				http.NotFound(rw, r)
+			}
+		}))
+
+		_, err := fetchChatDiffContents(ctx, client, chatID)
+		require.Error(t, err)
+	})
+
 	t.Run("ReturnsRemoteDiffWithoutDialingWatcher", func(t *testing.T) {
 		t.Parallel()
 
