This commit is contained in:
pphua committed 2026-08-24 14:14:15 +08:00
1 parent e67ff26e90
commit 58f23a2773
61 files changed
+11611 -443

No files matched your search

+177
View File
@@ -0,0 +1,177 @@
package server
import (
"encoding/json"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/db"
)
// 总览「约束管理」的人工 CRUD 接口 + 注入范围开关的解析。操作约束(allow/deny)与 agent 侧的
// set_constraints 工具写同一张 task_constraints 表;这里是人类在 UI 上直接增删改。约束仅是
// 提示上下文——增删改后【不】通知 planner,下一轮规划自然读库生效(按产品决策)。每个变更 handler
// 都走 beginTaskOperation/decInflight,避免与任务删除竞态(与目标/意图 CRUD 一致)。
// 注入范围开关的 settings key,默认都开(GetBool 第二参数 = true)。
const (
settingConstraintsInjectPlanner = "constraints_inject_planner"
settingConstraintsInjectWorker = "constraints_inject_worker"
)
// constraintInjectPlanner / constraintInjectWorker 报告是否把操作约束注入对应 agent 的
// 系统提示(默认开)。作为 resolver 传给 planner/worker,每轮读 → 改开关即时生效。
func (s *Server) constraintInjectPlanner() bool {
return s.m.pg.GetBool(settingConstraintsInjectPlanner, true)
}
func (s *Server) constraintInjectWorker() bool {
return s.m.pg.GetBool(settingConstraintsInjectWorker, true)
}
// listConstraints 返回本任务的全部操作约束(allow 在前、deny 在后)。
func (s *Server) listConstraints(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
rows, err := t.Store.ListConstraints()
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"constraints": constraintDTOs(rows)})
}
// addConstraint 人工新增一条操作约束(kind=allow|deny)。不通知 planner。
func (s *Server) addConstraint(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法新增约束")
return
}
defer s.engine.decInflight(t.ID)
var body struct {
Text string `json:"text"`
Kind string `json:"kind"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeErr(w, 400, "invalid JSON")
return
}
text := strings.TrimSpace(body.Text)
if text == "" {
writeErr(w, 400, "约束内容不能为空")
return
}
kind := normalizeConstraintKind(body.Kind)
if kind == "" {
writeErr(w, 400, "kind 必须是 allow 或 deny")
return
}
id, err := t.Store.AddConstraint(kind, text, "human")
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, ConstraintDTO{ID: strconv.FormatInt(id, 10), Kind: kind, Text: text, Origin: "human"})
}
// editConstraint 人工修改一条约束(kind + text)。不通知 planner。
func (s *Server) editConstraint(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法修改约束")
return
}
defer s.engine.decInflight(t.ID)
cid, err := strconv.ParseInt(r.PathValue("cid"), 10, 64)
if err != nil || cid <= 0 {
writeErr(w, 400, "bad constraint id")
return
}
var body struct {
Text string `json:"text"`
Kind string `json:"kind"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeErr(w, 400, "invalid JSON")
return
}
text := strings.TrimSpace(body.Text)
if text == "" {
writeErr(w, 400, "约束内容不能为空")
return
}
kind := normalizeConstraintKind(body.Kind)
if kind == "" {
writeErr(w, 400, "kind 必须是 allow 或 deny")
return
}
if err := t.Store.UpdateConstraint(cid, kind, text); err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, ConstraintDTO{ID: strconv.FormatInt(cid, 10), Kind: kind, Text: text})
}
// deleteConstraint 人工删除一条约束。不通知 planner。
func (s *Server) deleteConstraint(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法删除约束")
return
}
defer s.engine.decInflight(t.ID)
cid, err := strconv.ParseInt(r.PathValue("cid"), 10, 64)
if err != nil || cid <= 0 {
writeErr(w, 400, "bad constraint id")
return
}
if err := t.Store.DeleteConstraint(cid); err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]bool{"ok": true})
}
// normalizeConstraintKind lowercases + validates the kind; "" on invalid.
func normalizeConstraintKind(k string) string {
k = strings.TrimSpace(strings.ToLower(k))
if k == "allow" || k == "deny" {
return k
}
return ""
}
// constraintDTOs converts db rows to the frontend shape.
func constraintDTOs(in []db.Constraint) []ConstraintDTO {
out := make([]ConstraintDTO, 0, len(in))
for _, c := range in {
out = append(out, ConstraintDTO{
ID: strconv.FormatInt(c.ID, 10),
Kind: c.Kind,
Text: c.Text,
Origin: c.Origin,
TS: rfc3339(c.CreatedAt),
})
}
return out
}
+94
View File
@@ -30,6 +30,7 @@ func rawString(raw json.RawMessage) string {
// ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status.
type TaskDTO struct {
<<<<<<< Updated upstream
ID string `json:"id"`
ExplorationID int64 `json:"exploration_id"`
Description string `json:"description"`
@@ -45,6 +46,32 @@ type TaskDTO struct {
GoalsTotal int `json:"goals_total"`
GoalsMet int `json:"goals_met"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default
=======
ID string `json:"id"`
ExplorationID int64 `json:"exploration_id"`
Name string `json:"name"` // 可选任务名称;空=未命名
Description string `json:"description"`
Goal string `json:"goal"`
Status string `json:"status"` // created | running | paused | done | failed
CreatedAt string `json:"created_at"`
CreatedUnix int64 `json:"created_unix"` // created_at as unix seconds (for run-duration calc)
CompletedAt string `json:"completed_at"` // RFC3339 finish time (done/failed); "" if unfinished
CompletedUnix int64 `json:"completed_unix"` // completed_at as unix seconds (0 if unfinished)
LastActivity int64 `json:"last_activity_unix"` // unix seconds of the last activity (0 if none)
Paused bool `json:"paused"`
Queued bool `json:"queued"`
Tokens TokenTotalDTO `json:"tokens"` // whole-task token consumption
GoalsTotal int `json:"goals_total"`
GoalsMet int `json:"goals_met"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // LLM profile used for this task; nil = default
LLMProfileIDs []int64 `json:"llm_profile_ids"`
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
LLMFailoverState string `json:"llm_failover_state"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []string `json:"source_task_ids"`
CompanyIDs []int64 `json:"company_ids"`
CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定)
>>>>>>> Stashed changes
}
// TokenTotalDTO is a whole-task (all agents) token aggregate.
@@ -66,6 +93,7 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
func taskDTO(t *Task, status string) TaskDTO {
return TaskDTO{
<<<<<<< Updated upstream
ID: t.ID,
ExplorationID: t.ExpID,
Description: t.Description,
@@ -77,6 +105,28 @@ func taskDTO(t *Task, status string) TaskDTO {
CompletedUnix: t.CompletedAt,
Paused: t.Paused,
LLMProfileID: t.LLMProfileID,
=======
ID: t.ID,
ExplorationID: t.ExpID,
Name: t.Name,
Description: t.Description,
Goal: t.Goal,
Status: status,
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
CreatedUnix: t.CreatedAt,
CompletedAt: completedRFC(lifecycle.CompletedAt),
CompletedUnix: lifecycle.CompletedAt,
Paused: lifecycle.Paused,
Queued: lifecycle.Queued,
LLMProfileID: llmState.ProfileID,
LLMProfileIDs: profileIDs,
ActiveLLMProfileID: llmState.ActiveID,
LLMFailoverState: llmState.FailoverState,
LLMFailoverReason: llmState.FailoverReason,
SourceTaskIDs: sourceIDs,
CompanyIDs: lifecycle.CompanyIDs,
CoverageEnabled: t.CoverageEnabled,
>>>>>>> Stashed changes
}
}
@@ -141,6 +191,50 @@ func taskNodeDTO(n *db.Node) TaskNodeDTO {
}
}
// GoalDTO is a goal node with its payload unpacked into text/vulnclass — the shape
// the 总览「目标管理」UI works with (vs TaskNodeDTO which carries raw payload JSON).
type GoalDTO struct {
ID string `json:"id"`
Text string `json:"text"`
VulnClass string `json:"vulnclass,omitempty"`
State string `json:"state"`
Origin string `json:"origin,omitempty"`
TS string `json:"ts"`
}
func goalDTO(n *db.Node) GoalDTO {
var p struct {
Text string `json:"text"`
VulnClass string `json:"vulnclass"`
}
_ = json.Unmarshal(n.Payload, &p)
return GoalDTO{
ID: i64s(n.ID),
Text: p.Text,
VulnClass: p.VulnClass,
State: n.State,
Origin: n.Origin,
TS: rfc3339(n.CreatedAt),
}
}
func goalDTOs(in []*db.Node) []GoalDTO {
out := make([]GoalDTO, 0, len(in))
for _, n := range in {
out = append(out, goalDTO(n))
}
return out
}
// ConstraintDTO is one operation constraint (allow/deny) for the 总览「约束管理」UI.
type ConstraintDTO struct {
ID string `json:"id"`
Kind string `json:"kind"` // allow | deny
Text string `json:"text"`
Origin string `json:"origin,omitempty"`
TS string `json:"ts,omitempty"`
}
func taskNodeDTOs(in []*db.Node) []TaskNodeDTO {
out := make([]TaskNodeDTO, 0, len(in))
for _, n := range in {
+172
View File
@@ -0,0 +1,172 @@
package server
import (
"encoding/json"
"net/http"
"strconv"
"strings"
"github.com/Autumn-27/artex/db"
)
// 总览「目标管理」的人工 CRUD 接口。与 agent 侧的 set_goals 工具写同一批 goal 节点,
// 但入口是人类在 UI 上直接增删改;新增/修改后复用「复活任务」逻辑(admitTask resume:
// 终态→running、解除暂停、必要时排队),删除不复活(按产品决策)。每个变更 handler 都走
// beginTaskOperation/decInflight,避免与任务删除竞态(与意图 CRUD 一致)。
// listGoals 返回本任务的全部目标(text/vulnclass/state 拆好),供目标管理卡片渲染。
func (s *Server) listGoals(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
goals, err := t.Store.ListByKind(db.KindGoal, 10000)
if err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"goals": goalDTOs(goals)})
}
// addGoal 人工新增一个目标:落库(挂到任务根 spawns 下)→ 记一条「新增了目标」触发唤醒
// planner → 复活任务,让规划者据新目标重判是否达成。
func (s *Server) addGoal(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法新增目标")
return
}
defer s.engine.decInflight(t.ID)
var body struct {
Text string `json:"text"`
VulnClass string `json:"vulnclass"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeErr(w, 400, "invalid JSON")
return
}
text := strings.TrimSpace(body.Text)
if text == "" {
writeErr(w, 400, "目标内容不能为空")
return
}
payload := map[string]any{"text": text}
if vc := strings.TrimSpace(body.VulnClass); vc != "" {
payload["vulnclass"] = vc
}
id, err := t.Store.AddGoal(payload, "human")
if err != nil {
writeErr(w, 500, err.Error())
return
}
if of, _ := t.Store.OriginFactID(); of > 0 && id > 0 {
_ = t.Store.Link(of, db.RelSpawns, id) // goal descends from the task root (origin fact)
}
t.NotifyGoal([]string{text}) // 记「人新增了目标:…」触发并唤醒 planner
s.reviveTask(t) // 把已完成/暂停的任务拉回运行态继续跑
node, _ := t.Store.GetNode(id)
if node == nil {
writeErr(w, 500, "目标写入后读取失败")
return
}
writeJSON(w, 200, goalDTO(node))
}
// editGoal 人工修改一个目标文本(及 vulnclass):改库 → 记「用户修改了目标由 old 变为 new」
// 触发唤醒 planner → 复活任务,让规划者据新目标调整方向。
func (s *Server) editGoal(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法修改目标")
return
}
defer s.engine.decInflight(t.ID)
gid, err := strconv.ParseInt(r.PathValue("gid"), 10, 64)
if err != nil || gid <= 0 {
writeErr(w, 400, "bad goal id")
return
}
var body struct {
Text string `json:"text"`
VulnClass string `json:"vulnclass"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeErr(w, 400, "invalid JSON")
return
}
text := strings.TrimSpace(body.Text)
if text == "" {
writeErr(w, 400, "目标内容不能为空")
return
}
node, err := t.Store.GetNode(gid)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if node == nil || node.Kind != db.KindGoal {
writeErr(w, 404, "目标不存在")
return
}
oldText := goalDTO(node).Text
if err := t.Store.UpdateGoalPayload(gid, text, strings.TrimSpace(body.VulnClass)); err != nil {
writeErr(w, 500, err.Error())
return
}
t.NotifyGoalEdited(oldText, text) // 记「人修改了目标由 old 变为 new」触发并唤醒 planner
s.reviveTask(t) // 与新增一致:复活任务据新目标重判
updated, _ := t.Store.GetNode(gid)
if updated == nil {
writeErr(w, 500, "目标更新后读取失败")
return
}
writeJSON(w, 200, goalDTO(updated))
}
// deleteGoal 人工删除一个目标(硬删除,级联删边/锚点):删库 → 记「用户删除了该目标 X」
// 触发唤醒 planner 据此重判剩余目标。按产品决策,删除【不】复活任务。
func (s *Server) deleteGoal(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法删除目标")
return
}
defer s.engine.decInflight(t.ID)
gid, err := strconv.ParseInt(r.PathValue("gid"), 10, 64)
if err != nil || gid <= 0 {
writeErr(w, 400, "bad goal id")
return
}
node, err := t.Store.GetNode(gid)
if err != nil {
writeErr(w, 500, err.Error())
return
}
if node == nil || node.Kind != db.KindGoal {
writeErr(w, 404, "目标不存在")
return
}
text := goalDTO(node).Text
if err := t.Store.DeleteGoal(gid); err != nil {
writeErr(w, 500, err.Error())
return
}
t.NotifyGoalDeleted(text) // 记「人删除了该目标:…」触发并唤醒 planner(不复活任务)
writeJSON(w, 200, map[string]bool{"ok": true})
}
+90
View File
@@ -1,12 +1,16 @@
package server
import (
"context"
"fmt"
"net/http"
"regexp"
"strings"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/guard"
"github.com/Autumn-27/artex/intercept"
"github.com/Autumn-27/norma/llm"
)
// chatGuard returns a guard wired with the manager's interceptor, used for chat
@@ -15,6 +19,58 @@ func (s *Server) chatGuard() *guard.Guard {
return guard.NewWithInterceptor(s.m.interceptor)
}
// wireInterceptReviewer installs the LLM fallback judge into the interceptor. The
// judge runs only on tool calls that matched no rule (see intercept.Judge). It
// resolves the configured judge profile (0 → active/default), builds a provider,
// runs a one-shot single-line classification, and parses ALLOW/ASK/DENY.
func (s *Server) wireInterceptReviewer() {
s.m.interceptor.SetReviewer(func(ctx context.Context, profileID int64, prompt, tool, command string) (intercept.Decision, error) {
if profileID == 0 {
if p, err := s.m.pg.ActiveProfile(); err == nil && p != nil {
profileID = p.ID
}
}
if profileID == 0 {
return intercept.Decision{}, fmt.Errorf("未配置可用的裁判模型")
}
prov, _, ok := s.providerForProfile(profileID)
if !ok {
return intercept.Decision{}, fmt.Errorf("裁判模型 profile %d 不可用", profileID)
}
user := fmt.Sprintf("Tool: %s\nArguments:\n%s", tool, command)
text, err := streamCollectText(ctx, prov, prompt, user)
if err != nil {
return intercept.Decision{}, err
}
v := intercept.ParseVerdict(text)
return intercept.Decision{Action: v.Action, Message: v.Reason}, nil
})
}
// streamCollectText runs a single non-streaming-style completion (thinking off,
// low temperature, tiny output) and returns the concatenated text. Used by the
// LLM fallback judge, whose reply is one short line (ALLOW/ASK:.../DENY:...).
func streamCollectText(ctx context.Context, prov llm.Provider, system, user string) (string, error) {
temp := 0.0
req := llm.CompletionRequest{
System: []string{system},
Messages: []llm.Message{llm.UserText(user)},
MaxTokens: 128,
Temperature: &temp,
Thinking: "disabled",
}
var sb strings.Builder
for ev, err := range prov.Stream(ctx, req) {
if err != nil {
return "", err
}
if ev.Type == llm.SETextDelta {
sb.WriteString(ev.Text)
}
}
return sb.String(), nil
}
// --- intercept rule CRUD ---
func (s *Server) interceptListRules(w http.ResponseWriter, r *http.Request) {
@@ -261,6 +317,40 @@ func (s *Server) interceptSetToolConfig(w http.ResponseWriter, r *http.Request)
writeJSON(w, 200, map[string]any{"ok": true})
}
// --- LLM fallback judge config (全局模型兜底) ---
// interceptGetJudgeConfig returns the resolved judge configuration. Prompt is the
// effective prompt (built-in template when unset), so the UI can prefill it.
func (s *Server) interceptGetJudgeConfig(w http.ResponseWriter, r *http.Request) {
writeJSON(w, 200, s.m.interceptor.GetJudgeConfig())
}
// interceptSetJudgeConfig persists the judge configuration.
func (s *Server) interceptSetJudgeConfig(w http.ResponseWriter, r *http.Request) {
var req intercept.JudgeConfig
if err := decode(r, &req); err != nil {
writeErr(w, 400, err.Error())
return
}
switch req.FailAction {
case "allow", "ask", "deny":
default:
writeErr(w, 400, "fail_action 必须是 allow、ask 或 deny")
return
}
switch req.AskTimeoutAction {
case "allow", "deny":
default:
writeErr(w, 400, "ask_timeout_action 必须是 allow 或 deny")
return
}
if err := s.m.interceptor.SetJudgeConfig(req); err != nil {
writeErr(w, 500, err.Error())
return
}
writeJSON(w, 200, map[string]any{"ok": true})
}
// --- helpers ---
type interceptRuleReq struct {
+86 -1
View File
@@ -25,6 +25,7 @@ import (
// sharing the process-wide asset store. ID is the PG task id as a string; ExpID
// is the exploration the task owns.
type Task struct {
<<<<<<< Updated upstream
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Description string `json:"description"`
@@ -43,6 +44,43 @@ type Task struct {
Store *pgdb.ExplorationStore `json:"-"`
Guard *guard.Guard `json:"-"`
notify chan struct{}
=======
ID string `json:"id"`
ExpID int64 `json:"exploration_id"`
Name string `json:"name"` // 可选任务名称;空=未命名
Description string `json:"description"`
Goal string `json:"goal"`
CreatedAt int64 `json:"created_at"`
CompletedAt int64 `json:"completed_at,omitempty"` // 进入终态的 unix 秒;0=未完成
Paused bool `json:"paused"`
Queued bool `json:"queued"` // 因并发上限被挂起、等待空位自动启动;true=尚未开跑
// QueuedAt is an internal Unix-nanosecond ordering key. It is deliberately
// finer than CreatedAt so several tasks enqueued in the same second retain
// their real FIFO order.
QueuedAt int64 `json:"queued_at,omitempty"`
QueueMode string `json:"queue_mode,omitempty"`
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"`
ActiveLLMProfileID *int64 `json:"active_llm_profile_id,omitempty"`
LLMChainRevision int64 `json:"-"`
LLMFailoverState string `json:"llm_failover_state,omitempty"`
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
SourceTaskIDs []int64 `json:"source_task_ids,omitempty"`
CompanyIDs []int64 `json:"company_ids,omitempty"`
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态;空/其它则由运行态推导)
// 任务级超时(见 docs/任务级超时与收尾设计.md)。DeadlineAt/FirstRunAt 为 unix 秒,0=未设/未运行。
TimeoutSeconds int `json:"timeout_seconds"`
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒)
CoverageEnabled bool `json:"coverage_enabled"` // 资产覆盖度功能开关(创建时定,默认开)
FirstRunAt int64 `json:"first_run_at,omitempty"`
DeadlineAt int64 `json:"deadline_at,omitempty"`
Store *pgdb.ExplorationStore `json:"-"`
Guard *guard.Guard `json:"-"`
notify chan struct{}
lifecycleMu sync.RWMutex
llmMu sync.RWMutex
>>>>>>> Stashed changes
// pendingTriggers accumulates the concrete changes (worker done / finding) that
// fired planning rounds since the last one consumed them. The debounce coalesces
@@ -457,11 +495,18 @@ func unixOrZero(t *time.Time) int64 {
func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Interceptor) *Task {
return &Task{
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
<<<<<<< Updated upstream
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused,
=======
Name: pt.Name,
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued,
QueuedAt: unixNanoOrZero(pt.QueuedAt), QueueMode: pt.QueueMode,
>>>>>>> Stashed changes
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
LLMProfileID: pt.LLMProfileID,
TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds,
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
CoverageEnabled: pt.CoverageEnabled,
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
}
}
@@ -734,6 +779,46 @@ func (t *Task) NotifyGoal(texts []string) {
t.Notify()
}
// NotifyGoalDeleted records that the human deleted a goal (via 总览的目标管理), then
// wakes the planner so the next round spells out which goal was removed. The event
// survives an early-returning terminal round (drain happens after the gate).
func (t *Task) NotifyGoalDeleted(text string) {
text = strings.TrimSpace(text)
if text == "" {
return
}
t.trigMu.Lock()
t.pendingTriggers = append(t.pendingTriggers, agent.TriggerEvent{Kind: "goal_deleted", Detail: text})
t.trigMu.Unlock()
t.Notify()
}
// NotifyGoalEdited records that the human edited a goal (via 总览的目标管理), then wakes
// the planner so the next round spells out the old→new change. The event survives an
// early-returning terminal round (drain happens after the gate).
func (t *Task) NotifyGoalEdited(oldText, newText string) {
oldText, newText = strings.TrimSpace(oldText), strings.TrimSpace(newText)
if newText == "" {
return
}
t.trigMu.Lock()
t.pendingTriggers = append(t.pendingTriggers, agent.TriggerEvent{Kind: "goal_edited", OldGoal: oldText, NewGoal: newText})
t.trigMu.Unlock()
t.Notify()
}
// NotifyCancelled records that the human deleted intentID (reason = 删除原因), then
// wakes the planner so the next round spells out which intent was removed and why.
// The intent is stopped (not deleted); the reason is attached to it as a fact.
func (t *Task) NotifyCancelled(intentID int64, reason string) {
if intentID > 0 {
t.trigMu.Lock()
t.pendingTriggers = append(t.pendingTriggers, agent.TriggerEvent{Kind: "cancelled", IntentID: intentID, Detail: reason})
t.trigMu.Unlock()
}
t.Notify()
}
// drainTriggers returns and clears the trigger events accumulated since the last round.
func (t *Task) drainTriggers() []agent.TriggerEvent {
t.trigMu.Lock()
+39 -4
View File
@@ -445,6 +445,7 @@ func (s *Server) seedOrchestrationTools() {
s.seedPlannerDefaultBindings()
s.seedAutoReportFindingBinding()
s.unbindGoalMetDefault()
s.reseedGoalsPrompt() // goals 提示词加入「抽操作约束」步 → 旧库追加一版新默认(一次性)
s.seedReporterAgent() // 预置「报告撰写」agent + 工具绑定 + finding 触发器(一次性)
// 注:pentest 的默认工具绑定无需迁移——BuiltinToolSeeds 在全新初始化时就把
// list_assets/insert_assets/report_finding/list_findings/list_companies 连同
@@ -457,7 +458,7 @@ func (s *Server) seedOrchestrationTools() {
// reaches an old DB otherwise. Preserves each tool's agent binding + enabled flag.
// Bump the flag whenever these tools' schemas/descriptions change in code.
func (s *Server) refreshBuiltinToolSchemas() {
const flag = "tool_schema_refresh_v5_spawn_seed_intent"
const flag = "tool_schema_refresh_v6_insert_assets_related"
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
return
}
@@ -468,10 +469,14 @@ func (s *Server) refreshBuiltinToolSchemas() {
log.Printf("[tools] refresh %s schema failed: %v", t.Name(), err)
}
}
// 同时把 planner 的 goal_met 描述刷成代码默认:旧库 seed 的描述带“结束本轮规划”的
// 误导,会让 planner 把 goal_met 当成“结束空轮”的手段、刚开跑就误判整个任务完成。
// 同时把部分内置 agent 工具刷成代码默认:
// - goal_met:旧库 seed 的描述带“结束本轮规划”的误导,会让 planner 把它当成
// “结束空轮”的手段、刚开跑就误判整个任务完成。
// - insert_assets:新增 related 入参(标记资产是否与当前任务相关、决定是否入覆盖度),
// SeedTool 首插入only,旧库已 seed 的 schema 否则收不到这个新参数。
refreshBuiltin := map[string]bool{"goal_met": true, "insert_assets": true}
for _, sd := range agent.BuiltinToolSeeds() {
if sd.Key != "goal_met" {
if !refreshBuiltin[sd.Key] {
continue
}
schema, _ := json.Marshal(sd.Schema)
@@ -500,6 +505,36 @@ func (s *Server) unbindGoalMetDefault() {
_ = s.m.pg.SetSetting(flag, "true")
}
// reseedGoalsPrompt 把 goals 目标拆解器的提示词刷成【当前代码默认】——因为默认正文新增了
// 「先抽操作约束(set_constraints)再拆目标」这一步,而 SeedPromptIfEmpty 首插入only,旧库
// 已有的 version 1 收不到这步。这里用版本管理【追加一个新版本】并切过去(ResetPromptToDefault),
// 旧的版本仍保留在历史里,用户若自定义过可从版本记录找回。settings flag 守卫 → 只做一次;
// 以后默认再变就 bump 这个 flag。全新库无需处理(SeedPromptIfEmpty 已 seed 最新默认)。
func (s *Server) reseedGoalsPrompt() {
const flag = "goals_prompt_constraint_step_v1"
if v, _, _ := s.m.pg.GetSetting(flag); v == "true" {
return
}
defer func() { _ = s.m.pg.SetSetting(flag, "true") }() // 无论成功与否只尝试一次
a, err := s.m.pg.GetAgentByKey("goals")
if err != nil || a == nil {
return // 全新库尚未建 agent 行时,seedPrompts 会直接 seed 最新默认,无需此迁移
}
tmpl := agent.BuiltinPromptSeeds()["goals"]
if tmpl == "" {
return
}
// 全新库 seedPrompts 已 seed 最新默认 → 当前版本已等于代码默认,不必再追加重复版本。
if cur, err := s.m.pg.CurrentPrompt(a.ID); err == nil && cur == tmpl {
return
}
if _, err := s.m.pg.ResetPromptToDefault(a.ID, tmpl); err != nil {
log.Printf("[prompts] goals 提示词重刷为新默认失败: %v", err)
return
}
log.Printf("[prompts] goals 提示词已追加新默认版本(加入抽操作约束步,一次性)")
}
// seedReporterAgent 预置一个「报告撰写」自定义 agent(builtin=false,可在 UI 编辑/删除):
// 绑定 update_finding_report + 任务查询工具,并挂一个「report_finding 被调用即触发」的
// 触发器 —— 每登记一个漏洞就唤起它写详细报告。一次性(settings flag 守卫):用户删掉后不再重建。
+58
View File
@@ -0,0 +1,58 @@
package server
import (
"testing"
"github.com/Autumn-27/artex/db"
)
// A stored catalog entry that collides with a global runtime var (e.g. a legacy
// seeded goals."Now") must not produce a duplicate: the global wins and the name
// appears exactly once, so the UI never sees duplicate React keys.
func TestWithGlobalVarsDedupesCollidingNames(t *testing.T) {
t.Parallel()
stored := []db.PromptVar{
{Name: "EngagementDescription", Description: "task", Source: "exploration"},
{Name: "Now", Description: "stale per-agent copy", Source: "runtime"},
}
out := withGlobalVars(stored)
counts := map[string]int{}
var now db.PromptVar
for _, v := range out {
counts[v.Name]++
if v.Name == "Now" {
now = v
}
}
if counts["Now"] != 1 {
t.Fatalf("Now appeared %d times, want 1: %+v", counts["Now"], out)
}
if counts["EngagementDescription"] != 1 {
t.Fatalf("non-colliding stored var was dropped or duplicated: %+v", out)
}
// The surviving Now must be the authoritative global definition, not the stale
// stored one.
if now.Description == "stale per-agent copy" {
t.Fatalf("stored var shadowed the global instead of the other way around: %+v", now)
}
// Every global is present exactly once.
for _, g := range globalPromptVars {
if counts[g.Name] != 1 {
t.Fatalf("global %q present %d times, want 1", g.Name, counts[g.Name])
}
}
}
// The common case (no collisions) keeps every stored var and appends all globals.
func TestWithGlobalVarsKeepsDistinctVars(t *testing.T) {
t.Parallel()
stored := []db.PromptVar{{Name: "Goal", Source: "exploration"}}
out := withGlobalVars(stored)
if len(out) != len(stored)+len(globalPromptVars) {
t.Fatalf("len=%d, want %d: %+v", len(out), len(stored)+len(globalPromptVars), out)
}
if out[0].Name != "Goal" {
t.Fatalf("stored var order not preserved: %+v", out)
}
}
+185 -3
View File
@@ -230,6 +230,11 @@ func New(ctx context.Context, m *Manager, skillDir string, dataDir string, keyDi
if t := m.ActiveTask(); t != nil {
s.engine.Run(ctx, t)
}
<<<<<<< Updated upstream
=======
go s.reconcileConcurrency()
s.wireInterceptReviewer() // LLM 兜底审批:未命中拦截规则的命令交给模型判定
>>>>>>> Stashed changes
return s
}
@@ -344,12 +349,14 @@ func (s *Server) buildPlannerWorker(pinID *int64, gProv llm.Provider, gCfg agent
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
wk.SetWebSearch(s.webSearchFor("worker"))
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开;每轮读)
pProv, pCfg := s.providerForAgent("planner", pinID, gProv, gCfg)
pl := agent.NewPlanner(pProv, pCfg.Model, s.m.dir, tx, pCfg.CompactionWindow(), s.agentMaxTurns("planner"))
pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work
pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
pl.SetKillWork(s.engine.KillWork) // planner kill_work → terminate a running work
pl.SetSteerWork(s.engine.SteerWork) // planner steer_work → inject mid-run course-correction
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
pl.SetWebSearch(s.webSearchFor("planner"))
pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开;每轮读)
return pl, wk
}
@@ -378,6 +385,7 @@ func (s *Server) applyLLM(cfg agent.Config) error {
s.mainAgent = agent.NewMainAgent(mProv, mCfg.Model, s.m.dir, tx, mCfg.CompactionWindow(), s.agentMaxTurns("mainagent"))
s.mainAgent.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert()) // WebFetch through the recording proxy
s.mainAgent.SetWebSearch(s.webSearchFor("mainagent"))
s.mainAgent.SetSteerWork(s.engine.SteerWork) // steer_work:人对运行中 work 实时纠偏
// chat agent serves MANY custom agents by key → it holds the GLOBAL opts
// (backend/key) and gates Enabled per-conversation-agent at Chat time. 对话始终用激活配置。
s.chatAgent = agent.NewChatAgent(prov, cfg.Model, s.m.dir, tx, win) // chat page runner
@@ -581,6 +589,19 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/workspace/download", s.wsDownload)
mux.HandleFunc("POST /api/workspace/upload", s.wsUpload)
mux.HandleFunc("GET /api/tasks/{id}/scope", s.taskScopeList)
<<<<<<< Updated upstream
=======
mux.HandleFunc("POST /api/tasks/{id}/scope", s.taskScopeAdd)
mux.HandleFunc("DELETE /api/tasks/{id}/scope/{sid}", s.taskScopeDelete)
mux.HandleFunc("GET /api/tasks/{id}/goals", s.listGoals) // 目标管理:列出本任务全部目标
mux.HandleFunc("POST /api/tasks/{id}/goals", s.addGoal) // 目标管理:人工新增目标(复活任务)
mux.HandleFunc("PATCH /api/tasks/{id}/goals/{gid}", s.editGoal) // 目标管理:修改目标(复活任务)
mux.HandleFunc("DELETE /api/tasks/{id}/goals/{gid}", s.deleteGoal) // 目标管理:硬删除目标(不复活)
mux.HandleFunc("GET /api/tasks/{id}/constraints", s.listConstraints) // 约束管理:列出本任务操作约束
mux.HandleFunc("POST /api/tasks/{id}/constraints", s.addConstraint) // 约束管理:新增约束(不通知 planner)
mux.HandleFunc("PATCH /api/tasks/{id}/constraints/{cid}", s.editConstraint) // 约束管理:修改约束
mux.HandleFunc("DELETE /api/tasks/{id}/constraints/{cid}", s.deleteConstraint) // 约束管理:删除约束
>>>>>>> Stashed changes
mux.HandleFunc("POST /api/tasks/{id}/control", s.control)
mux.HandleFunc("POST /api/tasks/{id}/intents/{iid}/rerun", s.rerunIntent) // 重跑单条 blocked/exhausted/stopped 意图
mux.HandleFunc("POST /api/tasks/{id}/intents/rerun-blocked", s.rerunBlocked) // 批量重跑本任务全部 blocked 意图
@@ -627,6 +648,7 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("DELETE /api/traffic", s.deleteTraffic)
mux.HandleFunc("DELETE /api/traffic/hosts", s.deleteTrafficHosts)
mux.HandleFunc("GET /api/traffic/exchange", s.getTrafficExchange)
mux.HandleFunc("GET /api/traffic/blob", s.getTrafficBlob)
mux.HandleFunc("GET /api/commands", s.pgListCommands)
mux.HandleFunc("GET /api/llm/records", s.pgListLLMRecords)
mux.HandleFunc("DELETE /api/llm/records", s.pgDeleteLLMRecords)
@@ -734,6 +756,8 @@ func (s *Server) Handler() http.Handler {
mux.HandleFunc("GET /api/intercept/task/{taskID}", s.interceptListTaskItems)
mux.HandleFunc("GET /api/intercept/tool-config", s.interceptGetToolConfig)
mux.HandleFunc("PUT /api/intercept/tool-config", s.interceptSetToolConfig)
mux.HandleFunc("GET /api/intercept/judge", s.interceptGetJudgeConfig)
mux.HandleFunc("PUT /api/intercept/judge", s.interceptSetJudgeConfig)
// /api/* goes through CORS + JWT; everything else is served by the embedded
// frontend (public — auth is enforced client-side and on the API). With the
@@ -915,8 +939,53 @@ func (s *Server) control(w http.ResponseWriter, r *http.Request) {
if err := s.m.SetTaskPaused(t.ID, paused); err != nil { // persist (survives restart)
log.Printf("[control] persist paused %s: %v", t.ID, err)
}
<<<<<<< Updated upstream
log.Printf("[task] #%s %s", t.ID, map[bool]string{true: "已暂停", false: "已恢复"}[paused])
writeJSON(w, 200, map[string]any{"id": t.ID, "paused": paused})
=======
writeJSON(w, 200, result)
}
// controlIntent pauses, resumes, or cancels one local worker intent. A running
// worker is always stopped before state mutation/cleanup, preventing tool output
// that arrives after the user action from recreating deleted blackboard records.
func (s *Server) controlIntent(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
if !s.engine.beginTaskOperation(t.ID) {
writeErr(w, 409, "任务正在删除,无法控制意图")
return
}
defer s.engine.decInflight(t.ID)
iid, err := strconv.ParseInt(r.PathValue("iid"), 10, 64)
if err != nil || iid <= 0 {
writeErr(w, 400, "bad intent id")
return
}
var req struct {
Action string `json:"action"`
Reason string `json:"reason"` // cancel(删除)时必填:删除原因,挂为事实并告知 planner
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "bad json: "+err.Error())
return
}
if req.Action != "pause" && req.Action != "resume" && req.Action != "cancel" {
writeErr(w, 400, "action must be pause|resume|cancel")
return
}
result, err := s.applyIntentControl(r.Context(), t, iid, req.Action, req.Reason)
if err != nil {
writeErr(w, 409, err.Error())
return
}
writeJSON(w, 200, result)
>>>>>>> Stashed changes
}
// rerunIntent 重跑一条没跑成功的意图(blocked/exhausted/stopped):把它置回 open,worker
@@ -1073,12 +1142,26 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
}
type createTaskReq struct {
<<<<<<< Updated upstream
Description string `json:"description"`
Goal string `json:"goal"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置
TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600
SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。
=======
Name string `json:"name,omitempty"` // 可选任务名称;省略/空=未命名
Description string `json:"description"`
Goal string `json:"goal"`
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务的 LLM 配置;省略/null=用激活配置
LLMProfileIDs []int64 `json:"llm_profile_ids,omitempty"` // 有序任务级配置链;第一项初始生效
SourceTaskIDs []string `json:"source_task_ids,omitempty"` // 仅直接、只读继承的来源任务
CompanyIDs []int64 `json:"company_ids,omitempty"` // 关联企业资产范围;不复制资产或强制生成意图
TimeoutSeconds int `json:"timeout_seconds"` // 任务级超时(秒);0/省略=不限时
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒);0/省略=默认600(10min);下限=默认=600,低于自动抬到600
SeedFirstIntent *bool `json:"seed_first_intent,omitempty"` // 创建时直接下发一条种子意图(内容=描述+目标),让 worker 免等首轮 planner 直接开跑;省略/null=默认关闭,走标准先规划再执行。显式传 true 才开(CTF 常一 work 解决时可省掉开跑前的 planner 轮)。
CoverageEnabled *bool `json:"coverage_enabled,omitempty"` // 资产覆盖度功能;省略/null=默认开(true)。false=关闭覆盖度计算/展示/自动累积范围+隐藏 add_task_scope/list_untested_assets。company 关联不受影响。
>>>>>>> Stashed changes
}
func (s *Server) createTask(w http.ResponseWriter, r *http.Request) {
@@ -1103,6 +1186,24 @@ func (s *Server) createTask(w http.ResponseWriter, r *http.Request) {
}
t, err := s.m.CreateTask(req.Description, req.Goal, req.LLMProfileID, req.TimeoutSeconds, req.PlanHeartbeatSeconds)
if err != nil {
<<<<<<< Updated upstream
=======
writeErr(w, 400, fmt.Sprintf("关联企业无效:最多选择 %d 个有效企业", db.MaxTaskCompanyCount))
return
}
req.CompanyIDs = companyIDs
t, err := s.m.CreateTaskWithOptions(req.Description, req.Goal, db.TaskCreateOptions{
Name: strings.TrimSpace(req.Name),
SourceTaskIDs: sourceIDs, CompanyIDs: req.CompanyIDs, LLMProfileIDs: req.LLMProfileIDs,
TimeoutSeconds: req.TimeoutSeconds, PlanHeartbeatSeconds: req.PlanHeartbeatSeconds,
CoverageEnabled: req.CoverageEnabled,
})
if err != nil {
if errors.Is(err, db.ErrTaskCompanyIDsInvalid) || errors.Is(err, db.ErrTaskCompanyNotFound) {
writeErr(w, 400, "关联企业不存在或无效")
return
}
>>>>>>> Stashed changes
writeErr(w, 500, err.Error())
return
}
@@ -1258,12 +1359,18 @@ func (s *Server) taskCoverage(w http.ResponseWriter, r *http.Request) {
writeErr(w, 503, "asset store 未启用")
return
}
// 资产覆盖度功能关闭 → 短路返回 {enabled:false},前端据此隐藏覆盖度卡片/进度。
if !t.CoverageEnabled {
writeJSON(w, 200, &db.Coverage{Enabled: false, ByType: []db.CoverageByType{}})
return
}
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
cov, err := as.TaskCoverage(taskID, t.ExpID)
if err != nil {
writeErr(w, 500, err.Error())
return
}
cov.Enabled = true
writeJSON(w, 200, cov)
}
@@ -2135,6 +2242,35 @@ func (s *Server) getTrafficExchange(w http.ResponseWriter, r *http.Request) {
writeJSON(w, 200, map[string]any{"req": req, "resp": resp})
}
// getTrafficBlob streams one oversized body by its sha256. Bodies past the inline
// threshold are not carried by the exchange endpoint — it returns a preview plus
// an "@blob sha256:<hash>" pointer — so this is how the UI fetches them whole.
// Streamed rather than buffered: these are the bodies too large to hold in memory.
func (s *Server) getTrafficBlob(w http.ResponseWriter, r *http.Request) {
tr := s.m.Traffic()
if tr == nil {
writeErr(w, 404, "traffic disabled")
return
}
hash := strings.TrimSpace(r.URL.Query().Get("hash"))
if hash == "" {
writeErr(w, 400, "missing hash")
return
}
f, size, err := tr.Blob(hash)
if err != nil {
writeErr(w, 404, err.Error())
return
}
defer f.Close()
w.Header().Set("Content-Type", "application/octet-stream")
w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", hash+".bin"))
if _, err := io.Copy(w, f); err != nil {
log.Printf("[traffic] 下载 blob %s 中断:%v", hash, err)
}
}
// getSettings returns the runtime app settings the UI toggles. The Brave API key
// is returned as a boolean presence flag (brave_key_set), never the value itself,
// so the UI can show "configured" without echoing the secret back.
@@ -2146,6 +2282,7 @@ func (s *Server) settingsPayload() map[string]any {
on, backend, braveKey, tavilyKey, proxy := s.m.WebSearch()
pyStored, _, _ := s.m.pg.GetSetting(settingPythonInterp)
return map[string]any{
<<<<<<< Updated upstream
"traffic_capture": s.m.TrafficEnabled(),
"llm_record": s.m.LLMRecordEnabled(),
"web_search_enabled": on,
@@ -2155,6 +2292,26 @@ func (s *Server) settingsPayload() map[string]any {
"web_search_proxy": proxy, // 独立出口代理(http/https/socks5),空=直连
"python_interpreter": strings.TrimSpace(pyStored), // 用户/自动设的值(空=用运行时检测)
"workers": s.m.Workers(), // 并发工作 agent 数(默认3);对之后启动的任务生效
=======
"traffic_capture": s.m.TrafficEnabled(),
"llm_record": s.m.LLMRecordEnabled(),
"web_search_enabled": on,
"web_search_backend": backend,
"brave_key_set": strings.TrimSpace(braveKey) != "",
"tavily_key_set": strings.TrimSpace(tavilyKey) != "",
"web_search_proxy": proxy, // 独立出口代理(http/https/socks5),空=直连
"python_interpreter": strings.TrimSpace(pyStored), // 用户/自动设的值(空=用运行时检测)
"workers": s.m.Workers(), // 并发工作 agent 数(默认3);对之后启动的任务生效
"task_concurrency_enabled": concOn, // 任务并发上限开关(默认关)
"task_concurrency_limit": concLimit, // 同时运行任务上限(开启后默认5)
// LLM 轮询(故障转移)。默认关;开启后走全局激活配置的 agent 在当前配置不可用时
// 自动切到下一个配置。bind_fallback 仅在轮询开启时有意义(默认关)。
"llm_pool_enabled": s.m.LLMPoolEnabled(),
"llm_pool_bind_fallback": s.m.LLMPoolBindFallback(),
// 操作约束注入范围(默认都开):把本任务的 allow/deny 约束拼进对应 agent 的系统提示。
"constraints_inject_planner": s.constraintInjectPlanner(),
"constraints_inject_worker": s.constraintInjectWorker(),
>>>>>>> Stashed changes
}
}
@@ -2188,11 +2345,36 @@ func (s *Server) putSettings(w http.ResponseWriter, r *http.Request) {
WebSearchProxy *string `json:"web_search_proxy"` // 独立出口代理(http/https/socks5);null=不改,""=清空
PythonInterp *string `json:"python_interpreter"` // 自定义脚本工具的 python 解释器路径
Workers *int `json:"workers"` // 并发工作 agent 数(>0);对之后启动的任务生效
<<<<<<< Updated upstream
=======
// 任务并发上限:同时「运行中」的任务数上限。关闭=不限;开启后新建任务超限则排队,有空位自动启动。
ConcurrencyEnabled *bool `json:"task_concurrency_enabled"`
ConcurrencyLimit *int `json:"task_concurrency_limit"`
// LLM 轮询(故障转移)开关 + 「绑定配置失败也兜底回轮询链」开关。两者都需要
// 重建 provider 链才生效,走下面的 changed → applyLLM 路径。
LLMPoolEnabled *bool `json:"llm_pool_enabled"`
LLMPoolBindFallback *bool `json:"llm_pool_bind_fallback"`
// 操作约束注入范围开关(默认都开);即时生效(planner/worker 每轮读),无需重建 agent。
ConstraintsInjectPlanner *bool `json:"constraints_inject_planner"`
ConstraintsInjectWorker *bool `json:"constraints_inject_worker"`
>>>>>>> Stashed changes
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, err.Error())
return
}
if req.ConstraintsInjectPlanner != nil {
if err := s.m.pg.SetBool(settingConstraintsInjectPlanner, *req.ConstraintsInjectPlanner); err != nil {
writeErr(w, 500, err.Error())
return
}
}
if req.ConstraintsInjectWorker != nil {
if err := s.m.pg.SetBool(settingConstraintsInjectWorker, *req.ConstraintsInjectWorker); err != nil {
writeErr(w, 500, err.Error())
return
}
}
if req.Workers != nil {
if err := s.m.SetWorkers(*req.Workers); err != nil {
writeErr(w, 400, err.Error())
+15 -1
View File
@@ -1754,8 +1754,22 @@ var globalPromptVars = []db.PromptVar{
// withGlobalVars appends the universal runtime vars onto an agent's own catalog,
// so validation / the UI variable list / preview all recognize {{.Now}} etc.
// A stored catalog entry that collides with a global name is dropped: the global
// runtime var is authoritative (it's what the render path actually resolves), and
// this keeps the returned list name-unique so the UI never sees duplicate keys.
func withGlobalVars(vars []db.PromptVar) []db.PromptVar {
return append(append([]db.PromptVar{}, vars...), globalPromptVars...)
globalNames := make(map[string]bool, len(globalPromptVars))
for _, g := range globalPromptVars {
globalNames[g.Name] = true
}
out := make([]db.PromptVar, 0, len(vars)+len(globalPromptVars))
for _, v := range vars {
if globalNames[v.Name] {
continue // shadowed by the authoritative global runtime var
}
out = append(out, v)
}
return append(out, globalPromptVars...)
}
// validateTemplate parses the template and rejects any {{.Var}} not in the catalog.
+822
View File
@@ -0,0 +1,822 @@
package server
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
)
func newAdmissionTestServer(m *Manager, available func(*Task) bool) *Server {
ctx, cancel := context.WithCancel(context.Background())
cancel()
e := NewEngine(m)
planner, worker := new(agent.Planner), new(agent.Worker)
e.SetAuthoritativeAgentResolver(func(task *Task) (*agent.Planner, *agent.Worker) {
if available == nil || available(task) {
return planner, worker
}
return nil, nil
})
return &Server{m: m, engine: e, ctx: ctx, chatBusy: map[string]bool{}, chatCancel: map[string]context.CancelCauseFunc{}}
}
func restoreConcurrencySetting(t *testing.T, m *Manager) func() {
t.Helper()
enabled, limit := m.ConcurrencyLimit()
if limit < 1 {
limit = defaultConcurrencyLimit
}
return func() {
if err := m.SetConcurrency(enabled, limit); err != nil {
t.Errorf("restore concurrency setting: %v", err)
}
}
}
func TestAdmitAlreadyRunningTaskDoesNotQueueAfterLimitDecrease(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
oldEnabled, oldLimit := m.ConcurrencyLimit()
if oldLimit < 1 {
oldLimit = defaultConcurrencyLimit
}
defer m.SetConcurrency(oldEnabled, oldLimit)
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
running, err := m.CreateTask("already admitted", "accept follow-up", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("other admitted task", "occupy lowered limit", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(running.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
// Simulate a task whose long-lived loops were admitted before the limit was
// lowered. The other live task makes the current count exceed the new limit.
s.engine.started.Store(running.ID, true)
queued, err := s.admitTask(running, "resume")
if err != nil {
t.Fatal(err)
}
if queued || running.Queued {
t.Fatalf("already-running task was re-queued: queued=%v task=%+v", queued, running)
}
}
func TestAdmissionDoesNotOverwriteConcurrentTerminalStatus(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(false, defaultConcurrencyLimit); err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("stale admission", "preserve terminal writer", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
// Model an Engine terminal commit after admission captured its in-memory
// snapshot but before the scheduler UPDATE. Direct DB use intentionally leaves
// the live handle stale at its old status for this deterministic interleaving.
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
if err := m.pg.SetStatus(taskID, "done"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
if _, err := s.admitTask(task, "resume"); err == nil {
t.Fatal("stale admission overwrote a concurrent terminal transition")
}
stored, err := m.pg.GetTask(taskID)
if err != nil {
t.Fatal(err)
}
if stored == nil || stored.Status != "done" || stored.Queued || stored.Paused {
t.Fatalf("terminal state changed by stale admission: %+v", stored)
}
}
func TestAdmitKeepsQueuedBootstrapMode(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
oldEnabled, oldLimit := m.ConcurrencyLimit()
if oldLimit < 1 {
oldLimit = defaultConcurrencyLimit
}
defer m.SetConcurrency(oldEnabled, oldLimit)
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("queue holder", "occupy slot", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
queuedTask, err := m.CreateTask("queued bootstrap", "must decompose goals", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(queuedTask.ID, DeleteTaskOptions{})
}()
if err := m.EnqueueTask(queuedTask.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
firstQueuedAt := queuedTask.QueuedAt
s := newAdmissionTestServer(m, nil)
queued, err := s.admitTask(queuedTask, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || queuedTask.QueueMode != "bootstrap" || queuedTask.QueuedAt != firstQueuedAt {
t.Fatalf("re-admission changed bootstrap FIFO state: queued=%v task=%+v", queued, queuedTask)
}
}
func TestTerminalTaskQueuedByAdmissionKeepsExecutionBarrier(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("terminal follow-up", "must wait", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
intentID, err := task.Store.AddIntent(map[string]any{"summary": "follow-up"}, 10, nil, "human")
if err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
queued, err := s.admitTask(task, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || !task.Queued || task.Status != "running" || !s.engine.IsPaused(task.ID) {
t.Fatalf("queued terminal admission lost gate: queued=%v task=%+v enginePaused=%v", queued, task, s.engine.IsPaused(task.ID))
}
if execCtx := s.engine.execContextFor(context.Background(), task.ID); execCtx.Err() == nil {
t.Fatal("queued terminal task received a live execution context")
}
node, err := task.Store.GetNode(intentID)
if err != nil || node == nil || node.State != "open" {
t.Fatalf("queued intent changed before admission: node=%+v err=%v", node, err)
}
}
func TestTimedOutTaskRevivalResetsClockAndSettlingAcrossFIFO(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("timeout revival holder", "occupy slot", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("timed out follow-up", "run with a fresh clock", nil, 120, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
if _, err := m.StampTaskFirstRun(task.ID); err != nil {
t.Fatal(err)
}
oldClock := task.lifecycleSnapshot()
if oldClock.FirstRunAt == 0 || oldClock.DeadlineAt == 0 {
t.Fatalf("test setup did not stamp the original clock: %+v", oldClock)
}
if err := m.SetTaskStatus(task.ID, "timeout"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
s.engine.settling.Store(task.ID, true)
s.engine.deadline.Store(task.ID, oldClock.DeadlineAt)
s.engine.stamped.Store(task.ID, true)
s.engine.coordStarted.Store(task.ID, true)
queued, err := s.admitTask(task, "resume")
if err != nil {
t.Fatal(err)
}
queuedState := task.lifecycleSnapshot()
if !queued || !queuedState.Queued || queuedState.Status != "running" ||
queuedState.FirstRunAt != 0 || queuedState.DeadlineAt != 0 {
t.Fatalf("timed-out task was not reset before FIFO wait: queued=%v state=%+v", queued, queuedState)
}
if s.engine.isSettling(task.ID) {
t.Fatal("timed-out task retained the old settling gate while queued")
}
if _, ok := s.engine.deadline.Load(task.ID); ok {
t.Fatal("timed-out task retained the old in-memory deadline")
}
if _, ok := s.engine.stamped.Load(task.ID); ok {
t.Fatal("timed-out task retained the old first-run stamp")
}
stored, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if stored == nil || stored.TimeoutSeconds != 120 || stored.FirstRunAt != nil || stored.DeadlineAt != nil {
t.Fatalf("persisted timeout revival clock/config mismatch: %+v", stored)
}
// Release the only slot and promote the revived task. The test server uses a
// cancelled root context so no real Agent call races these state assertions.
s.engine.Pause(holder.ID, agent.AbortPausedByUser)
if err := m.ApplyTaskPause(holder.ID); err != nil {
t.Fatal(err)
}
s.reconcileConcurrency()
runningState := task.lifecycleSnapshot()
if runningState.Queued || runningState.Paused || s.engine.IsPaused(task.ID) || s.engine.isSettling(task.ID) {
t.Fatalf("revived task did not leave FIFO ready to execute: state=%+v enginePaused=%v settling=%v",
runningState, s.engine.IsPaused(task.ID), s.engine.isSettling(task.ID))
}
if execCtx := s.engine.execContextFor(context.Background(), task.ID); execCtx.Err() != nil {
t.Fatalf("revived task did not receive a live execution context: %v", context.Cause(execCtx))
}
if _, ok := s.engine.coordStarted.Load(task.ID); !ok {
t.Fatal("revived timeout task did not start a fresh deadline coordinator")
}
}
func TestQueuedTaskCanPauseAndResumeAtFIFOTail(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTask("pause queued", "resume later", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
if err := m.SetTaskPaused(task.ID, true); err != nil {
t.Fatal(err)
}
s.engine.Pause(task.ID, agent.AbortPausedByUser)
first, err := s.applyTaskControl(task, "resume")
if err != nil || !first.Queued || task.QueuedAt == 0 {
t.Fatalf("first resume: result=%+v task=%+v err=%v", first, task, err)
}
firstQueuedAt := task.QueuedAt
if _, err := s.applyTaskControl(task, "pause"); err != nil {
t.Fatal(err)
}
if !task.Paused || task.Queued || task.QueuedAt != 0 {
t.Fatalf("pause did not leave queue: %+v", task)
}
time.Sleep(2 * time.Millisecond)
second, err := s.applyTaskControl(task, "resume")
if err != nil || !second.Queued {
t.Fatalf("second resume: result=%+v task=%+v err=%v", second, task, err)
}
if task.QueuedAt <= firstQueuedAt {
t.Fatalf("resumed task did not move to FIFO tail: first=%d second=%d", firstQueuedAt, task.QueuedAt)
}
}
func TestReadyFIFOIsAdmittedBeforeNewTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 2); err != nil {
t.Fatal(err)
}
holder, err := m.CreateTask("slot holder", "hold", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
older, err := m.CreateTask("older queued", "first", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
incoming, err := m.CreateTask("new admission", "second", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(holder.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(older.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(incoming.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, nil)
if err := m.EnqueueTask(older.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(older.ID, agent.AbortPausedByOrchestrator)
queued, err := s.admitTask(incoming, "resume")
if err != nil {
t.Fatal(err)
}
if !queued || !incoming.Queued || incoming.QueuedAt <= older.QueuedAt {
t.Fatalf("new task bypassed ready FIFO: queued=%v older=%+v incoming=%+v", queued, older, incoming)
}
s.reconcileConcurrency()
if older.Queued || !incoming.Queued {
t.Fatalf("FIFO promotion order wrong: older=%+v incoming=%+v", older, incoming)
}
}
func TestUnavailableTaskReleasesSlotForReadyQueue(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
unavailable, err := m.CreateTask("unavailable llm", "park", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
ready, err := m.CreateTask("ready queued", "run", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(unavailable.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(ready.ID, DeleteTaskOptions{})
}()
s := newAdmissionTestServer(m, func(task *Task) bool { return task.ID != unavailable.ID })
s.engine.started.Store(unavailable.ID, true)
if err := m.EnqueueTask(ready.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(ready.ID, agent.AbortPausedByOrchestrator)
s.reconcileConcurrency()
if !unavailable.Queued || !s.engine.IsPaused(unavailable.ID) {
t.Fatalf("unavailable task still occupies slot: %+v", unavailable)
}
if ready.Queued || s.engine.IsPaused(ready.ID) {
t.Fatalf("ready task was not promoted: %+v", ready)
}
}
func TestUnavailableTaskWaitsForActiveLLMCallBeforeParking(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
defer restoreConcurrencySetting(t, m)()
if err := m.SetConcurrency(true, 1); err != nil {
t.Fatal(err)
}
active, err := m.CreateTask("active profile edit", "finish current call", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
waiting, err := m.CreateTask("waiting task", "respect active call", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(active.ID, DeleteTaskOptions{})
_, _ = m.DeleteTask(waiting.ID, DeleteTaskOptions{})
}()
available := map[string]bool{active.ID: true, waiting.ID: true}
s := newAdmissionTestServer(m, func(task *Task) bool { return available[task.ID] })
s.engine.started.Store(active.ID, true)
if err := m.EnqueueTask(waiting.ID, "resume"); err != nil {
t.Fatal(err)
}
s.engine.Pause(waiting.ID, agent.AbortPausedByOrchestrator)
// Model a manual chain edit that makes the task's next-call resolution
// unavailable while its old provider is still serving the current call.
s.engine.BeginLLMCall(active.ID)
available[active.ID] = false
s.reconcileConcurrency()
if active.Queued || s.engine.IsPaused(active.ID) {
t.Fatalf("active LLM call was cancelled/parked after profile edit: %+v", active)
}
if !waiting.Queued {
t.Fatalf("waiting task bypassed active provider call: %+v", waiting)
}
s.engine.EndLLMCall(active.ID)
s.reconcileConcurrency()
if !active.Queued || !s.engine.IsPaused(active.ID) {
t.Fatalf("unavailable task was not parked after its call ended: %+v", active)
}
if waiting.Queued || s.engine.IsPaused(waiting.ID) {
t.Fatalf("ready FIFO task was not admitted after active call ended: %+v", waiting)
}
}
func TestRerunAdmissionFailureRestoresIntentState(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("failed rerun admission", "restore intent", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
taskID, _ := strconv.ParseInt(task.ID, 10, 64)
defer func() {
_, _ = m.pg.Exec(`UPDATE tasks SET deleted_at=NULL WHERE id=$1`, taskID)
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
}()
intentID, err := task.Store.AddIntent(map[string]any{"summary": "retry"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
if err := task.Store.SetIntentBlockedReason(intentID, "network_unavailable"); err != nil {
t.Fatal(err)
}
if err := m.SetTaskStatus(task.ID, "done"); err != nil {
t.Fatal(err)
}
if _, err := m.pg.Exec(`UPDATE tasks SET deleted_at=now() WHERE id=$1`, taskID); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
req := httptest.NewRequest(http.MethodPost, "/api/tasks/1/intents/1/rerun", nil)
req.SetPathValue("id", task.ID)
req.SetPathValue("iid", fmt.Sprintf("%d", intentID))
rec := httptest.NewRecorder()
s.rerunIntent(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
node, err := task.Store.GetNode(intentID)
if err != nil || node == nil || node.State != "blocked" || node.BlockedReason != "network_unavailable" {
t.Fatalf("rerun rollback: node=%+v err=%v", node, err)
}
}
func TestTaskLLMResolutionRejectsInvalidExplicitProfile(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
profileID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("invalid-resolution-%d", time.Now().UnixNano()),
Format: "openai",
Model: "invalid-model",
APIKey: "test-key",
Proxy: "invalid-proxy-without-scheme",
})
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("invalid resolution", "report unavailable", db.TaskCreateOptions{
LLMProfileIDs: []int64{profileID},
})
if err != nil {
_ = m.pg.DeleteProfile(profileID)
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
_ = m.pg.DeleteProfile(profileID)
}()
s := &Server{m: m, provByProfile: map[int64]*provEntry{}}
got, err := s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if got.Available || got.Source != "task_chain" || got.ProfileID == nil || *got.ProfileID != profileID {
t.Fatalf("resolution=%+v, want unavailable explicit task profile", got)
}
}
// An Agent binding outranks the task's LLM chain; roles left unbound keep using
// the chain, so binding one role does not move the others off it.
func TestTaskLLMResolutionPrefersAgentBindingOverTaskChain(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
stamp := time.Now().UnixNano()
boundID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("bound-resolution-%d", stamp), Format: "openai",
Model: "bound-model", APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
chainID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("chain-resolution-%d", stamp), Format: "openai",
Model: "chain-model", APIKey: "test-key",
})
if err != nil {
_ = m.pg.DeleteProfile(boundID)
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("binding precedence", "prefer the bound profile", db.TaskCreateOptions{
LLMProfileIDs: []int64{chainID},
})
if err != nil {
_ = m.pg.DeleteProfile(boundID)
_ = m.pg.DeleteProfile(chainID)
t.Fatal(err)
}
defer func() {
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
_ = m.pg.SetAgentLLMProfile("worker", nil)
_ = m.pg.DeleteProfile(boundID)
_ = m.pg.DeleteProfile(chainID)
}()
s := &Server{m: m, provByProfile: map[int64]*provEntry{}}
got, err := s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if got.Source != "task_chain" || got.Model != "chain-model" {
t.Fatalf("resolution=%+v, want the task chain while no binding exists", got)
}
if err := m.pg.SetAgentLLMProfile("worker", &boundID); err != nil {
t.Fatal(err)
}
got, err = s.resolveTaskRoleLLM(task, "worker")
if err != nil {
t.Fatal(err)
}
if !got.Available || got.Source != "agent_binding" || got.Model != "bound-model" {
t.Fatalf("resolution=%+v, want the agent binding to outrank the task chain", got)
}
got, err = s.resolveTaskRoleLLM(task, "planner")
if err != nil {
t.Fatal(err)
}
if got.Source != "task_chain" || got.Model != "chain-model" {
t.Fatalf("resolution=%+v, want an unbound role to keep using the task chain", got)
}
}
func TestTaskLLMResolutionReportsDatabaseFailure(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("resolution storage failure", "return server error", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
if err := m.pg.Close(); err != nil {
t.Fatal(err)
}
s := &Server{m: m}
req := httptest.NewRequest(http.MethodGet, "/api/tasks/1/llm/resolution", nil)
req.SetPathValue("id", task.ID)
rec := httptest.NewRecorder()
s.taskLLMResolutionHandler(rec, req)
if rec.Code != http.StatusInternalServerError {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestTokenStatsWithoutTaskReturnsStableShape(t *testing.T) {
s := &Server{m: &Manager{tasks: map[string]*Task{}}}
req := httptest.NewRequest(http.MethodGet, "/api/exploration/tokens?task=missing", nil)
rec := httptest.NewRecorder()
s.tokenStats(rec, req)
if rec.Code != http.StatusOK {
t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String())
}
var payload struct {
Workers []db.TokenUsage `json:"workers"`
Sessions []db.SessionTokenUsage `json:"sessions"`
Total TokenTotalDTO `json:"total"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &payload); err != nil {
t.Fatal(err)
}
if payload.Workers == nil || payload.Sessions == nil {
t.Fatalf("empty collections must serialize as arrays: %s", rec.Body.String())
}
}
func TestBatchControlLimitCountsDeduplicatedIDs(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
s := newAdmissionTestServer(m, nil)
request := func(handler http.HandlerFunc, path string, pathValues map[string]string, payload any) *httptest.ResponseRecorder {
body, marshalErr := json.Marshal(payload)
if marshalErr != nil {
t.Fatal(marshalErr)
}
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(body))
for key, value := range pathValues {
req.SetPathValue(key, value)
}
rec := httptest.NewRecorder()
handler(rec, req)
return rec
}
rec := request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": []string{" 0009223372036854775807 ", "9223372036854775807", " bad-id ", "bad-id"},
"action": "pause",
})
if rec.Code != http.StatusOK {
t.Fatalf("normalized task ids status=%d body=%s", rec.Code, rec.Body.String())
}
var normalizedPayload struct {
Items []batchControlItem `json:"items"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &normalizedPayload); err != nil {
t.Fatal(err)
}
if len(normalizedPayload.Items) != 2 || normalizedPayload.Items[0].ID != "9223372036854775807" ||
normalizedPayload.Items[1].ID != "bad-id" || normalizedPayload.Items[1].Error != "bad task id" {
t.Fatalf("normalized task response=%+v", normalizedPayload.Items)
}
repeatedTasks := make([]string, maxBatchControlIDs+1)
for i := range repeatedTasks {
repeatedTasks[i] = "missing"
}
rec = request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": repeatedTasks,
"action": "pause",
})
if rec.Code != http.StatusOK {
t.Fatalf("deduplicated task ids status=%d body=%s", rec.Code, rec.Body.String())
}
var taskPayload struct {
Items []batchControlItem `json:"items"`
}
if err := json.Unmarshal(rec.Body.Bytes(), &taskPayload); err != nil {
t.Fatal(err)
}
if len(taskPayload.Items) != 1 || taskPayload.Items[0].ID != "missing" {
t.Fatalf("deduplicated task response=%+v", taskPayload.Items)
}
uniqueTasks := make([]string, maxBatchControlIDs+1)
for i := range uniqueTasks {
uniqueTasks[i] = fmt.Sprintf("missing-%d", i)
}
rec = request(s.controlTasksBatch, "/api/tasks/control/batch", nil, map[string]any{
"task_ids": uniqueTasks,
"action": "pause",
})
if rec.Code != http.StatusBadRequest {
t.Fatalf("unique task ids status=%d body=%s", rec.Code, rec.Body.String())
}
}
func TestPauseTaskToolPersistsAndDequeuesTask(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("orchestrated pause", "persist queued pause", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
if err := m.EnqueueTask(task.ID, "bootstrap"); err != nil {
t.Fatal(err)
}
s := newAdmissionTestServer(m, nil)
s.engine.Pause(task.ID, agent.Causef("queued_for_admission", "queued", "test queue barrier"))
input, err := json.Marshal(map[string]string{"task_id": task.ID})
if err != nil {
t.Fatal(err)
}
if _, err := s.toolPauseTask().Call(context.Background(), input, nil); err != nil {
t.Fatal(err)
}
if !task.Paused || task.Queued || task.QueueMode != "bootstrap" || !s.engine.IsPaused(task.ID) {
t.Fatalf("tool pause did not persist/dequeue task: %+v enginePaused=%v", task, s.engine.IsPaused(task.ID))
}
stored, err := m.pg.GetTask(mustTaskID(t, task.ID))
if err != nil {
t.Fatal(err)
}
if stored == nil || !stored.Paused || stored.Queued {
t.Fatalf("stored task=%+v, want paused and dequeued", stored)
}
}
func TestPauseTaskToolUsesOrchestratorCancellationCause(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
task, err := m.CreateTask("orchestrated pause cause", "keep audit reason", nil, 0, 0)
if err != nil {
t.Fatal(err)
}
defer func() { _, _ = m.DeleteTask(task.ID, DeleteTaskOptions{}) }()
s := newAdmissionTestServer(m, nil)
execCtx := s.engine.execContextFor(context.Background(), task.ID)
input, err := json.Marshal(map[string]string{"task_id": task.ID})
if err != nil {
t.Fatal(err)
}
if _, err := s.toolPauseTask().Call(context.Background(), input, nil); err != nil {
t.Fatal(err)
}
if code, _, _, ok := agent.AbortReason(execCtx); !ok || code != "paused_by_orchestrator" {
t.Fatalf("pause_task cancellation code=%q ok=%v, want paused_by_orchestrator", code, ok)
}
}
+246
View File
@@ -0,0 +1,246 @@
package server
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"strings"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
)
const maxBatchControlIDs = 100
type taskControlResult struct {
ID string `json:"id"`
Paused bool `json:"paused"`
Queued bool `json:"queued"`
Status string `json:"status"`
}
type intentControlResult struct {
ID int64 `json:"id"`
State string `json:"state"`
Deleted *db.IntentCleanup `json:"deleted,omitempty"`
}
type batchControlItem struct {
ID string `json:"id"`
OK bool `json:"ok"`
Status string `json:"status,omitempty"`
Queued bool `json:"queued,omitempty"`
Error string `json:"error,omitempty"`
}
func (s *Server) resumeAdmissionMode(t *Task) string {
if t == nil {
return "resume"
}
lifecycle := t.lifecycleSnapshot()
if lifecycle.QueueMode == "bootstrap" {
return "bootstrap"
}
if lifecycle.FirstRunAt == 0 {
goals, err := t.Store.ListByKind(db.KindGoal, 1)
if err == nil && len(goals) == 0 {
return "bootstrap"
}
}
return "resume"
}
// applyTaskControl is shared by the single and batch endpoints. Task resume is
// deliberately limited to paused tasks; reruns and finding follow-ups use
// admitTask directly when they need to revive a terminal task.
func (s *Server) applyTaskControl(t *Task, action string) (taskControlResult, error) {
return s.applyTaskControlWithCause(t, action, agent.AbortPausedByUser)
}
func (s *Server) applyTaskControlWithCause(t *Task, action string, pauseCause error) (taskControlResult, error) {
if t == nil {
return taskControlResult{}, fmt.Errorf("task not found")
}
out := taskControlResult{ID: t.ID}
switch action {
case "pause":
s.concMu.Lock()
defer s.concMu.Unlock()
current, exists := s.m.Task(t.ID)
if !exists || current != t || s.engine.IsDeleting(t.ID) {
return out, fmt.Errorf("任务正在删除,无法控制")
}
if !s.engine.beginTaskOperation(t.ID) {
return out, fmt.Errorf("任务正在删除,无法控制")
}
defer s.engine.decInflight(t.ID)
lifecycle := t.lifecycleSnapshot()
if isTerminalStatus(lifecycle.Status) {
return out, fmt.Errorf("终态任务不能执行暂停")
}
if lifecycle.Paused {
return out, fmt.Errorf("任务已经暂停")
}
wasQueued := lifecycle.Queued
wasEnginePaused := s.engine.IsPaused(t.ID)
if pauseCause == nil {
pauseCause = agent.AbortPausedByUser
}
s.engine.Pause(t.ID, pauseCause)
if err := s.m.ApplyTaskPause(t.ID); err != nil {
if !wasEnginePaused && !wasQueued {
s.engine.Resume(t)
}
return out, err
}
// Main Agent is independently cancellable. Only cancel its current turn
// after the persistent pause commits, so a failed control request is fully
// compensated and does not lose an otherwise valid conversation turn.
s.cancelTaskChat(t.ID, agent.AbortChatPausedWithTask)
out.Paused, out.Status = true, "paused"
go s.reconcileConcurrency()
case "resume":
queued, err := s.admitPausedTask(t)
if err != nil {
return out, err
}
out.Queued = queued
out.Status = map[bool]string{true: "queued", false: "running"}[queued]
default:
return out, fmt.Errorf("action must be pause|resume")
}
log.Printf("[task] #%s %s", t.ID, map[string]string{"pause": "已暂停", "resume": "已继续"}[action])
return out, nil
}
func (s *Server) applyIntentControl(ctx context.Context, t *Task, iid int64, action, reason string) (intentControlResult, error) {
out := intentControlResult{ID: iid}
node, err := t.Store.GetNode(iid)
if err != nil {
return out, err
}
if node == nil {
if inherited, sourceErr := t.Store.GetNodeWithSources(iid); sourceErr == nil && inherited != nil && inherited.Inherited {
return out, fmt.Errorf("继承意图为只读,不能控制")
}
return out, fmt.Errorf("intent not found")
}
if node.Kind != db.KindIntent {
return out, fmt.Errorf("node is not an intent")
}
switch action {
case "pause":
if node.State != "running" {
return out, fmt.Errorf("仅运行中的意图可以暂停")
}
if err := s.engine.ControlWork(ctx, iid, "pause"); err != nil {
return out, err
}
out.State = "paused"
case "resume":
if node.State != "paused" {
return out, fmt.Errorf("仅已暂停的意图可以恢复")
}
changed, err := t.Store.CompareAndSetIntentState(iid, "paused", "open")
if err != nil {
return out, err
}
if !changed {
return out, fmt.Errorf("%w: 意图不再是 paused 状态", db.ErrIntentStateConflict)
}
t.Notify()
out.State = "open"
case "cancel":
// 「删除」新语义:不销毁意图与产出,而是把意图停到 stopped、把用户填写的删除原因
// 作为一条事实挂到该意图上,并用 cancelled 触发告知 planner(意图内容 + 删除原因)。
if node.State != "running" && node.State != "paused" {
return out, fmt.Errorf("仅运行中或已暂停的意图可以删除")
}
reason = strings.TrimSpace(reason)
if reason == "" {
return out, fmt.Errorf("请填写删除原因")
}
if node.State == "running" {
if err := s.engine.ControlWork(ctx, iid, "cancel"); err != nil {
return out, err
}
}
if _, err := t.Store.StopIntentWithReason(iid, reason, "user"); err != nil {
return out, err
}
t.NotifyCancelled(iid, reason)
out.State = "stopped"
default:
return out, fmt.Errorf("action must be pause|resume|cancel")
}
return out, nil
}
func (s *Server) controlTasksBatch(w http.ResponseWriter, r *http.Request) {
r.Body = http.MaxBytesReader(w, r.Body, 256<<10)
var req struct {
TaskIDs []string `json:"task_ids"`
Action string `json:"action"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeErr(w, 400, "bad json: "+err.Error())
return
}
if req.Action != "pause" && req.Action != "resume" {
writeErr(w, 400, "action must be pause|resume")
return
}
type parsedTaskID struct {
id string
valid bool
}
seen := map[string]bool{}
seenInvalid := map[string]bool{}
taskIDs := make([]parsedTaskID, 0, len(req.TaskIDs))
for _, raw := range req.TaskIDs {
trimmed := strings.TrimSpace(raw)
id, valid := canonicalTaskID(trimmed)
if !valid {
if seenInvalid[trimmed] {
continue
}
seenInvalid[trimmed] = true
taskIDs = append(taskIDs, parsedTaskID{id: trimmed})
continue
}
if seen[id] {
continue
}
seen[id] = true
taskIDs = append(taskIDs, parsedTaskID{id: id, valid: true})
}
if len(taskIDs) == 0 || len(taskIDs) > maxBatchControlIDs {
writeErr(w, 400, fmt.Sprintf("task_ids 数量必须为 1-%d", maxBatchControlIDs))
return
}
items := make([]batchControlItem, 0, len(taskIDs))
for _, parsed := range taskIDs {
item := batchControlItem{ID: parsed.id}
if !parsed.valid {
item.Error = "bad task id"
items = append(items, item)
continue
}
t, ok := s.m.Task(parsed.id)
if !ok {
item.Error = "task not found"
items = append(items, item)
continue
}
result, err := s.applyTaskControl(t, req.Action)
if err != nil {
item.Error = err.Error()
} else {
item.OK, item.Status, item.Queued = true, result.Status, result.Queued
}
items = append(items, item)
}
writeJSON(w, 200, map[string]any{"items": items})
}
+554
View File
@@ -0,0 +1,554 @@
package server
import (
"context"
"encoding/json"
"errors"
"fmt"
"iter"
"log"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/artex/llmrec"
"github.com/Autumn-27/norma/llm"
"github.com/Autumn-27/norma/memory"
"github.com/Autumn-27/norma/transcript"
)
type taskAgentBundle struct {
runtime *taskLLMRuntime // goal decomposition runtime
plannerRuntime *taskLLMRuntime
workerRuntime *taskLLMRuntime
mainRuntime *taskLLMRuntime
pl *agent.Planner
wk *agent.Worker
main *agent.MainAgent
}
type llmAuditProfile struct {
ID int64 `json:"id"`
Name string `json:"name"`
Format string `json:"format"`
Model string `json:"model"`
}
type llmTransitionAudit struct {
Mode string `json:"mode"` // automatic | manual | exhausted
Reason string `json:"reason"`
Previous *llmAuditProfile `json:"previous,omitempty"`
Next *llmAuditProfile `json:"next,omitempty"`
}
type llmActivityMetadata struct {
LLMTransition llmTransitionAudit `json:"llm_transition"`
}
type taskLLMRuntime struct {
s *Server
taskID string
agentKey string
}
type taskLLMError struct {
taskID string
chainExhausted bool
cause error
}
func (e *taskLLMError) Error() string {
if e.chainExhausted {
return fmt.Sprintf("task %s LLM profile chain exhausted: %v", e.taskID, e.cause)
}
return e.cause.Error()
}
func (e *taskLLMError) Unwrap() error { return e.cause }
func isTaskLLMRuntimeError(err error) bool {
var target *taskLLMError
return errors.As(err, &target)
}
func isTaskLLMChainExhausted(err error) bool {
var target *taskLLMError
return errors.As(err, &target) && target.chainExhausted
}
// isQuotaExhaustedError is intentionally strict. A generic 429, auth error,
// network failure, or 5xx does not rotate providers; the response must explicitly
// identify quota, credits, billing balance, or payment exhaustion.
func isQuotaExhaustedError(err error) bool {
return err != nil && agent.IsQuotaExhaustedMessage(err.Error())
}
type taskLLMSelection struct {
task *Task
profileID int64
revision int64
provider llm.Provider
}
type taskLLMStreamHooks struct {
current func() (taskLLMSelection, error)
exhaust func(taskLLMSelection, error) (db.TaskLLMTransition, error)
transition func(taskLLMSelection, db.TaskLLMTransition, error)
}
// current resolves the provider this role runs on, by precedence:
// Agent 绑定 → 任务 LLM 配置链 → 全局/环境配置。
// 绑定优先于任务链:某个角色被显式指定了模型,就一直跑在那个模型上;绑定不存在或
// 构建失败时才降级到任务链,任务链为空时再降级到全局配置。
// 返回的 profile id 只有走任务链时才非零 —— streamTaskLLM 以此判断额度错误是否
// 应该推进任务的故障转移状态(绑定/全局路径不改任务链状态,沿用既有语义)。
func (r *taskLLMRuntime) current() (*Task, int64, int64, llm.Provider, error) {
taskNum, err := parseTaskID(r.taskID)
if err != nil {
return nil, 0, 0, nil, err
}
pt, err := r.s.m.pg.GetTask(taskNum)
if err != nil {
return nil, 0, 0, nil, err
}
if pt == nil {
return nil, 0, 0, nil, fmt.Errorf("task %s not found", r.taskID)
}
r.s.syncTaskLLMState(pt)
t, _ := r.s.m.Task(r.taskID)
if prov, _, ok := r.s.agentBindingProvider(r.agentKey); ok {
return t, 0, pt.LLMChainRevision, prov, nil
}
if len(pt.LLMProfileIDs) > 0 {
if pt.ActiveLLMProfileID == nil {
return t, 0, pt.LLMChainRevision, nil, &taskLLMError{taskID: r.taskID, chainExhausted: true, cause: errors.New("all selected profiles are quota exhausted")}
}
prov, _, ok := r.s.providerForProfile(*pt.ActiveLLMProfileID)
if !ok {
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, nil, fmt.Errorf("LLM profile #%d is missing or invalid", *pt.ActiveLLMProfileID)
}
return t, *pt.ActiveLLMProfileID, pt.LLMChainRevision, prov, nil
}
prov, _, ok := r.s.globalProvider()
if !ok {
return t, 0, pt.LLMChainRevision, nil, fmt.Errorf("task %s has no available fallback LLM provider", r.taskID)
}
return t, 0, pt.LLMChainRevision, prov, nil
}
func parseTaskID(id string) (int64, error) {
n, err := strconv.ParseInt(id, 10, 64)
if err != nil || n <= 0 {
return 0, fmt.Errorf("invalid task id %q", id)
}
return n, nil
}
func (r *taskLLMRuntime) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
ctx = llmrec.WithTaskID(ctx, r.taskID)
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
task, profileID, revision, provider, err := r.current()
return taskLLMSelection{task: task, profileID: profileID, revision: revision, provider: provider}, err
},
exhaust: func(selection taskLLMSelection, cause error) (db.TaskLLMTransition, error) {
taskNum, _ := parseTaskID(r.taskID)
transition, err := r.s.m.pg.MarkTaskLLMProfileQuotaExhaustedAtRevision(taskNum, selection.profileID, selection.revision, cause.Error())
if err != nil {
return transition, err
}
if pt, getErr := r.s.m.pg.GetTask(taskNum); getErr == nil && pt != nil {
r.s.syncTaskLLMState(pt)
}
return transition, nil
},
transition: func(selection taskLLMSelection, transition db.TaskLLMTransition, cause error) {
r.s.emitTaskLLMTransition(selection.task, transition, cause)
},
}
return streamTaskLLM(ctx, r.taskID, req, hooks)
}
func streamTaskLLM(ctx context.Context, taskID string, req llm.CompletionRequest, hooks taskLLMStreamHooks) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) {
for {
selection, err := hooks.current()
if err != nil {
yield(llm.StreamEvent{}, err)
return
}
committed := false
var pending []llm.StreamEvent
var streamErr error
// 同 provider 安全窗口重试:committed 之前(还没向调用方交付任何输出)
// 的瞬时失败可以原样重放,不会重复模型输出或工具执行。committed 之后、
// ctx 取消、或确定性/额度错误则跳出,交给下方原有的透传/故障转移逻辑。
for attempt := 0; ; attempt++ {
committed = false
pending = nil
streamErr = nil
for event, err := range selection.provider.Stream(ctx, req) {
if err != nil {
streamErr = err
break
}
if !committed && !streamEventCommitsOutput(event) {
pending = append(pending, event)
continue
}
if !committed {
for _, buffered := range pending {
if !yield(buffered, nil) {
return
}
}
pending = nil
committed = true
}
if !yield(event, nil) {
return
}
}
if streamErr != nil && !committed && ctx.Err() == nil &&
attempt < sameProviderStreamRetries && isRetryableStreamError(streamErr) {
backoff := sameProviderRetryBackoff(attempt)
log.Printf("[task-llm] task %s 提交前流失败,%v 后同 provider 重试 (%d/%d): %v",
taskID, backoff, attempt+1, sameProviderStreamRetries, streamErr)
if sleepCtx(ctx, backoff) {
break // 退避期间 ctx 取消 → 停止重试
}
continue
}
break
}
if streamErr == nil {
for _, buffered := range pending {
if !yield(buffered, nil) {
return
}
}
return
}
// profileID=0 means the explicit chain was cleared while this stable task
// bundle was still in use. Agent/global fallback errors follow the legacy
// behavior and never mutate task failover state.
if selection.profileID == 0 || !isQuotaExhaustedError(streamErr) {
for _, buffered := range pending {
if !yield(buffered, nil) {
return
}
}
yield(llm.StreamEvent{}, streamErr)
return
}
transition, markErr := hooks.exhaust(selection, streamErr)
if markErr != nil {
cause := fmt.Errorf("mark profile quota exhausted after %v: %w", streamErr, markErr)
if committed {
// Output may already have driven tool execution. Report the persistence
// failure, but classify it as router-handled so the worker does not
// replay the entire intent and duplicate those side effects.
yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, cause: cause})
} else {
yield(llm.StreamEvent{}, cause)
}
return
}
if transition.Advanced && !transition.Stale && hooks.transition != nil {
hooks.transition(selection, transition, streamErr)
}
if committed || (!transition.Stale && transition.NextProfileID == nil) {
yield(llm.StreamEvent{}, &taskLLMError{taskID: taskID, chainExhausted: transition.ChainExhausted, cause: streamErr})
return
}
// No event reached the caller, so replaying the same logical request on
// the next profile cannot duplicate model output or tool execution.
}
}
}
func streamEventCommitsOutput(event llm.StreamEvent) bool {
switch event.Type {
case llm.SETextDelta, llm.SEThinkingDelta, llm.SEToolInputJSON:
return event.Text != ""
case llm.SEToolUseStart, llm.SEMessageDelta, llm.SEMessageStop:
return true
default:
return false
}
}
// 提交前安全窗口内、对同一 provider 的重试次数。SDK 的 doStream 只重试建连阶段
// (拿到 200 之前);流一旦开始,中途断流 / overloaded / 流内 429 等瞬时故障会直接
// 冒泡成 model_error,零重试。只要一个 token 都还没交给调用方(!committed),重放
// 完全相同的请求就不会重复模型输出或工具副作用,因此这里补一层同 provider 退避重试,
// 把这类抖动挡在意图整体重跑之前。
const sameProviderStreamRetries = 2
// sameProviderRetryBackoff 是第 attempt 次重试前的退避(0.5s、1s…,上限 4s),
// 与 SDK backoffSleep 同风格但封顶更小,避免拖住 worker 的收尾/取消响应。
// 以变量形式暴露,便于测试将退避置零。
var sameProviderRetryBackoff = func(attempt int) time.Duration {
return min(500*time.Millisecond*(1<<attempt), 4*time.Second)
}
// isRetryableStreamError 判断「提交前的流失败」是否值得在同一 provider 上重放。
// 瞬时的传输中断 / 供应商过载 / 限流会自行恢复,可安全重试;而以下三类不重试:
// - 额度耗尽:交给 profile 故障转移处理,别在这里白烧重试
// - 上下文过长:相同请求重放也没用,交给 harness 的 reactive 压缩兜底
// - 4xx 确定性拒绝(400/401/403/404/422):到哪个 provider 都一样会失败
func isRetryableStreamError(err error) bool {
if err == nil {
return false
}
if isQuotaExhaustedError(err) {
return false
}
s := strings.ToLower(err.Error())
if strings.Contains(s, "too long") || strings.Contains(s, "context length") ||
strings.Contains(s, "context_length") || strings.Contains(s, "maximum context") ||
strings.Contains(s, "status 413") {
return false
}
for _, code := range []string{"status 400", "status 401", "status 403", "status 404", "status 422"} {
if strings.Contains(s, code) {
return false
}
}
// 其余(传输 reset/EOF/timeout、408/429/5xx、流内 error 事件如 anthropic
// overloaded_error 等)一律视为瞬时,允许重试。
return true
}
// 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 _, cfg, ok := r.s.agentBindingProvider(r.agentKey); ok {
return cfg.CompactionWindow()
}
taskNum, err := parseTaskID(r.taskID)
if err != nil {
return (agent.Config{}).CompactionWindow()
}
chain, err := r.s.m.pg.TaskLLMProfiles(taskNum)
if err != nil {
return (agent.Config{}).CompactionWindow()
}
if len(chain) == 0 {
if _, cfg, ok := r.s.globalProvider(); ok {
return cfg.CompactionWindow()
}
return (agent.Config{}).CompactionWindow()
}
minimum := 0
for _, entry := range chain {
if cfg, ok := r.s.loadProfileConfig(entry.ProfileID); ok {
window := cfg.CompactionWindow()
if minimum == 0 || window < minimum {
minimum = window
}
}
}
if minimum == 0 {
return (agent.Config{}).CompactionWindow()
}
return minimum
}
// agentBindingProvider resolves the profile a role is explicitly bound to
// (agents.llm_profile_id) — the highest-precedence level for task agents. ok=false
// when the role has no binding or the bound profile no longer builds, so callers
// fall through to the task chain. A bound profile stays exclusive unless
// llm_pool_bind_fallback is on, which is what poolForBinding encodes.
func (s *Server) agentBindingProvider(agentKey string) (llm.Provider, agent.Config, bool) {
id := s.effectiveProfileForAgent(agentKey, nil)
if id == nil {
return nil, agent.Config{}, false
}
prov, cfg, ok := s.providerForProfile(*id)
if !ok {
return nil, agent.Config{}, false
}
return s.poolForBinding(*id, prov, cfg), cfg, true
}
// globalProvider returns the process-wide provider (persisted active profile or
// environment config) — the last resort once a role has neither a binding nor a
// task chain.
func (s *Server) globalProvider() (llm.Provider, agent.Config, bool) {
s.cfgMu.Lock()
defer s.cfgMu.Unlock()
if !s.llmOn || s.llmProv == nil {
return nil, agent.Config{}, false
}
return s.llmProv, s.llmCfg, true
}
// taskRuntimeAvailable reports whether every listed role can resolve a provider
// under the runtime precedence in current(): the role's own binding first, then
// the task chain, then global. An exhausted chain is a hard stop for unbound
// roles rather than a silent fall through to global — same as current().
func (s *Server) taskRuntimeAvailable(t *Task, agentKeys ...string) bool {
if t == nil || len(agentKeys) == 0 {
return false
}
state := t.llmStateSnapshot()
unboundReady := false
if len(state.ProfileIDs) > 0 {
if state.ActiveID != nil {
_, _, unboundReady = s.providerForProfile(*state.ActiveID)
}
} else {
_, _, unboundReady = s.globalProvider()
}
for _, key := range agentKeys {
if _, _, ok := s.agentBindingProvider(key); ok {
continue
}
if !unboundReady {
return false
}
}
return true
}
func (s *Server) invalidateTaskAgents() {
s.taskAgentMu.Lock()
s.taskAgents = map[string]*taskAgentBundle{}
s.taskAgentMu.Unlock()
}
func (s *Server) agentsForTask(t *Task) *taskAgentBundle {
s.taskAgentMu.Lock()
defer s.taskAgentMu.Unlock()
if bundle := s.taskAgents[t.ID]; bundle != nil {
return bundle
}
goalRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "goals"}
plannerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "planner"}
workerRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "worker"}
mainRuntime := &taskLLMRuntime{s: s, taskID: t.ID, agentKey: "mainagent"}
tx := transcript.NewStore(filepath.Join(s.m.dir, "transcripts"))
window := workerRuntime.CompactionWindow()
wk := agent.NewWorker(workerRuntime, "task-router", s.m.dir, tx, window, s.agentMaxTurns("worker"))
wk.SetCompactionWindowResolver(workerRuntime.CompactionWindow)
wk.SetRunTimeout(time.Duration(s.agentRunSeconds("worker")) * time.Second)
wk.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
wk.SetMemory(memory.NewStore(filepath.Join(s.m.dir, "memory")))
wk.SetWebSearch(s.webSearchFor("worker"))
wk.SetConstraintInject(s.constraintInjectWorker) // 操作约束注入 worker(可配置,默认开)
pl := agent.NewPlanner(plannerRuntime, "task-router", s.m.dir, tx, plannerRuntime.CompactionWindow(), s.agentMaxTurns("planner"))
pl.SetCompactionWindowResolver(plannerRuntime.CompactionWindow)
pl.SetKillWork(s.engine.KillWork)
pl.SetSteerWork(s.engine.SteerWork)
pl.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
pl.SetWebSearch(s.webSearchFor("planner"))
pl.SetConstraintInject(s.constraintInjectPlanner) // 操作约束注入 planner(可配置,默认开)
main := agent.NewMainAgent(mainRuntime, "task-router", s.m.dir, tx, mainRuntime.CompactionWindow(), s.agentMaxTurns("mainagent"))
main.SetCompactionWindowResolver(mainRuntime.CompactionWindow)
main.SetProxy(s.m.ProxyAddr(), s.m.ProxyCACert())
main.SetWebSearch(s.webSearchFor("mainagent"))
main.SetSteerWork(s.engine.SteerWork) // steer_work:人对运行中 work 实时纠偏
bundle := &taskAgentBundle{
runtime: goalRuntime, plannerRuntime: plannerRuntime, workerRuntime: workerRuntime,
mainRuntime: mainRuntime, pl: pl, wk: wk, main: main,
}
s.taskAgents[t.ID] = bundle
return bundle
}
func (s *Server) syncTaskLLMState(pt *db.Task) {
if pt == nil {
return
}
id := fmt.Sprintf("%d", pt.ID)
s.m.mu.Lock()
if task := s.m.tasks[id]; task != nil {
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMChainRevision, pt.LLMFailoverState, pt.LLMFailoverReason)
}
s.m.mu.Unlock()
}
func (s *Server) emitTaskLLMTransition(t *Task, transition db.TaskLLMTransition, cause error) {
if t == nil {
return
}
previous := s.llmAuditProfile(transition.PreviousProfileID)
var next *llmAuditProfile
if transition.NextProfileID != nil {
next = s.llmAuditProfile(*transition.NextProfileID)
}
mode := "automatic"
kind := "llm_switch"
summary := fmt.Sprintf("%s 额度不足", llmAuditProfileLabel(previous))
if transition.NextProfileID != nil {
summary += fmt.Sprintf(",后续调用切换到 %s", llmAuditProfileLabel(next))
} else {
mode = "exhausted"
kind = "llm_failover"
summary += ",配置链已耗尽"
}
metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{
Mode: mode, Reason: cause.Error(), Previous: previous, Next: next,
}})
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: kind, IsError: transition.ChainExhausted, Summary: summary, Detail: cause.Error(), Metadata: metadata})
log.Printf("[llm-failover] task %s: %s", t.ID, summary)
}
func (s *Server) llmAuditProfile(id int64) *llmAuditProfile {
if id <= 0 || s.m == nil || s.m.pg == nil {
return nil
}
p, err := s.m.pg.ProfileByID(id)
if err != nil || p == nil {
return &llmAuditProfile{ID: id, Name: fmt.Sprintf("配置 #%d", id)}
}
return &llmAuditProfile{ID: p.ID, Name: p.Name, Format: p.Format, Model: p.Model}
}
func llmAuditProfileLabel(profile *llmAuditProfile) string {
if profile == nil {
return "默认配置"
}
name := profile.Name
if name == "" {
name = fmt.Sprintf("配置 #%d", profile.ID)
}
detail := []string{}
if profile.Format != "" {
detail = append(detail, profile.Format)
}
if profile.Model != "" {
detail = append(detail, profile.Model)
}
if len(detail) == 0 {
return name
}
return fmt.Sprintf("%s(%s)", name, strings.Join(detail, " / "))
}
func sameOptionalID(a, b *int64) bool {
if a == nil || b == nil {
return a == nil && b == nil
}
return *a == *b
}
func (s *Server) emitManualTaskLLMSwitch(t *Task, previousID, nextID *int64) db.Activity {
previous := (*llmAuditProfile)(nil)
next := (*llmAuditProfile)(nil)
if previousID != nil {
previous = s.llmAuditProfile(*previousID)
}
if nextID != nil {
next = s.llmAuditProfile(*nextID)
}
summary := fmt.Sprintf("用户手动将任务 LLM 从 %s 切换到 %s", llmAuditProfileLabel(previous), llmAuditProfileLabel(next))
metadata, _ := json.Marshal(llmActivityMetadata{LLMTransition: llmTransitionAudit{
Mode: "manual", Reason: "用户手动切换任务 LLM", Previous: previous, Next: next,
}})
return s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "llm_switch", Summary: summary, Detail: summary, Metadata: metadata})
}
+675
View File
@@ -0,0 +1,675 @@
package server
import (
"context"
"encoding/json"
"errors"
"fmt"
"iter"
"sync"
"testing"
"time"
"github.com/Autumn-27/artex/agent"
"github.com/Autumn-27/artex/db"
"github.com/Autumn-27/norma/harness"
"github.com/Autumn-27/norma/llm"
)
type scriptedLLMProvider struct {
events []llm.StreamEvent
err error
calls int
}
func (p *scriptedLLMProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) {
p.calls++
for _, event := range p.events {
if !yield(event, nil) {
return
}
}
if p.err != nil {
yield(llm.StreamEvent{}, p.err)
}
}
}
// withZeroRetryBackoff sets the same-provider retry backoff to zero for the
// duration of a (serial) test and returns a restore func for defer.
func withZeroRetryBackoff() func() {
prev := sameProviderRetryBackoff
sameProviderRetryBackoff = func(int) time.Duration { return 0 }
return func() { sameProviderRetryBackoff = prev }
}
// flakyThenOKProvider fails its first failCount stream attempts pre-commit
// (emitting only a non-committing SEMessageStart before the error, mirroring a
// gateway that returns 200 then drops), then serves okEvents.
type flakyThenOKProvider struct {
failCount int
failErr error
okEvents []llm.StreamEvent
calls int
}
func (p *flakyThenOKProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
return func(yield func(llm.StreamEvent, error) bool) {
p.calls++
if p.calls <= p.failCount {
if !yield(llm.StreamEvent{Type: llm.SEMessageStart}, nil) {
return
}
yield(llm.StreamEvent{}, p.failErr)
return
}
for _, ev := range p.okEvents {
if !yield(ev, nil) {
return
}
}
}
}
func collectTaskLLMStream(seq iter.Seq2[llm.StreamEvent, error]) ([]llm.StreamEvent, error) {
var events []llm.StreamEvent
var streamErr error
for event, err := range seq {
if err != nil {
streamErr = err
break
}
events = append(events, event)
}
return events, streamErr
}
func TestIsQuotaExhaustedError(t *testing.T) {
t.Parallel()
positive := []string{
`openai: code=insufficient_quota status=429`,
`provider error: quota_exceeded`,
`grpc code=RESOURCE_EXHAUSTED`,
`You exceeded your current quota, please check your plan and billing details.`,
`billing_hard_limit_reached`,
`billing_not_active`,
`HTTP status 402 Payment Required`,
`账户余额不足,请充值`,
`credit balance is too low`,
}
for _, message := range positive {
if !isQuotaExhaustedError(errors.New(message)) {
t.Errorf("expected quota classification for %q", message)
}
}
negative := []string{
`HTTP status 429: rate limit exceeded`,
`RESOURCE_EXHAUSTED: rate limit exceeded`,
`HTTP status 429: quota exceeded for quota metric GenerateRequestsPerMinutePerProjectPerBaseModel`,
`RESOURCE_EXHAUSTED: TPM quota exceeded`,
`HTTP status 401: invalid api key`,
`HTTP status 401: insufficient_quota`,
`HTTP status 403: forbidden`,
`HTTP status 403: billing_hard_limit_reached`,
`HTTP status 500: internal server error`,
`HTTP status 503: quota_exceeded`,
`dial tcp: network is unreachable`,
`context length exceeded`,
`invalid request parameter`,
}
for _, message := range negative {
if isQuotaExhaustedError(errors.New(message)) {
t.Errorf("unexpected quota classification for %q", message)
}
}
if isQuotaExhaustedError(nil) {
t.Error("nil error must not be classified as quota exhaustion")
}
}
func TestTaskDTOUsesEmptyArrays(t *testing.T) {
t.Parallel()
dto := taskDTO(&Task{ID: "7"}, "created")
raw, err := json.Marshal(dto)
if err != nil {
t.Fatal(err)
}
var payload map[string]any
if err := json.Unmarshal(raw, &payload); err != nil {
t.Fatal(err)
}
for _, field := range []string{"source_task_ids", "llm_profile_ids"} {
value, ok := payload[field].([]any)
if !ok || len(value) != 0 {
t.Fatalf("%s must serialize as [], payload=%s", field, raw)
}
}
}
func TestTaskLLMStateSnapshotIsConcurrentAndDetached(t *testing.T) {
t.Parallel()
task := &Task{ID: "7"}
var wg sync.WaitGroup
for writer := 0; writer < 2; writer++ {
wg.Add(1)
go func(offset int) {
defer wg.Done()
for i := 0; i < 500; i++ {
id := int64(offset*1000 + i + 1)
task.setLLMState(&id, &id, []int64{id, id + 1}, int64(i), "ready", "")
}
}(writer)
}
for reader := 0; reader < 4; reader++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 500; i++ {
state := task.llmStateSnapshot()
if len(state.ProfileIDs) > 0 {
state.ProfileIDs[0] = -1
}
_ = taskDTO(task, "running")
}
}()
}
wg.Wait()
state := task.llmStateSnapshot()
if len(state.ProfileIDs) != 2 || state.ProfileIDs[0] <= 0 {
t.Fatalf("snapshot mutation leaked into task state: %+v", state)
}
}
func TestTaskLLMStateRejectsOlderRevision(t *testing.T) {
t.Parallel()
task := &Task{ID: "7"}
current := int64(22)
stale := int64(11)
if !task.setLLMState(&current, &current, []int64{current}, 2, "ready", "") {
t.Fatal("initial LLM state was rejected")
}
if task.setLLMState(&stale, &stale, []int64{stale}, 1, "ready", "stale") {
t.Fatal("older LLM snapshot was accepted")
}
state := task.llmStateSnapshot()
if state.ChainRevision != 2 || state.ActiveID == nil || *state.ActiveID != current ||
len(state.ProfileIDs) != 1 || state.ProfileIDs[0] != current || state.FailoverReason != "" {
t.Fatalf("older LLM snapshot overwrote current state: %+v", state)
}
}
func TestTaskLLMStreamRetriesPreStreamQuotaOnNextProfile(t *testing.T) {
t.Parallel()
quota := errors.New("openai: status 402: insufficient_quota")
first := &scriptedLLMProvider{err: quota}
second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "fallback"}}}
active := int64(11)
transitions := 0
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
if active == 11 {
return taskLLMSelection{profileID: 11, provider: first}, nil
}
return taskLLMSelection{profileID: 22, provider: second}, nil
},
exhaust: func(selection taskLLMSelection, _ error) (db.TaskLLMTransition, error) {
if selection.profileID != 11 {
t.Fatalf("unexpected exhausted profile %d", selection.profileID)
}
active = 22
next := active
return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
},
transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ },
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err != nil {
t.Fatal(err)
}
if len(events) != 1 || events[0].Text != "fallback" {
t.Fatalf("unexpected fallback events: %+v", events)
}
if first.calls != 1 || second.calls != 1 || transitions != 1 {
t.Fatalf("calls first=%d second=%d transitions=%d", first.calls, second.calls, transitions)
}
}
func TestTaskLLMStreamRetriesAfterNonContentStartEvent(t *testing.T) {
t.Parallel()
quota := errors.New("anthropic: status 402: insufficient credit balance")
first := &scriptedLLMProvider{
events: []llm.StreamEvent{{Type: llm.SEMessageStart}},
err: quota,
}
second := &scriptedLLMProvider{events: []llm.StreamEvent{
{Type: llm.SEMessageStart},
{Type: llm.SETextDelta, Text: "fallback"},
}}
active := int64(11)
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
if active == 11 {
return taskLLMSelection{profileID: 11, provider: first}, nil
}
return taskLLMSelection{profileID: 22, provider: second}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
active = 22
next := active
return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
},
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err != nil {
t.Fatal(err)
}
if len(events) != 2 || events[0].Type != llm.SEMessageStart || events[1].Text != "fallback" {
t.Fatalf("non-content event from failed attempt leaked or fallback missing: %+v", events)
}
if first.calls != 1 || second.calls != 1 {
t.Fatalf("calls first=%d second=%d", first.calls, second.calls)
}
}
func TestTaskLLMStreamRetriesStalePreStreamFailureAgainstReplacementChain(t *testing.T) {
t.Parallel()
first := &scriptedLLMProvider{err: errors.New("status 402: insufficient_quota")}
second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "replacement"}}}
current := first
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
profileID := int64(11)
if current == second {
profileID = 22
}
return taskLLMSelection{profileID: profileID, provider: current}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
current = second
return db.TaskLLMTransition{PreviousProfileID: 11, Stale: true}, nil
},
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err != nil || len(events) != 1 || events[0].Text != "replacement" {
t.Fatalf("stale call did not continue on replacement chain: events=%+v err=%v", events, err)
}
}
func TestTaskLLMStreamDoesNotSwitchForOrdinaryErrors(t *testing.T) {
// Not parallel: overrides the package-level backoff so the retryable cases
// don't sleep. Serial tests never overlap the parallel batch, so this is safe.
defer withZeroRetryBackoff()()
tests := []string{
"openai: status 429: rate limit exceeded",
"openai: status 401: invalid api key",
"openai: status 500: internal server error",
"dial tcp: network is unreachable",
}
for _, message := range tests {
t.Run(message, func(t *testing.T) {
provider := &scriptedLLMProvider{err: errors.New(message)}
exhausted := 0
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
exhausted++
return db.TaskLLMTransition{}, nil
},
}
_, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err == nil || err.Error() != message {
t.Fatalf("error=%v, want %q", err, message)
}
if exhausted != 0 {
t.Fatalf("ordinary error changed failover state %d times", exhausted)
}
})
}
}
func TestTaskLLMStreamAdvancesAfterMidStreamQuotaWithoutReplay(t *testing.T) {
t.Parallel()
quota := errors.New("anthropic: status 402: credit balance is too low")
first := &scriptedLLMProvider{
events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "partial"}},
err: quota,
}
second := &scriptedLLMProvider{events: []llm.StreamEvent{{Type: llm.SETextDelta, Text: "next-call"}}}
active := int64(11)
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
if active == 11 {
return taskLLMSelection{profileID: 11, provider: first}, nil
}
return taskLLMSelection{profileID: 22, provider: second}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
active = 22
next := active
return db.TaskLLMTransition{PreviousProfileID: 11, NextProfileID: &next, Advanced: true}, nil
},
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if len(events) != 1 || events[0].Text != "partial" || err == nil {
t.Fatalf("mid-stream result events=%+v err=%v", events, err)
}
if first.calls != 1 || second.calls != 0 {
t.Fatalf("mid-stream call was replayed: first=%d second=%d", first.calls, second.calls)
}
nextEvents, nextErr := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if nextErr != nil || len(nextEvents) != 1 || nextEvents[0].Text != "next-call" || second.calls != 1 {
t.Fatalf("next call did not use fallback: events=%+v err=%v calls=%d", nextEvents, nextErr, second.calls)
}
}
func TestTaskLLMStreamMarkFailureAfterToolStartSkipsWorkerReplay(t *testing.T) {
t.Parallel()
quota := errors.New("openai: status 402: insufficient_quota")
markFailure := errors.New("database temporarily unavailable")
provider := &scriptedLLMProvider{
events: []llm.StreamEvent{{Type: llm.SEToolUseStart}},
err: quota,
}
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
return db.TaskLLMTransition{}, markFailure
},
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if len(events) != 1 || events[0].Type != llm.SEToolUseStart || err == nil {
t.Fatalf("committed tool stream events=%+v err=%v", events, err)
}
if !isTaskLLMRuntimeError(err) {
t.Fatalf("committed mark failure must be router-classified: %T %v", err, err)
}
if retryableWorkerModelError(harness.ReasonModelError, err) {
t.Fatal("worker must not replay after a committed tool stream when quota persistence fails")
}
if provider.calls != 1 {
t.Fatalf("provider calls=%d, want 1", provider.calls)
}
}
func TestTaskLLMRuntimeErrorSkipsWorkerReplay(t *testing.T) {
partialStreamErr := &taskLLMError{
taskID: "7",
chainExhausted: false,
cause: errors.New("insufficient_quota after partial output"),
}
if retryableWorkerModelError(harness.ReasonModelError, partialStreamErr) {
t.Fatal("worker must not replay an intent after the task router handled a mid-stream quota error")
}
if !retryableWorkerModelError(harness.ReasonModelError, errors.New("temporary network error")) {
t.Fatal("ordinary model errors should retain the existing worker retry behavior")
}
}
// A transient pre-commit stream failure (200-then-drop / overloaded / 5xx) is
// retried on the SAME provider, without switching profiles, and succeeds once the
// blip clears — instead of surfacing as a model_error.
func TestTaskLLMStreamRetriesTransientPreCommitOnSameProvider(t *testing.T) {
defer withZeroRetryBackoff()()
provider := &flakyThenOKProvider{
failCount: 2, // fail twice, succeed on the 3rd attempt (initial + 2 retries)
failErr: errors.New("anthropic: overloaded_error"),
okEvents: []llm.StreamEvent{{Type: llm.SEMessageStart}, {Type: llm.SETextDelta, Text: "ok"}},
}
exhausted := 0
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
exhausted++
return db.TaskLLMTransition{}, nil
},
}
events, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err != nil {
t.Fatalf("transient pre-commit failure should have been retried to success, got %v", err)
}
if provider.calls != 3 {
t.Fatalf("provider calls=%d, want 3 (1 initial + 2 retries)", provider.calls)
}
if exhausted != 0 {
t.Fatalf("same-provider retry must not change failover state (exhausted=%d)", exhausted)
}
if len(events) == 0 || events[len(events)-1].Text != "ok" {
t.Fatalf("recovered stream missing its committed output: %+v", events)
}
}
// Once retries are exhausted, the transient error is surfaced (becoming a
// model_error upstream) after exactly sameProviderStreamRetries+1 attempts.
func TestTaskLLMStreamSurfacesTransientAfterRetriesExhausted(t *testing.T) {
defer withZeroRetryBackoff()()
provider := &flakyThenOKProvider{
failCount: 99, // never recovers
failErr: errors.New("dial tcp: connection reset by peer"),
}
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
return db.TaskLLMTransition{}, nil
},
}
_, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err == nil {
t.Fatal("persistent transient failure must still surface an error")
}
if provider.calls != sameProviderStreamRetries+1 {
t.Fatalf("provider calls=%d, want %d", provider.calls, sameProviderStreamRetries+1)
}
}
// A deterministic 4xx rejection is NOT retried on the same provider — replaying
// the identical request everywhere fails the same way.
func TestTaskLLMStreamDoesNotRetryDeterministicRejection(t *testing.T) {
defer withZeroRetryBackoff()()
provider := &flakyThenOKProvider{
failCount: 99,
failErr: errors.New("anthropic: status 400: messages.1: invalid request"),
}
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
return db.TaskLLMTransition{}, nil
},
}
_, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if err == nil {
t.Fatal("expected the deterministic rejection to surface")
}
if provider.calls != 1 {
t.Fatalf("deterministic 4xx must not be retried: provider calls=%d, want 1", provider.calls)
}
}
func TestEngineReadyForTaskOverride(t *testing.T) {
t.Parallel()
e := &Engine{}
if e.Ready() {
t.Fatal("engine without a global provider must not be globally ready")
}
e.SetAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) {
return &agent.Planner{}, &agent.Worker{}
})
if !e.ReadyFor(&Task{ID: "7"}) {
t.Fatal("task-specific provider chain must make that task ready")
}
}
func TestResolvedTaskStatusKeepsExhaustedStartedTaskRunning(t *testing.T) {
t.Parallel()
engine := NewEngine(nil)
engine.started.Store("7", true)
task := &Task{ID: "7", Status: "running"}
profileID := int64(11)
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
s := &Server{engine: engine}
if got := s.resolvedTaskStatus(task); got != "running" {
t.Fatalf("exhausted started status=%q, want running", got)
}
task.Paused = true
if got := s.resolvedTaskStatus(task); got != "paused" {
t.Fatalf("paused exhausted status=%q, want paused", got)
}
}
// A role with no Agent binding falls through to the task chain, then to global —
// and an exhausted chain stays a hard stop instead of silently borrowing global.
func TestTaskRuntimeAvailableUnboundRoleFollowsChainThenGlobal(t *testing.T) {
t.Parallel()
// A nil pg means effectiveProfileForAgent finds no binding for any role, so
// every role here resolves through the chain/global levels.
s := &Server{m: &Manager{}}
task := &Task{ID: "7"}
if s.taskRuntimeAvailable(task, "worker") {
t.Fatal("no binding, no chain and no global provider must not be runnable")
}
s.llmOn, s.llmProv = true, &scriptedLLMProvider{}
if !s.taskRuntimeAvailable(task, "worker") {
t.Fatal("global provider must serve a role with no binding and no chain")
}
if s.taskRuntimeAvailable(task) {
t.Fatal("a task with no roles requested must not report as runnable")
}
profileID := int64(11)
task.setLLMState(&profileID, nil, []int64{profileID}, 1, "chain_exhausted", "quota exhausted")
if s.taskRuntimeAvailable(task, "worker") {
t.Fatal("exhausted chain must not fall back to the global provider")
}
}
func TestAuthoritativeTaskResolverDoesNotFallBackWhenChainUnavailable(t *testing.T) {
t.Parallel()
e := &Engine{}
e.UseLLM(&agent.Planner{}, &agent.Worker{})
if !e.ReadyFor(&Task{ID: "7"}) {
t.Fatal("global provider should initially make task ready")
}
e.SetAuthoritativeAgentResolver(func(*Task) (*agent.Planner, *agent.Worker) {
return nil, nil
})
if e.ReadyFor(&Task{ID: "7"}) {
t.Fatal("unavailable authoritative task chain must not fall back to global provider")
}
}
func TestTaskLLMStreamStopsWhenChainExhausted(t *testing.T) {
t.Parallel()
provider := &scriptedLLMProvider{err: errors.New("billing_not_active")}
transitions := 0
hooks := taskLLMStreamHooks{
current: func() (taskLLMSelection, error) {
return taskLLMSelection{profileID: 11, provider: provider}, nil
},
exhaust: func(taskLLMSelection, error) (db.TaskLLMTransition, error) {
return db.TaskLLMTransition{PreviousProfileID: 11, Advanced: true, ChainExhausted: true}, nil
},
transition: func(taskLLMSelection, db.TaskLLMTransition, error) { transitions++ },
}
_, err := collectTaskLLMStream(streamTaskLLM(context.Background(), "7", llm.CompletionRequest{}, hooks))
if !isTaskLLMChainExhausted(err) {
t.Fatalf("expected chain-exhausted error, got %v", err)
}
if provider.calls != 1 || transitions != 1 {
t.Fatalf("calls=%d transitions=%d", provider.calls, transitions)
}
}
func TestProfileDeleteRestoresQuotaBlockedIntentWhenFallbackAvailable(t *testing.T) {
m, err := NewManager(t.TempDir(), "")
if err != nil {
t.Skipf("postgres unavailable (%v) - skipping", err)
}
defer m.Close()
profileID, err := m.pg.SaveProfile(&db.LLMProfile{
Name: fmt.Sprintf("delete-quota-fallback-%d", time.Now().UnixNano()),
Format: "openai",
Model: "test-model",
APIKey: "test-key",
})
if err != nil {
t.Fatal(err)
}
task, err := m.CreateTaskWithOptions("quota fallback", "resume intent", db.TaskCreateOptions{
LLMProfileIDs: []int64{profileID},
})
if err != nil {
_ = m.pg.DeleteProfile(profileID)
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{})
_ = m.pg.DeleteProfile(profileID)
})
intentID, err := task.Store.AddIntent(map[string]any{"action": "retry after fallback"}, 1, nil, "planner")
if err != nil {
t.Fatal(err)
}
taskID, err := parseTaskID(task.ID)
if err != nil {
t.Fatal(err)
}
if transition, err := m.pg.MarkTaskLLMProfileQuotaExhausted(taskID, profileID, "insufficient_quota"); err != nil || !transition.ChainExhausted {
t.Fatalf("exhaust task chain: transition=%+v err=%v", transition, err)
}
if err := task.Store.SetIntentBlockedReason(intentID, db.IntentBlockedLLMQuota); err != nil {
t.Fatal(err)
}
task.Paused = true
if err := m.SetTaskPaused(task.ID, true); err != nil {
t.Fatal(err)
}
// Removing the final explicit entry restores the legacy/global fallback. The
// post-delete sync must reopen only quota-blocked work without unpausing it.
if err := m.pg.DeleteProfile(profileID); err != nil {
t.Fatal(err)
}
s := &Server{
m: m,
llmOn: true,
llmProv: &scriptedLLMProvider{},
provByProfile: map[int64]*provEntry{},
}
s.restoreTasksAfterProfileDelete(m.pg)
node, err := task.Store.GetNode(intentID)
if err != nil {
t.Fatal(err)
}
if node == nil || node.State != "open" || node.BlockedReason != "" {
t.Fatalf("quota-blocked intent was not reopened: %+v", node)
}
if !task.Paused {
t.Fatal("profile deletion must not resume a paused task")
}
state := task.llmStateSnapshot()
if len(state.ProfileIDs) != 0 || state.FailoverState != "default" {
t.Fatalf("task did not return to fallback state: %+v", state)
}
}
+123
View File
@@ -0,0 +1,123 @@
package server
import (
"fmt"
"net/http"
"github.com/Autumn-27/artex/db"
)
type taskLLMResolution struct {
ProfileID *int64 `json:"profile_id,omitempty"`
Name string `json:"name"`
Format string `json:"format"`
Model string `json:"model"`
Source string `json:"source"`
Available bool `json:"available"`
Reason string `json:"reason,omitempty"`
}
func (s *Server) resolutionFromProfile(p *db.LLMProfile, source string) taskLLMResolution {
if p == nil {
return taskLLMResolution{Source: source, Reason: "LLM 配置不存在"}
}
id := p.ID
result := taskLLMResolution{
ProfileID: &id,
Name: p.Name,
Format: p.Format,
Model: p.Model,
Source: source,
}
if p.APIKey == "" {
result.Reason = "LLM 配置未设置 API Key"
return result
}
if _, _, ok := s.providerForProfile(p.ID); !ok {
result.Reason = "LLM 配置格式或参数无效"
return result
}
result.Available = true
return result
}
// resolveTaskRoleLLM mirrors taskLLMRuntime.current without creating a provider:
// role Agent binding -> explicit task chain -> active/global environment config.
// Database failures are returned rather than represented as an unavailable model.
func (s *Server) resolveTaskRoleLLM(t *Task, agentKey string) (taskLLMResolution, error) {
if s == nil || s.m == nil || s.m.pg == nil {
return taskLLMResolution{}, fmt.Errorf("database unavailable")
}
a, err := s.m.pg.GetAgentByKey(agentKey)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load agent %q: %w", agentKey, err)
}
if a != nil && a.LLMProfileID != nil {
p, err := s.m.pg.ProfileByID(*a.LLMProfileID)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load agent LLM profile: %w", err)
}
if p != nil {
resolved := s.resolutionFromProfile(p, "agent_binding")
if resolved.Available {
return resolved, nil
}
}
}
state := t.llmStateSnapshot()
if len(state.ProfileIDs) > 0 {
if state.ActiveID == nil {
return taskLLMResolution{Source: "task_chain", Reason: "任务 LLM 配置链额度已耗尽"}, nil
}
p, err := s.m.pg.ProfileByID(*state.ActiveID)
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load task LLM profile: %w", err)
}
return s.resolutionFromProfile(p, "task_chain"), nil
}
p, err := s.m.pg.ActiveProfile()
if err != nil {
return taskLLMResolution{}, fmt.Errorf("load global LLM profile: %w", err)
}
if p != nil {
resolved := s.resolutionFromProfile(p, "global_profile")
if resolved.Available {
return resolved, nil
}
}
s.cfgMu.Lock()
on, name, cfg := s.llmOn, s.llmProf, s.llmCfg
s.cfgMu.Unlock()
if on {
if name == "" {
name = "全局配置"
}
return taskLLMResolution{
Name: name, Format: cfg.Provider(), Model: cfg.Model,
Source: "environment", Available: true,
}, nil
}
return taskLLMResolution{Source: "global", Reason: "没有可用的 LLM 配置"}, nil
}
func (s *Server) taskLLMResolutionHandler(w http.ResponseWriter, r *http.Request) {
t, ok := s.m.Task(r.PathValue("id"))
if !ok {
writeErr(w, 404, "task not found")
return
}
roles := map[string]taskLLMResolution{}
for responseKey, agentKey := range map[string]string{
"mainagent": "mainagent",
"planner": "planner",
"worker": "worker",
} {
resolved, err := s.resolveTaskRoleLLM(t, agentKey)
if err != nil {
writeErr(w, http.StatusInternalServerError, err.Error())
return
}
roles[responseKey] = resolved
}
writeJSON(w, http.StatusOK, roles)
}