diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 6fe7e722fd..e24d4196b5 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -17518,6 +17518,10 @@ const docTemplate = `{ "plan_mode": { "$ref": "#/definitions/codersdk.ChatPlanMode" }, + "queued_for_capacity": { + "description": "QueuedForCapacity reports that the chat is waiting for a concurrent\nagent slot. Single-chat reads derive it; list responses leave it false.", + "type": "boolean" + }, "root_chat_id": { "type": "string", "format": "uuid" diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index d658d05092..0b18135d1c 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15755,6 +15755,10 @@ "plan_mode": { "$ref": "#/definitions/codersdk.ChatPlanMode" }, + "queued_for_capacity": { + "description": "QueuedForCapacity reports that the chat is waiting for a concurrent\nagent slot. Single-chat reads derive it; list responses leave it false.", + "type": "boolean" + }, "root_chat_id": { "type": "string", "format": "uuid" diff --git a/coderd/coderd.go b/coderd/coderd.go index 69a6abe1d5..c7effe22b5 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -273,6 +273,8 @@ type Options struct { // Set by enterprise for HA deployments. Nil uses chatd's local // in-process channel dialer. ChatStreamPartsDialer chatd.StreamPartsDialer + // Nil keeps the default chat agent caps active. + ChatAgentCapacityUnlock chatd.AgentCapacityUnlock // ChatProviderAPIKeys overrides deployment-derived provider keys. // Test harnesses use this to route chat models to local providers. ChatProviderAPIKeys *chatprovider.ProviderAPIKeys @@ -941,6 +943,7 @@ func New(options *Options) *API { HookDispatcher: hookDispatcher, UsageTracker: options.WorkspaceUsageTracker, PrometheusRegistry: options.PrometheusRegistry, + AgentCapacityUnlock: options.ChatAgentCapacityUnlock, OIDCTokenSource: oidcMCPSrc, NotificationsEnqueuer: options.NotificationsEnqueuer, Auditor: &api.Auditor, diff --git a/coderd/database/db2sdk/db2sdk_test.go b/coderd/database/db2sdk/db2sdk_test.go index 69a1d7c9ba..6a9ac0f0d6 100644 --- a/coderd/database/db2sdk/db2sdk_test.go +++ b/coderd/database/db2sdk/db2sdk_test.go @@ -757,11 +757,8 @@ func TestChat_AllFieldsPopulated(t *testing.T) { v := reflect.ValueOf(got) typ := v.Type() - // HasUnread is populated by ChatRowsWithChildren (which joins the - // read-cursor query), not by Chat. Warnings is a transient - // field populated by handlers, not the converter. Both are - // expected to remain zero here. - skip := map[string]bool{"HasUnread": true, "Warnings": true} + // These fields are set outside db2sdk.Chat and intentionally remain zero. + skip := map[string]bool{"HasUnread": true, "Warnings": true, "QueuedForCapacity": true} for i := range typ.NumField() { field := typ.Field(i) if skip[field.Name] { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 321ec1e61d..743e51538b 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1964,6 +1964,20 @@ func (q *querier) CountAuditLogs(ctx context.Context, arg database.CountAuditLog return q.db.CountAuthorizedAuditLogs(ctx, arg, prep) } +func (q *querier) CountChatCapacityActiveByPool(ctx context.Context, arg database.CountChatCapacityActiveByPoolParams) (database.CountChatCapacityActiveByPoolRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil { + return database.CountChatCapacityActiveByPoolRow{}, err + } + return q.db.CountChatCapacityActiveByPool(ctx, arg) +} + +func (q *querier) CountChatCapacityQueuedByPool(ctx context.Context, staleSeconds int32) (database.CountChatCapacityQueuedByPoolRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil { + return database.CountChatCapacityQueuedByPoolRow{}, err + } + return q.db.CountChatCapacityQueuedByPool(ctx, staleSeconds) +} + func (q *querier) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { @@ -3483,6 +3497,15 @@ func (q *querier) GetChatPlanModeInstructions(ctx context.Context) (string, erro return q.db.GetChatPlanModeInstructions(ctx) } +func (q *querier) GetChatQueuedForCapacity(ctx context.Context, arg database.GetChatQueuedForCapacityParams) (bool, error) { + // The pool-fullness derivation counts other users' chats, so require + // deployment-wide chat read rather than per-chat authorization. + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil { + return false, err + } + return q.db.GetChatQueuedForCapacity(ctx, arg) +} + func (q *querier) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { _, err := q.GetChatByID(ctx, arg.ChatID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 1c6497a2e1..86665079cd 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -608,6 +608,23 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().GetChatWorkerAcquisitionCandidates(gomock.Any(), arg).Return([]database.GetChatWorkerAcquisitionCandidatesRow{row}, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.GetChatWorkerAcquisitionCandidatesRow{row}) })) + s.Run("CountChatCapacityActiveByPool", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.CountChatCapacityActiveByPoolParams{ExcludeChatID: uuid.New(), StaleSeconds: 30} + row := database.CountChatCapacityActiveByPoolRow{ActiveRootCount: 1, ActiveSubagentCount: 2} + dbm.EXPECT().CountChatCapacityActiveByPool(gomock.Any(), arg).Return(row, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionRead).Returns(row) + })) + s.Run("CountChatCapacityQueuedByPool", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + const staleSeconds = int32(30) + row := database.CountChatCapacityQueuedByPoolRow{QueuedRootCount: 3, QueuedSubagentCount: 4} + dbm.EXPECT().CountChatCapacityQueuedByPool(gomock.Any(), staleSeconds).Return(row, nil).AnyTimes() + check.Args(staleSeconds).Asserts(rbac.ResourceChat, policy.ActionRead).Returns(row) + })) + s.Run("GetChatQueuedForCapacity", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.GetChatQueuedForCapacityParams{ChatID: uuid.New(), StaleSeconds: 30, RootCapacity: 5, SubagentCapacity: 10} + dbm.EXPECT().GetChatQueuedForCapacity(gomock.Any(), arg).Return(true, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionRead).Returns(true) + })) s.Run("GetChatsByIDsForRunnerSync", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ids := []uuid.UUID{uuid.New(), uuid.New()} chat := testutil.Fake(s.T(), faker, database.Chat{ID: ids[0]}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 63c9c0f563..66d40bbb23 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -321,6 +321,22 @@ func (m queryMetricsStore) CountAuditLogs(ctx context.Context, arg database.Coun return r0, r1 } +func (m queryMetricsStore) CountChatCapacityActiveByPool(ctx context.Context, arg database.CountChatCapacityActiveByPoolParams) (database.CountChatCapacityActiveByPoolRow, error) { + start := time.Now() + r0, r1 := m.s.CountChatCapacityActiveByPool(ctx, arg) + m.queryLatencies.WithLabelValues("CountChatCapacityActiveByPool").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "CountChatCapacityActiveByPool").Inc() + return r0, r1 +} + +func (m queryMetricsStore) CountChatCapacityQueuedByPool(ctx context.Context, staleSeconds int32) (database.CountChatCapacityQueuedByPoolRow, error) { + start := time.Now() + r0, r1 := m.s.CountChatCapacityQueuedByPool(ctx, staleSeconds) + m.queryLatencies.WithLabelValues("CountChatCapacityQueuedByPool").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "CountChatCapacityQueuedByPool").Inc() + return r0, r1 +} + func (m queryMetricsStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { start := time.Now() r0, r1 := m.s.CountChatQueuedMessages(ctx, chatID) @@ -1705,6 +1721,14 @@ func (m queryMetricsStore) GetChatPlanModeInstructions(ctx context.Context) (str return r0, r1 } +func (m queryMetricsStore) GetChatQueuedForCapacity(ctx context.Context, arg database.GetChatQueuedForCapacityParams) (bool, error) { + start := time.Now() + r0, r1 := m.s.GetChatQueuedForCapacity(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatQueuedForCapacity").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedForCapacity").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessageByID(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 923416c68b..01d8960437 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -483,6 +483,36 @@ func (mr *MockStoreMockRecorder) CountAuthorizedConnectionLogs(ctx, arg, prepare return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountAuthorizedConnectionLogs", reflect.TypeOf((*MockStore)(nil).CountAuthorizedConnectionLogs), ctx, arg, prepared) } +// CountChatCapacityActiveByPool mocks base method. +func (m *MockStore) CountChatCapacityActiveByPool(ctx context.Context, arg database.CountChatCapacityActiveByPoolParams) (database.CountChatCapacityActiveByPoolRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CountChatCapacityActiveByPool", ctx, arg) + ret0, _ := ret[0].(database.CountChatCapacityActiveByPoolRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CountChatCapacityActiveByPool indicates an expected call of CountChatCapacityActiveByPool. +func (mr *MockStoreMockRecorder) CountChatCapacityActiveByPool(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountChatCapacityActiveByPool", reflect.TypeOf((*MockStore)(nil).CountChatCapacityActiveByPool), ctx, arg) +} + +// CountChatCapacityQueuedByPool mocks base method. +func (m *MockStore) CountChatCapacityQueuedByPool(ctx context.Context, staleSeconds int32) (database.CountChatCapacityQueuedByPoolRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CountChatCapacityQueuedByPool", ctx, staleSeconds) + ret0, _ := ret[0].(database.CountChatCapacityQueuedByPoolRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CountChatCapacityQueuedByPool indicates an expected call of CountChatCapacityQueuedByPool. +func (mr *MockStoreMockRecorder) CountChatCapacityQueuedByPool(ctx, staleSeconds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountChatCapacityQueuedByPool", reflect.TypeOf((*MockStore)(nil).CountChatCapacityQueuedByPool), ctx, staleSeconds) +} + // CountChatQueuedMessages mocks base method. func (m *MockStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { m.ctrl.T.Helper() @@ -3165,6 +3195,21 @@ func (mr *MockStoreMockRecorder) GetChatPlanModeInstructions(ctx any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatPlanModeInstructions", reflect.TypeOf((*MockStore)(nil).GetChatPlanModeInstructions), ctx) } +// GetChatQueuedForCapacity mocks base method. +func (m *MockStore) GetChatQueuedForCapacity(ctx context.Context, arg database.GetChatQueuedForCapacityParams) (bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatQueuedForCapacity", ctx, arg) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatQueuedForCapacity indicates an expected call of GetChatQueuedForCapacity. +func (mr *MockStoreMockRecorder) GetChatQueuedForCapacity(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedForCapacity", reflect.TypeOf((*MockStore)(nil).GetChatQueuedForCapacity), ctx, arg) +} + // GetChatQueuedMessageByID mocks base method. func (m *MockStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index d70220b91e..7365b8821b 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4842,7 +4842,7 @@ CREATE INDEX idx_chats_title_fts ON chats USING gin (to_tsvector('simple'::regco COMMENT ON INDEX idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.'; -CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (status, updated_at, id) WHERE (archived = false); +CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (((parent_chat_id IS NULL)), status, updated_at, id) WHERE (archived = false); CREATE INDEX idx_chats_workspace ON chats USING btree (workspace_id); diff --git a/coderd/database/lock.go b/coderd/database/lock.go index d2ec69293d..a9830336fe 100644 --- a/coderd/database/lock.go +++ b/coderd/database/lock.go @@ -17,6 +17,7 @@ const ( LockIDBoundaryUsageStats LockIDAIProvidersEnvSeed LockIDChatModelConfigWrites + LockIDChatCapacityAdmission ) // GenLockID generates a unique and consistent lock ID from a given string. diff --git a/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.down.sql b/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.down.sql new file mode 100644 index 0000000000..3629387476 --- /dev/null +++ b/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.down.sql @@ -0,0 +1,4 @@ +DROP INDEX idx_chats_worker_acquisition_candidates; +CREATE INDEX idx_chats_worker_acquisition_candidates ON chats + (status, updated_at, id) + WHERE archived = false; diff --git a/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.up.sql b/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.up.sql new file mode 100644 index 0000000000..ccb6691c6f --- /dev/null +++ b/coderd/database/migrations/000571_pool_aware_chat_acquisition_index.up.sql @@ -0,0 +1,4 @@ +DROP INDEX idx_chats_worker_acquisition_candidates; +CREATE INDEX idx_chats_worker_acquisition_candidates ON chats + ((parent_chat_id IS NULL), status, updated_at, id) + WHERE archived = false; diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 9b489e25a5..b7ccde5502 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -93,6 +93,9 @@ type sqlcQuerier interface { CleanupDeletedMCPServerIDsFromChats(ctx context.Context) error CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error) CountAuditLogs(ctx context.Context, arg CountAuditLogsParams) (int64, error) + // Excluding the candidate keeps ownership takeover capacity-neutral. + CountChatCapacityActiveByPool(ctx context.Context, arg CountChatCapacityActiveByPoolParams) (CountChatCapacityActiveByPoolRow, error) + CountChatCapacityQueuedByPool(ctx context.Context, staleSeconds int32) (CountChatCapacityQueuedByPoolRow, error) // Cheap queue-length check used by ChatMachine.Update when deciding // whether the chat is in a "1" sub-state. CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) @@ -477,6 +480,8 @@ type sqlcQuerier interface { // personal chat model overrides. It defaults to false when unset. GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) GetChatPlanModeInstructions(ctx context.Context) (string, error) + // Pool fullness distinguishes capacity waits from worker pickup delays. + GetChatQueuedForCapacity(ctx context.Context, arg GetChatQueuedForCapacityParams) (bool, error) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error) // Returns the queue head (lowest position, then lowest id). GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) @@ -507,17 +512,9 @@ type sqlcQuerier interface { // jsonb_array_elements never raises "cannot extract elements from a // scalar". Backed by idx_chat_messages_user_prompts. GetChatUserPromptsByChatID(ctx context.Context, arg GetChatUserPromptsByChatIDParams) ([]GetChatUserPromptsByChatIDRow, error) - // Returns chats that workers may try to acquire. Candidates must be: - // - in a worker-runnable execution status; - // - unarchived; and - // - missing ownership, carrying inconsistent ownership, or lacking a - // fresh heartbeat for the assigned runner. - // - // Missing ownership is worker_id IS NULL. Inconsistent ownership is - // runner_id IS NULL while worker_id is set. Stale ownership is no - // heartbeat row for (chat_id, runner_id), or one older than - // @stale_seconds by database time. Candidates are ordered by oldest - // updated_at first so workers drain stale runnable chats predictably. + // Returns a bounded, pool-interleaved set of chats that workers may acquire. + // Interrupting chats finish active work first. Requires-action chats follow so + // their runner can enforce the action deadline before new generations start. GetChatWorkerAcquisitionCandidates(ctx context.Context, arg GetChatWorkerAcquisitionCandidatesParams) ([]GetChatWorkerAcquisitionCandidatesRow, error) // Returns the global TTL for chat workspaces as a Go duration string. // Returns "0s" (disabled) when no value has been configured. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index c4a5ea782b..2649038095 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -7048,6 +7048,69 @@ func (q *sqlQuerier) BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUps return err } +const countChatCapacityActiveByPool = `-- name: CountChatCapacityActiveByPool :one +SELECT + COUNT(*) FILTER (WHERE c.parent_chat_id IS NULL)::bigint AS active_root_count, + COUNT(*) FILTER (WHERE c.parent_chat_id IS NOT NULL)::bigint AS active_subagent_count +FROM chat_heartbeats hb +JOIN chats c + ON c.id = hb.chat_id + AND c.runner_id = hb.runner_id +WHERE c.worker_id IS NOT NULL + AND c.id != $1::uuid + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * $2::int) +` + +type CountChatCapacityActiveByPoolParams struct { + ExcludeChatID uuid.UUID `db:"exclude_chat_id" json:"exclude_chat_id"` + StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` +} + +type CountChatCapacityActiveByPoolRow struct { + ActiveRootCount int64 `db:"active_root_count" json:"active_root_count"` + ActiveSubagentCount int64 `db:"active_subagent_count" json:"active_subagent_count"` +} + +// Excluding the candidate keeps ownership takeover capacity-neutral. +func (q *sqlQuerier) CountChatCapacityActiveByPool(ctx context.Context, arg CountChatCapacityActiveByPoolParams) (CountChatCapacityActiveByPoolRow, error) { + row := q.db.QueryRowContext(ctx, countChatCapacityActiveByPool, arg.ExcludeChatID, arg.StaleSeconds) + var i CountChatCapacityActiveByPoolRow + err := row.Scan(&i.ActiveRootCount, &i.ActiveSubagentCount) + return i, err +} + +const countChatCapacityQueuedByPool = `-- name: CountChatCapacityQueuedByPool :one +SELECT + COUNT(*) FILTER (WHERE c.parent_chat_id IS NULL)::bigint AS queued_root_count, + COUNT(*) FILTER (WHERE c.parent_chat_id IS NOT NULL)::bigint AS queued_subagent_count +FROM chats c +WHERE c.status = 'running'::chat_status + AND c.archived = false + AND ( + c.worker_id IS NULL + OR c.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats hb + WHERE hb.chat_id = c.id + AND hb.runner_id = c.runner_id + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) + ) + ) +` + +type CountChatCapacityQueuedByPoolRow struct { + QueuedRootCount int64 `db:"queued_root_count" json:"queued_root_count"` + QueuedSubagentCount int64 `db:"queued_subagent_count" json:"queued_subagent_count"` +} + +func (q *sqlQuerier) CountChatCapacityQueuedByPool(ctx context.Context, staleSeconds int32) (CountChatCapacityQueuedByPoolRow, error) { + row := q.db.QueryRowContext(ctx, countChatCapacityQueuedByPool, staleSeconds) + var i CountChatCapacityQueuedByPoolRow + err := row.Scan(&i.QueuedRootCount, &i.QueuedSubagentCount) + return i, err +} + const countChatQueuedMessages = `-- name: CountChatQueuedMessages :one SELECT COUNT(*)::bigint AS count FROM chat_queued_messages @@ -8507,6 +8570,62 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get return items, nil } +const getChatQueuedForCapacity = `-- name: GetChatQueuedForCapacity :one +WITH active AS ( + SELECT + COUNT(*) FILTER (WHERE a.parent_chat_id IS NULL)::bigint AS root_count, + COUNT(*) FILTER (WHERE a.parent_chat_id IS NOT NULL)::bigint AS subagent_count + FROM chat_heartbeats hb + JOIN chats a + ON a.id = hb.chat_id + AND a.runner_id = hb.runner_id + WHERE a.worker_id IS NOT NULL + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) +) +SELECT ( + c.status = 'running'::chat_status + AND c.archived = false + AND ( + c.worker_id IS NULL + OR c.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats hb + WHERE hb.chat_id = c.id + AND hb.runner_id = c.runner_id + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) + ) + ) + AND CASE + WHEN c.parent_chat_id IS NULL THEN active.root_count >= $2::bigint + ELSE active.subagent_count >= $3::bigint + END +)::boolean AS queued_for_capacity +FROM chats c +CROSS JOIN active +WHERE c.id = $4::uuid +` + +type GetChatQueuedForCapacityParams struct { + StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` + RootCapacity int64 `db:"root_capacity" json:"root_capacity"` + SubagentCapacity int64 `db:"subagent_capacity" json:"subagent_capacity"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` +} + +// Pool fullness distinguishes capacity waits from worker pickup delays. +func (q *sqlQuerier) GetChatQueuedForCapacity(ctx context.Context, arg GetChatQueuedForCapacityParams) (bool, error) { + row := q.db.QueryRowContext(ctx, getChatQueuedForCapacity, + arg.StaleSeconds, + arg.RootCapacity, + arg.SubagentCapacity, + arg.ChatID, + ) + var queued_for_capacity bool + err := row.Scan(&queued_for_capacity) + return queued_for_capacity, err +} + const getChatQueuedMessageByID = `-- name: GetChatQueuedMessageByID :one SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages WHERE id = $1::bigint AND chat_id = $2::uuid @@ -8759,108 +8878,80 @@ func (q *sqlQuerier) GetChatUserPromptsByChatID(ctx context.Context, arg GetChat } const getChatWorkerAcquisitionCandidates = `-- name: GetChatWorkerAcquisitionCandidates :many +WITH candidate_partitions AS ( + SELECT true AS is_root, 'interrupting'::chat_status AS status, 0 AS status_priority, 0 AS pool_priority + UNION ALL + SELECT false, 'interrupting'::chat_status, 0, 1 + UNION ALL + SELECT true, 'requires_action'::chat_status, 1, 0 + UNION ALL + SELECT false, 'requires_action'::chat_status, 1, 1 + UNION ALL + SELECT true, 'running'::chat_status, 2, 0 + UNION ALL + SELECT false, 'running'::chat_status, 2, 1 +), +candidates AS ( + SELECT + candidate.id, + candidate_partitions.status_priority, + candidate_partitions.pool_priority, + ROW_NUMBER() OVER ( + PARTITION BY candidate_partitions.status_priority, candidate_partitions.is_root + ORDER BY candidate.updated_at ASC, candidate.id ASC + ) AS pool_position + FROM candidate_partitions + CROSS JOIN LATERAL ( + SELECT chats.id, chats.updated_at + FROM chats + WHERE (chats.parent_chat_id IS NULL) = candidate_partitions.is_root + AND chats.status = candidate_partitions.status + AND chats.archived = false + AND ( + chats.worker_id IS NULL + OR chats.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats.id + AND current_lease.runner_id = chats.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * $2::int) + ) + ) + ORDER BY chats.updated_at ASC, chats.id ASC + LIMIT $1::int + ) candidate +) SELECT - chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.last_reasoning_effort, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.summary, chats_expanded.summary_generated_at, chats_expanded.snapshot_version, chats_expanded.history_version, chats_expanded.queue_version, chats_expanded.generation_attempt, chats_expanded.retry_state, chats_expanded.retry_state_version, chats_expanded.runner_id, chats_expanded.requires_action_deadline_at, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, - chat_heartbeats.heartbeat_at AS current_heartbeat_at, - NOT EXISTS ( - SELECT 1 - FROM chat_heartbeats current_lease - WHERE current_lease.chat_id = chats_expanded.id - AND current_lease.runner_id = chats_expanded.runner_id - AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) - ) AS heartbeat_stale -FROM chats_expanded -LEFT JOIN chat_heartbeats - ON chat_heartbeats.chat_id = chats_expanded.id - AND chat_heartbeats.runner_id = chats_expanded.runner_id -WHERE - chats_expanded.status IN ('running'::chat_status, 'interrupting'::chat_status, 'requires_action'::chat_status) - AND chats_expanded.archived = false - AND ( - chats_expanded.worker_id IS NULL - OR chats_expanded.runner_id IS NULL - OR NOT EXISTS ( - SELECT 1 - FROM chat_heartbeats current_lease - WHERE current_lease.chat_id = chats_expanded.id - AND current_lease.runner_id = chats_expanded.runner_id - AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) - ) - ) -ORDER BY chats_expanded.updated_at ASC, chats_expanded.id ASC -LIMIT $2::int + chats.id, + chats.status, + chats.parent_chat_id +FROM candidates +JOIN chats ON chats.id = candidates.id +ORDER BY + candidates.status_priority ASC, + candidates.pool_position ASC, + candidates.pool_priority ASC, + chats.id ASC +LIMIT $1::int ` type GetChatWorkerAcquisitionCandidatesParams struct { - StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` LimitCount int32 `db:"limit_count" json:"limit_count"` + StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` } type GetChatWorkerAcquisitionCandidatesRow struct { - ID uuid.UUID `db:"id" json:"id"` - OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` - WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` - Title string `db:"title" json:"title"` - Status ChatStatus `db:"status" json:"status"` - WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` - StartedAt sql.NullTime `db:"started_at" json:"started_at"` - HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` - CreatedAt time.Time `db:"created_at" json:"created_at"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` - ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` - RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` - LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` - LastReasoningEffort NullChatReasoningEffort `db:"last_reasoning_effort" json:"last_reasoning_effort"` - Archived bool `db:"archived" json:"archived"` - LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` - Mode NullChatMode `db:"mode" json:"mode"` - MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` - Labels StringMap `db:"labels" json:"labels"` - BuildID uuid.NullUUID `db:"build_id" json:"build_id"` - AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` - PinOrder int32 `db:"pin_order" json:"pin_order"` - LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` - DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` - OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` - PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` - ClientType ChatClientType `db:"client_type" json:"client_type"` - LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` - Summary sql.NullString `db:"summary" json:"summary"` - SummaryGeneratedAt sql.NullTime `db:"summary_generated_at" json:"summary_generated_at"` - SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` - HistoryVersion int64 `db:"history_version" json:"history_version"` - QueueVersion int64 `db:"queue_version" json:"queue_version"` - GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` - RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` - RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` - RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` - RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` - UserACL ChatACL `db:"user_acl" json:"user_acl"` - GroupACL ChatACL `db:"group_acl" json:"group_acl"` - OwnerUsername string `db:"owner_username" json:"owner_username"` - OwnerName string `db:"owner_name" json:"owner_name"` - ContextAggregateHash []byte `db:"context_aggregate_hash" json:"context_aggregate_hash"` - ContextDirtySince sql.NullTime `db:"context_dirty_since" json:"context_dirty_since"` - ContextDirtyResources pqtype.NullRawMessage `db:"context_dirty_resources" json:"context_dirty_resources"` - ContextError string `db:"context_error" json:"context_error"` - CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` - CurrentHeartbeatAt sql.NullTime `db:"current_heartbeat_at" json:"current_heartbeat_at"` - HeartbeatStale bool `db:"heartbeat_stale" json:"heartbeat_stale"` + ID uuid.UUID `db:"id" json:"id"` + Status ChatStatus `db:"status" json:"status"` + ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` } -// Returns chats that workers may try to acquire. Candidates must be: -// - in a worker-runnable execution status; -// - unarchived; and -// - missing ownership, carrying inconsistent ownership, or lacking a -// fresh heartbeat for the assigned runner. -// -// Missing ownership is worker_id IS NULL. Inconsistent ownership is -// runner_id IS NULL while worker_id is set. Stale ownership is no -// heartbeat row for (chat_id, runner_id), or one older than -// @stale_seconds by database time. Candidates are ordered by oldest -// updated_at first so workers drain stale runnable chats predictably. +// Returns a bounded, pool-interleaved set of chats that workers may acquire. +// Interrupting chats finish active work first. Requires-action chats follow so +// their runner can enforce the action deadline before new generations start. func (q *sqlQuerier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg GetChatWorkerAcquisitionCandidatesParams) ([]GetChatWorkerAcquisitionCandidatesRow, error) { - rows, err := q.db.QueryContext(ctx, getChatWorkerAcquisitionCandidates, arg.StaleSeconds, arg.LimitCount) + rows, err := q.db.QueryContext(ctx, getChatWorkerAcquisitionCandidates, arg.LimitCount, arg.StaleSeconds) if err != nil { return nil, err } @@ -8868,57 +8959,7 @@ func (q *sqlQuerier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg var items []GetChatWorkerAcquisitionCandidatesRow for rows.Next() { var i GetChatWorkerAcquisitionCandidatesRow - if err := rows.Scan( - &i.ID, - &i.OwnerID, - &i.WorkspaceID, - &i.Title, - &i.Status, - &i.WorkerID, - &i.StartedAt, - &i.HeartbeatAt, - &i.CreatedAt, - &i.UpdatedAt, - &i.ParentChatID, - &i.RootChatID, - &i.LastModelConfigID, - &i.LastReasoningEffort, - &i.Archived, - &i.LastError, - &i.Mode, - pq.Array(&i.MCPServerIDs), - &i.Labels, - &i.BuildID, - &i.AgentID, - &i.PinOrder, - &i.LastReadMessageID, - &i.DynamicTools, - &i.OrganizationID, - &i.PlanMode, - &i.ClientType, - &i.LastTurnSummary, - &i.Summary, - &i.SummaryGeneratedAt, - &i.SnapshotVersion, - &i.HistoryVersion, - &i.QueueVersion, - &i.GenerationAttempt, - &i.RetryState, - &i.RetryStateVersion, - &i.RunnerID, - &i.RequiresActionDeadlineAt, - &i.UserACL, - &i.GroupACL, - &i.OwnerUsername, - &i.OwnerName, - &i.ContextAggregateHash, - &i.ContextDirtySince, - &i.ContextDirtyResources, - &i.ContextError, - &i.CompactionRequestedAt, - &i.CurrentHeartbeatAt, - &i.HeartbeatStale, - ); err != nil { + if err := rows.Scan(&i.ID, &i.Status, &i.ParentChatID); err != nil { return nil, err } items = append(items, i) diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 1bb5549292..a320e0163d 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -2303,46 +2303,64 @@ WHERE chat_id = @chat_id::uuid AND content::jsonb @> '[{"type": "context-file"}]'; -- name: GetChatWorkerAcquisitionCandidates :many --- Returns chats that workers may try to acquire. Candidates must be: --- - in a worker-runnable execution status; --- - unarchived; and --- - missing ownership, carrying inconsistent ownership, or lacking a --- fresh heartbeat for the assigned runner. --- --- Missing ownership is worker_id IS NULL. Inconsistent ownership is --- runner_id IS NULL while worker_id is set. Stale ownership is no --- heartbeat row for (chat_id, runner_id), or one older than --- @stale_seconds by database time. Candidates are ordered by oldest --- updated_at first so workers drain stale runnable chats predictably. +-- Returns a bounded, pool-interleaved set of chats that workers may acquire. +-- Interrupting chats finish active work first. Requires-action chats follow so +-- their runner can enforce the action deadline before new generations start. +WITH candidate_partitions AS ( + SELECT true AS is_root, 'interrupting'::chat_status AS status, 0 AS status_priority, 0 AS pool_priority + UNION ALL + SELECT false, 'interrupting'::chat_status, 0, 1 + UNION ALL + SELECT true, 'requires_action'::chat_status, 1, 0 + UNION ALL + SELECT false, 'requires_action'::chat_status, 1, 1 + UNION ALL + SELECT true, 'running'::chat_status, 2, 0 + UNION ALL + SELECT false, 'running'::chat_status, 2, 1 +), +candidates AS ( + SELECT + candidate.id, + candidate_partitions.status_priority, + candidate_partitions.pool_priority, + ROW_NUMBER() OVER ( + PARTITION BY candidate_partitions.status_priority, candidate_partitions.is_root + ORDER BY candidate.updated_at ASC, candidate.id ASC + ) AS pool_position + FROM candidate_partitions + CROSS JOIN LATERAL ( + SELECT chats.id, chats.updated_at + FROM chats + WHERE (chats.parent_chat_id IS NULL) = candidate_partitions.is_root + AND chats.status = candidate_partitions.status + AND chats.archived = false + AND ( + chats.worker_id IS NULL + OR chats.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats.id + AND current_lease.runner_id = chats.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) + ) + ) + ORDER BY chats.updated_at ASC, chats.id ASC + LIMIT @limit_count::int + ) candidate +) SELECT - chats_expanded.*, - chat_heartbeats.heartbeat_at AS current_heartbeat_at, - NOT EXISTS ( - SELECT 1 - FROM chat_heartbeats current_lease - WHERE current_lease.chat_id = chats_expanded.id - AND current_lease.runner_id = chats_expanded.runner_id - AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) - ) AS heartbeat_stale -FROM chats_expanded -LEFT JOIN chat_heartbeats - ON chat_heartbeats.chat_id = chats_expanded.id - AND chat_heartbeats.runner_id = chats_expanded.runner_id -WHERE - chats_expanded.status IN ('running'::chat_status, 'interrupting'::chat_status, 'requires_action'::chat_status) - AND chats_expanded.archived = false - AND ( - chats_expanded.worker_id IS NULL - OR chats_expanded.runner_id IS NULL - OR NOT EXISTS ( - SELECT 1 - FROM chat_heartbeats current_lease - WHERE current_lease.chat_id = chats_expanded.id - AND current_lease.runner_id = chats_expanded.runner_id - AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) - ) - ) -ORDER BY chats_expanded.updated_at ASC, chats_expanded.id ASC + chats.id, + chats.status, + chats.parent_chat_id +FROM candidates +JOIN chats ON chats.id = candidates.id +ORDER BY + candidates.status_priority ASC, + candidates.pool_position ASC, + candidates.pool_priority ASC, + chats.id ASC LIMIT @limit_count::int; -- name: GetChatsByIDsForRunnerSync :many @@ -2800,3 +2818,71 @@ LEFT JOIN to_archive t ON t.id = a.id -- created_at ASC flows through to dbpurge's digest truncation; see -- buildDigestData in dbpurge.go for the tradeoff rationale. ORDER BY (a.root_chat_id IS NULL) DESC, a.owner_id ASC, a.created_at ASC, a.id ASC; + +-- name: CountChatCapacityActiveByPool :one +-- Excluding the candidate keeps ownership takeover capacity-neutral. +SELECT + COUNT(*) FILTER (WHERE c.parent_chat_id IS NULL)::bigint AS active_root_count, + COUNT(*) FILTER (WHERE c.parent_chat_id IS NOT NULL)::bigint AS active_subagent_count +FROM chat_heartbeats hb +JOIN chats c + ON c.id = hb.chat_id + AND c.runner_id = hb.runner_id +WHERE c.worker_id IS NOT NULL + AND c.id != @exclude_chat_id::uuid + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int); + +-- name: CountChatCapacityQueuedByPool :one +SELECT + COUNT(*) FILTER (WHERE c.parent_chat_id IS NULL)::bigint AS queued_root_count, + COUNT(*) FILTER (WHERE c.parent_chat_id IS NOT NULL)::bigint AS queued_subagent_count +FROM chats c +WHERE c.status = 'running'::chat_status + AND c.archived = false + AND ( + c.worker_id IS NULL + OR c.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats hb + WHERE hb.chat_id = c.id + AND hb.runner_id = c.runner_id + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) + ) + ); + +-- name: GetChatQueuedForCapacity :one +-- Pool fullness distinguishes capacity waits from worker pickup delays. +WITH active AS ( + SELECT + COUNT(*) FILTER (WHERE a.parent_chat_id IS NULL)::bigint AS root_count, + COUNT(*) FILTER (WHERE a.parent_chat_id IS NOT NULL)::bigint AS subagent_count + FROM chat_heartbeats hb + JOIN chats a + ON a.id = hb.chat_id + AND a.runner_id = hb.runner_id + WHERE a.worker_id IS NOT NULL + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) +) +SELECT ( + c.status = 'running'::chat_status + AND c.archived = false + AND ( + c.worker_id IS NULL + OR c.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats hb + WHERE hb.chat_id = c.id + AND hb.runner_id = c.runner_id + AND hb.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) + ) + ) + AND CASE + WHEN c.parent_chat_id IS NULL THEN active.root_count >= @root_capacity::bigint + ELSE active.subagent_count >= @subagent_capacity::bigint + END +)::boolean AS queued_for_capacity +FROM chats c +CROSS JOIN active +WHERE c.id = @chat_id::uuid; diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index b7a0c65346..5dcaad9e5c 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -1621,6 +1621,18 @@ func (api *API) getChat(rw http.ResponseWriter, r *http.Request) { sdkChat := db2sdk.Chat(chat, diffStatus, chatFiles) + if api.chatDaemon != nil { + queued, err := api.chatDaemon.ChatQueuedForCapacity(ctx, chat) + if err != nil { + api.Logger.Error(ctx, "failed to derive chat queued-for-capacity state", + slog.F("chat_id", chat.ID), + slog.Error(err), + ) + } else { + sdkChat.QueuedForCapacity = queued + } + } + // Enrich the lightweight context summary with the chat's pinned // resources (metadata only). This detail is computed on read and only // attached on the single-chat GET; list and watch payloads stay diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index bef7b57116..d9f1e16f8f 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -947,6 +947,10 @@ The abandon chat goroutine is responsible for abandoning the chat. It is spawned When the manager cleans up a runner, the runner must cancel all goroutines it has spawned and unsubscribe from pubsub. +## Concurrent agent limiter + +By default, chatd runs up to five top-level chats and ten subagent chats at once. Each limit applies across the entire deployment. Enterprise deployments can remove these limits when their plan permits it. Extra chats wait for capacity, but users can still interrupt active chats. + ## Auto-archive loop The worker periodically archives old, unused chats. diff --git a/coderd/x/chatd/agentadmission.go b/coderd/x/chatd/agentadmission.go new file mode 100644 index 0000000000..0809b6a008 --- /dev/null +++ b/coderd/x/chatd/agentadmission.go @@ -0,0 +1,85 @@ +package chatd + +import ( + "context" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" +) + +const ( + defaultMaxConcurrentRootAgents = int64(5) + defaultMaxConcurrentSubagents = int64(10) +) + +// AgentCapacityLimiter controls chat admission and reports the current per-pool limits. +type AgentCapacityLimiter interface { + // Admit runs inside the acquisition transaction so its serialization + // extends through the ownership write. Refused chats remain unowned. + Admit(ctx context.Context, store database.Store, chat database.Chat) (bool, error) + Limits() (limits AgentCapacityLimits, capped bool) +} + +// AgentCapacityUnlock reports whether the default chat agent caps are disabled. +type AgentCapacityUnlock interface { + Unlocked() bool +} + +// AgentCapacityLimits defines concurrent-agent limits for root and subagent pools. +type AgentCapacityLimits struct { + Root int64 + Subagent int64 +} + +type agentCapacityLimiter struct { + unlock AgentCapacityUnlock + + staleSeconds int32 + rootCapacity int64 + subagentCapacity int64 +} + +func newAgentCapacityLimiter(unlock AgentCapacityUnlock, staleSeconds int32) *agentCapacityLimiter { + return &agentCapacityLimiter{ + unlock: unlock, + staleSeconds: staleSeconds, + rootCapacity: defaultMaxConcurrentRootAgents, + subagentCapacity: defaultMaxConcurrentSubagents, + } +} + +func (a *agentCapacityLimiter) Admit(ctx context.Context, store database.Store, chat database.Chat) (bool, error) { + //nolint:gocritic // Capacity accounting is chatd-internal state. + ctx = dbauthz.AsChatd(ctx) + if a.unlocked() || chat.Status != database.ChatStatusRunning { + return true, nil + } + // The transaction lock remains held through the caller's ownership write, + // preventing replicas from over-admitting the pool. + if err := store.AcquireLock(ctx, database.LockIDChatCapacityAdmission); err != nil { + return false, err + } + counts, err := store.CountChatCapacityActiveByPool(ctx, database.CountChatCapacityActiveByPoolParams{ + ExcludeChatID: chat.ID, + StaleSeconds: a.staleSeconds, + }) + if err != nil { + return false, err + } + used, capacity := counts.ActiveRootCount, a.rootCapacity + if chat.ParentChatID.Valid { + used, capacity = counts.ActiveSubagentCount, a.subagentCapacity + } + return used < capacity, nil +} + +func (a *agentCapacityLimiter) Limits() (AgentCapacityLimits, bool) { + return AgentCapacityLimits{ + Root: a.rootCapacity, + Subagent: a.subagentCapacity, + }, !a.unlocked() +} + +func (a *agentCapacityLimiter) unlocked() bool { + return a.unlock != nil && a.unlock.Unlocked() +} diff --git a/coderd/x/chatd/agentadmission_internal_test.go b/coderd/x/chatd/agentadmission_internal_test.go new file mode 100644 index 0000000000..b567fa3464 --- /dev/null +++ b/coderd/x/chatd/agentadmission_internal_test.go @@ -0,0 +1,487 @@ +package chatd + +import ( + "context" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + promtestutil "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/testutil" +) + +type fakeAdmission struct { + mu sync.Mutex + refused map[uuid.UUID]bool + refuseFn func(database.Chat) bool + + admitCalls int + admitted []uuid.UUID + limits AgentCapacityLimits + uncapped bool +} + +func newFakeAdmission() *fakeAdmission { + return &fakeAdmission{ + refused: make(map[uuid.UUID]bool), + limits: AgentCapacityLimits{Root: 1, Subagent: 1}, + } +} + +func (f *fakeAdmission) Limits() (AgentCapacityLimits, bool) { + return f.limits, !f.uncapped +} + +func (f *fakeAdmission) refuse(chatID uuid.UUID) { + f.mu.Lock() + defer f.mu.Unlock() + f.refused[chatID] = true +} + +func (f *fakeAdmission) allow(chatID uuid.UUID) { + f.mu.Lock() + defer f.mu.Unlock() + delete(f.refused, chatID) +} + +func (f *fakeAdmission) Admit(_ context.Context, _ database.Store, chat database.Chat) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.admitCalls++ + if f.refused[chat.ID] || (f.refuseFn != nil && f.refuseFn(chat)) { + return false, nil + } + f.admitted = append(f.admitted, chat.ID) + return true, nil +} + +func (f *fakeAdmission) admittedOrder() []uuid.UUID { + f.mu.Lock() + defer f.mu.Unlock() + return append([]uuid.UUID(nil), f.admitted...) +} + +func (f *fakeAdmission) admitCallCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.admitCalls +} + +func TestWorker_AdmissionRefusalDoesNotAcquireChat(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + recording := newRecordingPubsub(f.pubsub) + starter := newRecordingTaskStarter() + admission := newFakeAdmission() + opts := testOptions(t, f, starter) + opts.Pubsub = recording + opts.AgentCapacityLimiter = admission + + chat := f.createRunningChat(t) + admission.refuse(chat.ID) + startWorker(t, opts) + + require.Eventually(t, func() bool { + return admission.admitCallCount() > 0 + }, testutil.WaitLong, testutil.IntervalFast) + starter.assertNoCall(t) + + // The recorder wraps only worker pubsub, so an ownership hint here would + // prove a refusal can wake workers into an immediate retry loop. + require.Empty(t, recording.ownershipMessages(t)) +} + +func TestWorker_InterruptingSortsBeforeRunning(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + running := []database.Chat{f.createRunningChat(t), f.createRunningChat(t)} + interrupting := f.createRunningChat(t) + interruptChat(t, f, interrupting.ID) + requiresAction := f.createRequiresActionChat(t) + _, err := f.sqlDB.ExecContext(ctx, ` + UPDATE chats + SET updated_at = NOW() - INTERVAL '1 hour' + WHERE id IN ($1, $2) + `, running[0].ID, running[1].ID) + require.NoError(t, err) + + rows, err := f.db.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: 30, + LimitCount: 2, + }) + require.NoError(t, err) + require.Len(t, rows, 2) + require.Equal(t, interrupting.ID, rows[0].ID) + require.Equal(t, requiresAction.ID, rows[1].ID) +} + +func TestWorker_AcquisitionCandidatesInterleavePools(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + rootOlder := f.createRunningChat(t) + rootNewer := f.createRunningChat(t) + subOlder := f.createRunningSubagentChat(t, rootOlder.ID) + subNewer := f.createRunningSubagentChat(t, rootOlder.ID) + _, err := f.sqlDB.ExecContext(ctx, ` + UPDATE chats + SET updated_at = CASE id + WHEN $1 THEN NOW() - INTERVAL '4 hours' + WHEN $2 THEN NOW() - INTERVAL '3 hours' + WHEN $3 THEN NOW() - INTERVAL '2 hours' + WHEN $4 THEN NOW() - INTERVAL '1 hour' + END + WHERE id IN ($1, $2, $3, $4) + `, rootOlder.ID, subOlder.ID, rootNewer.ID, subNewer.ID) + require.NoError(t, err) + + rows, err := f.db.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: 30, + LimitCount: 4, + }) + require.NoError(t, err) + require.Len(t, rows, 4) + require.Equal(t, []uuid.UUID{rootOlder.ID, subOlder.ID, rootNewer.ID, subNewer.ID}, []uuid.UUID{ + rows[0].ID, + rows[1].ID, + rows[2].ID, + rows[3].ID, + }) +} + +func TestWorker_MessageBumpSendsChatToQueueBack(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + older := f.createRunningChat(t) + newer := f.createRunningChat(t) + _, err := f.sqlDB.ExecContext(ctx, ` + UPDATE chats + SET updated_at = CASE id + WHEN $1 THEN NOW() - INTERVAL '1 hour' + WHEN $2 THEN NOW() - INTERVAL '30 minutes' + END + WHERE id IN ($1, $2) + `, older.ID, newer.ID) + require.NoError(t, err) + + machine := chatstate.NewChatMachine(f.db, f.pubsub, older.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage(t, "move me", f.user.ID, f.model.ID, f.apiKey.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + + rows, err := f.db.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: 30, + LimitCount: 10, + }) + require.NoError(t, err) + require.GreaterOrEqual(t, len(rows), 2) + require.Equal(t, newer.ID, rows[0].ID) + require.Equal(t, older.ID, rows[1].ID) +} + +func newRootRefusingAdmission() *fakeAdmission { + admission := newFakeAdmission() + admission.refuseFn = func(chat database.Chat) bool { + return chat.Status == database.ChatStatusRunning && !chat.ParentChatID.Valid + } + return admission +} + +func TestWorker_FullPoolDoesNotStarveOtherPool(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + opts.AgentCapacityLimiter = newRootRefusingAdmission() + + roots := make([]database.Chat, 0, 2*int(opts.AcquisitionBatchSize)+5) + for range cap(roots) { + roots = append(roots, f.createRunningChat(t)) + } + sub := f.createRunningSubagentChat(t, roots[0].ID) + + startWorker(t, opts) + + call := starter.waitCall(t, taskKindGeneration, sub.ID) + require.Equal(t, sub.ID, call.input.ChatID) +} + +func TestWorker_BatchSizeOneCannotHideAPool(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + opts.AgentCapacityLimiter = newRootRefusingAdmission() + opts.AcquisitionBatchSize = 1 + + roots := []database.Chat{f.createRunningChat(t), f.createRunningChat(t)} + sub := f.createRunningSubagentChat(t, roots[0].ID) + + startWorker(t, opts) + + call := starter.waitCall(t, taskKindGeneration, sub.ID) + require.Equal(t, sub.ID, call.input.ChatID) +} + +func TestWorker_FullPoolSkipsRefusalsAfterFirst(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + admission := newRootRefusingAdmission() + opts.AgentCapacityLimiter = admission + opts.AcquisitionBatchSize = 2 + + for range 5 { + f.createRunningChat(t) + } + startWorker(t, opts) + + require.Eventually(t, func() bool { + return admission.admitCallCount() == 1 + }, testutil.WaitLong, testutil.IntervalFast) + starter.assertNoCall(t) + require.Equal(t, 1, admission.admitCallCount(), + "a full pool must be skipped after one refusal, not re-refused per chat") +} + +func TestWorker_AdmissionAdmitsInUpdatedAtOrder(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + admission := newFakeAdmission() + opts.AgentCapacityLimiter = admission + + older := f.createRunningChat(t) + newer := f.createRunningChat(t) + // Back-to-back inserts can collide at timestamp resolution, which would + // leave FIFO order to the random UUID tiebreak. + ctx := testutil.Context(t, testutil.WaitLong) + _, err := f.sqlDB.ExecContext(ctx, + "UPDATE chats SET updated_at = NOW() - INTERVAL '1 hour' WHERE id = $1", older.ID) + require.NoError(t, err) + admission.refuse(older.ID) + admission.refuse(newer.ID) + worker := startWorker(t, opts) + + require.Eventually(t, func() bool { + return admission.admitCallCount() == 1 + }, testutil.WaitLong, testutil.IntervalFast) + + admission.allow(older.ID) + admission.allow(newer.ID) + worker.Wake() + + // Runner goroutines race task starts, so wait for both without + // ordering and assert the worker's serial admission order instead. + starter.waitCall(t, taskKindGeneration, uuid.Nil) + starter.waitCall(t, taskKindGeneration, uuid.Nil) + require.Equal(t, []uuid.UUID{older.ID, newer.ID}, admission.admittedOrder(), + "the longer-waiting chat must admit first") +} + +func TestWorker_InterruptClaimsCapacityQueuedChat(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + admission := newFakeAdmission() + admission.refuseFn = func(chat database.Chat) bool { + return chat.Status == database.ChatStatusRunning + } + opts.AgentCapacityLimiter = admission + + chat := f.createRunningChat(t) + worker := startWorker(t, opts) + + require.Eventually(t, func() bool { + return admission.admitCallCount() > 0 + }, testutil.WaitLong, testutil.IntervalFast) + + interruptChat(t, f, chat.ID) + worker.Wake() + + call := starter.waitCall(t, taskKindInterrupt, chat.ID) + require.Equal(t, chat.ID, call.input.ChatID) +} + +func TestWorker_CapacityMetricsUseFreshOwnership(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + metrics := newCapacityMetrics(prometheus.NewRegistry()) + opts := testOptions(t, f, newRecordingTaskStarter()) + opts.CapacityMetrics = metrics + opts.AgentCapacityLimiter = newFakeAdmission() + ctx := testutil.Context(t, testutil.WaitLong) + + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + f.createRunningChat(t) + + worker, err := newChatWorker(newUnstartedServer(t, f.pubsub, f.db), opts) + require.NoError(t, err) + worker.refreshCapacityMetrics(ctx) + + require.Equal(t, float64(1), promtestutil.ToFloat64(metrics.active.WithLabelValues("root"))) + require.Equal(t, float64(1), promtestutil.ToFloat64(metrics.queued.WithLabelValues("root"))) + + forceExecutionState(t, f, occupied.ID, database.ChatStatusWaiting, true) + worker.refreshCapacityMetrics(ctx) + require.Equal(t, float64(1), promtestutil.ToFloat64(metrics.active.WithLabelValues("root"))) + require.Equal(t, float64(1), promtestutil.ToFloat64(metrics.queued.WithLabelValues("root"))) +} + +func TestGetChatQueuedForCapacity(t *testing.T) { + t.Parallel() + + queued := func(t *testing.T, f *workerTestFixture, chatID uuid.UUID, rootCap, subagentCap int64) bool { + t.Helper() + ctx := testutil.Context(t, testutil.WaitLong) + got, err := f.db.GetChatQueuedForCapacity(ctx, database.GetChatQueuedForCapacityParams{ + ChatID: chatID, + StaleSeconds: 30, + RootCapacity: rootCap, + SubagentCapacity: subagentCap, + }) + require.NoError(t, err) + return got + } + + t.Run("PoolNotFull", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + require.False(t, queued(t, f, chat.ID, 1, 1)) + }) + + t.Run("PoolFull", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + chat := f.createRunningChat(t) + require.True(t, queued(t, f, chat.ID, 1, 1)) + }) + + t.Run("IncompleteOwnershipIsQueued", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + chat := f.createRunningChat(t) + acquireChat(t, f, chat.ID, uuid.New(), uuid.New()) + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitLong), `UPDATE chats SET worker_id = NULL WHERE id = $1`, chat.ID) + require.NoError(t, err) + require.True(t, queued(t, f, chat.ID, 1, 1)) + }) + + t.Run("OwnedChatIsNotQueued", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + require.False(t, queued(t, f, occupied.ID, 1, 1)) + }) + + t.Run("NonRunningChatIsNotQueued", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + requiresAction := f.createRequiresActionChat(t) + require.False(t, queued(t, f, requiresAction.ID, 1, 1)) + }) + + t.Run("PoolsAreIndependent", func(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + sub := f.createRunningSubagentChat(t, occupied.ID) + require.False(t, queued(t, f, sub.ID, 1, 1), + "a full root pool must not mark subagents queued") + }) +} + +func TestServer_ChatQueuedForCapacity(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + occupied := f.createRunningChat(t) + acquireChat(t, f, occupied.ID, uuid.New(), uuid.New()) + for range 4 { + chat := f.createRunningChat(t) + acquireChat(t, f, chat.ID, uuid.New(), uuid.New()) + } + waiting := f.createRunningChat(t) + + server := newUnstartedServer(t, f.pubsub, f.db) + + queued, err := server.ChatQueuedForCapacity(ctx, waiting) + require.NoError(t, err) + require.True(t, queued, "AGPL deployments must enforce the default root capacity") + + uncapped := newFakeAdmission() + uncapped.uncapped = true + server.agentCapacityLimiter = uncapped + queued, err = server.ChatQueuedForCapacity(ctx, waiting) + require.NoError(t, err) + require.False(t, queued, "uncapped deployments must never report queued") + + server.agentCapacityLimiter = newFakeAdmission() + queued, err = server.ChatQueuedForCapacity(ctx, waiting) + require.NoError(t, err) + require.True(t, queued) + + queued, err = server.ChatQueuedForCapacity(ctx, occupied) + require.NoError(t, err) + require.False(t, queued, "owned chats are active, not queued") +} + +func TestChatCapacityCountsByPool(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + owned := f.createRunningChat(t) + acquireChat(t, f, owned.ID, uuid.New(), uuid.New()) + f.createRunningChat(t) + f.createRunningSubagentChat(t, owned.ID) + incomplete := f.createRunningChat(t) + acquireChat(t, f, incomplete.ID, uuid.New(), uuid.New()) + _, err := f.sqlDB.ExecContext(ctx, `UPDATE chats SET worker_id = NULL WHERE id = $1`, incomplete.ID) + require.NoError(t, err) + + active, err := f.db.CountChatCapacityActiveByPool(ctx, database.CountChatCapacityActiveByPoolParams{StaleSeconds: 30}) + require.NoError(t, err) + require.EqualValues(t, 1, active.ActiveRootCount) + require.EqualValues(t, 0, active.ActiveSubagentCount) + + queued, err := f.db.CountChatCapacityQueuedByPool(ctx, 30) + require.NoError(t, err) + require.EqualValues(t, 2, queued.QueuedRootCount) + require.EqualValues(t, 1, queued.QueuedSubagentCount) + + active, err = f.db.CountChatCapacityActiveByPool(ctx, database.CountChatCapacityActiveByPoolParams{ + ExcludeChatID: owned.ID, + StaleSeconds: 30, + }) + require.NoError(t, err) + require.EqualValues(t, 0, active.ActiveRootCount, "the excluded chat must not count as active") +} diff --git a/coderd/x/chatd/agentcapacitylimiter_internal_test.go b/coderd/x/chatd/agentcapacitylimiter_internal_test.go new file mode 100644 index 0000000000..8936441dde --- /dev/null +++ b/coderd/x/chatd/agentcapacitylimiter_internal_test.go @@ -0,0 +1,266 @@ +package chatd + +import ( + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/testutil" +) + +type admissionFixture struct { + db database.Store + ps pubsub.Pubsub + owner database.User + org database.Organization + modelConfig database.ChatModelConfig +} + +func newAdmissionFixture(t *testing.T) admissionFixture { + t.Helper() + db, ps := dbtestutil.NewDB(t) + owner := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + return admissionFixture{db: db, ps: ps, owner: owner, org: org, modelConfig: modelConfig} +} + +func (f admissionFixture) chat(t *testing.T, seed database.Chat) database.Chat { + t.Helper() + seed.OwnerID = f.owner.ID + seed.OrganizationID = f.org.ID + seed.LastModelConfigID = f.modelConfig.ID + if seed.Status == "" { + seed.Status = database.ChatStatusRunning + } + return dbgen.Chat(t, f.db, seed) +} + +func (f admissionFixture) occupy(t *testing.T, chatID uuid.UUID) { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + machine := chatstate.NewChatMachine(f.db, f.ps, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: uuid.New(), RunnerID: uuid.New()}) + return err + })) +} + +func (f admissionFixture) occupiedRoot(t *testing.T) database.Chat { + t.Helper() + chat := f.chat(t, database.Chat{}) + f.occupy(t, chat.ID) + return chat +} + +func (f admissionFixture) occupiedSubagent(t *testing.T, root database.Chat) database.Chat { + t.Helper() + chat := f.chat(t, database.Chat{ + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + f.occupy(t, chat.ID) + return chat +} + +func testAdmission() *agentCapacityLimiter { + a := newAgentCapacityLimiter(nil, 30) + a.rootCapacity = 2 + a.subagentCapacity = 2 + return a +} + +func TestAdmission_RootPoolCap(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + root := f.occupiedRoot(t) + f.occupiedRoot(t) + + admitted, err := a.Admit(ctx, f.db, f.chat(t, database.Chat{})) + require.NoError(t, err) + require.False(t, admitted, "third root must be refused at capacity 2") + + subagent := f.chat(t, database.Chat{ + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + admitted, err = a.Admit(ctx, f.db, subagent) + require.NoError(t, err) + require.True(t, admitted) +} + +func TestAdmission_SubagentPoolCap(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + root := f.occupiedRoot(t) + f.occupiedSubagent(t, root) + f.occupiedSubagent(t, root) + + subagent := f.chat(t, database.Chat{ + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + admitted, err := a.Admit(ctx, f.db, subagent) + require.NoError(t, err) + require.False(t, admitted, "third subagent must be refused at capacity 2") + + admitted, err = a.Admit(ctx, f.db, f.chat(t, database.Chat{})) + require.NoError(t, err) + require.True(t, admitted, "a full subagent pool must not refuse roots") +} + +func TestAdmission_InterruptingBypassesCap(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + f.occupiedRoot(t) + f.occupiedRoot(t) + + interrupting := f.chat(t, database.Chat{Status: database.ChatStatusInterrupting}) + admitted, err := a.Admit(ctx, f.db, interrupting) + require.NoError(t, err) + require.True(t, admitted, "interrupting chats must always be acquirable") +} + +func TestAdmission_RequiresActionBypassesCap(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + f.occupiedRoot(t) + f.occupiedRoot(t) + + requiresAction := f.chat(t, database.Chat{Status: database.ChatStatusRequiresAction}) + admitted, err := a.Admit(ctx, f.db, requiresAction) + require.NoError(t, err) + require.True(t, admitted, "requires_action chats hold no slot and need their runner") +} + +func TestAdmission_TakeoverOfCountedChatIsCapacityNeutral(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + counted := f.occupiedRoot(t) + f.occupiedRoot(t) + + chat, err := f.db.GetChatByID(ctx, counted.ID) + require.NoError(t, err) + admitted, err := a.Admit(ctx, f.db, chat) + require.NoError(t, err) + require.True(t, admitted, "an already-counted chat must re-admit for takeover at full capacity") +} + +// The single transaction verifies that the admission lock covers the +// ownership write. +func TestAdmission_ConcurrentAdmitNeverOverAdmits(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + a := testAdmission() + + const attempts = 8 + chats := make([]database.Chat, attempts) + for i := range chats { + chats[i] = f.chat(t, database.Chat{}) + } + + errRefused := xerrors.New("refused") + var ( + admitted atomic.Int64 + unexpected atomic.Int64 + wg sync.WaitGroup + ) + for _, chat := range chats { + wg.Go(func() { + machine := chatstate.NewChatMachine(f.db, f.ps, chat.ID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + ok, err := a.Admit(ctx, store, chat) + if err != nil { + return err + } + if !ok { + return errRefused + } + _, err = tx.Acquire(chatstate.AcquireInput{WorkerID: uuid.New(), RunnerID: uuid.New()}) + return err + }) + switch { + case err == nil: + admitted.Add(1) + case errors.Is(err, errRefused): + default: + unexpected.Add(1) + } + }) + } + wg.Wait() + + require.EqualValues(t, 0, unexpected.Load(), "admission attempts must not error") + require.EqualValues(t, 2, admitted.Load(), "exactly rootCapacity chats must admit") +} + +func TestAdmission_StaleHeartbeatsFreeSlots(t *testing.T) { + t.Parallel() + f := newAdmissionFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + + f.occupiedRoot(t) + f.occupiedRoot(t) + + // A zero staleness window makes every heartbeat stale. + a := newAgentCapacityLimiter(nil, 0) + a.rootCapacity = 2 + a.subagentCapacity = 2 + admitted, err := a.Admit(ctx, f.db, f.chat(t, database.Chat{})) + require.NoError(t, err) + require.True(t, admitted) +} + +type staticAgentCapacityUnlock bool + +func (u staticAgentCapacityUnlock) Unlocked() bool { + return bool(u) +} + +func TestAdmission_UnlockBypassesCaps(t *testing.T) { + t.Parallel() + a := newAgentCapacityLimiter(staticAgentCapacityUnlock(true), 30) + + admitted, err := a.Admit(t.Context(), nil, database.Chat{Status: database.ChatStatusRunning}) + require.NoError(t, err) + require.True(t, admitted) + + _, capped := a.Limits() + require.False(t, capped) +} + +func TestAdmission_LimitsReportsCaps(t *testing.T) { + t.Parallel() + a := newAgentCapacityLimiter(nil, 30) + + limits, capped := a.Limits() + require.True(t, capped) + require.EqualValues(t, defaultMaxConcurrentRootAgents, limits.Root) + require.EqualValues(t, defaultMaxConcurrentSubagents, limits.Subagent) +} diff --git a/coderd/x/chatd/capacity.go b/coderd/x/chatd/capacity.go new file mode 100644 index 0000000000..1ebcf8669a --- /dev/null +++ b/coderd/x/chatd/capacity.go @@ -0,0 +1,84 @@ +package chatd + +import ( + "context" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + + "github.com/coder/coder/v2/coderd/database" +) + +type capacityMetrics struct { + active *prometheus.GaugeVec + queued *prometheus.GaugeVec +} + +func newCapacityMetrics(registerer prometheus.Registerer) *capacityMetrics { + m := &capacityMetrics{ + active: prometheus.NewGaugeVec(prometheus.GaugeOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "agents_active", + Help: "Deployment-wide number of chats holding a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum.", + }, []string{"pool"}), + queued: prometheus.NewGaugeVec(prometheus.GaugeOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "agents_queued_for_capacity", + Help: "Deployment-wide number of chats waiting for a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum.", + }, []string{"pool"}), + } + registerer.MustRegister(m.active, m.queued) + return m +} + +func (w *chatWorker) capacityMetricsLoop(ctx context.Context) { + ticker := w.opts.Clock.NewTicker(w.opts.CapacityMetricsInterval, "chatworker", "capacity-metrics") + defer ticker.Stop() + for { + select { + case <-ticker.C: + case <-ctx.Done(): + return + } + w.refreshCapacityMetrics(ctx) + } +} + +func (w *chatWorker) refreshCapacityMetrics(ctx context.Context) { + active, err := w.opts.Store.CountChatCapacityActiveByPool(ctx, database.CountChatCapacityActiveByPoolParams{ + ExcludeChatID: uuid.Nil, + StaleSeconds: w.opts.HeartbeatStaleSeconds, + }) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker count active capacity chats failed", slogError(err)) + } + return + } + + limits, capped := w.opts.AgentCapacityLimiter.Limits() + var queuedRoot, queuedSubagent int64 + if capped && (active.ActiveRootCount >= limits.Root || active.ActiveSubagentCount >= limits.Subagent) { + queued, err := w.opts.Store.CountChatCapacityQueuedByPool(ctx, w.opts.HeartbeatStaleSeconds) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker count queued capacity chats failed", slogError(err)) + } + return + } + if active.ActiveRootCount >= limits.Root { + queuedRoot = queued.QueuedRootCount + } + if active.ActiveSubagentCount >= limits.Subagent { + queuedSubagent = queued.QueuedSubagentCount + } + } + + metrics := w.opts.CapacityMetrics + metrics.active.WithLabelValues("root").Set(float64(active.ActiveRootCount)) + metrics.active.WithLabelValues("subagent").Set(float64(active.ActiveSubagentCount)) + metrics.queued.WithLabelValues("root").Set(float64(queuedRoot)) + metrics.queued.WithLabelValues("subagent").Set(float64(queuedSubagent)) +} diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 792e72bcd6..41966f45da 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -189,13 +189,14 @@ type Server struct { configCacheUnsubscribe func() providerCacheUnsubscribe func() - usageTracker *workspacestats.UsageTracker - clock quartz.Clock - metrics *chatloop.Metrics - chatWorker *chatWorker - messagePartBuffer *messagepartbuffer.Buffer - streamSyncPoller *streamSyncPoller - recordingSem chan struct{} + usageTracker *workspacestats.UsageTracker + clock quartz.Clock + metrics *chatloop.Metrics + chatWorker *chatWorker + messagePartBuffer *messagepartbuffer.Buffer + streamSyncPoller *streamSyncPoller + recordingSem chan struct{} + agentCapacityLimiter AgentCapacityLimiter aibridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] experiments codersdk.Experiments @@ -3050,6 +3051,8 @@ type Config struct { PrometheusRegistry prometheus.Registerer + AgentCapacityUnlock AgentCapacityUnlock + // OIDCTokenSource resolves the calling user's OIDC access // token for MCP servers configured with auth_type=user_oidc. // May be nil if the deployment has no OIDC provider; servers @@ -3181,6 +3184,15 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { p.streamPartsDialer = streamPartsDialerForServer(workerID, localStreamPartsDialer, cfg.StreamPartsDialer) p.streamSyncPoller = newStreamSyncPoller(ctx, cfg.Database, clk, cfg.Logger.Named("chatstream")) p.streamSyncPoller.Start() + agentCapacityLimiter := newAgentCapacityLimiter( + cfg.AgentCapacityUnlock, + int32(inFlightChatStaleAfter.Seconds()), + ) + var agentCapacityMetrics *capacityMetrics + if cfg.PrometheusRegistry != nil { + agentCapacityMetrics = newCapacityMetrics(cfg.PrometheusRegistry) + } + p.agentCapacityLimiter = agentCapacityLimiter chatWorker, err := newChatWorker(p, chatWorkerOptions{ WorkerID: workerID, Store: cfg.Database, @@ -3188,6 +3200,8 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { Logger: cfg.Logger.Named("chatworker"), Clock: clk, MessagePartBuffer: p.messagePartBuffer, + AgentCapacityLimiter: agentCapacityLimiter, + CapacityMetrics: agentCapacityMetrics, AcquisitionInterval: pendingChatAcquireInterval, AcquisitionBatchSize: maxChatsPerAcquire, HeartbeatInterval: chatHeartbeatInterval, @@ -3302,9 +3316,6 @@ func chatWatchEventSDKChat(chat database.Chat, diffStatus *codersdk.ChatDiffStat // publishChatPubsubEvent broadcasts a chat lifecycle event via PostgreSQL // pubsub so that all replicas can push updates to watching clients. func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWatchEventKind, diffStatus *codersdk.ChatDiffStatus) { - if p.pubsub == nil { - return - } event := codersdk.ChatWatchEvent{ Kind: kind, Chat: chatWatchEventSDKChat(chat, diffStatus), @@ -3326,6 +3337,27 @@ func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWa } } +// ChatQueuedForCapacity reports whether the chat is waiting for a +// concurrent-agent capacity slot. Uncapped deployments always return false. +func (p *Server) ChatQueuedForCapacity(ctx context.Context, chat database.Chat) (bool, error) { + limits, capped := p.agentCapacityLimiter.Limits() + if !capped { + return false, nil + } + if chat.Archived || chat.Status != database.ChatStatusRunning { + return false, nil + } + // The pool count spans other users' chats, which the requester cannot + // read directly. + //nolint:gocritic // Capacity accounting is chatd-internal state. + return p.db.GetChatQueuedForCapacity(dbauthz.AsChatd(ctx), database.GetChatQueuedForCapacityParams{ + ChatID: chat.ID, + StaleSeconds: int32(p.inFlightChatStaleAfter.Seconds()), + RootCapacity: limits.Root, + SubagentCapacity: limits.Subagent, + }) +} + // PublishDiffStatusChange broadcasts a diff_status_change event for // the given chat so that watching clients know to re-fetch the diff // status. This is called from the HTTP layer after the diff status diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 4d16d0b5f3..ae5b9bcd8f 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -187,7 +187,7 @@ func TestStoreSubagentReportSummary(t *testing.T) { ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) - server := &Server{db: db} + server := &Server{db: db, pubsub: dbpubsub.NewInMemory()} chat := database.Chat{ ID: uuid.New(), OwnerID: uuid.New(), @@ -216,7 +216,7 @@ func TestStoreSubagentReportSummary(t *testing.T) { ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) - server := &Server{db: db} + server := &Server{db: db, pubsub: dbpubsub.NewInMemory()} chat := database.Chat{ ID: uuid.New(), OwnerID: uuid.New(), @@ -237,7 +237,7 @@ func TestStoreSubagentReportSummary(t *testing.T) { ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) - server := &Server{db: db} + server := &Server{db: db, pubsub: dbpubsub.NewInMemory()} chat := database.Chat{ ID: uuid.New(), OwnerID: uuid.New(), @@ -268,7 +268,7 @@ func TestStoreSubagentReportSummary(t *testing.T) { ctrl := gomock.NewController(t) db := dbmock.NewMockStore(ctrl) - server := &Server{db: db} + server := &Server{db: db, pubsub: dbpubsub.NewInMemory()} chat := database.Chat{ ID: uuid.New(), OwnerID: uuid.New(), diff --git a/coderd/x/chatd/helpers_test.go b/coderd/x/chatd/helpers_test.go index 18e4ec3f0e..6747ff6628 100644 --- a/coderd/x/chatd/helpers_test.go +++ b/coderd/x/chatd/helpers_test.go @@ -148,6 +148,25 @@ func (f *workerTestFixture) createRunningChat(t *testing.T) database.Chat { return res.Chat } +func (f *workerTestFixture) createRunningSubagentChat(t *testing.T, parentID uuid.UUID) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + res, err := chatstate.CreateChat(ctx, f.db, f.pubsub, chatstate.CreateChatInput{ + OrganizationID: f.org.ID, + OwnerID: f.user.ID, + LastModelConfigID: f.model.ID, + Title: "subagent", + ClientType: database.ChatClientTypeApi, + ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentID, Valid: true}, + InitialMessages: []chatstate.Message{ + userTextMessage(t, "hello", f.user.ID, f.model.ID, f.apiKey.ID), + }, + }) + require.NoError(t, err) + return res.Chat +} + func (f *workerTestFixture) createRequiresActionChat(t *testing.T) database.Chat { t.Helper() ctx := testutil.Context(t, testutil.WaitShort) diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index f7570a9035..b887b928ba 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -21,12 +21,13 @@ import ( ) const ( - defaultAcquisitionInterval = 30 * time.Second - defaultAcquisitionBatchSize = int32(10) - defaultRunnerSyncInterval = 15 * time.Second - defaultHeartbeatInterval = 9 * time.Second - defaultHeartbeatCleanupEvery = 30 * time.Second - defaultHeartbeatStaleSeconds = int32(30) + defaultAcquisitionInterval = 30 * time.Second + defaultAcquisitionBatchSize = int32(10) + defaultCapacityMetricsInterval = 30 * time.Second + defaultRunnerSyncInterval = 15 * time.Second + defaultHeartbeatInterval = 9 * time.Second + defaultHeartbeatCleanupEvery = 30 * time.Second + defaultHeartbeatStaleSeconds = int32(30) // The archive cutoff is based on UTC start-of-day and only moves // once per day, so hourly runs are more than enough to keep up // while still catching chats that cross the threshold shortly @@ -192,7 +193,11 @@ type chatWorkerOptions struct { Auditor *atomic.Pointer[audit.Auditor] AutoArchiveRecords prometheus.Counter + AgentCapacityLimiter AgentCapacityLimiter + CapacityMetrics *capacityMetrics + AcquisitionInterval time.Duration + CapacityMetricsInterval time.Duration AcquisitionBatchSize int32 ArchiveInterval time.Duration ArchiveBatchSize int32 @@ -226,6 +231,9 @@ func (o chatWorkerOptions) withDefaults() (chatWorkerOptions, error) { if o.AcquisitionInterval <= 0 { o.AcquisitionInterval = defaultAcquisitionInterval } + if o.CapacityMetricsInterval <= 0 { + o.CapacityMetricsInterval = defaultCapacityMetricsInterval + } if o.AcquisitionBatchSize <= 0 { o.AcquisitionBatchSize = defaultAcquisitionBatchSize } @@ -250,6 +258,9 @@ func (o chatWorkerOptions) withDefaults() (chatWorkerOptions, error) { if o.HeartbeatStaleSeconds <= 0 { o.HeartbeatStaleSeconds = defaultHeartbeatStaleSeconds } + if o.AgentCapacityLimiter == nil { + o.AgentCapacityLimiter = newAgentCapacityLimiter(nil, o.HeartbeatStaleSeconds) + } if o.StateChannelSize <= 0 { o.StateChannelSize = defaultStateChannelSize } diff --git a/coderd/x/chatd/quickgen_internal_test.go b/coderd/x/chatd/quickgen_internal_test.go index a761bebc96..7d0bef5def 100644 --- a/coderd/x/chatd/quickgen_internal_test.go +++ b/coderd/x/chatd/quickgen_internal_test.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/coderd/x/chatd/chattest" @@ -580,7 +581,7 @@ func TestMaybeGenerateChatTitlePreservesUpdatedAt(t *testing.T) { logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) generated := &generatedChatTitle{} - server := &Server{db: db} + server := &Server{db: db, pubsub: dbpubsub.NewInMemory()} server.maybeGenerateChatTitle( ctx, chat, diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 9f1e55c54f..2774d90f9c 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -293,6 +293,7 @@ func TestCreateChildSubagentChatDispatchesUserPromptSubmit(t *testing.T) { t.Cleanup(consumer.Close) server := &Server{ db: db, + pubsub: pubsub.NewInMemory(), logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), hooks: chathooks.NewTrigger(dispatch.New( slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), diff --git a/coderd/x/chatd/title_override_internal_test.go b/coderd/x/chatd/title_override_internal_test.go index 36ebef7226..e109ca7652 100644 --- a/coderd/x/chatd/title_override_internal_test.go +++ b/coderd/x/chatd/title_override_internal_test.go @@ -24,6 +24,7 @@ import ( "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbmock" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" @@ -709,6 +710,7 @@ func titleOverrideTestServer(db database.Store, logger slog.Logger) *Server { })} return &Server{ db: db, + pubsub: dbpubsub.NewInMemory(), logger: logger, configCache: newChatConfigCache(context.Background(), db, quartz.NewReal()), aibridgeTransportFactory: aibridgeTestFactoryPointer(factory), diff --git a/coderd/x/chatd/worker.go b/coderd/x/chatd/worker.go index 5aed8e5713..83123f48cc 100644 --- a/coderd/x/chatd/worker.go +++ b/coderd/x/chatd/worker.go @@ -97,6 +97,11 @@ func (w *chatWorker) Start(ctx context.Context) error { w.wg.Go(func() { w.archiveLoop(workerCtx) }) + if w.opts.CapacityMetrics != nil { + w.wg.Go(func() { + w.capacityMetricsLoop(workerCtx) + }) + } wake(wakeCh) return nil } @@ -187,49 +192,65 @@ func (w *chatWorker) acquisitionLoop( } func (w *chatWorker) acquireOnce(ctx context.Context, workerID uuid.UUID, manager *runnerManager) { - attempted := make(map[uuid.UUID]struct{}) - for { - rows, err := w.opts.Store.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ - StaleSeconds: w.opts.HeartbeatStaleSeconds, - LimitCount: w.opts.AcquisitionBatchSize, - }) + // Fetch twice the budget so one full pool cannot hide candidates in the other. + rows, err := w.opts.Store.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: w.opts.HeartbeatStaleSeconds, + LimitCount: w.opts.AcquisitionBatchSize * 2, + }) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker acquisition query failed", slogError(err)) + } + return + } + + acquired := int32(0) + rootPoolRefused := false + subagentPoolRefused := false + for _, row := range rows { + if acquired >= w.opts.AcquisitionBatchSize { + return + } + // Interrupting and requires-action chats bypass capacity so their runners + // can finish work or enforce the action deadline. + isSubagent := row.ParentChatID.Valid + if row.Status == database.ChatStatusRunning && + ((isSubagent && subagentPoolRefused) || (!isSubagent && rootPoolRefused)) { + continue + } + candidateAcquired, err := w.acquireCandidateSafely(ctx, workerID, manager, row.ID) + if errors.Is(err, errCapacityRefused) { + if isSubagent { + subagentPoolRefused = true + } else { + rootPoolRefused = true + } + continue + } if err != nil { - if ctx.Err() == nil { - w.opts.Logger.Warn(ctx, "chatworker acquisition query failed", slogError(err)) + if ctx.Err() != nil { + return } - return + w.opts.Logger.Warn(ctx, "chatworker acquisition candidate failed", slogError(err)) + continue } - if len(rows) == 0 { - return - } - newRows := 0 - for _, row := range rows { - if _, ok := attempted[row.ID]; ok { - continue - } - attempted[row.ID] = struct{}{} - newRows++ - if err := w.acquireCandidateSafely(ctx, workerID, manager, row.ID); err != nil { - if ctx.Err() != nil { - return - } - w.opts.Logger.Warn(ctx, "chatworker acquisition candidate failed", slogError(err)) - } - } - if len(rows) < int(w.opts.AcquisitionBatchSize) || newRows == 0 { - return + if candidateAcquired { + acquired++ } } } -var errSkipAcquire = xerrors.New("skip acquire") +var ( + errSkipAcquire = xerrors.New("skip acquire") + errCapacityRefused = xerrors.New("capacity refused") +) func (w *chatWorker) acquireCandidateSafely( ctx context.Context, workerID uuid.UUID, manager *runnerManager, chatID uuid.UUID, -) (err error) { +) (acquired bool, err error) { defer func() { if recovered := recover(); recovered != nil { err = xerrors.Errorf("chatworker acquisition panic: %v", recovered) @@ -243,7 +264,7 @@ func (w *chatWorker) acquireCandidate( workerID uuid.UUID, manager *runnerManager, chatID uuid.UUID, -) error { +) (bool, error) { runnerID := uuid.New() machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID) err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { @@ -274,22 +295,34 @@ func (w *chatWorker) acquireCandidate( return errSkipAcquire } } + admitted, err := w.opts.AgentCapacityLimiter.Admit(ctx, store, chat) + if err != nil { + return xerrors.Errorf("agent admission: %w", err) + } + if !admitted { + // Roll back to suppress the ownership hint, which would wake every + // worker into an immediate retry of this unowned chat. + return errCapacityRefused + } _, err = tx.Acquire(chatstate.AcquireInput{WorkerID: workerID, RunnerID: runnerID}) return err }) + if errors.Is(err, errCapacityRefused) { + return false, errCapacityRefused + } if errors.Is(err, errSkipAcquire) || errors.Is(err, chatstate.ErrChatNotFound) { - return nil + return false, nil } if err != nil { - return err + return false, err } if err := manager.Spawn(ctx, spawnRunnerRequest{ChatID: chatID, WorkerID: workerID, RunnerID: runnerID}); err != nil { if errAbandon := w.abandonAcquiredChat(ctx, workerID, runnerID, chatID); errAbandon != nil { - return errors.Join(err, errAbandon) + return false, errors.Join(err, errAbandon) } - return err + return false, err } - return nil + return true, nil } func (w *chatWorker) abandonAcquiredChat(ctx context.Context, workerID uuid.UUID, runnerID uuid.UUID, chatID uuid.UUID) error { diff --git a/coderd/x/chatd/worker_internal_test.go b/coderd/x/chatd/worker_internal_test.go index a7bae0d7dd..ff98c332c1 100644 --- a/coderd/x/chatd/worker_internal_test.go +++ b/coderd/x/chatd/worker_internal_test.go @@ -156,7 +156,7 @@ func TestWorker_TwoWorkersRaceSingleOwner(t *testing.T) { require.Equal(t, call.input.RunnerID, latest.RunnerID.UUID) } -func TestWorker_DrainsMultipleRunnableChatsOnWake(t *testing.T) { +func TestWorker_AcquisitionBatchSizeLimitsSuccessfulAcquisitions(t *testing.T) { t.Parallel() f := newWorkerTestFixture(t) first := f.createRunningChat(t) @@ -165,12 +165,14 @@ func TestWorker_DrainsMultipleRunnableChatsOnWake(t *testing.T) { starter := newRecordingTaskStarter() opts := testOptions(t, f, starter) opts.AcquisitionBatchSize = 1 - startWorker(t, opts) + worker := startWorker(t, opts) want := map[uuid.UUID]bool{first.ID: true, second.ID: true, third.ID: true} for range 3 { call := starter.waitCall(t, taskKindGeneration, uuid.Nil) delete(want, call.input.ChatID) + starter.assertNoCall(t) + worker.Wake() } require.Empty(t, want) } diff --git a/codersdk/chats.go b/codersdk/chats.go index e55986342d..de4f7696d0 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -141,9 +141,12 @@ type Chat struct { // Context reports the chat's pinned workspace-context state and // whether it has drifted from the agent's latest pushed snapshot. // Nil when the chat has no pinned context yet. - Context *ChatContext `json:"context,omitempty"` - Warnings []string `json:"warnings,omitempty"` - ClientType ChatClientType `json:"client_type"` + Context *ChatContext `json:"context,omitempty"` + // QueuedForCapacity reports that the chat is waiting for a concurrent + // agent slot. Single-chat reads derive it; list responses leave it false. + QueuedForCapacity bool `json:"queued_for_capacity,omitempty"` + Warnings []string `json:"warnings,omitempty"` + ClientType ChatClientType `json:"client_type"` // Children holds child (subagent) chats nested under this root // chat. Always initialized to an empty slice so the JSON field // is present as []. Child chats cannot create their own diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index 5d8551d991..fad03a9440 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -236,6 +236,8 @@ The `coder_ai_gateway_cost_control_*` metrics are exported only by `coderd`. | `coderd_authz_prepare_authorize_duration_seconds` | histogram | Duration of the 'PrepareAuthorize' call in seconds. | | | `coderd_build_info` | gauge | Describes the current build/version of the Coder server. Value is always 1. | `revision` `version` | | `coderd_chat_auto_archive_records_archived_total` | counter | Total number of chats archived by the auto-archive job (counting both roots and cascaded children). | | +| `coderd_chatd_agents_active` | gauge | Deployment-wide number of chats holding a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum. | `pool` | +| `coderd_chatd_agents_queued_for_capacity` | gauge | Deployment-wide number of chats waiting for a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum. | `pool` | | `coderd_chatd_chats` | gauge | Number of chats being processed, by state. | `state` | | `coderd_chatd_compaction_total` | counter | Total compaction outcomes (only recorded when compaction was triggered or failed). | `model` `provider` `result` | | `coderd_chatd_hook_context_size_bytes` | histogram | Lifecycle hook model context response size in bytes. | `event` | diff --git a/docs/ai-coder/agents/getting-started.md b/docs/ai-coder/agents/getting-started.md index 4579c72819..080ef8936b 100644 --- a/docs/ai-coder/agents/getting-started.md +++ b/docs/ai-coder/agents/getting-started.md @@ -221,6 +221,15 @@ token volume. Consider: - Capping spend with [AI Gateway budgets](./platform-controls/spend-management.md). - Monitoring provider dashboards for usage trends during the evaluation. +### Plan for concurrency limits + +Community licenses run up to 5 agents at once. +Additional agents queue and start automatically when capacity frees. +Premium licenses with Agent Hours do not impose a concurrency limit unless the Agent Hours hard limit is reached. +If the Agent Hours allocation is exhausted without a configured hard limit, Coder warns about usage but does not impose a concurrency limit. +When the Agent Hours hard limit is reached, additional agents queue under the concurrency limit. +Refer to [Concurrent agents](./platform-controls/index.md#concurrent-agents) for details. + ### Pilot with a small group Identify 3–5 developers and a few concrete use cases for the initial rollout. diff --git a/docs/ai-coder/agents/platform-controls/index.md b/docs/ai-coder/agents/platform-controls/index.md index 432050909e..0a9237d0ab 100644 --- a/docs/ai-coder/agents/platform-controls/index.md +++ b/docs/ai-coder/agents/platform-controls/index.md @@ -112,6 +112,22 @@ This setting is available under **Agents** > **Settings** > days. When disabled, workspaces follow their template's autostop rules (or none, if the template does not define any). +### Concurrent agents + +Community licenses support up to 5 concurrently active agents. +Coder doesn't limit how long those agents can run or how many tasks they complete over time. +Additional agents queue until an agent session becomes available. +With concurrent agents, individuals and small teams can experiment with Coder Agents at no cost. + +Queued agents show a banner in the chat and start automatically when capacity frees. +Subtasks delegated by an agent don't count toward this limit. +Those subtasks run in a separate pool of up to 10 concurrent subtasks. + +Premium deployments can purchase Agent Hours with their Premium license. +Agent Hours are shared across the deployment, and agents can run concurrently unless the Agent Hours hard limit is reached. +If the Agent Hours allocation is exhausted without a configured hard limit, Coder warns about usage but does not impose a concurrency limit. +When the Agent Hours hard limit is reached, additional agents queue under the concurrency limit. + ### Spend management AI Gateway budgets cap each user's AI spend, including Coder Agents chats, over a monthly period. diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 7a51212c7f..82567f1166 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -127,6 +127,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -224,6 +225,7 @@ Status Code **200** | `» parent_chat_id` | string(uuid) | false | | | | `» pin_order` | integer | false | | | | `» plan_mode` | [codersdk.ChatPlanMode](schemas.md#codersdkchatplanmode) | false | | | +| `» queued_for_capacity` | boolean | false | | Queued for capacity reports that the chat is waiting for a concurrent agent slot. Single-chat reads derive it; list responses leave it false. | | `» root_chat_id` | string(uuid) | false | | | | `» shared` | boolean | false | | Shared is true when this chat's root chat has explicit user or group ACL entries. | | `» status` | [codersdk.ChatStatus](schemas.md#codersdkchatstatus) | false | | | @@ -404,6 +406,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -498,6 +501,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -749,6 +753,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -897,6 +902,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -991,6 +997,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -1176,6 +1183,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -1270,6 +1278,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -1505,6 +1514,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -1599,6 +1609,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -2520,6 +2531,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -2614,6 +2626,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -3122,6 +3135,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -3216,6 +3230,7 @@ Experimental: this endpoint is subject to change. "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index b7c446a934..b00a428094 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -2280,6 +2280,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -2374,6 +2375,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", @@ -2416,6 +2418,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in | `parent_chat_id` | string | false | | | | `pin_order` | integer | false | | | | `plan_mode` | [codersdk.ChatPlanMode](#codersdkchatplanmode) | false | | | +| `queued_for_capacity` | boolean | false | | Queued for capacity reports that the chat is waiting for a concurrent agent slot. Single-chat reads derive it; list responses leave it false. | | `root_chat_id` | string | false | | | | `shared` | boolean | false | | Shared is true when this chat's root chat has explicit user or group ACL entries. | | `status` | [codersdk.ChatStatus](#codersdkchatstatus) | false | | | @@ -4203,6 +4206,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", "pin_order": 0, "plan_mode": "plan", + "queued_for_capacity": true, "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", "shared": true, "status": "waiting", diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index bb58222b49..adbdec8168 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -222,6 +222,8 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { }, }) + options.Options.ChatAgentCapacityUnlock = entchatd.NewAgentCapacityUnlock(options.Entitlements) + api.AGPL = coderd.New(options.Options) api.aiSeatTracker = aiseats.New(options.Database, api.Logger.Named("aiseats"), quartz.NewReal(), &api.AGPL.Auditor) api.AGPL.AISeatTracker = api.aiSeatTracker diff --git a/enterprise/coderd/x/chatd/agentadmission.go b/enterprise/coderd/x/chatd/agentadmission.go new file mode 100644 index 0000000000..c12b0456d6 --- /dev/null +++ b/enterprise/coderd/x/chatd/agentadmission.go @@ -0,0 +1,29 @@ +package chatd + +import ( + "github.com/coder/coder/v2/coderd/entitlements" + osschatd "github.com/coder/coder/v2/coderd/x/chatd" + "github.com/coder/coder/v2/codersdk" +) + +// NewAgentCapacityUnlock returns an unlock that tracks entitlement changes. +func NewAgentCapacityUnlock(set *entitlements.Set) osschatd.AgentCapacityUnlock { + return &agentCapacityUnlock{entitlements: set} +} + +type agentCapacityUnlock struct { + entitlements *entitlements.Set +} + +// The Agent Hours allocation is advisory; the hard limit restores concurrency caps. +func (u *agentCapacityUnlock) Unlocked() bool { + f, ok := u.entitlements.Feature(codersdk.FeatureAgentRuntimeHours) + if !ok || !f.Enabled { + return false + } + if f.HardLimit == nil || f.Actual == nil { + // Missing hard limits or usage measurements leave concurrency uncapped. + return true + } + return *f.Actual < *f.HardLimit +} diff --git a/enterprise/coderd/x/chatd/agentadmission_internal_test.go b/enterprise/coderd/x/chatd/agentadmission_internal_test.go new file mode 100644 index 0000000000..44a24addae --- /dev/null +++ b/enterprise/coderd/x/chatd/agentadmission_internal_test.go @@ -0,0 +1,147 @@ +package chatd + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/entitlements" + "github.com/coder/coder/v2/coderd/util/ptr" + "github.com/coder/coder/v2/codersdk" +) + +func TestAgentCapacityUnlock(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + feature *codersdk.Feature + unlocked bool + }{ + { + name: "EnabledWithoutUsage", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + }, + unlocked: true, + }, + { + name: "RemainingHours", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + Actual: ptr.Ref(int64(40)), + }, + unlocked: true, + }, + { + name: "AtAllocationWithoutHardLimit", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + Actual: ptr.Ref(int64(100)), + }, + unlocked: true, + }, + { + name: "OverAllocationWithoutHardLimit", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + Actual: ptr.Ref(int64(150)), + }, + unlocked: true, + }, + { + name: "BelowHardLimit", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + HardLimit: ptr.Ref(int64(120)), + Actual: ptr.Ref(int64(119)), + }, + unlocked: true, + }, + { + name: "AtHardLimit", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + HardLimit: ptr.Ref(int64(120)), + Actual: ptr.Ref(int64(120)), + }, + }, + { + name: "AboveHardLimit", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + HardLimit: ptr.Ref(int64(120)), + Actual: ptr.Ref(int64(121)), + }, + }, + { + name: "HardLimitEqualsAllocation", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + Limit: ptr.Ref(int64(100)), + HardLimit: ptr.Ref(int64(100)), + Actual: ptr.Ref(int64(100)), + }, + }, + { + name: "DisabledZeroAllocation", + feature: &codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: false, + Limit: ptr.Ref(int64(0)), + }, + }, + {name: "NoFeature"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + set := entitlements.New() + if tc.feature != nil { + set.Modify(func(ents *codersdk.Entitlements) { + ents.Features[codersdk.FeatureAgentRuntimeHours] = *tc.feature + }) + } + require.Equal(t, tc.unlocked, NewAgentCapacityUnlock(set).Unlocked()) + }) + } +} + +func TestAgentCapacityUnlockTracksEntitlementUpdates(t *testing.T) { + t.Parallel() + + set := entitlements.New() + unlock := NewAgentCapacityUnlock(set) + require.False(t, unlock.Unlocked()) + + set.Modify(func(ents *codersdk.Entitlements) { + ents.Features[codersdk.FeatureAgentRuntimeHours] = codersdk.Feature{ + Entitlement: codersdk.EntitlementEntitled, + Enabled: true, + HardLimit: ptr.Ref(int64(120)), + Actual: ptr.Ref(int64(119)), + } + }) + require.True(t, unlock.Unlocked()) + + set.Modify(func(ents *codersdk.Entitlements) { + feature := ents.Features[codersdk.FeatureAgentRuntimeHours] + feature.Actual = ptr.Ref(int64(120)) + ents.Features[codersdk.FeatureAgentRuntimeHours] = feature + }) + require.False(t, unlock.Unlocked()) +} diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index 7b5cd266e6..a0ad8d2a58 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -271,6 +271,12 @@ coderd_build_info{version="",revision=""} 0 # HELP coderd_chat_auto_archive_records_archived_total Total number of chats archived by the auto-archive job (counting both roots and cascaded children). # TYPE coderd_chat_auto_archive_records_archived_total counter coderd_chat_auto_archive_records_archived_total 0 +# HELP coderd_chatd_agents_active Deployment-wide number of chats holding a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum. +# TYPE coderd_chatd_agents_active gauge +coderd_chatd_agents_active{pool=""} 0 +# HELP coderd_chatd_agents_queued_for_capacity Deployment-wide number of chats waiting for a concurrent-agent capacity slot. Every replica reports the same database-derived value; aggregate with max, not sum. +# TYPE coderd_chatd_agents_queued_for_capacity gauge +coderd_chatd_agents_queued_for_capacity{pool=""} 0 # HELP coderd_chatd_chats Number of chats being processed, by state. # TYPE coderd_chatd_chats gauge coderd_chatd_chats{state=""} 0 diff --git a/site/src/@types/storybook.d.ts b/site/src/@types/storybook.d.ts index 76166ba53c..3c34fdc352 100644 --- a/site/src/@types/storybook.d.ts +++ b/site/src/@types/storybook.d.ts @@ -1,6 +1,7 @@ import type { DeploymentValues, Experiments, + Feature, FeatureName, Organization, SerpentOption, @@ -15,7 +16,7 @@ declare module "@storybook/react-vite" { | { event: "message"; data: string } | { event: "open" | "error" | "close" }; interface Parameters { - features?: FeatureName[]; + features?: (FeatureName | ({ name: FeatureName } & Partial))[]; experiments?: Experiments; showOrganizations?: boolean; organizations?: Organization[]; diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 05d0397d5d..1715d440d2 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -44,6 +44,7 @@ import { deleteChatQueuedMessage, editChatMessage, getChatListQueryString, + getOpenChatPollInterval, infiniteChats, interruptChat, invalidateChatACL, @@ -58,6 +59,7 @@ import { invalidateChatsByWorkspace, mergeWatchedChatIntoCaches, mergeWatchedChatSummary, + openChat, patchChatEntity, patchChatMessages, pinChat, @@ -2466,6 +2468,43 @@ describe("mergeWatchedChatSummary", () => { }); }); + it("preserves queued_for_capacity while the chat remains running", () => { + const cachedChat = makeChat("chat-1", { + status: "running", + updated_at: "2025-01-01T00:00:00.000Z", + queued_for_capacity: true, + }); + const watchedChat = makeChat("chat-1", { + status: "running", + updated_at: "2025-01-01T00:01:00.000Z", + queued_for_capacity: false, + }); + + expect( + mergeWatchedChatSummary(cachedChat, watchedChat, { + eventKind: "status_change", + }).queued_for_capacity, + ).toBe(true); + }); + + it("clears queued_for_capacity when the chat stops running", () => { + const cachedChat = makeChat("chat-1", { + status: "running", + updated_at: "2025-01-01T00:00:00.000Z", + queued_for_capacity: true, + }); + const watchedChat = makeChat("chat-1", { + status: "waiting", + updated_at: "2025-01-01T00:01:00.000Z", + }); + + expect( + mergeWatchedChatSummary(cachedChat, watchedChat, { + eventKind: "status_change", + }).queued_for_capacity, + ).toBe(false); + }); + it("leaves context untouched for non-context events", () => { const context = { dirty: true, dirty_since: "2025-01-02T00:00:00.000Z" }; const cachedChat = makeChat("chat-1", { @@ -3322,8 +3361,6 @@ describe("semantic cache operations: prefix invalidations", () => { }); describe(shouldInvalidateChatsByWorkspace.name, () => { - // created/deleted have their own watch branches; title, summary, - // diff, and context events do not move updated_at ordering. const expectedByKind: Record = { action_required: true, chat_summary_change: false, @@ -3403,11 +3440,6 @@ describe("semantic cache operations: prefix invalidations", () => { }); describe(shouldInvalidateChatSearches.name, () => { - // Search results render title, status, diff status, and the - // action-required badge. Summary and context events are excluded: - // stale last_turn_summary subtitles are accepted until - // reconciliation lands. The created and deleted kinds are handled - // by their own watch branches before the merge path runs. const expectedByKind: Record = { action_required: true, chat_summary_change: false, @@ -3426,6 +3458,36 @@ describe("semantic cache operations: prefix invalidations", () => { }); }); +describe("openChat", () => { + it("does not poll in the background", () => { + expect(openChat("chat-1").refetchIntervalInBackground).toBe(false); + }); + + it("polls while the open chat is running", () => { + expect( + getOpenChatPollInterval(makeChat("chat-1", { status: "running" })), + ).toBe(5_000); + }); + + it("stops polling after the chat leaves running", () => { + expect( + getOpenChatPollInterval(makeChat("chat-1", { status: "waiting" })), + ).toBe(false); + }); + + it("does not poll archived chats", () => { + expect( + getOpenChatPollInterval( + makeChat("chat-1", { status: "running", archived: true }), + ), + ).toBe(false); + }); + + it("does not poll before the chat loads", () => { + expect(getOpenChatPollInterval(undefined)).toBe(false); + }); +}); + describe("semantic cache operations: cancellation", () => { it("cancelChatListQueries cancels unconditionally across the list family", async () => { const queryClient = createTestQueryClient(); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index e323151392..198573ab74 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -548,6 +548,10 @@ export const mergeWatchedChatSummary = ( isContextDirtyEvent && watchedChat.context ? { ...cachedChat.context, ...watchedChat.context } : cachedChat.context; + const nextQueuedForCapacity = + isStatusEvent && nextStatus !== "running" + ? false + : (cachedChat.queued_for_capacity ?? false); const nextWorkspaceId = isFreshEnough ? (watchedChat.workspace_id ?? cachedChat.workspace_id) : cachedChat.workspace_id; @@ -594,7 +598,8 @@ export const mergeWatchedChatSummary = ( nextSummary === cachedChat.summary && nextHasUnread === cachedChat.has_unread && nextUpdatedAt === cachedChat.updated_at && - nextContext === cachedChat.context + nextContext === cachedChat.context && + nextQueuedForCapacity === (cachedChat.queued_for_capacity ?? false) ) { return cachedChat; } @@ -612,6 +617,7 @@ export const mergeWatchedChatSummary = ( has_unread: nextHasUnread, updated_at: nextUpdatedAt, context: nextContext, + queued_for_capacity: nextQueuedForCapacity, }; }; @@ -1149,6 +1155,18 @@ export const chat = (chatId: string) => ({ queryFn: () => API.experimental.getChat(chatId), }); +export const getOpenChatPollInterval = ( + data: TypesGen.Chat | undefined, +): number | false => + data?.status === "running" && !data.archived ? 5_000 : false; + +export const openChat = (chatId: string) => + queryOptions({ + ...chat(chatId), + refetchInterval: ({ state }) => getOpenChatPollInterval(state.data), + refetchIntervalInBackground: false, + }); + export const chatACLKey = (chatId: string) => [...chatEntityKey(chatId), "acl"] as const; diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 6059173172..da8a217888 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1979,6 +1979,11 @@ export interface Chat { * Nil when the chat has no pinned context yet. */ readonly context?: ChatContext; + /** + * QueuedForCapacity reports that the chat is waiting for a concurrent + * agent slot. Single-chat reads derive it; list responses leave it false. + */ + readonly queued_for_capacity?: boolean; readonly warnings?: readonly string[]; readonly client_type: ChatClientType; /** diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index 41168f472b..983ab28cda 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -1359,6 +1359,49 @@ export const Loading: Story = { }, }; +const capacityPollingChat: TypesGen.Chat = { + id: CHAT_ID, + ...baseChatFields, + title: "Capacity polling", + status: "running", + queued_for_capacity: false, +}; + +export const QueuedForCapacityAfterPolling: Story = { + parameters: { + queries: [ + { key: chatEntityKey(CHAT_ID), data: capacityPollingChat }, + { + key: chatMessagesKey(CHAT_ID), + data: { + pages: [{ messages: [], queued_messages: [], has_more: false }], + pageParams: [undefined], + }, + }, + ], + }, + beforeEach: () => { + spyOn(API.experimental, "getChat").mockResolvedValue({ + ...capacityPollingChat, + queued_for_capacity: true, + }); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.queryByText(/This agent is queued and will start automatically/), + ).not.toBeInTheDocument(); + + const callout = await canvas.findByRole("alert", undefined, { + timeout: 7_000, + }); + expect(API.experimental.getChat).toHaveBeenCalledWith(CHAT_ID); + expect(callout).toHaveTextContent( + "This agent is queued and will start automatically when capacity is available.", + ); + }, +}; + export const OtherUserChatReadOnly: Story = { parameters: { queries: buildQueries( diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index efe81b3a18..eb9bfe1656 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -40,9 +40,11 @@ import { createChatMessage, deleteChatQueuedMessage, editChatMessage, + getOpenChatPollInterval, interruptChat, invalidateChatEntity, mcpServerConfigs, + openChat, patchChatEntity, promoteChatQueuedMessage, updateChatPlanMode, @@ -934,12 +936,18 @@ const AgentChatPage: FC = () => { }; const chatQuery = useQuery({ - ...chat(agentId ?? ""), + ...openChat(agentId ?? ""), enabled: Boolean(agentId), - // Poll while the binding is unresolved: repair happens on chat reads - // and watch events cannot be relied on for retries because an idle - // workspace publishes none. + // Poll while the chat runs (this override replaces openChat's + // interval, and queued_for_capacity depends on the poll) or while + // the binding is unresolved: repair happens on chat reads and watch + // events cannot be relied on for retries because an idle workspace + // publishes none. refetchInterval: ({ state }) => { + const openPollMs = getOpenChatPollInterval(state.data); + if (openPollMs !== false) { + return openPollMs; + } const workspaceId = state.data?.workspace_id; const workspace = workspaceId ? queryClient.getQueryData( @@ -2098,6 +2106,7 @@ const AgentChatPage: FC = () => { onMCPSelectionChange={handleMCPSelectionChange} onMCPAuthComplete={handleMCPAuthComplete} chatContext={chatQuery.data?.context} + queuedForCapacity={chatQuery.data?.queued_for_capacity ?? false} workspaceSkills={workspaceSkillsFromChat(chatQuery.data)} /> ); diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index 748c7cefc6..932ab83f12 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -33,6 +33,7 @@ import { withProxyProvider, withWebSocket, } from "#/testHelpers/storybook"; +import { docs } from "#/utils/docs"; import { AgentChatPageLoadingView, AgentChatPageNotFoundView, @@ -329,6 +330,135 @@ export const ArchivedOtherUserChat: Story = { }, }; +export const QueuedForCapacityCommunityAdmin: Story = { + parameters: { + permissions: { viewAllLicenses: true }, + }, + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const callout = within(canvas.getByRole("alert")); + const message = callout.getByText( + /reached the Community license limit for active agents/, + ); + expect(message).toBeVisible(); + expect(message).toHaveTextContent( + "This agent is queued and will start automatically when capacity is available.", + ); + const trialLink = canvas.getByRole("link", { + name: /start an unlimited trial/i, + }); + expect(trialLink).toHaveAttribute("href", "https://coder.com/trial"); + const learnMoreLink = canvas.getByRole("link", { name: /learn more/i }); + expect(learnMoreLink).toHaveAttribute( + "href", + docs("/ai-coder/agents/platform-controls#concurrent-agents"), + ); + }, +}; + +export const QueuedForCapacityCommunityMember: Story = { + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const message = canvas.getByText( + /reached the Community license limit for active agents/, + ); + expect(message).toBeVisible(); + expect( + canvas.queryByRole("link", { name: /start an unlimited trial/i }), + ).not.toBeInTheDocument(); + const learnMoreLink = canvas.getByRole("link", { name: /learn more/i }); + expect(learnMoreLink).toHaveAttribute( + "href", + docs("/ai-coder/agents/platform-controls#concurrent-agents"), + ); + }, +}; + +export const QueuedForCapacityPremiumAdmin: Story = { + parameters: { + features: ["multiple_organizations"], + permissions: { viewAllLicenses: true }, + }, + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const message = canvas.getByText( + /reached your license’s limit for active agents/, + ); + expect(message).toBeVisible(); + expect(message).toHaveTextContent( + "Contact your Coder account team or sales@coder.com to upgrade to unlimited concurrent agents.", + ); + const salesLink = canvas.getByRole("link", { name: /sales@coder\.com/ }); + expect(salesLink).toHaveAttribute("href", "mailto:sales@coder.com"); + expect( + canvas.queryByRole("link", { name: /learn more/i }), + ).not.toBeInTheDocument(); + }, +}; + +export const QueuedForCapacityPremiumMember: Story = { + parameters: { + features: ["multiple_organizations"], + }, + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const message = canvas.getByText( + /reached your license’s limit for active agents/, + ); + expect(message).toBeVisible(); + expect( + canvas.queryByRole("link", { name: /sales@coder\.com/ }), + ).not.toBeInTheDocument(); + const learnMoreLink = canvas.getByRole("link", { name: /learn more/i }); + expect(learnMoreLink).toHaveAttribute( + "href", + docs("/ai-coder/agents/platform-controls#concurrent-agents"), + ); + }, +}; + +export const QueuedForCapacityPremiumHardLimit: Story = { + parameters: { + features: [ + "multiple_organizations", + { + name: "agent_runtime_hours", + limit: 3000, + hard_limit: 4000, + actual: 4000, + }, + ], + permissions: { viewAllLicenses: true }, + }, + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const message = canvas.getByText( + /reached the 4000-hour Agent Hours hard limit/, + ); + expect(message).toBeVisible(); + expect(message).toHaveTextContent( + "This agent is queued and will start automatically when capacity is available.", + ); + const salesLink = canvas.getByRole("link", { name: /sales@coder\.com/ }); + expect(salesLink).toHaveAttribute("href", "mailto:sales@coder.com"); + }, +}; + +export const NotQueuedForCapacity: Story = { + render: () => , + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect( + canvas.queryByText(/limit for active agents/), + ).not.toBeInTheDocument(); + }, +}; + /** Shows the parent chat link in the top bar when a parent exists. */ export const WithParentChat: Story = { render: () => ( diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index 8c24c0631d..10ab490e35 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -18,12 +18,14 @@ import type { ChatMessagePart, } from "#/api/typesGenerated"; import { useProxy } from "#/contexts/ProxyContext"; +import { useAuthenticated } from "#/hooks/useAuthenticated"; import { getAgentBrowserApp, isWorkspaceAppEmbeddable, } from "#/modules/apps/apps"; import { WorkspaceAppFrame } from "#/modules/apps/WorkspaceAppFrame"; import { findWorkspaceAppWithAgent } from "#/modules/apps/workspaceApps"; +import { useDashboard } from "#/modules/dashboard/useDashboard"; import { cn } from "#/utils/cn"; import { pageTitle } from "#/utils/page"; import { findWorkspaceAgent } from "#/utils/workspace"; @@ -37,6 +39,7 @@ import { } from "./components/AgentsSkeletons"; import type { ChatDetailError } from "./components/ChatConversation/chatError"; import type { useChatStore } from "./components/ChatConversation/chatStore"; +import { QueuedForCapacityCallout } from "./components/ChatConversation/QueuedForCapacityCallout"; import type { ModelSelectorOption } from "./components/ChatElements"; import { DesktopPanelContext } from "./components/ChatElements/tools/DesktopPanelContext"; import type { SkillMetadata } from "./components/ChatMessageInput/SkillsTriggerMenu"; @@ -122,6 +125,7 @@ interface AgentChatPageViewProps { isArchived: boolean; isSharedChat: boolean; chatOwner: ChatOwnerInfo | undefined; + queuedForCapacity?: boolean; canShareChat: boolean; workspaceAgent?: TypesGen.WorkspaceAgent; workspace?: TypesGen.Workspace; @@ -324,6 +328,7 @@ export const AgentChatPageView: FC = ({ isArchived, isSharedChat, chatOwner, + queuedForCapacity, canShareChat, workspaceAgent, workspace, @@ -396,6 +401,8 @@ export const AgentChatPageView: FC = ({ }) => { const queryClient = useQueryClient(); const { proxy } = useProxy(); + const { entitlements } = useDashboard(); + const { permissions } = useAuthenticated(); const wildcardHostname = proxy.preferredWildcardHostname; const canOpenChatSharing = canShareChat && organizationId !== undefined; @@ -830,6 +837,17 @@ export const AgentChatPageView: FC = ({ ? `This chat is owned by ${chatOwnerLabel}. It is read-only.` : undefined; + const hasLicense = entitlements.has_license; + const canManageLicenses = permissions.viewAllLicenses; + const runtimeHours = entitlements.features.agent_runtime_hours; + const agentHoursHardLimit = + runtimeHours.enabled && + runtimeHours.hard_limit !== undefined && + runtimeHours.actual !== undefined && + runtimeHours.actual >= runtimeHours.hard_limit + ? runtimeHours.hard_limit + : undefined; + const titleElement = ( {chatTitle ? pageTitle(chatTitle, "Agents") : pageTitle("Agents")} @@ -943,6 +961,15 @@ export const AgentChatPageView: FC<AgentChatPageViewProps> = ({ onSendAskUserQuestionResponse={ isOtherUserReadOnly ? undefined : canSendAskUserQuestionResponse } + footer={ + queuedForCapacity ? ( + <QueuedForCapacityCallout + hasLicense={hasLicense} + canManageLicenses={canManageLicenses} + agentHoursHardLimit={agentHoursHardLimit} + /> + ) : undefined + } /> <div className="shrink-0 overflow-y-auto px-4 pb-3 md:pb-0 [scrollbar-gutter:stable] [scrollbar-width:thin]"> <ChatPageInput diff --git a/site/src/pages/AgentsPage/components/ChatConversation/QueuedForCapacityCallout.tsx b/site/src/pages/AgentsPage/components/ChatConversation/QueuedForCapacityCallout.tsx new file mode 100644 index 0000000000..db9adb9adf --- /dev/null +++ b/site/src/pages/AgentsPage/components/ChatConversation/QueuedForCapacityCallout.tsx @@ -0,0 +1,87 @@ +import type { FC, ReactNode } from "react"; +import { Alert, AlertDescription } from "#/components/Alert/Alert"; +import { Link } from "#/components/Link/Link"; +import { docs } from "#/utils/docs"; + +const concurrencyDocsUrl = docs( + "/ai-coder/agents/platform-controls#concurrent-agents", +); + +interface QueuedForCapacityCalloutProps { + hasLicense: boolean; + canManageLicenses: boolean; + agentHoursHardLimit?: number; +} + +export const QueuedForCapacityCallout: FC<QueuedForCapacityCalloutProps> = ({ + hasLicense, + canManageLicenses, + agentHoursHardLimit, +}) => { + let limitMessage = + "Your team has reached the Community license limit for active agents."; + if (hasLicense) { + limitMessage = + "Your team has reached your license’s limit for active agents."; + } + if (agentHoursHardLimit !== undefined) { + limitMessage = `Your team has reached the ${agentHoursHardLimit}-hour Agent Hours hard limit.`; + } + + let action: ReactNode = ( + <> + <Link + href={concurrencyDocsUrl} + target="_blank" + rel="noreferrer" + size="sm" + > + Learn more + </Link> + . + </> + ); + if (canManageLicenses && hasLicense) { + action = ( + <> + Contact your Coder account team or{" "} + <Link href="mailto:sales@coder.com" size="sm" showExternalIcon={false}> + sales@coder.com + </Link>{" "} + to upgrade to unlimited concurrent agents. + </> + ); + } else if (canManageLicenses) { + action = ( + <> + <Link + href="https://coder.com/trial" + target="_blank" + rel="noreferrer" + size="sm" + > + Start an unlimited trial + </Link>{" "} + or{" "} + <Link + href={concurrencyDocsUrl} + target="_blank" + rel="noreferrer" + size="sm" + > + learn more + </Link> + . + </> + ); + } + + return ( + <Alert severity="warning" className="mt-2"> + <AlertDescription> + {limitMessage} This agent is queued and will start automatically when + capacity is available. {action} + </AlertDescription> + </Alert> + ); +}; diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 8038cec91b..f04515f3ce 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -105,6 +105,7 @@ interface ChatPageTimelineProps { onSendAskUserQuestionResponse?: (message: string) => Promise<void> | void; urlTransform?: UrlTransform; mcpServers?: readonly TypesGen.MCPServerConfig[]; + footer?: ReactNode; } export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({ @@ -122,6 +123,7 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({ onSendAskUserQuestionResponse, urlTransform, mcpServers, + footer, }) => { const [chatFullWidth] = useChatFullWidth(); const messagesByID = useChatSelector(store, selectMessagesByID); @@ -221,6 +223,7 @@ export const ChatPageTimeline: FC<ChatPageTimelineProps> = ({ isTranscriptEmpty={parsedMessages.length === 0} liveStatus={liveStatus} /> + {footer} </div> </Profiler> ); diff --git a/site/src/testHelpers/storybook.tsx b/site/src/testHelpers/storybook.tsx index 6a885776d1..fbb80d54a6 100644 --- a/site/src/testHelpers/storybook.tsx +++ b/site/src/testHelpers/storybook.tsx @@ -43,10 +43,13 @@ export const withDashboardProvider = ( has_license: features.length > 0, features: withDefaultFeatures( Object.fromEntries( - features.map((feature) => [ - feature, - { enabled: true, entitlement: "entitled" }, - ]), + features.map((feature) => { + if (typeof feature === "string") { + return [feature, { enabled: true, entitlement: "entitled" }]; + } + const { name, ...values } = feature; + return [name, { enabled: true, entitlement: "entitled", ...values }]; + }), ), ), };