feat: add chat context pinning storage and push trigger (#26385)

Foundation for the Workspace Context Sources RFC (phase 3). The agent push (#25983) and coderd snapshot storage (#26145) already persist per-agent context snapshots; this PR lands the **chat-side storage** plus the **`agentapi` push trigger** that a follow-up will use to read them. It does **not** touch `chatd` and changes no behavior — nothing wires an implementation yet. ## What changed - Adds four nullable columns to `chats` — `context_aggregate_hash`, `context_dirty_since`, `context_dirty_resources`, and `context_error` — and rebuilds the `chats_expanded` view. - Adds three queries — `SetChatContextSnapshot`, `HydrateAgentChatsContext`, `MarkChatsContextDirtyByAgent` — with `dbauthz` wrappers and `audit` entries. They are store-interface methods covered by a Postgres test (`TestChatContextHydration`). - Adds the `agentapi.ContextDirtyMarker` interface and invokes it inside the `PushContextState` transaction, publishing collected events only after commit. ## Intentionally inert There are **no production callers** of the three queries and **no implementation** wired for `ContextDirtyMarker`, so the push trigger is dormant. This is deliberate: the PR is the durable storage/query foundation only. The actual integration — the `chatd` implementation that hydrates/dirties chats and backs a refresh endpoint, consuming the pinned context in prompt building, the rich SDK types + UI, and retiring the live per-turn pull — lands as a single follow-up PR. Splitting this way keeps the schema/query layer reviewable on its own and keeps the integration whole in one place. Refs #25983, #26145. <details> <summary>Decision log</summary> - **Columns over a side table.** The four `chats` columns are the durable model (accepting the one-time `chats_expanded` view/CTE churn). `last_injected_context` is deliberately left untouched — it is load-bearing for the live per-turn context pull. - **Keep `agentapi`, drop `chatd`.** The earlier revision wired the hydrate/dirty implementation through `chatd` and added a `PUT /chats/{chat}/context` refresh endpoint. Those were removed so this PR is pure foundation; `agentapi` defines the trigger + interface (it does not import `chatd`), and the `chatd` implementation arrives with the full integration. - **No new experiment flag.** The columns are dark and unread by prompt building. - **Authz.** The new query wrappers authorize chat updates under the chat RBAC object / `ResourceChat`, consistent with the existing system chat mutators. </details> --- 🤖 Generated by Coder Agents on behalf of @kylecarbs.
2026-09-24 15:04:27 +08:00 · 2026-06-15 14:41:00 -07:00
parent 89f200872b
commit 210261b143
18 changed files with 1116 additions and 140 deletions
@@ -75,12 +75,15 @@ type Options struct {
 	OrganizationID    uuid.UUID
 	TemplateVersionID uuid.UUID

-	AuthenticatedCtx                  context.Context
-	Log                               slog.Logger
-	Clock                             quartz.Clock
-	Database                          database.Store
-	NotificationsEnqueuer             notifications.Enqueuer
-	Pubsub                            pubsub.Pubsub
+	AuthenticatedCtx      context.Context
+	Log                   slog.Logger
+	Clock                 quartz.Clock
+	Database              database.Store
+	NotificationsEnqueuer notifications.Enqueuer
+	Pubsub                pubsub.Pubsub
+	// ContextDirtyMarker is the chatd-backed hydrate/dirty fan-out invoked
+	// from PushContextState. Nil when chatd is disabled.
+	ContextDirtyMarker                ContextDirtyMarker
 	ConnectionLogger                  *atomic.Pointer[connectionlog.ConnectionLogger]
 	DerpMapFn                         func() *tailcfg.DERPMap
 	TailnetCoordinator                *atomic.Pointer[tailnet.Coordinator]
@@ -248,11 +251,12 @@ func New(opts Options, workspace database.Workspace, agent database.WorkspaceAge
 	}

 	api.ContextAPI = &ContextAPI{
-		AgentID:   agent.ID,
-		Workspace: api.cachedWorkspaceFields,
-		Log:       opts.Log,
-		Clock:     opts.Clock,
-		Database:  opts.Database,
+		AgentID:     agent.ID,
+		Workspace:   api.cachedWorkspaceFields,
+		Log:         opts.Log,
+		Clock:       opts.Clock,
+		Database:    opts.Database,
+		DirtyMarker: opts.ContextDirtyMarker,
 	}

 	// Start background cache refresh loop to handle workspace changes
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"math"
 	"sort"
+	"time"

 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -62,6 +63,25 @@ type ContextAPI struct {
 	Log       slog.Logger
 	Clock     quartz.Clock
 	Database  database.Store
+	// DirtyMarker hydrates chats from, and marks chats dirty against, the
+	// snapshot persisted by a push. It is nil when chatd is not running,
+	// in which case PushContextState stays a pure write path.
+	DirtyMarker ContextDirtyMarker
+}
+
+// ContextDirtyMarker hydrates chats from, and marks chats dirty against, a
+// freshly persisted agent context snapshot. It is implemented by chatd and
+// injected at coderd construction so this package neither imports the chat
+// domain nor performs chat-authorized writes directly.
+type ContextDirtyMarker interface {
+	// HydrateAndMarkChatsDirty runs inside the PushContextState
+	// transaction using the supplied store. It hydrates chats for the
+	// agent that have no pinned hash yet (no dirty event) and flips
+	// already-pinned chats whose hash differs from aggregateHash. It
+	// returns a callback that publishes the resulting dirty watch events;
+	// the caller invokes it only after the transaction commits. The
+	// callback is nil when nothing transitioned to dirty.
+	HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store, agentID uuid.UUID, aggregateHash []byte, snapshotError string, now time.Time) (publishDirty func(), err error)
 }

 // PushContextState persists a snapshot pushed by the workspace
@@ -120,10 +140,15 @@ func (a *ContextAPI) PushContextState(ctx context.Context, req *agentproto.PushC
 	sort.Strings(activeSources)

 	var accepted bool
+	// publishDirty is captured from the final (committed) attempt and
+	// invoked after the transaction commits; ReadModifyUpdate may re-run
+	// the closure on serialization conflicts.
+	var publishDirty func()
 	err = database.ReadModifyUpdate(a.Database, func(tx database.Store) error {
 		// The closure re-runs on serialization conflicts; reset any
 		// state carried over from a rolled-back attempt.
 		accepted = false
+		publishDirty = nil

 		existing, err := tx.GetLatestWorkspaceAgentContextSnapshot(ctx, a.AgentID)
 		switch {
@@ -171,6 +196,16 @@ func (a *ContextAPI) PushContextState(ctx context.Context, req *agentproto.PushC
 			return xerrors.Errorf("delete stale resources: %w", err)
 		}

+		// Hydrate and dirty chats against the snapshot just written, in the
+		// same transaction so a concurrent refresh cannot interleave with
+		// the version gate. Events are published only after commit.
+		if a.DirtyMarker != nil {
+			publishDirty, err = a.DirtyMarker.HydrateAndMarkChatsDirty(ctx, tx, a.AgentID, req.AggregateHash, req.SnapshotError, now)
+			if err != nil {
+				return xerrors.Errorf("hydrate and mark chats dirty: %w", err)
+			}
+		}
+
 		accepted = true
 		return nil
 	})
@@ -187,6 +222,12 @@ func (a *ContextAPI) PushContextState(ctx context.Context, req *agentproto.PushC
 		return &agentproto.PushContextStateResponse{Accepted: false}, nil
 	}

+	// The snapshot committed; fan out dirty watch events to chats whose
+	// pinned context drifted from this push.
+	if publishDirty != nil {
+		publishDirty()
+	}
+
 	a.Log.Debug(ctx, "PushContextState accepted",
 		slog.F("agent_id", a.AgentID),
 		slog.F("version", req.Version),
@@ -88,6 +88,70 @@ func TestPushContextState(t *testing.T) {
 		require.True(t, resp.GetAccepted())
 	})

+	t.Run("DirtyMarkerInvokedAfterCommit", func(t *testing.T) {
+		t.Parallel()
+
+		api, dbm := makeAPI(t)
+		marker := &fakeDirtyMarker{}
+		api.DirtyMarker = marker
+		expectInTx(dbm)
+
+		dbm.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), agentID).
+			Return(database.WorkspaceAgentContextSnapshot{}, errNoRows())
+		dbm.EXPECT().UpsertWorkspaceAgentContextSnapshot(gomock.Any(), gomock.Any()).
+			Return(database.WorkspaceAgentContextSnapshot{}, nil)
+		dbm.EXPECT().UpsertWorkspaceAgentContextResource(gomock.Any(), gomock.Any()).
+			Return(database.WorkspaceAgentContextResource{}, nil).Times(1)
+		dbm.EXPECT().DeleteStaleWorkspaceAgentContextResources(gomock.Any(), gomock.Any()).
+			Return(nil)
+
+		resp, err := api.PushContextState(context.Background(), &agentproto.PushContextStateRequest{
+			Version:       1,
+			AggregateHash: []byte{0xaa, 0xbb},
+			SnapshotError: "watcher degraded",
+			Initial:       true,
+			Resources: []*agentproto.ContextResource{
+				instructionResource("/home/coder/AGENTS.md", "hello"),
+			},
+		})
+		require.NoError(t, err)
+		require.True(t, resp.GetAccepted())
+		// The marker runs inside the push transaction and its returned
+		// callback publishes only after the transaction commits.
+		require.Equal(t, 1, marker.called)
+		require.Equal(t, 1, marker.published)
+		require.Equal(t, agentID, marker.gotAgent)
+		require.Equal(t, []byte{0xaa, 0xbb}, marker.gotHash)
+		require.Equal(t, "watcher degraded", marker.gotErr)
+	})
+
+	t.Run("DirtyMarkerSkippedOnDrop", func(t *testing.T) {
+		t.Parallel()
+
+		api, dbm := makeAPI(t)
+		marker := &fakeDirtyMarker{}
+		api.DirtyMarker = marker
+		expectInTx(dbm)
+
+		// A non-initial push at a version not strictly greater than the
+		// stored one is dropped before any write; hydration and the
+		// dirty fan-out must not run.
+		dbm.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), agentID).
+			Return(database.WorkspaceAgentContextSnapshot{Version: 5}, nil)
+
+		resp, err := api.PushContextState(context.Background(), &agentproto.PushContextStateRequest{
+			Version:       2,
+			AggregateHash: []byte{0x01},
+			Resources: []*agentproto.ContextResource{
+				instructionResource("/home/coder/AGENTS.md", "hello"),
+			},
+		})
+		require.NoError(t, err)
+		require.False(t, resp.GetAccepted())
+		require.Equal(t, 0, marker.called)
+		require.Equal(t, 0, marker.published)
+	})
+
 	t.Run("RejectsEmptyAndDuplicateSources", func(t *testing.T) {
 		t.Parallel()

@@ -598,3 +662,22 @@ func mcpServerResource(source, serverName, description string) *agentproto.Conte
 		},
 	}
 }
+
+// fakeDirtyMarker is a test double for agentapi.ContextDirtyMarker. It records
+// the in-transaction call and counts callback invocations so tests can assert
+// the marker runs inside the push transaction and publishes only after commit.
+type fakeDirtyMarker struct {
+	called    int
+	published int
+	gotAgent  uuid.UUID
+	gotHash   []byte
+	gotErr    string
+}
+
+func (f *fakeDirtyMarker) HydrateAndMarkChatsDirty(_ context.Context, _ database.Store, agentID uuid.UUID, aggregateHash []byte, snapshotError string, _ time.Time) (func(), error) {
+	f.called++
+	f.gotAgent = agentID
+	f.gotHash = aggregateHash
+	f.gotErr = snapshotError
+	return func() { f.published++ }, nil
+}