diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 5a18c5fb8f3..e3abd19503b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5793,12 +5793,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) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index d3263972e6b..02dc4a46829 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -548,8 +548,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()} diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 16ba3084ebd..017a0522cb6 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3921,12 +3921,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) { diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 67e42764a88..aa8507ab088 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7334,11 +7334,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. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index cd0a2bfa7fb..e63ecbe96c0 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1010,11 +1010,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). diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 8b526f8389e..c590dc06407 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -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() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index a65852cf8de..d85ccf8b10a 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -9444,41 +9444,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 @@ -9487,13 +9490,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 @@ -12296,83 +12319,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( diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 51bd02e34bd..8131fe50d67 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -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 diff --git a/coderd/x/chatd/context_hydration.go b/coderd/x/chatd/context_hydration.go index 147dbf36c75..4e53f99013c 100644 --- a/coderd/x/chatd/context_hydration.go +++ b/coderd/x/chatd/context_hydration.go @@ -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 diff --git a/coderd/x/chatd/context_hydration_internal_test.go b/coderd/x/chatd/context_hydration_internal_test.go index 7e71c053e31..5155af536d8 100644 --- a/coderd/x/chatd/context_hydration_internal_test.go +++ b/coderd/x/chatd/context_hydration_internal_test.go @@ -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) { diff --git a/coderd/x/chatd/context_prompt_internal_test.go b/coderd/x/chatd/context_prompt_internal_test.go index 278d8ea1e17..ff5184ad14e 100644 --- a/coderd/x/chatd/context_prompt_internal_test.go +++ b/coderd/x/chatd/context_prompt_internal_test.go @@ -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") diff --git a/coderd/x/chatd/context_rebind_internal_test.go b/coderd/x/chatd/context_rebind_internal_test.go index 4c7d62e92ff..1c8dc55502d 100644 --- a/coderd/x/chatd/context_rebind_internal_test.go +++ b/coderd/x/chatd/context_rebind_internal_test.go @@ -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") diff --git a/codersdk/chats.go b/codersdk/chats.go index aea0e0b4d7e..7e4cf105b76 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1866,8 +1866,11 @@ const ( ChatWatchEventKindDiffStatusChange ChatWatchEventKind = "diff_status_change" ChatWatchEventKindActionRequired ChatWatchEventKind = "action_required" // ChatWatchEventKindContextDirty signals that the chat's pinned - // workspace context drifted from the agent's latest pushed snapshot. - // The chat stays usable; a refresh re-pins it to the latest snapshot. + // workspace context changed: it drifted from the agent's latest + // pushed snapshot, or hydration first populated it (a first-turn + // pin or an agent push reaching a not-yet-pinned chat). The chat + // stays usable; a refresh re-pins a drifted chat to the latest + // snapshot. ChatWatchEventKindContextDirty ChatWatchEventKind = "context_dirty" ) diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 7789d4877be..19e3f92fc23 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -966,7 +966,7 @@ describe("mutation invalidation scope", () => { } }); - it("createChatMessage invalidates only debug runs, not chat detail or messages", async () => { + it("createChatMessage invalidates debug runs and chat detail, not messages", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; seedAllActiveQueries(queryClient, chatId); @@ -980,10 +980,9 @@ describe("mutation invalidation scope", () => { ).toBe(true); const chatState = queryClient.getQueryState(chatKey(chatId)); - expect( - chatState?.isInvalidated, - "chatKey should NOT be invalidated", - ).not.toBe(true); + expect(chatState?.isInvalidated, "chatKey should be invalidated").toBe( + true, + ); const messagesState = queryClient.getQueryState(chatMessagesKey(chatId)); expect( diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 281f04234b1..51f63b9fe1b 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1337,6 +1337,10 @@ export const createChatMessage = ( API.experimental.createChatMessage(chatId, req), onSuccess: () => { void invalidateChatDebugRuns(queryClient, chatId); + void queryClient.invalidateQueries({ + queryKey: chatKey(chatId), + exact: true, + }); void queryClient.invalidateQueries({ queryKey: chatPromptsKey(chatId), exact: true, diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 90364d68f3e..b38bc4380e0 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -85,6 +85,7 @@ import { } from "./components/ChatConversation/chatStore"; import { useChatToolInvalidations } from "./components/ChatConversation/useChatToolInvalidations"; import type { PendingAttachment } from "./components/ChatPageContent"; +import { workspaceSkillsFromChat } from "./components/ChatPageContent"; import { getDefaultMCPSelection, getSavedMCPSelection, @@ -1756,6 +1757,7 @@ const AgentChatPage: FC = () => { onMCPSelectionChange={handleMCPSelectionChange} onMCPAuthComplete={handleMCPAuthComplete} chatContext={chatQuery.data?.context} + workspaceSkills={workspaceSkillsFromChat(chatQuery.data)} /> ); }; diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index e61fcd12305..cbe48cc83a4 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -36,6 +36,7 @@ import { import type { useChatStore } from "./components/ChatConversation/chatStore"; import type { ModelSelectorOption } from "./components/ChatElements"; import { DesktopPanelContext } from "./components/ChatElements/tools/DesktopPanelContext"; +import type { SkillMetadata } from "./components/ChatMessageInput/SkillsTriggerMenu"; import type { PendingAttachment } from "./components/ChatPageContent"; import { ChatPageInput, ChatPageTimeline } from "./components/ChatPageContent"; import { ChatScrollContainer } from "./components/ChatScrollContainer"; @@ -222,6 +223,7 @@ interface AgentChatPageViewProps { desktopChatId?: string; chatContext?: TypesGen.ChatContext; + workspaceSkills?: readonly SkillMetadata[]; } const UnavailableTabMessage: FC<{ message: string }> = ({ message }) => ( @@ -386,6 +388,7 @@ export const AgentChatPageView: FC = ({ onMCPAuthComplete, desktopChatId, chatContext, + workspaceSkills, }) => { const queryClient = useQueryClient(); const { proxy } = useProxy(); @@ -972,6 +975,7 @@ export const AgentChatPageView: FC = ({ onMCPSelectionChange={onMCPSelectionChange} onMCPAuthComplete={onMCPAuthComplete} chatContext={chatContext} + workspaceSkills={workspaceSkills} workspace={workspace} workspaceAgent={workspaceAgent} chatId={agentId} diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 335434dacfb..fec0e6a7f7f 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -83,6 +83,7 @@ import { ChatMessageInput, type ChatMessageInputRef, } from "./ChatMessageInput/ChatMessageInput"; +import type { SkillMetadata } from "./ChatMessageInput/SkillsTriggerMenu"; import type { AgentContextUsage } from "./ContextUsageIndicator"; import { ContextUsageIndicator } from "./ContextUsageIndicator"; import { ImageLightbox } from "./ImageLightbox"; @@ -190,6 +191,7 @@ interface AgentChatInputProps { selectedMCPServerIds?: readonly string[]; onMCPSelectionChange?: (ids: string[]) => void; onMCPAuthComplete?: (serverId: string) => void; + workspaceSkills?: readonly SkillMetadata[]; workspace?: TypesGen.Workspace; workspaceAgent?: TypesGen.WorkspaceAgent; chatId?: string; @@ -398,6 +400,7 @@ export const AgentChatInput: FC = ({ selectedMCPServerIds, onMCPSelectionChange, onMCPAuthComplete, + workspaceSkills, workspace, workspaceAgent, chatId, @@ -582,6 +585,12 @@ export const AgentChatInput: FC = ({ }); }; + // Only a chat-bound workspace counts: an unbound selection (new chat + // form, or a picked workspace before the first send) has no pinned + // context to resolve, so treating it as a workspace would leave the + // menu in the loading state forever. + const hasSkillsWorkspace = Boolean(attachedWorkspace?.id ?? workspace?.id); + const selectedWorkspace = workspaceOptions?.find( (ws) => ws.id === selectedWorkspaceId, ); @@ -1210,6 +1219,8 @@ export const AgentChatInput: FC = ({ onEnter={handleSubmit} sendShortcut={sendShortcut} disabled={isDisabled || isLoading} + hasWorkspace={hasSkillsWorkspace} + workspaceSkills={workspaceSkills} autoFocus /> {/* Warn about invisible Unicode in the message text. diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.stories.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.stories.tsx index 70816ee8d14..a79e613b54b 100644 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.stories.tsx @@ -3,6 +3,7 @@ import { type PropsWithChildren, useEffect } from "react"; import { expect, fn, userEvent, waitFor, within } from "storybook/test"; import type * as TypesGen from "#/api/typesGenerated"; import { ChatMessageInput } from "./ChatMessageInput"; +import type { SkillMetadata } from "./SkillsTriggerMenu"; import { expectNoVisibleText, findVisibleText, @@ -10,6 +11,18 @@ import { MockSkills, } from "./storyHelpers"; +// Override props keep skill menu stories deterministic without network calls. +const mockWorkspaceSkills: SkillMetadata[] = [ + { + name: "test-runner", + description: "Run the workspace test command.", + }, + { + name: "workspace-docs", + description: "Use repository documentation conventions.", + }, +]; + const meta: Meta = { title: "components/ChatMessageInput/ChatMessageInput", component: ChatMessageInput, @@ -156,6 +169,110 @@ export const ClickSelectsSkill: Story = { }, }; +export const OpensWithPersonalAndWorkspaceSkills: Story = { + args: { + hasWorkspace: true, + workspaceSkills: mockWorkspaceSkills, + }, + play: async ({ canvasElement }) => { + await typeInEditor(canvasElement, "/"); + expect(await findVisibleText("Personal skills")).toBeDefined(); + expect(await findVisibleText("Workspace skills")).toBeDefined(); + expect(await findVisibleText("/reviewer")).toBeDefined(); + expect(await findVisibleText("/workspace/test-runner")).toBeDefined(); + }, +}; + +export const ArrowDownSelectsWorkspaceSkill: Story = { + args: { + hasWorkspace: true, + workspaceSkills: mockWorkspaceSkills, + }, + play: async ({ canvasElement }) => { + const editor = await typeInEditor(canvasElement, "/"); + await findVisibleText("/workspace/test-runner"); + await userEvent.keyboard("{ArrowDown}{ArrowDown}{ArrowDown}{Enter}"); + await waitFor(() => { + expect(editor.textContent).toBe("/workspace/test-runner"); + }); + }, +}; + +export const CollidingPersonalSkillInsertsQualifiedTrigger: Story = { + args: { + hasWorkspace: true, + workspaceSkills: [ + { name: "reviewer", description: "Workspace review process." }, + ], + }, + play: async ({ canvasElement }) => { + const editor = await typeInEditor(canvasElement, "/rev"); + expect(await findVisibleText("/personal/reviewer")).toBeDefined(); + expect(await findVisibleText("/workspace/reviewer")).toBeDefined(); + await userEvent.click(await findVisibleText("/personal/reviewer")); + await waitFor(() => { + expect(editor.textContent).toBe("/personal/reviewer"); + }); + }, +}; + +export const PersonalTriggersQualifiedWhileWorkspaceSkillsUnknown: Story = { + args: { + // No workspaceSkills: the chat detail has not resolved, so + // collisions are unknown. + hasWorkspace: true, + }, + play: async ({ canvasElement }) => { + await typeInEditor(canvasElement, "/rev"); + expect(await findVisibleText("/personal/reviewer")).toBeDefined(); + }, +}; + +export const EmptyPersonalKeepsMenuOpenWhileWorkspaceSkillsUnknown: Story = { + args: { + personalSkillsOverride: [], + // No workspaceSkills: closing the menu here would record the slash + // as dismissed, so skills arriving later could never reopen it. + hasWorkspace: true, + }, + play: async ({ canvasElement }) => { + await typeInEditor(canvasElement, "/"); + expect(await findVisibleText("Loading workspace skills...")).toBeDefined(); + }, +}; + +export const QualifiedPersonalQueryMatchesBareTrigger: Story = { + args: { + hasWorkspace: true, + // Workspace skills resolve without collisions, so personal items + // display bare triggers while the typed query stays qualified. + workspaceSkills: mockWorkspaceSkills, + }, + play: async ({ canvasElement }) => { + const editor = await typeInEditor(canvasElement, "/personal/rev"); + expect(await findVisibleText("/reviewer")).toBeDefined(); + await userEvent.keyboard("{Enter}"); + await waitFor(() => { + expect(editor.textContent).toBe("/reviewer"); + }); + }, +}; + +export const UniqueWorkspaceQualifiedPrefixStaysSearchable: Story = { + args: { + hasWorkspace: true, + workspaceSkills: mockWorkspaceSkills, + }, + play: async ({ canvasElement }) => { + const editor = await typeInEditor(canvasElement, "/workspace/t"); + expect(await findVisibleText("/workspace/test-runner")).toBeDefined(); + await userEvent.keyboard("{Enter}"); + await waitFor(() => { + expect(editor.textContent).toBe("/workspace/test-runner"); + }); + }, +}; + export const EmptyDescriptionInsertsNameOnly: Story = { play: async ({ canvasElement }) => { const editor = await typeInEditor(canvasElement, "/pla"); @@ -174,18 +291,6 @@ export const SlashInsideUrlDoesNotOpen: Story = { }, }; -export const BackspaceClosesWithoutEmptyStateFlash: Story = { - play: async ({ canvasElement }) => { - const editor = await typeInEditor(canvasElement, "/"); - await findVisibleText("/reviewer"); - await userEvent.keyboard("{Backspace}"); - - expect(editor.textContent).toBe(""); - expectNoVisibleTextImmediately("No personal skills found."); - await expectNoVisibleText("/reviewer"); - }, -}; - export const EscapeClosesWithoutReplacing: Story = { play: async ({ canvasElement }) => { const editor = await typeInEditor(canvasElement, "/"); @@ -230,8 +335,8 @@ export const OutsideClickDismissesTriggerOnRefocus: Story = { }, }; -// Stories below verify that on mobile viewports, the personal skills -// popup sits directly above the chat input rather than being clipped +// Stories below verify that on mobile viewports, the skills popup +// sits directly above the chat input rather than being clipped // above the visible viewport. const MOBILE_MEDIA_QUERY = "(max-width: 767px)"; diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx index 0a7605bb801..ba7e9957e61 100644 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.test.tsx @@ -6,14 +6,15 @@ import { useRef, useState, } from "react"; -import { QueryClientProvider } from "react-query"; -import { describe, expect, it } from "vitest"; +import { type QueryClient, QueryClientProvider } from "react-query"; +import { afterEach, beforeAll, describe, expect, it, vi } from "vitest"; import { createTestQueryClient } from "#/testHelpers/renderHelpers"; import { ChatMessageInput, type ChatMessageInputRef } from "./ChatMessageInput"; -const renderWithQueryClient = (children: ReactNode) => { - const queryClient = createTestQueryClient(); - +const renderWithQueryClient = ( + children: ReactNode, + queryClient: QueryClient = createTestQueryClient(), +) => { return render( {children}, ); @@ -65,7 +66,18 @@ const QueuedReplacementHarness: FC<{ ); }; +beforeAll(() => { + Object.defineProperty(Range.prototype, "getBoundingClientRect", { + configurable: true, + value: () => new DOMRect(0, 0, 1, 16), + }); +}); + describe("ChatMessageInput", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + it("returns the initial draft before the editor visually hydrates", async () => { renderWithQueryClient( , diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx index 545bdff455c..c7586d184d9 100644 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/ChatMessageInput.tsx @@ -43,16 +43,14 @@ import { } from "../../utils/agentChatSendShortcut"; import { isChatAttachmentFile } from "../../utils/chatAttachments"; import { - filterPersonalSkills, + filterSkillsByQuery, isPersonalSkillTriggerToken, - personalSkillTriggerText, } from "../../utils/personalSkills"; import { $createFileReferenceNode, FileReferenceNode, } from "./FileReferenceNode"; import { IOSBackspacePlugin } from "./iosBackspace"; -import { PersonalSkillsTriggerMenu } from "./PersonalSkillsTriggerMenu"; import { createPasteFile, getPasteDataTransfer, @@ -60,6 +58,12 @@ import { isLargePaste, type PasteCommandEvent, } from "./pasteHelpers"; +import { + createSkillMenuItem, + type SkillMenuItem, + type SkillMetadata, + SkillsTriggerMenu, +} from "./SkillsTriggerMenu"; import { type ActiveSkillsTrigger, SkillsTriggerPlugin, @@ -501,7 +505,21 @@ interface ChatMessageInputProps allowTextAttachmentPaste?: boolean; disabled?: boolean; autoFocus?: boolean; + /** + * True when the chat has a bound workspace, so workspace skills may + * exist even while workspaceSkills is still undefined. + */ + hasWorkspace?: boolean; + /** + * Story and test seam for deterministic personal skill menu data. + */ personalSkillsOverride?: readonly TypesGen.UserSkillMetadata[]; + /** + * Workspace skill menu data from the chat's pinned context, so the + * menu matches read_skill resolution. Undefined while the chat + * detail is still loading (or when no chat exists yet). + */ + workspaceSkills?: readonly SkillMetadata[]; "aria-label"?: string; } @@ -567,7 +585,9 @@ const ChatMessageInput = ({ allowTextAttachmentPaste, disabled, autoFocus, + hasWorkspace, personalSkillsOverride, + workspaceSkills, "aria-label": ariaLabel, ref, ...props @@ -597,30 +617,65 @@ const ChatMessageInput = ({ const [skillsMenuSelectedIndex, setSkillsMenuSelectedIndex] = useState(0); const hasSkillsTrigger = Boolean(skillsTrigger); const hasPersonalSkillsOverride = personalSkillsOverride !== undefined; + const personalSkillsQueryEnabled = + hasSkillsTrigger && !hasPersonalSkillsOverride; const skillsQuery = useQuery({ ...userSkills(), - enabled: hasSkillsTrigger && !hasPersonalSkillsOverride, + enabled: personalSkillsQueryEnabled, // Avoid refetching on each trigger toggle from caret movement. staleTime: 60_000, }); const personalSkills = personalSkillsOverride ?? skillsQuery.data ?? []; + const loadedWorkspaceSkills = workspaceSkills ?? []; + // Until the chat detail resolves, workspace skills are unknown: keep + // personal triggers qualified (a qualified alias always resolves) and + // treat the workspace list as still loading. + const workspaceSkillsKnown = !hasWorkspace || workspaceSkills !== undefined; // A stale empty cache with a refetch in flight must not dismiss the menu. - const isResolvedEmptySkillsList = hasPersonalSkillsOverride + const isResolvedEmptyPersonalSkills = hasPersonalSkillsOverride ? personalSkills.length === 0 : skillsQuery.isSuccess && !skillsQuery.isFetching && personalSkills.length === 0; - // When the loaded skills list is empty, "/" is plain text. When only + // Unknown workspace skills must not close the menu: the trigger plugin + // records a closed trigger as dismissed, so skills arriving later could + // never reopen it. + const isResolvedEmptyWorkspaceSkills = + workspaceSkillsKnown && loadedWorkspaceSkills.length === 0; + // When both skills lists resolve empty, "/" is plain text. When only // the filtered result is empty, keep the menu open for the no-match // message. - const skillsMenuOpen = hasSkillsTrigger && !isResolvedEmptySkillsList; - const filteredPersonalSkills = skillsTrigger - ? filterPersonalSkills(personalSkills, skillsTrigger.query) - : []; + const skillsMenuOpen = + hasSkillsTrigger && + !(isResolvedEmptyPersonalSkills && isResolvedEmptyWorkspaceSkills); + const skillsSearchQuery = skillsTrigger?.query ?? ""; + const workspaceSkillNames = new Set( + loadedWorkspaceSkills.map((skill) => skill.name), + ); + const personalSkillItems: readonly SkillMenuItem[] = filterSkillsByQuery( + personalSkills.map((skill) => + createSkillMenuItem( + "personal", + skill, + !workspaceSkillsKnown || workspaceSkillNames.has(skill.name), + ), + ), + skillsSearchQuery, + ); + const workspaceSkillItems: readonly SkillMenuItem[] = filterSkillsByQuery( + loadedWorkspaceSkills.map((skill) => + createSkillMenuItem("workspace", skill), + ), + skillsSearchQuery, + ); + const allFilteredSkills: readonly SkillMenuItem[] = [ + ...personalSkillItems, + ...workspaceSkillItems, + ]; const selectedSkillIndex = - filteredPersonalSkills.length === 0 + allFilteredSkills.length === 0 ? -1 - : Math.min(skillsMenuSelectedIndex, filteredPersonalSkills.length - 1); + : Math.min(skillsMenuSelectedIndex, allFilteredSkills.length - 1); const handleSkillsTriggerChange = (trigger: ActiveSkillsTrigger | null) => { if ( @@ -640,7 +695,7 @@ const ChatMessageInput = ({ setSkillsTrigger(trigger); }; - const replaceActiveSkillsTrigger = (skill: TypesGen.UserSkillMetadata) => { + const replaceActiveSkillsTrigger = (skill: SkillMenuItem) => { const editor = editorRef.current; const trigger = skillsTrigger; if (!editor || !trigger) { @@ -680,7 +735,7 @@ const ChatMessageInput = ({ selection.anchor.set(trigger.nodeKey, trigger.slashOffset, "text"); selection.focus.set(trigger.nodeKey, caretOffset, "text"); - selection.insertText(personalSkillTriggerText(skill)); + selection.insertText(skill.triggerText); }); setSkillsTrigger(null); setSkillsMenuSelectedIndex(0); @@ -890,7 +945,7 @@ const ChatMessageInput = ({ {autoFocus && } - handleSkillsTriggerChange(null)} /> diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.stories.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.stories.tsx deleted file mode 100644 index c1c5f1214a0..00000000000 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.stories.tsx +++ /dev/null @@ -1,113 +0,0 @@ -import type { Meta, StoryObj } from "@storybook/react-vite"; -import { expect, fn, userEvent } from "storybook/test"; -import { filterPersonalSkills } from "../../utils/personalSkills"; -import { PersonalSkillsTriggerMenu } from "./PersonalSkillsTriggerMenu"; -import { - expectNoVisibleText, - findVisibleText, - MockSkills, -} from "./storyHelpers"; - -const meta: Meta = { - title: "components/ChatMessageInput/PersonalSkillsTriggerMenu", - component: PersonalSkillsTriggerMenu, - args: { - open: true, - anchorRect: { top: 120, left: 80, height: 20 }, - query: "", - skills: MockSkills, - onSelectedIndexChange: fn(), - selectedIndex: 0, - onSelect: fn(), - onClose: fn(), - }, - decorators: [ - (Story) => ( -
-

- The menu is anchored to a mock caret position. -

- -
- ), - ], -}; - -export default meta; -type Story = StoryObj; - -export const Open: Story = { - play: async () => { - expect(await findVisibleText("/reviewer")).toBeDefined(); - expect( - await findVisibleText("Review changed files and suggest fixes."), - ).toBeDefined(); - }, -}; - -export const Loading: Story = { - args: { - isLoading: true, - skills: [], - }, - play: async () => { - expect(await findVisibleText("Loading personal skills...")).toBeDefined(); - }, -}; - -export const ErrorState: Story = { - args: { - isError: true, - skills: [], - }, - play: async () => { - expect( - await findVisibleText( - "Could not load personal skills. Close and type / again to retry.", - ), - ).toBeDefined(); - }, -}; - -export const Empty: Story = { - args: { - skills: [], - }, - play: async () => { - expect(await findVisibleText("No personal skills found.")).toBeDefined(); - }, -}; - -export const FilteredEmpty: Story = { - args: { - query: "xyz", - skills: [], - }, - play: async () => { - expect( - await findVisibleText("No personal skills match that query."), - ).toBeDefined(); - }, -}; - -export const Filtered: Story = { - args: { - query: "rev", - skills: filterPersonalSkills(MockSkills, "rev"), - }, - play: async () => { - expect(await findVisibleText("/reviewer")).toBeDefined(); - await expectNoVisibleText("/docs"); - }, -}; - -export const SelectsByClick: Story = { - args: { - onSelect: fn(), - }, - play: async ({ args }) => { - await userEvent.click(await findVisibleText("/reviewer")); - expect(args.onSelect).toHaveBeenCalledTimes(1); - expect(args.onSelect).toHaveBeenCalledWith(MockSkills[0]); - }, -}; diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.tsx deleted file mode 100644 index ae0b509bbd5..00000000000 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/PersonalSkillsTriggerMenu.tsx +++ /dev/null @@ -1,188 +0,0 @@ -import { useLayoutEffect, useState } from "react"; -import type * as TypesGen from "#/api/typesGenerated"; -import { - Command, - CommandEmpty, - CommandGroup, - CommandItem, - CommandList, -} from "#/components/Command/Command"; -import { - Popover, - PopoverAnchor, - PopoverContent, -} from "#/components/Popover/Popover"; -import { personalSkillTriggerText } from "../../utils/personalSkills"; - -// Prevent zero-height anchors when the browser returns a degenerate caret rect. -const MIN_ANCHOR_HEIGHT_PX = 16; - -export type CaretAnchorRect = { - top: number; - left: number; - height: number; -}; - -type PersonalSkillsTriggerMenuProps = { - open: boolean; - anchorRect: CaretAnchorRect | null; - query: string; - skills: readonly TypesGen.UserSkillMetadata[]; - isLoading?: boolean; - isError?: boolean; - selectedIndex: number; - onSelectedIndexChange: (index: number) => void; - onSelect: (skill: TypesGen.UserSkillMetadata) => void; - onClose: () => void; -}; - -type PersonalSkillsMenuState = { - anchorRect: CaretAnchorRect; - query: string; - skills: readonly TypesGen.UserSkillMetadata[]; - isLoading?: boolean; - isError?: boolean; - selectedIndex: number; -}; - -export const PersonalSkillsTriggerMenu = ({ - open, - anchorRect, - query, - skills, - isLoading, - isError, - selectedIndex, - onSelectedIndexChange, - onSelect, - onClose, -}: PersonalSkillsTriggerMenuProps) => { - const [lastOpenMenuState, setLastOpenMenuState] = - useState(null); - const isAnchoredOpen = open && anchorRect !== null; - const activeMenuState: PersonalSkillsMenuState | null = isAnchoredOpen - ? { - anchorRect, - query, - skills, - isLoading, - isError, - selectedIndex, - } - : null; - const menuState = activeMenuState ?? lastOpenMenuState; - const menuAnchorRect = menuState?.anchorRect ?? null; - const menuSkills = menuState?.skills ?? []; - const menuSelectedIndex = menuState?.selectedIndex ?? -1; - - useLayoutEffect(() => { - if (!isAnchoredOpen) { - return; - } - setLastOpenMenuState({ - anchorRect, - query, - skills, - isLoading, - isError, - selectedIndex, - }); - }, [ - anchorRect, - isAnchoredOpen, - isError, - isLoading, - query, - selectedIndex, - skills, - ]); - - const handleHighlightedValueChange = (value: string) => { - const nextIndex = menuSkills.findIndex((skill) => skill.name === value); - if (nextIndex >= 0) { - onSelectedIndexChange(nextIndex); - } - }; - - return ( - { - if (!nextOpen) { - onClose(); - } - }} - > - {menuAnchorRect && ( - - - )} - event.preventDefault()} - onOpenAutoFocus={(event) => event.preventDefault()} - onCloseAutoFocus={(event) => event.preventDefault()} - > - - - {menuState?.isLoading ? ( - - Loading personal skills... - - ) : menuState?.isError ? ( - - Could not load personal skills. Close and type / again to retry. - - ) : menuSkills.length === 0 ? ( - - {menuState?.query - ? "No personal skills match that query." - : "No personal skills found."} - - ) : ( - - {menuSkills.map((skill) => ( - onSelect(skill)} - > -
-
- {personalSkillTriggerText(skill)} -
- {skill.description.trim() && ( -
- {skill.description} -
- )} -
- - ))} - - )} - - - - - ); -}; diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.stories.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.stories.tsx new file mode 100644 index 00000000000..c3e2bbf52f0 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.stories.tsx @@ -0,0 +1,155 @@ +import type { Meta, StoryObj } from "@storybook/react-vite"; +import { expect, fn, userEvent } from "storybook/test"; +import { filterSkillsByQuery } from "../../utils/personalSkills"; +import { + createSkillMenuItem, + type SkillMetadata, + SkillsTriggerMenu, +} from "./SkillsTriggerMenu"; +import { + expectNoVisibleText, + findVisibleText, + MockSkills, +} from "./storyHelpers"; + +const mockWorkspaceSkills: SkillMetadata[] = [ + { + name: "test-runner", + description: "Run the workspace test command.", + }, + { + name: "workspace-docs", + description: "Use repository documentation conventions.", + }, +]; + +const mockPersonalSkillItems = MockSkills.map((skill) => + createSkillMenuItem("personal", skill), +); +const mockWorkspaceSkillItems = mockWorkspaceSkills.map((skill) => + createSkillMenuItem("workspace", skill), +); + +const meta: Meta = { + title: "components/ChatMessageInput/SkillsTriggerMenu", + component: SkillsTriggerMenu, + args: { + open: true, + anchorRect: { top: 120, left: 80, height: 20 }, + query: "", + personalSkills: mockPersonalSkillItems, + workspaceSkills: [], + workspaceSkillsEnabled: false, + onSelectedIndexChange: fn(), + selectedIndex: 0, + onSelect: fn(), + onClose: fn(), + }, + decorators: [ + (Story) => ( +
+

+ The menu is anchored to a mock caret position. +

+ +
+ ), + ], +}; + +export default meta; +type Story = StoryObj; + +export const PersonalOnly: Story = { + play: async () => { + expect(await findVisibleText("Personal skills")).toBeDefined(); + expect(await findVisibleText("/reviewer")).toBeDefined(); + expect( + await findVisibleText("Review changed files and suggest fixes."), + ).toBeDefined(); + await expectNoVisibleText("Workspace skills"); + }, +}; + +export const BothGroups: Story = { + args: { + workspaceSkills: mockWorkspaceSkillItems, + workspaceSkillsEnabled: true, + }, + play: async () => { + expect(await findVisibleText("Personal skills")).toBeDefined(); + expect(await findVisibleText("Workspace skills")).toBeDefined(); + expect(await findVisibleText("/reviewer")).toBeDefined(); + expect(await findVisibleText("/workspace/test-runner")).toBeDefined(); + }, +}; + +export const Loading: Story = { + args: { + isPersonalLoading: true, + personalSkills: [], + }, + play: async () => { + expect(await findVisibleText("Loading personal skills...")).toBeDefined(); + }, +}; + +export const WorkspaceLoading: Story = { + args: { + personalSkills: [], + workspaceSkills: [], + workspaceSkillsEnabled: true, + isWorkspaceLoading: true, + }, + play: async () => { + expect(await findVisibleText("Loading workspace skills...")).toBeDefined(); + }, +}; + +export const EmptyWithWorkspace: Story = { + args: { + personalSkills: [], + workspaceSkills: [], + workspaceSkillsEnabled: true, + }, + play: async () => { + expect( + await findVisibleText("No personal or workspace skills found."), + ).toBeDefined(); + }, +}; + +export const Empty: Story = { + args: { + personalSkills: [], + workspaceSkills: [], + }, + play: async () => { + expect(await findVisibleText("No personal skills found.")).toBeDefined(); + }, +}; + +export const Filtered: Story = { + args: { + query: "rev", + personalSkills: filterSkillsByQuery(mockPersonalSkillItems, "rev"), + workspaceSkills: filterSkillsByQuery(mockWorkspaceSkillItems, "rev"), + workspaceSkillsEnabled: true, + }, + play: async () => { + expect(await findVisibleText("/reviewer")).toBeDefined(); + await expectNoVisibleText("/docs"); + await expectNoVisibleText("/workspace/test-runner"); + }, +}; + +export const SelectsByClick: Story = { + args: { + onSelect: fn(), + }, + play: async ({ args }) => { + await userEvent.click(await findVisibleText("/reviewer")); + expect(args.onSelect).toHaveBeenCalledTimes(1); + expect(args.onSelect).toHaveBeenCalledWith(mockPersonalSkillItems[0]); + }, +}; diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.tsx new file mode 100644 index 00000000000..9e9a2ac4435 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerMenu.tsx @@ -0,0 +1,238 @@ +import { + Command, + CommandEmpty, + CommandGroup, + CommandItem, + CommandList, +} from "#/components/Command/Command"; +import { + Popover, + PopoverAnchor, + PopoverContent, +} from "#/components/Popover/Popover"; +import { cn } from "#/utils/cn"; + +// Prevent zero-height anchors when the browser returns a degenerate caret rect. +const MIN_ANCHOR_HEIGHT_PX = 16; + +export type CaretAnchorRect = { + top: number; + left: number; + height: number; +}; + +type SkillSource = "personal" | "workspace"; + +export type SkillMetadata = { + name: string; + description: string; +}; + +export type SkillMenuItem = SkillMetadata & { + source: SkillSource; + triggerText: string; + // The qualified alias stays searchable even when the displayed + // trigger is bare, so a typed qualified query keeps matching after + // collision state changes mid-trigger. + altTriggerText: string; +}; + +export const createSkillMenuItem = ( + source: SkillSource, + skill: SkillMetadata, + // Bare personal names are ambiguous to read_skill when a workspace + // skill shares the name, so colliding triggers must stay qualified. + qualifyTrigger = source === "workspace", +): SkillMenuItem => ({ + name: skill.name, + description: skill.description, + source, + triggerText: qualifyTrigger ? `/${source}/${skill.name}` : `/${skill.name}`, + altTriggerText: `/${source}/${skill.name}`, +}); + +type SkillsTriggerMenuProps = { + open: boolean; + anchorRect: CaretAnchorRect | null; + query: string; + personalSkills: readonly SkillMenuItem[]; + workspaceSkills: readonly SkillMenuItem[]; + workspaceSkillsEnabled?: boolean; + isPersonalLoading?: boolean; + isPersonalError?: boolean; + isWorkspaceLoading?: boolean; + selectedIndex: number; + onSelectedIndexChange: (index: number) => void; + onSelect: (skill: SkillMenuItem) => void; + onClose: () => void; +}; + +const getEmptyMessage = (query: string, workspaceSkillsEnabled: boolean) => { + if (query) { + return workspaceSkillsEnabled + ? "No skills match that query." + : "No personal skills match that query."; + } + return workspaceSkillsEnabled + ? "No personal or workspace skills found." + : "No personal skills found."; +}; + +const SkillCommandItem = ({ + skill, + value, + selected, + onSelect, +}: { + skill: SkillMenuItem; + value: string; + selected: boolean; + onSelect: (skill: SkillMenuItem) => void; +}) => { + const handleSelect = () => onSelect(skill); + + return ( + +
+
+ {skill.triggerText} +
+ {skill.description.trim() && ( +
+ {skill.description} +
+ )} +
+ + ); +}; + +export const SkillsTriggerMenu = ({ + open, + anchorRect, + query, + personalSkills, + workspaceSkills, + workspaceSkillsEnabled, + isPersonalLoading, + isPersonalError, + isWorkspaceLoading, + selectedIndex, + onSelectedIndexChange, + onSelect, + onClose, +}: SkillsTriggerMenuProps) => { + const allSkills = [...personalSkills, ...workspaceSkills]; + const statusItems = [ + isPersonalLoading && personalSkills.length === 0 + ? "Loading personal skills..." + : undefined, + isPersonalError && personalSkills.length === 0 + ? "Could not load personal skills. Close and type / again to retry." + : undefined, + isWorkspaceLoading && workspaceSkills.length === 0 + ? "Loading workspace skills..." + : undefined, + ].filter((item) => item !== undefined); + const shouldRender = open && anchorRect; + const shouldShowEmpty = allSkills.length === 0 && statusItems.length === 0; + const selectedValue = selectedIndex >= 0 ? String(selectedIndex) : ""; + + const handleHighlightedValueChange = (value: string) => { + const nextIndex = Number(value); + if ( + Number.isInteger(nextIndex) && + nextIndex >= 0 && + nextIndex < allSkills.length + ) { + onSelectedIndexChange(nextIndex); + } + }; + + const renderSkill = (skill: SkillMenuItem, index: number) => ( + + ); + + return ( + { + if (!nextOpen) { + onClose(); + } + }} + > + {shouldRender && ( + + + )} + event.preventDefault()} + onOpenAutoFocus={(event) => event.preventDefault()} + onCloseAutoFocus={(event) => event.preventDefault()} + > + + + {personalSkills.length > 0 && ( + + {personalSkills.map((skill, index) => + renderSkill(skill, index), + )} + + )} + {workspaceSkills.length > 0 && ( + + {workspaceSkills.map((skill, index) => + renderSkill(skill, personalSkills.length + index), + )} + + )} + {statusItems.map((message) => ( + + {message} + + ))} + {shouldShowEmpty && ( + + {getEmptyMessage(query, Boolean(workspaceSkillsEnabled))} + + )} + + + + + ); +}; diff --git a/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerPlugin.tsx b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerPlugin.tsx index cc230d9535a..9d314789426 100644 --- a/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerPlugin.tsx +++ b/site/src/pages/AgentsPage/components/ChatMessageInput/SkillsTriggerPlugin.tsx @@ -12,9 +12,8 @@ import { type NodeKey, } from "lexical"; import { useEffect, useEffectEvent, useLayoutEffect, useRef } from "react"; -import type * as TypesGen from "#/api/typesGenerated"; import { parsePersonalSkillTrigger } from "../../utils/personalSkills"; -import type { CaretAnchorRect } from "./PersonalSkillsTriggerMenu"; +import type { CaretAnchorRect, SkillMenuItem } from "./SkillsTriggerMenu"; export type ActiveSkillsTrigger = { nodeKey: NodeKey; @@ -30,11 +29,11 @@ type DismissedSkillsTrigger = Pick< type SkillsTriggerPluginProps = { open: boolean; - skills: readonly TypesGen.UserSkillMetadata[]; + skills: readonly SkillMenuItem[]; selectedIndex: number; onSelectedIndexChange: (index: number) => void; onTriggerChange: (trigger: ActiveSkillsTrigger | null) => void; - onSkillSelect: (skill: TypesGen.UserSkillMetadata) => void; + onSkillSelect: (skill: SkillMenuItem) => void; }; const currentCaretRect = (): CaretAnchorRect | null => { @@ -181,8 +180,11 @@ export const SkillsTriggerPlugin = ({ if (count === 0) { return true; } - const currentIndex = Math.max(0, selectedIndex); - onSelectedIndexChange((currentIndex + delta + count) % count); + if (selectedIndex < 0) { + onSelectedIndexChange(delta > 0 ? 0 : count - 1); + return true; + } + onSelectedIndexChange((selectedIndex + delta + count) % count); return true; }, ); @@ -192,7 +194,7 @@ export const SkillsTriggerPlugin = ({ return false; } event?.preventDefault(); - const skill = skills[selectedIndex]; + const skill = selectedIndex >= 0 ? skills[selectedIndex] : undefined; if (skill) { onSkillSelect(skill); } @@ -203,7 +205,7 @@ export const SkillsTriggerPlugin = ({ if (!open) { return false; } - const skill = skills[selectedIndex]; + const skill = selectedIndex >= 0 ? skills[selectedIndex] : undefined; if (!skill) { return false; } diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.test.ts b/site/src/pages/AgentsPage/components/ChatPageContent.test.ts new file mode 100644 index 00000000000..eb6b86cc965 --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatPageContent.test.ts @@ -0,0 +1,95 @@ +import { describe, expect, it } from "vitest"; +import type * as TypesGen from "#/api/typesGenerated"; +import { MockChat } from "#/testHelpers/chatEntities"; +import { workspaceSkillsFromChat } from "./ChatPageContent"; + +const skillResource = ( + name: string, + overrides: Partial = {}, +): TypesGen.ChatContextResource => ({ + source: `/workspace/.agents/skills/${name}`, + kind: "skill", + size_bytes: 128, + skill_name: name, + skill_description: `${name} description`, + status: "ok", + ...overrides, +}); + +const instructionResource = (): TypesGen.ChatContextResource => ({ + source: "/workspace/AGENTS.md", + kind: "instruction_file", + size_bytes: 64, + status: "ok", +}); + +const chatWithContext = ( + context: TypesGen.ChatContext | undefined, +): TypesGen.Chat => ({ ...MockChat, context }); + +describe("workspaceSkillsFromChat", () => { + it("returns undefined while the chat detail is unresolved", () => { + expect(workspaceSkillsFromChat(undefined)).toBeUndefined(); + }); + + it("returns an empty authoritative list for a resolved unpinned chat", () => { + expect(workspaceSkillsFromChat(chatWithContext(undefined))).toEqual([]); + expect(workspaceSkillsFromChat(chatWithContext({ dirty: false }))).toEqual( + [], + ); + }); + + it("maps healthy skill resources to workspace skills", () => { + const chat = chatWithContext({ + dirty: false, + resources: [ + instructionResource(), + skillResource("reviewer"), + skillResource("docs"), + ], + }); + expect(workspaceSkillsFromChat(chat)).toEqual([ + { name: "reviewer", description: "reviewer description" }, + { name: "docs", description: "docs description" }, + ]); + }); + + it("keeps the first resource for duplicate skill names, matching read_skill", () => { + const chat = chatWithContext({ + dirty: false, + resources: [ + skillResource("reviewer", { + source: "/workspace/.agents/skills/reviewer", + }), + skillResource("reviewer", { + source: "/workspace/other/skills/reviewer", + skill_description: "shadowed duplicate", + }), + ], + }); + expect(workspaceSkillsFromChat(chat)).toEqual([ + { name: "reviewer", description: "reviewer description" }, + ]); + }); + + it("omits non-ok skill resources", () => { + const chat = chatWithContext({ + dirty: true, + resources: [ + skillResource("reviewer"), + skillResource("broken", { status: "unreadable", skill_name: "" }), + ], + }); + expect(workspaceSkillsFromChat(chat)).toEqual([ + { name: "reviewer", description: "reviewer description" }, + ]); + }); + + it("returns an empty authoritative list when pinned context has no skills", () => { + const chat = chatWithContext({ + dirty: false, + resources: [instructionResource()], + }); + expect(workspaceSkillsFromChat(chat)).toEqual([]); + }); +}); diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 02d5590e1ba..bf35260abf5 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -39,6 +39,7 @@ import { } from "./ChatConversation/messageParsing"; import { useOnRenderProfiler } from "./ChatConversation/useOnRenderProfiler"; import type { ModelSelectorOption } from "./ChatElements"; +import type { SkillMetadata } from "./ChatMessageInput/SkillsTriggerMenu"; type ChatStoreHandle = ReturnType["store"]; @@ -46,6 +47,33 @@ const isChatMessage = ( message: TypesGen.ChatMessage | undefined, ): message is TypesGen.ChatMessage => Boolean(message); +// A resolved chat with no context (unpinned) or no resources authoritatively +// has no workspace skills; only an unresolved chat leaves them unknown. +// Duplicate names keep the first resource to match read_skill resolution, +// which also collapses duplicates first-wins in resource order. +export const workspaceSkillsFromChat = ( + chat: TypesGen.Chat | undefined, +): SkillMetadata[] | undefined => { + if (!chat) { + return undefined; + } + const skills = new Map(); + for (const resource of chat.context?.resources ?? []) { + if ( + resource.kind !== "skill" || + resource.status !== "ok" || + skills.has(resource.skill_name ?? "") + ) { + continue; + } + skills.set(resource.skill_name ?? "", { + name: resource.skill_name ?? "", + description: resource.skill_description ?? "", + }); + } + return [...skills.values()]; +}; + interface ChatPageTimelineProps { store: ChatStoreHandle; persistedError: ChatDetailError | undefined; @@ -215,6 +243,9 @@ interface ChatPageInputProps { // Pinned workspace-context state for the chat, surfaced by the // context indicator (dirty marker and pinned resources). chatContext?: TypesGen.ChatContext; + // Workspace skill menu data derived from the resolved chat detail; + // undefined while the chat is still loading. + workspaceSkills?: readonly SkillMetadata[]; workspaceOptions: readonly TypesGen.Workspace[]; chatOrganizationId?: string; selectedWorkspaceId: string | null; @@ -273,6 +304,7 @@ export const ChatPageInput: FC = ({ onMCPSelectionChange, onMCPAuthComplete, chatContext, + workspaceSkills, workspaceOptions, chatOrganizationId, selectedWorkspaceId, @@ -521,6 +553,7 @@ export const ChatPageInput: FC = ({ selectedMCPServerIds={selectedMCPServerIds} onMCPSelectionChange={onMCPSelectionChange} onMCPAuthComplete={onMCPAuthComplete} + workspaceSkills={workspaceSkills} workspace={workspace} workspaceAgent={workspaceAgent} chatId={chatId} diff --git a/site/src/pages/AgentsPage/utils/personalSkills.test.ts b/site/src/pages/AgentsPage/utils/personalSkills.test.ts index 881362f0db8..bd4ccb828d6 100644 --- a/site/src/pages/AgentsPage/utils/personalSkills.test.ts +++ b/site/src/pages/AgentsPage/utils/personalSkills.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest"; import type * as TypesGen from "#/api/typesGenerated"; import { buildPersonalSkillMarkdown, - filterPersonalSkills, + filterSkillsByQuery, getPersonalSkillContentSizeBytes, isPersonalSkillTriggerToken, isValidPersonalSkillDescription, @@ -28,7 +28,7 @@ const skill = ( updated_at: now, }); -describe("filterPersonalSkills", () => { +describe("filterSkillsByQuery", () => { const skills = [ skill("deploy", "Ship reviewed production changes", 0), skill("reviewer", "Review changed files", 1), @@ -37,7 +37,7 @@ describe("filterPersonalSkills", () => { ]; it("sorts unfiltered skills by name", () => { - expect(filterPersonalSkills(skills, "").map(({ name }) => name)).toEqual([ + expect(filterSkillsByQuery(skills, "").map(({ name }) => name)).toEqual([ "api-review", "deploy", "docs", @@ -46,9 +46,11 @@ describe("filterPersonalSkills", () => { }); it("ranks prefix, name substring, then description matches", () => { - expect(filterPersonalSkills(skills, "rev").map(({ name }) => name)).toEqual( - ["reviewer", "api-review", "deploy"], - ); + expect(filterSkillsByQuery(skills, "rev").map(({ name }) => name)).toEqual([ + "reviewer", + "api-review", + "deploy", + ]); }); it("matches names and descriptions case-insensitively", () => { @@ -58,12 +60,48 @@ describe("filterPersonalSkills", () => { ]; expect( - filterPersonalSkills(mixedCaseSkills, "DEP").map(({ name }) => name), + filterSkillsByQuery(mixedCaseSkills, "DEP").map(({ name }) => name), ).toEqual(["deploy-bot"]); expect( - filterPersonalSkills(mixedCaseSkills, "changes").map(({ name }) => name), + filterSkillsByQuery(mixedCaseSkills, "changes").map(({ name }) => name), ).toEqual(["deploy-bot"]); }); + + it("matches trigger text", () => { + const skillsWithTriggerText = [ + { + name: "reviewer", + description: "Review changed files", + triggerText: "/reviewer", + }, + { + name: "test-runner", + description: "Run tests", + triggerText: "/workspace/test-runner", + }, + ]; + + expect( + filterSkillsByQuery(skillsWithTriggerText, "workspace/t").map( + ({ triggerText }) => triggerText, + ), + ).toEqual(["/workspace/test-runner"]); + }); + + it("matches the alternate qualified trigger of a bare item", () => { + const skills = [ + { + name: "reviewer", + description: "Review changed files", + triggerText: "/reviewer", + altTriggerText: "/personal/reviewer", + }, + ]; + + expect( + filterSkillsByQuery(skills, "personal/rev").map(({ name }) => name), + ).toEqual(["reviewer"]); + }); }); describe("personal skill slash triggers", () => { diff --git a/site/src/pages/AgentsPage/utils/personalSkills.ts b/site/src/pages/AgentsPage/utils/personalSkills.ts index 905fa4efe9c..435951bc852 100644 --- a/site/src/pages/AgentsPage/utils/personalSkills.ts +++ b/site/src/pages/AgentsPage/utils/personalSkills.ts @@ -1,5 +1,4 @@ import frontMatter from "front-matter"; -import type * as TypesGen from "#/api/typesGenerated"; export const PERSONAL_SKILL_MAX_SIZE_BYTES = 64 * 1024; const PERSONAL_SKILL_MAX_NAME_BYTES = 256; @@ -15,15 +14,23 @@ export type PersonalSkillFormValues = { body: string; }; -type RankedPersonalSkill = { - skill: TypesGen.UserSkillMetadata; +type SkillSearchMetadata = { + name: string; + description: string; + triggerText?: string; + // Alternate trigger form that stays searchable regardless of which + // form is displayed, e.g. the qualified /source/name alias. + altTriggerText?: string; +}; + +type RankedSkill = { + skill: T; rank: number; index: number; }; -export const personalSkillTriggerText = ( - skill: TypesGen.UserSkillMetadata, -): string => `/${skill.name}`; +export const personalSkillTriggerText = (skill: { name: string }): string => + `/${skill.name}`; type PersonalSkillTriggerMatch = { slashOffset: number; @@ -48,26 +55,41 @@ export const isPersonalSkillTriggerToken = (token: string): boolean => /^\/\S*$/.test(token); /** - * Filters personal skills by name and description. Matches are ranked by - * name prefix, name substring, then description substring. + * Filters skills by name, trigger text, and description. Matches are ranked + * by name or trigger text prefix, name or trigger text substring, then + * description substring. */ -export const filterPersonalSkills = ( - skills: readonly TypesGen.UserSkillMetadata[], +export const filterSkillsByQuery = ( + skills: readonly T[], query: string, -): TypesGen.UserSkillMetadata[] => { +): T[] => { const normalizedQuery = query.toLocaleLowerCase("en-US"); if (!normalizedQuery) { return skills.toSorted((a, b) => a.name.localeCompare(b.name, "en-US")); } - const rankedSkills: RankedPersonalSkill[] = []; + const rankedSkills: RankedSkill[] = []; for (const [index, skill] of skills.entries()) { const name = skill.name.toLocaleLowerCase("en-US"); + const triggerText = skill.triggerText + ?.replace(/^\//, "") + .toLocaleLowerCase("en-US"); + const altTriggerText = skill.altTriggerText + ?.replace(/^\//, "") + .toLocaleLowerCase("en-US"); const description = skill.description.toLocaleLowerCase("en-US"); let rank: number | undefined; - if (name.startsWith(normalizedQuery)) { + if ( + name.startsWith(normalizedQuery) || + triggerText?.startsWith(normalizedQuery) || + altTriggerText?.startsWith(normalizedQuery) + ) { rank = 0; - } else if (name.includes(normalizedQuery)) { + } else if ( + name.includes(normalizedQuery) || + triggerText?.includes(normalizedQuery) || + altTriggerText?.includes(normalizedQuery) + ) { rank = 1; } else if (description.includes(normalizedQuery)) { rank = 2;