同步
This commit is contained in:
1 parent
e67ff26e90
commit
58f23a2773
61 files changed
+11611
-443
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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 守卫):用户删掉后不再重建。
|
||||
|
||||
@@ -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
@@ -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
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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(¤t, ¤t, []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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
Reference in new issue
Block a user