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/apidoc/docs.go

@@ -12783,6 +12783,41 @@ const docTemplate = `{
                 ]
             }
         },
+        "/api/v2/workspaces/{workspace}/agent-connection-watch": {
+            "get": {
+                "produces": [
+                    "application/json"
+                ],
+                "tags": [
+                    "Workspaces"
+                ],
+                "summary": "Workspace Agent Connection Watch",
+                "operationId": "workspace-agent-connection-watch",
+                "parameters": [
+                    {
+                        "type": "string",
+                        "format": "uuid",
+                        "description": "Workspace ID",
+                        "name": "workspace",
+                        "in": "path",
+                        "required": true
+                    }
+                ],
+                "responses": {
+                    "101": {
+                        "description": "Switching Protocols",
+                        "schema": {
+                            "$ref": "#/definitions/workspacesdk.ConnectionWatchEvent"
+                        }
+                    }
+                },
+                "security": [
+                    {
+                        "CoderSessionToken": []
+                    }
+                ]
+            }
+        },
         "/api/v2/workspaces/{workspace}/autostart": {
             "put": {
                 "consumes": [
@@ -27964,6 +27999,93 @@ const docTemplate = `{
                 }
             }
         },
+        "workspacesdk.AgentUpdate": {
+            "type": "object",
+            "properties": {
+                "id": {
+                    "type": "string",
+                    "format": "uuid"
+                },
+                "lifecycle": {
+                    "$ref": "#/definitions/codersdk.WorkspaceAgentLifecycle"
+                }
+            }
+        },
+        "workspacesdk.BuildUpdate": {
+            "type": "object",
+            "properties": {
+                "job_status": {
+                    "$ref": "#/definitions/codersdk.ProvisionerJobStatus"
+                },
+                "transition": {
+                    "$ref": "#/definitions/codersdk.WorkspaceTransition"
+                }
+            }
+        },
+        "workspacesdk.ConnectionWatchEvent": {
+            "type": "object",
+            "properties": {
+                "agent_update": {
+                    "$ref": "#/definitions/workspacesdk.AgentUpdate"
+                },
+                "build_update": {
+                    "$ref": "#/definitions/workspacesdk.BuildUpdate"
+                },
+                "error": {
+                    "$ref": "#/definitions/workspacesdk.WatchError"
+                }
+            }
+        },
+        "workspacesdk.WatchError": {
+            "type": "object",
+            "properties": {
+                "code": {
+                    "$ref": "#/definitions/workspacesdk.WatchErrorCode"
+                },
+                "details": {
+                    "type": "string"
+                },
+                "message": {
+                    "type": "string"
+                },
+                "retryable": {
+                    "type": "boolean"
+                }
+            }
+        },
+        "workspacesdk.WatchErrorCode": {
+            "type": "integer",
+            "enum": [
+                0,
+                1,
+                2,
+                3,
+                4,
+                5,
+                6
+            ],
+            "x-enum-comments": {
+                "_": "Ensure that zero value is not a valid code"
+            },
+            "x-enum-descriptions": [
+                "Ensure that zero value is not a valid code",
+                "",
+                "",
+                "",
+                "",
+                "",
+                ""
+            ],
+            "x-enum-varnames": [
+                "_",
+                "WatchErrorTooManyAgents",
+                "WatchErrorNameNotFound",
+                "WatchErrorNoAgents",
+                "WatchErrorServerShutdown",
+                "WatchErrorDatabase",
+                "WatchErrorInternal"
+            ]
+        },
         "wsproxysdk.CryptoKeysResponse": {
             "type": "object",
             "properties": {

coderd/apidoc/swagger.json

@@ -11343,6 +11343,37 @@
 				]
 			}
 		},
+		"/api/v2/workspaces/{workspace}/agent-connection-watch": {
+			"get": {
+				"produces": ["application/json"],
+				"tags": ["Workspaces"],
+				"summary": "Workspace Agent Connection Watch",
+				"operationId": "workspace-agent-connection-watch",
+				"parameters": [
+					{
+						"type": "string",
+						"format": "uuid",
+						"description": "Workspace ID",
+						"name": "workspace",
+						"in": "path",
+						"required": true
+					}
+				],
+				"responses": {
+					"101": {
+						"description": "Switching Protocols",
+						"schema": {
+							"$ref": "#/definitions/workspacesdk.ConnectionWatchEvent"
+						}
+					}
+				},
+				"security": [
+					{
+						"CoderSessionToken": []
+					}
+				]
+			}
+		},
 		"/api/v2/workspaces/{workspace}/autostart": {
 			"put": {
 				"consumes": ["application/json"],
@@ -25823,6 +25854,85 @@
 				}
 			}
 		},
+		"workspacesdk.AgentUpdate": {
+			"type": "object",
+			"properties": {
+				"id": {
+					"type": "string",
+					"format": "uuid"
+				},
+				"lifecycle": {
+					"$ref": "#/definitions/codersdk.WorkspaceAgentLifecycle"
+				}
+			}
+		},
+		"workspacesdk.BuildUpdate": {
+			"type": "object",
+			"properties": {
+				"job_status": {
+					"$ref": "#/definitions/codersdk.ProvisionerJobStatus"
+				},
+				"transition": {
+					"$ref": "#/definitions/codersdk.WorkspaceTransition"
+				}
+			}
+		},
+		"workspacesdk.ConnectionWatchEvent": {
+			"type": "object",
+			"properties": {
+				"agent_update": {
+					"$ref": "#/definitions/workspacesdk.AgentUpdate"
+				},
+				"build_update": {
+					"$ref": "#/definitions/workspacesdk.BuildUpdate"
+				},
+				"error": {
+					"$ref": "#/definitions/workspacesdk.WatchError"
+				}
+			}
+		},
+		"workspacesdk.WatchError": {
+			"type": "object",
+			"properties": {
+				"code": {
+					"$ref": "#/definitions/workspacesdk.WatchErrorCode"
+				},
+				"details": {
+					"type": "string"
+				},
+				"message": {
+					"type": "string"
+				},
+				"retryable": {
+					"type": "boolean"
+				}
+			}
+		},
+		"workspacesdk.WatchErrorCode": {
+			"type": "integer",
+			"enum": [0, 1, 2, 3, 4, 5, 6],
+			"x-enum-comments": {
+				"_": "Ensure that zero value is not a valid code"
+			},
+			"x-enum-descriptions": [
+				"Ensure that zero value is not a valid code",
+				"",
+				"",
+				"",
+				"",
+				"",
+				""
+			],
+			"x-enum-varnames": [
+				"_",
+				"WatchErrorTooManyAgents",
+				"WatchErrorNameNotFound",
+				"WatchErrorNoAgents",
+				"WatchErrorServerShutdown",
+				"WatchErrorDatabase",
+				"WatchErrorInternal"
+			]
+		},
 		"wsproxysdk.CryptoKeysResponse": {
 			"type": "object",
 			"properties": {

coderd/coderd.go

@@ -92,6 +92,7 @@ import (
 	"github.com/coder/coder/v2/coderd/webpush"
 	"github.com/coder/coder/v2/coderd/workspaceapps"
 	"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
+	"github.com/coder/coder/v2/coderd/workspaceconnwatcher"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/coderd/wsbuilder"
 	"github.com/coder/coder/v2/coderd/x/chatd"
@@ -923,6 +924,8 @@ func New(options *Options) *API {
 		APIKeyEncryptionKeycache: options.AppEncryptionKeyCache,
 	})
 
+	api.workspaceAgentConnWatcher = workspaceconnwatcher.New(api.ctx, options.Logger, options.Pubsub, options.Database)
+
 	apiKeyMiddleware := httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
 		DB:                            options.Database,
 		ActivateDormantUser:           ActivateDormantUser(options.Logger, &api.Auditor, options.Database),
@@ -1820,6 +1823,7 @@ func New(options *Options) *API {
 					r.Patch("/", api.patchWorkspaceACL)
 					r.Delete("/", api.deleteWorkspaceACL)
 				})
+				r.Get("/agent-connection-watch", api.workspaceAgentConnWatcher.WorkspaceAgentConnectionWatch)
 			})
 		})
 		r.Route("/workspacebuilds/{workspacebuild}", func(r chi.Router) {
@@ -2238,6 +2242,8 @@ type API struct {
 	// profile collection (via /debug/profile) can run at a time. The CPU
 	// profiler is process-global, so concurrent collections would fail.
 	ProfileCollecting atomic.Bool
+
+	workspaceAgentConnWatcher *workspaceconnwatcher.Watcher
 }
 
 // Close waits for all WebSocket connections to drain before returning.
@@ -2301,6 +2307,7 @@ func (api *API) Close() error {
 	_ = api.AppSigningKeyCache.Close()
 	_ = api.AppEncryptionKeyCache.Close()
 	_ = api.UpdatesProvider.Close()
+	api.workspaceAgentConnWatcher.Close()
 
 	if current := api.PrebuildsReconciler.Load(); current != nil {
 		ctx, giveUp := context.WithTimeoutCause(context.Background(), time.Second*30, xerrors.New("gave up waiting for reconciler to stop before shutdown"))

coderd/coderdtest/database.go

@@ -0,0 +1,28 @@
+package coderdtest
+
+import (
+	"sync/atomic"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"go.uber.org/mock/gomock"
+
+	"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/dbmock"
+	"github.com/coder/coder/v2/coderd/rbac"
+)
+
+func MockedDatabaseWithAuthz(t testing.TB, logger slog.Logger) (*gomock.Controller, *dbmock.MockStore, database.Store, rbac.Authorizer) {
+	ctrl := gomock.NewController(t)
+	mDB := dbmock.NewMockStore(ctrl)
+	auth := rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry())
+	accessControlStore := &atomic.Pointer[dbauthz.AccessControlStore]{}
+	var acs dbauthz.AccessControlStore = dbauthz.AGPLTemplateAccessControlStore{}
+	accessControlStore.Store(&acs)
+	// dbauthz will call Wrappers() to check for wrapped databases
+	mDB.EXPECT().Wrappers().Return([]string{}).AnyTimes()
+	authDB := dbauthz.New(mDB, auth, logger, accessControlStore)
+	return ctrl, mDB, authDB, auth
+}

coderd/coderdtest/subjects.go

@@ -0,0 +1,31 @@
+package coderdtest
+
+import (
+	"github.com/google/uuid"
+
+	"github.com/coder/coder/v2/coderd/rbac"
+	"github.com/coder/coder/v2/coderd/rbac/rolestore"
+)
+
+func MemberSubject(userID, orgID uuid.UUID) rbac.Subject {
+	memberRole, err := rbac.RoleByName(rbac.RoleMember())
+	if err != nil {
+		panic(err)
+	}
+	orgMember, err := rolestore.TestingGetSystemRole(
+		rbac.RoleOrgMember(),
+		orgID,
+		rbac.OrgSettings{ShareableWorkspaceOwners: rbac.ShareableWorkspaceOwnersNone},
+	)
+	if err != nil {
+		panic(err)
+	}
+	return rbac.Subject{
+		FriendlyName: "coderdtest-member",
+		Email:        "member@coderd.test",
+		Type:         rbac.SubjectTypeUser,
+		ID:           userID.String(),
+		Roles:        rbac.Roles{memberRole, orgMember},
+		Scope:        rbac.ScopeAll,
+	}.WithCachedASTValue()
+}

coderd/coderdtest/swagger_test.go

@@ -16,7 +16,7 @@ import (
 func TestEndpointsDocumented(t *testing.T) {
 	t.Parallel()
 
-	swaggerComments, err := coderdtest.ParseSwaggerComments("..")
+	swaggerComments, err := coderdtest.ParseSwaggerComments("..", "../workspaceconnwatcher")
 	require.NoError(t, err, "can't parse swagger comments")
 	require.NotEmpty(t, swaggerComments, "swagger comments must be present")
 

coderd/httpmw/workspaceparam.go

@@ -54,3 +54,7 @@ func ExtractWorkspaceParam(db database.Store) func(http.Handler) http.Handler {
 		})
 	}
 }
+
+func WithWorkspaceParam(ctx context.Context, workspace database.Workspace) context.Context {
+	return context.WithValue(ctx, workspaceParamContextKey{}, workspace)
+}

coderd/rbac/rolestore/rolestore.go

@@ -170,6 +170,25 @@ var systemRoles = map[string]permissionsFunc{
 	rbac.RoleOrgServiceAccount(): rbac.OrgServiceAccountPermissions,
 }
 
+func TestingGetSystemRole(name string, orgID uuid.UUID, settings rbac.OrgSettings) (rbac.Role, error) {
+	f, ok := systemRoles[name]
+	if !ok {
+		return rbac.Role{}, xerrors.Errorf("role %q not found", name)
+	}
+	perms := f(settings)
+	return rbac.Role{
+		Identifier:  rbac.RoleIdentifier{Name: name, OrganizationID: orgID},
+		DisplayName: "",
+		Site:        nil,
+		ByOrgID: map[string]rbac.OrgPermissions{
+			orgID.String(): {
+				Org:    perms.Org,
+				Member: perms.Member,
+			},
+		},
+	}, nil
+}
+
 // permissionsFunc produces the desired permissions for a system role
 // given organization settings.
 type permissionsFunc func(rbac.OrgSettings) rbac.OrgRolePermissions

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)
+}