diff --git a/server/engine.go b/server/engine.go index b5643b2..9a8c07b 100644 --- a/server/engine.go +++ b/server/engine.go @@ -148,11 +148,18 @@ type Engine struct { // (for example an exhausted failover chain) — there is no global-pair fallback. resolve func(t *Task) (*agent.Planner, *agent.Worker) resolveAuthoritative bool + workerResolve func(t *Task, slot WorkerSlot) (*agent.Worker, error) + slotState sync.Map // task/slot -> slotAvailability // readiness reports whether a global LLM provider is configured — the signal behind // Ready()/the llm_configured indicator. Wired once at startup; nil → not ready. readiness func() bool } +type slotAvailability struct { + available bool + lastEmit time.Time +} + type taskRuntime struct { ctx context.Context cancel context.CancelFunc @@ -285,6 +292,13 @@ func (e *Engine) StopTask(taskID string) { e.stamped.Delete(taskID) e.inflight.Delete(taskID) e.coordStarted.Delete(taskID) + e.slotState.Range(func(key, _ any) bool { + keyText, ok := key.(string) + if ok && strings.HasPrefix(keyText, taskID+"/") { + e.slotState.Delete(key) + } + return true + }) e.deleteMu.Lock() e.deleting.Delete(taskID) e.deleteMu.Unlock() @@ -715,6 +729,37 @@ func (e *Engine) SetAuthoritativeAgentResolver(fn func(t *Task) (*agent.Planner, e.resolveAuthoritative = true } +func (e *Engine) SetWorkerResolver(fn func(t *Task, slot WorkerSlot) (*agent.Worker, error)) { + e.workerResolve = fn +} + +func (e *Engine) reportSlotAvailability(t *Task, slot WorkerSlot, err error) { + key := t.ID + "/" + slot.Name + now := time.Now() + value, loaded := e.slotState.LoadOrStore(key, slotAvailability{available: err == nil, lastEmit: now}) + state := value.(slotAvailability) + if err == nil { + if loaded && !state.available { + e.slotState.Store(key, slotAvailability{available: true, lastEmit: now}) + e.emitActivity(t, db.Activity{Worker: slot.Name, Kind: "text", Summary: slot.Name + " 绑定的 LLM 配置已恢复，可以继续领取任务"}) + } + return + } + if loaded && state.available { + state.available = false + state.lastEmit = now + e.slotState.Store(key, state) + } else if loaded && now.Sub(state.lastEmit) < 5*time.Minute { + return + } else { + state = slotAvailability{available: false, lastEmit: now} + e.slotState.Store(key, state) + } + summary := fmt.Sprintf("%s 绑定的 LLM 配置 #%d 不可用（已删除或配置无效），该会话暂停领取任务，不会回退到默认模型", slot.Name, slot.ProfileID) + e.emitActivity(t, db.Activity{Worker: slot.Name, Kind: "text", IsError: true, Summary: summary, Detail: err.Error()}) + log.Printf("[worker %s] task %s pinned profile #%d unavailable: %v", slot.Name, t.ID, slot.ProfileID, err) +} + // snapshotFor returns the planner/worker a task should run on, from the task-router // resolver. nil,nil means the task is deliberately unavailable (e.g. an exhausted // failover chain); there is no global-pair fallback. @@ -746,7 +791,11 @@ func (e *Engine) ReadyFor(t *Task) bool { // but no-op until an LLM is configured (so a task created while idle picks up // automatically once LLM is set from the UI). func (e *Engine) Run(ctx context.Context, t *Task) { - workers := e.m.Workers() + slots, err := ExpandWorkerSlots(e.m.WorkerPools(), e.m.Workers()) + if err != nil { + log.Printf("[engine] task %s worker pool invalid, using legacy workers: %v", t.ID, err) + slots, _ = ExpandWorkerSlots(nil, e.m.Workers()) + } e.deleteMu.RLock() if e.IsDeleting(t.ID) { e.deleteMu.RUnlock() @@ -757,13 +806,13 @@ func (e *Engine) Run(ctx context.Context, t *Task) { t.Notify() // already running — just nudge a planning round return } - rt := e.registerTaskRoutines(ctx, t.ID, 1+workers) + rt := e.registerTaskRoutines(ctx, t.ID, 1+len(slots)) e.deleteMu.RUnlock() e.touch(t.ID) runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) }) - for i := 0; i < workers; i++ { - name := fmt.Sprintf("work#%d", i+1) - runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, name) }) + for _, slot := range slots { + slot := slot + runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, slot) }) } e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重) // 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已 @@ -913,14 +962,32 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) { } } -func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) { +func (e *Engine) workerLoop(ctx context.Context, t *Task, slot WorkerSlot) { + name := slot.Name for { select { case <-ctx.Done(): return default: } - _, worker := e.snapshotFor(t) + var worker *agent.Worker + var resolveErr error + if slot.ProfileID > 0 { + if e.workerResolve == nil { + resolveErr = errors.New("pinned worker resolver unavailable") + } else { + worker, resolveErr = e.workerResolve(t, slot) + if worker == nil && resolveErr == nil { + resolveErr = errors.New("pinned worker resolver returned no worker") + } + } + e.reportSlotAvailability(t, slot, resolveErr) + if resolveErr != nil { + worker = nil + } + } else { + _, worker = e.snapshotFor(t) + } if worker == nil { if sleepCtx(ctx, 1500*time.Millisecond) { return diff --git a/server/manager.go b/server/manager.go index 09fdcd3..76f72ae 100644 --- a/server/manager.go +++ b/server/manager.go @@ -265,6 +265,7 @@ const ( // only routes the search backend) and the per-profile LLM proxy. settingGlobalProxy = "global_proxy" settingWorkers = "workers" + settingWorkerPools = "worker_pools" settingLLMRecord = "llm_record" // LLM 轮询(故障转移)。默认关闭——开启后走「全局激活配置」的 agent 在当前配置 // 不可用(余额不足/key 失效/限流/服务异常)时自动切到下一个配置。 @@ -291,6 +292,102 @@ const ( defaultConcurrencyLimit = 5 ) +// WorkerPool binds a number of worker slots to one LLM profile. Profile IDs are +// encoded as JSON strings to match the frontend LLMProfile API. +type WorkerPool struct { + ProfileID int64 `json:"profile_id,string"` + Workers int `json:"workers"` +} + +func (p *WorkerPool) UnmarshalJSON(data []byte) error { + var raw struct { + ProfileID json.RawMessage `json:"profile_id"` + Workers int `json:"workers"` + } + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + var id int64 + if err := json.Unmarshal(raw.ProfileID, &id); err != nil { + var text string + if err := json.Unmarshal(raw.ProfileID, &text); err != nil { + return fmt.Errorf("profile_id 必须是整数或数字字符串") + } + parsed, err := strconv.ParseInt(strings.TrimSpace(text), 10, 64) + if err != nil { + return fmt.Errorf("profile_id 必须是整数或数字字符串") + } + id = parsed + } + p.ProfileID, p.Workers = id, raw.Workers + return nil +} + +// WorkerSlot is the concrete routing plan for one task run. +type WorkerSlot struct { + Name string + ProfileID int64 // zero means legacy task-router resolution +} + +// ValidateWorkerPools validates pool shape without requiring a database. Pool +// entries with duplicate profile IDs are rejected so each profile has one +// unambiguous slot group in the saved order. +func ValidateWorkerPools(pools []WorkerPool) error { + if len(pools) == 0 { + return nil + } + seen := make(map[int64]bool, len(pools)) + total := 0 + for _, pool := range pools { + if pool.ProfileID <= 0 { + return fmt.Errorf("worker_pools 中 profile_id 必须大于 0") + } + if pool.Workers <= 0 { + return fmt.Errorf("worker_pools 中 workers 必须大于 0") + } + if seen[pool.ProfileID] { + return fmt.Errorf("worker_pools 中 profile_id %d 重复", pool.ProfileID) + } + seen[pool.ProfileID] = true + if total > 50-pool.Workers { + return fmt.Errorf("worker_pools 的 Worker 总数必须在 1 到 50 之间") + } + total += pool.Workers + } + if total < 1 || total > 50 { + return fmt.Errorf("worker_pools 的 Worker 总数必须在 1 到 50 之间") + } + return nil +} + +// ExpandWorkerSlots expands pools in their saved order; invalid or empty legacy +// settings use the historical worker count (default 3). +func ExpandWorkerSlots(pools []WorkerPool, legacy int) ([]WorkerSlot, error) { + if err := ValidateWorkerPools(pools); err != nil { + return nil, err + } + if len(pools) == 0 { + if legacy <= 0 { + legacy = defaultWorkers + } + slots := make([]WorkerSlot, legacy) + for i := range slots { + slots[i] = WorkerSlot{Name: fmt.Sprintf("work#%d", i+1)} + } + return slots, nil + } + slots := make([]WorkerSlot, 0) + for _, pool := range pools { + for i := 0; i < pool.Workers; i++ { + slots = append(slots, WorkerSlot{ProfileID: pool.ProfileID}) + } + } + for i := range slots { + slots[i].Name = fmt.Sprintf("work#%d", i+1) + } + return slots, nil +} + // ConcurrencyLimit returns whether the simultaneous-running-task cap is enabled and // its limit (default 5 when enabled but unset). limit is always >=1 when enabled. func (m *Manager) ConcurrencyLimit() (enabled bool, limit int) { @@ -332,6 +429,58 @@ func (m *Manager) Workers() int { return n } +func (m *Manager) WorkerPools() []WorkerPool { + v, ok, err := m.pg.GetSetting(settingWorkerPools) + if err != nil || !ok || strings.TrimSpace(v) == "" { + return []WorkerPool{} + } + var pools []WorkerPool + if json.Unmarshal([]byte(v), &pools) != nil || ValidateWorkerPools(pools) != nil { + return []WorkerPool{} + } + if pools == nil { + return []WorkerPool{} + } + return pools +} + +// SetWorkerPools validates profile references at save time and synchronizes the +// legacy workers setting with the pool total. An empty slice restores legacy mode. +func (m *Manager) SetWorkerPools(pools []WorkerPool) error { + if pools == nil { + pools = []WorkerPool{} + } + if err := ValidateWorkerPools(pools); err != nil { + return err + } + if len(pools) > 0 { + for _, pool := range pools { + profile, err := m.pg.ProfileByID(pool.ProfileID) + if err != nil { + return err + } + if profile == nil { + return fmt.Errorf("worker_pools 中 profile_id %d 不存在", pool.ProfileID) + } + } + } + b, err := json.Marshal(pools) + if err != nil { + return err + } + if err := m.pg.SetSetting(settingWorkerPools, string(b)); err != nil { + return err + } + if len(pools) > 0 { + total := 0 + for _, pool := range pools { + total += pool.Workers + } + return m.pg.SetSetting(settingWorkers, strconv.Itoa(total)) + } + return nil +} + // SetWorkers persists the concurrent work-agent count. Values <=0 are rejected. func (m *Manager) SetWorkers(n int) error { if n <= 0 { diff --git a/server/server.go b/server/server.go index 969216c..fc3e0b5 100644 --- a/server/server.go +++ b/server/server.go @@ -177,6 +177,9 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi b := s.agentsForTask(t) return b.pl, b.wk }) + s.engine.SetWorkerResolver(func(t *Task, slot WorkerSlot) (*agent.Worker, error) { + return s.workerForTaskSlot(t, slot) + }) // Global readiness (Ready()/llm_configured): a global active LLM provider is // installed. Task-level runnability is separate (ReadyFor → the resolver above). s.engine.SetReadiness(func() bool { @@ -575,7 +578,8 @@ func (s *Server) chatUnavailableReason() string { // providerForProfile returns a cached provider+cfg for a profile id, so every agent // bound/pinned to the same profile shares one provider instance (one rate limiter). -// ok=false when the profile is missing/invalid → caller falls back to the global pair. +// ok=false when the profile is missing/invalid; pinned callers surface that error +// while legacy callers retain their existing fallback behavior. func (s *Server) providerForProfile(id int64) (llm.Provider, agent.Config, bool) { s.provCacheMu.Lock() if e := s.provByProfile[id]; e != nil { @@ -3432,6 +3436,7 @@ func (s *Server) settingsPayload() map[string]any { "global_proxy": s.m.GlobalProxy(), // 全局出口代理(http/https/socks5)，所有目标流量走它，空=直连 "python_interpreter": strings.TrimSpace(pyStored), // 用户/自动设的值(空=用运行时检测) "workers": s.m.Workers(), // 并发工作 agent 数(默认3)；对之后启动的任务生效 + "worker_pools": s.m.WorkerPools(), // 非空时按池顺序绑定 worker 会话 "task_concurrency_enabled": concOn, // 任务并发上限开关(默认关) "task_concurrency_limit": concLimit, // 同时运行任务上限(开启后默认5) // LLM 轮询(故障转移)。默认关；开启后走全局激活配置的 agent 在当前配置不可用时 @@ -3499,14 +3504,15 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { LLMRecord *bool `json:"llm_record"` // LLM 录制开关（默认关）；即时生效，无需重建 agent // Web search. WebSearchEnabled/Backend toggle the tool + backend; BraveKey/TavilyKey // are optional — omit (null) to leave a stored key untouched, send "" to clear. - WebSearchEnabled *bool `json:"web_search_enabled"` - WebSearchBackend *string `json:"web_search_backend"` - BraveKey *string `json:"brave_search_api_key"` - TavilyKey *string `json:"tavily_search_api_key"` - WebSearchProxy *string `json:"web_search_proxy"` // 独立出口代理(http/https/socks5)；null=不改，""=清空 - GlobalProxy *string `json:"global_proxy"` // 全局出口代理(http/https/socks5)；null=不改，""=清空(直连) - PythonInterp *string `json:"python_interpreter"` // 自定义脚本工具的 python 解释器路径 - Workers *int `json:"workers"` // 并发工作 agent 数(>0)；对之后启动的任务生效 + WebSearchEnabled *bool `json:"web_search_enabled"` + WebSearchBackend *string `json:"web_search_backend"` + BraveKey *string `json:"brave_search_api_key"` + TavilyKey *string `json:"tavily_search_api_key"` + WebSearchProxy *string `json:"web_search_proxy"` // 独立出口代理(http/https/socks5)；null=不改，""=清空 + GlobalProxy *string `json:"global_proxy"` // 全局出口代理(http/https/socks5)；null=不改，""=清空(直连) + PythonInterp *string `json:"python_interpreter"` // 自定义脚本工具的 python 解释器路径 + Workers *int `json:"workers"` // 并发工作 agent 数(>0)；对之后启动的任务生效 + WorkerPools *[]WorkerPool `json:"worker_pools"` // 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。 ConcurrencyEnabled *bool `json:"task_concurrency_enabled"` ConcurrencyLimit *int `json:"task_concurrency_limit"` @@ -3579,11 +3585,34 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) { return } } - if req.Workers != nil { - if err := s.m.SetWorkers(*req.Workers); err != nil { + // A non-empty pool list takes precedence when both fields are present and + // synchronizes workers to its total. An empty list clears pools; when workers + // is also present it sets the legacy count for the newly restored legacy mode. + if req.WorkerPools != nil { + if len(*req.WorkerPools) == 0 && req.Workers != nil && *req.Workers <= 0 { + writeErr(w, 400, "workers 必须 >0") + return + } + if err := s.m.SetWorkerPools(*req.WorkerPools); err != nil { writeErr(w, 400, err.Error()) return } + if len(*req.WorkerPools) == 0 && req.Workers != nil { + if err := s.m.SetWorkers(*req.Workers); err != nil { + writeErr(w, 400, err.Error()) + return + } + } + } else if req.Workers != nil { + // An omitted worker_pools field leaves the saved mode unchanged. If a + // non-empty pool configuration is already saved, its total remains the + // authoritative concurrency and workers is deliberately ignored. + if len(s.m.WorkerPools()) == 0 { + if err := s.m.SetWorkers(*req.Workers); err != nil { + writeErr(w, 400, err.Error()) + return + } + } } if req.ConcurrencyEnabled != nil || req.ConcurrencyLimit != nil { // 部分 PUT:未给的字段用当前值兜底,避免只改一个把另一个重置。 diff --git a/server/task_llm.go b/server/task_llm.go index 2cae9ce..8063a3e 100644 --- a/server/task_llm.go +++ b/server/task_llm.go @@ -27,6 +27,7 @@ type taskAgentBundle struct { pl *agent.Planner wk *agent.Worker main *agent.MainAgent + pinnedWorkers map[int64]*agent.Worker } type llmAuditProfile struct { @@ -48,9 +49,10 @@ type llmActivityMetadata struct { } type taskLLMRuntime struct { - s *Server - taskID string - agentKey string + s *Server + taskID string + agentKey string + pinProfileID *int64 } type taskLLMError struct { @@ -101,13 +103,24 @@ type taskLLMStreamHooks struct { transition func(taskLLMSelection, db.TaskLLMTransition, error) } -// current resolves the provider this role runs on, by precedence: -// Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置。 +// current resolves the provider this role runs on. Pinned workers use only their +// assigned profile; legacy task agents use Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置. // 绑定优先于任务链：某个角色被显式指定了模型，就一直跑在那个模型上；绑定不存在或 // 构建失败时才降级到任务链，任务链为空时再降级到全局配置。 // 返回的 profile id 只有走任务链时才非零 —— streamTaskLLM 以此判断额度错误是否 // 应该推进任务的故障转移状态（绑定/全局路径不改任务链状态，沿用既有语义）。 func (r *taskLLMRuntime) current() (taskLLMSelection, error) { + if r.pinProfileID != nil { + // Pinned worker runtimes are exclusive: resolve their profile before + // touching bindings or the task failover state. A missing or invalid + // profile is an error, never an invitation to use another model. + prov, cfg, ok := r.s.providerForProfile(*r.pinProfileID) + if !ok { + return taskLLMSelection{}, fmt.Errorf("LLM profile #%d is missing or invalid", *r.pinProfileID) + } + t, _ := r.s.m.Task(r.taskID) + return taskLLMSelection{task: t, provider: prov, retry: cfg.Retry}, nil + } taskNum, err := parseTaskID(r.taskID) if err != nil { return taskLLMSelection{}, err @@ -152,6 +165,10 @@ func (r *taskLLMRuntime) current() (taskLLMSelection, error) { // the caller. If failover switches profiles, the change takes effect on the next // agent run (a fresh Session is built per run in captureRun). func (r *taskLLMRuntime) activeCfg() (agent.Config, bool) { + if r.pinProfileID != nil { + _, cfg, ok := r.s.providerForProfile(*r.pinProfileID) + return cfg, ok + } taskNum, err := parseTaskID(r.taskID) if err != nil { return agent.Config{}, false @@ -454,6 +471,12 @@ func isRetryableStreamError(err error) bool { // CompactionWindow mirrors current()'s precedence so the context window always // matches the provider the role will actually stream on. func (r *taskLLMRuntime) CompactionWindow() int { + if r.pinProfileID != nil { + if _, cfg, ok := r.s.providerForProfile(*r.pinProfileID); ok { + return cfg.CompactionWindow() + } + return (agent.Config{}).CompactionWindow() + } if _, cfg, ok := r.s.agentBindingProvider(r.agentKey); ok { return cfg.CompactionWindow() } @@ -597,12 +620,54 @@ func (s *Server) agentsForTask(t *Task) *taskAgentBundle { main.SetSteerWork(s.engine.SteerWork) // steer_work：人对运行中 work 实时纠偏 bundle := &taskAgentBundle{ runtime: goalRuntime, plannerRuntime: plannerRuntime, workerRuntime: workerRuntime, - mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, + mainRuntime: mainRuntime, pl: pl, wk: wk, main: main, pinnedWorkers: map[int64]*agent.Worker{}, } s.taskAgents[t.ID] = bundle return bundle } +func (s *Server) workerForTaskSlot(t *Task, slot WorkerSlot) (*agent.Worker, error) { + if slot.ProfileID <= 0 { + return nil, errors.New("worker slot is not pinned") + } + _, cfg, ok := s.providerForProfile(slot.ProfileID) + if !ok { + return nil, fmt.Errorf("LLM profile #%d is missing or invalid", slot.ProfileID) + } + s.taskAgentMu.Lock() + bundle := s.taskAgents[t.ID] + s.taskAgentMu.Unlock() + if bundle == nil { + bundle = s.agentsForTask(t) + } + s.taskAgentMu.Lock() + defer s.taskAgentMu.Unlock() + if current := s.taskAgents[t.ID]; current != nil { + bundle = current + } + if bundle.pinnedWorkers == nil { + bundle.pinnedWorkers = make(map[int64]*agent.Worker) + } + if worker := bundle.pinnedWorkers[slot.ProfileID]; worker != nil { + return worker, nil + } + id := slot.ProfileID + runtime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "worker", pinProfileID: &id} + tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts")) + worker := agent.NewWorker(runtime, cfg.Model, s.m.dir, tx, runtime.CompactionWindow(), s.agentMaxTurns("worker")) + worker.SetFindingRecorder(s.evidenceStore()) + worker.SetCompactionWindowResolver(runtime.CompactionWindow) + worker.SetNonStreaming(runtime.nonStreaming) + worker.SetMaxTokens(runtime.maxTokens) + worker.SetNoaEnabled(s.m.NoaCompactionEnabled) + worker.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second) + worker.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) + worker.SetWebSearch(s.webSearchFor("worker")) + worker.SetConstraintInject(s.constraintInjectWorker) + bundle.pinnedWorkers[id] = worker + return worker, nil +} + func (s *Server) syncTaskLLMState(pt *db.Task) { if pt == nil { return diff --git a/server/worker_pools_test.go b/server/worker_pools_test.go new file mode 100644 index 0000000..36f40aa --- /dev/null +++ b/server/worker_pools_test.go @@ -0,0 +1,263 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Autumn-27/artex/agent" + "github.com/Autumn-27/artex/db" +) + +func TestExpandWorkerSlotsLegacy(t *testing.T) { + for _, tc := range []struct { + name string + count int + want int + }{ + {"configured", 2, 2}, + {"zero defaults", 0, 3}, + {"negative defaults", -1, 3}, + } { + t.Run(tc.name, func(t *testing.T) { + slots, err := ExpandWorkerSlots(nil, tc.count) + if err != nil || len(slots) != tc.want { + t.Fatalf("slots=%v err=%v, want %d", slots, err, tc.want) + } + for i, slot := range slots { + if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != 0 { + t.Fatalf("slot[%d]=%+v", i, slot) + } + } + }) + } +} + +func TestExpandWorkerSlotsPoolOrder(t *testing.T) { + slots, err := ExpandWorkerSlots([]WorkerPool{{ProfileID: 7, Workers: 2}, {ProfileID: 12, Workers: 3}}, 99) + if err != nil || len(slots) != 5 { + t.Fatalf("slots=%v err=%v", slots, err) + } + for i, slot := range slots { + want := int64(7) + if i >= 2 { + want = 12 + } + if slot.Name != "work#"+strconv.Itoa(i+1) || slot.ProfileID != want { + t.Fatalf("slot[%d]=%+v want profile %d", i, slot, want) + } + } +} + +func TestValidateWorkerPools(t *testing.T) { + for _, tc := range []struct { + name string + pools []WorkerPool + valid bool + }{ + {"empty", nil, true}, + {"bad profile", []WorkerPool{{ProfileID: 0, Workers: 1}}, false}, + {"bad workers", []WorkerPool{{ProfileID: 1, Workers: 0}}, false}, + {"duplicate", []WorkerPool{{ProfileID: 1, Workers: 1}, {ProfileID: 1, Workers: 1}}, false}, + {"too many", []WorkerPool{{ProfileID: 1, Workers: 51}}, false}, + {"valid", []WorkerPool{{ProfileID: 1, Workers: 1}}, true}, + } { + t.Run(tc.name, func(t *testing.T) { + if got := ValidateWorkerPools(tc.pools) == nil; got != tc.valid { + t.Fatalf("valid=%v want %v", got, tc.valid) + } + }) + } +} + +func TestWorkerPoolSettingsRoundTrip(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + pID, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pool-test", Format: "openai", Model: "pool-test", APIKey: "pool-test-key"}) + if err != nil { + t.Fatal(err) + } + defer m.pg.DeleteProfile(pID) + oldPools, _, _ := m.pg.GetSetting(settingWorkerPools) + oldWorkers, workersSet, _ := m.pg.GetSetting(settingWorkers) + defer func() { + if workersSet { + _ = m.pg.SetSetting(settingWorkers, oldWorkers) + } else { + _, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkers) + } + if oldPools != "" { + _ = m.pg.SetSetting(settingWorkerPools, oldPools) + } else { + _, _ = m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools) + } + }() + if err := m.SetWorkerPools([]WorkerPool{{ProfileID: 9223372036854775807, Workers: 1}}); err == nil { + t.Fatal("nonexistent profile accepted") + } + if err := m.SetWorkerPools([]WorkerPool{{ProfileID: pID, Workers: 2}}); err != nil { + t.Fatal(err) + } + got := m.WorkerPools() + if len(got) != 1 || got[0].ProfileID != pID || got[0].Workers != 2 || m.Workers() != 2 { + t.Fatalf("round trip pools=%+v workers=%d", got, m.Workers()) + } + if err := m.SetWorkerPools(nil); err != nil { + t.Fatal(err) + } + if got := m.WorkerPools(); got == nil || len(got) != 0 { + t.Fatalf("cleared pools=%+v", got) + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir()) + if _, err := m.pg.Exec("DELETE FROM settings WHERE key=$1", settingWorkerPools); err != nil { + t.Fatal(err) + } + res := httptest.NewRecorder() + s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil)) + var unset map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &unset); err != nil { + t.Fatal(err) + } + if pools, ok := unset["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 { + t.Fatalf("GET unset worker_pools=%#v, want empty array", unset["worker_pools"]) + } + put := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":2,"worker_pools":[{"profile_id":"`+strconv.FormatInt(pID, 10)+`","workers":2}]}`)) + putResult := httptest.NewRecorder() + s.putSettings(putResult, put) + if putResult.Code != http.StatusOK { + t.Fatalf("PUT worker pools status=%d body=%s", putResult.Code, putResult.Body.String()) + } + if m.Workers() != 2 || len(m.WorkerPools()) != 1 { + t.Fatalf("PUT worker pools did not sync settings: workers=%d pools=%+v", m.Workers(), m.WorkerPools()) + } + workersOnly := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":9}`)) + workersOnlyResult := httptest.NewRecorder() + s.putSettings(workersOnlyResult, workersOnly) + if workersOnlyResult.Code != http.StatusOK || m.Workers() != 2 { + t.Fatalf("workers-only PUT changed pool total: status=%d workers=%d", workersOnlyResult.Code, m.Workers()) + } + req := httptest.NewRequest(http.MethodGet, "/api/settings", nil) + res = httptest.NewRecorder() + s.getSettings(res, req) + if res.Code != http.StatusOK { + t.Fatalf("GET settings status=%d", res.Code) + } + var payload map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &payload); err != nil { + t.Fatal(err) + } + pools, ok := payload["worker_pools"].([]any) + if !ok || len(pools) != 1 || pools[0].(map[string]any)["profile_id"] != strconv.FormatInt(pID, 10) { + t.Fatalf("GET worker_pools=%#v, want string profile id", payload["worker_pools"]) + } + clear := httptest.NewRequest(http.MethodPut, "/api/settings", strings.NewReader(`{"workers":4,"worker_pools":[]}`)) + clearResult := httptest.NewRecorder() + s.putSettings(clearResult, clear) + if clearResult.Code != http.StatusOK || m.Workers() != 4 || len(m.WorkerPools()) != 0 { + t.Fatalf("PUT legacy clear status=%d body=%s workers=%d pools=%+v", clearResult.Code, clearResult.Body.String(), m.Workers(), m.WorkerPools()) + } + res = httptest.NewRecorder() + s.getSettings(res, httptest.NewRequest(http.MethodGet, "/api/settings", nil)) + var cleared map[string]any + if err := json.Unmarshal(res.Body.Bytes(), &cleared); err != nil { + t.Fatal(err) + } + if pools, ok := cleared["worker_pools"].([]any); !ok || pools == nil || len(pools) != 0 { + t.Fatalf("GET cleared worker_pools=%#v, want empty array", cleared["worker_pools"]) + } +} + +func TestPinnedRuntimeMissingProfileDoesNotFallback(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned runtime test", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + fallback, err := m.pg.SaveProfile(&db.LLMProfile{Name: "pinned-fallback", Format: "openai", Model: "fallback", APIKey: "fallback-key"}) + if err != nil { + t.Fatal(err) + } + defer m.pg.DeleteProfile(fallback) + if _, err := m.ReplaceTaskLLMProfiles(task.ID, []int64{fallback}, fallback); err != nil { + t.Fatal(err) + } + missing := int64(9223372036854775807) + s := &Server{m: m, llmOn: true, llmProv: &scriptedLLMProvider{}} + runtime := &taskLLMRuntime{s: s, taskID: task.ID, agentKey: "worker", pinProfileID: &missing} + if _, err := runtime.current(); err == nil { + t.Fatal("missing pinned profile silently fell back") + } + if _, ok := runtime.activeCfg(); ok || runtime.nonStreaming() || runtime.maxTokens() != 0 { + t.Fatal("missing pinned profile exposed fallback runtime settings") + } +} + +func TestPinnedWorkerLoopDoesNotConsultTaskRouterAndReportsFailureOnce(t *testing.T) { + m, err := NewManager(t.TempDir(), "") + if err != nil { + t.Skipf("postgres unavailable: %v", err) + } + defer m.Close() + task, err := m.CreateTask("pinned worker loop", "test", nil, 0, 0) + if err != nil { + t.Fatal(err) + } + defer m.DeleteTask(task.ID, DeleteTaskOptions{}) + + e := NewEngine(m) + var routerCalls atomic.Int32 + e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) { + routerCalls.Add(1) + return nil, nil + }) + e.SetWorkerResolver(func(*Task, WorkerSlot) (*agent.Worker, error) { + return nil, errors.New("profile unavailable") + }) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + e.workerLoop(ctx, task, WorkerSlot{Name: "work#4", ProfileID: 12}) + close(done) + }() + time.Sleep(100 * time.Millisecond) + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("worker loop did not stop") + } + if routerCalls.Load() != 0 { + t.Fatalf("pinned worker consulted task router %d times", routerCalls.Load()) + } + activities, _, err := task.Store.ActivityList(nil, 0, 20) + if err != nil { + t.Fatal(err) + } + var failures int + for _, activity := range activities { + if activity.Worker == "work#4" && activity.IsError { + failures++ + } + } + if failures != 1 { + t.Fatalf("pinned failure activities=%d, want one throttled entry: %+v", failures, activities) + } +} diff --git a/web/src/app/(main)/system/settings/page.tsx b/web/src/app/(main)/system/settings/page.tsx index fea869f..84bfe45 100644 --- a/web/src/app/(main)/system/settings/page.tsx +++ b/web/src/app/(main)/system/settings/page.tsx @@ -14,7 +14,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { Switch } from "@/components/ui/switch"; import { api } from "@/lib/api"; import { CHAT_SEND_MODE_OPTIONS, type ChatSendMode, setChatSendMode, useChatSendMode } from "@/lib/chat-send-mode"; -import type { Settings } from "@/lib/types"; +import type { LLMProfile, Settings, WorkerPool } from "@/lib/types"; import { BasicAuthCard } from "./_components/basic-auth-card"; import { UpdateCard } from "./_components/update-card"; @@ -45,6 +45,8 @@ export default function SystemSettingsPage() { const [pyInterp, setPyInterp] = React.useState(""); const [workers, setWorkers] = React.useState("3"); const [savingWorkers, setSavingWorkers] = React.useState(false); + const [workerPools, setWorkerPools] = React.useState([]); + const [llmProfiles, setLlmProfiles] = React.useState([]); // 操作约束注入范围(默认都开)。 const [injectPlanner, setInjectPlanner] = React.useState(true); const [injectWorker, setInjectWorker] = React.useState(true); @@ -67,6 +69,7 @@ export default function SystemSettingsPage() { setGlobalProxyInput(s.global_proxy ?? ""); setPyInterp(s.python_interpreter ?? ""); setWorkers(String(s.workers ?? 3)); + setWorkerPools(s.worker_pools); setInjectPlanner(s.constraints_inject_planner !== false); setInjectWorker(s.constraints_inject_worker !== false); setNoaCompaction(!!s.noa_compaction); @@ -74,13 +77,46 @@ export default function SystemSettingsPage() { const saveWorkers = () => { const n = Number(workers); + if (workerPools.length > 0) { + const total = workerPools.reduce((sum, pool) => sum + Number(pool.workers), 0); + if (workerPools.some((pool) => !pool.profile_id)) { + toast.error("请选择每个模型池的 LLM 配置"); + return; + } + if (workerPools.some((pool) => !Number.isInteger(Number(pool.workers)) || Number(pool.workers) <= 0)) { + toast.error("每个模型池的 Worker 数必须是大于 0 的整数"); + return; + } + if (total > 50) { + toast.error("Worker 总数不能超过 50"); + return; + } + const ids = workerPools.map((pool) => pool.profile_id); + if (new Set(ids).size !== ids.length) { + toast.error("不能重复选择同一个 LLM 配置"); + return; + } + setSavingWorkers(true); + api + .setSettings({ + workers: total, + worker_pools: workerPools.map((pool) => ({ ...pool, workers: Number(pool.workers) })), + }) + .then((s) => { + apply(s); + toast.success("已保存多模型 Worker 池（对之后启动的任务生效）"); + }) + .catch((e) => toast.error(`保存失败：${(e as Error).message}`)) + .finally(() => setSavingWorkers(false)); + return; + } if (!Number.isInteger(n) || n <= 0) { toast.error("并发数必须是大于 0 的整数"); return; } setSavingWorkers(true); api - .setSettings({ workers: n }) + .setSettings({ workers: n, worker_pools: [] }) .then((s) => { apply(s); toast.success("已保存并发工作 agent 数（对之后启动的任务生效）"); @@ -89,6 +125,8 @@ export default function SystemSettingsPage() { .finally(() => setSavingWorkers(false)); }; + const addWorkerPool = () => setWorkerPools((pools) => [...pools, { profile_id: "", workers: 1 }]); + const savePython = () => { setSaving(true); api @@ -117,6 +155,13 @@ export default function SystemSettingsPage() { .finally(() => setLoaded(true)); }, [apply]); + React.useEffect(() => { + api + .llmProfiles() + .then(setLlmProfiles) + .catch(() => setLlmProfiles([])); + }, []); + const toggleTraffic = (v: boolean) => { setTrafficCapture(v); // optimistic setSaving(true); @@ -756,6 +801,73 @@ export default function SystemSettingsPage() { +
+
+ + +
+ {workerPools.map((pool, index) => { + const selected = llmProfiles.find((profile) => profile.id === pool.profile_id); + const placeholder = + selected?.name || (pool.profile_id ? `配置 #${pool.profile_id} 已删除/不可用` : "选择 LLM 配置"); + return ( + // biome-ignore lint/suspicious/noArrayIndexKey: pool rows are edited in place and only deleted. +
+ + + setWorkerPools((current) => + current.map((item, itemIndex) => + itemIndex === index ? { ...item, workers: Number(event.target.value) } : item, + ), + ) + } + /> + +
+ ); + })} + {workerPools.length > 0 ? ( +
+ 当前合计：{workerPools.reduce((sum, pool) => sum + (Number(pool.workers) || 0), 0)} 个 Worker 会话 +
+ ) : null} +
0} onChange={(e) => setWorkers(e.target.value)} /> + {workerPools.length > 0 ? ( + 已配置模型池时由各池数量决定并发。 + ) : null} diff --git a/web/src/lib/mock/data.ts b/web/src/lib/mock/data.ts index 531df7b..9b8a69c 100644 --- a/web/src/lib/mock/data.ts +++ b/web/src/lib/mock/data.ts @@ -3338,6 +3338,7 @@ export const settings: Settings = { global_proxy: "", python_interpreter: "/usr/bin/python3", workers: 3, + worker_pools: [], llm_pool_enabled: true, llm_pool_bind_fallback: false, noa_compaction: false, diff --git a/web/src/lib/mock/handler.ts b/web/src/lib/mock/handler.ts index 2f15a69..185d5a7 100644 --- a/web/src/lib/mock/handler.ts +++ b/web/src/lib/mock/handler.ts @@ -2249,7 +2249,15 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re if (path === "/traffic") return D.traffic; if (path === "/traffic/exchange") return D.trafficDetail; if (path === "/settings" && m === "GET") return D.settings; - if (path === "/settings" && m === "PUT") return { ...D.settings, ...b }; + if (path === "/settings" && m === "PUT") { + const next = { ...D.settings, ...b }; + if (!Array.isArray(next.worker_pools)) next.worker_pools = D.settings.worker_pools; + if (next.worker_pools.length > 0) { + next.workers = next.worker_pools.reduce((sum, pool) => sum + Number(pool.workers), 0); + } + Object.assign(D.settings, next); + return next; + } if (path === "/settings/basic-auth" && m === "GET") return { enabled: false, username: "", password_set: false }; if (path === "/settings/basic-auth" && m === "PUT") throw new Error("演示模式不支持修改 HTTP Basic Auth，请在实际部署中配置"); @@ -2467,7 +2475,8 @@ function route(m: string, path: string, seg: string[], q: URLSearchParams, b: Re const from = q.get("from") || ""; const to = q.get("to") || ""; if (status && !["pending", "allowed", "denied", "timeout"].includes(status)) throw new Error("无效审批状态"); - if (decisionSource && !["model", "rule", "manual", "unknown"].includes(decisionSource)) throw new Error("无效判定来源"); + if (decisionSource && !["model", "rule", "manual", "unknown"].includes(decisionSource)) + throw new Error("无效判定来源"); if ([...search].length > 200) throw new Error("搜索关键词不能超过 200 个字符"); if (search.startsWith("#") && !/^#[1-9]\d*$/.test(search)) throw new Error("审批编号必须为 # 后接正整数"); for (const value of [from, to]) { diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts index 54e26a6..e080c12 100644 --- a/web/src/lib/types.ts +++ b/web/src/lib/types.ts @@ -968,6 +968,7 @@ export interface Settings { global_proxy?: string; python_interpreter?: string; // 自定义脚本工具的 python 解释器路径(空=运行时检测) workers?: number; // 并发工作 agent 数(默认3)；对之后启动的任务生效 + worker_pools: WorkerPool[]; // 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。 task_concurrency_enabled?: boolean; // 默认 false task_concurrency_limit?: number; // 开启后默认 5 @@ -989,6 +990,11 @@ export interface Settings { notify_digest_interval_min?: number; // 汇总模式周期（分钟），默认 30 } +export interface WorkerPool { + profile_id: string; + workers: number; +} + // ---- 漏洞 IM 推送 ---- // NotificationFilter 是渠道的过滤条件，字段全部可选，缺省即不过滤。