diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 4057544d75..fabb470cb2 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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 diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index f8200119ba..181ce84c8f 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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, diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index cafb9fdc9e..d760236922 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 565c3a6051..81389ab649 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 2c4d46556a..55db066a02 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -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) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3730a6c448..e6a49465dd 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -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, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 10f9f993b2..6c98afc337 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -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 diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index a10a69abd9..33fc87f33b 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -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 } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 59cecf4c27..801cd095ec 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -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() diff --git a/coderd/x/chatd/context_hydration.go b/coderd/x/chatd/context_hydration.go index 85621524a5..5e409a8f65 100644 --- a/coderd/x/chatd/context_hydration.go +++ b/coderd/x/chatd/context_hydration.go @@ -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) diff --git a/coderd/x/chatd/context_hydration_internal_test.go b/coderd/x/chatd/context_hydration_internal_test.go index 90f315de97..1a6a1a82de 100644 --- a/coderd/x/chatd/context_hydration_internal_test.go +++ b/coderd/x/chatd/context_hydration_internal_test.go @@ -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) diff --git a/coderd/x/chatd/context_integration_test.go b/coderd/x/chatd/context_integration_test.go index 1ac9f4bc5e..75c8e197be 100644 --- a/coderd/x/chatd/context_integration_test.go +++ b/coderd/x/chatd/context_integration_test.go @@ -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) diff --git a/coderd/x/chatd/context_rebind_internal_test.go b/coderd/x/chatd/context_rebind_internal_test.go new file mode 100644 index 0000000000..4c7d62e92f --- /dev/null +++ b/coderd/x/chatd/context_rebind_internal_test.go @@ -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 +}