Skip to content

Commit de7a285

Browse files
sreyasreya
authored
fix: add flag to disable workspace agent context sync (cherry-pick #28522) (#28525)
Cherry-pick of #28522 (`26de6140fb`) onto `release/2.35`. Adds `CODER_DISABLE_WORKSPACE_AGENT_CONTEXT_SYNC` / `--disable-workspace-agent-context-sync`. When set, `PushContextState` rejects agent context pushes with a dRPC `Unimplemented` code before any validation or database work; deployed agents translate that into `ErrPushUnimplemented` and stop their push loop for the life of the connection. This gives large deployments a server-only kill switch for context sync database write load, with no agent updates or workspace restarts required. Two conflicts resolved relative to the original commit: - `coderd/workspaceagentsrpc.go`: `release/2.35` assigns `ContextDirtyMarker: api.chatDaemon` directly (the nil-guarded local was introduced later on `main`); kept the branch's form and added `ContextSyncDisabled` alongside it. - `docs/admin/setup/configuration-reference.md`: dropped; this generated doc does not exist on `release/2.35`. Validated on this branch: `go build`, `TestPushContextState` (unit), `TestWorkspaceAgentPushContextState*` (end-to-end over real dRPC), CLI and enterprise golden-file tests, full pre-commit hooks.
1 parent 1e68cef commit de7a285

15 files changed

Lines changed: 146 additions & 2 deletions

File tree

‎cli/testdata/coder_server_--help.golden‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,13 @@ OPTIONS:
5555
the workspace serves malicious JavaScript. This is recommended for
5656
security purposes if a --wildcard-access-url is configured.
5757

58+
--disable-workspace-agent-context-sync bool, $CODER_DISABLE_WORKSPACE_AGENT_CONTEXT_SYNC
59+
Stop persisting workspace agent context snapshots (instructions,
60+
skills, and MCP state used for pinned chat context). When set, coderd
61+
rejects agent context pushes as unimplemented and agents stop sending
62+
them; chats cannot pin workspace context. Use this to shed the
63+
database write load of context sync on large deployments.
64+
5865
--disable-workspace-sharing bool, $CODER_DISABLE_WORKSPACE_SHARING
5966
Disable workspace sharing. Workspace ACL checking is disabled and only
6067
owners can have ssh, apps and terminal access to workspaces. Access

‎cli/testdata/server-config.yaml.golden‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -559,6 +559,13 @@ disableWorkspaceSharing: false
559559
# their chats.
560560
# (default: <unset>, type: bool)
561561
disableChatSharing: false
562+
# Stop persisting workspace agent context snapshots (instructions, skills, and MCP
563+
# state used for pinned chat context). When set, coderd rejects agent context
564+
# pushes as unimplemented and agents stop sending them; chats cannot pin workspace
565+
# context. Use this to shed the database write load of context sync on large
566+
# deployments.
567+
# (default: <unset>, type: bool)
568+
disableWorkspaceAgentContextSync: false
562569
# These options change the behavior of how clients interact with the Coder.
563570
# Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI.
564571
client:

‎coderd/agentapi/api.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,10 @@ type Options struct {
8383
Pubsub pubsub.Pubsub
8484
// ContextDirtyMarker is the chatd-backed hydrate/dirty fan-out invoked
8585
// from PushContextState. Nil when chatd is disabled.
86-
ContextDirtyMarker ContextDirtyMarker
86+
ContextDirtyMarker ContextDirtyMarker
87+
// ContextSyncDisabled makes PushContextState reject pushes with a dRPC
88+
// Unimplemented code so agents stop sending context snapshots.
89+
ContextSyncDisabled bool
8790
ConnectionLogger *atomic.Pointer[connectionlog.ConnectionLogger]
8891
DerpMapFn func() *tailcfg.DERPMap
8992
TailnetCoordinator *atomic.Pointer[tailnet.Coordinator]
@@ -257,6 +260,7 @@ func New(opts Options, workspace database.Workspace, agent database.WorkspaceAge
257260
Clock: opts.Clock,
258261
Database: opts.Database,
259262
DirtyMarker: opts.ContextDirtyMarker,
263+
Disabled: opts.ContextSyncDisabled,
260264
}
261265

262266
// Start background cache refresh loop to handle workspace changes

‎coderd/agentapi/context.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"golang.org/x/xerrors"
1313
"google.golang.org/protobuf/encoding/protojson"
1414
"google.golang.org/protobuf/proto"
15+
"storj.io/drpc/drpcerr"
1516

1617
"cdr.dev/slog/v3"
1718
agentproto "github.com/coder/coder/v2/agent/proto"
@@ -67,6 +68,13 @@ type ContextAPI struct {
6768
// snapshot persisted by a push. It is nil when chatd is not running,
6869
// in which case PushContextState stays a pure write path.
6970
DirtyMarker ContextDirtyMarker
71+
// Disabled rejects every push with a dRPC Unimplemented code. The
72+
// agent's DRPCPusher translates that code into ErrPushUnimplemented,
73+
// which terminates its RunPush loop for the life of the connection,
74+
// exactly as if coderd predated the v2.10 Agent API. This is the
75+
// deployment-wide kill switch for context sync write load
76+
// (CODER_DISABLE_WORKSPACE_AGENT_CONTEXT_SYNC).
77+
Disabled bool
7078
}
7179

7280
// ContextDirtyMarker hydrates chats from, and marks chats dirty against, a
@@ -103,6 +111,14 @@ type ContextDirtyMarker interface {
103111
// authorizes the actor (the agent's token subject) against the
104112
// workspace that owns the agent.
105113
func (a *ContextAPI) PushContextState(ctx context.Context, req *agentproto.PushContextStateRequest) (*agentproto.PushContextStateResponse, error) {
114+
if a.Disabled {
115+
// The Unimplemented code (not a plain error) is what tells the
116+
// agent to stop pushing instead of retrying with backoff.
117+
return nil, drpcerr.WithCode(
118+
xerrors.New("agentapi: workspace agent context sync is disabled on this deployment"),
119+
drpcerr.Unimplemented,
120+
)
121+
}
106122
if req == nil {
107123
return nil, xerrors.New("agentapi: PushContextState request is nil")
108124
}

‎coderd/agentapi/context_test.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"github.com/lib/pq"
1414
"github.com/stretchr/testify/require"
1515
"go.uber.org/mock/gomock"
16+
"storj.io/drpc/drpcerr"
1617

1718
"cdr.dev/slog/v3"
1819
"cdr.dev/slog/v3/sloggers/slogtest"
@@ -58,6 +59,27 @@ func TestPushContextState(t *testing.T) {
5859
)
5960
}
6061

62+
t.Run("DisabledReturnsUnimplemented", func(t *testing.T) {
63+
t.Parallel()
64+
65+
// No InTx or query expectations: a disabled push must return
66+
// before touching the store. The Unimplemented dRPC code is
67+
// load-bearing; the agent's DRPCPusher translates it into
68+
// ErrPushUnimplemented, which stops its RunPush loop instead
69+
// of retrying with backoff.
70+
api, _ := makeAPI(t)
71+
api.Disabled = true
72+
73+
resp, err := api.PushContextState(context.Background(), &agentproto.PushContextStateRequest{
74+
Version: 1,
75+
AggregateHash: []byte{0x01, 0x02},
76+
Initial: true,
77+
})
78+
require.Error(t, err)
79+
require.Nil(t, resp)
80+
require.EqualValues(t, drpcerr.Unimplemented, drpcerr.Code(err))
81+
})
82+
6183
t.Run("AcceptsInitialPush", func(t *testing.T) {
6284
t.Parallel()
6385

‎coderd/apidoc/docs.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/apidoc/swagger.json‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/workspaceagents_test.go‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import (
3535
"github.com/coder/coder/v2/agent/agentcontainers"
3636
"github.com/coder/coder/v2/agent/agentcontainers/acmock"
3737
"github.com/coder/coder/v2/agent/agentcontainers/watcher"
38+
"github.com/coder/coder/v2/agent/agentcontext"
3839
"github.com/coder/coder/v2/agent/agenttest"
3940
agentproto "github.com/coder/coder/v2/agent/proto"
4041
"github.com/coder/coder/v2/coderd/agentapi/metadatabatcher"
@@ -3240,6 +3241,54 @@ func TestWorkspaceAgentPushContextState(t *testing.T) {
32403241
require.False(t, resp.GetAccepted())
32413242
}
32423243

3244+
// TestWorkspaceAgentPushContextStateDisabled verifies the
3245+
// --disable-workspace-agent-context-sync kill switch end to end over a
3246+
// real dRPC connection: the handler's Unimplemented code must survive
3247+
// the transport and be translated by the agent's DRPCPusher into
3248+
// ErrPushUnimplemented, which is what terminates the agent's RunPush
3249+
// loop instead of retrying with backoff. Nothing may be persisted.
3250+
func TestWorkspaceAgentPushContextStateDisabled(t *testing.T) {
3251+
t.Parallel()
3252+
3253+
dv := coderdtest.DeploymentValues(t)
3254+
dv.DisableWorkspaceAgentContextSync = true
3255+
client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
3256+
DeploymentValues: dv,
3257+
})
3258+
user := coderdtest.CreateFirstUser(t, client)
3259+
r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
3260+
OrganizationID: user.OrganizationID,
3261+
OwnerID: user.UserID,
3262+
}).WithAgent().Do()
3263+
require.Len(t, r.Agents, 1)
3264+
agentID := r.Agents[0].ID
3265+
3266+
ctx := testutil.Context(t, testutil.WaitLong)
3267+
3268+
agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
3269+
aAPI, _, err := agentClient.ConnectRPC210(ctx)
3270+
require.NoError(t, err)
3271+
defer func() {
3272+
cErr := aAPI.DRPCConn().Close()
3273+
require.NoError(t, cErr)
3274+
}()
3275+
3276+
// Push through the same adapter the agent's RunPush loop uses so
3277+
// the test breaks if either side of the Unimplemented contract
3278+
// changes.
3279+
pusher := agentcontext.NewDRPCPusher(aAPI)
3280+
resp, err := pusher.PushContextState(ctx, &agentcontext.PushRequest{
3281+
Version: 1,
3282+
Initial: true,
3283+
})
3284+
require.ErrorIs(t, err, agentcontext.ErrPushUnimplemented)
3285+
require.Nil(t, resp)
3286+
3287+
// The rejected push must not have persisted anything.
3288+
_, err = db.GetLatestWorkspaceAgentContextSnapshot(dbauthz.AsSystemRestricted(ctx), agentID) //nolint:gocritic // Test assertions read agent-pushed rows directly from the store.
3289+
require.ErrorIs(t, err, sql.ErrNoRows)
3290+
}
3291+
32433292
func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest {
32443293
mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
32453294
require.NoError(t, err)

‎coderd/workspaceagentsrpc.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,8 @@ func (api *API) workspaceAgentRPC(rw http.ResponseWriter, r *http.Request) {
182182
UpdateAgentMetricsFn: api.UpdateAgentMetrics,
183183
// chatDaemon is always constructed (only its worker is gated), so
184184
// this is non-nil; agentapi treats a nil marker as "chatd absent".
185-
ContextDirtyMarker: api.chatDaemon,
185+
ContextDirtyMarker: api.chatDaemon,
186+
ContextSyncDisabled: api.DeploymentValues.DisableWorkspaceAgentContextSync.Value(),
186187
}, workspace, workspaceAgent)
187188

188189
streamID := tailnet.StreamID{

‎codersdk/deployment.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -670,6 +670,7 @@ type DeploymentValues struct {
670670
DisableOwnerWorkspaceExec serpent.Bool `json:"disable_owner_workspace_exec,omitempty" typescript:",notnull"`
671671
DisableWorkspaceSharing serpent.Bool `json:"disable_workspace_sharing,omitempty" typescript:",notnull"`
672672
DisableChatSharing serpent.Bool `json:"disable_chat_sharing,omitempty" typescript:",notnull"`
673+
DisableWorkspaceAgentContextSync serpent.Bool `json:"disable_workspace_agent_context_sync,omitempty" typescript:",notnull"`
673674
ProxyHealthStatusInterval serpent.Duration `json:"proxy_health_status_interval,omitempty" typescript:",notnull"`
674675
EnableTerraformDebugMode serpent.Bool `json:"enable_terraform_debug_mode,omitempty" typescript:",notnull"`
675676
UserQuietHoursSchedule UserQuietHoursScheduleConfig `json:"user_quiet_hours_schedule,omitempty" typescript:",notnull"`
@@ -3680,6 +3681,15 @@ communicating directly.`,
36803681
Value: &c.DisableChatSharing,
36813682
YAML: "disableChatSharing",
36823683
},
3684+
{
3685+
Name: "Disable Workspace Agent Context Sync",
3686+
Description: "Stop persisting workspace agent context snapshots (instructions, skills, and MCP state used for pinned chat context). When set, coderd rejects agent context pushes as unimplemented and agents stop sending them; chats cannot pin workspace context. Use this to shed the database write load of context sync on large deployments.",
3687+
Flag: "disable-workspace-agent-context-sync",
3688+
Env: "CODER_DISABLE_WORKSPACE_AGENT_CONTEXT_SYNC",
3689+
3690+
Value: &c.DisableWorkspaceAgentContextSync,
3691+
YAML: "disableWorkspaceAgentContextSync",
3692+
},
36833693
{
36843694
Name: "Session Duration",
36853695
Description: "The token expiry duration for browser sessions. Sessions may last longer if they are actively making requests, but this functionality can be disabled via --disable-session-expiry-refresh.",

0 commit comments

Comments
 (0)