feat: integrate agent context snapshots into chats (#26389)

Makes the chat context foundation from #26385 live. That PR added the
storage columns, writer queries, and a dormant
`agentapi.ContextDirtyMarker` trigger with no production callers; this
PR wires them together end to end.

When a workspace agent pushes a context snapshot, bound chats now
hydrate to that snapshot's hash, and a later push with a different hash
flips already-pinned chats to dirty (emitting a `context_dirty` watch
event after the transaction commits). Chat creation pins the agent's
latest snapshot when one already exists. The experimental chat API
exposes this as `Chat.Context` (`*ChatContext` with `dirty`,
`dirty_since`, `error`), and a new `PUT
/api/experimental/chats/{chat}/context` endpoint re-pins the agent's
latest snapshot and clears the dirty marker.

`context_dirty_resources` stays NULL (the resource-level diff is
deferred to the UI phase) and the live per-turn context pull is
unchanged.

The end-to-end test provisions a workspace agent via the echo
provisioner, connects it over the Agent API v2.10, and exercises the
full path: an initial push hydrates a bound chat (clean), a second push
with a different hash marks it dirty, the API reports the dirty state,
and the refresh endpoint clears it.

<details>
<summary>Decision log</summary>

- **API shape — sub-struct.** Dirty state is surfaced as
`codersdk.Chat.Context *ChatContext { Dirty bool; DirtySince *time.Time;
Error string }` rather than flat fields, matching the RFC's named
`ChatContext` type and leaving room for future fields (resource diff,
sources). `db2sdk.Chat` populates it when the chat is context-tracked
(`len(ContextAggregateHash) > 0`), dirty, or carries a snapshot error,
and leaves it nil (`omitempty`) otherwise. `Dirty` mirrors
`context_dirty_since` being set.
- **Marker wiring.** The chat daemon is injected directly as the
`agentapi.ContextDirtyMarker`. It is unconditionally constructed (only
its background worker is gated), so the marker is always non-nil and the
wiring matches every other `api.chatDaemon` call site. `agentapi` still
treats a nil marker as "chatd absent", so `PushContextState` stays a
pure write path for any future caller that does not wire chatd in.
- **Refresh is atomic.** `RefreshChatContext` reads the agent's latest
snapshot and re-pins the chat in one repeatable-read transaction, so a
concurrent push cannot land between the read and the write and leave the
chat pinned to a stale hash with the dirty marker cleared.
- **Hydrate + dirty run inside the push transaction.** The fan-out
shares the push's transaction so a concurrent refresh cannot interleave
with the version gate; `context_dirty` watch events publish only after
commit. The pinned hash on dirtied chats is intentionally left unchanged
— the refresh endpoint re-pins it.
- **Dirtied chats keep their pinned hash.** Drift is advisory: a dirty
chat stays usable, and refreshing is the only path that advances the
pinned hash.
- **Test binds `chats.agent_id` directly.** In production the binding is
set lazily during a chat turn (`chatd.persistBuildAgentBinding`); the
test sets it via `dbgen` so it exercises the context flow rather than
turn resolution.

Plan: `coderd/x/chatd` context integration + E2E (sub-struct API,
create-time + push-time hydration, refresh endpoint;
`context_dirty_resources` and the per-turn pull untouched).

</details>

🤖 Generated by Coder Agents on behalf of @kylecarbs

coderd/x/chatd/chatd.go

@@ -1457,6 +1457,11 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
 	// committed and emitted its own state-machine notifications. The
 	// watch endpoint is maintained separately from chatstate notifications.
 	p.publishChatPubsubEvent(chat, codersdk.ChatWatchEventKindCreated, nil)
+
+	// Pin the chat to the agent's latest context snapshot if one exists.
+	// Best-effort: a chat created before its agent has pushed is hydrated
+	// by that agent's next push.
+	p.hydrateChatContextOnCreate(ctx, chat)
 	return chat, nil
 }
 

coderd/x/chatd/context_hydration.go

@@ -0,0 +1,169 @@
+package chatd
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"time"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog/v3"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/codersdk"
+)
+
+// latestAgentSnapshot looks up an agent's pinned context snapshot; ok is false
+// (with a nil error) when the agent has not pushed one yet.
+func latestAgentSnapshot(ctx context.Context, db database.Store, agentID uuid.UUID) (aggregateHash []byte, snapshotError string, ok bool, err error) {
+	snapshot, err := db.GetLatestWorkspaceAgentContextSnapshot(ctx, agentID)
+	switch {
+	case errors.Is(err, sql.ErrNoRows):
+		return nil, "", false, nil
+	case err != nil:
+		return nil, "", false, xerrors.Errorf("get latest snapshot: %w", err)
+	default:
+		return snapshot.AggregateHash, snapshot.SnapshotError, true, nil
+	}
+}
+
+// HydrateAndMarkChatsDirty implements agentapi.ContextDirtyMarker. It runs
+// inside the PushContextState transaction: it stamps the pushed snapshot hash
+// on chats for the agent that have not been hydrated yet (no dirty event),
+// then flips already-pinned chats whose hash differs to dirty. It returns a
+// callback that publishes the dirty watch events; the caller invokes it only
+// after the transaction commits, and the callback is a no-op when nothing
+// transitioned to dirty.
+//
+// The pinned hash on dirtied chats is intentionally left unchanged; the
+// refresh endpoint re-pins it.
+func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store, agentID uuid.UUID, aggregateHash []byte, snapshotError string, now time.Time) (func(), error) {
+	//nolint:gocritic // An agent does not own the chats bound to it.
+	ctx = dbauthz.AsChatd(ctx)
+
+	// Chats created before the agent's first push land with a NULL pinned
+	// hash. Stamp them now so they start clean; this is their first
+	// hydration, so no dirty event is emitted.
+	if err := tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
+		AgentID:       agentID,
+		AggregateHash: aggregateHash,
+		ContextError:  snapshotError,
+	}); err != nil {
+		return nil, xerrors.Errorf("hydrate agent chats context: %w", err)
+	}
+
+	dirtied, err := tx.MarkChatsContextDirtyByAgent(ctx, database.MarkChatsContextDirtyByAgentParams{
+		AgentID:       agentID,
+		AggregateHash: aggregateHash,
+		DirtySince:    sql.NullTime{Time: now, Valid: true},
+	})
+	if err != nil {
+		return nil, xerrors.Errorf("mark chats context dirty: %w", err)
+	}
+	if len(dirtied) == 0 {
+		return func() {}, nil
+	}
+
+	// Read the dirtied chats inside the transaction and capture their rows so
+	// the post-commit callback needs no database access: the published payload
+	// reflects the just-committed dirty state (no re-read a concurrent refresh
+	// could race), and the callback does not depend on the request-scoped
+	// context surviving past commit. Only the transitioned chats are read.
+	dirtyChats := make([]database.Chat, 0, len(dirtied))
+	for _, d := range dirtied {
+		chat, err := tx.GetChatByID(ctx, d.ID)
+		if err != nil {
+			return nil, xerrors.Errorf("get dirtied chat %s: %w", d.ID, err)
+		}
+		dirtyChats = append(dirtyChats, chat)
+	}
+
+	return func() {
+		p.publishChatPubsubEvents(dirtyChats, codersdk.ChatWatchEventKindContextDirty)
+	}, nil
+}
+
+// hydrateChatContextOnCreate pins a newly created chat to its agent's latest
+// context snapshot when one already exists. Best-effort: a chat whose agent
+// has not pushed yet is hydrated later by that agent's next push. Failures
+// are logged and swallowed so they never block chat creation.
+//
+// A concurrent push that already hydrated the chat is not clobbered with a
+// stale hash.
+func (p *Server) hydrateChatContextOnCreate(ctx context.Context, chat database.Chat) {
+	if !chat.AgentID.Valid {
+		return
+	}
+	//nolint:gocritic // Chatd stamps chats it does not own as the daemon subject.
+	ctx = dbauthz.AsChatd(ctx)
+
+	aggregateHash, snapshotError, ok, err := latestAgentSnapshot(ctx, p.db, chat.AgentID.UUID)
+	if err != nil {
+		p.logger.Warn(ctx, "hydrate chat context on create: get latest snapshot",
+			slog.F("chat_id", chat.ID), slog.Error(err))
+		return
+	}
+	if !ok {
+		return
+	}
+	if err := p.db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
+		AgentID:       chat.AgentID.UUID,
+		AggregateHash: aggregateHash,
+		ContextError:  snapshotError,
+	}); err != nil {
+		p.logger.Warn(ctx, "hydrate chat context on create: stamp chats",
+			slog.F("chat_id", chat.ID), slog.Error(err))
+	}
+}
+
+// RefreshChatContext re-pins a chat to its agent's latest context snapshot and
+// clears the dirty marker. It backs PUT /chats/{chat}/context (no body). A
+// chat with no bound agent, or whose agent has no snapshot, simply has its
+// pinned hash and dirty marker cleared.
+//
+// The snapshot read and the re-pin run in one repeatable-read transaction so a
+// concurrent push cannot land between them and leave the chat pinned to a
+// stale hash with the dirty marker cleared.
+func (p *Server) RefreshChatContext(ctx context.Context, chat database.Chat) (database.Chat, error) {
+	//nolint:gocritic // Chatd re-pins the chat as the daemon subject.
+	ctx = dbauthz.AsChatd(ctx)
+
+	var updated database.Chat
+	err := database.ReadModifyUpdate(p.db, func(tx database.Store) error {
+		var (
+			aggregateHash []byte
+			snapshotError string
+		)
+		if chat.AgentID.Valid {
+			hash, snapErr, ok, err := latestAgentSnapshot(ctx, tx, chat.AgentID.UUID)
+			if err != nil {
+				return err
+			}
+			if ok {
+				aggregateHash = hash
+				snapshotError = snapErr
+			}
+		}
+
+		if err := tx.SetChatContextSnapshot(ctx, database.SetChatContextSnapshotParams{
+			ID:            chat.ID,
+			AggregateHash: aggregateHash,
+			ContextError:  snapshotError,
+		}); err != nil {
+			return xerrors.Errorf("set chat context snapshot: %w", err)
+		}
+
+		got, err := tx.GetChatByID(ctx, chat.ID)
+		if err != nil {
+			return xerrors.Errorf("get chat after refresh: %w", err)
+		}
+		updated = got
+		return nil
+	})
+	if err != nil {
+		return database.Chat{}, err
+	}
+	return updated, nil
+}

coderd/x/chatd/context_hydration_internal_test.go

@@ -0,0 +1,81 @@
+package chatd
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/google/uuid"
+	"go.uber.org/mock/gomock"
+
+	"cdr.dev/slog/v3/sloggers/slogtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbmock"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestHydrateChatContextOnCreate covers the create-time pinning path, which the
+// end-to-end test cannot reach: chats there are inserted directly, bypassing
+// CreateChat. It pins to the agent's latest snapshot via the NULL-guarded
+// HydrateAgentChatsContext so a concurrent push is never clobbered, and is a
+// best-effort no-op when there is no agent or no snapshot.
+func TestHydrateChatContextOnCreate(t *testing.T) {
+	t.Parallel()
+
+	t.Run("PinsWhenSnapshotExists", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{db: db, logger: slogtest.Make(t, nil)}
+
+		agentID := uuid.New()
+		chat := database.Chat{ID: uuid.New(), AgentID: uuid.NullUUID{UUID: agentID, Valid: true}}
+		snapshot := database.WorkspaceAgentContextSnapshot{
+			WorkspaceAgentID: agentID,
+			AggregateHash:    []byte{0x0a, 0x0b},
+			SnapshotError:    "one source failed",
+		}
+
+		db.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), agentID).
+			Return(snapshot, nil)
+		// The guarded agent-scoped stamp, not an unconditional SetChatContextSnapshot,
+		// so a concurrent push that already hydrated the chat wins.
+		db.EXPECT().HydrateAgentChatsContext(gomock.Any(), database.HydrateAgentChatsContextParams{
+			AgentID:       agentID,
+			AggregateHash: snapshot.AggregateHash,
+			ContextError:  snapshot.SnapshotError,
+		}).Return(nil)
+
+		server.hydrateChatContextOnCreate(ctx, chat)
+	})
+
+	t.Run("SkipsWhenAgentless", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		// No EXPECT calls: a chat with no agent must touch the database zero times.
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{db: db, logger: slogtest.Make(t, nil)}
+
+		server.hydrateChatContextOnCreate(ctx, database.Chat{ID: uuid.New()})
+	})
+
+	t.Run("SkipsWhenNoSnapshot", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
+		ctrl := gomock.NewController(t)
+		db := dbmock.NewMockStore(ctrl)
+		server := &Server{db: db, logger: slogtest.Make(t, nil)}
+
+		agentID := uuid.New()
+		// ErrNoRows means the agent has not pushed yet; no stamp is written
+		// (HydrateAgentChatsContext has no EXPECT, so a call would fail the test).
+		db.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), agentID).
+			Return(database.WorkspaceAgentContextSnapshot{}, sql.ErrNoRows)
+
+		server.hydrateChatContextOnCreate(ctx, database.Chat{
+			ID:      uuid.New(),
+			AgentID: uuid.NullUUID{UUID: agentID, Valid: true},
+		})
+	})
+}

coderd/x/chatd/context_integration_test.go

@@ -0,0 +1,165 @@
+package chatd_test
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	agentproto "github.com/coder/coder/v2/agent/proto"
+	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+	"github.com/coder/coder/v2/provisioner/echo"
+	"github.com/coder/coder/v2/testutil"
+)
+
+// TestChatContextDirtyFromAgentPush is an end-to-end check of the chat
+// context integration. An echo-provisioned workspace agent pushes a context
+// snapshot that hydrates a bound chat; a later push with a different hash
+// marks the chat dirty; the experimental API reports the dirty state and the
+// snapshot error; the refresh endpoint re-pins the latest snapshot and clears
+// it; and a re-push of the now-pinned hash stays clean. A second chat bound to
+// no agent stays untouched throughout, guarding the agent-scoped queries.
+func TestChatContextDirtyFromAgentPush(t *testing.T) {
+	t.Parallel()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{
+		DeploymentValues:         directChatRoutingDeploymentValues(t),
+		IncludeProvisionerDaemon: true,
+	})
+	user := coderdtest.CreateFirstUser(t, client)
+	expClient := codersdk.NewExperimentalClient(client)
+
+	// Build a workspace with an agent via the echo provisioner.
+	agentToken := uuid.NewString()
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse:          echo.ParseComplete,
+		ProvisionPlan:  echo.PlanComplete,
+		ProvisionApply: echo.ApplyComplete,
+		ProvisionGraph: echo.ProvisionGraphWithAgent(agentToken),
+	})
+	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, template.ID)
+	coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+
+	ws, err := client.Workspace(ctx, workspace.ID)
+	require.NoError(t, err)
+	require.Len(t, ws.LatestBuild.Resources, 1)
+	require.Len(t, ws.LatestBuild.Resources[0].Agents, 1)
+	agentID := ws.LatestBuild.Resources[0].Agents[0].ID
+
+	// A chat bound to the agent. In production agent_id is set lazily during
+	// a workspace turn (chatd.persistBuildAgentBinding); bind it directly here
+	// so the test exercises the context flow rather than turn resolution.
+	// dbgen.ChatModelConfig provisions an AI provider as needed so the chat's
+	// last_model_config_id foreign key is satisfied.
+	model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{})
+	chat := dbgen.Chat(t, db, database.Chat{
+		OrganizationID:    user.OrganizationID,
+		OwnerID:           user.UserID,
+		WorkspaceID:       uuid.NullUUID{UUID: workspace.ID, Valid: true},
+		AgentID:           uuid.NullUUID{UUID: agentID, Valid: true},
+		LastModelConfigID: model.ID,
+		Status:            database.ChatStatusWaiting,
+	})
+
+	// An unrelated chat bound to no agent. The hydrate and dirty queries
+	// scope by agent_id, so this chat must stay untouched by every push
+	// below; it guards against the scoping clause silently breaking.
+	otherChat := dbgen.Chat(t, db, database.Chat{
+		OrganizationID:    user.OrganizationID,
+		OwnerID:           user.UserID,
+		LastModelConfigID: model.ID,
+		Status:            database.ChatStatusWaiting,
+	})
+
+	// Before any push there is no pinned context.
+	got, err := expClient.GetChat(ctx, chat.ID)
+	require.NoError(t, err)
+	require.Nil(t, got.Context, "no pinned context before the first push")
+
+	requireChatContextNil := func(id uuid.UUID, msg string) {
+		t.Helper()
+		unrelated, err := expClient.GetChat(ctx, id)
+		require.NoError(t, err)
+		require.Nil(t, unrelated.Context, msg)
+	}
+	requireChatContextNil(otherChat.ID, "agent-less chat has no pinned context")
+
+	// Connect as the agent and push the initial snapshot. The push runs the
+	// hydrate/dirty fan-out synchronously inside its transaction, so the chat
+	// reflects the change by the time the RPC returns.
+	agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(agentToken))
+	aAPI, _, err := agentClient.ConnectRPC210(ctx)
+	require.NoError(t, err)
+	defer func() { _ = aAPI.DRPCConn().Close() }()
+
+	hashA := []byte{0x01, 0x02, 0x03}
+	resp, err := aAPI.PushContextState(ctx, &agentproto.PushContextStateRequest{
+		Version:       1,
+		Initial:       true,
+		AggregateHash: hashA,
+	})
+	require.NoError(t, err)
+	require.True(t, resp.GetAccepted())
+
+	// The initial push hydrates the chat to a clean (not dirty) context.
+	got, err = expClient.GetChat(ctx, chat.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got.Context, "chat should be hydrated after the initial push")
+	require.False(t, got.Context.Dirty, "initial hydration is clean")
+	require.Nil(t, got.Context.DirtySince)
+
+	// The agent refreshes its context and pushes a different hash carrying a
+	// snapshot-level error, which drifts from the pinned hash and marks the
+	// chat dirty.
+	hashB := []byte{0x04, 0x05, 0x06}
+	const snapshotError = "two sources failed to resolve"
+	resp, err = aAPI.PushContextState(ctx, &agentproto.PushContextStateRequest{
+		Version:       2,
+		AggregateHash: hashB,
+		SnapshotError: snapshotError,
+	})
+	require.NoError(t, err)
+	require.True(t, resp.GetAccepted())
+
+	got, err = expClient.GetChat(ctx, chat.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got.Context)
+	require.True(t, got.Context.Dirty, "drift should mark the chat dirty")
+	require.NotNil(t, got.Context.DirtySince)
+	require.Empty(t, got.Context.Error, "dirty marking leaves the pinned hash and error unchanged")
+	requireChatContextNil(otherChat.ID, "agent-less chat unaffected by the dirty fan-out")
+
+	// Refreshing re-pins the latest snapshot (hash and error) and clears the
+	// dirty marker.
+	refreshed, err := expClient.RefreshChatContext(ctx, chat.ID)
+	require.NoError(t, err)
+	require.NotNil(t, refreshed.Context)
+	require.False(t, refreshed.Context.Dirty, "refresh clears the dirty marker")
+	require.Equal(t, snapshotError, refreshed.Context.Error, "refresh re-pins the snapshot error")
+
+	got, err = expClient.GetChat(ctx, chat.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got.Context)
+	require.False(t, got.Context.Dirty)
+
+	// Re-pushing the now-pinned hash proves the refresh advanced the pin to
+	// hashB: a matching hash must not re-dirty the chat.
+	resp, err = aAPI.PushContextState(ctx, &agentproto.PushContextStateRequest{
+		Version:       3,
+		AggregateHash: hashB,
+	})
+	require.NoError(t, err)
+	require.True(t, resp.GetAccepted())
+
+	got, err = expClient.GetChat(ctx, chat.ID)
+	require.NoError(t, err)
+	require.NotNil(t, got.Context)
+	require.False(t, got.Context.Dirty, "re-push of the pinned hash stays clean")
+}