feat: add flag to disable workspace agent context sync (#28522)

coderd/agentapi/api.go

@@ -83,7 +83,10 @@ type Options struct {
 	Pubsub                pubsub.Pubsub
 	// ContextDirtyMarker is the chatd-backed hydrate/dirty fan-out invoked
 	// from PushContextState. Nil when chatd is disabled.
-	ContextDirtyMarker                ContextDirtyMarker
+	ContextDirtyMarker ContextDirtyMarker
+	// ContextSyncDisabled makes PushContextState reject pushes with a dRPC
+	// Unimplemented code so agents stop sending context snapshots.
+	ContextSyncDisabled               bool
 	ConnectionLogger                  *atomic.Pointer[connectionlog.ConnectionLogger]
 	DerpMapFn                         func() *tailcfg.DERPMap
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
@@ -257,6 +260,7 @@ func New(opts Options, workspace database.Workspace, agent database.WorkspaceAge
 		Clock:       opts.Clock,
 		Database:    opts.Database,
 		DirtyMarker: opts.ContextDirtyMarker,
+		Disabled:    opts.ContextSyncDisabled,
 	}
 
 	// Start background cache refresh loop to handle workspace changes

coderd/agentapi/context.go

@@ -12,6 +12,7 @@ import (
 	"golang.org/x/xerrors"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
+	"storj.io/drpc/drpcerr"
 
 	"cdr.dev/slog/v3"
 	agentproto "github.com/coder/coder/v2/agent/proto"
@@ -67,6 +68,13 @@ type ContextAPI struct {
 	// snapshot persisted by a push. It is nil when chatd is not running,
 	// in which case PushContextState stays a pure write path.
 	DirtyMarker ContextDirtyMarker
+	// Disabled rejects every push with a dRPC Unimplemented code. The
+	// agent's DRPCPusher translates that code into ErrPushUnimplemented,
+	// which terminates its RunPush loop for the life of the connection,
+	// exactly as if coderd predated the v2.10 Agent API. This is the
+	// deployment-wide kill switch for context sync write load
+	// (CODER_DISABLE_WORKSPACE_AGENT_CONTEXT_SYNC).
+	Disabled bool
 }
 
 // ContextDirtyMarker hydrates chats from, and marks chats dirty against, a
@@ -103,6 +111,14 @@ type ContextDirtyMarker interface {
 // authorizes the actor (the agent's token subject) against the
 // workspace that owns the agent.
 func (a *ContextAPI) PushContextState(ctx context.Context, req *agentproto.PushContextStateRequest) (*agentproto.PushContextStateResponse, error) {
+	if a.Disabled {
+		// The Unimplemented code (not a plain error) is what tells the
+		// agent to stop pushing instead of retrying with backoff.
+		return nil, drpcerr.WithCode(
+			xerrors.New("agentapi: workspace agent context sync is disabled on this deployment"),
+			drpcerr.Unimplemented,
+		)
+	}
 	if req == nil {
 		return nil, xerrors.New("agentapi: PushContextState request is nil")
 	}

coderd/agentapi/context_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
+	"storj.io/drpc/drpcerr"
 
 	"cdr.dev/slog/v3"
 	"cdr.dev/slog/v3/sloggers/slogtest"
@@ -58,6 +59,27 @@ func TestPushContextState(t *testing.T) {
 		)
 	}
 
+	t.Run("DisabledReturnsUnimplemented", func(t *testing.T) {
+		t.Parallel()
+
+		// No InTx or query expectations: a disabled push must return
+		// before touching the store. The Unimplemented dRPC code is
+		// load-bearing; the agent's DRPCPusher translates it into
+		// ErrPushUnimplemented, which stops its RunPush loop instead
+		// of retrying with backoff.
+		api, _ := makeAPI(t)
+		api.Disabled = true
+
+		resp, err := api.PushContextState(context.Background(), &agentproto.PushContextStateRequest{
+			Version:       1,
+			AggregateHash: []byte{0x01, 0x02},
+			Initial:       true,
+		})
+		require.Error(t, err)
+		require.Nil(t, resp)
+		require.EqualValues(t, drpcerr.Unimplemented, drpcerr.Code(err))
+	})
+
 	t.Run("AcceptsInitialPush", func(t *testing.T) {
 		t.Parallel()
 

coderd/apidoc/docs.go

@@ -21719,6 +21719,9 @@ const docTemplate = `{
                 "disable_path_apps": {
                     "type": "boolean"
                 },
+                "disable_workspace_agent_context_sync": {
+                    "type": "boolean"
+                },
                 "disable_workspace_sharing": {
                     "type": "boolean"
                 },

coderd/apidoc/swagger.json

@@ -19758,6 +19758,9 @@
 				"disable_path_apps": {
 					"type": "boolean"
 				},
+				"disable_workspace_agent_context_sync": {
+					"type": "boolean"
+				},
 				"disable_workspace_sharing": {
 					"type": "boolean"
 				},

coderd/workspaceagents_test.go

@@ -38,6 +38,7 @@ import (
 	"github.com/coder/coder/v2/agent/agentcontainers"
 	"github.com/coder/coder/v2/agent/agentcontainers/acmock"
 	"github.com/coder/coder/v2/agent/agentcontainers/watcher"
+	"github.com/coder/coder/v2/agent/agentcontext"
 	"github.com/coder/coder/v2/agent/agenttest"
 	agentproto "github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/coderd/agentapi/metadatabatcher"
@@ -3374,6 +3375,54 @@ func TestWorkspaceAgentPushContextState(t *testing.T) {
 	require.False(t, resp.GetAccepted())
 }
 
+// TestWorkspaceAgentPushContextStateDisabled verifies the
+// --disable-workspace-agent-context-sync kill switch end to end over a
+// real dRPC connection: the handler's Unimplemented code must survive
+// the transport and be translated by the agent's DRPCPusher into
+// ErrPushUnimplemented, which is what terminates the agent's RunPush
+// loop instead of retrying with backoff. Nothing may be persisted.
+func TestWorkspaceAgentPushContextStateDisabled(t *testing.T) {
+	t.Parallel()
+
+	dv := coderdtest.DeploymentValues(t)
+	dv.DisableWorkspaceAgentContextSync = true
+	client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
+		DeploymentValues: dv,
+	})
+	user := coderdtest.CreateFirstUser(t, client)
+	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OrganizationID: user.OrganizationID,
+		OwnerID:        user.UserID,
+	}).WithAgent().Do()
+	require.Len(t, r.Agents, 1)
+	agentID := r.Agents[0].ID
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
+	aAPI, _, err := agentClient.ConnectRPC210(ctx)
+	require.NoError(t, err)
+	defer func() {
+		cErr := aAPI.DRPCConn().Close()
+		require.NoError(t, cErr)
+	}()
+
+	// Push through the same adapter the agent's RunPush loop uses so
+	// the test breaks if either side of the Unimplemented contract
+	// changes.
+	pusher := agentcontext.NewDRPCPusher(aAPI)
+	resp, err := pusher.PushContextState(ctx, &agentcontext.PushRequest{
+		Version: 1,
+		Initial: true,
+	})
+	require.ErrorIs(t, err, agentcontext.ErrPushUnimplemented)
+	require.Nil(t, resp)
+
+	// The rejected push must not have persisted anything.
+	_, err = db.GetLatestWorkspaceAgentContextSnapshot(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // Test assertions read agent-pushed rows directly from the store.
+	require.ErrorIs(t, err, sql.ErrNoRows)
+}
+
 func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest {
 	mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
 	require.NoError(t, err)

coderd/workspaceagentsrpc.go

@@ -192,6 +192,7 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
 		// Optional:
 		UpdateAgentMetricsFn: api.UpdateAgentMetrics,
 		ContextDirtyMarker:   contextDirtyMarker,
+		ContextSyncDisabled:  api.DeploymentValues.DisableWorkspaceAgentContextSync.Value(),
 	}, workspace, workspaceAgent)
 
 	streamID := tailnet.StreamID{