mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add workspace skills to agent chat slash menu (#25600)
> This Pull Request was updated by Mux working on behalf of Mike. Adds workspace skills to the agent chat slash menu, sourced entirely from the chat's pinned context resources (the single-chat GET response the page already fetches), the same inventory `read_skill` resolves from. No new API endpoint is introduced. Personal entries insert `/name`, or `/personal/name` when the name collides with a workspace skill or the chat's pinned context has not resolved yet; workspace entries insert `/workspace/name`. Qualified aliases stay searchable even when the displayed trigger is bare. Before a chat binds a workspace (new chat form, or a selected but unbound workspace), the menu lists personal skills only. Sending a message invalidates the chat detail query, and chatd broadcasts a context watch event when a first-turn bind pins the chat, so the menu picks up newly pinned context without a reload. Makes `UpdateChatWorkspaceBinding` a no-op when the requested workspace/build/agent binding is unchanged, preserving `updated_at` so chat list ordering and watch events stay stable. Includes regression coverage for the no-op binding guard, pinned-context skill mapping, collision qualification, and skills menu behavior. Refs [CODAGT-474](https://linear.app/codercom/issue/CODAGT-474/ux-improvements-for-coder-agents) (skills autocompleting in the editor).
This commit is contained in:
@@ -5828,12 +5828,12 @@ func (q *querier) HasTemplateVersionsUsingCachedModuleFileInOrg(ctx context.Cont
|
||||
return q.db.HasTemplateVersionsUsingCachedModuleFileInOrg(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error {
|
||||
func (q *querier) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) {
|
||||
// System-level operation: an agent context push fans hydration out
|
||||
// across every not-yet-pinned chat for the agent, so it authorizes at
|
||||
// the resource level rather than per-chat.
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return err
|
||||
return nil, err
|
||||
}
|
||||
return q.db.HydrateAgentChatsContext(ctx, arg)
|
||||
}
|
||||
|
||||
@@ -562,8 +562,9 @@ func (s *MethodTestSuite) TestConnectionLogs() {
|
||||
func (s *MethodTestSuite) TestChats() {
|
||||
s.Run("HydrateAgentChatsContext", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.HydrateAgentChatsContextParams{AgentID: uuid.New()}
|
||||
dbm.EXPECT().HydrateAgentChatsContext(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate)
|
||||
hydrated := []uuid.UUID{uuid.New()}
|
||||
dbm.EXPECT().HydrateAgentChatsContext(gomock.Any(), arg).Return(hydrated, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(hydrated)
|
||||
}))
|
||||
s.Run("MarkChatsContextDirtyByAgent", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.MarkChatsContextDirtyByAgentParams{AgentID: uuid.New()}
|
||||
|
||||
+3
-3
@@ -3937,12 +3937,12 @@ func (m queryMetricsStore) HasTemplateVersionsUsingCachedModuleFileInOrg(ctx con
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error {
|
||||
func (m queryMetricsStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) {
|
||||
start := time.Now()
|
||||
r0 := m.s.HydrateAgentChatsContext(ctx, arg)
|
||||
r0, r1 := m.s.HydrateAgentChatsContext(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("HydrateAgentChatsContext").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "HydrateAgentChatsContext").Inc()
|
||||
return r0
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) {
|
||||
|
||||
Generated
+4
-3
@@ -7364,11 +7364,12 @@ func (mr *MockStoreMockRecorder) HasTemplateVersionsUsingCachedModuleFileInOrg(c
|
||||
}
|
||||
|
||||
// HydrateAgentChatsContext mocks base method.
|
||||
func (m *MockStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error {
|
||||
func (m *MockStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "HydrateAgentChatsContext", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
ret0, _ := ret[0].([]uuid.UUID)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// HydrateAgentChatsContext indicates an expected call of HydrateAgentChatsContext.
|
||||
|
||||
Generated
+4
-2
@@ -1012,11 +1012,13 @@ type sqlcQuerier interface {
|
||||
// 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
|
||||
// without a dirty marker. 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
|
||||
// Returns the hydrated chat IDs so callers can notify watchers of every
|
||||
// chat the statement pinned.
|
||||
HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) ([]uuid.UUID, error)
|
||||
// Increments generation_attempt and returns the resulting value.
|
||||
IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error)
|
||||
// Adds cost_micros to the spend for (user_id, effective_group_id, day).
|
||||
|
||||
@@ -1296,11 +1296,14 @@ func TestChatContextHydration(t *testing.T) {
|
||||
_, err := db.ArchiveChatByID(ctx, chatArchived.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Hydrate stamps only the NULL-hash chat for this agent.
|
||||
require.NoError(t, db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
// Hydrate stamps only the NULL-hash chat for this agent and returns
|
||||
// exactly the chats it pinned.
|
||||
hydrated, err := db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: agent.ID,
|
||||
AggregateHash: hashH,
|
||||
}))
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []uuid.UUID{chatNull.ID}, hydrated)
|
||||
gotNull, err := db.GetChatByID(ctx, chatNull.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, hashH, gotNull.ContextAggregateHash, "NULL-hash chat is hydrated")
|
||||
@@ -13505,6 +13508,92 @@ func TestUpdateChatLastTurnSummary(t *testing.T) {
|
||||
require.NotEqual(t, chat.HistoryVersion, fetched.HistoryVersion)
|
||||
}
|
||||
|
||||
func TestUpdateChatWorkspaceBindingNoOp(t *testing.T) {
|
||||
t.Parallel()
|
||||
if testing.Short() {
|
||||
t.SkipNow()
|
||||
}
|
||||
|
||||
sqlDB := testSQLDB(t)
|
||||
err := migrations.Up(sqlDB)
|
||||
require.NoError(t, err)
|
||||
db := database.New(sqlDB)
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
owner := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: owner.ID, OrganizationID: org.ID})
|
||||
|
||||
dbgen.ChatProvider(t, db, database.ChatProvider{
|
||||
Provider: "openai",
|
||||
DisplayName: "OpenAI",
|
||||
APIKey: "test-key",
|
||||
Enabled: true,
|
||||
CentralApiKeyEnabled: true,
|
||||
})
|
||||
|
||||
modelCfg, err := insertChatModelConfigForTest(ctx, t, db, "openai", database.InsertChatModelConfigParams{
|
||||
Model: "test-model",
|
||||
DisplayName: "Test Model",
|
||||
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
||||
Enabled: true,
|
||||
IsDefault: true,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 80,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
chat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||
OrganizationID: org.ID,
|
||||
Status: database.ChatStatusWaiting,
|
||||
ClientType: database.ChatClientTypeUi,
|
||||
OwnerID: owner.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: "binding-chat",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
template := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: owner.ID,
|
||||
})
|
||||
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OwnerID: owner.ID,
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: template.ID,
|
||||
})
|
||||
workspaceID := workspace.ID
|
||||
|
||||
bound, err := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
|
||||
ID: chat.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, workspaceID, bound.WorkspaceID.UUID)
|
||||
require.False(t, bound.UpdatedAt.Before(chat.UpdatedAt))
|
||||
|
||||
// Rebinding to the same workspace/build/agent is a no-op and must
|
||||
// preserve updated_at so chat list ordering and watch events stay
|
||||
// stable.
|
||||
rebound, err := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
|
||||
ID: chat.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, workspaceID, rebound.WorkspaceID.UUID)
|
||||
require.Equal(t, bound.UpdatedAt, rebound.UpdatedAt)
|
||||
|
||||
// Clearing the binding is a real change and must advance updated_at.
|
||||
cleared, err := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{
|
||||
ID: chat.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.False(t, cleared.WorkspaceID.Valid)
|
||||
require.True(t, cleared.UpdatedAt.After(bound.UpdatedAt))
|
||||
}
|
||||
|
||||
func TestDeleteChatDebugDataAfterMessageIDIncludesTriggeredRuns(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Generated
+126
-82
@@ -9489,41 +9489,44 @@ func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGrou
|
||||
return limit_micros, err
|
||||
}
|
||||
|
||||
const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :exec
|
||||
const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :many
|
||||
WITH hydrated AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
context_aggregate_hash = $2,
|
||||
context_error = $3
|
||||
WHERE agent_id = $1::uuid
|
||||
context_aggregate_hash = $1,
|
||||
context_error = $2
|
||||
WHERE agent_id = $3::uuid
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL
|
||||
RETURNING id
|
||||
),
|
||||
copied AS (
|
||||
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 = $3::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()
|
||||
)
|
||||
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()
|
||||
SELECT id FROM hydrated
|
||||
`
|
||||
|
||||
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
|
||||
@@ -9532,13 +9535,33 @@ type HydrateAgentChatsContextParams struct {
|
||||
// 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
|
||||
// without a dirty marker. 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.AgentID, arg.AggregateHash, arg.ContextError)
|
||||
return err
|
||||
// Returns the hydrated chat IDs so callers can notify watchers of every
|
||||
// chat the statement pinned.
|
||||
func (q *sqlQuerier) HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) ([]uuid.UUID, error) {
|
||||
rows, err := q.db.QueryContext(ctx, hydrateAgentChatsContext, arg.AggregateHash, arg.ContextError, arg.AgentID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []uuid.UUID
|
||||
for rows.Next() {
|
||||
var id uuid.UUID
|
||||
if err := rows.Scan(&id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, id)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const incrementChatGenerationAttempt = `-- name: IncrementChatGenerationAttempt :one
|
||||
@@ -12341,83 +12364,104 @@ func (q *sqlQuerier) UpdateChatTitleByID(ctx context.Context, arg UpdateChatTitl
|
||||
}
|
||||
|
||||
const updateChatWorkspaceBinding = `-- name: UpdateChatWorkspaceBinding :one
|
||||
WITH updated_chat AS (
|
||||
UPDATE chats SET
|
||||
workspace_id = $1::uuid,
|
||||
build_id = $2::uuid,
|
||||
agent_id = $3::uuid,
|
||||
updated_at = NOW()
|
||||
WHERE id = $4::uuid
|
||||
RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort
|
||||
WITH current_chat AS (
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort
|
||||
FROM chats
|
||||
WHERE id = $1::uuid
|
||||
),
|
||||
binding_changed AS (
|
||||
SELECT
|
||||
workspace_id IS DISTINCT FROM $2::uuid
|
||||
OR build_id IS DISTINCT FROM $3::uuid
|
||||
OR agent_id IS DISTINCT FROM $4::uuid AS changed
|
||||
FROM current_chat
|
||||
),
|
||||
changed_chat AS (
|
||||
UPDATE chats SET
|
||||
workspace_id = $2::uuid,
|
||||
build_id = $3::uuid,
|
||||
agent_id = $4::uuid,
|
||||
updated_at = NOW()
|
||||
WHERE id = $1::uuid
|
||||
AND (SELECT changed FROM binding_changed)
|
||||
RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort
|
||||
),
|
||||
result_chat AS (
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort
|
||||
FROM changed_chat
|
||||
UNION ALL
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort
|
||||
FROM current_chat
|
||||
WHERE NOT (SELECT changed FROM binding_changed)
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
updated_chat.id,
|
||||
updated_chat.owner_id,
|
||||
updated_chat.workspace_id,
|
||||
updated_chat.title,
|
||||
updated_chat.status,
|
||||
updated_chat.worker_id,
|
||||
updated_chat.started_at,
|
||||
updated_chat.heartbeat_at,
|
||||
updated_chat.created_at,
|
||||
updated_chat.updated_at,
|
||||
updated_chat.parent_chat_id,
|
||||
updated_chat.root_chat_id,
|
||||
updated_chat.last_model_config_id,
|
||||
updated_chat.last_reasoning_effort,
|
||||
updated_chat.archived,
|
||||
updated_chat.last_error,
|
||||
updated_chat.mode,
|
||||
updated_chat.mcp_server_ids,
|
||||
updated_chat.labels,
|
||||
updated_chat.build_id,
|
||||
updated_chat.agent_id,
|
||||
updated_chat.pin_order,
|
||||
updated_chat.last_read_message_id,
|
||||
updated_chat.dynamic_tools,
|
||||
updated_chat.organization_id,
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
result_chat.id,
|
||||
result_chat.owner_id,
|
||||
result_chat.workspace_id,
|
||||
result_chat.title,
|
||||
result_chat.status,
|
||||
result_chat.worker_id,
|
||||
result_chat.started_at,
|
||||
result_chat.heartbeat_at,
|
||||
result_chat.created_at,
|
||||
result_chat.updated_at,
|
||||
result_chat.parent_chat_id,
|
||||
result_chat.root_chat_id,
|
||||
result_chat.last_model_config_id,
|
||||
result_chat.last_reasoning_effort,
|
||||
result_chat.archived,
|
||||
result_chat.last_error,
|
||||
result_chat.mode,
|
||||
result_chat.mcp_server_ids,
|
||||
result_chat.labels,
|
||||
result_chat.build_id,
|
||||
result_chat.agent_id,
|
||||
result_chat.pin_order,
|
||||
result_chat.last_read_message_id,
|
||||
result_chat.dynamic_tools,
|
||||
result_chat.organization_id,
|
||||
result_chat.plan_mode,
|
||||
result_chat.client_type,
|
||||
result_chat.last_turn_summary,
|
||||
result_chat.snapshot_version,
|
||||
result_chat.history_version,
|
||||
result_chat.queue_version,
|
||||
result_chat.generation_attempt,
|
||||
result_chat.retry_state,
|
||||
result_chat.retry_state_version,
|
||||
result_chat.runner_id,
|
||||
result_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, result_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, result_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name,
|
||||
updated_chat.context_aggregate_hash,
|
||||
updated_chat.context_dirty_since,
|
||||
updated_chat.context_dirty_resources,
|
||||
updated_chat.context_error
|
||||
result_chat.context_aggregate_hash,
|
||||
result_chat.context_dirty_since,
|
||||
result_chat.context_dirty_resources,
|
||||
result_chat.context_error
|
||||
FROM
|
||||
updated_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = updated_chat.owner_id
|
||||
result_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(result_chat.root_chat_id, result_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = result_chat.owner_id
|
||||
)
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error
|
||||
FROM chats_expanded
|
||||
`
|
||||
|
||||
type UpdateChatWorkspaceBindingParams struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"`
|
||||
BuildID uuid.NullUUID `db:"build_id" json:"build_id"`
|
||||
AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"`
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpdateChatWorkspaceBinding(ctx context.Context, arg UpdateChatWorkspaceBindingParams) (Chat, error) {
|
||||
row := q.db.QueryRowContext(ctx, updateChatWorkspaceBinding,
|
||||
arg.ID,
|
||||
arg.WorkspaceID,
|
||||
arg.BuildID,
|
||||
arg.AgentID,
|
||||
arg.ID,
|
||||
)
|
||||
var i Chat
|
||||
err := row.Scan(
|
||||
|
||||
@@ -1264,65 +1264,86 @@ SELECT *
|
||||
FROM chats_expanded;
|
||||
|
||||
-- name: UpdateChatWorkspaceBinding :one
|
||||
WITH updated_chat AS (
|
||||
UPDATE chats SET
|
||||
workspace_id = sqlc.narg('workspace_id')::uuid,
|
||||
build_id = sqlc.narg('build_id')::uuid,
|
||||
agent_id = sqlc.narg('agent_id')::uuid,
|
||||
updated_at = NOW()
|
||||
WHERE id = @id::uuid
|
||||
RETURNING *
|
||||
WITH current_chat AS (
|
||||
SELECT *
|
||||
FROM chats
|
||||
WHERE id = @id::uuid
|
||||
),
|
||||
binding_changed AS (
|
||||
SELECT
|
||||
workspace_id IS DISTINCT FROM sqlc.narg('workspace_id')::uuid
|
||||
OR build_id IS DISTINCT FROM sqlc.narg('build_id')::uuid
|
||||
OR agent_id IS DISTINCT FROM sqlc.narg('agent_id')::uuid AS changed
|
||||
FROM current_chat
|
||||
),
|
||||
changed_chat AS (
|
||||
UPDATE chats SET
|
||||
workspace_id = sqlc.narg('workspace_id')::uuid,
|
||||
build_id = sqlc.narg('build_id')::uuid,
|
||||
agent_id = sqlc.narg('agent_id')::uuid,
|
||||
updated_at = NOW()
|
||||
WHERE id = @id::uuid
|
||||
AND (SELECT changed FROM binding_changed)
|
||||
RETURNING *
|
||||
),
|
||||
result_chat AS (
|
||||
SELECT *
|
||||
FROM changed_chat
|
||||
UNION ALL
|
||||
SELECT *
|
||||
FROM current_chat
|
||||
WHERE NOT (SELECT changed FROM binding_changed)
|
||||
),
|
||||
chats_expanded AS (
|
||||
SELECT
|
||||
updated_chat.id,
|
||||
updated_chat.owner_id,
|
||||
updated_chat.workspace_id,
|
||||
updated_chat.title,
|
||||
updated_chat.status,
|
||||
updated_chat.worker_id,
|
||||
updated_chat.started_at,
|
||||
updated_chat.heartbeat_at,
|
||||
updated_chat.created_at,
|
||||
updated_chat.updated_at,
|
||||
updated_chat.parent_chat_id,
|
||||
updated_chat.root_chat_id,
|
||||
updated_chat.last_model_config_id,
|
||||
updated_chat.last_reasoning_effort,
|
||||
updated_chat.archived,
|
||||
updated_chat.last_error,
|
||||
updated_chat.mode,
|
||||
updated_chat.mcp_server_ids,
|
||||
updated_chat.labels,
|
||||
updated_chat.build_id,
|
||||
updated_chat.agent_id,
|
||||
updated_chat.pin_order,
|
||||
updated_chat.last_read_message_id,
|
||||
updated_chat.dynamic_tools,
|
||||
updated_chat.organization_id,
|
||||
updated_chat.plan_mode,
|
||||
updated_chat.client_type,
|
||||
updated_chat.last_turn_summary,
|
||||
updated_chat.snapshot_version,
|
||||
updated_chat.history_version,
|
||||
updated_chat.queue_version,
|
||||
updated_chat.generation_attempt,
|
||||
updated_chat.retry_state,
|
||||
updated_chat.retry_state_version,
|
||||
updated_chat.runner_id,
|
||||
updated_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl,
|
||||
result_chat.id,
|
||||
result_chat.owner_id,
|
||||
result_chat.workspace_id,
|
||||
result_chat.title,
|
||||
result_chat.status,
|
||||
result_chat.worker_id,
|
||||
result_chat.started_at,
|
||||
result_chat.heartbeat_at,
|
||||
result_chat.created_at,
|
||||
result_chat.updated_at,
|
||||
result_chat.parent_chat_id,
|
||||
result_chat.root_chat_id,
|
||||
result_chat.last_model_config_id,
|
||||
result_chat.last_reasoning_effort,
|
||||
result_chat.archived,
|
||||
result_chat.last_error,
|
||||
result_chat.mode,
|
||||
result_chat.mcp_server_ids,
|
||||
result_chat.labels,
|
||||
result_chat.build_id,
|
||||
result_chat.agent_id,
|
||||
result_chat.pin_order,
|
||||
result_chat.last_read_message_id,
|
||||
result_chat.dynamic_tools,
|
||||
result_chat.organization_id,
|
||||
result_chat.plan_mode,
|
||||
result_chat.client_type,
|
||||
result_chat.last_turn_summary,
|
||||
result_chat.snapshot_version,
|
||||
result_chat.history_version,
|
||||
result_chat.queue_version,
|
||||
result_chat.generation_attempt,
|
||||
result_chat.retry_state,
|
||||
result_chat.retry_state_version,
|
||||
result_chat.runner_id,
|
||||
result_chat.requires_action_deadline_at,
|
||||
COALESCE(root.user_acl, result_chat.user_acl) AS user_acl,
|
||||
COALESCE(root.group_acl, result_chat.group_acl) AS group_acl,
|
||||
owner.username AS owner_username,
|
||||
owner.name AS owner_name,
|
||||
updated_chat.context_aggregate_hash,
|
||||
updated_chat.context_dirty_since,
|
||||
updated_chat.context_dirty_resources,
|
||||
updated_chat.context_error
|
||||
result_chat.context_aggregate_hash,
|
||||
result_chat.context_dirty_since,
|
||||
result_chat.context_dirty_resources,
|
||||
result_chat.context_error
|
||||
FROM
|
||||
updated_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = updated_chat.owner_id
|
||||
result_chat
|
||||
LEFT JOIN chats root ON root.id = COALESCE(result_chat.root_chat_id, result_chat.parent_chat_id)
|
||||
JOIN visible_users owner ON owner.id = result_chat.owner_id
|
||||
)
|
||||
SELECT *
|
||||
FROM chats_expanded;
|
||||
@@ -1485,17 +1506,19 @@ SET
|
||||
context_dirty_since = NULL
|
||||
WHERE id = @id::uuid;
|
||||
|
||||
-- name: HydrateAgentChatsContext :exec
|
||||
-- name: HydrateAgentChatsContext :many
|
||||
-- Stamps the pinned hash and error on every not-yet-hydrated chat for
|
||||
-- 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
|
||||
-- without a dirty marker. 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.
|
||||
-- Returns the hydrated chat IDs so callers can notify watchers of every
|
||||
-- chat the statement pinned.
|
||||
WITH hydrated AS (
|
||||
UPDATE chats
|
||||
SET
|
||||
@@ -1505,25 +1528,28 @@ WITH hydrated AS (
|
||||
AND archived = false
|
||||
AND context_aggregate_hash IS NULL
|
||||
RETURNING id
|
||||
),
|
||||
copied AS (
|
||||
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()
|
||||
)
|
||||
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();
|
||||
SELECT id FROM hydrated;
|
||||
|
||||
-- name: MarkChatsContextDirtyByAgent :many
|
||||
-- Flips active, already-hydrated chats for an agent to dirty when the
|
||||
|
||||
@@ -31,11 +31,13 @@ func latestAgentSnapshot(ctx context.Context, db database.Store, agentID uuid.UU
|
||||
|
||||
// HydrateAndMarkChatsDirty implements agentapi.ContextDirtyMarker. It runs
|
||||
// inside the PushContextState transaction: it stamps the pushed snapshot hash
|
||||
// on chats for the agent that have not been hydrated yet (no dirty event),
|
||||
// then flips already-pinned chats whose hash differs to dirty. It returns a
|
||||
// callback that publishes the dirty watch events; the caller invokes it only
|
||||
// after the transaction commits, and the callback is a no-op when nothing
|
||||
// transitioned to dirty.
|
||||
// on chats for the agent that have not been hydrated yet, then flips
|
||||
// already-pinned chats whose hash differs to dirty. It returns a callback
|
||||
// that publishes a context watch event for every chat it touched; the caller
|
||||
// invokes it only after the transaction commits, and the callback is a no-op
|
||||
// when no chat was hydrated or dirtied. Hydrated chats start clean (no dirty
|
||||
// marker), but still need the event: watching clients cached their details
|
||||
// without pinned resources and refetch only on context events.
|
||||
//
|
||||
// The pinned hash on dirtied chats is intentionally left unchanged; the
|
||||
// refresh endpoint re-pins it.
|
||||
@@ -44,13 +46,13 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store
|
||||
ctx = dbauthz.AsChatd(ctx)
|
||||
|
||||
// Chats created before the agent's first push land with a NULL pinned
|
||||
// hash. Stamp them now so they start clean; this is their first
|
||||
// hydration, so no dirty event is emitted.
|
||||
if err := tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
// hash. Stamp them now so they start clean.
|
||||
hydrated, err := tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: agentID,
|
||||
AggregateHash: aggregateHash,
|
||||
ContextError: snapshotError,
|
||||
}); err != nil {
|
||||
})
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("hydrate agent chats context: %w", err)
|
||||
}
|
||||
|
||||
@@ -62,26 +64,33 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("mark chats context dirty: %w", err)
|
||||
}
|
||||
if len(dirtied) == 0 {
|
||||
// Hydrated chats had a NULL hash and dirtied chats a non-NULL one, so
|
||||
// the two sets never overlap.
|
||||
touched := make([]uuid.UUID, 0, len(hydrated)+len(dirtied))
|
||||
touched = append(touched, hydrated...)
|
||||
for _, d := range dirtied {
|
||||
touched = append(touched, d.ID)
|
||||
}
|
||||
if len(touched) == 0 {
|
||||
return func() {}, nil
|
||||
}
|
||||
|
||||
// Read the dirtied chats inside the transaction and capture their rows so
|
||||
// Read the touched chats inside the transaction and capture their rows so
|
||||
// the post-commit callback needs no database access: the published payload
|
||||
// reflects the just-committed dirty state (no re-read a concurrent refresh
|
||||
// reflects the just-committed state (no re-read a concurrent refresh
|
||||
// could race), and the callback does not depend on the request-scoped
|
||||
// context surviving past commit. Only the transitioned chats are read.
|
||||
dirtyChats := make([]database.Chat, 0, len(dirtied))
|
||||
for _, d := range dirtied {
|
||||
chat, err := tx.GetChatByID(ctx, d.ID)
|
||||
touchedChats := make([]database.Chat, 0, len(touched))
|
||||
for _, id := range touched {
|
||||
chat, err := tx.GetChatByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("get dirtied chat %s: %w", d.ID, err)
|
||||
return nil, xerrors.Errorf("get touched chat %s: %w", id, err)
|
||||
}
|
||||
dirtyChats = append(dirtyChats, chat)
|
||||
touchedChats = append(touchedChats, chat)
|
||||
}
|
||||
|
||||
return func() {
|
||||
p.publishChatPubsubEvents(dirtyChats, codersdk.ChatWatchEventKindContextDirty)
|
||||
p.publishChatPubsubEvents(touchedChats, codersdk.ChatWatchEventKindContextDirty)
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -94,22 +103,31 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store
|
||||
// HydrateAgentChatsContext only touches NULL-hash chats (a concurrent push that
|
||||
// already hydrated the chat is not clobbered), and snapshot-gated so it does
|
||||
// nothing when the agent has not pushed yet, never stamping empty state that
|
||||
// would keep a later push from hydrating.
|
||||
func (p *Server) hydrateAgentChatsFromSnapshot(ctx context.Context, agentID uuid.UUID) error {
|
||||
return database.ReadModifyUpdate(p.db, func(tx database.Store) error {
|
||||
// would keep a later push from hydrating. It returns the IDs of the chats it
|
||||
// pinned; empty when the agent has no snapshot or every chat was already
|
||||
// pinned.
|
||||
func (p *Server) hydrateAgentChatsFromSnapshot(ctx context.Context, agentID uuid.UUID) ([]uuid.UUID, error) {
|
||||
var hydrated []uuid.UUID
|
||||
err := database.ReadModifyUpdate(p.db, func(tx database.Store) error {
|
||||
aggregateHash, snapshotError, ok, err := latestAgentSnapshot(ctx, tx, agentID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !ok {
|
||||
hydrated = nil
|
||||
return nil
|
||||
}
|
||||
return tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
hydrated, err = tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: agentID,
|
||||
AggregateHash: aggregateHash,
|
||||
ContextError: snapshotError,
|
||||
})
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return hydrated, nil
|
||||
}
|
||||
|
||||
// hydrateChatContextOnCreate pins a newly created chat to its agent's latest
|
||||
@@ -125,7 +143,7 @@ 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)
|
||||
if err := p.hydrateAgentChatsFromSnapshot(ctx, chat.AgentID.UUID); err != nil {
|
||||
if _, err := p.hydrateAgentChatsFromSnapshot(ctx, chat.AgentID.UUID); err != nil {
|
||||
p.logger.Warn(ctx, "hydrate chat context on create",
|
||||
slog.F("chat_id", chat.ID), slog.Error(err))
|
||||
}
|
||||
@@ -140,20 +158,40 @@ func (p *Server) hydrateChatContextOnCreate(ctx context.Context, chat database.C
|
||||
// exist. It reuses the create-path hydration, which is idempotent and
|
||||
// snapshot-gated, so it never clobbers an already-pinned chat and never stamps
|
||||
// empty state. The NULL-hash gate also leaves dirtied chats alone: their stale
|
||||
// pinned hash is non-NULL until the refresh endpoint re-pins. Best-effort:
|
||||
// failures are logged and swallowed so they never fail the turn.
|
||||
// pinned hash is non-NULL until the refresh endpoint re-pins. Hydration
|
||||
// pins every unpinned chat bound to the agent in one statement, so a
|
||||
// context watch event is published for each pinned chat: watching clients
|
||||
// cached those chats' details without pinned resources and need to
|
||||
// refetch. Best-effort: failures are logged and swallowed so they never
|
||||
// fail the turn.
|
||||
func (p *Server) ensureChatContextPinnedOnFirstTurn(ctx context.Context, chat database.Chat) {
|
||||
if !chat.AgentID.Valid || chat.ContextAggregateHash != nil {
|
||||
return
|
||||
}
|
||||
//nolint:gocritic // Chatd stamps chats it does not own as the daemon subject.
|
||||
ctx = dbauthz.AsChatd(ctx)
|
||||
if err := p.hydrateAgentChatsFromSnapshot(ctx, chat.AgentID.UUID); err != nil {
|
||||
hydrated, err := p.hydrateAgentChatsFromSnapshot(ctx, chat.AgentID.UUID)
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "ensure chat context pinned on first turn",
|
||||
slog.F("chat_id", chat.ID),
|
||||
slog.F("agent_id", chat.AgentID.UUID),
|
||||
slog.Error(err))
|
||||
return
|
||||
}
|
||||
pinnedChats := make([]database.Chat, 0, len(hydrated))
|
||||
for _, chatID := range hydrated {
|
||||
// Re-read each chat so the event payload carries the pinned
|
||||
// state, not the pre-hydration row.
|
||||
pinned, err := p.db.GetChatByID(ctx, chatID)
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "read chat after first-turn context pin",
|
||||
slog.F("chat_id", chatID),
|
||||
slog.Error(err))
|
||||
continue
|
||||
}
|
||||
pinnedChats = append(pinnedChats, pinned)
|
||||
}
|
||||
p.publishChatPubsubEvents(pinnedChats, codersdk.ChatWatchEventKindContextDirty)
|
||||
}
|
||||
|
||||
// repinChatContext re-pins a single chat to its agent's latest context
|
||||
|
||||
@@ -1,15 +1,21 @@
|
||||
package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbmock"
|
||||
dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
@@ -46,7 +52,7 @@ func TestHydrateChatContextOnCreate(t *testing.T) {
|
||||
AgentID: agentID,
|
||||
AggregateHash: snapshot.AggregateHash,
|
||||
ContextError: snapshot.SnapshotError,
|
||||
}).Return(nil)
|
||||
}).Return([]uuid.UUID{chat.ID}, nil)
|
||||
|
||||
server.hydrateChatContextOnCreate(ctx, chat)
|
||||
})
|
||||
@@ -84,6 +90,62 @@ func TestHydrateChatContextOnCreate(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
// TestHydrateAndMarkChatsDirtyPublishesForHydratedAndDirtied covers the
|
||||
// agent-push path: a chat hydrated by the push (first pin, no dirty marker)
|
||||
// and a chat flipped to dirty must both get a context watch event, because
|
||||
// watching clients refetch pinned resources only on those events.
|
||||
func TestHydrateAndMarkChatsDirtyPublishesForHydratedAndDirtied(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
ps := dbpubsub.NewInMemory()
|
||||
server := &Server{db: db, logger: slogtest.Make(t, nil), pubsub: ps}
|
||||
|
||||
ownerID := uuid.New()
|
||||
agentID := uuid.New()
|
||||
hash := []byte{0x01}
|
||||
now := time.Now()
|
||||
|
||||
hydratedChat := database.Chat{ID: uuid.New(), OwnerID: ownerID, ContextAggregateHash: hash}
|
||||
dirtiedChat := database.Chat{ID: uuid.New(), OwnerID: ownerID, ContextAggregateHash: []byte{0x99}}
|
||||
|
||||
events := make(chan codersdk.ChatWatchEvent, 2)
|
||||
cancelSub, err := ps.SubscribeWithErr(
|
||||
coderdpubsub.ChatWatchEventChannel(ownerID),
|
||||
coderdpubsub.HandleChatWatchEvent(func(_ context.Context, payload codersdk.ChatWatchEvent, err error) {
|
||||
require.NoError(t, err)
|
||||
events <- payload
|
||||
}),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
defer cancelSub()
|
||||
|
||||
db.EXPECT().HydrateAgentChatsContext(gomock.Any(), database.HydrateAgentChatsContextParams{
|
||||
AgentID: agentID,
|
||||
AggregateHash: hash,
|
||||
}).Return([]uuid.UUID{hydratedChat.ID}, nil)
|
||||
db.EXPECT().MarkChatsContextDirtyByAgent(gomock.Any(), database.MarkChatsContextDirtyByAgentParams{
|
||||
AgentID: agentID,
|
||||
AggregateHash: hash,
|
||||
DirtySince: sql.NullTime{Time: now, Valid: true},
|
||||
}).Return([]database.MarkChatsContextDirtyByAgentRow{{ID: dirtiedChat.ID, OwnerID: ownerID}}, nil)
|
||||
db.EXPECT().GetChatByID(gomock.Any(), hydratedChat.ID).Return(hydratedChat, nil)
|
||||
db.EXPECT().GetChatByID(gomock.Any(), dirtiedChat.ID).Return(dirtiedChat, nil)
|
||||
|
||||
publish, err := server.HydrateAndMarkChatsDirty(ctx, db, agentID, hash, "", now)
|
||||
require.NoError(t, err)
|
||||
publish()
|
||||
|
||||
gotChatIDs := make([]uuid.UUID, 0, 2)
|
||||
for range 2 {
|
||||
event := testutil.RequireReceive(ctx, t, events)
|
||||
require.Equal(t, codersdk.ChatWatchEventKindContextDirty, event.Kind)
|
||||
gotChatIDs = append(gotChatIDs, event.Chat.ID)
|
||||
}
|
||||
require.ElementsMatch(t, []uuid.UUID{hydratedChat.ID, dirtiedChat.ID}, gotChatIDs)
|
||||
}
|
||||
|
||||
// TestEnsureChatContextPinnedOnFirstTurn covers the lazy-bind pinning path. An
|
||||
// API-created chat carries no agent at create, binds its agent on the first
|
||||
// turn, and must pin the agent's already-pushed snapshot then. This is the
|
||||
@@ -99,14 +161,34 @@ func TestEnsureChatContextPinnedOnFirstTurn(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
server := &Server{db: db, logger: slogtest.Make(t, nil)}
|
||||
ps := dbpubsub.NewInMemory()
|
||||
server := &Server{db: db, logger: slogtest.Make(t, nil), pubsub: ps}
|
||||
|
||||
ownerID := uuid.New()
|
||||
agentID := uuid.New()
|
||||
chat := database.Chat{ID: uuid.New(), AgentID: uuid.NullUUID{UUID: agentID, Valid: true}}
|
||||
chat := database.Chat{ID: uuid.New(), OwnerID: ownerID, AgentID: uuid.NullUUID{UUID: agentID, Valid: true}}
|
||||
// A second unpinned chat bound to the same agent is hydrated by the
|
||||
// same statement and must get its own watch event.
|
||||
siblingChat := database.Chat{ID: uuid.New(), OwnerID: ownerID, AgentID: chat.AgentID}
|
||||
snapshot := database.WorkspaceAgentContextSnapshot{
|
||||
WorkspaceAgentID: agentID,
|
||||
AggregateHash: []byte{0x0a, 0x0b},
|
||||
}
|
||||
pinnedChat := chat
|
||||
pinnedChat.ContextAggregateHash = snapshot.AggregateHash
|
||||
pinnedSibling := siblingChat
|
||||
pinnedSibling.ContextAggregateHash = snapshot.AggregateHash
|
||||
|
||||
events := make(chan codersdk.ChatWatchEvent, 2)
|
||||
cancelSub, err := ps.SubscribeWithErr(
|
||||
coderdpubsub.ChatWatchEventChannel(ownerID),
|
||||
coderdpubsub.HandleChatWatchEvent(func(_ context.Context, payload codersdk.ChatWatchEvent, err error) {
|
||||
require.NoError(t, err)
|
||||
events <- payload
|
||||
}),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
defer cancelSub()
|
||||
|
||||
db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(f func(database.Store) error, _ *database.TxOptions) error { return f(db) })
|
||||
@@ -118,9 +200,43 @@ func TestEnsureChatContextPinnedOnFirstTurn(t *testing.T) {
|
||||
AgentID: agentID,
|
||||
AggregateHash: snapshot.AggregateHash,
|
||||
ContextError: snapshot.SnapshotError,
|
||||
}).Return(nil)
|
||||
}).Return([]uuid.UUID{chat.ID, siblingChat.ID}, nil)
|
||||
db.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(pinnedChat, nil)
|
||||
db.EXPECT().GetChatByID(gomock.Any(), siblingChat.ID).Return(pinnedSibling, nil)
|
||||
|
||||
server.ensureChatContextPinnedOnFirstTurn(ctx, chat)
|
||||
|
||||
// Watching clients cached both details without pinned resources, so
|
||||
// every hydrated chat must broadcast a context event.
|
||||
gotChatIDs := make([]uuid.UUID, 0, 2)
|
||||
for range 2 {
|
||||
event := testutil.RequireReceive(ctx, t, events)
|
||||
require.Equal(t, codersdk.ChatWatchEventKindContextDirty, event.Kind)
|
||||
gotChatIDs = append(gotChatIDs, event.Chat.ID)
|
||||
}
|
||||
require.ElementsMatch(t, []uuid.UUID{chat.ID, siblingChat.ID}, gotChatIDs)
|
||||
})
|
||||
|
||||
t.Run("SkipsPublishWhenNoSnapshot", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
server := &Server{db: db, logger: slogtest.Make(t, nil)}
|
||||
|
||||
agentID := uuid.New()
|
||||
// ErrNoRows means the agent has not pushed yet: nothing is stamped
|
||||
// and no event is published (GetChatByID has no EXPECT, so a
|
||||
// post-hydration read would fail the test).
|
||||
db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn(
|
||||
func(f func(database.Store) error, _ *database.TxOptions) error { return f(db) })
|
||||
db.EXPECT().GetLatestWorkspaceAgentContextSnapshot(gomock.Any(), agentID).
|
||||
Return(database.WorkspaceAgentContextSnapshot{}, sql.ErrNoRows)
|
||||
|
||||
server.ensureChatContextPinnedOnFirstTurn(ctx, database.Chat{
|
||||
ID: uuid.New(),
|
||||
AgentID: uuid.NullUUID{UUID: agentID, Valid: true},
|
||||
})
|
||||
})
|
||||
|
||||
t.Run("SkipsWhenAlreadyPinned", func(t *testing.T) {
|
||||
|
||||
@@ -399,10 +399,11 @@ func TestPinnedWorkspaceContextFromHydratedPin(t *testing.T) {
|
||||
AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true},
|
||||
Status: database.ChatStatusWaiting,
|
||||
})
|
||||
require.NoError(t, db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
_, err := db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: agent.ID,
|
||||
AggregateHash: hash,
|
||||
}))
|
||||
})
|
||||
require.NoError(t, err)
|
||||
rows, err := db.ListChatContextResourcesByChatID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rows, 2, "the pin holds the agent's instruction file and skill")
|
||||
|
||||
@@ -48,10 +48,11 @@ func TestPersistBuildAgentBindingRepinsContext(t *testing.T) {
|
||||
// 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{
|
||||
_, err := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: fix.agentA,
|
||||
AggregateHash: fix.hashA,
|
||||
}))
|
||||
})
|
||||
require.NoError(t, err)
|
||||
preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, preRes, 1)
|
||||
@@ -121,10 +122,11 @@ func TestPersistBuildAgentBindingRepinsContext(t *testing.T) {
|
||||
AgentID: uuid.NullUUID{UUID: fix.agentA, Valid: true},
|
||||
Status: database.ChatStatusWaiting,
|
||||
})
|
||||
require.NoError(t, fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{
|
||||
_, err := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{
|
||||
AgentID: fix.agentA,
|
||||
AggregateHash: fix.hashA,
|
||||
}))
|
||||
})
|
||||
require.NoError(t, err)
|
||||
preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, preRes, 1, "chat starts pinned to agent A")
|
||||
|
||||
Reference in New Issue
Block a user