chore: fix disconnect bug and add agentcontainers test

DanielleMaywood · DanielleMaywood · commit 975ef8bfcaa7 · 2025-07-08T15:25:10.000Z
diff --git a/agent/agentcontainers/api.go b/agent/agentcontainers/api.go
@@ -560,14 +560,20 @@ func (api *API) watchContainers(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	ctx = api.ctx
+
 	go httpapi.Heartbeat(ctx, conn)
 	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
 
 	encoder := wsjson.NewEncoder[codersdk.WorkspaceAgentListContainersResponse](conn, websocket.MessageText)
 	defer encoder.Close(websocket.StatusNormalClosure)
 
-	updateCh := make(chan struct{})
-	defer close(updateCh)
+	updateCh := make(chan struct{}, 1)
+	defer func() {
+		api.mu.Lock()
+		close(updateCh)
+		api.mu.Unlock()
+	}()
 
 	api.mu.Lock()
 	api.updateChans = append(api.updateChans, updateCh)
@@ -644,7 +650,10 @@ func (api *API) updateContainers(ctx context.Context) error {
 
 	// Broadcast our updates
 	for _, ch := range api.updateChans {
-		ch <- struct{}{}
+		select {
+		case ch <- struct{}{}:
+		default:
+		}
 	}
 
 	api.logger.Debug(ctx, "containers updated successfully", slog.F("container_count", len(api.containers.Containers)), slog.F("warning_count", len(api.containers.Warnings)), slog.F("devcontainer_count", len(api.knownDevcontainers)))
diff --git a/agent/agentcontainers/api_test.go b/agent/agentcontainers/api_test.go
@@ -36,6 +36,7 @@ import (
 	"github.com/coder/coder/v2/pty"
 	"github.com/coder/coder/v2/testutil"
 	"github.com/coder/quartz"
+	"github.com/coder/websocket"
 )
 
 // fakeContainerCLI implements the agentcontainers.ContainerCLI interface for
@@ -441,6 +442,73 @@ func TestAPI(t *testing.T) {
 		logbuf.Reset()
 	})
 
+	t.Run("Watch", func(t *testing.T) {
+		t.Parallel()
+
+		fakeContainer1 := fakeContainer(t)
+		fakeContainer2 := fakeContainer(t)
+		fakeContainer3 := fakeContainer(t)
+
+		makeResponse := func(cts ...codersdk.WorkspaceAgentContainer) codersdk.WorkspaceAgentListContainersResponse {
+			return codersdk.WorkspaceAgentListContainersResponse{Containers: cts}
+		}
+
+		var (
+			ctx               = testutil.Context(t, testutil.WaitShort)
+			mClock            = quartz.NewMock(t)
+			updaterTickerTrap = mClock.Trap().TickerFunc("updaterLoop")
+			mCtrl             = gomock.NewController(t)
+			mLister           = acmock.NewMockContainerCLI(mCtrl)
+			logger            = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
+		)
+
+		mLister.EXPECT().List(gomock.Any()).Return(makeResponse(), nil)
+
+		api := agentcontainers.NewAPI(logger,
+			agentcontainers.WithClock(mClock),
+			agentcontainers.WithContainerCLI(mLister),
+			agentcontainers.WithContainerLabelIncludeFilter("this.label.does.not.exist.ignore.devcontainers", "true"),
+		)
+		api.Start()
+		defer api.Close()
+
+		srv := httptest.NewServer(api.Routes())
+		defer srv.Close()
+
+		updaterTickerTrap.MustWait(ctx).MustRelease(ctx)
+		defer updaterTickerTrap.Close()
+
+		client, _, err := websocket.Dial(ctx, srv.URL+"/watch", nil)
+		require.NoError(t, err)
+
+		for _, mockResponse := range []codersdk.WorkspaceAgentListContainersResponse{
+			makeResponse(),
+			makeResponse(fakeContainer1),
+			makeResponse(fakeContainer1, fakeContainer2),
+			makeResponse(fakeContainer1, fakeContainer2, fakeContainer3),
+			makeResponse(fakeContainer1, fakeContainer2),
+			makeResponse(fakeContainer1),
+			makeResponse(),
+		} {
+			mLister.EXPECT().List(gomock.Any()).Return(mockResponse, nil)
+
+			// Given: We allow the update loop to progress
+			_, aw := mClock.AdvanceNext()
+			aw.MustWait(ctx)
+
+			// When: We attempt to read a message from the socket.
+			mt, msg, err := client.Read(ctx)
+			require.NoError(t, err)
+			require.Equal(t, websocket.MessageText, mt)
+
+			// Then: We expect the receieved message matches the mocked response.
+			var got codersdk.WorkspaceAgentListContainersResponse
+			err = json.Unmarshal(msg, &got)
+			require.NoError(t, err)
+			require.Equal(t, mockResponse, got)
+		}
+	})
+
 	// List tests the API.getContainers method using a mock
 	// implementation. It specifically tests caching behavior.
 	t.Run("List", func(t *testing.T) {
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -865,6 +865,8 @@ func (api *API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r *http.Re
 		return
 	}
 
+	ctx = api.ctx
+
 	go httpapi.Heartbeat(ctx, conn)
 	defer conn.Close(websocket.StatusNormalClosure, "connection closed")
 

Original file line number	Diff line number	Diff line change
`@@ -865,6 +865,8 @@ func (api API) watchWorkspaceAgentContainers(rw http.ResponseWriter, r http.Re`
`865`	`865`	`return`
`866`	`866`	`}`
`867`	`867`
	`868`	`+ ctx = api.ctx`
	`869`	`+`
`868`	`870`	`go httpapi.Heartbeat(ctx, conn)`
`869`	`871`	`defer conn.Close(websocket.StatusNormalClosure, "connection closed")`
`870`	`872`