mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +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:
Kyle Carberry
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,
|
||||
|
||||
Reference in New Issue
Block a user