feat(tasks): add inherited context and resilient LLM execution
Add read-only task context inheritance, ordered task-level LLM failover, live Worker controls, cascading task cleanup, FIFO task concurrency, cancellation diagnostics, UI updates, tests, build script, and changelog.
This commit is contained in:
1 parent
7090b034f8
commit
9793b522d9
82 files changed
+11094
-1712
No files matched your search
+2
-5
@@ -217,7 +217,6 @@ func (s *Server) reattribute(w http.ResponseWriter, r *http.Request) {
|
||||
// GET /api/assets
|
||||
// =====================================================================
|
||||
|
||||
// defaultAssetPageSize is the page size used when the caller omits ?limit.
|
||||
const defaultAssetPageSize = 50
|
||||
|
||||
func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -230,8 +229,6 @@ func (s *Server) listAssets(w http.ResponseWriter, r *http.Request) {
|
||||
typ := q.Get("type")
|
||||
limit, _ := strconv.Atoi(q.Get("limit"))
|
||||
offset, _ := strconv.Atoi(q.Get("offset"))
|
||||
// limit/offset are page controls, not a cap: `total` always reports the full
|
||||
// match count so callers can page through everything.
|
||||
if limit <= 0 {
|
||||
limit = defaultAssetPageSize
|
||||
}
|
||||
@@ -351,8 +348,8 @@ func (s *Server) insertAssets(w http.ResponseWriter, r *http.Request) {
|
||||
Assets []struct {
|
||||
Type string `json:"type"`
|
||||
// root_domain / subdomain
|
||||
Domain string `json:"domain"`
|
||||
ICP string `json:"icp"`
|
||||
Domain string `json:"domain"`
|
||||
ICP string `json:"icp"`
|
||||
RecordType string `json:"record_type"`
|
||||
RecordValue []string `json:"record_value"`
|
||||
// ip
|
||||
|
||||
@@ -45,9 +45,11 @@ type chatAttachment struct {
|
||||
// 文件先落这里,前端按返回的 abs 绝对路径写进任务描述)
|
||||
func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) {
|
||||
var sub string
|
||||
taskScoped := false
|
||||
switch r.URL.Query().Get("scope") {
|
||||
case "task":
|
||||
sub = "tasks"
|
||||
taskScoped = true
|
||||
case "session":
|
||||
sub = "sessions"
|
||||
case "staging":
|
||||
@@ -61,6 +63,17 @@ func (s *Server) chatUpload(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, "非法 id")
|
||||
return
|
||||
}
|
||||
if taskScoped {
|
||||
if s.m.ResolveTask(id) == nil {
|
||||
writeErr(w, 404, "task not found")
|
||||
return
|
||||
}
|
||||
if !s.engine.beginTaskOperation(id) {
|
||||
writeErr(w, http.StatusConflict, "任务正在删除,无法上传附件")
|
||||
return
|
||||
}
|
||||
defer s.engine.decInflight(id)
|
||||
}
|
||||
dir := filepath.Join(s.m.dir, sub, id, "uploads")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
writeErr(w, 500, "建目录失败: "+err.Error())
|
||||
|
||||
+75
-3
@@ -6,6 +6,8 @@ import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -92,15 +94,85 @@ func TestCoreTaskLifecyclePG(t *testing.T) {
|
||||
t.Fatalf("list tasks: %d", code)
|
||||
}
|
||||
|
||||
// delete → cascade removes the exploration subgraph
|
||||
code, _ = doRetry("DELETE", "/api/tasks/"+id, nil)
|
||||
// The task detail header consumes the top-level engine mode while some clients
|
||||
// read the active-task snapshot. Keep both representations in sync.
|
||||
code, out = doRetry("GET", "/api/stats?task="+id, nil)
|
||||
if code != 200 {
|
||||
t.Fatalf("task stats: %d (%v)", code, out)
|
||||
}
|
||||
activeTask, ok := out["active_task"].(map[string]any)
|
||||
if !ok || activeTask["engine_mode"] != out["engine_mode"] {
|
||||
t.Fatalf("engine mode mismatch: top=%v active_task=%v", out["engine_mode"], activeTask)
|
||||
}
|
||||
|
||||
// delete → cascade removes the exploration subgraph and selected related data
|
||||
taskDir := filepath.Join(m.dir, "tasks", id)
|
||||
if err := os.MkdirAll(taskDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(taskDir, "artifact.txt"), []byte("test"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
transcriptBase := "exp" + strconv.FormatInt(expID, 10) + "-main"
|
||||
transcriptPath := filepath.Join(m.dir, "transcripts", transcriptBase+".jsonl")
|
||||
sidechainPath := filepath.Join(m.dir, "transcripts", transcriptBase)
|
||||
if err := os.MkdirAll(filepath.Join(sidechainPath, "subagents"), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(transcriptPath, []byte("{}\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(sidechainPath, "subagents", "child.jsonl"), []byte("{}\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
nid, _ := strconv.ParseInt(id, 10, 64)
|
||||
findingID, err := m.pg.AddFinding(nid, 0, "__core_task_delete__", "", db.SeverityHigh, "summary", "evidence", "tester", nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := m.pg.EnsureLLMRecordsTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := m.pg.InsertLLMRecord(&db.LLMRecord{TaskID: id, SessionID: "delete-test", Status: "ok", RequestBody: "secret"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// A non-canonical path still resolves to the one canonical task key for the
|
||||
// barrier, workspace, in-memory registry and DB delete.
|
||||
code, out = doRetry("DELETE", "/api/tasks/000"+id, map[string]bool{
|
||||
"delete_files": true,
|
||||
"delete_findings": true,
|
||||
"delete_llm_records": true,
|
||||
})
|
||||
if code != 200 {
|
||||
t.Fatalf("delete task: %d", code)
|
||||
}
|
||||
nid, _ := strconv.ParseInt(id, 10, 64)
|
||||
if deleted, _ := out["files_deleted"].(bool); !deleted {
|
||||
t.Fatalf("expected files_deleted response, got %v", out)
|
||||
}
|
||||
if deleted, _ := out["findings_deleted"].(float64); deleted != 1 {
|
||||
t.Fatalf("expected one deleted finding, got %v", out)
|
||||
}
|
||||
if deleted, _ := out["llm_records_deleted"].(float64); deleted != 1 {
|
||||
t.Fatalf("expected one deleted LLM record, got %v", out)
|
||||
}
|
||||
if _, err := os.Stat(taskDir); !os.IsNotExist(err) {
|
||||
t.Fatalf("task directory should be deleted, stat err=%v", err)
|
||||
}
|
||||
for _, path := range []string{transcriptPath, sidechainPath} {
|
||||
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Fatalf("task transcript should be deleted (%s), stat err=%v", path, err)
|
||||
}
|
||||
}
|
||||
if got, _ := m.pg.GetTask(nid); got != nil {
|
||||
t.Fatalf("task should be deleted from PG")
|
||||
}
|
||||
if finding, err := m.pg.GetFinding(findingID); err != nil || finding != nil {
|
||||
t.Fatalf("finding should be deleted, got finding=%+v err=%v", finding, err)
|
||||
}
|
||||
var llmRecords int
|
||||
if err := m.pg.QueryRow(`SELECT count(*) FROM llm_records WHERE task_id=$1`, id).Scan(&llmRecords); err != nil || llmRecords != 0 {
|
||||
t.Fatalf("task LLM records should be deleted, count=%d err=%v", llmRecords, err)
|
||||
}
|
||||
var n int
|
||||
m.pg.QueryRow(`SELECT count(*) FROM exploration_nodes WHERE exploration_id=$1`, expID).Scan(&n)
|
||||
if n != 0 {
|
||||
|
||||
+114
-48
@@ -30,21 +30,27 @@ func rawString(raw json.RawMessage) string {
|
||||
|
||||
// ---- Task (frontend "Task") ---- created_at as RFC3339, plus a derived status.
|
||||
type TaskDTO struct {
|
||||
ID string `json:"id"`
|
||||
ExplorationID int64 `json:"exploration_id"`
|
||||
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"`
|
||||
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
|
||||
ID string `json:"id"`
|
||||
ExplorationID int64 `json:"exploration_id"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// TokenTotalDTO is a whole-task (all agents) token aggregate.
|
||||
@@ -65,18 +71,30 @@ func tokenTotalDTO(u db.TokenUsage) TokenTotalDTO {
|
||||
}
|
||||
|
||||
func taskDTO(t *Task, status string) TaskDTO {
|
||||
llmState := t.llmStateSnapshot()
|
||||
sourceIDs := make([]string, 0, len(t.SourceTaskIDs))
|
||||
for _, id := range t.SourceTaskIDs {
|
||||
sourceIDs = append(sourceIDs, i64s(id))
|
||||
}
|
||||
profileIDs := append(make([]int64, 0, len(llmState.ProfileIDs)), llmState.ProfileIDs...)
|
||||
return TaskDTO{
|
||||
ID: t.ID,
|
||||
ExplorationID: t.ExpID,
|
||||
Description: t.Description,
|
||||
Goal: t.Goal,
|
||||
Status: status,
|
||||
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
|
||||
CreatedUnix: t.CreatedAt,
|
||||
CompletedAt: completedRFC(t.CompletedAt),
|
||||
CompletedUnix: t.CompletedAt,
|
||||
Paused: t.Paused,
|
||||
LLMProfileID: t.LLMProfileID,
|
||||
ID: t.ID,
|
||||
ExplorationID: t.ExpID,
|
||||
Description: t.Description,
|
||||
Goal: t.Goal,
|
||||
Status: status,
|
||||
CreatedAt: rfc3339(time.Unix(t.CreatedAt, 0)),
|
||||
CreatedUnix: t.CreatedAt,
|
||||
CompletedAt: completedRFC(t.CompletedAt),
|
||||
CompletedUnix: t.CompletedAt,
|
||||
Paused: t.Paused,
|
||||
Queued: t.Queued,
|
||||
LLMProfileID: llmState.ProfileID,
|
||||
LLMProfileIDs: profileIDs,
|
||||
ActiveLLMProfileID: llmState.ActiveID,
|
||||
LLMFailoverState: llmState.FailoverState,
|
||||
LLMFailoverReason: llmState.FailoverReason,
|
||||
SourceTaskIDs: sourceIDs,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,17 +138,19 @@ func trafficDTOs(ex []traffic.ExchangeMeta) []TrafficExchangeDTO {
|
||||
// ---- TaskNode (frontend "TaskNode") — frontier, intents, exploration graph nodes ----
|
||||
|
||||
type TaskNodeDTO struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // db Node.Kind
|
||||
Payload string `json:"payload,omitempty"`
|
||||
Priority int `json:"priority"`
|
||||
State string `json:"state"`
|
||||
Origin string `json:"origin"`
|
||||
TS string `json:"ts"`
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // db Node.Kind
|
||||
Payload string `json:"payload,omitempty"`
|
||||
Priority int `json:"priority"`
|
||||
State string `json:"state"`
|
||||
Origin string `json:"origin"`
|
||||
TS string `json:"ts"`
|
||||
SourceTaskID string `json:"source_task_id,omitempty"`
|
||||
Inherited bool `json:"inherited,omitempty"`
|
||||
}
|
||||
|
||||
func taskNodeDTO(n *db.Node) TaskNodeDTO {
|
||||
return TaskNodeDTO{
|
||||
d := TaskNodeDTO{
|
||||
ID: i64s(n.ID),
|
||||
Type: n.Kind,
|
||||
Payload: rawString(n.Payload),
|
||||
@@ -139,6 +159,11 @@ func taskNodeDTO(n *db.Node) TaskNodeDTO {
|
||||
Origin: n.Origin,
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
if n.SourceTaskID > 0 {
|
||||
d.SourceTaskID = i64s(n.SourceTaskID)
|
||||
}
|
||||
d.Inherited = n.Inherited
|
||||
return d
|
||||
}
|
||||
|
||||
func taskNodeDTOs(in []*db.Node) []TaskNodeDTO {
|
||||
@@ -169,6 +194,27 @@ func edgeDTOs(in []db.Edge) []EdgeDTO {
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Coverage asset references ----
|
||||
|
||||
type CoverageAssetRefDTO struct {
|
||||
ID int64 `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
State string `json:"state"`
|
||||
Summary string `json:"summary"`
|
||||
SourceTaskID string `json:"source_task_id,omitempty"`
|
||||
Inherited bool `json:"inherited,omitempty"`
|
||||
}
|
||||
|
||||
func coverageAssetRefDTO(ref db.AssetRef) CoverageAssetRefDTO {
|
||||
out := CoverageAssetRefDTO{
|
||||
ID: ref.ID, Kind: ref.Kind, State: ref.State, Summary: ref.Summary, Inherited: ref.Inherited,
|
||||
}
|
||||
if ref.SourceTaskID > 0 {
|
||||
out.SourceTaskID = i64s(ref.SourceTaskID)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- Finding (frontend "Finding") ----
|
||||
|
||||
type FindingDTO struct {
|
||||
@@ -186,6 +232,8 @@ type FindingDTO struct {
|
||||
ParamID string `json:"param_id,omitempty"`
|
||||
TaskID string `json:"task_id,omitempty"`
|
||||
TaskDescription string `json:"task_description,omitempty"`
|
||||
SourceTaskID string `json:"source_task_id,omitempty"`
|
||||
Inherited bool `json:"inherited,omitempty"`
|
||||
Assets []FindingAssetDTO `json:"assets,omitempty"`
|
||||
TS string `json:"ts"`
|
||||
}
|
||||
@@ -241,7 +289,7 @@ type findingPayload struct {
|
||||
func findingDTO(n *db.Node) FindingDTO {
|
||||
var p findingPayload
|
||||
_ = json.Unmarshal(n.Payload, &p)
|
||||
return FindingDTO{
|
||||
d := FindingDTO{
|
||||
ID: i64s(n.ID),
|
||||
VulnClass: p.VulnClass,
|
||||
Name: p.Name,
|
||||
@@ -251,6 +299,11 @@ func findingDTO(n *db.Node) FindingDTO {
|
||||
Evidence: rawString(p.Evidence),
|
||||
TS: rfc3339(n.CreatedAt),
|
||||
}
|
||||
if n.SourceTaskID > 0 {
|
||||
d.SourceTaskID = i64s(n.SourceTaskID)
|
||||
}
|
||||
d.Inherited = n.Inherited
|
||||
return d
|
||||
}
|
||||
|
||||
// findingDTOsForTask converts a task's finding nodes to DTOs, stamping each with
|
||||
@@ -260,11 +313,15 @@ func findingDTO(n *db.Node) FindingDTO {
|
||||
// nodes with no row keep the 'pending' default and no finding_id (not editable).
|
||||
// assets pre-resolves the anchored asset rows for label rendering.
|
||||
func findingDTOsForTask(t *Task, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO {
|
||||
return findingDTOsForOwner(t.ID, t.Description, in, meta, assets)
|
||||
}
|
||||
|
||||
func findingDTOsForOwner(taskID, description string, in []*db.Node, meta map[int64]db.FindingMeta, assets map[int64]*db.Asset) []FindingDTO {
|
||||
out := make([]FindingDTO, 0, len(in))
|
||||
for _, n := range in {
|
||||
d := findingDTO(n)
|
||||
d.TaskID = t.ID
|
||||
d.TaskDescription = t.Description
|
||||
d.TaskID = taskID
|
||||
d.TaskDescription = description
|
||||
if m, ok := meta[n.ID]; ok {
|
||||
d.FindingID = i64s(m.ID)
|
||||
d.Status = m.Status
|
||||
@@ -320,16 +377,19 @@ func findingFromDB(f *db.DBFinding, assets map[int64]*db.Asset) FindingDTO {
|
||||
// ---- Activity (frontend "Activity") ----
|
||||
|
||||
type ActivityDTO struct {
|
||||
Seq int64 `json:"seq"` // db Activity.ID
|
||||
IntentID string `json:"intent_id,omitempty"` // db NodeID
|
||||
Worker string `json:"worker"`
|
||||
TS string `json:"ts"` // db CreatedAt
|
||||
Kind string `json:"kind"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
IsError bool `json:"is_error"`
|
||||
Summary string `json:"summary"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
Seq int64 `json:"seq"` // db Activity.ID
|
||||
IntentID string `json:"intent_id,omitempty"` // db NodeID
|
||||
Worker string `json:"worker"`
|
||||
TS string `json:"ts"` // db CreatedAt
|
||||
Kind string `json:"kind"`
|
||||
Tool string `json:"tool,omitempty"`
|
||||
ToolUseID string `json:"tool_use_id,omitempty"`
|
||||
IsError bool `json:"is_error"`
|
||||
Summary string `json:"summary"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
Metadata json.RawMessage `json:"metadata,omitempty"`
|
||||
SourceTaskID string `json:"source_task_id,omitempty"`
|
||||
Inherited bool `json:"inherited,omitempty"`
|
||||
// token usage (set only on kind='result'); used for per-session token totals.
|
||||
InputTokens *int `json:"input_tokens,omitempty"`
|
||||
OutputTokens *int `json:"output_tokens,omitempty"`
|
||||
@@ -342,7 +402,7 @@ func activityDTO(a db.Activity) ActivityDTO {
|
||||
if a.NodeID != nil {
|
||||
intent = i64s(*a.NodeID)
|
||||
}
|
||||
return ActivityDTO{
|
||||
d := ActivityDTO{
|
||||
Seq: a.ID,
|
||||
IntentID: intent,
|
||||
Worker: a.Worker,
|
||||
@@ -353,11 +413,17 @@ func activityDTO(a db.Activity) ActivityDTO {
|
||||
IsError: a.IsError,
|
||||
Summary: a.Summary,
|
||||
Detail: a.Detail, // list endpoint leaves this empty (lazy)
|
||||
Metadata: a.Metadata,
|
||||
InputTokens: a.InputTokens,
|
||||
OutputTokens: a.OutputTokens,
|
||||
CacheReadTokens: a.CacheReadTokens,
|
||||
CacheWriteTokens: a.CacheWriteTokens,
|
||||
}
|
||||
if a.SourceTaskID > 0 {
|
||||
d.SourceTaskID = i64s(a.SourceTaskID)
|
||||
}
|
||||
d.Inherited = a.Inherited
|
||||
return d
|
||||
}
|
||||
|
||||
func activityDTOs(in []db.Activity) []ActivityDTO {
|
||||
|
||||
+427
-137
@@ -73,6 +73,13 @@ const (
|
||||
modelErrorRetryBackoff = 3 * time.Second // 每次重试前的退避
|
||||
)
|
||||
|
||||
// retryableWorkerModelError excludes errors already handled by the task router.
|
||||
// In particular, a quota error after partial streaming advances the task cursor
|
||||
// for the next LLM call but must not replay this whole intent on the backup.
|
||||
func retryableWorkerModelError(reason harness.TerminalReason, err error) bool {
|
||||
return reason == harness.ReasonModelError && !isTaskLLMRuntimeError(err)
|
||||
}
|
||||
|
||||
// Engine drives the event-driven exploration loop with real LLM agents
|
||||
// (docs §4.3/§4.4): on asset/exploration-graph change (debounced) it wakes the
|
||||
// planner, which reads the route, queries assets, judges goals and emits intents;
|
||||
@@ -86,28 +93,43 @@ type Engine struct {
|
||||
|
||||
bc *Broadcaster // live activity pub/sub (SSE)
|
||||
|
||||
mu sync.RWMutex
|
||||
planner *agent.Planner
|
||||
worker *agent.Worker
|
||||
started sync.Map // taskID -> bool, so Run is idempotent per task
|
||||
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
||||
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
|
||||
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
|
||||
mu sync.RWMutex
|
||||
planner *agent.Planner
|
||||
worker *agent.Worker
|
||||
started sync.Map // taskID -> bool, so Run is idempotent per task
|
||||
lastAct sync.Map // taskID -> int64 unix, last planner/worker activity (heartbeat)
|
||||
llmCalls sync.Map // taskID -> *int64, actual planner/worker/main-agent LLM calls
|
||||
paused sync.Map // taskID -> bool, user-paused (planner + workers idle but loops alive)
|
||||
deleting sync.Map // taskID -> bool, delete barrier (no new task-owned writes)
|
||||
// deleteWasPaused remembers whether BeginDelete temporarily paused an already
|
||||
// running task. AbortDelete uses it to restore the exact pre-delete state.
|
||||
deleteWasPaused sync.Map // taskID -> bool
|
||||
dropCnt sync.Map // taskID -> *int64, running count of dropped (unpersistable) activity records
|
||||
|
||||
// deleteMu makes installing the delete barrier atomic with registering a new
|
||||
// task operation. Once BeginDelete returns, every admitted writer is reflected
|
||||
// in inflight and every later writer is rejected.
|
||||
deleteMu sync.RWMutex
|
||||
|
||||
// Every long-lived task goroutine (planner, workers and deadline coordinator)
|
||||
// runs under one task-scoped context. Successful deletion cancels that context,
|
||||
// waits for all goroutines, then releases every task-level Engine reference.
|
||||
runtimeMu sync.Mutex
|
||||
runtimes map[string]*taskRuntime
|
||||
|
||||
// per-task execution context: each planner.Plan / worker.Execute runs under it,
|
||||
// so pausing can CANCEL an in-flight run (not just skip the next one). Recreated
|
||||
// on resume since cancelling is one-shot.
|
||||
// CancelCauseFunc (not CancelFunc): every cancel site attaches a named
|
||||
// *agent.AbortCause so the run's trace can say WHICH of the ~10 cancel
|
||||
// paths fired instead of one catch-all "运行被中断" line.
|
||||
// on resume since cancelling is one-shot. Every cancellation carries a named
|
||||
// cause so the activity trace can identify the initiating control path.
|
||||
execMu sync.Mutex
|
||||
execCancel map[string]context.CancelCauseFunc
|
||||
execCtx map[string]context.Context
|
||||
|
||||
// per-work (per running intent) cancel, so the planner's kill_work can stop a
|
||||
// single worker without pausing the whole task. Keyed by globally-unique intent id.
|
||||
workMu sync.Mutex
|
||||
workCancel map[int64]context.CancelCauseFunc
|
||||
// Per-work control lets the planner kill a worker and lets the UI pause/cancel
|
||||
// one intent without pausing the whole task. The done channel closes only after
|
||||
// runWorkerStep has stopped writing and committed its final state.
|
||||
workMu sync.Mutex
|
||||
work map[int64]*workExecution
|
||||
|
||||
// steerBox queues planner course-corrections for a running work (keyed by intent
|
||||
// id). The worker's PreToolUse hook drains it before its next tool call and hands
|
||||
@@ -124,9 +146,23 @@ type Engine struct {
|
||||
inflight sync.Map // taskID -> *int64, 在跑的 planner.Plan + worker.Execute 计数(用于 drain)
|
||||
coordStarted sync.Map // taskID -> bool, deadline 协调器是否已启动(Run/reload 去重)
|
||||
|
||||
// resolve, if set, returns a task's dedicated planner/worker when it pins a
|
||||
// specific LLM profile (nil,nil = use the global active pair). Set once at startup.
|
||||
resolve func(t *Task) (*agent.Planner, *agent.Worker)
|
||||
// resolve, if set, returns a task's dedicated planner/worker. In legacy mode
|
||||
// nil,nil falls back to the global pair; an authoritative resolver uses nil,nil
|
||||
// to mean this task is deliberately unavailable (for example an exhausted chain).
|
||||
resolve func(t *Task) (*agent.Planner, *agent.Worker)
|
||||
resolveAuthoritative bool
|
||||
}
|
||||
|
||||
type taskRuntime struct {
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
type workExecution struct {
|
||||
cancel context.CancelCauseFunc
|
||||
done chan struct{}
|
||||
action string // user action: pause | cancel
|
||||
}
|
||||
|
||||
// nextPlannerRound returns the next planner round number for a task (1-based).
|
||||
@@ -139,17 +175,132 @@ func (e *Engine) nextPlannerRound(taskID string) int {
|
||||
|
||||
// Pause stops a task: marks it paused AND cancels any in-flight planner/worker run
|
||||
// for it (a long worker.Execute would otherwise keep going until it finishes).
|
||||
// cause says WHY (user button / orchestrator tool / delete / reload) — it surfaces
|
||||
// verbatim in the cancelled run's trace, so pass the matching agent.Abort* value.
|
||||
func (e *Engine) Pause(taskID string, cause error) {
|
||||
e.paused.Store(taskID, true)
|
||||
e.cancelExec(taskID, cause)
|
||||
}
|
||||
|
||||
// BeginDelete installs an execution barrier before task data/files are removed.
|
||||
// The temporary pause is not a user pause: AbortDelete restores the exact state
|
||||
// observed here when cleanup fails.
|
||||
func (e *Engine) BeginDelete(taskID string) bool {
|
||||
e.deleteMu.Lock()
|
||||
if _, loaded := e.deleting.LoadOrStore(taskID, true); loaded {
|
||||
e.deleteMu.Unlock()
|
||||
return false
|
||||
}
|
||||
wasPaused := e.IsPaused(taskID)
|
||||
e.deleteWasPaused.Store(taskID, wasPaused)
|
||||
e.paused.Store(taskID, true)
|
||||
e.deleteMu.Unlock()
|
||||
e.cancelExec(taskID, agent.AbortTaskDeleted)
|
||||
return true
|
||||
}
|
||||
|
||||
func (e *Engine) AbortDelete(taskID string) {
|
||||
e.deleteMu.Lock()
|
||||
if !e.IsDeleting(taskID) {
|
||||
e.deleteMu.Unlock()
|
||||
return
|
||||
}
|
||||
wasPaused := false
|
||||
if v, ok := e.deleteWasPaused.Load(taskID); ok {
|
||||
wasPaused, _ = v.(bool)
|
||||
}
|
||||
e.deleteWasPaused.Delete(taskID)
|
||||
e.deleting.Delete(taskID)
|
||||
if !wasPaused {
|
||||
e.paused.Delete(taskID)
|
||||
}
|
||||
e.deleteMu.Unlock()
|
||||
if !wasPaused && e.m != nil {
|
||||
if t, ok := e.m.Task(taskID); ok {
|
||||
t.Notify()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) IsDeleting(taskID string) bool {
|
||||
_, ok := e.deleting.Load(taskID)
|
||||
return ok
|
||||
}
|
||||
|
||||
// registerTaskRoutines reserves count goroutines in the task runtime. Callers
|
||||
// hold deleteMu for reading so StopTask cannot race WaitGroup.Add with Wait.
|
||||
func (e *Engine) registerTaskRoutines(parent context.Context, taskID string, count int) *taskRuntime {
|
||||
e.runtimeMu.Lock()
|
||||
defer e.runtimeMu.Unlock()
|
||||
rt := e.runtimes[taskID]
|
||||
if rt == nil {
|
||||
ctx, cancel := context.WithCancel(parent)
|
||||
rt = &taskRuntime{ctx: ctx, cancel: cancel}
|
||||
e.runtimes[taskID] = rt
|
||||
}
|
||||
rt.wg.Add(count)
|
||||
return rt
|
||||
}
|
||||
|
||||
func runTaskRoutine(rt *taskRuntime, fn func(context.Context)) {
|
||||
go func() {
|
||||
defer rt.wg.Done()
|
||||
fn(rt.ctx)
|
||||
}()
|
||||
}
|
||||
|
||||
// StopTask permanently stops every long-lived goroutine and removes all Engine
|
||||
// state for a successfully deleted task. The delete barrier remains installed
|
||||
// until cleanup finishes, so no new task operation can race the teardown.
|
||||
func (e *Engine) StopTask(taskID string) {
|
||||
e.deleteMu.Lock()
|
||||
e.deleting.Store(taskID, true)
|
||||
e.deleteMu.Unlock()
|
||||
|
||||
e.cancelExec(taskID, agent.AbortTaskDeleted)
|
||||
e.runtimeMu.Lock()
|
||||
rt := e.runtimes[taskID]
|
||||
if rt != nil {
|
||||
rt.cancel()
|
||||
}
|
||||
e.runtimeMu.Unlock()
|
||||
if rt != nil {
|
||||
rt.wg.Wait()
|
||||
}
|
||||
|
||||
e.execMu.Lock()
|
||||
if cancel := e.execCancel[taskID]; cancel != nil {
|
||||
cancel(agent.AbortTaskDeleted)
|
||||
}
|
||||
delete(e.execCancel, taskID)
|
||||
delete(e.execCtx, taskID)
|
||||
e.execMu.Unlock()
|
||||
|
||||
e.runtimeMu.Lock()
|
||||
if e.runtimes[taskID] == rt {
|
||||
delete(e.runtimes, taskID)
|
||||
}
|
||||
e.runtimeMu.Unlock()
|
||||
|
||||
e.started.Delete(taskID)
|
||||
e.lastAct.Delete(taskID)
|
||||
e.llmCalls.Delete(taskID)
|
||||
e.paused.Delete(taskID)
|
||||
e.dropCnt.Delete(taskID)
|
||||
e.plannerRound.Delete(taskID)
|
||||
e.settling.Delete(taskID)
|
||||
e.deadline.Delete(taskID)
|
||||
e.stamped.Delete(taskID)
|
||||
e.inflight.Delete(taskID)
|
||||
e.coordStarted.Delete(taskID)
|
||||
e.deleteWasPaused.Delete(taskID)
|
||||
|
||||
e.deleteMu.Lock()
|
||||
e.deleting.Delete(taskID)
|
||||
e.deleteMu.Unlock()
|
||||
}
|
||||
|
||||
// cancelExec cancels a task's current per-task exec context (any in-flight
|
||||
// planner.Plan / worker.Execute), if present. Shared by Pause and the settle
|
||||
// sequence's hard-drain backstop. cause is recovered by the aborted run via
|
||||
// context.Cause and printed in its result record.
|
||||
// sequence's hard-drain backstop.
|
||||
func (e *Engine) cancelExec(taskID string, cause error) {
|
||||
e.execMu.Lock()
|
||||
if cancel := e.execCancel[taskID]; cancel != nil {
|
||||
@@ -206,32 +357,94 @@ func (e *Engine) LastActivity(taskID string) int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// BeginLLMCall/EndLLMCall track actual provider calls separately from the
|
||||
// scheduler's task-operation counter. A task can have live loops while all of
|
||||
// them are waiting for a trigger; that state must remain idle in the UI.
|
||||
func (e *Engine) BeginLLMCall(taskID string) {
|
||||
v, _ := e.llmCalls.LoadOrStore(taskID, new(int64))
|
||||
atomic.AddInt64(v.(*int64), 1)
|
||||
}
|
||||
|
||||
func (e *Engine) EndLLMCall(taskID string) {
|
||||
if v, ok := e.llmCalls.Load(taskID); ok {
|
||||
p := v.(*int64)
|
||||
if atomic.AddInt64(p, -1) <= 0 {
|
||||
atomic.StoreInt64(p, 0)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) ActiveLLMCalls(taskID string) int64 {
|
||||
if v, ok := e.llmCalls.Load(taskID); ok {
|
||||
return atomic.LoadInt64(v.(*int64))
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (e *Engine) touch(taskID string) { e.lastAct.Store(taskID, time.Now().Unix()) }
|
||||
|
||||
func NewEngine(m *Manager) *Engine {
|
||||
return &Engine{m: m, debounce: 800 * time.Millisecond, bc: NewBroadcaster(),
|
||||
execCancel: map[string]context.CancelCauseFunc{}, execCtx: map[string]context.Context{},
|
||||
workCancel: map[int64]context.CancelCauseFunc{}, steerBox: map[int64][]string{}}
|
||||
work: map[int64]*workExecution{}, steerBox: map[int64][]string{},
|
||||
runtimes: map[string]*taskRuntime{}}
|
||||
}
|
||||
|
||||
// registerWork records the cancel for the work currently running intentID.
|
||||
func (e *Engine) registerWork(intentID int64, cancel context.CancelCauseFunc) {
|
||||
e.workMu.Lock()
|
||||
e.workCancel[intentID] = cancel
|
||||
e.work[intentID] = &workExecution{cancel: cancel, done: make(chan struct{})}
|
||||
e.workMu.Unlock()
|
||||
}
|
||||
|
||||
// unregisterWork removes and releases the work's cancel (called when Execute returns).
|
||||
func (e *Engine) unregisterWork(intentID int64) {
|
||||
// detachWork removes the live control handle once Execute has returned. complete
|
||||
// must be called after the final intent state write so a waiting cancel handler can
|
||||
// safely delete the worker's blackboard output without racing a late write.
|
||||
func (e *Engine) detachWork(intentID int64) (action string, complete func()) {
|
||||
e.workMu.Lock()
|
||||
if c := e.workCancel[intentID]; c != nil {
|
||||
delete(e.workCancel, intentID)
|
||||
c(agent.AbortWorkFinished) // release context resources (no-op if already cancelled)
|
||||
run := e.work[intentID]
|
||||
if run != nil {
|
||||
delete(e.work, intentID)
|
||||
action = run.action
|
||||
run.cancel(agent.AbortWorkFinished) // release resources (no-op if already cancelled)
|
||||
}
|
||||
e.workMu.Unlock()
|
||||
e.steerMu.Lock()
|
||||
delete(e.steerBox, intentID) // drop any undelivered steering for a finished work
|
||||
e.steerMu.Unlock()
|
||||
if run == nil {
|
||||
return action, func() {}
|
||||
}
|
||||
return action, func() { close(run.done) }
|
||||
}
|
||||
|
||||
// ControlWork requests a user-visible pause or cancellation and waits until the
|
||||
// worker has fully stopped writing. Cancellation cleanup is performed by the API
|
||||
// handler after this returns; pause state is committed by runWorkerStep itself.
|
||||
func (e *Engine) ControlWork(intentID int64, action string) error {
|
||||
if action != "pause" && action != "cancel" {
|
||||
return fmt.Errorf("unsupported work action %q", action)
|
||||
}
|
||||
e.workMu.Lock()
|
||||
run := e.work[intentID]
|
||||
if run == nil {
|
||||
e.workMu.Unlock()
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
||||
}
|
||||
if run.action != "" && run.action != action {
|
||||
e.workMu.Unlock()
|
||||
return fmt.Errorf("意图 %d 正在执行 %s 操作", intentID, run.action)
|
||||
}
|
||||
run.action = action
|
||||
done := run.done
|
||||
cause := error(agent.AbortWorkPausedByUser)
|
||||
if action == "cancel" {
|
||||
cause = agent.AbortWorkCancelledByUser
|
||||
}
|
||||
run.cancel(cause)
|
||||
e.workMu.Unlock()
|
||||
<-done
|
||||
return nil
|
||||
}
|
||||
|
||||
// SteerWork queues a mid-run course-correction for the work running intentID (the
|
||||
@@ -242,7 +455,7 @@ func (e *Engine) SteerWork(intentID int64, msg string) error {
|
||||
return fmt.Errorf("纠偏消息不能为空")
|
||||
}
|
||||
e.workMu.Lock()
|
||||
running := e.workCancel[intentID] != nil
|
||||
running := e.work[intentID] != nil
|
||||
e.workMu.Unlock()
|
||||
if !running {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
||||
@@ -307,12 +520,12 @@ func (h steerHooks) Stop(ctx context.Context, messages []llm.Message) (bool, []s
|
||||
// The work's agent-core session honors ctx cancellation and aborts promptly.
|
||||
func (e *Engine) KillWork(intentID int64) error {
|
||||
e.workMu.Lock()
|
||||
c := e.workCancel[intentID]
|
||||
run := e.work[intentID]
|
||||
e.workMu.Unlock()
|
||||
if c == nil {
|
||||
if run == nil {
|
||||
return fmt.Errorf("意图 %d 当前没有运行中的 work(可能已结束或未被领取)", intentID)
|
||||
}
|
||||
c(agent.AbortKilledByPlanner)
|
||||
run.cancel(agent.AbortKilledByPlanner)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -321,7 +534,7 @@ func (e *Engine) Broadcaster() *Broadcaster { return e.bc }
|
||||
|
||||
// emitActivity persists one captured step AND fans it out to live subscribers,
|
||||
// from a single point so storage and the SSE stream never diverge.
|
||||
func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
||||
func (e *Engine) emitActivity(t *Task, r db.Activity) db.Activity {
|
||||
id, err := e.appendActivity(t, r)
|
||||
if err != nil {
|
||||
// NO LONGER SILENT: dropping a record breaks command↔result pairing in the
|
||||
@@ -346,7 +559,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
||||
log.Printf("[activity] task %s 丢弃活动记录(该任务累计第 %d 条) worker=%s kind=%s tool=%s tuid=%s reason=%s summary=%q: %v%s",
|
||||
t.ID, n, r.Worker, r.Kind, r.Tool, r.ToolUseID, dropReason(err), preview(r.Summary, 80), err, diag)
|
||||
e.touch(t.ID)
|
||||
return
|
||||
return r
|
||||
}
|
||||
r.ID = id
|
||||
if r.CreatedAt.IsZero() {
|
||||
@@ -354,6 +567,7 @@ func (e *Engine) emitActivity(t *Task, r db.Activity) {
|
||||
}
|
||||
e.bc.Publish(t.ID, r)
|
||||
e.touch(t.ID)
|
||||
return r
|
||||
}
|
||||
|
||||
// appendActivity persists one activity row, retrying briefly on write failure.
|
||||
@@ -396,6 +610,15 @@ func (e *Engine) snapshot() (*agent.Planner, *agent.Worker) {
|
||||
// Called once at startup before any task loop runs, so no lock is needed on reads.
|
||||
func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) {
|
||||
e.resolve = fn
|
||||
e.resolveAuthoritative = false
|
||||
}
|
||||
|
||||
// SetAuthoritativeAgentResolver installs a resolver whose nil result must not
|
||||
// fall through to the global provider. Task-level failover chains use this so a
|
||||
// fully exhausted chain cannot silently bypass its configured boundary.
|
||||
func (e *Engine) SetAuthoritativeAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worker)) {
|
||||
e.resolve = fn
|
||||
e.resolveAuthoritative = true
|
||||
}
|
||||
|
||||
// snapshotFor returns the planner/worker a task should run on: the task's dedicated
|
||||
@@ -403,7 +626,8 @@ func (e *Engine) SetAgentResolver(fn func(t *Task) (*agent.Planner, *agent.Worke
|
||||
// pair. The resolver returning nil means "no override" (fall back to global).
|
||||
func (e *Engine) snapshotFor(t *Task) (*agent.Planner, *agent.Worker) {
|
||||
if e.resolve != nil {
|
||||
if p, w := e.resolve(t); p != nil && w != nil {
|
||||
p, w := e.resolve(t)
|
||||
if (p != nil && w != nil) || e.resolveAuthoritative {
|
||||
return p, w
|
||||
}
|
||||
}
|
||||
@@ -416,6 +640,14 @@ func (e *Engine) Ready() bool {
|
||||
return p != nil && w != nil
|
||||
}
|
||||
|
||||
// ReadyFor reports whether a specific task can resolve a planner/worker pair.
|
||||
// An explicit task profile chain can be runnable even when no global default
|
||||
// provider is configured, so task status must not rely on Ready alone.
|
||||
func (e *Engine) ReadyFor(t *Task) bool {
|
||||
p, w := e.snapshotFor(t)
|
||||
return p != nil && w != nil
|
||||
}
|
||||
|
||||
// Mode reports "llm" when ready, else "idle".
|
||||
func (e *Engine) Mode() string {
|
||||
if e.Ready() {
|
||||
@@ -428,15 +660,24 @@ func (e *Engine) Mode() string {
|
||||
// but no-op until an LLM is configured (so a task created while idle picks up
|
||||
// automatically once LLM is set from the UI).
|
||||
func (e *Engine) Run(ctx context.Context, t *Task) {
|
||||
workers := e.m.Workers()
|
||||
e.deleteMu.RLock()
|
||||
if e.IsDeleting(t.ID) {
|
||||
e.deleteMu.RUnlock()
|
||||
return
|
||||
}
|
||||
if _, loaded := e.started.LoadOrStore(t.ID, true); loaded {
|
||||
e.deleteMu.RUnlock()
|
||||
t.Notify() // already running — just nudge a planning round
|
||||
return
|
||||
}
|
||||
rt := e.registerTaskRoutines(ctx, t.ID, 1+workers)
|
||||
e.deleteMu.RUnlock()
|
||||
e.touch(t.ID)
|
||||
go e.plannerLoop(ctx, t)
|
||||
workers := e.m.Workers()
|
||||
runTaskRoutine(rt, func(loopCtx context.Context) { e.plannerLoop(loopCtx, t) })
|
||||
for i := 0; i < workers; i++ {
|
||||
go e.workerLoop(ctx, t, fmt.Sprintf("work#%d", i+1))
|
||||
name := fmt.Sprintf("work#%d", i+1)
|
||||
runTaskRoutine(rt, func(loopCtx context.Context) { e.workerLoop(loopCtx, t, name) })
|
||||
}
|
||||
e.startDeadlineCoordinator(ctx, t) // 任务级超时定时器(仅 timeout>0;去重)
|
||||
// 仅在「完全没有活动意图(open+running)」时才 kick 首轮规划。带种子意图的任务:种子已
|
||||
@@ -497,6 +738,9 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
||||
if e.IsPaused(t.ID) {
|
||||
return // user-paused: don't plan
|
||||
}
|
||||
if e.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
// terminal task (goals all met → done, or failed): the run is over. A
|
||||
// resume/nudge — e.g. auto-resume of the active task on restart — must NOT
|
||||
// re-plan (it would burn an LLM round and re-confirm a settled result).
|
||||
@@ -508,10 +752,17 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
||||
if e.isSettling(t.ID) {
|
||||
return
|
||||
}
|
||||
if !e.beginTaskOperation(t.ID) {
|
||||
return
|
||||
}
|
||||
defer e.decInflight(t.ID)
|
||||
e.stampFirstRun(t) // 首次真正规划 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
if ectx.Err() != nil || e.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
log.Printf("[planner] task %s 规划中…(%s 触发)", t.ID, src)
|
||||
// round marker: each Plan() is one planner round; emit a boundary so the
|
||||
// UI can separate rounds in the transcript (kind='round').
|
||||
@@ -520,10 +771,10 @@ func (e *Engine) plannerLoop(ctx context.Context, t *Task) {
|
||||
// what fired this round (worker done / finding; may be several — debounce
|
||||
// coalesces a burst; empty for time/heartbeat wakes).
|
||||
triggers := t.drainTriggers()
|
||||
e.incInflight(t.ID)
|
||||
taskIDInt, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.BeginLLMCall(t.ID)
|
||||
met, reason, err := planner.Plan(ectx, taskIDInt, e.m.assets, t.Store, t.Goal, triggers, emit)
|
||||
e.decInflight(t.ID)
|
||||
e.EndLLMCall(t.ID)
|
||||
switch {
|
||||
case err != nil && ectx.Err() == nil:
|
||||
log.Printf("[planner] task %s 规划出错: %v", t.ID, err)
|
||||
@@ -578,6 +829,9 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) {
|
||||
}
|
||||
continue // user-paused: don't claim/execute intents
|
||||
}
|
||||
if e.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
if e.isSettling(t.ID) {
|
||||
if sleepCtx(ctx, 1000*time.Millisecond) {
|
||||
return
|
||||
@@ -590,108 +844,144 @@ func (e *Engine) workerLoop(ctx context.Context, t *Task, name string) {
|
||||
}
|
||||
continue // 任务已终态(done/failed/timeout):停止领取遗留意图,别在完成后空跑 frontier
|
||||
}
|
||||
intent := e.claimNext(t, name)
|
||||
if intent == nil {
|
||||
if sleepCtx(ctx, 800*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID)
|
||||
e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
// per-work child context so the planner's kill_work can stop just this work.
|
||||
workCtx, workCancel := context.WithCancelCause(ectx)
|
||||
e.registerWork(intent.ID, workCancel)
|
||||
// wrap the guard hooks so steer_work can inject a mid-run course-correction
|
||||
// for THIS intent (drained before the worker's next tool call).
|
||||
iid := intent.ID
|
||||
taskEmit := func(a db.Activity) {
|
||||
nid := iid
|
||||
a.NodeID, a.Worker = &nid, name
|
||||
emit(a)
|
||||
}
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit)
|
||||
hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }}
|
||||
wTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.incInflight(t.ID) // 计入在跑,供收尾时序 drain 等待
|
||||
reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding)
|
||||
// model_error 收场 → 额外重跑几次(退避后再试)。仅在意图仍属本 work、任务
|
||||
// 未暂停/未终止/未取消【且未进入收尾】时重试;否则让位给对应分支处理(收尾期不
|
||||
// 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。
|
||||
for attempt := 1; attempt <= modelErrorRetries &&
|
||||
reason == harness.ReasonModelError &&
|
||||
workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ {
|
||||
log.Printf("[worker %s] task %s 意图 #%d model_error 收场,%v 后重试 (%d/%d)",
|
||||
name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries)
|
||||
if sleepCtx(workCtx, modelErrorRetryBackoff) {
|
||||
break // 退避期间被取消(终止/暂停)→ 交给下方分支处理
|
||||
}
|
||||
reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding)
|
||||
if !e.beginTaskOperation(t.ID) {
|
||||
return
|
||||
}
|
||||
claimed := e.runWorkerStep(ctx, t, name, worker)
|
||||
e.decInflight(t.ID)
|
||||
// CAPTURE kill state BEFORE unregisterWork cancels workCtx. kill = this work's
|
||||
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
||||
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
|
||||
// unregister would always be true (unregister cancels it) → every completed
|
||||
// work would be wrongly marked stopped.
|
||||
killed := workCtx.Err() != nil && ectx.Err() == nil
|
||||
e.unregisterWork(intent.ID)
|
||||
// if a pause cancelled this run mid-flight, return the intent to the frontier
|
||||
// so it is re-claimed (and re-run from scratch) on resume — not marked done.
|
||||
if ectx.Err() != nil && e.IsPaused(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "open")
|
||||
continue
|
||||
if !claimed && sleepCtx(ctx, 800*time.Millisecond) {
|
||||
return
|
||||
}
|
||||
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
|
||||
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
|
||||
if ectx.Err() != nil && e.isSettling(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "exhausted")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
|
||||
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
|
||||
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
continue
|
||||
}
|
||||
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
|
||||
if killed {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
t.Notify()
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
log.Printf("[worker %s] intent %d: %v", name, intent.ID, err)
|
||||
}
|
||||
// terminal分流:撞步数上限 ≠ 完成。max_turns→exhausted(规划者据此知道这个方向
|
||||
// 试过但没真正做完、需换角度,而非当成已覆盖永久跳过);出错→blocked;正常→done。
|
||||
state := "done"
|
||||
switch {
|
||||
case err != nil:
|
||||
state = "blocked"
|
||||
case reason == harness.ReasonMaxTurns:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 撞步数上限(exhausted),本次写回 %s", name, intent.ID, wrote)
|
||||
case reason == harness.ReasonTimeout:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 运行超时(exhausted),收尾后写回 %s", name, intent.ID, wrote)
|
||||
}
|
||||
_ = t.Store.SetIntentState(intent.ID, state)
|
||||
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote)
|
||||
e.touch(t.ID)
|
||||
t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id)
|
||||
}
|
||||
}
|
||||
|
||||
// runWorkerStep claims and fully settles at most one intent. Its caller holds one
|
||||
// task-operation admission for the whole claim/execute/state-write sequence so a
|
||||
// delete cannot observe quiescence between the LLM return and the final DB writes.
|
||||
func (e *Engine) runWorkerStep(ctx context.Context, t *Task, name string, worker *agent.Worker) bool {
|
||||
intent := e.claimNext(t, name)
|
||||
if intent == nil {
|
||||
return false
|
||||
}
|
||||
log.Printf("[worker %s] task %s 领取意图 #%d", name, t.ID, intent.ID)
|
||||
e.stampFirstRun(t) // 首次真正执行 → 盖 first_run_at + 算 deadline(仅带 timeout 的任务)
|
||||
e.touch(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
ectx := e.clockCtx(e.execContextFor(ctx, t.ID), t, false) // cancellable by Pause; 带任务 deadline
|
||||
if ectx.Err() != nil || e.IsDeleting(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "open")
|
||||
return true
|
||||
}
|
||||
// per-work child context so the planner's kill_work can stop just this work.
|
||||
workCtx, workCancel := context.WithCancelCause(ectx)
|
||||
e.registerWork(intent.ID, workCancel)
|
||||
// wrap the guard hooks so steer_work can inject a mid-run course-correction
|
||||
// for THIS intent (drained before the worker's next tool call).
|
||||
iid := intent.ID
|
||||
taskEmit := func(a db.Activity) {
|
||||
nid := iid
|
||||
a.NodeID, a.Worker = &nid, name
|
||||
emit(a)
|
||||
}
|
||||
workCtx = intercept.WithTaskContext(workCtx, t.ID, fmt.Sprintf("%s · #%d", name, iid), taskEmit)
|
||||
hooks := steerHooks{inner: t.Guard.Hooks(), drain: func() (string, bool) { return e.drainSteer(iid) }}
|
||||
wTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.BeginLLMCall(t.ID)
|
||||
reason, wrote, err := worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding)
|
||||
e.EndLLMCall(t.ID)
|
||||
// model_error 收场 → 额外重跑几次(退避后再试)。仅在意图仍属本 work、任务
|
||||
// 未暂停/未终止/未取消【且未进入收尾】时重试;否则让位给对应分支处理(收尾期不
|
||||
// 再重试,避免退避挤占其他 worker 的优雅收尾窗口)。
|
||||
for attempt := 1; attempt <= modelErrorRetries &&
|
||||
retryableWorkerModelError(reason, err) &&
|
||||
workCtx.Err() == nil && ectx.Err() == nil && !e.IsPaused(t.ID) && !e.isSettling(t.ID); attempt++ {
|
||||
log.Printf("[worker %s] task %s 意图 #%d model_error 收场,%v 后重试 (%d/%d)",
|
||||
name, t.ID, intent.ID, modelErrorRetryBackoff, attempt, modelErrorRetries)
|
||||
if sleepCtx(workCtx, modelErrorRetryBackoff) {
|
||||
break // 退避期间被取消(终止/暂停)→ 交给下方分支处理
|
||||
}
|
||||
e.BeginLLMCall(t.ID)
|
||||
reason, wrote, err = worker.Execute(workCtx, name, wTaskID, e.m.assets, t.Store, intent, hooks, emit, e.m.enrich, t.NotifyFinding)
|
||||
e.EndLLMCall(t.ID)
|
||||
}
|
||||
// Capture kill state before detachWork cancels workCtx. kill = this work's
|
||||
// ctx was cancelled (planner kill_work) while the TASK ctx kept running; a
|
||||
// pause cancels the task ctx (ectx) instead. Checking workCtx.Err() AFTER
|
||||
// unregister would always be true (unregister cancels it) → every completed
|
||||
// work would be wrongly marked stopped.
|
||||
killed := workCtx.Err() != nil && ectx.Err() == nil
|
||||
action, completeWork := e.detachWork(intent.ID)
|
||||
defer completeWork()
|
||||
if action == "pause" {
|
||||
_ = t.Store.SetIntentState(intent.ID, "paused")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 已暂停", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
return true
|
||||
}
|
||||
if action == "cancel" {
|
||||
// Leave the node in running until the waiting API handler deletes the intent
|
||||
// and its direct outputs transactionally. No task Notify is emitted here.
|
||||
log.Printf("[worker %s] task %s 意图 #%d 已停止,等待取消清理", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
return true
|
||||
}
|
||||
// if a pause cancelled this run mid-flight, return the intent to the frontier
|
||||
// so it is re-claimed (and re-run from scratch) on resume — not marked done.
|
||||
if ectx.Err() != nil && e.IsPaused(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "open")
|
||||
return true
|
||||
}
|
||||
// 任务超时收尾的硬兜底 cancel(非 pause、非 kill)取消了本 run → 归为 exhausted(已收尾),
|
||||
// 不要误标 blocked。此时 worker 通常已在 settlement 阶段把结果写回。
|
||||
if ectx.Err() != nil && e.isSettling(t.ID) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "exhausted")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务超时收尾结束(exhausted),写回 %s", name, t.ID, intent.ID, wrote)
|
||||
e.touch(t.ID)
|
||||
return true
|
||||
}
|
||||
// 任务已判完成(done via 常规路径)→ 上面 cancelExec 取消了本 run。意图结果已无意义,
|
||||
// 标 stopped(不是 blocked),别污染已完成任务的意图状态。
|
||||
if ectx.Err() != nil && isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 因任务已完成而取消(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
return true
|
||||
}
|
||||
// killed by the planner: mark stopped (don't write back results, don't auto-reclaim).
|
||||
if killed {
|
||||
_ = t.Store.SetIntentState(intent.ID, "stopped")
|
||||
log.Printf("[worker %s] task %s 意图 #%d 被终止(stopped)", name, t.ID, intent.ID)
|
||||
e.touch(t.ID)
|
||||
t.Notify()
|
||||
return true
|
||||
}
|
||||
if err != nil {
|
||||
log.Printf("[worker %s] intent %d: %v", name, intent.ID, err)
|
||||
}
|
||||
// terminal分流:撞步数上限 ≠ 完成。max_turns→exhausted(规划者据此知道这个方向
|
||||
// 试过但没真正做完、需换角度,而非当成已覆盖永久跳过);出错→blocked;正常→done。
|
||||
state := "done"
|
||||
switch {
|
||||
case err != nil:
|
||||
state = "blocked"
|
||||
case reason == harness.ReasonMaxTurns:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 撞步数上限(exhausted),本次写回 %s", name, intent.ID, wrote)
|
||||
case reason == harness.ReasonTimeout:
|
||||
state = "exhausted"
|
||||
log.Printf("[worker %s] intent %d 运行超时(exhausted),收尾后写回 %s", name, intent.ID, wrote)
|
||||
}
|
||||
if state == "blocked" && isTaskLLMChainExhausted(err) {
|
||||
_ = t.Store.SetIntentBlockedReason(intent.ID, db.IntentBlockedLLMQuota)
|
||||
} else {
|
||||
_ = t.Store.SetIntentState(intent.ID, state)
|
||||
}
|
||||
log.Printf("[worker %s] task %s 意图 #%d 结束: %s (写回 %s)", name, t.ID, intent.ID, state, wrote)
|
||||
e.touch(t.ID)
|
||||
t.NotifyDone(intent.ID) // results changed the graph -> wake the planner (with the just-finished intent id)
|
||||
return true
|
||||
}
|
||||
|
||||
func sleepCtx(ctx context.Context, d time.Duration) (done bool) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
|
||||
@@ -7,90 +7,103 @@ import (
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
)
|
||||
|
||||
// The engine's cancel paths must hand a NAMED cause to the run being cancelled —
|
||||
// that cause is the whole content of the "运行被中断" line in the trace. These
|
||||
// exercise the context machinery only (no Manager / DB needed).
|
||||
|
||||
func TestPauseCancelsWithNamedCause(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
ectx := e.execContextFor(context.Background(), "t1")
|
||||
// the per-work child the worker loop actually runs under
|
||||
workCtx, workCancel := context.WithCancelCause(ectx)
|
||||
execCtx := e.execContextFor(context.Background(), "t1")
|
||||
workCtx, workCancel := context.WithCancelCause(execCtx)
|
||||
defer workCancel(nil)
|
||||
|
||||
e.Pause("t1", agent.AbortPausedByUser)
|
||||
|
||||
if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "paused_by_user" {
|
||||
t.Fatalf("work ctx 拿到的原因 code=%q ok=%v,期望 paused_by_user", code, ok)
|
||||
t.Fatalf("code=%q ok=%v, want paused_by_user", code, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKillWorkCancelsOnlyThatWorkWithNamedCause(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
ectx := e.execContextFor(context.Background(), "t1")
|
||||
workCtx, workCancel := context.WithCancelCause(ectx)
|
||||
execCtx := e.execContextFor(context.Background(), "t1")
|
||||
workCtx, workCancel := context.WithCancelCause(execCtx)
|
||||
e.registerWork(42, workCancel)
|
||||
|
||||
if err := e.KillWork(42); err != nil {
|
||||
t.Fatalf("KillWork: %v", err)
|
||||
t.Fatal(err)
|
||||
}
|
||||
if code, _, _, ok := agent.AbortReason(workCtx); !ok || code != "killed_by_planner" {
|
||||
t.Fatalf("code=%q ok=%v,期望 killed_by_planner", code, ok)
|
||||
t.Fatalf("code=%q ok=%v, want killed_by_planner", code, ok)
|
||||
}
|
||||
// the task ctx must survive — kill_work stops one work, not the whole task
|
||||
if ectx.Err() != nil {
|
||||
t.Error("kill_work 不该取消整个任务的 exec ctx")
|
||||
}
|
||||
if err := e.KillWork(999); err == nil {
|
||||
t.Error("对没有在跑的意图 KillWork 应该报错")
|
||||
if execCtx.Err() != nil {
|
||||
t.Fatal("kill_work cancelled the entire task context")
|
||||
}
|
||||
_, complete := e.detachWork(42)
|
||||
complete()
|
||||
}
|
||||
|
||||
func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) {
|
||||
func TestControlWorkCarriesUserActionCause(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
cause *agent.AbortCause
|
||||
want string
|
||||
action string
|
||||
code string
|
||||
}{
|
||||
{"settle", agent.AbortSettleDrainTimeout, "settle_drain_timeout"},
|
||||
{"goalMet", agent.AbortGoalMet, "goal_met"},
|
||||
{action: "pause", code: "work_paused_by_user"},
|
||||
{action: "cancel", code: "work_cancelled_by_user"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Run(tc.action, func(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
ectx := e.execContextFor(context.Background(), "t1")
|
||||
e.cancelExec("t1", tc.cause)
|
||||
if code, _, _, _ := agent.AbortReason(ectx); code != tc.want {
|
||||
t.Fatalf("code=%q,期望 %q", code, tc.want)
|
||||
ctx, cancel := context.WithCancelCause(context.Background())
|
||||
e.registerWork(42, cancel)
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- e.ControlWork(42, tc.action) }()
|
||||
<-ctx.Done()
|
||||
if code, _, _, _ := agent.AbortReason(ctx); code != tc.code {
|
||||
t.Fatalf("code=%q, want %q", code, tc.code)
|
||||
}
|
||||
action, complete := e.detachWork(42)
|
||||
if action != tc.action {
|
||||
t.Fatalf("action=%q, want %q", action, tc.action)
|
||||
}
|
||||
complete()
|
||||
if err := <-done; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// asking for an exec ctx while paused must return an already-cancelled ctx that
|
||||
// still SAYS why — this is the claim→Execute race guard, easily mistaken for a
|
||||
// mystery abort.
|
||||
func TestSettleDrainAndGoalMetCausesAreDistinct(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
cause *agent.AbortCause
|
||||
code string
|
||||
}{
|
||||
{agent.AbortSettleDrainTimeout, "settle_drain_timeout"},
|
||||
{agent.AbortGoalMet, "goal_met"},
|
||||
} {
|
||||
e := NewEngine(nil)
|
||||
ctx := e.execContextFor(context.Background(), "t1")
|
||||
e.cancelExec("t1", tc.cause)
|
||||
if code, _, _, _ := agent.AbortReason(ctx); code != tc.code {
|
||||
t.Fatalf("code=%q, want %q", code, tc.code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecContextWhilePausedCarriesRaceGuardCause(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
e.Pause("t1", agent.AbortPausedByUser)
|
||||
ctx := e.execContextFor(context.Background(), "t1")
|
||||
if ctx.Err() == nil {
|
||||
t.Fatal("暂停期间不该发出可用的 exec ctx")
|
||||
t.Fatal("paused task received a live execution context")
|
||||
}
|
||||
if code, _, _, _ := agent.AbortReason(ctx); code != "paused_race_guard" {
|
||||
t.Fatalf("code=%q,期望 paused_race_guard", code)
|
||||
t.Fatalf("code=%q, want paused_race_guard", code)
|
||||
}
|
||||
}
|
||||
|
||||
// Resume must hand out a live context again (cancelling is one-shot).
|
||||
func TestResumeYieldsFreshContext(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
old := e.execContextFor(context.Background(), "t1")
|
||||
e.Pause("t1", agent.AbortPausedByUser)
|
||||
if old.Err() == nil {
|
||||
t.Fatal("Pause 没取消旧 ctx")
|
||||
t.Fatal("Pause did not cancel old context")
|
||||
}
|
||||
e.paused.Delete("t1") // Resume() without the Task/Notify plumbing
|
||||
e.paused.Delete("t1")
|
||||
if fresh := e.execContextFor(context.Background(), "t1"); fresh.Err() != nil {
|
||||
t.Fatal("恢复后应拿到全新可用的 exec ctx")
|
||||
t.Fatal("resume did not create a fresh execution context")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestEngineActiveLLMCallsConcurrent(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
const taskID = "42"
|
||||
const calls = 64
|
||||
|
||||
var started sync.WaitGroup
|
||||
var release sync.WaitGroup
|
||||
var finished sync.WaitGroup
|
||||
started.Add(calls)
|
||||
finished.Add(calls)
|
||||
release.Add(1)
|
||||
for range calls {
|
||||
go func() {
|
||||
defer finished.Done()
|
||||
e.BeginLLMCall(taskID)
|
||||
started.Done()
|
||||
release.Wait()
|
||||
e.EndLLMCall(taskID)
|
||||
}()
|
||||
}
|
||||
started.Wait()
|
||||
if got := e.ActiveLLMCalls(taskID); got != calls {
|
||||
t.Fatalf("active calls = %d, want %d", got, calls)
|
||||
}
|
||||
release.Done()
|
||||
finished.Wait()
|
||||
if got := e.ActiveLLMCalls(taskID); got != 0 {
|
||||
t.Fatalf("active calls after completion = %d, want 0", got)
|
||||
}
|
||||
|
||||
// A defensive extra end must never expose a negative UI count.
|
||||
e.EndLLMCall(taskID)
|
||||
if got := e.ActiveLLMCalls(taskID); got != 0 {
|
||||
t.Fatalf("active calls after extra end = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
@@ -43,7 +43,20 @@ func (e *Engine) inflightCounter(taskID string) *int64 {
|
||||
v, _ := e.inflight.LoadOrStore(taskID, new(int64))
|
||||
return v.(*int64)
|
||||
}
|
||||
func (e *Engine) incInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), 1) }
|
||||
|
||||
// beginTaskOperation atomically registers a task-owned operation unless deletion
|
||||
// has already installed its barrier. The delete handler can therefore wait for
|
||||
// inflight==0 without a check-then-start race recreating files after cleanup.
|
||||
func (e *Engine) beginTaskOperation(taskID string) bool {
|
||||
e.deleteMu.RLock()
|
||||
defer e.deleteMu.RUnlock()
|
||||
if e.IsDeleting(taskID) {
|
||||
return false
|
||||
}
|
||||
atomic.AddInt64(e.inflightCounter(taskID), 1)
|
||||
return true
|
||||
}
|
||||
|
||||
func (e *Engine) decInflight(taskID string) { atomic.AddInt64(e.inflightCounter(taskID), -1) }
|
||||
func (e *Engine) inflightCount(taskID string) int64 {
|
||||
return atomic.LoadInt64(e.inflightCounter(taskID))
|
||||
@@ -106,10 +119,18 @@ func (e *Engine) startDeadlineCoordinator(ctx context.Context, t *Task) {
|
||||
if t == nil || t.TimeoutSeconds <= 0 {
|
||||
return
|
||||
}
|
||||
if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded {
|
||||
e.deleteMu.RLock()
|
||||
if e.IsDeleting(t.ID) {
|
||||
e.deleteMu.RUnlock()
|
||||
return
|
||||
}
|
||||
go e.deadlineCoordinator(ctx, t)
|
||||
if _, loaded := e.coordStarted.LoadOrStore(t.ID, true); loaded {
|
||||
e.deleteMu.RUnlock()
|
||||
return
|
||||
}
|
||||
rt := e.registerTaskRoutines(ctx, t.ID, 1)
|
||||
e.deleteMu.RUnlock()
|
||||
runTaskRoutine(rt, func(loopCtx context.Context) { e.deadlineCoordinator(loopCtx, t) })
|
||||
}
|
||||
|
||||
// deadlineCoordinator waits until the task's absolute deadline, then runs the settle
|
||||
@@ -170,6 +191,10 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) {
|
||||
|
||||
// ⑤ 终局一轮 planner(任务超时词,最后目标判定,不产新意图)。
|
||||
met := e.runFinalPlannerRound(ctx, t)
|
||||
if !e.beginTaskOperation(t.ID) {
|
||||
return
|
||||
}
|
||||
defer e.decInflight(t.ID)
|
||||
|
||||
// ⑥ 定终态(带守卫):met → done(completed);否则 timeout。若常规路径已先落 done,
|
||||
// 守卫(SetTaskStatusGuarded)会拒绝覆盖,保留 completed 语义。
|
||||
@@ -192,6 +217,9 @@ func (e *Engine) settleTask(ctx context.Context, t *Task) {
|
||||
// task-timeout planner words (final goal judgment; no new intents). Waits for the
|
||||
// LLM to be ready (bounded by ctx) so a completable task isn't mis-judged timeout.
|
||||
func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) {
|
||||
if e.IsDeleting(t.ID) {
|
||||
return false
|
||||
}
|
||||
planner, _ := e.snapshotFor(t)
|
||||
for planner == nil {
|
||||
if sleepCtx(ctx, deadlinePollInterval) {
|
||||
@@ -200,17 +228,24 @@ func (e *Engine) runFinalPlannerRound(ctx context.Context, t *Task) (met bool) {
|
||||
if isTerminalStatus(e.m.TaskStatus(t.ID)) {
|
||||
return false
|
||||
}
|
||||
if e.IsDeleting(t.ID) {
|
||||
return false
|
||||
}
|
||||
planner, _ = e.snapshotFor(t)
|
||||
}
|
||||
// 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。
|
||||
fctx := e.clockCtx(ctx, t, true)
|
||||
if !e.beginTaskOperation(t.ID) {
|
||||
return false
|
||||
}
|
||||
defer e.decInflight(t.ID)
|
||||
emit := func(r db.Activity) { e.emitActivity(t, r) }
|
||||
e.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: fmt.Sprintf("任务超时收尾·终局判定(第 %d 轮)", e.nextPlannerRound(t.ID))})
|
||||
// 独立 ctx(不挂 execCancel,避免 pause/硬 cancel 打断这最后一轮),带 Final 注入任务超时词。
|
||||
fctx := e.clockCtx(ctx, t, true)
|
||||
e.incInflight(t.ID)
|
||||
tTaskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
e.BeginLLMCall(t.ID)
|
||||
met, reason, err := planner.Plan(fctx, tTaskID, e.m.assets, t.Store, t.Goal, t.drainTriggers(), emit)
|
||||
e.decInflight(t.ID)
|
||||
e.EndLLMCall(t.ID)
|
||||
if err != nil {
|
||||
log.Printf("[deadline] task %s 终局规划出错: %v", t.ID, err)
|
||||
} else if met {
|
||||
|
||||
+63
-87
@@ -10,7 +10,6 @@ import (
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
type goalSpec struct {
|
||||
@@ -27,25 +26,36 @@ type goalSpec struct {
|
||||
//
|
||||
// 异步(goroutine)所以调用方立即返回,两条路径行为一致:秒建任务、后台拆目标。
|
||||
func (s *Server) launchTask(t *Task, seedText string, seedFirstIntent bool) {
|
||||
if !s.engine.beginTaskOperation(t.ID) {
|
||||
return
|
||||
}
|
||||
s.seed(t, seedText)
|
||||
if seedFirstIntent {
|
||||
s.seedFirstIntent(t)
|
||||
}
|
||||
// 并发上限:满了就挂起排队(不做目标拆解、不启动引擎),由 reconcileConcurrency 补位启动。
|
||||
if s.queueIfAtCapacity(t) {
|
||||
s.engine.decInflight(t.ID)
|
||||
return
|
||||
}
|
||||
go s.startTaskEngine(t)
|
||||
go func() {
|
||||
defer s.engine.decInflight(t.ID)
|
||||
s.startTaskEngine(t)
|
||||
}()
|
||||
}
|
||||
|
||||
// startTaskEngine does the actual running work: 第0轮目标拆解 + 启动引擎循环。
|
||||
// Shared by fresh launches and by the concurrency reconciler promoting a queued task.
|
||||
func (s *Server) startTaskEngine(t *Task) {
|
||||
ctx := s.engine.execContextFor(s.ctx, t.ID)
|
||||
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "round",
|
||||
Summary: "第 0 轮目标拆解"})
|
||||
goals := s.createGoals(s.ctx, t, func(r db.Activity) {
|
||||
goals := s.createGoals(ctx, t, func(r db.Activity) {
|
||||
s.engine.emitActivity(t, r)
|
||||
})
|
||||
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
for _, g := range goals {
|
||||
summary := g.Text
|
||||
if g.VulnClass != "" {
|
||||
@@ -53,37 +63,26 @@ func (s *Server) startTaskEngine(t *Task) {
|
||||
}
|
||||
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text", Summary: summary})
|
||||
}
|
||||
if ctx.Err() != nil || s.engine.IsDeleting(t.ID) {
|
||||
return
|
||||
}
|
||||
s.engine.Run(s.ctx, t)
|
||||
}
|
||||
|
||||
// occupiesSlot reports whether a task counts against the concurrency cap: it is
|
||||
// live (not terminal), not queued, and not paused. A just-created / decomposing
|
||||
// task counts too (it will run momentarily).
|
||||
func (s *Server) occupiesSlot(t *Task) bool {
|
||||
if t == nil || t.Queued || isTerminalStatus(t.Status) {
|
||||
return false
|
||||
}
|
||||
return !t.Paused && !s.engine.IsPaused(t.ID)
|
||||
func (s *Server) occupiesConcurrencySlot(t *Task) bool {
|
||||
return t != nil && !t.Queued && !t.Paused && !s.engine.IsPaused(t.ID) && !isTerminalStatus(t.Status)
|
||||
}
|
||||
|
||||
// runningTaskCount counts tasks occupying a concurrency slot, excluding excludeID
|
||||
// (pass the candidate's id when deciding admission so it doesn't count itself).
|
||||
func (s *Server) runningTaskCount(excludeID string) int {
|
||||
n := 0
|
||||
for _, t := range s.m.List() {
|
||||
if t.ID == excludeID {
|
||||
continue
|
||||
}
|
||||
if s.occupiesSlot(t) {
|
||||
n++
|
||||
count := 0
|
||||
for _, task := range s.m.List() {
|
||||
if task.ID != excludeID && s.occupiesConcurrencySlot(task) {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return n
|
||||
return count
|
||||
}
|
||||
|
||||
// queueIfAtCapacity holds a task in the queued state when the concurrency cap is
|
||||
// enabled and already full. Returns true if the task was queued (caller must not
|
||||
// start it). No-op (false) when the feature is off or a slot is free.
|
||||
func (s *Server) queueIfAtCapacity(t *Task) bool {
|
||||
s.concMu.Lock()
|
||||
defer s.concMu.Unlock()
|
||||
@@ -93,52 +92,49 @@ func (s *Server) queueIfAtCapacity(t *Task) bool {
|
||||
}
|
||||
if err := s.m.SetTaskQueued(t.ID, true); err != nil {
|
||||
log.Printf("[concurrency] task %s 置排队失败: %v", t.ID, err)
|
||||
return false // 落库失败宁可放它跑,也别静默丢任务
|
||||
return false
|
||||
}
|
||||
t.Queued = true
|
||||
s.engine.emitActivity(t, db.Activity{Worker: "planner", Kind: "text",
|
||||
s.engine.emitActivity(t, db.Activity{Worker: "system", Kind: "text",
|
||||
Summary: fmt.Sprintf("已排队:达到并发上限 %d,等待空位后自动开始", limit)})
|
||||
log.Printf("[concurrency] task %s 排队(并发上限 %d)", t.ID, limit)
|
||||
return true
|
||||
}
|
||||
|
||||
// reconcileConcurrency promotes queued tasks into free slots. Runs on every
|
||||
// scheduler tick. When the feature is disabled it releases ALL queued tasks;
|
||||
// when enabled it starts the oldest-queued tasks up to (limit - running).
|
||||
func (s *Server) reconcileConcurrency() {
|
||||
s.concMu.Lock()
|
||||
defer s.concMu.Unlock()
|
||||
var queued []*Task
|
||||
for _, t := range s.m.List() {
|
||||
if t.Queued && !isTerminalStatus(t.Status) {
|
||||
queued = append(queued, t)
|
||||
queued := []*Task{}
|
||||
for _, task := range s.m.List() {
|
||||
if task.Queued && !isTerminalStatus(task.Status) {
|
||||
queued = append(queued, task)
|
||||
}
|
||||
}
|
||||
if len(queued) == 0 {
|
||||
return
|
||||
}
|
||||
// 最早排队的先启动。
|
||||
sort.Slice(queued, func(i, j int) bool { return queued[i].CreatedAt < queued[j].CreatedAt })
|
||||
|
||||
enabled, limit := s.m.ConcurrencyLimit()
|
||||
slots := len(queued) // feature off → 全部放行
|
||||
slots := len(queued)
|
||||
if enabled {
|
||||
if !s.engine.Ready() {
|
||||
return // 引擎未就绪(无 LLM):继续挂起,别空跑
|
||||
}
|
||||
slots = limit - s.runningTaskCount("")
|
||||
}
|
||||
for _, t := range queued {
|
||||
for _, task := range queued {
|
||||
if slots <= 0 {
|
||||
break
|
||||
}
|
||||
if err := s.m.SetTaskQueued(t.ID, false); err != nil {
|
||||
log.Printf("[concurrency] task %s 解除排队失败: %v", t.ID, err)
|
||||
// Keep FIFO order, but do not consume a concurrency slot for a task
|
||||
// whose explicit chain/global provider is unavailable. It will be retried
|
||||
// after the user configures or resets its LLM chain.
|
||||
if task.Paused || !s.engine.ReadyFor(task) {
|
||||
continue
|
||||
}
|
||||
t.Queued = false
|
||||
log.Printf("[concurrency] task %s 出队启动", t.ID)
|
||||
go s.startTaskEngine(t)
|
||||
if !s.engine.beginTaskOperation(task.ID) {
|
||||
continue
|
||||
}
|
||||
if err := s.m.SetTaskQueued(task.ID, false); err != nil {
|
||||
s.engine.decInflight(task.ID)
|
||||
continue
|
||||
}
|
||||
go func(t *Task) {
|
||||
defer s.engine.decInflight(t.ID)
|
||||
s.startTaskEngine(t)
|
||||
}(task)
|
||||
slots--
|
||||
}
|
||||
}
|
||||
@@ -153,7 +149,6 @@ func (s *Server) reviveTask(t *Task) {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
// 排队中的任务不复活启动——它本就等待并发空位,交给 reconcileConcurrency。
|
||||
if t.Queued {
|
||||
return
|
||||
}
|
||||
@@ -184,24 +179,25 @@ func (s *Server) reviveTask(t *Task) {
|
||||
// Returns the seeded specs so callers can emit activity records for them.
|
||||
// emit, when non-nil, is forwarded to DecomposeGoals so LLM steps are visible in the UI.
|
||||
func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity)) []goalSpec {
|
||||
if t == nil {
|
||||
return nil
|
||||
}
|
||||
// Use the SAME LLM the task runs on (its pinned profile, else the active profile),
|
||||
// NOT agent.FromEnv() — the LLM is configured via the UI (DB profile), not env vars,
|
||||
// so FromEnv returned empty and every task silently fell back to the crude rule
|
||||
// splitter (which shredded URLs / made meaningless 2-way splits).
|
||||
var specs []goalSpec
|
||||
if prov, ok := s.goalsProvider(t); ok {
|
||||
var as *db.AssetStore
|
||||
if s.m != nil {
|
||||
as = s.m.Assets()
|
||||
}
|
||||
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
// DecomposeGoals persists each decomposed goal via the set_goals tool straight
|
||||
// into t.Store; it returns the specs it wrote so we can emit activity + spot the
|
||||
// empty (LLM error / no provider) case below. No write-back needed here anymore.
|
||||
for _, g := range agent.DecomposeGoals(ctx, prov, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit) {
|
||||
if strings.TrimSpace(g.Text) != "" {
|
||||
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
||||
}
|
||||
var as *db.AssetStore
|
||||
if s.m != nil {
|
||||
as = s.m.Assets()
|
||||
}
|
||||
taskID, _ := strconv.ParseInt(t.ID, 10, 64)
|
||||
s.engine.BeginLLMCall(t.ID)
|
||||
decomposed := agent.DecomposeGoalsWithProvider(ctx, s.agentsForTask(t).runtime, s.m.dir, t.Goal, t.Description, as, t.Store, taskID, emit)
|
||||
s.engine.EndLLMCall(t.ID)
|
||||
for _, g := range decomposed {
|
||||
if strings.TrimSpace(g.Text) != "" {
|
||||
specs = append(specs, goalSpec{Text: g.Text, VulnClass: g.VulnClass})
|
||||
}
|
||||
}
|
||||
if len(specs) == 0 {
|
||||
@@ -220,23 +216,3 @@ func (s *Server) createGoals(ctx context.Context, t *Task, emit func(db.Activity
|
||||
}
|
||||
return specs
|
||||
}
|
||||
|
||||
// goalsProvider resolves the provider for goal decomposition by the standard
|
||||
// precedence: the goals agent's own binding → the task's pinned profile → the
|
||||
// global active provider (the very instance the engine runs on, so decomposition
|
||||
// shares its rate limiter, gets recorded, and follows LLM failover).
|
||||
func (s *Server) goalsProvider(t *Task) (llm.Provider, bool) {
|
||||
var pin *int64
|
||||
if t != nil {
|
||||
pin = t.LLMProfileID
|
||||
}
|
||||
if eff := s.effectiveProfileForAgent("goals", pin); eff != nil {
|
||||
if prov, cfg, ok := s.providerForProfile(*eff); ok {
|
||||
return s.poolForBinding(*eff, prov, cfg), true
|
||||
}
|
||||
}
|
||||
s.cfgMu.Lock()
|
||||
prov, on := s.llmProv, s.llmOn
|
||||
s.cfgMu.Unlock()
|
||||
return prov, on && prov != nil
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCreateGoalsNilTask(t *testing.T) {
|
||||
if got := (&Server{}).createGoals(context.Background(), nil, nil); len(got) != 0 {
|
||||
t.Fatalf("nil task produced goals: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
func TestInheritedActivityDetailAndRelationDeletion(t *testing.T) {
|
||||
m, err := NewManager(t.TempDir(), "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
defer m.Close()
|
||||
|
||||
source, err := m.CreateTaskWithOptions("detail source", "source goal", db.TaskCreateOptions{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sourceID := mustTaskID(t, source.ID)
|
||||
current, err := m.CreateTaskWithOptions("detail current", "current goal", db.TaskCreateOptions{SourceTaskIDs: []int64{sourceID}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = m.DeleteTask(current.ID, DeleteTaskOptions{})
|
||||
_, _ = m.DeleteTask(source.ID, DeleteTaskOptions{})
|
||||
})
|
||||
|
||||
plannerStep, err := source.Store.AppendActivity(db.Activity{
|
||||
Worker: "planner", Kind: "text", Summary: "private planner", Detail: "private planner detail",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
intentID, err := source.Store.AddIntent(map[string]any{"summary": "completed source work"}, 1, nil, "planner")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := source.Store.SetIntentState(intentID, "done"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
workerStep, err := source.Store.AppendActivity(db.Activity{
|
||||
NodeID: &intentID, Worker: "worker", Kind: "result", Summary: "shared result", Detail: "shared worker detail",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
s := New(context.Background(), m, t.TempDir(), t.TempDir(), t.TempDir())
|
||||
token, err := signJWT(s.jwtKey)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fetchDetail := func(stepID int64) string {
|
||||
req := httptest.NewRequest(http.MethodGet,
|
||||
fmt.Sprintf("/api/exploration/activity/%d?task=%s", stepID, current.ID), nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
rec := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(rec, req)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("activity detail %d: status=%d body=%s", stepID, rec.Code, rec.Body.String())
|
||||
}
|
||||
var body struct {
|
||||
Detail string `json:"detail"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return body.Detail
|
||||
}
|
||||
if got := fetchDetail(plannerStep); got != "" {
|
||||
t.Fatalf("source planner transcript leaked through API: %q", got)
|
||||
}
|
||||
if got := fetchDetail(workerStep); got != "shared worker detail" {
|
||||
t.Fatalf("terminal inherited worker detail missing: %q", got)
|
||||
}
|
||||
|
||||
if _, err := m.DeleteTask(source.ID, DeleteTaskOptions{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
currentAfter, ok := m.Task(current.ID)
|
||||
if !ok || len(currentAfter.SourceTaskIDs) != 0 {
|
||||
t.Fatalf("deleted source remained in live task DTO state: %+v", currentAfter)
|
||||
}
|
||||
}
|
||||
|
||||
func mustTaskID(t *testing.T, raw string) int64 {
|
||||
t.Helper()
|
||||
var id int64
|
||||
if _, err := fmt.Sscan(raw, &id); err != nil || id <= 0 {
|
||||
t.Fatalf("bad task id %q: %v", raw, err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
)
|
||||
|
||||
func TestInheritedDTOProvenance(t *testing.T) {
|
||||
createdAt := time.Date(2026, time.August, 19, 12, 0, 0, 0, time.UTC)
|
||||
node := &db.Node{
|
||||
ID: 41,
|
||||
Kind: db.KindFinding,
|
||||
Payload: json.RawMessage(`{"vulnclass":"xss","severity":"high","summary":"source finding"}`),
|
||||
State: "confirmed",
|
||||
CreatedAt: createdAt,
|
||||
SourceTaskID: 7,
|
||||
Inherited: true,
|
||||
}
|
||||
|
||||
nodeDTO := taskNodeDTO(node)
|
||||
if !nodeDTO.Inherited || nodeDTO.SourceTaskID != "7" {
|
||||
t.Fatalf("node provenance lost: %+v", nodeDTO)
|
||||
}
|
||||
finding := findingDTOsForOwner("7", "source", []*db.Node{node}, nil, nil)
|
||||
if len(finding) != 1 || !finding[0].Inherited || finding[0].SourceTaskID != "7" || finding[0].TaskID != "7" {
|
||||
t.Fatalf("finding provenance lost: %+v", finding)
|
||||
}
|
||||
|
||||
intentID := int64(44)
|
||||
activity := activityDTO(db.Activity{
|
||||
ID: 45,
|
||||
NodeID: &intentID,
|
||||
CreatedAt: createdAt,
|
||||
SourceTaskID: 7,
|
||||
Inherited: true,
|
||||
})
|
||||
if !activity.Inherited || activity.SourceTaskID != "7" || activity.IntentID != "44" {
|
||||
t.Fatalf("activity provenance lost: %+v", activity)
|
||||
}
|
||||
assetRef := coverageAssetRefDTO(db.AssetRef{
|
||||
ID: 46, Kind: db.KindFact, State: "confirmed", Summary: "source fact", SourceTaskID: 7, Inherited: true,
|
||||
})
|
||||
if !assetRef.Inherited || assetRef.SourceTaskID != "7" {
|
||||
t.Fatalf("asset ref provenance lost: %+v", assetRef)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInheritedIntentResultOnlyAllowsHistory(t *testing.T) {
|
||||
for _, state := range []string{"done", "blocked", "exhausted", "stopped"} {
|
||||
if !inheritedIntentResult(state) {
|
||||
t.Errorf("terminal state %q should be inherited history", state)
|
||||
}
|
||||
}
|
||||
for _, state := range []string{"open", "running", ""} {
|
||||
if inheritedIntentResult(state) {
|
||||
t.Errorf("live state %q must not enter inherited history", state)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInheritedGraphSnapshotFiltersLiveIntentsAndEdges(t *testing.T) {
|
||||
nodes := []*db.Node{
|
||||
{ID: 1, Kind: db.KindFact, State: "confirmed"},
|
||||
{ID: 2, Kind: db.KindIntent, State: "running"},
|
||||
{ID: 3, Kind: db.KindIntent, State: "done"},
|
||||
{ID: 4, Kind: db.KindFinding, State: "confirmed"},
|
||||
}
|
||||
edges := []db.Edge{
|
||||
{From: 1, Rel: db.RelDerivedFrom, To: 2},
|
||||
{From: 1, Rel: db.RelDerivedFrom, To: 3},
|
||||
{From: 3, Rel: db.RelYields, To: 4},
|
||||
}
|
||||
|
||||
gotNodes, gotEdges := inheritedGraphSnapshot(nodes, edges, 77)
|
||||
if len(gotNodes) != 3 {
|
||||
t.Fatalf("expected fact, terminal intent and finding; got %+v", gotNodes)
|
||||
}
|
||||
for _, node := range gotNodes {
|
||||
if node.ID == 2 {
|
||||
t.Fatal("running inherited intent leaked into graph snapshot")
|
||||
}
|
||||
if !node.Inherited || node.SourceTaskID != 77 {
|
||||
t.Fatalf("source provenance missing: %+v", node)
|
||||
}
|
||||
}
|
||||
if len(gotEdges) != 2 {
|
||||
t.Fatalf("edge referencing hidden live intent was not removed: %+v", gotEdges)
|
||||
}
|
||||
for _, edge := range gotEdges {
|
||||
if edge.From == 2 || edge.To == 2 {
|
||||
t.Fatalf("live intent id leaked through edge: %+v", edge)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFindingProvenanceInTask(t *testing.T) {
|
||||
task := &Task{ID: "10", SourceTaskIDs: []int64{7, 8}}
|
||||
|
||||
local := int64(10)
|
||||
if source, inherited, allowed := findingProvenanceInTask(task, &local); !allowed || inherited || source != "" {
|
||||
t.Fatalf("local finding: source=%q inherited=%v allowed=%v", source, inherited, allowed)
|
||||
}
|
||||
sourceID := int64(7)
|
||||
if source, inherited, allowed := findingProvenanceInTask(task, &sourceID); !allowed || !inherited || source != "7" {
|
||||
t.Fatalf("direct source finding: source=%q inherited=%v allowed=%v", source, inherited, allowed)
|
||||
}
|
||||
unrelated := int64(6)
|
||||
if _, _, allowed := findingProvenanceInTask(task, &unrelated); allowed {
|
||||
t.Fatal("unrelated finding must not be readable through task context")
|
||||
}
|
||||
if _, _, allowed := findingProvenanceInTask(task, nil); allowed {
|
||||
t.Fatal("detached finding must not be readable through task context")
|
||||
}
|
||||
}
|
||||
+365
-29
@@ -2,6 +2,7 @@ package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
@@ -12,38 +13,44 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/enrich"
|
||||
"github.com/Autumn-27/artex/guard"
|
||||
"github.com/Autumn-27/artex/intercept"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
actool "github.com/Autumn-27/norma/tool"
|
||||
)
|
||||
|
||||
// Task is one engagement: a description + goal + its own exploration store,
|
||||
// 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 {
|
||||
ID string `json:"id"`
|
||||
ExpID int64 `json:"exploration_id"`
|
||||
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=尚未开跑
|
||||
ParentRef string `json:"parent_ref,omitempty"` // 父任务 id(编排 spawn 记录)
|
||||
LLMProfileID *int64 `json:"llm_profile_id,omitempty"` // 指定运行本任务 planner/worker 的 LLM 配置;nil=用全局激活配置
|
||||
Status string `json:"status"` // persisted lifecycle status (done/failed/timeout 为终态;空/其它则由运行态推导)
|
||||
ID string `json:"id"`
|
||||
ExpID int64 `json:"exploration_id"`
|
||||
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=尚未开跑
|
||||
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"`
|
||||
LLMFailoverState string `json:"llm_failover_state,omitempty"`
|
||||
LLMFailoverReason string `json:"llm_failover_reason,omitempty"`
|
||||
SourceTaskIDs []int64 `json:"source_task_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 心跳触发间隔(秒)
|
||||
FirstRunAt int64 `json:"first_run_at,omitempty"`
|
||||
DeadlineAt int64 `json:"deadline_at,omitempty"`
|
||||
Store *pgdb.ExplorationStore `json:"-"`
|
||||
Guard *guard.Guard `json:"-"`
|
||||
notify chan struct{}
|
||||
TimeoutSeconds int `json:"timeout_seconds"`
|
||||
PlanHeartbeatSeconds int `json:"plan_heartbeat_seconds"` // planner 心跳触发间隔(秒)
|
||||
FirstRunAt int64 `json:"first_run_at,omitempty"`
|
||||
DeadlineAt int64 `json:"deadline_at,omitempty"`
|
||||
Store *pgdb.ExplorationStore `json:"-"`
|
||||
Guard *guard.Guard `json:"-"`
|
||||
notify chan struct{}
|
||||
llmMu sync.RWMutex
|
||||
|
||||
// pendingTriggers accumulates the concrete changes (worker done / finding) that
|
||||
// fired planning rounds since the last one consumed them. The debounce coalesces
|
||||
@@ -52,14 +59,74 @@ type Task struct {
|
||||
pendingTriggers []agent.TriggerEvent
|
||||
}
|
||||
|
||||
type taskLLMState struct {
|
||||
ProfileID *int64
|
||||
ProfileIDs []int64
|
||||
ActiveID *int64
|
||||
FailoverState string
|
||||
FailoverReason string
|
||||
}
|
||||
|
||||
func cloneInt64Ptr(value *int64) *int64 {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
copy := *value
|
||||
return ©
|
||||
}
|
||||
|
||||
func (t *Task) llmStateSnapshot() taskLLMState {
|
||||
t.llmMu.RLock()
|
||||
defer t.llmMu.RUnlock()
|
||||
return taskLLMState{
|
||||
ProfileID: cloneInt64Ptr(t.LLMProfileID),
|
||||
ProfileIDs: append(make([]int64, 0, len(t.LLMProfileIDs)), t.LLMProfileIDs...),
|
||||
ActiveID: cloneInt64Ptr(t.ActiveLLMProfileID),
|
||||
FailoverState: t.LLMFailoverState,
|
||||
FailoverReason: t.LLMFailoverReason,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) setLLMState(profileID, activeID *int64, profileIDs []int64, state, reason string) {
|
||||
t.llmMu.Lock()
|
||||
t.LLMProfileID = cloneInt64Ptr(profileID)
|
||||
t.ActiveLLMProfileID = cloneInt64Ptr(activeID)
|
||||
t.LLMProfileIDs = append(t.LLMProfileIDs[:0], profileIDs...)
|
||||
t.LLMFailoverState = state
|
||||
t.LLMFailoverReason = reason
|
||||
t.llmMu.Unlock()
|
||||
}
|
||||
|
||||
// DeleteTaskOptions controls cleanup of data stored outside the task's own
|
||||
// exploration graph. All options default to false for backward compatibility.
|
||||
type DeleteTaskOptions struct {
|
||||
DeleteAssets bool `json:"delete_assets"`
|
||||
DeleteTraffic bool `json:"delete_traffic"`
|
||||
DeleteFiles bool `json:"delete_files"`
|
||||
DeleteFindings bool `json:"delete_findings"`
|
||||
DeleteLLMRecords bool `json:"delete_llm_records"`
|
||||
}
|
||||
|
||||
// DeleteTaskResult makes destructive cleanup auditable to API callers.
|
||||
type DeleteTaskResult struct {
|
||||
Deleted string `json:"deleted"`
|
||||
AssetsDeleted int64 `json:"assets_deleted"`
|
||||
AssetsDetached int64 `json:"assets_detached"`
|
||||
TrafficDeleted int64 `json:"traffic_deleted"`
|
||||
FilesDeleted bool `json:"files_deleted"`
|
||||
FindingsDeleted int64 `json:"findings_deleted"`
|
||||
LLMRecordsDeleted int64 `json:"llm_records_deleted"`
|
||||
CleanupWarning string `json:"cleanup_warning,omitempty"`
|
||||
}
|
||||
|
||||
// Manager owns the PostgreSQL data source (asset graph + every task's exploration
|
||||
// graph + config) and the in-memory set of task handles.
|
||||
type Manager struct {
|
||||
dir string
|
||||
pg *pgdb.DB
|
||||
assets *pgdb.AssetStore
|
||||
traffic *traffic.Traffic // process-wide recording proxy (may be nil)
|
||||
enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP)
|
||||
traffic *traffic.Traffic // process-wide recording proxy (may be nil)
|
||||
enrich *enrich.Engine // engine-side asset auto-completion (DNS/HTTP)
|
||||
interceptor *intercept.Interceptor // user-configured tool-call interception rules
|
||||
|
||||
mu sync.RWMutex
|
||||
@@ -529,7 +596,10 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter
|
||||
ID: strconv.FormatInt(pt.ID, 10), ExpID: pt.ExplorationID,
|
||||
Description: pt.Description, Goal: pt.Goal, CreatedAt: pt.CreatedAt.Unix(), Paused: pt.Paused, Queued: pt.Queued,
|
||||
CompletedAt: unixOrZero(pt.CompletedAt), Status: pt.Status, ParentRef: pt.ParentRef,
|
||||
LLMProfileID: pt.LLMProfileID,
|
||||
LLMProfileID: pt.LLMProfileID,
|
||||
LLMProfileIDs: append([]int64(nil), pt.LLMProfileIDs...), ActiveLLMProfileID: pt.ActiveLLMProfileID,
|
||||
LLMFailoverState: pt.LLMFailoverState, LLMFailoverReason: pt.LLMFailoverReason,
|
||||
SourceTaskIDs: append([]int64(nil), pt.SourceTaskIDs...),
|
||||
TimeoutSeconds: pt.TimeoutSeconds, PlanHeartbeatSeconds: pt.PlanHeartbeatSeconds,
|
||||
FirstRunAt: unixOrZero(pt.FirstRunAt), DeadlineAt: unixOrZero(pt.DeadlineAt),
|
||||
Store: store, Guard: guard.NewWithInterceptor(ic), notify: make(chan struct{}, 1),
|
||||
@@ -539,7 +609,17 @@ func taskFromPG(pt *pgdb.Task, store *pgdb.ExplorationStore, ic *intercept.Inter
|
||||
// CreateTask creates a task + its exploration and makes it active.
|
||||
// timeoutSeconds is the task-level wall-clock budget (0 = 不限时).
|
||||
func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, timeoutSeconds, planHeartbeatSeconds int) (*Task, error) {
|
||||
pt, err := m.pg.CreateTask(description, goal, llmProfileID, timeoutSeconds, planHeartbeatSeconds)
|
||||
var ids []int64
|
||||
if llmProfileID != nil {
|
||||
ids = []int64{*llmProfileID}
|
||||
}
|
||||
return m.CreateTaskWithOptions(description, goal, pgdb.TaskCreateOptions{
|
||||
LLMProfileIDs: ids, TimeoutSeconds: timeoutSeconds, PlanHeartbeatSeconds: planHeartbeatSeconds,
|
||||
})
|
||||
}
|
||||
|
||||
func (m *Manager) CreateTaskWithOptions(description, goal string, opts pgdb.TaskCreateOptions) (*Task, error) {
|
||||
pt, err := m.pg.CreateTaskWithOptions(description, goal, opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -551,6 +631,35 @@ func (m *Manager) CreateTask(description, goal string, llmProfileID *int64, time
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// ReplaceTaskLLMProfiles resets a non-terminal task's ordered provider chain and
|
||||
// mirrors the committed state onto the live task handle.
|
||||
func (m *Manager) ReplaceTaskLLMProfiles(id string, profileIDs []int64, activeProfileID int64) (int64, error) {
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if err := m.pg.ReplaceTaskLLMProfiles(n, profileIDs, activeProfileID); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
pt, err := m.pg.GetTask(n)
|
||||
if err != nil || pt == nil {
|
||||
return 0, err
|
||||
}
|
||||
m.mu.Lock()
|
||||
if task := m.tasks[id]; task != nil {
|
||||
task.setLLMState(pt.LLMProfileID, pt.ActiveLLMProfileID, pt.LLMProfileIDs, pt.LLMFailoverState, pt.LLMFailoverReason)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
if task, ok := m.Task(id); ok {
|
||||
reopened, reopenErr := task.Store.ReopenIntentsByBlockedReason(pgdb.IntentBlockedLLMQuota)
|
||||
if reopenErr != nil {
|
||||
return reopened, reopenErr
|
||||
}
|
||||
return reopened, nil
|
||||
}
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// LoadExisting rebuilds in-memory task handles from the PG task registry.
|
||||
func (m *Manager) LoadExisting() []*Task {
|
||||
pts, err := m.pg.ListTasks()
|
||||
@@ -706,17 +815,123 @@ func (m *Manager) SetTaskStatus(id, status string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteTask removes a task (and its exploration subgraph) from PG + memory.
|
||||
func (m *Manager) DeleteTask(id string) error {
|
||||
// DeleteTask removes a task and optionally its related global data. Traffic has
|
||||
// no task-id column, so related exchanges are resolved by exact hosts from the
|
||||
// task's asset rows. Files are staged before the database operation; traffic is
|
||||
// staged while PostgreSQL excludes asset/anchor writers. Both are restored on a
|
||||
// database failure and purged only after its commit.
|
||||
func (m *Manager) DeleteTask(id string, opts DeleteTaskOptions) (DeleteTaskResult, error) {
|
||||
result := DeleteTaskResult{Deleted: id}
|
||||
n, err := strconv.ParseInt(id, 10, 64)
|
||||
if err != nil {
|
||||
return err
|
||||
return result, err
|
||||
}
|
||||
if err := m.pg.DeleteTask(n); err != nil {
|
||||
return err
|
||||
registered, err := m.pg.GetTask(n)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if registered == nil {
|
||||
m.forgetTask(id, n)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
var fileStage *taskFileDeleteStage
|
||||
if opts.DeleteFiles {
|
||||
fileStage, err = stageTaskFiles(m.dir, id, registered.ExplorationID)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
result.FilesDeleted = fileStage.deleted
|
||||
}
|
||||
|
||||
var trafficStage *traffic.HostDeleteStage
|
||||
var prepare func(pgdb.TaskDeletePreparation) error
|
||||
if opts.DeleteTraffic && m.traffic != nil {
|
||||
prepare = func(p pgdb.TaskDeletePreparation) error {
|
||||
if len(p.TrafficHosts) == 0 {
|
||||
return nil
|
||||
}
|
||||
trafficStage, err = m.traffic.StageDeleteHostsExact(p.TrafficHosts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
result.TrafficDeleted = trafficStage.Deleted()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
dbResult, err := m.pg.DeleteTaskCascadePrepared(
|
||||
n, opts.DeleteAssets, opts.DeleteFindings, opts.DeleteLLMRecords, prepare,
|
||||
)
|
||||
if err != nil {
|
||||
return result, rollbackTaskDelete(err, trafficStage, fileStage)
|
||||
}
|
||||
result.AssetsDeleted = dbResult.AssetsDeleted
|
||||
result.AssetsDetached = dbResult.AssetsDetached
|
||||
result.FindingsDeleted = dbResult.FindingsDeleted
|
||||
result.LLMRecordsDeleted = dbResult.LLMRecordsDeleted
|
||||
|
||||
// PostgreSQL is now authoritative: finalize the staged external deletion and
|
||||
// forget the live task even if a final purge reports an error. Such errors are
|
||||
// typed so the HTTP layer can still tear down the task runtime instead of
|
||||
// incorrectly reviving a task whose database row is already gone.
|
||||
var finalizeErrs []error
|
||||
if trafficStage != nil {
|
||||
if err := trafficStage.Commit(); err != nil {
|
||||
finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize traffic deletion: %w", err))
|
||||
}
|
||||
}
|
||||
if fileStage != nil {
|
||||
if err := fileStage.commit(); err != nil {
|
||||
finalizeErrs = append(finalizeErrs, fmt.Errorf("finalize task file deletion: %w", err))
|
||||
}
|
||||
}
|
||||
m.forgetTask(id, n)
|
||||
if err := errors.Join(finalizeErrs...); err != nil {
|
||||
return result, &taskDeleteCommittedError{err: err}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// taskDeleteCommittedError means PostgreSQL deletion succeeded but purging one
|
||||
// of the recoverable staging directories failed. The task must stay deleted.
|
||||
type taskDeleteCommittedError struct{ err error }
|
||||
|
||||
func (e *taskDeleteCommittedError) Error() string {
|
||||
return "task deletion committed; external cleanup incomplete: " + e.err.Error()
|
||||
}
|
||||
|
||||
func (e *taskDeleteCommittedError) Unwrap() error { return e.err }
|
||||
|
||||
func rollbackTaskDelete(cause error, trafficStage *traffic.HostDeleteStage, fileStage *taskFileDeleteStage) error {
|
||||
errs := []error{cause}
|
||||
// Reverse the preparation order. Both restorations are attempted even if the
|
||||
// first one fails, and errors.Join preserves the original PostgreSQL error.
|
||||
if trafficStage != nil {
|
||||
if err := trafficStage.Rollback(); err != nil {
|
||||
errs = append(errs, fmt.Errorf("restore traffic after task delete failure: %w", err))
|
||||
}
|
||||
}
|
||||
if fileStage != nil {
|
||||
if err := fileStage.rollback(); err != nil {
|
||||
errs = append(errs, fmt.Errorf("restore task files after task delete failure: %w", err))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func (m *Manager) forgetTask(id string, numericID int64) {
|
||||
m.mu.Lock()
|
||||
delete(m.tasks, id)
|
||||
for _, task := range m.tasks {
|
||||
kept := task.SourceTaskIDs[:0]
|
||||
for _, sourceID := range task.SourceTaskIDs {
|
||||
if sourceID != numericID {
|
||||
kept = append(kept, sourceID)
|
||||
}
|
||||
}
|
||||
task.SourceTaskIDs = kept
|
||||
}
|
||||
if m.active == id {
|
||||
m.active = ""
|
||||
for _, t := range m.tasks {
|
||||
@@ -725,7 +940,128 @@ func (m *Manager) DeleteTask(id string) error {
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
type stagedTaskPath struct {
|
||||
source string
|
||||
staged string
|
||||
}
|
||||
|
||||
type taskFileDeleteStage struct {
|
||||
stageDir string
|
||||
moves []stagedTaskPath
|
||||
deleted bool
|
||||
done bool
|
||||
}
|
||||
|
||||
// stageTaskFiles atomically renames the task workspace and owned transcripts to
|
||||
// a same-filesystem staging directory. The trailing dash in the transcript
|
||||
// prefix is significant: exploration 12 must not match exploration 123.
|
||||
func stageTaskFiles(dataDir, taskID string, explorationID int64) (*taskFileDeleteStage, error) {
|
||||
stage := &taskFileDeleteStage{}
|
||||
var targets []string
|
||||
taskDir := filepath.Join(dataDir, "tasks", taskID)
|
||||
if _, err := os.Lstat(taskDir); err == nil {
|
||||
targets = append(targets, taskDir)
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
transcriptDir := filepath.Join(dataDir, "transcripts")
|
||||
entries, err := os.ReadDir(transcriptDir)
|
||||
if err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
return nil, err
|
||||
}
|
||||
entries = nil
|
||||
}
|
||||
prefix := fmt.Sprintf("exp%d-", explorationID)
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if !strings.HasPrefix(name, prefix) || (!entry.IsDir() && !strings.HasSuffix(name, ".jsonl")) {
|
||||
continue
|
||||
}
|
||||
targets = append(targets, filepath.Join(transcriptDir, name))
|
||||
}
|
||||
if len(targets) == 0 {
|
||||
stage.done = true
|
||||
return stage, nil
|
||||
}
|
||||
|
||||
parent := filepath.Join(dataDir, ".delete-staging")
|
||||
if err := os.MkdirAll(parent, 0o700); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
stage.stageDir, err = os.MkdirTemp(parent, "task-"+taskID+"-")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, source := range targets {
|
||||
staged := filepath.Join(stage.stageDir, fmt.Sprintf("%d-%s", len(stage.moves), filepath.Base(source)))
|
||||
if err := os.Rename(source, staged); err != nil {
|
||||
cause := fmt.Errorf("stage task file %s: %w", source, err)
|
||||
if restoreErr := stage.rollback(); restoreErr != nil {
|
||||
return nil, errors.Join(cause, fmt.Errorf("restore partially staged task files: %w", restoreErr))
|
||||
}
|
||||
return nil, cause
|
||||
}
|
||||
stage.moves = append(stage.moves, stagedTaskPath{source: source, staged: staged})
|
||||
}
|
||||
stage.deleted = true
|
||||
return stage, nil
|
||||
}
|
||||
|
||||
func (s *taskFileDeleteStage) commit() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
err := os.RemoveAll(s.stageDir)
|
||||
s.done = true
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *taskFileDeleteStage) rollback() error {
|
||||
if s == nil || s.done {
|
||||
return nil
|
||||
}
|
||||
var errs []error
|
||||
for i := len(s.moves) - 1; i >= 0; i-- {
|
||||
move := s.moves[i]
|
||||
if _, err := os.Lstat(move.source); err == nil {
|
||||
errs = append(errs, fmt.Errorf("restore destination already exists: %s", move.source))
|
||||
continue
|
||||
} else if !os.IsNotExist(err) {
|
||||
errs = append(errs, fmt.Errorf("inspect restore destination %s: %w", move.source, err))
|
||||
continue
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(move.source), 0o755); err != nil {
|
||||
errs = append(errs, fmt.Errorf("create restore parent for %s: %w", move.source, err))
|
||||
continue
|
||||
}
|
||||
if err := os.Rename(move.staged, move.source); err != nil {
|
||||
errs = append(errs, fmt.Errorf("restore %s: %w", move.source, err))
|
||||
}
|
||||
}
|
||||
if len(errs) == 0 && s.stageDir != "" {
|
||||
if err := os.RemoveAll(s.stageDir); err != nil {
|
||||
errs = append(errs, fmt.Errorf("remove task file stage: %w", err))
|
||||
}
|
||||
}
|
||||
s.done = true
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
// deleteTaskFiles retains the standalone helper contract used by focused tests.
|
||||
func deleteTaskFiles(dataDir, taskID string, explorationID int64) (bool, error) {
|
||||
stage, err := stageTaskFiles(dataDir, taskID, explorationID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
deleted := stage.deleted
|
||||
if err := stage.commit(); err != nil {
|
||||
return deleted, err
|
||||
}
|
||||
return deleted, nil
|
||||
}
|
||||
|
||||
func (m *Manager) Task(id string) (*Task, bool) {
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
pgdb "github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/artex/traffic"
|
||||
)
|
||||
|
||||
func TestSeedAssociatesTargetAssetWithTask(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("seed ownership", "seed ownership", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true})
|
||||
})
|
||||
|
||||
host := fmt.Sprintf("seed-%s.example.test", task.ID)
|
||||
s := &Server{m: m}
|
||||
s.seed(task, "https://"+host)
|
||||
taskID := mustTaskID(t, task.ID)
|
||||
assets, err := m.Assets().QueryByTask(taskID, "", 20, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, asset := range assets {
|
||||
if asset.Domain == host || asset.URL == "https://"+host {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("seed target was not associated with task %s: %+v", task.ID, assets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteTaskFilesRemovesOnlyOwnedWorkspaceAndTranscripts(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
|
||||
removed := []struct {
|
||||
path string
|
||||
isDir bool
|
||||
}{
|
||||
{path: filepath.Join("tasks", "77"), isDir: true},
|
||||
{path: filepath.Join("transcripts", "exp12-main.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp12-planner.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp12-worker-i9.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp12-main"), isDir: true},
|
||||
{path: filepath.Join("transcripts", "exp12-worker-i9"), isDir: true},
|
||||
}
|
||||
kept := []struct {
|
||||
path string
|
||||
isDir bool
|
||||
}{
|
||||
{path: filepath.Join("tasks", "78"), isDir: true},
|
||||
{path: filepath.Join("transcripts", "exp123-main.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp123-main"), isDir: true},
|
||||
{path: filepath.Join("transcripts", "exp1-main.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp12main.jsonl")},
|
||||
{path: filepath.Join("transcripts", "exp12-main.jsonl.bak")},
|
||||
{path: filepath.Join("transcripts", "exp12-main.txt")},
|
||||
{path: filepath.Join("transcripts", "conv-12.jsonl")},
|
||||
{path: filepath.Join("transcripts", "unrelated-session"), isDir: true},
|
||||
}
|
||||
|
||||
for _, fixture := range append(removed, kept...) {
|
||||
createDeleteFixture(t, dataDir, fixture.path, fixture.isDir)
|
||||
}
|
||||
|
||||
deleted, err := deleteTaskFiles(dataDir, "77", 12)
|
||||
if err != nil {
|
||||
t.Fatalf("delete task files: %v", err)
|
||||
}
|
||||
if !deleted {
|
||||
t.Fatal("expected deletion to be reported")
|
||||
}
|
||||
for _, fixture := range removed {
|
||||
assertPathMissing(t, filepath.Join(dataDir, fixture.path))
|
||||
}
|
||||
for _, fixture := range kept {
|
||||
assertPathExists(t, filepath.Join(dataDir, fixture.path))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteTaskFilesReportsTranscriptOnlyDeletion(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
path := filepath.Join("transcripts", "exp44-main.jsonl")
|
||||
createDeleteFixture(t, dataDir, path, false)
|
||||
|
||||
deleted, err := deleteTaskFiles(dataDir, "missing-task", 44)
|
||||
if err != nil {
|
||||
t.Fatalf("delete task files: %v", err)
|
||||
}
|
||||
if !deleted {
|
||||
t.Fatal("expected transcript deletion to be reported")
|
||||
}
|
||||
assertPathMissing(t, filepath.Join(dataDir, path))
|
||||
}
|
||||
|
||||
func TestDeleteTaskFilesMissingTargetsIsIdempotent(t *testing.T) {
|
||||
deleted, err := deleteTaskFiles(t.TempDir(), "77", 12)
|
||||
if err != nil {
|
||||
t.Fatalf("delete missing task files: %v", err)
|
||||
}
|
||||
if deleted {
|
||||
t.Fatal("missing targets must not be reported as deleted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStageTaskFilesRollbackRestoresWorkspaceAndTranscripts(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
paths := []string{
|
||||
filepath.Join("tasks", "77", "notes.txt"),
|
||||
filepath.Join("transcripts", "exp12-main.jsonl"),
|
||||
filepath.Join("transcripts", "exp12-worker-i9", "sidechain.jsonl"),
|
||||
}
|
||||
for _, path := range paths {
|
||||
createDeleteFixture(t, dataDir, path, false)
|
||||
}
|
||||
|
||||
stage, err := stageTaskFiles(dataDir, "77", 12)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !stage.deleted {
|
||||
t.Fatal("expected files to be staged")
|
||||
}
|
||||
for _, path := range paths {
|
||||
assertPathMissing(t, filepath.Join(dataDir, path))
|
||||
}
|
||||
if err := stage.rollback(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, path := range paths {
|
||||
assertPathExists(t, filepath.Join(dataDir, path))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStageTaskFilesRollbackReportsRestoreFailure(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
workspace := filepath.Join(dataDir, "tasks", "77")
|
||||
createDeleteFixture(t, dataDir, filepath.Join("tasks", "77"), true)
|
||||
|
||||
stage, err := stageTaskFiles(dataDir, "77", 12)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(workspace, []byte("conflict"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cause := errors.New("forced PostgreSQL delete failure")
|
||||
err = rollbackTaskDelete(cause, nil, stage)
|
||||
if err == nil || !errors.Is(err, cause) || !strings.Contains(err.Error(), "destination already exists") {
|
||||
t.Fatalf("rollback err=%v, want original and restore errors", err)
|
||||
}
|
||||
if _, statErr := os.Stat(stage.stageDir); statErr != nil {
|
||||
t.Fatalf("staging was removed after failed restore: %v", statErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerDeleteTaskRestoresFilesAndTrafficWhenPostgresDeleteFails(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
m, err := NewManager(dataDir, "")
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v) - skipping", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = m.Close() })
|
||||
|
||||
task, err := m.CreateTask("delete rollback", "delete rollback", nil, 0, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
taskID := mustTaskID(t, task.ID)
|
||||
triggerName := fmt.Sprintf("test_fail_task_delete_%d", taskID)
|
||||
functionName := triggerName + "_fn"
|
||||
cleanupDB := func() {
|
||||
_, _ = m.pg.Exec(fmt.Sprintf(`DROP TRIGGER IF EXISTS %s ON tasks`, triggerName))
|
||||
_, _ = m.pg.Exec(fmt.Sprintf(`DROP FUNCTION IF EXISTS %s()`, functionName))
|
||||
_, _ = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteAssets: true, DeleteTraffic: true, DeleteFiles: true})
|
||||
}
|
||||
t.Cleanup(cleanupDB)
|
||||
|
||||
host := fmt.Sprintf("delete-rollback-%d.example.test", taskID)
|
||||
if _, err := m.Assets().UpsertHTTPService(pgdb.UpsertHTTPServiceReq{URL: "https://" + host, TaskID: taskID}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
createDeleteFixture(t, dataDir, filepath.Join("tasks", task.ID), true)
|
||||
transcriptPath := filepath.Join(dataDir, "transcripts", fmt.Sprintf("exp%d-main.jsonl", task.ExpID))
|
||||
createDeleteFixture(t, dataDir, filepath.Join("transcripts", filepath.Base(transcriptPath)), false)
|
||||
|
||||
tr, err := traffic.Open(filepath.Join(dataDir, "traffic"), "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m.traffic = tr
|
||||
trafficTree := filepath.Join(dataDir, "traffic", host, "GET", "1-0001")
|
||||
if err := os.MkdirAll(trafficTree, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
trafficMarker := filepath.Join(trafficTree, "request.http")
|
||||
if err := os.WriteFile(trafficMarker, []byte("original traffic"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := tr.DB().Exec(`INSERT INTO exchanges(id,ts,host,method,url_template,url,status,content_type,req_len,resp_len,path)
|
||||
VALUES(?,?,?,?,?,?,?,?,?,?,?)`, "1-0001", 1, host, "GET", "/", "https://"+host+"/", 200, "text/plain", 0, 0, host+"/GET/1-0001"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if _, err := m.pg.Exec(fmt.Sprintf(`CREATE FUNCTION %s() RETURNS trigger LANGUAGE plpgsql AS $body$
|
||||
BEGIN RAISE EXCEPTION 'forced task delete failure'; END $body$`, functionName)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := m.pg.Exec(fmt.Sprintf(`CREATE TRIGGER %s BEFORE DELETE ON tasks
|
||||
FOR EACH ROW WHEN (OLD.id = %d) EXECUTE FUNCTION %s()`, triggerName, taskID, functionName)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, err = m.DeleteTask(task.ID, DeleteTaskOptions{DeleteTraffic: true, DeleteFiles: true})
|
||||
if err == nil || !strings.Contains(err.Error(), "forced task delete failure") {
|
||||
t.Fatalf("delete err=%v, want injected PostgreSQL failure", err)
|
||||
}
|
||||
if got, err := m.pg.GetTask(taskID); err != nil || got == nil {
|
||||
t.Fatalf("task row was lost after failed delete: task=%+v err=%v", got, err)
|
||||
}
|
||||
assertPathExists(t, filepath.Join(dataDir, "tasks", task.ID))
|
||||
assertPathExists(t, transcriptPath)
|
||||
if got, err := os.ReadFile(trafficMarker); err != nil || string(got) != "original traffic" {
|
||||
t.Fatalf("traffic tree was not restored: content=%q err=%v", got, err)
|
||||
}
|
||||
var trafficRows int
|
||||
if err := tr.DB().QueryRow(`SELECT COUNT(*) FROM exchanges WHERE host=?`, host).Scan(&trafficRows); err != nil || trafficRows != 1 {
|
||||
t.Fatalf("traffic index was not restored: rows=%d err=%v", trafficRows, err)
|
||||
}
|
||||
}
|
||||
|
||||
func createDeleteFixture(t *testing.T, root, relative string, isDir bool) {
|
||||
t.Helper()
|
||||
path := filepath.Join(root, relative)
|
||||
if isDir {
|
||||
if err := os.MkdirAll(filepath.Join(path, "subagents"), 0o755); err != nil {
|
||||
t.Fatalf("create fixture directory %s: %v", path, err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(path, "subagents", "agent-test.jsonl"), []byte("test"), 0o644); err != nil {
|
||||
t.Fatalf("create fixture sidechain %s: %v", path, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatalf("create fixture parent %s: %v", path, err)
|
||||
}
|
||||
if err := os.WriteFile(path, []byte("test"), 0o644); err != nil {
|
||||
t.Fatalf("create fixture file %s: %v", path, err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertPathMissing(t *testing.T, path string) {
|
||||
t.Helper()
|
||||
if _, err := os.Lstat(path); !os.IsNotExist(err) {
|
||||
t.Fatalf("expected %s to be removed, stat err=%v", path, err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertPathExists(t *testing.T, path string) {
|
||||
t.Helper()
|
||||
if _, err := os.Lstat(path); err != nil {
|
||||
t.Fatalf("expected %s to remain: %v", path, err)
|
||||
}
|
||||
}
|
||||
@@ -199,12 +199,13 @@ func (s *Server) toolListTasks() actool.CoreTool {
|
||||
if t.ParentRef != "" {
|
||||
row["parent_ref"] = t.ParentRef
|
||||
}
|
||||
if t.LLMProfileID == nil {
|
||||
llmState := t.llmStateSnapshot()
|
||||
if llmState.ProfileID == nil {
|
||||
row["llm_profile"] = "(激活配置)"
|
||||
} else if n, ok := profName[*t.LLMProfileID]; ok {
|
||||
} else if n, ok := profName[*llmState.ProfileID]; ok {
|
||||
row["llm_profile"] = n
|
||||
} else {
|
||||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *t.LLMProfileID)
|
||||
row["llm_profile"] = fmt.Sprintf("#%d(已删除)", *llmState.ProfileID)
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
@@ -305,6 +306,7 @@ func (s *Server) toolPauseTask() actool.CoreTool {
|
||||
return actool.Errorf("task 不存在: " + a.TaskID), nil
|
||||
}
|
||||
s.engine.Pause(a.TaskID, agent.AbortPausedByOrchestrator)
|
||||
s.cancelTaskChat(a.TaskID, agent.AbortChatPausedWithTask)
|
||||
return actool.Text("task paused: " + a.TaskID), nil
|
||||
})
|
||||
}
|
||||
@@ -419,7 +421,7 @@ func (s *Server) deriveTaskStatus(t *Task) string {
|
||||
return t.Status
|
||||
case t.Paused || s.engine.IsPaused(t.ID):
|
||||
return "paused"
|
||||
case s.engine.Ready() && s.engine.Started(t.ID):
|
||||
case s.engine.ReadyFor(t) && s.engine.Started(t.ID):
|
||||
return "running"
|
||||
}
|
||||
return "created"
|
||||
|
||||
+658
-116
File diff suppressed because it is too large.
Load diff
+129
-53
@@ -5,6 +5,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
@@ -66,61 +67,97 @@ func decode(r *http.Request, v any) error { return json.NewDecoder(r.Body).Decod
|
||||
|
||||
// ---------- tasks (delete) ----------
|
||||
|
||||
// pgDeleteTask removes a task and, per opt-in flags in the (optional) JSON body,
|
||||
// its associated data: assets (task-owned), findings, recorded traffic (by the
|
||||
// task's hosts) and the per-task working directory. Empty body → all flags false →
|
||||
// only the task + its exploration subgraph are removed (legacy behavior).
|
||||
const taskDeleteDrainTimeout = 10 * time.Second
|
||||
|
||||
func canonicalTaskID(raw string) (string, bool) {
|
||||
n, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || n <= 0 {
|
||||
return "", false
|
||||
}
|
||||
return strconv.FormatInt(n, 10), true
|
||||
}
|
||||
|
||||
func (s *Server) pgDeleteTask(w http.ResponseWriter, r *http.Request) {
|
||||
id := r.PathValue("id")
|
||||
var opts struct {
|
||||
Assets bool `json:"assets"`
|
||||
Findings bool `json:"findings"`
|
||||
Traffic bool `json:"traffic"`
|
||||
Files bool `json:"files"`
|
||||
id, ok := canonicalTaskID(r.PathValue("id"))
|
||||
if !ok {
|
||||
writeErr(w, http.StatusBadRequest, "任务 id 无效")
|
||||
return
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&opts) // optional body; ignore EOF/parse errors
|
||||
|
||||
result := map[string]any{"deleted": id}
|
||||
nid, _ := strconv.ParseInt(id, 10, 64)
|
||||
as := s.m.Assets()
|
||||
|
||||
// Traffic first: it needs the task's hosts, which come from assets that a later
|
||||
// asset-delete would remove. Traffic is host-keyed (no per-task tag), so this
|
||||
// purges ALL recorded traffic for those hosts, not just this task's.
|
||||
if opts.Traffic && nid > 0 && as != nil {
|
||||
if tr := s.m.Traffic(); tr != nil {
|
||||
if hosts, err := as.HostsByTask(nid); err == nil && len(hosts) > 0 {
|
||||
if n, err := tr.DeleteHostsExact(hosts); err == nil {
|
||||
result["traffic_deleted"] = n
|
||||
}
|
||||
}
|
||||
var opts DeleteTaskOptions
|
||||
if err := decode(r, &opts); err != nil && err != io.EOF {
|
||||
writeErr(w, 400, "invalid JSON: "+err.Error())
|
||||
return
|
||||
}
|
||||
if !s.engine.BeginDelete(id) {
|
||||
writeErr(w, http.StatusConflict, "任务正在删除")
|
||||
return
|
||||
}
|
||||
deleted := false
|
||||
defer func() {
|
||||
if !deleted {
|
||||
s.engine.AbortDelete(id)
|
||||
}
|
||||
}
|
||||
if opts.Findings && nid > 0 {
|
||||
if n, err := s.m.pg.DeleteFindingsByTask(nid); err == nil {
|
||||
result["findings_deleted"] = n
|
||||
}
|
||||
}
|
||||
if opts.Assets && nid > 0 && as != nil {
|
||||
if n, err := as.DeleteByTaskID(nid); err == nil {
|
||||
result["assets_deleted"] = n
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
s.engine.Pause(id, agent.AbortTaskDeleted) // cancel running loops before removing
|
||||
if err := s.m.DeleteTask(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
// Main-agent runs use a separate context from planner/workers. Cancel it, then
|
||||
// wait until both execution domains have returned before removing transcripts.
|
||||
s.cancelTaskChat(id, agent.AbortTaskDeleted)
|
||||
drainCtx, cancelDrain := context.WithTimeout(r.Context(), taskDeleteDrainTimeout)
|
||||
defer cancelDrain()
|
||||
if err := s.waitTaskQuiescent(drainCtx, id); err != nil {
|
||||
writeErr(w, http.StatusConflict, "任务仍有运行中的 Agent,删除已取消")
|
||||
return
|
||||
}
|
||||
|
||||
// Files last: the per-task working dir <workDir>/tasks/<id> (worker artifacts +
|
||||
// uploads). Confined to the work dir; id is numeric so no traversal risk.
|
||||
if opts.Files && nid > 0 {
|
||||
if err := os.RemoveAll(filepath.Join(s.m.dir, "tasks", id)); err == nil {
|
||||
result["files_deleted"] = true
|
||||
result, err := s.m.DeleteTask(id, opts)
|
||||
if err != nil {
|
||||
var committed *taskDeleteCommittedError
|
||||
if errors.As(err, &committed) {
|
||||
// PostgreSQL is already gone. Complete runtime teardown and return the
|
||||
// auditable counts together with the post-commit cleanup warning.
|
||||
s.engine.StopTask(id)
|
||||
s.taskAgentMu.Lock()
|
||||
delete(s.taskAgents, id)
|
||||
s.taskAgentMu.Unlock()
|
||||
deleted = true
|
||||
writeCommittedTaskDelete(w, result, err)
|
||||
return
|
||||
}
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
// Manager has removed the task from the registry, so no new API operation can
|
||||
// resolve it. Now stop and join every task-owned Engine goroutine and clear its
|
||||
// lifecycle maps before releasing the delete barrier.
|
||||
s.engine.StopTask(id)
|
||||
s.taskAgentMu.Lock()
|
||||
delete(s.taskAgents, id)
|
||||
s.taskAgentMu.Unlock()
|
||||
deleted = true
|
||||
writeJSON(w, 200, result)
|
||||
}
|
||||
|
||||
func writeCommittedTaskDelete(w http.ResponseWriter, result DeleteTaskResult, err error) {
|
||||
result.CleanupWarning = err.Error()
|
||||
writeJSON(w, http.StatusOK, result)
|
||||
}
|
||||
|
||||
func (s *Server) waitTaskQuiescent(ctx context.Context, taskID string) error {
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
s.chatMu.Lock()
|
||||
chatBusy := s.chatBusy[taskID]
|
||||
s.chatMu.Unlock()
|
||||
if s.engine.inflightCount(taskID) == 0 && !chatBusy {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
writeJSON(w, 200, result)
|
||||
}
|
||||
|
||||
// ---------- agents ----------
|
||||
@@ -352,6 +389,9 @@ func (s *Server) pgSaveAgentConfig(w http.ResponseWriter, r *http.Request) {
|
||||
if profileChanged {
|
||||
s.invalidateProfileAgents()
|
||||
}
|
||||
// Task bundles capture operational Agent settings (turn/time budgets, tools,
|
||||
// web search). Rebuild them even when no global profile is configured.
|
||||
s.invalidateTaskAgents()
|
||||
// rebuild the live agents so the new max_turns/run_seconds/binding apply without a restart.
|
||||
s.cfgMu.Lock()
|
||||
cfg, on := s.llmCfg, s.llmOn
|
||||
@@ -396,8 +436,8 @@ func (s *Server) pgGetAgent(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"agent": agentDTO(a), "prompt": cur, "variables": vars, "versions": vers,
|
||||
"visibility": map[string]any{"mcp": mcp, "skill": sk},
|
||||
"llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选
|
||||
"visibility": map[string]any{"mcp": mcp, "skill": sk},
|
||||
"llm_profiles": llmProfiles, // 可绑定的 LLM 配置候选
|
||||
|
||||
"wrapup_prompt": a.WrapupPrompt, // 已保存的收尾提示词(空=用内置默认)
|
||||
"wrapup_default": agent.WrapupDefault(a.Key), // 内置默认(供占位/恢复默认)
|
||||
@@ -1607,7 +1647,9 @@ func (s *Server) pgSaveProfile(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
// Editing a profile rebuilds any task pinned to it on its next round.
|
||||
// Editing a profile rebuilds any task pinned to it on its next round. Reapply
|
||||
// the active profile too: the global fallback and explicit task chains must
|
||||
// repopulate the same provider-cache entry and therefore share one limiter.
|
||||
s.invalidateProfileAgents()
|
||||
// Hot-apply. Not just when the ACTIVE profile changed: with failover on, any
|
||||
// profile's key/model/priority/pool_exclude edit reshapes the chain, so the
|
||||
@@ -1622,16 +1664,50 @@ func (s *Server) pgDeleteProfile(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
id, _ := pathInt(r, "id")
|
||||
if err := pg.DeleteProfile(id); err != nil {
|
||||
writeErr(w, 500, err.Error())
|
||||
if err := pg.DeleteProfileContext(r.Context(), id); err != nil {
|
||||
switch {
|
||||
case errors.Is(err, db.ErrActiveLLMProfileDelete):
|
||||
writeErr(w, 409, "当前激活的 LLM 配置不能删除,请先激活其他配置")
|
||||
case errors.Is(err, db.ErrLLMProfileReferencesChanged):
|
||||
writeErr(w, 409, "LLM 配置正在被任务或会话修改,请重试")
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
writeErr(w, 409, "等待 LLM 配置引用释放超时,请重试")
|
||||
case errors.Is(err, db.ErrLLMProfileNotFound):
|
||||
writeErr(w, 404, err.Error())
|
||||
default:
|
||||
writeErr(w, 500, err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
s.invalidateProfileAgents() // drop cached agents for the removed profile
|
||||
s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away)
|
||||
s.reapplyActiveProfile() // the deleted profile may have been in the failover chain
|
||||
s.llmHealth.Reset(id) // its breaker state is meaningless now (row is FK-cascaded away)
|
||||
s.restoreTasksAfterProfileDelete(pg)
|
||||
// Cache invalidation also removed the active profile's shared provider entry.
|
||||
// Reapply after task-state sync so the wake-up observes the post-delete chain.
|
||||
s.reapplyActiveProfile()
|
||||
writeJSON(w, 200, map[string]any{"deleted": id})
|
||||
}
|
||||
|
||||
func (s *Server) restoreTasksAfterProfileDelete(pg *db.DB) {
|
||||
for _, task := range s.m.List() {
|
||||
if taskID, err := strconv.ParseInt(task.ID, 10, 64); err == nil {
|
||||
if pt, err := pg.GetTask(taskID); err == nil && pt != nil {
|
||||
s.syncTaskLLMState(pt)
|
||||
// Deleting the active entry may advance the task to a ready successor,
|
||||
// or clear the explicit chain and restore its Agent/global fallback.
|
||||
// Resume only intents that were blocked by the exhausted chain; a
|
||||
// paused task remains paused and terminal tasks remain immutable.
|
||||
if !isTerminalStatus(task.Status) && s.taskRuntimeAvailable(task, "planner", "worker") {
|
||||
if _, reopenErr := task.Store.ReopenIntentsByBlockedReason(db.IntentBlockedLLMQuota); reopenErr != nil {
|
||||
log.Printf("[llm-profile] task %s reopen quota-blocked intents after profile delete: %v", task.ID, reopenErr)
|
||||
}
|
||||
}
|
||||
task.Notify()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) pgActivateProfile(w http.ResponseWriter, r *http.Request) {
|
||||
pg := s.pg(w)
|
||||
if pg == nil {
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Autumn-27/artex/agent"
|
||||
)
|
||||
|
||||
func TestEngineDeleteBarrierRejectsNewTaskOperations(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
const taskID = "42"
|
||||
|
||||
if !e.beginTaskOperation(taskID) {
|
||||
t.Fatal("first task operation should be admitted")
|
||||
}
|
||||
if !e.BeginDelete(taskID) {
|
||||
t.Fatal("first delete should install the barrier")
|
||||
}
|
||||
if !e.IsDeleting(taskID) {
|
||||
t.Fatal("delete barrier should be visible")
|
||||
}
|
||||
if !e.IsPaused(taskID) {
|
||||
t.Fatal("delete should temporarily pause task execution")
|
||||
}
|
||||
if e.beginTaskOperation(taskID) {
|
||||
t.Fatal("operation started after BeginDelete returned")
|
||||
}
|
||||
if got := e.inflightCount(taskID); got != 1 {
|
||||
t.Fatalf("existing operation must remain drainable, inflight=%d", got)
|
||||
}
|
||||
e.decInflight(taskID)
|
||||
|
||||
if e.BeginDelete(taskID) {
|
||||
t.Fatal("duplicate delete must not acquire the barrier")
|
||||
}
|
||||
e.AbortDelete(taskID)
|
||||
if e.IsPaused(taskID) {
|
||||
t.Fatal("aborted delete should restore a previously running task")
|
||||
}
|
||||
if !e.beginTaskOperation(taskID) {
|
||||
t.Fatal("aborting deletion should reopen operation admission")
|
||||
}
|
||||
e.decInflight(taskID)
|
||||
}
|
||||
|
||||
func TestAbortDeletePreservesExistingPause(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
e.Pause("42", agent.AbortPausedByUser)
|
||||
if !e.BeginDelete("42") {
|
||||
t.Fatal("delete barrier should be installed")
|
||||
}
|
||||
e.AbortDelete("42")
|
||||
if !e.IsPaused("42") {
|
||||
t.Fatal("aborted delete must preserve a pre-existing user pause")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopTaskCancelsRuntimeAndClearsLifecycleState(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
const taskID = "42"
|
||||
e.deleteMu.RLock()
|
||||
rt := e.registerTaskRoutines(context.Background(), taskID, 1)
|
||||
e.started.Store(taskID, true)
|
||||
e.lastAct.Store(taskID, int64(1))
|
||||
e.paused.Store(taskID, true)
|
||||
e.plannerRound.Store(taskID, 3)
|
||||
e.coordStarted.Store(taskID, true)
|
||||
e.deleteMu.RUnlock()
|
||||
|
||||
done := make(chan struct{})
|
||||
runTaskRoutine(rt, func(ctx context.Context) {
|
||||
<-ctx.Done()
|
||||
close(done)
|
||||
})
|
||||
e.StopTask(taskID)
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
default:
|
||||
t.Fatal("task runtime was not cancelled before StopTask returned")
|
||||
}
|
||||
if e.Started(taskID) || e.IsPaused(taskID) || e.IsDeleting(taskID) {
|
||||
t.Fatalf("task lifecycle maps were not cleared: started=%v paused=%v deleting=%v",
|
||||
e.Started(taskID), e.IsPaused(taskID), e.IsDeleting(taskID))
|
||||
}
|
||||
if _, ok := e.lastAct.Load(taskID); ok {
|
||||
t.Fatal("last activity state was retained")
|
||||
}
|
||||
if _, ok := e.plannerRound.Load(taskID); ok {
|
||||
t.Fatal("planner round state was retained")
|
||||
}
|
||||
if _, ok := e.coordStarted.Load(taskID); ok {
|
||||
t.Fatal("deadline coordinator state was retained")
|
||||
}
|
||||
e.runtimeMu.Lock()
|
||||
_, retained := e.runtimes[taskID]
|
||||
e.runtimeMu.Unlock()
|
||||
if retained {
|
||||
t.Fatal("task runtime retained after StopTask")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanonicalTaskID(t *testing.T) {
|
||||
for raw, want := range map[string]string{"42": "42", "00042": "42", "+42": "42"} {
|
||||
if got, ok := canonicalTaskID(raw); !ok || got != want {
|
||||
t.Fatalf("canonicalTaskID(%q)=(%q,%v), want (%q,true)", raw, got, ok, want)
|
||||
}
|
||||
}
|
||||
for _, raw := range []string{"", "0", "-1", "abc", "42x"} {
|
||||
if got, ok := canonicalTaskID(raw); ok {
|
||||
t.Fatalf("canonicalTaskID(%q)=(%q,true), want invalid", raw, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelChatThenWaitTaskQuiescent(t *testing.T) {
|
||||
e := NewEngine(nil)
|
||||
chatCtx, chatCancel := context.WithCancelCause(context.Background())
|
||||
s := &Server{
|
||||
engine: e,
|
||||
chatBusy: map[string]bool{"7": true},
|
||||
chatCancel: map[string]context.CancelCauseFunc{"7": chatCancel},
|
||||
}
|
||||
if !e.beginTaskOperation("7") {
|
||||
t.Fatal("task operation should be admitted")
|
||||
}
|
||||
|
||||
cancelObserved := make(chan struct{})
|
||||
releaseChat := make(chan struct{})
|
||||
go func() {
|
||||
<-chatCtx.Done()
|
||||
close(cancelObserved)
|
||||
<-releaseChat
|
||||
s.finishTaskChat("7", chatCancel)
|
||||
}()
|
||||
if !s.cancelTaskChat("7", agent.AbortTaskDeleted) {
|
||||
t.Fatal("expected active chat to be cancelled")
|
||||
}
|
||||
<-cancelObserved
|
||||
|
||||
waitDone := make(chan error, 1)
|
||||
waitCtx, waitCancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer waitCancel()
|
||||
go func() { waitDone <- s.waitTaskQuiescent(waitCtx, "7") }()
|
||||
|
||||
select {
|
||||
case err := <-waitDone:
|
||||
t.Fatalf("wait returned while chat and engine operation were active: %v", err)
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
close(releaseChat)
|
||||
select {
|
||||
case err := <-waitDone:
|
||||
t.Fatalf("wait returned while engine operation was active: %v", err)
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
e.decInflight("7")
|
||||
if err := <-waitDone; err != nil {
|
||||
t.Fatalf("wait for quiescence: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWaitTaskQuiescentHonorsContext(t *testing.T) {
|
||||
s := &Server{
|
||||
engine: NewEngine(nil),
|
||||
chatBusy: map[string]bool{"9": true},
|
||||
chatCancel: map[string]context.CancelCauseFunc{},
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if err := s.waitTaskQuiescent(ctx, "9"); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("expected deadline exceeded, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskChatUploadRejectsDeleteBarrierBeforeCreatingDirectory(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
m := &Manager{dir: dataDir, tasks: map[string]*Task{"7": {ID: "7"}}}
|
||||
s := &Server{m: m, engine: NewEngine(m)}
|
||||
if !s.engine.BeginDelete("7") {
|
||||
t.Fatal("delete barrier should be installed")
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/chat/upload?scope=task&id=7", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
s.chatUpload(rec, req)
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("expected 409 while deleting, got %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dataDir, "tasks", "7")); !os.IsNotExist(err) {
|
||||
t.Fatalf("upload must not recreate task directory, stat err=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCommittedTaskDeleteReturnsSuccessWithCleanupWarning(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
result := DeleteTaskResult{Deleted: "7", FilesDeleted: true}
|
||||
err := &taskDeleteCommittedError{err: errors.New("remove staged task files: permission denied")}
|
||||
|
||||
writeCommittedTaskDelete(rec, result, err)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s, want 200", rec.Code, rec.Body.String())
|
||||
}
|
||||
var got DeleteTaskResult
|
||||
if decodeErr := json.NewDecoder(rec.Body).Decode(&got); decodeErr != nil {
|
||||
t.Fatal(decodeErr)
|
||||
}
|
||||
if got.Deleted != "7" || !got.FilesDeleted || !strings.Contains(got.CleanupWarning, "cleanup incomplete") {
|
||||
t.Fatalf("unexpected response: %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,462 @@
|
||||
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)
|
||||
}
|
||||
|
||||
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 len(pt.LLMProfileIDs) == 0 {
|
||||
prov, _, ok := r.s.fallbackProviderForAgent(r.agentKey)
|
||||
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
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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
|
||||
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 {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
func (r *taskLLMRuntime) CompactionWindow() int {
|
||||
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.fallbackProviderForAgent(r.agentKey); 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
|
||||
}
|
||||
|
||||
// fallbackProviderForAgent resolves the pre-existing precedence used by tasks
|
||||
// without an explicit chain: the Agent's saved binding first, then the current
|
||||
// global provider (which may come from the default profile or environment).
|
||||
func (s *Server) fallbackProviderForAgent(agentKey string) (llm.Provider, agent.Config, bool) {
|
||||
if id := s.effectiveProfileForAgent(agentKey, nil); id != nil {
|
||||
if prov, cfg, ok := s.providerForProfile(*id); ok {
|
||||
return s.poolForBinding(*id, prov, cfg), cfg, true
|
||||
}
|
||||
}
|
||||
s.cfgMu.Lock()
|
||||
defer s.cfgMu.Unlock()
|
||||
if !s.llmOn || s.llmProv == nil {
|
||||
return nil, agent.Config{}, false
|
||||
}
|
||||
return s.llmProv, s.llmCfg, true
|
||||
}
|
||||
|
||||
func (s *Server) taskRuntimeAvailable(t *Task, agentKeys ...string) bool {
|
||||
if t == nil {
|
||||
return false
|
||||
}
|
||||
state := t.llmStateSnapshot()
|
||||
if len(state.ProfileIDs) > 0 {
|
||||
if state.ActiveID == nil {
|
||||
return false
|
||||
}
|
||||
_, _, ok := s.providerForProfile(*state.ActiveID)
|
||||
return ok
|
||||
}
|
||||
s.cfgMu.Lock()
|
||||
globallyReady := s.llmOn && s.llmProv != nil
|
||||
s.cfgMu.Unlock()
|
||||
if globallyReady {
|
||||
return true
|
||||
}
|
||||
for _, key := range agentKeys {
|
||||
if _, _, ok := s.fallbackProviderForAgent(key); !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return len(agentKeys) > 0
|
||||
}
|
||||
|
||||
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"))
|
||||
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"))
|
||||
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"))
|
||||
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.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,507 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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}, "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 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) {
|
||||
t.Parallel()
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
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}, "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)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user