Files
coder/coderd/x/chatd/context_hydration_internal_test.go
T
Kyle Carberry 1c78bd84b7 feat(coderd): copy agent context resources into the per-chat pin (#26438)
## What

Populates `chat_context_resources` (the per-chat pinned copy added in
#26430) by copying from `workspace_agent_context_resources` at the
points where a chat's `context_aggregate_hash` is set, in the same
transaction, so the pinned hash and pinned bodies always agree. No
prompt-building change yet; consuming the pinned copy in
`prepareGeneration` is a later, experiment-gated PR.

## How

- `HydrateAgentChatsContext` now hydrates NULL-hash chats **and** copies
the agent's resources onto them in one statement (a data-modifying CTE),
so the chat-create and agent-push paths need no Go change.
- New queries `InsertAgentContextResourcesIntoChat`,
`DeleteChatContextResources`, `ListChatContextResources`, each with a
hand-written dbauthz wrapper (per-chat update/read) and a
`MethodTestSuite` entry.
- `RefreshChatContext` re-pins resources via a shared `repinChatContext`
helper (clear-then-copy in a transaction). A dirty chat keeps its old
bodies until refresh.
- On agent rebind (e.g. a workspace rebuild produces a new agent), the
chat's context is re-pinned to the new agent so it stops injecting the
previous agent's resources. Best-effort: a context error never fails the
binding.

## Invariant

A chat's `chat_context_resources` always correspond to its
`context_aggregate_hash`. Bodies are (re)written only when the hash is
set (hydrate, refresh, rebind); a dirty chat keeps its old bodies until
refresh.

## Testing

Extends the context integration test to push real resources and assert
the copy across hydrate, dirty (no re-copy), and refresh. The dbauthz
`MethodTestSuite` covers the three new methods.

<details>
<summary>Why clear-then-copy (two statements)</summary>

The refresh/rebind re-pin clears the chat's rows then inserts the
agent's. It uses two sequential statements inside the transaction rather
than a single `WITH cleared AS (DELETE ...) INSERT ...`, because a
data-modifying CTE cannot see its own delete under snapshot isolation,
so overlapping sources (the common case: the same files re-pinned) would
collide on the `(chat_id, source)` primary key. The hydrate path inserts
into never-pinned (NULL-hash) chats and uses `ON CONFLICT DO UPDATE`
defensively.

</details>

<details>
<summary>Follow-ups</summary>

- `prepareGeneration` consuming the pinned instructions and skills
(experiment-gated).
- `codersdk.ChatContext` resources plus changed diff, and the frontend
indicator/refresh.
- Removing the per-turn pull and `last_injected_context`.

</details>

---

*This PR was created by Coder Agents on behalf of @kylecarbs.* Builds on
#26430.
2026-06-17 00:10:07 -07:00

86 lines
3.1 KiB
Go

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().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
func(f func(database.Store) error, _ *database.TxOptions) error { return f(db) })
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().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
func(f func(database.Store) error, _ *database.TxOptions) error { return f(db) })
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},
})
})
}