mirror of
https://github.com/coder/coder.git
synced 2026-11-01 04:47:17 +08:00
Michael Suchacz
119f2b1dd9
Limits concurrent chat generation on capped deployments to 5 root chats and 10 delegated subagent chats. The pools are deployment-wide and independent, so delegated work can continue while root capacity is full. The default caps live in AGPL code. Enterprise contributes only a licensing unlock, so unlicensed deployments stay capped and cannot fail open. Licensed deployments are uncapped while Agent Hours usage stays below an explicit hard limit. Deployments without a hard limit remain uncapped, and reaching the Agent Hours allocation only triggers warnings. Admission happens before a worker takes chat ownership. Capped deployments serialize admission across replicas with a transaction-scoped advisory lock and derive active and queued state from current ownership plus fresh runner heartbeats, rather than persisted queue markers or per-replica state. The acquisition query returns a bounded, pool-interleaved candidate set instead of ranking the whole backlog; a migration replaces the acquisition index with a pool-aware one. Refused chats stay running but unowned, and interrupt requests bypass admission so users can stop queued or over-cap chats. The single-chat API derives `queued_for_capacity` from live pool state; list endpoints do not report it. The UI polls that value every 5 seconds while a chat is running and shows a callout when the chat is waiting for capacity. Updates the administrator documentation and deployment-wide Prometheus gauges for active and queued agents. Replica-level values must be aggregated with `max`, not `sum`. > Mux updated this PR on Mike's behalf.
86 lines
2.5 KiB
Go
86 lines
2.5 KiB
Go
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()
|
|
}
|