chore: add agent-connection-watch for workspaces (#24507)

<!--

If you have used AI to produce some or all of this PR, please ensure you have read our [AI Contribution guidelines](https://coder.com/docs/about/contributing/AI_CONTRIBUTING) before submitting.

-->

relates to GRU-18  
  
Adds basic implementation for Workspace Agent Connection Watch and tests.  
  
Missing are handling of logs.

coderd/workspaceconnwatcher/watcher.go

@@ -0,0 +1,333 @@
+package workspaceconnwatcher
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"net/http"
+	"sync"
+
+	"github.com/google/uuid"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/coderd/httpapi"
+	"github.com/coder/coder/v2/coderd/httpmw"
+	"github.com/coder/coder/v2/coderd/wspubsub"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+	"github.com/coder/coder/v2/codersdk/wsjson"
+	"github.com/coder/websocket"
+)
+
+type Watcher struct {
+	logger slog.Logger
+	sub    pubsub.Subscriber
+	db     database.Store
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	mu     sync.Mutex
+	wg     sync.WaitGroup
+	closed bool
+}
+
+type event struct {
+	sync    bool
+	wsEvent *wspubsub.WorkspaceEvent
+}
+
+func New(ctx context.Context, logger slog.Logger, sub pubsub.Subscriber, db database.Store) *Watcher {
+	ctx, cancel := context.WithCancel(ctx)
+	w := &Watcher{
+		logger: logger.Named("wsconnwatcher"),
+		ctx:    ctx,
+		cancel: cancel,
+		sub:    sub,
+		db:     db,
+	}
+	go func() {
+		<-ctx.Done()
+		w.Close()
+	}()
+	return w
+}
+
+// @Summary Workspace Agent Connection Watch
+// @ID workspace-agent-connection-watch
+// @Security CoderSessionToken
+// @Produce json
+// @Tags Workspaces
+// @Param workspace path string true "Workspace ID" format(uuid)
+// @Success 101 {object} workspacesdk.ConnectionWatchEvent
+// @Router /api/v2/workspaces/{workspace}/agent-connection-watch [get]
+func (w *Watcher) WorkspaceAgentConnectionWatch(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+	workspace := httpmw.WorkspaceParam(r)
+	agentName := r.URL.Query().Get("agent_name")
+
+	filteredEvents := make(chan event, 1)
+	filteredEvents <- event{sync: true} // init sync
+	cancelWorkspaceSubscribe, err := w.sub.SubscribeWithErr(wspubsub.WorkspaceEventChannel(workspace.OwnerID),
+		wspubsub.HandleWorkspaceEvent(
+			func(ctx context.Context, payload wspubsub.WorkspaceEvent, err error) {
+				if err != nil {
+					// subscription error, resync
+					select {
+					case filteredEvents <- event{sync: true}:
+					case <-ctx.Done():
+					}
+					return
+				}
+				if payload.WorkspaceID != workspace.ID {
+					return
+				}
+				select {
+				case filteredEvents <- event{wsEvent: &payload}:
+				case <-ctx.Done():
+				}
+			}))
+	if err != nil {
+		w.logger.Error(ctx, "failed to subscribe to workspace events",
+			slog.Error(err), slog.F("owner_id", workspace.OwnerID))
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error setting up workspace event subscription",
+			// Don't include the error in case it leaks infra details about the pubsub
+		})
+		return
+	}
+	defer cancelWorkspaceSubscribe()
+
+	closed := false
+	w.mu.Lock()
+	closed = w.closed
+	if !closed {
+		w.wg.Add(1)
+	}
+	w.mu.Unlock()
+	if closed {
+		w.logger.Debug(ctx, "server is closed, writing error")
+		httpapi.Write(ctx, rw, http.StatusServiceUnavailable, codersdk.Response{
+			Message: "Server instance is shutting down",
+		})
+		return
+	}
+	defer w.wg.Done()
+
+	conn, err := websocket.Accept(rw, r, nil)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Failed to accept WebSocket.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	// CloseRead starts a goroutine to read and discard messages from the client,
+	// including Pong messages sent in response to our Ping heartbeats.
+	_ = conn.CloseRead(ctx)
+
+	ctx, cancel := context.WithCancel(ctx)
+	go httpapi.HeartbeatClose(ctx, w.logger, cancel, conn)
+	defer cancel()
+
+	u := &updater{
+		db:          w.db,
+		watcherCtx:  w.ctx,
+		connCtx:     ctx,
+		conn:        conn,
+		workspaceID: workspace.ID,
+		events:      filteredEvents,
+		agentName:   agentName,
+	}
+	u.run()
+}
+
+func (w *Watcher) Close() {
+	w.mu.Lock()
+	w.closed = true
+	w.mu.Unlock()
+
+	w.cancel()
+	w.wg.Wait()
+}
+
+type updater struct {
+	db          database.Store
+	watcherCtx  context.Context
+	connCtx     context.Context
+	conn        *websocket.Conn
+	enc         *wsjson.Encoder[workspacesdk.ConnectionWatchEvent]
+	workspaceID uuid.UUID
+	events      <-chan event
+	agentName   string
+
+	lastBuild database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow
+}
+
+func (u *updater) run() {
+	u.enc = wsjson.NewEncoder[workspacesdk.ConnectionWatchEvent](u.conn, websocket.MessageText)
+	defer func() {
+		// this is a no-op if we have already closed for some other reason.
+		_ = u.enc.Close(websocket.StatusNormalClosure)
+	}()
+
+	for {
+		select {
+		case <-u.watcherCtx.Done():
+			u.errorThenClose(workspacesdk.WatchError{
+				Code:      workspacesdk.WatchErrorServerShutdown,
+				Retryable: true,
+				Message:   "server is shutting down",
+			})
+			return
+		case <-u.connCtx.Done():
+			return
+		case e := <-u.events:
+			if e.sync {
+				// zero this out so we'll send a full update
+				u.lastBuild = database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{}
+				if !u.buildUpdate() {
+					return
+				}
+			}
+			if e.wsEvent != nil {
+				switch e.wsEvent.Kind {
+				case wspubsub.WorkspaceEventKindStateChange:
+					if !u.buildUpdate() {
+						return
+					}
+				case wspubsub.WorkspaceEventKindAgentLifecycleUpdate:
+					if !u.maybeSendAgentUpdate() {
+						return
+					}
+				}
+			}
+		}
+	}
+}
+
+func (u *updater) buildUpdate() bool {
+	build, err := u.db.GetLatestWorkspaceBuildWithStatusByWorkspaceID(u.connCtx, u.workspaceID)
+	if err != nil {
+		retryable := true
+		details := err.Error()
+		if errors.Is(err, sql.ErrNoRows) {
+			// There is no build (unlikely), or the workspace was deleted. In both cases, retrying won't help.
+			retryable = false
+		}
+		if dbauthz.IsNotAuthorizedError(err) {
+			retryable = false
+			details = "unauthorized" // security: don't leak internal authz details
+		}
+		u.errorThenClose(workspacesdk.WatchError{
+			Code:      workspacesdk.WatchErrorDatabase,
+			Retryable: retryable,
+			Message:   "failed to fetch latest workspace build",
+			Details:   details,
+		})
+		return false
+	}
+
+	if build.BuildNumber != u.lastBuild.BuildNumber ||
+		build.JobStatus != u.lastBuild.JobStatus ||
+		build.Transition != u.lastBuild.Transition {
+		u.lastBuild = build
+		err = u.enc.Encode(workspacesdk.ConnectionWatchEvent{BuildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransition(build.Transition),
+			JobStatus:  codersdk.ProvisionerJobStatus(build.JobStatus),
+		}})
+		if err != nil {
+			// probably this is just that the connection is closed, but in case there is some actual JSON serialization
+			// error, send a close frame.
+			_ = u.conn.Close(websocket.StatusInternalError, "failed to encode build update")
+			return false
+		}
+		return u.maybeSendAgentUpdate()
+	}
+	return true
+}
+
+func (u *updater) maybeSendAgentUpdate() (ok bool) {
+	if u.lastBuild.Transition != database.WorkspaceTransitionStart ||
+		u.lastBuild.JobStatus != database.ProvisionerJobStatusSucceeded {
+		// only send agent updates for successfully started workspaces
+		return true
+	}
+
+	agents, err := u.db.GetWorkspaceAgentsByWorkspaceAndBuildNumber(u.connCtx,
+		database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
+			WorkspaceID: u.workspaceID,
+			BuildNumber: u.lastBuild.BuildNumber,
+		})
+	if err != nil && !errors.Is(err, sql.ErrNoRows) {
+		details := err.Error()
+		retryable := true
+		if dbauthz.IsNotAuthorizedError(err) {
+			retryable = false
+			details = "unauthorized"
+		}
+		u.errorThenClose(workspacesdk.WatchError{
+			Code:      workspacesdk.WatchErrorDatabase,
+			Retryable: retryable,
+			Message:   "failed to fetch workspace agents",
+			Details:   details,
+		})
+		return false
+	}
+	if len(agents) == 0 {
+		u.errorThenClose(workspacesdk.WatchError{
+			Code:      workspacesdk.WatchErrorNoAgents,
+			Retryable: false,
+			Message:   "no agents found for workspace",
+		})
+		return false
+	}
+	if len(agents) > 1 && u.agentName == "" {
+		u.errorThenClose(workspacesdk.WatchError{
+			Code:      workspacesdk.WatchErrorTooManyAgents,
+			Retryable: false,
+			Message:   "more than one agent on workspace and target not specified",
+		})
+		return false
+	}
+	var agent database.WorkspaceAgent
+	if u.agentName == "" {
+		agent = agents[0]
+	} else {
+		for _, a := range agents {
+			if a.Name == u.agentName {
+				agent = a
+				break
+			}
+		}
+		if agent.ID == uuid.Nil {
+			u.errorThenClose(workspacesdk.WatchError{
+				Code:      workspacesdk.WatchErrorNameNotFound,
+				Retryable: false,
+				Message:   "target agent not found by name",
+			})
+			return false
+		}
+	}
+
+	err = u.enc.Encode(workspacesdk.ConnectionWatchEvent{AgentUpdate: &workspacesdk.AgentUpdate{
+		Lifecycle: codersdk.WorkspaceAgentLifecycle(agent.LifecycleState),
+		ID:        agent.ID,
+	}})
+	if err != nil {
+		// probably this is just that the connection is closed, but in case there is some actual JSON serialization
+		// error, send a close frame.
+		_ = u.conn.Close(websocket.StatusInternalError, "failed to encode agent update")
+		return false
+	}
+	return true
+}
+
+func (u *updater) errorThenClose(err workspacesdk.WatchError) {
+	_ = u.enc.Encode(workspacesdk.ConnectionWatchEvent{Error: &err})
+	// ignore encoding errors above because in any case, we are going to close the connection.
+	_ = u.conn.Close(websocket.StatusNormalClosure, "error")
+}

coderd/workspaceconnwatcher/watcher_test.go

@@ -0,0 +1,474 @@
+package workspaceconnwatcher_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbmock"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/coderd/httpmw"
+	"github.com/coder/coder/v2/coderd/workspaceconnwatcher"
+	"github.com/coder/coder/v2/coderd/wspubsub"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+	"github.com/coder/coder/v2/codersdk/wsjson"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/websocket"
+)
+
+var (
+	workspaceID = uuid.UUID{1}
+	userID      = uuid.UUID{2}
+	orgID       = uuid.UUID{3}
+	agentID     = uuid.UUID{4}
+)
+
+type harness struct {
+	db      *dbmock.MockStore
+	watcher *workspaceconnwatcher.Watcher
+	pub     pubsub.Publisher
+	logger  slog.Logger
+
+	// Initialized, but overridable before Dial()
+	workspace     database.Workspace
+	userID, orgID uuid.UUID
+}
+
+func newHarness(ctx context.Context, t *testing.T, logger slog.Logger) *harness {
+	h := &harness{
+		workspace: database.Workspace{
+			ID:             workspaceID,
+			OrganizationID: orgID,
+			OwnerID:        userID,
+		},
+		orgID:  orgID,
+		userID: userID,
+		logger: logger,
+	}
+	ps := pubsub.NewInMemory()
+	h.pub = ps
+
+	var authzDB database.Store
+	_, h.db, authzDB, _ = coderdtest.MockedDatabaseWithAuthz(t, logger)
+	h.watcher = workspaceconnwatcher.New(ctx, logger.Named("watcher"), ps, authzDB)
+	t.Cleanup(h.watcher.Close)
+	return h
+}
+
+func (h *harness) Dial(ctx context.Context, url string) (*wsjson.Decoder[workspacesdk.ConnectionWatchEvent], error) {
+	rt := testutil.InMemWebsocketRoundTripper{
+		Handler: http.HandlerFunc(h.watcher.WorkspaceAgentConnectionWatch),
+		CtxMutator: func(ctx context.Context) context.Context {
+			ctx = httpmw.WithWorkspaceParam(ctx, h.workspace)
+			ctx = dbauthz.As(ctx, coderdtest.MemberSubject(userID, orgID))
+			return ctx
+		},
+		Logger: h.logger.Named("roundtripper"),
+	}
+	// nolint: bodyclose
+	clientSock, resp, err := websocket.Dial(ctx, url, &websocket.DialOptions{
+		HTTPClient: &http.Client{Transport: rt},
+	})
+	if err != nil {
+		if resp.StatusCode != http.StatusSwitchingProtocols {
+			return nil, codersdk.ReadBodyAsError(resp)
+		}
+		return nil, err
+	}
+
+	dec := wsjson.NewDecoder[workspacesdk.ConnectionWatchEvent](
+		clientSock, websocket.MessageText, h.logger.Named("decoder"))
+	return dec, nil
+}
+
+func TestWatcher_Agents(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name                   string
+		agents                 []database.WorkspaceAgent
+		agentDBError           error
+		url                    string
+		expectedAgentUpdate    *workspacesdk.AgentUpdate
+		expectedErrorCode      workspacesdk.WatchErrorCode
+		expectedErrorRetryable bool
+	}{
+		{
+			name: "noNameSingleAgent",
+			agents: []database.WorkspaceAgent{
+				{
+					Name:           "test",
+					ID:             agentID,
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+			},
+			url: "wss://local.test/",
+			expectedAgentUpdate: &workspacesdk.AgentUpdate{
+				Lifecycle: codersdk.WorkspaceAgentLifecycleCreated,
+				ID:        agentID,
+			},
+		},
+		{
+			name: "noNameMultiAgent",
+			agents: []database.WorkspaceAgent{
+				{
+					Name:           "agent0",
+					ID:             agentID,
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+				{
+					Name:           "agent1",
+					ID:             uuid.UUID{77},
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+			},
+			url:                    "wss://local.test/",
+			expectedErrorCode:      workspacesdk.WatchErrorTooManyAgents,
+			expectedErrorRetryable: false,
+		},
+		{
+			name: "namedAgentMultiAgent",
+			agents: []database.WorkspaceAgent{
+				{
+					Name:           "agent0",
+					ID:             agentID,
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+				{
+					Name:           "agent1",
+					ID:             uuid.UUID{77},
+					LifecycleState: database.WorkspaceAgentLifecycleStateReady,
+				},
+			},
+			url: "wss://local.test/?agent_name=agent0",
+			expectedAgentUpdate: &workspacesdk.AgentUpdate{
+				Lifecycle: codersdk.WorkspaceAgentLifecycleCreated,
+				ID:        agentID,
+			},
+		},
+		{
+			name: "namedAgentNonexistent",
+			agents: []database.WorkspaceAgent{
+				{
+					Name:           "agent0",
+					ID:             agentID,
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+				{
+					Name:           "agent1",
+					ID:             uuid.UUID{77},
+					LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+				},
+			},
+			url:                    "wss://local.test/?agent_name=agent2",
+			expectedErrorCode:      workspacesdk.WatchErrorNameNotFound,
+			expectedErrorRetryable: false,
+		},
+		{
+			name:                   "dbError",
+			agentDBError:           xerrors.New("a bad thing happened"),
+			url:                    "wss://local.test/",
+			expectedErrorCode:      workspacesdk.WatchErrorDatabase,
+			expectedErrorRetryable: true,
+		},
+		{
+			name:                   "unauthorized",
+			agentDBError:           dbauthz.NotAuthorizedError{Err: xerrors.New("not allowed")},
+			url:                    "wss://local.test/",
+			expectedErrorCode:      workspacesdk.WatchErrorDatabase,
+			expectedErrorRetryable: false,
+		},
+		{
+			name:                   "noAgents",
+			agents:                 []database.WorkspaceAgent{},
+			url:                    "wss://local.test/",
+			expectedErrorCode:      workspacesdk.WatchErrorNoAgents,
+			expectedErrorRetryable: false,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+			logger := testutil.Logger(t)
+			h := newHarness(ctx, t, logger)
+
+			h.db.EXPECT().GetLatestWorkspaceBuildWithStatusByWorkspaceID(gomock.Any(), h.workspace.ID).
+				Times(1).
+				Return(database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{
+					Transition:  database.WorkspaceTransitionStart,
+					BuildNumber: 1,
+					JobStatus:   database.ProvisionerJobStatusSucceeded,
+					WorkspaceTable: database.WorkspaceTable{
+						ID:             h.workspace.ID,
+						OwnerID:        userID,
+						OrganizationID: orgID,
+					},
+				}, nil)
+			// RBAC check for agent query
+			h.db.EXPECT().GetWorkspaceByID(gomock.Any(), h.workspace.ID).
+				Times(1).
+				Return(h.workspace, nil)
+			h.db.EXPECT().GetWorkspaceAgentsByWorkspaceAndBuildNumber(
+				gomock.Any(),
+				database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
+					WorkspaceID: h.workspace.ID,
+					BuildNumber: 1,
+				}).
+				Times(1).
+				Return(tc.agents, tc.agentDBError)
+
+			dec, err := h.Dial(ctx, tc.url)
+			require.NoError(t, err)
+			defer dec.Close()
+			events := dec.Chan()
+			e0 := testutil.RequireReceive(ctx, t, events)
+			require.Equal(t, workspacesdk.ConnectionWatchEvent{
+				BuildUpdate: &workspacesdk.BuildUpdate{
+					Transition: codersdk.WorkspaceTransitionStart,
+					JobStatus:  codersdk.ProvisionerJobSucceeded,
+				},
+			}, e0)
+
+			e1 := testutil.RequireReceive(ctx, t, events)
+			if tc.expectedAgentUpdate != nil {
+				require.Equal(t, workspacesdk.ConnectionWatchEvent{AgentUpdate: tc.expectedAgentUpdate}, e1)
+			} else {
+				require.NotNil(t, e1.Error)
+				require.Equal(t, tc.expectedErrorRetryable, e1.Error.Retryable)
+				require.Equal(t, tc.expectedErrorCode, e1.Error.Code)
+			}
+		})
+	}
+}
+
+func TestWatcher_LostAccess(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := testutil.Logger(t)
+	h := newHarness(ctx, t, logger)
+
+	h.db.EXPECT().GetLatestWorkspaceBuildWithStatusByWorkspaceID(gomock.Any(), h.workspace.ID).
+		Times(1).
+		Return(database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{
+			Transition:  database.WorkspaceTransitionStart,
+			BuildNumber: 1,
+			JobStatus:   database.ProvisionerJobStatusSucceeded,
+			WorkspaceTable: database.WorkspaceTable{
+				ID:             h.workspace.ID,
+				OwnerID:        uuid.UUID{99}, // workspace gets a new owner, e.g.
+				OrganizationID: orgID,
+			},
+		}, nil)
+
+	dec, err := h.Dial(ctx, "wss://local.test/")
+	require.NoError(t, err)
+	defer func() {
+		err := dec.Close()
+		require.NoError(t, err)
+	}()
+	events := dec.Chan()
+	e0 := testutil.RequireReceive(ctx, t, events)
+	require.NotNil(t, e0.Error)
+	require.Equal(t, workspacesdk.WatchErrorDatabase, e0.Error.Code)
+	require.False(t, e0.Error.Retryable)
+	require.Equal(t, "unauthorized", e0.Error.Details, "should not leak internal auth details")
+}
+
+func TestWatcher_PublishChanges(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := testutil.Logger(t)
+	h := newHarness(ctx, t, logger)
+
+	// Initial build update, job is running.
+	build0 := h.db.EXPECT().GetLatestWorkspaceBuildWithStatusByWorkspaceID(gomock.Any(), h.workspace.ID).
+		Times(1).
+		Return(database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{
+			Transition:  database.WorkspaceTransitionStart,
+			BuildNumber: 1,
+			JobStatus:   database.ProvisionerJobStatusRunning,
+			WorkspaceTable: database.WorkspaceTable{
+				ID:             h.workspace.ID,
+				OwnerID:        userID,
+				OrganizationID: orgID,
+			},
+		}, nil)
+
+	dec, err := h.Dial(ctx, "wss://local.test/")
+	require.NoError(t, err)
+	defer func() {
+		err := dec.Close()
+		require.NoError(t, err)
+	}()
+	events := dec.Chan()
+
+	e0 := testutil.RequireReceive(ctx, t, events)
+	require.Equal(t, workspacesdk.ConnectionWatchEvent{
+		BuildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransitionStart,
+			JobStatus:  codersdk.ProvisionerJobRunning,
+		},
+	}, e0)
+
+	// Since job is still running, we don't immediately query for agents. Next we set up the db queries and send in an
+	// update over the pubsub to kick  a new query.
+	build1 := h.db.EXPECT().GetLatestWorkspaceBuildWithStatusByWorkspaceID(gomock.Any(), h.workspace.ID).
+		After(build0).
+		Times(1).
+		Return(database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{
+			Transition:  database.WorkspaceTransitionStart,
+			BuildNumber: 1,
+			JobStatus:   database.ProvisionerJobStatusSucceeded,
+			WorkspaceTable: database.WorkspaceTable{
+				ID:             h.workspace.ID,
+				OwnerID:        userID,
+				OrganizationID: orgID,
+			},
+		}, nil)
+	// RBAC check for agent query
+	h.db.EXPECT().GetWorkspaceByID(gomock.Any(), h.workspace.ID).
+		After(build1).
+		Times(2). // these queries are identical between the initial and the update below
+		Return(h.workspace, nil)
+	agent0 := h.db.EXPECT().GetWorkspaceAgentsByWorkspaceAndBuildNumber(
+		gomock.Any(),
+		database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
+			WorkspaceID: h.workspace.ID,
+			BuildNumber: 1,
+		}).
+		After(build1).
+		Times(1).
+		Return([]database.WorkspaceAgent{
+			{
+				Name:           "test",
+				ID:             agentID,
+				LifecycleState: database.WorkspaceAgentLifecycleStateCreated,
+			},
+		}, nil)
+	changeMsg := wspubsub.WorkspaceEvent{
+		Kind:        wspubsub.WorkspaceEventKindStateChange,
+		WorkspaceID: h.workspace.ID,
+	}
+	changeBytes, err := json.Marshal(changeMsg)
+	require.NoError(t, err)
+	err = h.pub.Publish(wspubsub.WorkspaceEventChannel(h.workspace.OwnerID), changeBytes)
+	require.NoError(t, err)
+
+	e1 := testutil.RequireReceive(ctx, t, events)
+	require.Equal(t, workspacesdk.ConnectionWatchEvent{
+		BuildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransitionStart,
+			JobStatus:  codersdk.ProvisionerJobSucceeded,
+		},
+	}, e1)
+	e2 := testutil.RequireReceive(ctx, t, events)
+	require.Equal(t, workspacesdk.ConnectionWatchEvent{AgentUpdate: &workspacesdk.AgentUpdate{
+		ID:        agentID,
+		Lifecycle: codersdk.WorkspaceAgentLifecycleCreated,
+	}}, e2)
+
+	// Finally, send in a change event for the agent. But first, program the mock for the expected query.
+	h.db.EXPECT().GetWorkspaceAgentsByWorkspaceAndBuildNumber(
+		gomock.Any(),
+		database.GetWorkspaceAgentsByWorkspaceAndBuildNumberParams{
+			WorkspaceID: h.workspace.ID,
+			BuildNumber: 1,
+		}).
+		After(agent0).
+		Times(1).
+		Return([]database.WorkspaceAgent{
+			{
+				Name:           "test",
+				ID:             agentID,
+				LifecycleState: database.WorkspaceAgentLifecycleStateReady,
+			},
+		}, nil)
+	changeMsg = wspubsub.WorkspaceEvent{
+		Kind:        wspubsub.WorkspaceEventKindAgentLifecycleUpdate,
+		WorkspaceID: h.workspace.ID,
+		AgentID:     &agentID,
+	}
+	changeBytes, err = json.Marshal(changeMsg)
+	require.NoError(t, err)
+	err = h.pub.Publish(wspubsub.WorkspaceEventChannel(h.workspace.OwnerID), changeBytes)
+	require.NoError(t, err)
+
+	e3 := testutil.RequireReceive(ctx, t, events)
+	require.Equal(t, workspacesdk.ConnectionWatchEvent{AgentUpdate: &workspacesdk.AgentUpdate{
+		ID:        agentID,
+		Lifecycle: codersdk.WorkspaceAgentLifecycleReady,
+	}}, e3)
+}
+
+func TestWatcher_ClosedBeforeDial(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := testutil.Logger(t)
+	h := newHarness(ctx, t, logger)
+	h.watcher.Close()
+	_, err := h.Dial(ctx, "wss://local.test/")
+	var sdkError *codersdk.Error
+	require.True(t, errors.As(err, &sdkError))
+	require.Equal(t, http.StatusServiceUnavailable, sdkError.StatusCode())
+}
+
+func TestWatcher_ClosedAfterDial(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitShort)
+	logger := testutil.Logger(t)
+	h := newHarness(ctx, t, logger)
+
+	h.db.EXPECT().GetLatestWorkspaceBuildWithStatusByWorkspaceID(gomock.Any(), h.workspace.ID).
+		Times(1).
+		Return(database.GetLatestWorkspaceBuildWithStatusByWorkspaceIDRow{
+			Transition:  database.WorkspaceTransitionStop,
+			BuildNumber: 1,
+			JobStatus:   database.ProvisionerJobStatusSucceeded,
+			WorkspaceTable: database.WorkspaceTable{
+				ID:             h.workspace.ID,
+				OwnerID:        userID,
+				OrganizationID: orgID,
+			},
+		}, nil)
+
+	dec, err := h.Dial(ctx, "wss://local.test/")
+	require.NoError(t, err)
+	events := dec.Chan()
+	_ = testutil.RequireReceive(ctx, t, events)
+
+	closed := make(chan struct{})
+	go func() {
+		defer close(closed)
+		h.watcher.Close()
+	}()
+
+	e := testutil.RequireReceive(ctx, t, events)
+	require.NotNil(t, e.Error)
+	require.Equal(t, workspacesdk.WatchErrorServerShutdown, e.Error.Code)
+	require.True(t, e.Error.Retryable)
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("context timed out")
+	case _, ok := <-events:
+		require.False(t, ok, "socket not closed")
+	}
+	testutil.TryReceive(ctx, t, closed)
+}