chore: send workspace pubsub events by owner id (#14964)

We currently send empty payloads to pubsub channels of the form `workspace:<workspace_id>` to notify listeners of updates to workspaces (such as for refreshing the workspace dashboard).

To support https://github.com/coder/coder/issues/14716, we'll instead send `WorkspaceEvent` payloads to pubsub channels of the form `workspace_owner:<owner_id>`. This enables a listener to receive events for all workspaces owned by a user.
This PR replaces the usage of the old channels without modifying any existing behaviors.

```
type WorkspaceEvent struct {
	Kind        WorkspaceEventKind `json:"kind"`
	WorkspaceID uuid.UUID          `json:"workspace_id" format:"uuid"`
	// AgentID is only set for WorkspaceEventKindAgent* events
	// (excluding AgentTimeout)
	AgentID *uuid.UUID `json:"agent_id,omitempty" format:"uuid"`
}
```

We've defined `WorkspaceEventKind`s based on how the old channel was used, but it's not yet necessary to inspect the types of any of the events, as the existing listeners are designed to fire off any of them.

```
WorkspaceEventKindStateChange     WorkspaceEventKind = "state_change"
WorkspaceEventKindStatsUpdate     WorkspaceEventKind = "stats_update"
WorkspaceEventKindMetadataUpdate  WorkspaceEventKind = "mtd_update"
WorkspaceEventKindAppHealthUpdate WorkspaceEventKind = "app_health"

WorkspaceEventKindAgentLifecycleUpdate  WorkspaceEventKind = "agt_lifecycle_update"
WorkspaceEventKindAgentLogsUpdate       WorkspaceEventKind = "agt_logs_update"
WorkspaceEventKindAgentConnectionUpdate WorkspaceEventKind = "agt_connection_update"
WorkspaceEventKindAgentLogsOverflow     WorkspaceEventKind = "agt_logs_overflow"
WorkspaceEventKindAgentTimeout          WorkspaceEventKind = "agt_timeout"
```

coderd/agentapi/api.go

@@ -24,6 +24,7 @@ import (
 	"github.com/coder/coder/v2/coderd/prometheusmetrics"
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/coderd/workspacestats"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/tailnet"
@@ -45,14 +46,15 @@ type API struct {
 	*ScriptsAPI
 	*tailnet.DRPCService
 
-	mu                sync.Mutex
-	cachedWorkspaceID uuid.UUID
+	mu sync.Mutex
 }
 
 var _ agentproto.DRPCAgentServer = &API{}
 
 type Options struct {
-	AgentID uuid.UUID
+	AgentID     uuid.UUID
+	OwnerID     uuid.UUID
+	WorkspaceID uuid.UUID
 
 	Ctx                               context.Context
 	Log                               slog.Logger
@@ -62,7 +64,7 @@ type Options struct {
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
 	StatsReporter                     *workspacestats.Reporter
 	AppearanceFetcher                 *atomic.Pointer[appearance.Fetcher]
-	PublishWorkspaceUpdateFn          func(ctx context.Context, workspaceID uuid.UUID)
+	PublishWorkspaceUpdateFn          func(ctx context.Context, userID uuid.UUID, event wspubsub.WorkspaceEvent)
 	PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
 	NetworkTelemetryHandler           func(batch []*tailnetproto.TelemetryEvent)
 
@@ -75,18 +77,13 @@ type Options struct {
 	ExternalAuthConfigs       []*externalauth.Config
 	Experiments               codersdk.Experiments
 
-	// Optional:
-	// WorkspaceID avoids a future lookup to find the workspace ID by setting
-	// the cache in advance.
-	WorkspaceID          uuid.UUID
 	UpdateAgentMetricsFn func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
 }
 
 func New(opts Options) *API {
 	api := &API{
-		opts:              opts,
-		mu:                sync.Mutex{},
-		cachedWorkspaceID: opts.WorkspaceID,
+		opts: opts,
+		mu:   sync.Mutex{},
 	}
 
 	api.ManifestAPI = &ManifestAPI{
@@ -98,16 +95,7 @@ func New(opts Options) *API {
 		AgentFn:                  api.agent,
 		Database:                 opts.Database,
 		DerpMapFn:                opts.DerpMapFn,
-		WorkspaceIDFn: func(ctx context.Context, wa *database.WorkspaceAgent) (uuid.UUID, error) {
-			if opts.WorkspaceID != uuid.Nil {
-				return opts.WorkspaceID, nil
-			}
-			ws, err := opts.Database.GetWorkspaceByAgentID(ctx, wa.ID)
-			if err != nil {
-				return uuid.Nil, err
-			}
-			return ws.ID, nil
-		},
+		WorkspaceID:              opts.WorkspaceID,
 	}
 
 	api.AnnouncementBannerAPI = &AnnouncementBannerAPI{
@@ -125,7 +113,7 @@ func New(opts Options) *API {
 
 	api.LifecycleAPI = &LifecycleAPI{
 		AgentFn:                  api.agent,
-		WorkspaceIDFn:            api.workspaceID,
+		WorkspaceID:              opts.WorkspaceID,
 		Database:                 opts.Database,
 		Log:                      opts.Log,
 		PublishWorkspaceUpdateFn: api.publishWorkspaceUpdate,
@@ -209,39 +197,11 @@ func (a *API) agent(ctx context.Context) (database.WorkspaceAgent, error) {
 	return agent, nil
 }
 
-func (a *API) workspaceID(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-	a.mu.Lock()
-	if a.cachedWorkspaceID != uuid.Nil {
-		id := a.cachedWorkspaceID
-		a.mu.Unlock()
-		return id, nil
-	}
-
-	if agent == nil {
-		agnt, err := a.agent(ctx)
-		if err != nil {
-			return uuid.Nil, err
-		}
-		agent = &agnt
-	}
-
-	getWorkspaceAgentByIDRow, err := a.opts.Database.GetWorkspaceByAgentID(ctx, agent.ID)
-	if err != nil {
-		return uuid.Nil, xerrors.Errorf("get workspace by agent id %q: %w", agent.ID, err)
-	}
-
-	a.mu.Lock()
-	a.cachedWorkspaceID = getWorkspaceAgentByIDRow.ID
-	a.mu.Unlock()
-	return getWorkspaceAgentByIDRow.ID, nil
-}
-
-func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent) error {
-	workspaceID, err := a.workspaceID(ctx, agent)
-	if err != nil {
-		return err
-	}
-
-	a.opts.PublishWorkspaceUpdateFn(ctx, workspaceID)
+func (a *API) publishWorkspaceUpdate(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
+	a.opts.PublishWorkspaceUpdateFn(ctx, a.opts.OwnerID, wspubsub.WorkspaceEvent{
+		Kind:        kind,
+		WorkspaceID: a.opts.WorkspaceID,
+		AgentID:     &agent.ID,
+	})
 	return nil
 }

coderd/agentapi/apps.go

@@ -9,13 +9,14 @@ import (
 	"cdr.dev/slog"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 )
 
 type AppsAPI struct {
 	AgentFn                  func(context.Context) (database.WorkspaceAgent, error)
 	Database                 database.Store
 	Log                      slog.Logger
-	PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
+	PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
 }
 
 func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
@@ -96,7 +97,7 @@ func (a *AppsAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.Bat
 	}
 
 	if a.PublishWorkspaceUpdateFn != nil && len(newApps) > 0 {
-		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
+		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAppHealthUpdate)
 		if err != nil {
 			return nil, xerrors.Errorf("publish workspace update: %w", err)
 		}

coderd/agentapi/apps_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/coder/coder/v2/coderd/agentapi"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 )
 
 func TestBatchUpdateAppHealths(t *testing.T) {
@@ -62,7 +63,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},
@@ -100,7 +101,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},
@@ -139,7 +140,7 @@ func TestBatchUpdateAppHealths(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},

coderd/agentapi/lifecycle.go

@@ -15,6 +15,7 @@ import (
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 )
 
 type contextKeyAPIVersion struct{}
@@ -25,10 +26,10 @@ func WithAPIVersion(ctx context.Context, version string) context.Context {
 
 type LifecycleAPI struct {
 	AgentFn                  func(context.Context) (database.WorkspaceAgent, error)
-	WorkspaceIDFn            func(context.Context, *database.WorkspaceAgent) (uuid.UUID, error)
+	WorkspaceID              uuid.UUID
 	Database                 database.Store
 	Log                      slog.Logger
-	PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent) error
+	PublishWorkspaceUpdateFn func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
 
 	TimeNowFn func() time.Time // defaults to dbtime.Now()
 }
@@ -45,13 +46,9 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
 	if err != nil {
 		return nil, err
 	}
-	workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
-	if err != nil {
-		return nil, err
-	}
 
 	logger := a.Log.With(
-		slog.F("workspace_id", workspaceID),
+		slog.F("workspace_id", a.WorkspaceID),
 		slog.F("payload", req),
 	)
 	logger.Debug(ctx, "workspace agent state report")
@@ -122,7 +119,7 @@ func (a *LifecycleAPI) UpdateLifecycle(ctx context.Context, req *agentproto.Upda
 	}
 
 	if a.PublishWorkspaceUpdateFn != nil {
-		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
+		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentLifecycleUpdate)
 		if err != nil {
 			return nil, xerrors.Errorf("publish workspace update: %w", err)
 		}
@@ -140,15 +137,11 @@ func (a *LifecycleAPI) UpdateStartup(ctx context.Context, req *agentproto.Update
 	if err != nil {
 		return nil, err
 	}
-	workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
-	if err != nil {
-		return nil, err
-	}
 
 	a.Log.Debug(
 		ctx,
 		"post workspace agent version",
-		slog.F("workspace_id", workspaceID),
+		slog.F("workspace_id", a.WorkspaceID),
 		slog.F("agent_version", req.Startup.Version),
 	)
 

coderd/agentapi/lifecycle_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 )
 
 func TestUpdateLifecycle(t *testing.T) {
@@ -69,12 +70,10 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agentCreated, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
+			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},
@@ -111,11 +110,9 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agentStarting, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
 			// Test that nil publish fn works.
 			PublishWorkspaceUpdateFn: nil,
 		}
@@ -156,12 +153,10 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agentCreated, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
+			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},
@@ -204,9 +199,7 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agentCreated, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
+			WorkspaceID:              workspaceID,
 			Database:                 dbM,
 			Log:                      slogtest.Make(t, nil),
 			PublishWorkspaceUpdateFn: nil,
@@ -239,12 +232,10 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
+			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				atomic.AddInt64(&publishCalled, 1)
 				return nil
 			},
@@ -314,12 +305,10 @@ func TestUpdateLifecycle(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agentCreated, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent) error {
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
+			PublishWorkspaceUpdateFn: func(ctx context.Context, agent *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishCalled = true
 				return nil
 			},
@@ -354,11 +343,9 @@ func TestUpdateStartup(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
 			// Not used by UpdateStartup.
 			PublishWorkspaceUpdateFn: nil,
 		}
@@ -402,11 +389,9 @@ func TestUpdateStartup(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
 			// Not used by UpdateStartup.
 			PublishWorkspaceUpdateFn: nil,
 		}
@@ -435,11 +420,9 @@ func TestUpdateStartup(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, agent *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspaceID, nil
-			},
-			Database: dbM,
-			Log:      slogtest.Make(t, nil),
+			WorkspaceID: workspaceID,
+			Database:    dbM,
+			Log:         slogtest.Make(t, nil),
 			// Not used by UpdateStartup.
 			PublishWorkspaceUpdateFn: nil,
 		}

coderd/agentapi/logs.go

@@ -11,6 +11,7 @@ import (
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
 
@@ -18,7 +19,7 @@ type LogsAPI struct {
 	AgentFn                           func(context.Context) (database.WorkspaceAgent, error)
 	Database                          database.Store
 	Log                               slog.Logger
-	PublishWorkspaceUpdateFn          func(context.Context, *database.WorkspaceAgent) error
+	PublishWorkspaceUpdateFn          func(context.Context, *database.WorkspaceAgent, wspubsub.WorkspaceEventKind) error
 	PublishWorkspaceAgentLogsUpdateFn func(ctx context.Context, workspaceAgentID uuid.UUID, msg agentsdk.LogsNotifyMessage)
 
 	TimeNowFn func() time.Time // defaults to dbtime.Now()
@@ -123,7 +124,7 @@ func (a *LogsAPI) BatchCreateLogs(ctx context.Context, req *agentproto.BatchCrea
 		}
 
 		if a.PublishWorkspaceUpdateFn != nil {
-			err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
+			err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentLogsOverflow)
 			if err != nil {
 				return nil, xerrors.Errorf("publish workspace update: %w", err)
 			}
@@ -143,7 +144,7 @@ func (a *LogsAPI) BatchCreateLogs(ctx context.Context, req *agentproto.BatchCrea
 	if workspaceAgent.LogsLength == 0 && a.PublishWorkspaceUpdateFn != nil {
 		// If these are the first logs being appended, we publish a UI update
 		// to notify the UI that logs are now available.
-		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent)
+		err = a.PublishWorkspaceUpdateFn(ctx, &workspaceAgent, wspubsub.WorkspaceEventKindAgentFirstLogs)
 		if err != nil {
 			return nil, xerrors.Errorf("publish workspace update: %w", err)
 		}

coderd/agentapi/logs_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
 
@@ -50,7 +51,7 @@ func TestBatchCreateLogs(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishWorkspaceUpdateCalled = true
 				return nil
 			},
@@ -154,7 +155,7 @@ func TestBatchCreateLogs(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishWorkspaceUpdateCalled = true
 				return nil
 			},
@@ -202,7 +203,7 @@ func TestBatchCreateLogs(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishWorkspaceUpdateCalled = true
 				return nil
 			},
@@ -295,7 +296,7 @@ func TestBatchCreateLogs(t *testing.T) {
 				},
 				Database: dbM,
 				Log:      slogtest.Make(t, nil),
-				PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+				PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 					publishWorkspaceUpdateCalled = true
 					return nil
 				},
@@ -339,7 +340,7 @@ func TestBatchCreateLogs(t *testing.T) {
 				},
 				Database: dbM,
 				Log:      slogtest.Make(t, nil),
-				PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+				PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 					publishWorkspaceUpdateCalled = true
 					return nil
 				},
@@ -386,7 +387,7 @@ func TestBatchCreateLogs(t *testing.T) {
 			},
 			Database: dbM,
 			Log:      slogtest.Make(t, nil),
-			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent) error {
+			PublishWorkspaceUpdateFn: func(ctx context.Context, wa *database.WorkspaceAgent, kind wspubsub.WorkspaceEventKind) error {
 				publishWorkspaceUpdateCalled = true
 				return nil
 			},

coderd/agentapi/manifest.go

@@ -29,11 +29,11 @@ type ManifestAPI struct {
 	ExternalAuthConfigs      []*externalauth.Config
 	DisableDirectConnections bool
 	DerpForceWebSockets      bool
+	WorkspaceID              uuid.UUID
 
-	AgentFn       func(context.Context) (database.WorkspaceAgent, error)
-	WorkspaceIDFn func(context.Context, *database.WorkspaceAgent) (uuid.UUID, error)
-	Database      database.Store
-	DerpMapFn     func() *tailcfg.DERPMap
+	AgentFn   func(context.Context) (database.WorkspaceAgent, error)
+	Database  database.Store
+	DerpMapFn func() *tailcfg.DERPMap
 }
 
 func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
@@ -41,11 +41,6 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
 	if err != nil {
 		return nil, err
 	}
-	workspaceID, err := a.WorkspaceIDFn(ctx, &workspaceAgent)
-	if err != nil {
-		return nil, err
-	}
-
 	var (
 		dbApps    []database.WorkspaceApp
 		scripts   []database.WorkspaceAgentScript
@@ -75,7 +70,7 @@ func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifest
 		return err
 	})
 	eg.Go(func() (err error) {
-		workspace, err = a.Database.GetWorkspaceByID(ctx, workspaceID)
+		workspace, err = a.Database.GetWorkspaceByID(ctx, a.WorkspaceID)
 		if err != nil {
 			return xerrors.Errorf("getting workspace by id: %w", err)
 		}

coderd/agentapi/manifest_test.go

@@ -288,11 +288,9 @@ func TestGetManifest(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, _ *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspace.ID, nil
-			},
-			Database:  mDB,
-			DerpMapFn: derpMapFn,
+			WorkspaceID: workspace.ID,
+			Database:    mDB,
+			DerpMapFn:   derpMapFn,
 		}
 
 		mDB.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return(apps, nil)
@@ -355,11 +353,9 @@ func TestGetManifest(t *testing.T) {
 			AgentFn: func(ctx context.Context) (database.WorkspaceAgent, error) {
 				return agent, nil
 			},
-			WorkspaceIDFn: func(ctx context.Context, _ *database.WorkspaceAgent) (uuid.UUID, error) {
-				return workspace.ID, nil
-			},
-			Database:  mDB,
-			DerpMapFn: derpMapFn,
+			WorkspaceID: workspace.ID,
+			Database:    mDB,
+			DerpMapFn:   derpMapFn,
 		}
 
 		mDB.EXPECT().GetWorkspaceAppsByAgentID(gomock.Any(), agent.ID).Return(apps, nil)

coderd/agentapi/stats_test.go

@@ -23,6 +23,7 @@ import (
 	"github.com/coder/coder/v2/coderd/schedule"
 	"github.com/coder/coder/v2/coderd/workspacestats"
 	"github.com/coder/coder/v2/coderd/workspacestats/workspacestatstest"
+	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -153,12 +154,19 @@ func TestUpdateStates(t *testing.T) {
 		}).Return(nil)
 
 		// Ensure that pubsub notifications are sent.
-		notifyDescription := make(chan []byte)
-		ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
-			go func() {
-				notifyDescription <- description
-			}()
-		})
+		notifyDescription := make(chan struct{})
+		ps.SubscribeWithErr(wspubsub.WorkspaceEventChannel(workspace.OwnerID),
+			wspubsub.HandleWorkspaceEvent(
+				func(_ context.Context, e wspubsub.WorkspaceEvent, err error) {
+					if err != nil {
+						return
+					}
+					if e.Kind == wspubsub.WorkspaceEventKindStatsUpdate && e.WorkspaceID == workspace.ID {
+						go func() {
+							notifyDescription <- struct{}{}
+						}()
+					}
+				}))
 
 		resp, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)
@@ -183,8 +191,7 @@ func TestUpdateStates(t *testing.T) {
 		select {
 		case <-ctx.Done():
 			t.Error("timed out while waiting for pubsub notification")
-		case description := <-notifyDescription:
-			require.Equal(t, description, []byte{})
+		case <-notifyDescription:
 		}
 		require.True(t, updateAgentMetricsFnCalled)
 	})
@@ -495,12 +502,19 @@ func TestUpdateStates(t *testing.T) {
 		dbM.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
 
 		// Ensure that pubsub notifications are sent.
-		notifyDescription := make(chan []byte)
-		ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, description []byte) {
-			go func() {
-				notifyDescription <- description
-			}()
-		})
+		notifyDescription := make(chan struct{})
+		ps.SubscribeWithErr(wspubsub.WorkspaceEventChannel(workspace.OwnerID),
+			wspubsub.HandleWorkspaceEvent(
+				func(_ context.Context, e wspubsub.WorkspaceEvent, err error) {
+					if err != nil {
+						return
+					}
+					if e.Kind == wspubsub.WorkspaceEventKindStatsUpdate && e.WorkspaceID == workspace.ID {
+						go func() {
+							notifyDescription <- struct{}{}
+						}()
+					}
+				}))
 
 		resp, err := api.UpdateStats(context.Background(), req)
 		require.NoError(t, err)
@@ -523,8 +537,7 @@ func TestUpdateStates(t *testing.T) {
 		select {
 		case <-ctx.Done():
 			t.Error("timed out while waiting for pubsub notification")
-		case description := <-notifyDescription:
-			require.Equal(t, description, []byte{})
+		case <-notifyDescription:
 		}
 		require.True(t, updateAgentMetricsFnCalled)
 	})