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.

codersdk/workspacesdk/tunneler/integration_test.go

@@ -0,0 +1,100 @@
+package tunneler_test
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
+
+	"github.com/coder/coder/v2/agent/agenttest"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk/tunneler"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestTunneler_Integration is an integration test using coderdtest. It should be removed when we integrate the Tunneler
+// into coder ssh and those integration test cover this functionality.
+func TestTunneler_Integration(t *testing.T) {
+	t.Parallel()
+	ctx := testutil.Context(t, testutil.WaitLong)
+	client, store := coderdtest.NewWithDatabase(t, nil)
+	logger := testutil.Logger(t)
+	client.SetLogger(logger.Named("client"))
+	first := coderdtest.CreateFirstUser(t, client)
+	userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
+		r.Username = "myuser"
+	})
+	userClient.SetLogger(logger.Named("userclient"))
+	r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
+		Name:           "myworkspace",
+		OrganizationID: first.OrganizationID,
+		OwnerID:        user.ID,
+	}).WithAgent().Do()
+	wsSDKClient := workspacesdk.New(userClient)
+	logs := &bytes.Buffer{}
+
+	app := &sshApplication{
+		t:    t,
+		ctx:  ctx,
+		done: make(chan struct{}),
+	}
+
+	tun := tunneler.NewTunneler(wsSDKClient, tunneler.Config{
+		WorkspaceID:      r.Workspace.ID,
+		App:              app,
+		WorkspaceStarter: nil,
+		AgentName:        "",
+		LogWriter:        logs,
+		DebugLogger:      logger.Named("tunneler"),
+	})
+
+	testAgent := agenttest.New(t, client.URL, r.AgentToken)
+	defer testAgent.Close()
+
+	testutil.TryReceive(ctx, t, app.done)
+	require.Equal(t, app.result, "foo\n")
+
+	err := tun.GracefulShutdown(ctx)
+	require.NoError(t, err)
+}
+
+type sshApplication struct {
+	t      *testing.T
+	ctx    context.Context
+	client *ssh.Client
+	done   chan struct{}
+	result string
+}
+
+func (s *sshApplication) Close() error {
+	return s.client.Close()
+}
+
+func (s *sshApplication) Start(conn workspacesdk.AgentConn) error {
+	var err error
+	s.client, err = conn.SSHClient(s.ctx)
+	if err != nil {
+		s.t.Error(err)
+		return err
+	}
+	go func() {
+		defer close(s.done)
+		sess, err := s.client.NewSession()
+		if err != nil {
+			s.t.Error("failed to create session", err)
+		}
+		defer sess.Close()
+		out, err := sess.Output("echo foo")
+		if err != nil {
+			s.t.Error("failed to echo", err)
+		}
+		s.result = string(out)
+	}()
+	return nil
+}

codersdk/workspacesdk/tunneler/tunneler.go

@@ -11,6 +11,7 @@ import (
 	"cdr.dev/slog/v3"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
+	"github.com/coder/coder/v2/codersdk/wsjson"
 )
 
 type state int
@@ -33,6 +34,11 @@ type WorkspaceStarter interface {
 
 type Client interface {
 	DialAgent(dialCtx context.Context, agentID uuid.UUID, options *workspacesdk.DialAgentOptions) (workspacesdk.AgentConn, error)
+	WorkspaceAgentConnectionWatch(
+		dialCtx context.Context, workspaceID uuid.UUID, agentName string,
+	) (
+		dec *wsjson.Decoder[workspacesdk.ConnectionWatchEvent], err error,
+	)
 }
 
 const (
@@ -127,9 +133,9 @@ type Config struct {
 // ordering.
 type tunnelerEvent struct {
 	shutdownSignal    *shutdownSignal
-	buildUpdate       *buildUpdate
+	buildUpdate       *workspacesdk.BuildUpdate
 	provisionerJobLog *codersdk.ProvisionerJobLog
-	agentUpdate       *agentUpdate
+	agentUpdate       *workspacesdk.AgentUpdate
 	agentLog          *codersdk.WorkspaceAgentLog
 	appUpdate         *networkedApplicationUpdate
 	tailnetUpdate     *tailnetUpdate
@@ -137,16 +143,6 @@ type tunnelerEvent struct {
 
 type shutdownSignal struct{}
 
-type buildUpdate struct {
-	transition codersdk.WorkspaceTransition
-	jobStatus  codersdk.ProvisionerJobStatus
-}
-
-type agentUpdate struct {
-	lifecycle codersdk.WorkspaceAgentLifecycle
-	id        uuid.UUID
-}
-
 type networkedApplicationUpdate struct {
 	// up is true if the application is up. False if it is down.
 	up  bool
@@ -174,10 +170,64 @@ func NewTunneler(client Client, config Config) *Tunneler {
 	return t
 }
 
+func (t *Tunneler) GracefulShutdown(ctx context.Context) error {
+	select {
+	case t.events <- tunnelerEvent{shutdownSignal: &shutdownSignal{}}:
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-t.ctx.Done():
+	}
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		t.wg.Wait()
+	}()
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 func (t *Tunneler) start() {
 	defer t.wg.Done()
-	// here we would subscribe to updates.
-	// t.client.AgentConnectionWatch(t.config.WorkspaceID, t.config.AgentName)
+	d, err := t.client.WorkspaceAgentConnectionWatch(t.ctx, t.config.WorkspaceID, t.config.AgentName)
+	// TODO: handle retries
+	if err != nil {
+		return
+	}
+	defer d.Close()
+	c := d.Chan()
+	for {
+		select {
+		case <-t.ctx.Done():
+			return
+		case event, ok := <-c:
+			if !ok {
+				t.config.DebugLogger.Error(t.ctx, "watch closed")
+			}
+			if event.Error != nil {
+				t.config.DebugLogger.Error(t.ctx, "workspace agent connection watch error", slog.Error(event.Error))
+			}
+			if !ok || event.Error != nil {
+				// TODO: handle retries
+				select {
+				case t.events <- tunnelerEvent{shutdownSignal: &shutdownSignal{}}:
+				case <-t.ctx.Done():
+				}
+				return
+			}
+			select {
+			case <-t.ctx.Done():
+				return
+			case t.events <- tunnelerEvent{
+				buildUpdate: event.BuildUpdate,
+				agentUpdate: event.AgentUpdate,
+			}:
+			}
+		}
+	}
 }
 
 func (t *Tunneler) eventLoop() {
@@ -235,13 +285,13 @@ func (t *Tunneler) handleSignal() {
 	}
 }
 
-func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
+func (t *Tunneler) handleBuildUpdate(update *workspacesdk.BuildUpdate) {
 	if t.state == shutdownTailnet || t.state == shutdownApplication || t.state == exit {
 		return // no-op
 	}
 
 	var canMakeProgress, jobUnhealthy bool
-	switch update.jobStatus {
+	switch update.JobStatus {
 	case codersdk.ProvisionerJobPending, codersdk.ProvisionerJobRunning:
 		canMakeProgress = true
 	case codersdk.ProvisionerJobSucceeded:
@@ -249,21 +299,21 @@ func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
 		jobUnhealthy = true
 	}
 
-	if update.transition == codersdk.WorkspaceTransitionDelete {
-		t.config.DebugLogger.Info(t.ctx, "workspace is being deleted", slog.F("job_status", update.jobStatus))
+	if update.Transition == codersdk.WorkspaceTransitionDelete {
+		t.config.DebugLogger.Info(t.ctx, "workspace is being deleted", slog.F("job_status", update.JobStatus))
 		// treat same as signal
 		t.handleSignal()
 		return
 	}
 	if jobUnhealthy {
-		t.config.DebugLogger.Info(t.ctx, "build job is in unhealthy state", slog.F("job_status", update.jobStatus))
+		t.config.DebugLogger.Info(t.ctx, "build job is in unhealthy state", slog.F("job_status", update.JobStatus))
 		// treat same as signal
 		t.handleSignal()
 		return
 	}
 
-	if update.transition == codersdk.WorkspaceTransitionStart && canMakeProgress {
-		t.config.DebugLogger.Debug(t.ctx, "workspace is starting", slog.F("job_status", update.jobStatus))
+	if update.Transition == codersdk.WorkspaceTransitionStart && canMakeProgress {
+		t.config.DebugLogger.Debug(t.ctx, "workspace is starting", slog.F("job_status", update.JobStatus))
 		switch t.state {
 		// new build after we have already connected
 		case establishTailnet: // we are starting the tailnet
@@ -279,8 +329,8 @@ func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
 		}
 		return
 	}
-	if update.transition == codersdk.WorkspaceTransitionStart && update.jobStatus == codersdk.ProvisionerJobSucceeded {
-		t.config.DebugLogger.Debug(t.ctx, "workspace is started", slog.F("job_status", update.jobStatus))
+	if update.Transition == codersdk.WorkspaceTransitionStart && update.JobStatus == codersdk.ProvisionerJobSucceeded {
+		t.config.DebugLogger.Debug(t.ctx, "workspace is started", slog.F("job_status", update.JobStatus))
 		switch t.state {
 		case establishTailnet, applicationUp, tailnetUp:
 			// no-op. Later agent updates will tell us whether the tailnet connection is current.
@@ -290,7 +340,7 @@ func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
 		return
 	}
 
-	if update.transition == codersdk.WorkspaceTransitionStop {
+	if update.Transition == codersdk.WorkspaceTransitionStop {
 		// these cases take effect regardless of whether the transition is complete or not
 		switch t.state {
 		// all 3 of these mean a new build after we have already started connecting
@@ -312,7 +362,7 @@ func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
 			t.state = exit
 			return
 		}
-		if update.jobStatus == codersdk.ProvisionerJobSucceeded {
+		if update.JobStatus == codersdk.ProvisionerJobSucceeded {
 			switch t.state {
 			case stateInit, waitToStart, waitForAgent:
 				t.wg.Add(1)
@@ -335,29 +385,29 @@ func (t *Tunneler) handleBuildUpdate(update *buildUpdate) {
 	}
 	// unhittable
 	t.config.DebugLogger.Critical(t.ctx, "unhandled build update",
-		slog.F("job_status", update.jobStatus), slog.F("transition", update.transition), slog.F("state", t.state))
+		slog.F("job_status", update.JobStatus), slog.F("transition", update.Transition), slog.F("state", t.state))
 }
 
 func (*Tunneler) handleProvisionerJobLog(*codersdk.ProvisionerJobLog) {
 }
 
-func (t *Tunneler) handleAgentUpdate(update *agentUpdate) {
+func (t *Tunneler) handleAgentUpdate(update *workspacesdk.AgentUpdate) {
 	t.config.DebugLogger.Debug(t.ctx, "handling agent update",
 		slog.F("state", t.state),
-		slog.F("lifecycle", update.lifecycle),
-		slog.F("agent_id", update.id))
+		slog.F("lifecycle", update.Lifecycle),
+		slog.F("agent_id", update.ID))
 	if t.state != waitForAgent {
 		return
 	}
 	doConnect := func() {
 		t.wg.Add(1)
 		t.state = establishTailnet
-		go t.connectTailnet(update.id)
+		go t.connectTailnet(update.ID)
 	}
 	// consequence of ignoring updates if we are not waiting for the agent is that we MUST receive
 	// the start build succeeded update BEFORE we get the Agent connected / ready update.  We should keep this
 	// in mind when implementing the watch in Coderd.
-	switch update.lifecycle {
+	switch update.Lifecycle {
 	case codersdk.WorkspaceAgentLifecycleReady:
 		doConnect()
 		return
@@ -376,7 +426,7 @@ func (t *Tunneler) handleAgentUpdate(update *agentUpdate) {
 	default:
 		// unhittable, unless new states are added. We structure this with the switch and all cases covered to ensure
 		// we cover all cases.
-		t.config.DebugLogger.Critical(t.ctx, "unhandled agent update", slog.F("lifecycle", update.lifecycle))
+		t.config.DebugLogger.Critical(t.ctx, "unhandled agent update", slog.F("lifecycle", update.Lifecycle))
 	}
 }
 

codersdk/workspacesdk/tunneler/tunneler_internal_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
+	"github.com/coder/coder/v2/codersdk/wsjson"
 	"github.com/coder/coder/v2/testutil"
 )
 
@@ -28,7 +29,7 @@ func TestHandleBuildUpdate_Coverage(t *testing.T) {
 						t.Run(fmt.Sprintf("%d_%s_%s_%t_%t", s, trans, jobStatus, noAutostart, noWaitForScripts), func(t *testing.T) {
 							t.Parallel()
 							coverUpdate(t, workspaceID, noAutostart, noWaitForScripts, s, func(uut *Tunneler) {
-								uut.handleBuildUpdate(&buildUpdate{transition: trans, jobStatus: jobStatus})
+								uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: trans, JobStatus: jobStatus})
 							})
 						})
 					}
@@ -105,31 +106,31 @@ func TestBuildUpdatesStoppedWorkspace(t *testing.T) {
 		state:  stateInit,
 	}
 
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStop, jobStatus: codersdk.ProvisionerJobPending})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStop, JobStatus: codersdk.ProvisionerJobPending})
 	require.Equal(t, waitToStart, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 	require.False(t, fWorkspaceStarter.started)
 
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStop, jobStatus: codersdk.ProvisionerJobRunning})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStop, JobStatus: codersdk.ProvisionerJobRunning})
 	require.Equal(t, waitToStart, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 	require.False(t, fWorkspaceStarter.started)
 
 	// when stop job succeeds, we start the workspace
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStop, jobStatus: codersdk.ProvisionerJobSucceeded})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStop, JobStatus: codersdk.ProvisionerJobSucceeded})
 	require.Equal(t, waitForWorkspaceStarted, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 	require.True(t, fWorkspaceStarter.started)
 
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStart, jobStatus: codersdk.ProvisionerJobPending})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStart, JobStatus: codersdk.ProvisionerJobPending})
 	require.Equal(t, waitForWorkspaceStarted, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStart, jobStatus: codersdk.ProvisionerJobRunning})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStart, JobStatus: codersdk.ProvisionerJobRunning})
 	require.Equal(t, waitForWorkspaceStarted, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStart, jobStatus: codersdk.ProvisionerJobSucceeded})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStart, JobStatus: codersdk.ProvisionerJobSucceeded})
 	require.Equal(t, waitForAgent, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 }
@@ -157,7 +158,7 @@ func TestBuildUpdatesNewBuildWhileWaiting(t *testing.T) {
 	}
 
 	// New build comes in while we are waiting for the agent to start. We roll back to waiting for the workspace to start.
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStart, jobStatus: codersdk.ProvisionerJobRunning})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStart, JobStatus: codersdk.ProvisionerJobRunning})
 	require.Equal(t, waitForWorkspaceStarted, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 	require.False(t, fWorkspaceStarter.started)
@@ -193,12 +194,12 @@ func TestBuildUpdatesBadJobs(t *testing.T) {
 				state:  stateInit,
 			}
 
-			uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStart, jobStatus: codersdk.ProvisionerJobRunning})
+			uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStart, JobStatus: codersdk.ProvisionerJobRunning})
 			require.Equal(t, waitForWorkspaceStarted, uut.state)
 			waitForGoroutines(testCtx, t, uut)
 			require.False(t, fWorkspaceStarter.started)
 
-			uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStop, jobStatus: jobStatus})
+			uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStop, JobStatus: jobStatus})
 			require.Equal(t, exit, uut.state)
 			waitForGoroutines(testCtx, t, uut)
 			require.False(t, fWorkspaceStarter.started)
@@ -233,7 +234,7 @@ func TestBuildUpdatesNoAutostart(t *testing.T) {
 	}
 
 	// when stop job succeeds, we exit because autostart is disabled
-	uut.handleBuildUpdate(&buildUpdate{transition: codersdk.WorkspaceTransitionStop, jobStatus: codersdk.ProvisionerJobSucceeded})
+	uut.handleBuildUpdate(&workspacesdk.BuildUpdate{Transition: codersdk.WorkspaceTransitionStop, JobStatus: codersdk.ProvisionerJobSucceeded})
 	require.Equal(t, exit, uut.state)
 	waitForGoroutines(testCtx, t, uut)
 	require.False(t, fWorkspaceStarter.started)
@@ -254,7 +255,7 @@ func TestAgentUpdate_Coverage(t *testing.T) {
 					t.Run(fmt.Sprintf("%d_%s_%t_%t", s, lifecycle, noAutostart, noWaitForScripts), func(t *testing.T) {
 						t.Parallel()
 						coverUpdate(t, workspaceID, noAutostart, noWaitForScripts, s, func(uut *Tunneler) {
-							uut.handleAgentUpdate(&agentUpdate{lifecycle: lifecycle, id: agentID})
+							uut.handleAgentUpdate(&workspacesdk.AgentUpdate{Lifecycle: lifecycle, ID: agentID})
 						})
 					})
 				}
@@ -288,7 +289,7 @@ func TestAgentUpdateReady(t *testing.T) {
 		client: fClient,
 	}
 
-	uut.handleAgentUpdate(&agentUpdate{lifecycle: codersdk.WorkspaceAgentLifecycleReady, id: agentID})
+	uut.handleAgentUpdate(&workspacesdk.AgentUpdate{Lifecycle: codersdk.WorkspaceAgentLifecycleReady, ID: agentID})
 	require.Equal(t, establishTailnet, uut.state)
 	event := testutil.RequireReceive(testCtx, t, uut.events)
 	require.NotNil(t, event.tailnetUpdate)
@@ -323,7 +324,7 @@ func TestAgentUpdateNoWait(t *testing.T) {
 		client: fClient,
 	}
 
-	uut.handleAgentUpdate(&agentUpdate{lifecycle: codersdk.WorkspaceAgentLifecycleStarting, id: agentID})
+	uut.handleAgentUpdate(&workspacesdk.AgentUpdate{Lifecycle: codersdk.WorkspaceAgentLifecycleStarting, ID: agentID})
 	require.Equal(t, establishTailnet, uut.state)
 	event := testutil.RequireReceive(testCtx, t, uut.events)
 	require.NotNil(t, event.tailnetUpdate)
@@ -526,27 +527,27 @@ func TestTunneler_EventLoop_Signal(t *testing.T) {
 	go uut.eventLoop()
 
 	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
-		buildUpdate: &buildUpdate{
-			transition: codersdk.WorkspaceTransitionStart,
-			jobStatus:  codersdk.ProvisionerJobPending,
+		buildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransitionStart,
+			JobStatus:  codersdk.ProvisionerJobPending,
 		},
 	})
 	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
-		buildUpdate: &buildUpdate{
-			transition: codersdk.WorkspaceTransitionStart,
-			jobStatus:  codersdk.ProvisionerJobRunning,
+		buildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransitionStart,
+			JobStatus:  codersdk.ProvisionerJobRunning,
 		},
 	})
 	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
-		buildUpdate: &buildUpdate{
-			transition: codersdk.WorkspaceTransitionStart,
-			jobStatus:  codersdk.ProvisionerJobSucceeded,
+		buildUpdate: &workspacesdk.BuildUpdate{
+			Transition: codersdk.WorkspaceTransitionStart,
+			JobStatus:  codersdk.ProvisionerJobSucceeded,
 		},
 	})
 	testutil.RequireSend(testCtx, t, uut.events, tunnelerEvent{
-		agentUpdate: &agentUpdate{
-			lifecycle: codersdk.WorkspaceAgentLifecycleReady,
-			id:        agentID,
+		agentUpdate: &workspacesdk.AgentUpdate{
+			Lifecycle: codersdk.WorkspaceAgentLifecycleReady,
+			ID:        agentID,
 		},
 	})
 
@@ -658,6 +659,11 @@ type fakeClient struct {
 	dialed bool
 }
 
+func (*fakeClient) WorkspaceAgentConnectionWatch(context.Context, uuid.UUID, string) (dec *wsjson.Decoder[workspacesdk.ConnectionWatchEvent], err error) {
+	// TODO implement me
+	panic("implement me")
+}
+
 func (f *fakeClient) DialAgent(
 	_ context.Context, id uuid.UUID, _ *workspacesdk.DialAgentOptions,
 ) (

codersdk/workspacesdk/workspaceagentconnwatch.go

@@ -0,0 +1,86 @@
+package workspacesdk
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/wsjson"
+	"github.com/coder/websocket"
+)
+
+type WatchErrorCode int
+
+const (
+	_ WatchErrorCode = iota // Ensure that zero value is not a valid code
+	WatchErrorTooManyAgents
+	WatchErrorNameNotFound
+	WatchErrorNoAgents
+	WatchErrorServerShutdown
+	WatchErrorDatabase
+	WatchErrorInternal
+)
+
+type ConnectionWatchEvent struct {
+	Error       *WatchError  `json:"error"`
+	BuildUpdate *BuildUpdate `json:"build_update,omitempty"`
+	AgentUpdate *AgentUpdate `json:"agent_update,omitempty"`
+}
+
+type WatchError struct {
+	Code      WatchErrorCode `json:"code"`
+	Retryable bool           `json:"retryable"`
+	Message   string         `json:"message"`
+	Details   string         `json:"details,omitempty"`
+}
+
+func (e *WatchError) Error() string {
+	if e.Details != "" {
+		return fmt.Sprintf("%s: %s", e.Message, e.Details)
+	}
+	return e.Message
+}
+
+type BuildUpdate struct {
+	Transition codersdk.WorkspaceTransition  `json:"transition"`
+	JobStatus  codersdk.ProvisionerJobStatus `json:"job_status"`
+}
+
+type AgentUpdate struct {
+	Lifecycle codersdk.WorkspaceAgentLifecycle `json:"lifecycle"`
+	ID        uuid.UUID                        `json:"id" format:"uuid"`
+}
+
+func (c *Client) WorkspaceAgentConnectionWatch(
+	dialCtx context.Context, workspaceID uuid.UUID, agentName string,
+) (
+	dec *wsjson.Decoder[ConnectionWatchEvent], err error,
+) {
+	wsOptions := &websocket.DialOptions{
+		HTTPClient: c.client.HTTPClient,
+		// Need to disable compression to avoid a data-race.
+		CompressionMode: websocket.CompressionDisabled,
+	}
+	c.client.SessionTokenProvider.SetDialOption(wsOptions)
+
+	watchURL, err := c.client.URL.Parse(fmt.Sprintf("/api/v2/workspaces/%s/agent-connection-watch", workspaceID))
+	if err != nil {
+		return nil, xerrors.Errorf("parse url: %w", err)
+	}
+	if agentName != "" {
+		q := watchURL.Query()
+		q.Set("agent_name", agentName)
+		watchURL.RawQuery = q.Encode()
+	}
+
+	// nolint:bodyclose
+	conn, res, err := websocket.Dial(dialCtx, watchURL.String(), wsOptions)
+	if err != nil {
+		bodyErr := codersdk.ReadBodyAsError(res)
+		return nil, bodyErr
+	}
+	return wsjson.NewDecoder[ConnectionWatchEvent](conn, websocket.MessageText, c.client.Logger()), nil
+}