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")

}