mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +08:00
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.
This commit is contained in:
@@ -2065,6 +2065,17 @@ func (q *querier) DeleteApplicationConnectAPIKeysByUserID(ctx context.Context, u
|
||||
return q.db.DeleteApplicationConnectAPIKeysByUserID(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) error {
|
||||
chat, err := q.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.DeleteChatContextResourcesByChatID(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatDebugDataAfterMessageID(ctx context.Context, arg database.DeleteChatDebugDataAfterMessageIDParams) (int64, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
@@ -5767,6 +5778,17 @@ func (q *querier) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyPar
|
||||
q.db.InsertAPIKey)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertAgentContextResourcesIntoChat(ctx context.Context, arg database.InsertAgentContextResourcesIntoChatParams) error {
|
||||
chat, err := q.db.GetChatByID(ctx, arg.ChatID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.InsertAgentContextResourcesIntoChat(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertAllUsersGroup(ctx context.Context, organizationID uuid.UUID) (database.Group, error) {
|
||||
// This method creates a new group.
|
||||
return insert(q.log, q.auth, rbac.ResourceGroup.InOrg(organizationID), q.db.InsertAllUsersGroup)(ctx, organizationID)
|
||||
@@ -6549,6 +6571,17 @@ func (q *querier) ListBoundaryLogsBySessionID(ctx context.Context, arg database.
|
||||
return q.db.ListBoundaryLogsBySessionID(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatContextResource, error) {
|
||||
chat, err := q.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, chat); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.ListChatContextResourcesByChatID(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -551,6 +551,26 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().SetChatContextSnapshot(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("InsertAgentContextResourcesIntoChat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
arg := database.InsertAgentContextResourcesIntoChatParams{ChatID: chat.ID, AgentID: uuid.New()}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().InsertAgentContextResourcesIntoChat(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(chat, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("DeleteChatContextResourcesByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteChatContextResourcesByChatID(gomock.Any(), chat.ID).Return(nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("ListChatContextResourcesByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
rows := []database.ChatContextResource{testutil.Fake(s.T(), faker, database.ChatContextResource{ChatID: chat.ID})}
|
||||
dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes()
|
||||
dbm.EXPECT().ListChatContextResourcesByChatID(gomock.Any(), chat.ID).Return(rows, nil).AnyTimes()
|
||||
check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(rows)
|
||||
}))
|
||||
s.Run("GetChatWorkerAcquisitionCandidates", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetChatWorkerAcquisitionCandidatesParams{
|
||||
StaleSeconds: 30,
|
||||
|
||||
+24
@@ -498,6 +498,14 @@ func (m queryMetricsStore) DeleteApplicationConnectAPIKeysByUserID(ctx context.C
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteChatContextResourcesByChatID(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("DeleteChatContextResourcesByChatID").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatContextResourcesByChatID").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatDebugDataAfterMessageID(ctx context.Context, arg database.DeleteChatDebugDataAfterMessageIDParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteChatDebugDataAfterMessageID(ctx, arg)
|
||||
@@ -3978,6 +3986,14 @@ func (m queryMetricsStore) InsertAPIKey(ctx context.Context, arg database.Insert
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) InsertAgentContextResourcesIntoChat(ctx context.Context, arg database.InsertAgentContextResourcesIntoChatParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.InsertAgentContextResourcesIntoChat(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("InsertAgentContextResourcesIntoChat").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertAgentContextResourcesIntoChat").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) InsertAllUsersGroup(ctx context.Context, organizationID uuid.UUID) (database.Group, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.InsertAllUsersGroup(ctx, organizationID)
|
||||
@@ -4674,6 +4690,14 @@ func (m queryMetricsStore) ListBoundaryLogsBySessionID(ctx context.Context, arg
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatContextResource, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ListChatContextResourcesByChatID(ctx, chatID)
|
||||
m.queryLatencies.WithLabelValues("ListChatContextResourcesByChatID").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatContextResourcesByChatID").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ListChatUsageLimitGroupOverrides(ctx)
|
||||
|
||||
Generated
+43
@@ -805,6 +805,20 @@ func (mr *MockStoreMockRecorder) DeleteApplicationConnectAPIKeysByUserID(ctx, us
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteApplicationConnectAPIKeysByUserID", reflect.TypeOf((*MockStore)(nil).DeleteApplicationConnectAPIKeysByUserID), ctx, userID)
|
||||
}
|
||||
|
||||
// DeleteChatContextResourcesByChatID mocks base method.
|
||||
func (m *MockStore) DeleteChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteChatContextResourcesByChatID", ctx, chatID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteChatContextResourcesByChatID indicates an expected call of DeleteChatContextResourcesByChatID.
|
||||
func (mr *MockStoreMockRecorder) DeleteChatContextResourcesByChatID(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatContextResourcesByChatID", reflect.TypeOf((*MockStore)(nil).DeleteChatContextResourcesByChatID), ctx, chatID)
|
||||
}
|
||||
|
||||
// DeleteChatDebugDataAfterMessageID mocks base method.
|
||||
func (m *MockStore) DeleteChatDebugDataAfterMessageID(ctx context.Context, arg database.DeleteChatDebugDataAfterMessageIDParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -7451,6 +7465,20 @@ func (mr *MockStoreMockRecorder) InsertAPIKey(ctx, arg any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAPIKey", reflect.TypeOf((*MockStore)(nil).InsertAPIKey), ctx, arg)
|
||||
}
|
||||
|
||||
// InsertAgentContextResourcesIntoChat mocks base method.
|
||||
func (m *MockStore) InsertAgentContextResourcesIntoChat(ctx context.Context, arg database.InsertAgentContextResourcesIntoChatParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "InsertAgentContextResourcesIntoChat", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// InsertAgentContextResourcesIntoChat indicates an expected call of InsertAgentContextResourcesIntoChat.
|
||||
func (mr *MockStoreMockRecorder) InsertAgentContextResourcesIntoChat(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAgentContextResourcesIntoChat", reflect.TypeOf((*MockStore)(nil).InsertAgentContextResourcesIntoChat), ctx, arg)
|
||||
}
|
||||
|
||||
// InsertAllUsersGroup mocks base method.
|
||||
func (m *MockStore) InsertAllUsersGroup(ctx context.Context, organizationID uuid.UUID) (database.Group, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -8801,6 +8829,21 @@ func (mr *MockStoreMockRecorder) ListBoundaryLogsBySessionID(ctx, arg any) *gomo
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListBoundaryLogsBySessionID", reflect.TypeOf((*MockStore)(nil).ListBoundaryLogsBySessionID), ctx, arg)
|
||||
}
|
||||
|
||||
// ListChatContextResourcesByChatID mocks base method.
|
||||
func (m *MockStore) ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatContextResource, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ListChatContextResourcesByChatID", ctx, chatID)
|
||||
ret0, _ := ret[0].([]database.ChatContextResource)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ListChatContextResourcesByChatID indicates an expected call of ListChatContextResourcesByChatID.
|
||||
func (mr *MockStoreMockRecorder) ListChatContextResourcesByChatID(ctx, chatID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatContextResourcesByChatID", reflect.TypeOf((*MockStore)(nil).ListChatContextResourcesByChatID), ctx, chatID)
|
||||
}
|
||||
|
||||
// ListChatUsageLimitGroupOverrides mocks base method.
|
||||
func (m *MockStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
+21
-3
@@ -134,6 +134,10 @@ type sqlcQuerier interface {
|
||||
// be recreated.
|
||||
DeleteAllWebpushSubscriptions(ctx context.Context) error
|
||||
DeleteApplicationConnectAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error
|
||||
// Clears a chat's pinned context resources. Used as the first half of a
|
||||
// clear-then-copy re-pin, and on its own when the chat's current agent
|
||||
// has no snapshot.
|
||||
DeleteChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) error
|
||||
// Deletes debug runs (and their cascaded steps) whose message IDs
|
||||
// exceed the cutoff. The started_before bound prevents retried
|
||||
// cleanup from deleting runs created by a replacement turn that
|
||||
@@ -1002,9 +1006,15 @@ type sqlcQuerier interface {
|
||||
GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]GetWorkspacesEligibleForTransitionRow, error)
|
||||
GetWorkspacesForWorkspaceMetrics(ctx context.Context) ([]GetWorkspacesForWorkspaceMetricsRow, error)
|
||||
// Stamps the pinned hash and error on every not-yet-hydrated chat for
|
||||
// an agent (context_aggregate_hash IS NULL). Runs as a side effect of
|
||||
// an agent push so chats created before the agent was ready pick up the
|
||||
// snapshot without a dirty event. Does not bump updated_at.
|
||||
// an agent (context_aggregate_hash IS NULL) and copies the agent's
|
||||
// current context resources onto those chats in the same statement, so
|
||||
// a chat's pinned hash and pinned bodies are always written together.
|
||||
// Runs as a side effect of an agent push and of chat-create hydration,
|
||||
// so chats created before the agent was ready pick up the snapshot
|
||||
// without a dirty event. The ON CONFLICT upsert is defensive: a
|
||||
// not-yet-hydrated chat has no pinned rows, so it normally inserts.
|
||||
// Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch
|
||||
// sets chat_context_resources.updated_at on the rows it rewrites.
|
||||
HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) error
|
||||
// Increments generation_attempt and returns the resulting value.
|
||||
IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error)
|
||||
@@ -1017,6 +1027,11 @@ type sqlcQuerier interface {
|
||||
InsertAIProvider(ctx context.Context, arg InsertAIProviderParams) (AIProvider, error)
|
||||
InsertAIProviderKey(ctx context.Context, arg InsertAIProviderKeyParams) (AIProviderKey, error)
|
||||
InsertAPIKey(ctx context.Context, arg InsertAPIKeyParams) (APIKey, error)
|
||||
// Copies an agent's current context resources onto a single chat. Pair
|
||||
// with DeleteChatContextResourcesByChatID (clear-then-copy, in a
|
||||
// transaction) to re-pin a chat to its agent's latest snapshot from the
|
||||
// refresh endpoint and on agent rebinding.
|
||||
InsertAgentContextResourcesIntoChat(ctx context.Context, arg InsertAgentContextResourcesIntoChatParams) error
|
||||
// We use the organization_id as the id
|
||||
// for simplicity since all users is
|
||||
// every member of the org.
|
||||
@@ -1164,6 +1179,9 @@ type sqlcQuerier interface {
|
||||
// Supports optional exclusive sequence number bounds (seq_after, seq_before)
|
||||
// for fetching events between two known interceptions.
|
||||
ListBoundaryLogsBySessionID(ctx context.Context, arg ListBoundaryLogsBySessionIDParams) ([]BoundaryLog, error)
|
||||
// Lists a chat's pinned context resources, ordered deterministically by
|
||||
// source.
|
||||
ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatContextResource, error)
|
||||
ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error)
|
||||
ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error)
|
||||
ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
|
||||
|
||||
Generated
+120
-12
@@ -6622,6 +6622,19 @@ func (q *sqlQuerier) DeleteAllChatQueuedMessagesReturningCount(ctx context.Conte
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
const deleteChatContextResourcesByChatID = `-- name: DeleteChatContextResourcesByChatID :exec
|
||||
DELETE FROM chat_context_resources
|
||||
WHERE chat_id = $1::uuid
|
||||
`
|
||||
|
||||
// Clears a chat's pinned context resources. Used as the first half of a
|
||||
// clear-then-copy re-pin, and on its own when the chat's current agent
|
||||
// has no snapshot.
|
||||
func (q *sqlQuerier) DeleteChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) error {
|
||||
_, err := q.db.ExecContext(ctx, deleteChatContextResourcesByChatID, chatID)
|
||||
return err
|
||||
}
|
||||
|
||||
const deleteChatQueuedMessage = `-- name: DeleteChatQueuedMessage :exec
|
||||
DELETE FROM chat_queued_messages WHERE id = $1 AND chat_id = $2
|
||||
`
|
||||
@@ -9782,27 +9795,54 @@ func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGrou
|
||||
}
|
||||
|
||||
const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :exec
|
||||
UPDATE chats
|
||||
SET
|
||||
context_aggregate_hash = $1,
|
||||
context_error = $2
|
||||
WHERE agent_id = $3::uuid
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL
|
||||
WITH hydrated AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
context_aggregate_hash = $2,
|
||||
context_error = $3
|
||||
WHERE agent_id = $1::uuid
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL
|
||||
RETURNING id
|
||||
)
|
||||
INSERT INTO chat_context_resources (
|
||||
chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path
|
||||
)
|
||||
SELECT
|
||||
hydrated.id, r.source, r.body_kind, r.body, r.content_hash,
|
||||
r.size_bytes, r.status, r.error, r.source_path
|
||||
FROM hydrated
|
||||
CROSS JOIN workspace_agent_context_resources r
|
||||
WHERE r.workspace_agent_id = $1::uuid
|
||||
ON CONFLICT (chat_id, source) DO UPDATE SET
|
||||
body_kind = EXCLUDED.body_kind,
|
||||
body = EXCLUDED.body,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
size_bytes = EXCLUDED.size_bytes,
|
||||
status = EXCLUDED.status,
|
||||
error = EXCLUDED.error,
|
||||
source_path = EXCLUDED.source_path,
|
||||
updated_at = now()
|
||||
`
|
||||
|
||||
type HydrateAgentChatsContextParams struct {
|
||||
AgentID uuid.UUID `db:"agent_id" json:"agent_id"`
|
||||
AggregateHash []byte `db:"aggregate_hash" json:"aggregate_hash"`
|
||||
ContextError string `db:"context_error" json:"context_error"`
|
||||
AgentID uuid.UUID `db:"agent_id" json:"agent_id"`
|
||||
}
|
||||
|
||||
// Stamps the pinned hash and error on every not-yet-hydrated chat for
|
||||
// an agent (context_aggregate_hash IS NULL). Runs as a side effect of
|
||||
// an agent push so chats created before the agent was ready pick up the
|
||||
// snapshot without a dirty event. Does not bump updated_at.
|
||||
// an agent (context_aggregate_hash IS NULL) and copies the agent's
|
||||
// current context resources onto those chats in the same statement, so
|
||||
// a chat's pinned hash and pinned bodies are always written together.
|
||||
// Runs as a side effect of an agent push and of chat-create hydration,
|
||||
// so chats created before the agent was ready pick up the snapshot
|
||||
// without a dirty event. The ON CONFLICT upsert is defensive: a
|
||||
// not-yet-hydrated chat has no pinned rows, so it normally inserts.
|
||||
// Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch
|
||||
// sets chat_context_resources.updated_at on the rows it rewrites.
|
||||
func (q *sqlQuerier) HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) error {
|
||||
_, err := q.db.ExecContext(ctx, hydrateAgentChatsContext, arg.AggregateHash, arg.ContextError, arg.AgentID)
|
||||
_, err := q.db.ExecContext(ctx, hydrateAgentChatsContext, arg.AgentID, arg.AggregateHash, arg.ContextError)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -9821,6 +9861,31 @@ func (q *sqlQuerier) IncrementChatGenerationAttempt(ctx context.Context, id uuid
|
||||
return generation_attempt, err
|
||||
}
|
||||
|
||||
const insertAgentContextResourcesIntoChat = `-- name: InsertAgentContextResourcesIntoChat :exec
|
||||
INSERT INTO chat_context_resources (
|
||||
chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path
|
||||
)
|
||||
SELECT
|
||||
$1::uuid, r.source, r.body_kind, r.body, r.content_hash,
|
||||
r.size_bytes, r.status, r.error, r.source_path
|
||||
FROM workspace_agent_context_resources r
|
||||
WHERE r.workspace_agent_id = $2::uuid
|
||||
`
|
||||
|
||||
type InsertAgentContextResourcesIntoChatParams struct {
|
||||
ChatID uuid.UUID `db:"chat_id" json:"chat_id"`
|
||||
AgentID uuid.UUID `db:"agent_id" json:"agent_id"`
|
||||
}
|
||||
|
||||
// Copies an agent's current context resources onto a single chat. Pair
|
||||
// with DeleteChatContextResourcesByChatID (clear-then-copy, in a
|
||||
// transaction) to re-pin a chat to its agent's latest snapshot from the
|
||||
// refresh endpoint and on agent rebinding.
|
||||
func (q *sqlQuerier) InsertAgentContextResourcesIntoChat(ctx context.Context, arg InsertAgentContextResourcesIntoChatParams) error {
|
||||
_, err := q.db.ExecContext(ctx, insertAgentContextResourcesIntoChat, arg.ChatID, arg.AgentID)
|
||||
return err
|
||||
}
|
||||
|
||||
const insertChat = `-- name: InsertChat :one
|
||||
WITH inserted_chat AS (
|
||||
INSERT INTO chats (
|
||||
@@ -10332,6 +10397,49 @@ func (q *sqlQuerier) LinkChatFiles(ctx context.Context, arg LinkChatFilesParams)
|
||||
return rejected_new_files, err
|
||||
}
|
||||
|
||||
const listChatContextResourcesByChatID = `-- name: ListChatContextResourcesByChatID :many
|
||||
SELECT chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path, created_at, updated_at FROM chat_context_resources
|
||||
WHERE chat_id = $1::uuid
|
||||
ORDER BY source ASC
|
||||
`
|
||||
|
||||
// Lists a chat's pinned context resources, ordered deterministically by
|
||||
// source.
|
||||
func (q *sqlQuerier) ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatContextResource, error) {
|
||||
rows, err := q.db.QueryContext(ctx, listChatContextResourcesByChatID, chatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []ChatContextResource
|
||||
for rows.Next() {
|
||||
var i ChatContextResource
|
||||
if err := rows.Scan(
|
||||
&i.ChatID,
|
||||
&i.Source,
|
||||
&i.BodyKind,
|
||||
&i.Body,
|
||||
&i.ContentHash,
|
||||
&i.SizeBytes,
|
||||
&i.Status,
|
||||
&i.Error,
|
||||
&i.SourcePath,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, i)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const listChatUsageLimitGroupOverrides = `-- name: ListChatUsageLimitGroupOverrides :many
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
|
||||
@@ -1493,16 +1493,43 @@ WHERE id = @id::uuid;
|
||||
|
||||
-- name: HydrateAgentChatsContext :exec
|
||||
-- Stamps the pinned hash and error on every not-yet-hydrated chat for
|
||||
-- an agent (context_aggregate_hash IS NULL). Runs as a side effect of
|
||||
-- an agent push so chats created before the agent was ready pick up the
|
||||
-- snapshot without a dirty event. Does not bump updated_at.
|
||||
UPDATE chats
|
||||
SET
|
||||
context_aggregate_hash = @aggregate_hash,
|
||||
context_error = @context_error
|
||||
WHERE agent_id = @agent_id::uuid
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL;
|
||||
-- an agent (context_aggregate_hash IS NULL) and copies the agent's
|
||||
-- current context resources onto those chats in the same statement, so
|
||||
-- a chat's pinned hash and pinned bodies are always written together.
|
||||
-- Runs as a side effect of an agent push and of chat-create hydration,
|
||||
-- so chats created before the agent was ready pick up the snapshot
|
||||
-- without a dirty event. The ON CONFLICT upsert is defensive: a
|
||||
-- not-yet-hydrated chat has no pinned rows, so it normally inserts.
|
||||
-- Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch
|
||||
-- sets chat_context_resources.updated_at on the rows it rewrites.
|
||||
WITH hydrated AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
context_aggregate_hash = @aggregate_hash,
|
||||
context_error = @context_error
|
||||
WHERE agent_id = @agent_id::uuid
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL
|
||||
RETURNING id
|
||||
)
|
||||
INSERT INTO chat_context_resources (
|
||||
chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path
|
||||
)
|
||||
SELECT
|
||||
hydrated.id, r.source, r.body_kind, r.body, r.content_hash,
|
||||
r.size_bytes, r.status, r.error, r.source_path
|
||||
FROM hydrated
|
||||
CROSS JOIN workspace_agent_context_resources r
|
||||
WHERE r.workspace_agent_id = @agent_id::uuid
|
||||
ON CONFLICT (chat_id, source) DO UPDATE SET
|
||||
body_kind = EXCLUDED.body_kind,
|
||||
body = EXCLUDED.body,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
size_bytes = EXCLUDED.size_bytes,
|
||||
status = EXCLUDED.status,
|
||||
error = EXCLUDED.error,
|
||||
source_path = EXCLUDED.source_path,
|
||||
updated_at = now();
|
||||
|
||||
-- name: MarkChatsContextDirtyByAgent :many
|
||||
-- Flips active, already-hydrated chats for an agent to dirty when the
|
||||
@@ -1520,6 +1547,34 @@ WHERE agent_id = @agent_id::uuid
|
||||
AND context_dirty_since IS NULL
|
||||
RETURNING id, owner_id;
|
||||
|
||||
-- name: InsertAgentContextResourcesIntoChat :exec
|
||||
-- Copies an agent's current context resources onto a single chat. Pair
|
||||
-- with DeleteChatContextResourcesByChatID (clear-then-copy, in a
|
||||
-- transaction) to re-pin a chat to its agent's latest snapshot from the
|
||||
-- refresh endpoint and on agent rebinding.
|
||||
INSERT INTO chat_context_resources (
|
||||
chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path
|
||||
)
|
||||
SELECT
|
||||
@chat_id::uuid, r.source, r.body_kind, r.body, r.content_hash,
|
||||
r.size_bytes, r.status, r.error, r.source_path
|
||||
FROM workspace_agent_context_resources r
|
||||
WHERE r.workspace_agent_id = @agent_id::uuid;
|
||||
|
||||
-- name: DeleteChatContextResourcesByChatID :exec
|
||||
-- Clears a chat's pinned context resources. Used as the first half of a
|
||||
-- clear-then-copy re-pin, and on its own when the chat's current agent
|
||||
-- has no snapshot.
|
||||
DELETE FROM chat_context_resources
|
||||
WHERE chat_id = @chat_id::uuid;
|
||||
|
||||
-- name: ListChatContextResourcesByChatID :many
|
||||
-- Lists a chat's pinned context resources, ordered deterministically by
|
||||
-- source.
|
||||
SELECT * FROM chat_context_resources
|
||||
WHERE chat_id = @chat_id::uuid
|
||||
ORDER BY source ASC;
|
||||
|
||||
-- name: LinkChatFiles :one
|
||||
-- LinkChatFiles inserts file associations into the chat_file_links
|
||||
-- join table with deduplication (ON CONFLICT DO NOTHING). The INSERT
|
||||
|
||||
@@ -707,6 +707,25 @@ func (c *turnWorkspaceContext) persistBuildAgentBinding(
|
||||
"update chat build/agent binding: %w", err,
|
||||
)
|
||||
}
|
||||
|
||||
// If the chat was rebound to a different agent (e.g. a workspace rebuild
|
||||
// produced a new agent), re-pin its context to the new agent so it stops
|
||||
// injecting the previous agent's resources. Best-effort: a context error
|
||||
// must never fail the binding. The pinned context fields on updatedChat
|
||||
// are background state, reloaded on the next snapshot fetch.
|
||||
if chatSnapshot.AgentID.Valid && chatSnapshot.AgentID.UUID != agentID {
|
||||
//nolint:gocritic // Chatd re-pins chats it does not own as the daemon subject.
|
||||
repinCtx := dbauthz.AsChatd(ctx)
|
||||
if repinErr := database.ReadModifyUpdate(c.server.db, func(tx database.Store) error {
|
||||
return repinChatContext(repinCtx, tx, chatSnapshot.ID, uuid.NullUUID{UUID: agentID, Valid: true})
|
||||
}); repinErr != nil {
|
||||
c.server.logger.Warn(ctx, "re-pin chat context after agent rebind",
|
||||
slog.F("chat_id", chatSnapshot.ID),
|
||||
slog.F("agent_id", agentID),
|
||||
slog.Error(repinErr))
|
||||
}
|
||||
}
|
||||
|
||||
c.setCurrentChat(updatedChat)
|
||||
return updatedChat, nil
|
||||
}
|
||||
|
||||
@@ -1814,12 +1814,27 @@ func TestTurnWorkspaceContext_NullBindingLazyBind(t *testing.T) {
|
||||
require.Equal(t, workspaceAgent, gotAgent)
|
||||
}
|
||||
|
||||
// expectBestEffortContextRepin lets persistBuildAgentBinding's best-effort
|
||||
// context re-pin run against a mock store. The re-pin fires whenever a turn
|
||||
// rebinds a chat to a different agent; these agent-switch tests set up no
|
||||
// context snapshot, so it takes the no-snapshot clear path. The re-pin
|
||||
// behavior itself is covered by TestPersistBuildAgentBindingRepinsContext.
|
||||
func expectBestEffortContextRepin(db *dbmock.MockStore) {
|
||||
db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(f func(database.Store) error, _ *database.TxOptions) error { return f(db) }).AnyTimes()
|
||||
db.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), gomock.Any()).
|
||||
Return(database.WorkspaceAgentContextSnapshot{}, sql.ErrNoRows).AnyTimes()
|
||||
db.EXPECT().SetChatContextSnapshot(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
db.EXPECT().DeleteChatContextResourcesByChatID(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
}
|
||||
|
||||
func TestTurnWorkspaceContext_StaleBindingRepair(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := context.Background()
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
expectBestEffortContextRepin(db)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
staleAgentID := uuid.New()
|
||||
@@ -1875,6 +1890,7 @@ func TestTurnWorkspaceContextGetWorkspaceConnLazyValidationSwitchesWorkspaceAgen
|
||||
ctx := context.Background()
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
expectBestEffortContextRepin(db)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
staleAgentID := uuid.New()
|
||||
@@ -3023,6 +3039,7 @@ func TestGetWorkspaceConn_StaleAgentRecovery(t *testing.T) {
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
expectBestEffortContextRepin(db)
|
||||
|
||||
workspaceID := uuid.New()
|
||||
oldAgentID := uuid.New()
|
||||
|
||||
@@ -99,29 +99,84 @@ func (p *Server) hydrateChatContextOnCreate(ctx context.Context, chat database.C
|
||||
//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,
|
||||
// Read the snapshot hash and copy the agent's resources in one
|
||||
// repeatable-read transaction so a concurrent push cannot commit between
|
||||
// the two and leave the chat stamped with one snapshot's hash but another
|
||||
// snapshot's resources. The NULL-hash guard inside the statement still
|
||||
// keeps a concurrent push that already hydrated the chat from being
|
||||
// clobbered.
|
||||
if err := database.ReadModifyUpdate(p.db, func(tx database.Store) error {
|
||||
aggregateHash, snapshotError, ok, err := latestAgentSnapshot(ctx, tx, chat.AgentID.UUID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return tx.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",
|
||||
p.logger.Warn(ctx, "hydrate chat context on create",
|
||||
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.
|
||||
// repinChatContext re-pins a single chat to its agent's latest context
|
||||
// snapshot: it sets the pinned hash and error and rewrites the chat's pinned
|
||||
// resources (clear-then-copy) so the two always agree. A chat with no bound
|
||||
// agent, or whose agent has no snapshot, has its pinned hash, dirty marker,
|
||||
// and resources cleared. Callers run this inside a transaction.
|
||||
func repinChatContext(ctx context.Context, db database.Store, chatID uuid.UUID, agentID uuid.NullUUID) error {
|
||||
var (
|
||||
aggregateHash []byte
|
||||
snapshotError string
|
||||
hasSnapshot bool
|
||||
)
|
||||
if agentID.Valid {
|
||||
hash, snapErr, ok, err := latestAgentSnapshot(ctx, db, agentID.UUID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if ok {
|
||||
aggregateHash = hash
|
||||
snapshotError = snapErr
|
||||
hasSnapshot = true
|
||||
}
|
||||
}
|
||||
|
||||
if err := db.SetChatContextSnapshot(ctx, database.SetChatContextSnapshotParams{
|
||||
ID: chatID,
|
||||
AggregateHash: aggregateHash,
|
||||
ContextError: snapshotError,
|
||||
}); err != nil {
|
||||
return xerrors.Errorf("set chat context snapshot: %w", err)
|
||||
}
|
||||
|
||||
// Clear-then-copy so the pinned resources always match the pinned hash.
|
||||
// A single delete+insert statement cannot see its own delete under
|
||||
// snapshot isolation, so overlapping sources would collide.
|
||||
if err := db.DeleteChatContextResourcesByChatID(ctx, chatID); err != nil {
|
||||
return xerrors.Errorf("clear chat context resources: %w", err)
|
||||
}
|
||||
if hasSnapshot {
|
||||
if err := db.InsertAgentContextResourcesIntoChat(ctx, database.InsertAgentContextResourcesIntoChatParams{
|
||||
ChatID: chatID,
|
||||
AgentID: agentID.UUID,
|
||||
}); err != nil {
|
||||
return xerrors.Errorf("copy agent context resources: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RefreshChatContext re-pins a chat to its agent's latest context snapshot
|
||||
// (hash, error, and resource bodies) 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, dirty marker, and
|
||||
// resources 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
|
||||
@@ -132,29 +187,17 @@ func (p *Server) RefreshChatContext(ctx context.Context, chat database.Chat) (da
|
||||
|
||||
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
|
||||
}
|
||||
// Re-read the chat inside the transaction so a serialization-conflict
|
||||
// retry re-pins against the chat's current agent. Using the AgentID
|
||||
// captured before the transaction would re-pin to a stale agent if a
|
||||
// concurrent rebind landed between that read and the retry.
|
||||
current, err := tx.GetChatByID(ctx, chat.ID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get chat for refresh: %w", err)
|
||||
}
|
||||
|
||||
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)
|
||||
if err := repinChatContext(ctx, tx, current.ID, current.AgentID); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
got, err := tx.GetChatByID(ctx, chat.ID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get chat after refresh: %w", err)
|
||||
|
||||
@@ -36,6 +36,8 @@ func TestHydrateChatContextOnCreate(t *testing.T) {
|
||||
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,
|
||||
@@ -70,6 +72,8 @@ func TestHydrateChatContextOnCreate(t *testing.T) {
|
||||
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)
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
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/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/agentsdk"
|
||||
@@ -91,6 +92,47 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
}
|
||||
requireChatContextNil(otherChat.ID, "agent-less chat has no pinned context")
|
||||
|
||||
// Resource builders and a reader for the per-chat pinned copy. The agent
|
||||
// pushes these; hydration and refresh copy them onto the bound chat.
|
||||
agentsSource := "/home/coder/workspace/AGENTS.md"
|
||||
skillSource := "/home/coder/workspace/.agents/skills/example/SKILL.md"
|
||||
agentsV1Hash := []byte{0x11}
|
||||
agentsV2Hash := []byte{0x22}
|
||||
skillHash := []byte{0x33}
|
||||
instructionResource := func(source, content string, hash []byte) *agentproto.ContextResource {
|
||||
return &agentproto.ContextResource{
|
||||
Source: source,
|
||||
ContentHash: hash,
|
||||
SizeBytes: uint64(len(content)),
|
||||
Status: agentproto.ContextResource_OK,
|
||||
Body: &agentproto.ContextResource_InstructionFile{
|
||||
InstructionFile: &agentproto.InstructionFileBody{Content: []byte(content)},
|
||||
},
|
||||
}
|
||||
}
|
||||
skillResource := func(source string, hash []byte) *agentproto.ContextResource {
|
||||
return &agentproto.ContextResource{
|
||||
Source: source,
|
||||
ContentHash: hash,
|
||||
SizeBytes: 16,
|
||||
Status: agentproto.ContextResource_OK,
|
||||
Body: &agentproto.ContextResource_Skill{
|
||||
Skill: &agentproto.SkillMetaBody{Meta: []byte("---\nname: example\n---"), Name: "example", Description: "demo skill"},
|
||||
},
|
||||
}
|
||||
}
|
||||
pinnedResources := func(id uuid.UUID) map[string]database.ChatContextResource {
|
||||
t.Helper()
|
||||
//nolint:gocritic // Test reads the chat-owned rows as the chatd subject; ctx carries no per-user actor.
|
||||
rows, lerr := db.ListChatContextResourcesByChatID(dbauthz.AsChatd(ctx), id)
|
||||
require.NoError(t, lerr)
|
||||
out := make(map[string]database.ChatContextResource, len(rows))
|
||||
for _, r := range rows {
|
||||
out[r.Source] = r
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// 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.
|
||||
@@ -104,6 +146,9 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
Version: 1,
|
||||
Initial: true,
|
||||
AggregateHash: hashA,
|
||||
Resources: []*agentproto.ContextResource{
|
||||
instructionResource(agentsSource, "hello-v1", agentsV1Hash),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.True(t, resp.GetAccepted())
|
||||
@@ -115,6 +160,14 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
require.False(t, got.Context.Dirty, "initial hydration is clean")
|
||||
require.Nil(t, got.Context.DirtySince)
|
||||
|
||||
// The initial push also copied the agent's resources onto the chat.
|
||||
pinned := pinnedResources(chat.ID)
|
||||
require.Len(t, pinned, 1, "initial hydration copies the agent's resources")
|
||||
require.Equal(t, agentsV1Hash, pinned[agentsSource].ContentHash)
|
||||
require.Equal(t, database.WorkspaceAgentContextBodyKindInstructionFile, pinned[agentsSource].BodyKind)
|
||||
require.Equal(t, database.WorkspaceAgentContextResourceStatusOk, pinned[agentsSource].Status)
|
||||
require.Empty(t, pinnedResources(otherChat.ID), "agent-less chat has no pinned resources")
|
||||
|
||||
// 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.
|
||||
@@ -124,6 +177,10 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
Version: 2,
|
||||
AggregateHash: hashB,
|
||||
SnapshotError: snapshotError,
|
||||
Resources: []*agentproto.ContextResource{
|
||||
instructionResource(agentsSource, "hello-v2", agentsV2Hash),
|
||||
skillResource(skillSource, skillHash),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.True(t, resp.GetAccepted())
|
||||
@@ -136,6 +193,12 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
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")
|
||||
|
||||
// The dirty fan-out must NOT re-copy resources: the chat keeps the bodies
|
||||
// from its pinned (hashA) snapshot until it is refreshed.
|
||||
pinned = pinnedResources(chat.ID)
|
||||
require.Len(t, pinned, 1, "dirty marking does not re-copy resources")
|
||||
require.Equal(t, agentsV1Hash, pinned[agentsSource].ContentHash, "chat keeps the pinned snapshot's resources while dirty")
|
||||
|
||||
// Refreshing re-pins the latest snapshot (hash and error) and clears the
|
||||
// dirty marker.
|
||||
refreshed, err := expClient.RefreshChatContext(ctx, chat.ID)
|
||||
@@ -144,6 +207,13 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) {
|
||||
require.False(t, refreshed.Context.Dirty, "refresh clears the dirty marker")
|
||||
require.Equal(t, snapshotError, refreshed.Context.Error, "refresh re-pins the snapshot error")
|
||||
|
||||
// Refresh re-pinned the agent's current resources (the hashB set).
|
||||
pinned = pinnedResources(chat.ID)
|
||||
require.Len(t, pinned, 2, "refresh re-pins the agent's current resources")
|
||||
require.Equal(t, agentsV2Hash, pinned[agentsSource].ContentHash)
|
||||
require.Equal(t, skillHash, pinned[skillSource].ContentHash)
|
||||
require.Equal(t, database.WorkspaceAgentContextBodyKindSkill, pinned[skillSource].BodyKind)
|
||||
|
||||
got, err = expClient.GetChat(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got.Context)
|
||||
|
||||
@@ -0,0 +1,319 @@
|
||||
package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtime"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
// TestPersistBuildAgentBindingRepinsContext covers the agent-rebind re-pin path
|
||||
// in persistBuildAgentBinding, which the end-to-end context test does not
|
||||
// reach: the shared repinChatContext helper is proven through the refresh
|
||||
// endpoint, but the rebind-specific wiring (the AgentID change guard, the
|
||||
// AsChatd escalation, the ReadModifyUpdate transaction, and the best-effort
|
||||
// error swallow) is exercised only here.
|
||||
func TestPersistBuildAgentBindingRepinsContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// RebindsToNewAgent: a chat pinned to agent A is rebound to agent B, so
|
||||
// its pinned hash and pinned resources must switch to B's snapshot.
|
||||
t.Run("RebindsToNewAgent", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
fix := newRebindFixture(t)
|
||||
|
||||
chat := dbgen.Chat(t, fix.db, database.Chat{
|
||||
OwnerID: fix.user.ID,
|
||||
OrganizationID: fix.org.ID,
|
||||
LastModelConfigID: fix.model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: fix.ws.ID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: fix.agentA, Valid: true},
|
||||
Status: database.ChatStatusWaiting,
|
||||
})
|
||||
|
||||
// Pin the chat to agent A through the production hydrate path so it
|
||||
// starts with A's hash and A's resources, exactly as an agent push
|
||||
// would leave it.
|
||||
require.NoError(t, fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: fix.agentA,
|
||||
AggregateHash: fix.hashA,
|
||||
}))
|
||||
preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, preRes, 1)
|
||||
require.Equal(t, fix.srcA, preRes[0].Source)
|
||||
|
||||
wc := newRebindTurnContext(t, fix.db, chat)
|
||||
updated, err := wc.persistBuildAgentBinding(fix.ctx, chat, fix.buildID, fix.agentB)
|
||||
require.NoError(t, err)
|
||||
require.True(t, updated.AgentID.Valid)
|
||||
require.Equal(t, fix.agentB, updated.AgentID.UUID, "the binding commits the new agent")
|
||||
|
||||
// The re-pin runs in its own transaction after the binding row is
|
||||
// written, so re-read the chat to observe the new pinned state.
|
||||
post, err := fix.db.GetChatByID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, fix.hashB, post.ContextAggregateHash, "rebind re-pins the new agent's hash")
|
||||
|
||||
postRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, postRes, 1, "rebind swaps the pinned resources to the new agent's set")
|
||||
require.Equal(t, fix.srcB, postRes[0].Source)
|
||||
require.Equal(t, fix.hashB, postRes[0].ContentHash)
|
||||
require.JSONEq(t, string(fix.bodyB), string(postRes[0].Body), "the new agent's resource body is copied verbatim")
|
||||
})
|
||||
|
||||
// SkipsRepinWhenNoPriorAgent: binding a chat that had no agent must not
|
||||
// re-pin here (the create/push path owns first-time pinning), so the guard
|
||||
// leaves the chat's context untouched.
|
||||
t.Run("SkipsRepinWhenNoPriorAgent", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
fix := newRebindFixture(t)
|
||||
|
||||
chat := dbgen.Chat(t, fix.db, database.Chat{
|
||||
OwnerID: fix.user.ID,
|
||||
OrganizationID: fix.org.ID,
|
||||
LastModelConfigID: fix.model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: fix.ws.ID, Valid: true},
|
||||
Status: database.ChatStatusWaiting,
|
||||
})
|
||||
require.False(t, chat.AgentID.Valid, "chat starts with no bound agent")
|
||||
|
||||
wc := newRebindTurnContext(t, fix.db, chat)
|
||||
updated, err := wc.persistBuildAgentBinding(fix.ctx, chat, fix.buildID, fix.agentB)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, fix.agentB, updated.AgentID.UUID, "the binding commits the new agent")
|
||||
|
||||
post, err := fix.db.GetChatByID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, post.ContextAggregateHash, "first-time binding does not re-pin a hash")
|
||||
postRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, postRes, "first-time binding copies no resources via the rebind path")
|
||||
})
|
||||
|
||||
// ClearsContextWhenNewAgentHasNoSnapshot: rebinding to an agent that has
|
||||
// not pushed a snapshot yet clears the chat's pinned hash and resources
|
||||
// (repinChatContext's no-snapshot branch).
|
||||
t.Run("ClearsContextWhenNewAgentHasNoSnapshot", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
fix := newRebindFixture(t)
|
||||
|
||||
chat := dbgen.Chat(t, fix.db, database.Chat{
|
||||
OwnerID: fix.user.ID,
|
||||
OrganizationID: fix.org.ID,
|
||||
LastModelConfigID: fix.model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: fix.ws.ID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: fix.agentA, Valid: true},
|
||||
Status: database.ChatStatusWaiting,
|
||||
})
|
||||
require.NoError(t, fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: fix.agentA,
|
||||
AggregateHash: fix.hashA,
|
||||
}))
|
||||
preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, preRes, 1, "chat starts pinned to agent A")
|
||||
|
||||
wc := newRebindTurnContext(t, fix.db, chat)
|
||||
updated, err := wc.persistBuildAgentBinding(fix.ctx, chat, fix.buildID, fix.agentNoSnap)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, fix.agentNoSnap, updated.AgentID.UUID)
|
||||
|
||||
post, err := fix.db.GetChatByID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, post.ContextAggregateHash, "rebinding to an agent with no snapshot clears the pinned hash")
|
||||
postRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, postRes, "rebinding to an agent with no snapshot clears the pinned resources")
|
||||
})
|
||||
|
||||
// SwallowsRepinError: a re-pin failure is logged and swallowed so it never
|
||||
// fails the agent binding itself.
|
||||
t.Run("SwallowsRepinError", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
dbm := dbmock.NewMockStore(ctrl)
|
||||
server := &Server{db: dbm, logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})}
|
||||
|
||||
chatID := uuid.New()
|
||||
priorAgent := uuid.New()
|
||||
newAgent := uuid.New()
|
||||
buildID := uuid.New()
|
||||
boundChat := database.Chat{
|
||||
ID: chatID,
|
||||
BuildID: uuid.NullUUID{UUID: buildID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: newAgent, Valid: true},
|
||||
}
|
||||
|
||||
dbm.EXPECT().UpdateChatBuildAgentBinding(gomock.Any(), database.UpdateChatBuildAgentBindingParams{
|
||||
ID: chatID,
|
||||
BuildID: uuid.NullUUID{UUID: buildID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: newAgent, Valid: true},
|
||||
}).Return(boundChat, nil)
|
||||
// The re-pin runs inside ReadModifyUpdate; drive the closure and fail
|
||||
// its first read so repinChatContext returns an error.
|
||||
dbm.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(f func(database.Store) error, _ *database.TxOptions) error {
|
||||
return f(dbm)
|
||||
})
|
||||
dbm.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), newAgent).
|
||||
Return(database.WorkspaceAgentContextSnapshot{}, xerrors.New("boom"))
|
||||
|
||||
cur := database.Chat{ID: chatID}
|
||||
wc := turnWorkspaceContext{
|
||||
server: server,
|
||||
chatStateMu: &sync.Mutex{},
|
||||
currentChat: &cur,
|
||||
loadChatSnapshot: func(context.Context, uuid.UUID) (database.Chat, error) { return database.Chat{}, nil },
|
||||
}
|
||||
t.Cleanup(wc.close)
|
||||
|
||||
prior := database.Chat{ID: chatID, AgentID: uuid.NullUUID{UUID: priorAgent, Valid: true}}
|
||||
updated, err := wc.persistBuildAgentBinding(ctx, prior, buildID, newAgent)
|
||||
require.NoError(t, err, "a re-pin failure must not fail the binding")
|
||||
require.Equal(t, newAgent, updated.AgentID.UUID)
|
||||
require.Equal(t, boundChat, wc.currentChatSnapshot(), "the binding still commits the new agent")
|
||||
})
|
||||
}
|
||||
|
||||
type rebindFixture struct {
|
||||
db database.Store
|
||||
ctx context.Context
|
||||
org database.Organization
|
||||
user database.User
|
||||
model database.ChatModelConfig
|
||||
ws database.WorkspaceTable
|
||||
buildID uuid.UUID
|
||||
agentA uuid.UUID
|
||||
agentB uuid.UUID
|
||||
agentNoSnap uuid.UUID
|
||||
hashA []byte
|
||||
hashB []byte
|
||||
srcA string
|
||||
srcB string
|
||||
bodyB json.RawMessage
|
||||
}
|
||||
|
||||
// newRebindFixture seeds a workspace with agents A and B that each have a
|
||||
// pushed context snapshot and one resource (so a chat can be pinned to A and
|
||||
// rebound to B), plus a third agent that never pushes a snapshot (for the
|
||||
// clear-only re-pin path). The agents share one build/resource because the
|
||||
// rebind guard keys on the agent, not the build.
|
||||
func newRebindFixture(t *testing.T) rebindFixture {
|
||||
t.Helper()
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
})
|
||||
tmpl := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
ActiveVersionID: tv.ID,
|
||||
CreatedBy: user.ID,
|
||||
})
|
||||
ws := dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: tmpl.ID,
|
||||
})
|
||||
pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
|
||||
OrganizationID: org.ID,
|
||||
CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()},
|
||||
})
|
||||
build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
|
||||
WorkspaceID: ws.ID,
|
||||
TemplateVersionID: tv.ID,
|
||||
JobID: pj.ID,
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
})
|
||||
res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
JobID: pj.ID,
|
||||
})
|
||||
agentA := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
|
||||
agentB := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
|
||||
// A third agent that never pushes a snapshot, for the clear-only re-pin path.
|
||||
agentNoSnap := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID})
|
||||
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{})
|
||||
|
||||
fix := rebindFixture{
|
||||
db: db,
|
||||
ctx: ctx,
|
||||
org: org,
|
||||
user: user,
|
||||
model: model,
|
||||
ws: ws,
|
||||
buildID: build.ID,
|
||||
agentA: agentA.ID,
|
||||
agentB: agentB.ID,
|
||||
agentNoSnap: agentNoSnap.ID,
|
||||
hashA: []byte{0xa1, 0xa2},
|
||||
hashB: []byte{0xb1, 0xb2},
|
||||
srcA: "/home/coder/workspace/AGENTS.md",
|
||||
srcB: "/home/coder/workspace/.agents/skills/example/SKILL.md",
|
||||
bodyB: json.RawMessage(`{"skill":{"name":"example"}}`),
|
||||
}
|
||||
|
||||
seedAgentContext(ctx, t, db, fix.agentA, fix.srcA, fix.hashA,
|
||||
database.WorkspaceAgentContextBodyKindInstructionFile, json.RawMessage(`{"instruction_file":{"content":"agent-a"}}`))
|
||||
seedAgentContext(ctx, t, db, fix.agentB, fix.srcB, fix.hashB,
|
||||
database.WorkspaceAgentContextBodyKindSkill, fix.bodyB)
|
||||
return fix
|
||||
}
|
||||
|
||||
func seedAgentContext(ctx context.Context, t *testing.T, db database.Store, agentID uuid.UUID, source string, hash []byte, kind database.WorkspaceAgentContextBodyKind, body json.RawMessage) {
|
||||
t.Helper()
|
||||
now := dbtime.Now()
|
||||
_, err := db.UpsertWorkspaceAgentContextSnapshot(ctx, database.UpsertWorkspaceAgentContextSnapshotParams{
|
||||
WorkspaceAgentID: agentID,
|
||||
Version: 1,
|
||||
AggregateHash: hash,
|
||||
ReceivedAt: now,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertWorkspaceAgentContextResource(ctx, database.UpsertWorkspaceAgentContextResourceParams{
|
||||
WorkspaceAgentID: agentID,
|
||||
Source: source,
|
||||
BodyKind: kind,
|
||||
Body: body,
|
||||
ContentHash: hash,
|
||||
SizeBytes: int64(len(body)),
|
||||
Status: database.WorkspaceAgentContextResourceStatusOk,
|
||||
Now: now,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func newRebindTurnContext(t *testing.T, db database.Store, chat database.Chat) *turnWorkspaceContext {
|
||||
t.Helper()
|
||||
server := &Server{db: db, logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})}
|
||||
cur := chat
|
||||
wc := &turnWorkspaceContext{
|
||||
server: server,
|
||||
chatStateMu: &sync.Mutex{},
|
||||
currentChat: &cur,
|
||||
loadChatSnapshot: db.GetChatByID,
|
||||
}
|
||||
t.Cleanup(wc.close)
|
||||
return wc
|
||||
}
|
||||
Reference in New Issue
Block a user