feat: reinitialize agents when a prebuilt workspace is claimed (#17475)

This pull request allows coder workspace agents to be reinitialized when a prebuilt workspace is claimed by a user. This facilitates the transfer of ownership between the anonymous prebuilds system user and the new owner of the workspace. Only a single agent per prebuilt workspace is supported for now, but plumbing has already been done to facilitate the seamless transition to multi-agent support. --------- Signed-off-by: Danny Kopping <dannykopping@gmail.com> Co-authored-by: Danny Kopping <dannykopping@gmail.com>
2025-07-12 00:14:10 +00:00 · 2025-05-14 14:15:36 +02:00
parent fcbdd1a28e
commit 425ee6fa55
38 changed files with 2184 additions and 449 deletions
--- a/codersdk/agentsdk/agentsdk_test.go
+++ b/codersdk/agentsdk/agentsdk_test.go
@ -0,0 +1,122 @@
+package agentsdk_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestStreamAgentReinitEvents(t *testing.T) {
+	t.Parallel()
+
+	t.Run("transmitted events are received", func(t *testing.T) {
+		t.Parallel()
+
+		eventToSend := agentsdk.ReinitializationEvent{
+			WorkspaceID: uuid.New(),
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
+
+		events := make(chan agentsdk.ReinitializationEvent, 1)
+		events <- eventToSend
+
+		transmitCtx := testutil.Context(t, testutil.WaitShort)
+		transmitErrCh := make(chan error, 1)
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			transmitter := agentsdk.NewSSEAgentReinitTransmitter(slogtest.Make(t, nil), w, r)
+			transmitErrCh <- transmitter.Transmit(transmitCtx, events)
+		}))
+		defer srv.Close()
+
+		requestCtx := testutil.Context(t, testutil.WaitShort)
+		req, err := http.NewRequestWithContext(requestCtx, "GET", srv.URL, nil)
+		require.NoError(t, err)
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		receiveCtx := testutil.Context(t, testutil.WaitShort)
+		receiver := agentsdk.NewSSEAgentReinitReceiver(resp.Body)
+		sentEvent, receiveErr := receiver.Receive(receiveCtx)
+		require.Nil(t, receiveErr)
+		require.Equal(t, eventToSend, *sentEvent)
+	})
+
+	t.Run("doesn't transmit events if the transmitter context is canceled", func(t *testing.T) {
+		t.Parallel()
+
+		eventToSend := agentsdk.ReinitializationEvent{
+			WorkspaceID: uuid.New(),
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
+
+		events := make(chan agentsdk.ReinitializationEvent, 1)
+		events <- eventToSend
+
+		transmitCtx, cancelTransmit := context.WithCancel(testutil.Context(t, testutil.WaitShort))
+		cancelTransmit()
+		transmitErrCh := make(chan error, 1)
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			transmitter := agentsdk.NewSSEAgentReinitTransmitter(slogtest.Make(t, nil), w, r)
+			transmitErrCh <- transmitter.Transmit(transmitCtx, events)
+		}))
+
+		defer srv.Close()
+
+		requestCtx := testutil.Context(t, testutil.WaitShort)
+		req, err := http.NewRequestWithContext(requestCtx, "GET", srv.URL, nil)
+		require.NoError(t, err)
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		receiveCtx := testutil.Context(t, testutil.WaitShort)
+		receiver := agentsdk.NewSSEAgentReinitReceiver(resp.Body)
+		sentEvent, receiveErr := receiver.Receive(receiveCtx)
+		require.Nil(t, sentEvent)
+		require.ErrorIs(t, receiveErr, io.EOF)
+	})
+
+	t.Run("does not receive events if the receiver context is canceled", func(t *testing.T) {
+		t.Parallel()
+
+		eventToSend := agentsdk.ReinitializationEvent{
+			WorkspaceID: uuid.New(),
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
+
+		events := make(chan agentsdk.ReinitializationEvent, 1)
+		events <- eventToSend
+
+		transmitCtx := testutil.Context(t, testutil.WaitShort)
+		transmitErrCh := make(chan error, 1)
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			transmitter := agentsdk.NewSSEAgentReinitTransmitter(slogtest.Make(t, nil), w, r)
+			transmitErrCh <- transmitter.Transmit(transmitCtx, events)
+		}))
+		defer srv.Close()
+
+		requestCtx := testutil.Context(t, testutil.WaitShort)
+		req, err := http.NewRequestWithContext(requestCtx, "GET", srv.URL, nil)
+		require.NoError(t, err)
+		resp, err := http.DefaultClient.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		receiveCtx, cancelReceive := context.WithCancel(context.Background())
+		cancelReceive()
+		receiver := agentsdk.NewSSEAgentReinitReceiver(resp.Body)
+		sentEvent, receiveErr := receiver.Receive(receiveCtx)
+		require.Nil(t, sentEvent)
+		require.ErrorIs(t, receiveErr, context.Canceled)
+	})
+}